From feb891c508e843c32fede66cfb7ea49e8631cff5 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Fri, 28 Feb 2025 14:44:13 -0800
Subject: [PATCH] [executorch][runtime] Add get_named_data_map to Program

Add to the program interface, to allow users to retrieve the NDM.

Differential Revision: [D70276106](https://our.internmc.facebook.com/intern/diff/D70276106/)

[ghstack-poisoned]
---
 runtime/executor/program.cpp           |  7 +++++++
 runtime/executor/program.h             |  6 ++++++
 runtime/executor/test/program_test.cpp | 12 ++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
index b303575593f..6fe7e8186a8 100644
--- a/runtime/executor/program.cpp
+++ b/runtime/executor/program.cpp
@@ -373,6 +373,13 @@ Result<const void*> Program::get_constant_buffer_data(
   }
 }
 
+Result<const NamedDataMap*> Program::get_named_data_map() const {
+  if (core_data_map_.has_value()) {
+    return &core_data_map_.value();
+  }
+  return Error::NotFound;
+}
+
 Result<const char*> Program::get_output_flattening_encoding(
     const char* method_name) const {
   auto plan = get_execution_plan(internal_program_, method_name);
diff --git a/runtime/executor/program.h b/runtime/executor/program.h
index ff04b630716..55506a1bcbe 100644
--- a/runtime/executor/program.h
+++ b/runtime/executor/program.h
@@ -107,6 +107,12 @@ class Program final {
   Result<const void*> get_constant_buffer_data(size_t buffer_idx, size_t nbytes)
       const;
 
+  /**
+   * Get the named data map from the program.
+   * @return The named data map.
+   */
+  Result<const NamedDataMap*> get_named_data_map() const;
+
   /**
    * Returns the number of methods in the program.
    */
diff --git a/runtime/executor/test/program_test.cpp b/runtime/executor/test/program_test.cpp
index 80f91f1af6a..d82bfb50778 100644
--- a/runtime/executor/test/program_test.cpp
+++ b/runtime/executor/test/program_test.cpp
@@ -371,6 +371,18 @@ TEST_F(ProgramTest, getMethods) {
   EXPECT_EQ(strcmp(res2.get(), "forward2"), 0);
 }
 
+TEST_F(ProgramTest, GetNamedDataMap_Fail) {
+  Result<Program> program =
+      Program::load(add_loader_.get(), kDefaultVerification);
+  ASSERT_EQ(program.error(), Error::Ok);
+
+  // Get the named data map. Expect to fail, as add.pte does not have any
+  // named data segments.
+  Result<const executorch::runtime::NamedDataMap*> named_data_map =
+      program->get_named_data_map();
+  EXPECT_EQ(named_data_map.error(), Error::NotFound);
+}
+
 // Test that the deprecated Load method (capital 'L') still works.
 TEST_F(ProgramTest, DEPRECATEDLoad) {
   // Parse the Program from the data.