Add new FlutterEngineAOTData argument to FlutterProjectArgs (#18146)

Added a new `FlutterEngineAOTData` argument to `FlutterProjectArgs`. Embedders can instantiate and destroy this object via the new `FlutterEngineCreateAOTData` and `FlutterEngineCollectAOTData` methods provided.

If an embedder provides more than one source of AOT data to `FlutterEngineInitialize` or `FlutterEngineRun` (e.g. snapshots as well as `FlutterEngineAOTData`), the engine will error out.

Resolves: https://github.com/flutter/flutter/issues/50778
diff --git a/shell/platform/embedder/BUILD.gn b/shell/platform/embedder/BUILD.gn
index 09547a3..f4b5be5 100644
--- a/shell/platform/embedder/BUILD.gn
+++ b/shell/platform/embedder/BUILD.gn
@@ -71,6 +71,7 @@
       "//flutter/shell/common",
       "//flutter/third_party/tonic",
       "//third_party/dart/runtime/bin:dart_io_api",
+      "//third_party/dart/runtime/bin:elf_loader",
       "//third_party/skia",
     ]
 
diff --git a/shell/platform/embedder/embedder.cc b/shell/platform/embedder/embedder.cc
index 102b4f4..bfbcb00 100644
--- a/shell/platform/embedder/embedder.cc
+++ b/shell/platform/embedder/embedder.cc
@@ -11,6 +11,7 @@
 #include "flutter/fml/closure.h"
 #include "flutter/fml/make_copyable.h"
 #include "flutter/fml/native_library.h"
+#include "third_party/dart/runtime/bin/elf_loader.h"
 #include "third_party/dart/runtime/include/dart_native_api.h"
 
 #if OS_WIN
@@ -532,6 +533,83 @@
   fml::RefPtr<flutter::PlatformMessage> message;
 };
 
+struct LoadedElfDeleter {
+  void operator()(Dart_LoadedElf* elf) {
+    if (elf) {
+      ::Dart_UnloadELF(elf);
+    }
+  }
+};
+
+using UniqueLoadedElf = std::unique_ptr<Dart_LoadedElf, LoadedElfDeleter>;
+
+struct _FlutterEngineAOTData {
+  UniqueLoadedElf loaded_elf = nullptr;
+  const uint8_t* vm_snapshot_data = nullptr;
+  const uint8_t* vm_snapshot_instrs = nullptr;
+  const uint8_t* vm_isolate_data = nullptr;
+  const uint8_t* vm_isolate_instrs = nullptr;
+};
+
+FlutterEngineResult FlutterEngineCreateAOTData(
+    const FlutterEngineAOTDataSource* source,
+    FlutterEngineAOTData* data_out) {
+  if (!flutter::DartVM::IsRunningPrecompiledCode()) {
+    return LOG_EMBEDDER_ERROR(kInvalidArguments,
+                              "AOT data can only be created in AOT mode.");
+  } else if (!source) {
+    return LOG_EMBEDDER_ERROR(kInvalidArguments, "Null source specified.");
+  } else if (!data_out) {
+    return LOG_EMBEDDER_ERROR(kInvalidArguments, "Null data_out specified.");
+  }
+
+  switch (source->type) {
+    case kFlutterEngineAOTDataSourceTypeElfPath: {
+      if (!source->elf_path || !fml::IsFile(source->elf_path)) {
+        return LOG_EMBEDDER_ERROR(kInvalidArguments,
+                                  "Invalid ELF path specified.");
+      }
+
+      auto aot_data = std::make_unique<_FlutterEngineAOTData>();
+      const char* error = nullptr;
+
+      Dart_LoadedElf* loaded_elf = Dart_LoadELF(
+          source->elf_path,               // file path
+          0,                              // file offset
+          &error,                         // error (out)
+          &aot_data->vm_snapshot_data,    // vm snapshot data (out)
+          &aot_data->vm_snapshot_instrs,  // vm snapshot instr (out)
+          &aot_data->vm_isolate_data,     // vm isolate data (out)
+          &aot_data->vm_isolate_instrs    // vm isolate instr (out)
+      );
+
+      if (loaded_elf == nullptr) {
+        return LOG_EMBEDDER_ERROR(kInvalidArguments, error);
+      }
+
+      aot_data->loaded_elf.reset(loaded_elf);
+
+      *data_out = aot_data.release();
+      return kSuccess;
+    }
+  }
+
+  return LOG_EMBEDDER_ERROR(
+      kInvalidArguments,
+      "Invalid FlutterEngineAOTDataSourceType type specified.");
+}
+
+FlutterEngineResult FlutterEngineCollectAOTData(FlutterEngineAOTData data) {
+  if (data) {
+    data->loaded_elf = nullptr;
+    data->vm_snapshot_data = nullptr;
+    data->vm_snapshot_instrs = nullptr;
+    data->vm_isolate_data = nullptr;
+    data->vm_isolate_instrs = nullptr;
+  }
+  return kSuccess;
+}
+
 void PopulateSnapshotMappingCallbacks(const FlutterProjectArgs* args,
                                       flutter::Settings& settings) {
   // There are no ownership concerns here as all mappings are owned by the
@@ -543,6 +621,20 @@
   };
 
   if (flutter::DartVM::IsRunningPrecompiledCode()) {
+    if (SAFE_ACCESS(args, aot_data, nullptr) != nullptr) {
+      settings.vm_snapshot_data =
+          make_mapping_callback(args->aot_data->vm_snapshot_data, 0);
+
+      settings.vm_snapshot_instr =
+          make_mapping_callback(args->aot_data->vm_snapshot_instrs, 0);
+
+      settings.isolate_snapshot_data =
+          make_mapping_callback(args->aot_data->vm_isolate_data, 0);
+
+      settings.isolate_snapshot_instr =
+          make_mapping_callback(args->aot_data->vm_isolate_instrs, 0);
+    }
+
     if (SAFE_ACCESS(args, vm_snapshot_data, nullptr) != nullptr) {
       settings.vm_snapshot_data = make_mapping_callback(
           args->vm_snapshot_data, SAFE_ACCESS(args, vm_snapshot_data_size, 0));
@@ -659,6 +751,18 @@
 
   flutter::Settings settings = flutter::SettingsFromCommandLine(command_line);
 
+  if (SAFE_ACCESS(args, aot_data, nullptr)) {
+    if (SAFE_ACCESS(args, vm_snapshot_data, nullptr) ||
+        SAFE_ACCESS(args, vm_snapshot_instructions, nullptr) ||
+        SAFE_ACCESS(args, isolate_snapshot_data, nullptr) ||
+        SAFE_ACCESS(args, isolate_snapshot_instructions, nullptr)) {
+      return LOG_EMBEDDER_ERROR(
+          kInvalidArguments,
+          "Multiple AOT sources specified. Embedders should provide either "
+          "*_snapshot_* buffers or aot_data, not both.");
+    }
+  }
+
   PopulateSnapshotMappingCallbacks(args, settings);
 
   settings.icu_data_path = icu_data_path;
diff --git a/shell/platform/embedder/embedder.h b/shell/platform/embedder/embedder.h
index cab0341..67c5437 100644
--- a/shell/platform/embedder/embedder.h
+++ b/shell/platform/embedder/embedder.h
@@ -959,6 +959,57 @@
 typedef void (*FlutterNativeThreadCallback)(FlutterNativeThreadType type,
                                             void* user_data);
 
+/// AOT data source type.
+typedef enum {
+  kFlutterEngineAOTDataSourceTypeElfPath
+} FlutterEngineAOTDataSourceType;
+
+/// This struct specifies one of the various locations the engine can look for
+/// AOT data sources.
+typedef struct {
+  FlutterEngineAOTDataSourceType type;
+  union {
+    /// Absolute path to an ELF library file.
+    const char* elf_path;
+  };
+} FlutterEngineAOTDataSource;
+
+/// An opaque object that describes the AOT data that can be used to launch a
+/// FlutterEngine instance in AOT mode.
+typedef struct _FlutterEngineAOTData* FlutterEngineAOTData;
+
+//------------------------------------------------------------------------------
+/// @brief      Creates the necessary data structures to launch a Flutter Dart
+///             application in AOT mode. The data may only be collected after
+///             all FlutterEngine instances launched using this data have been
+///             terminated.
+///
+/// @param[in]  source    The source of the AOT data.
+/// @param[out] data_out  The AOT data on success. Unchanged on failure.
+///
+/// @return     Returns if the AOT data could be successfully resolved.
+///
+FLUTTER_EXPORT
+FlutterEngineResult FlutterEngineCreateAOTData(
+    const FlutterEngineAOTDataSource* source,
+    FlutterEngineAOTData* data_out);
+
+//------------------------------------------------------------------------------
+/// @brief      Collects the AOT data.
+///
+/// @warning    The embedder must ensure that this call is made only after all
+///             FlutterEngine instances launched using this data have been
+///             terminated, and that all of those instances were launched with
+///             the FlutterProjectArgs::shutdown_dart_vm_when_done flag set to
+///             true.
+///
+/// @param[in]  data   The data to collect.
+///
+/// @return     Returns if the AOT data was successfully collected.
+///
+FLUTTER_EXPORT
+FlutterEngineResult FlutterEngineCollectAOTData(FlutterEngineAOTData data);
+
 typedef struct {
   /// The size of this struct. Must be sizeof(FlutterProjectArgs).
   size_t struct_size;
@@ -1146,6 +1197,14 @@
   /// See also:
   /// https://github.com/dart-lang/sdk/blob/ca64509108b3e7219c50d6c52877c85ab6a35ff2/runtime/vm/flag_list.h#L150
   int64_t dart_old_gen_heap_size;
+
+  /// The AOT data to be used in AOT operation.
+  ///
+  /// Embedders should instantiate and destroy this object via the
+  /// FlutterEngineCreateAOTData and FlutterEngineCollectAOTData methods.
+  ///
+  /// Embedders can provide either snapshot buffers or aot_data, but not both.
+  FlutterEngineAOTData aot_data;
 } FlutterProjectArgs;
 
 //------------------------------------------------------------------------------
diff --git a/shell/platform/embedder/tests/embedder_config_builder.cc b/shell/platform/embedder/tests/embedder_config_builder.cc
index 567729a..72cc3d4 100644
--- a/shell/platform/embedder/tests/embedder_config_builder.cc
+++ b/shell/platform/embedder/tests/embedder_config_builder.cc
@@ -4,6 +4,7 @@
 
 #include "flutter/shell/platform/embedder/tests/embedder_config_builder.h"
 
+#include "flutter/runtime/dart_vm.h"
 #include "flutter/shell/platform/embedder/embedder.h"
 #include "third_party/skia/include/core/SkBitmap.h"
 
@@ -75,12 +76,20 @@
   // to do this manually.
   AddCommandLineArgument("embedder_unittest");
 
-  if (preference == InitializationPreference::kInitialize) {
+  if (preference != InitializationPreference::kNoInitialize) {
     SetAssetsPath();
-    SetSnapshots();
     SetIsolateCreateCallbackHook();
     SetSemanticsCallbackHooks();
     AddCommandLineArgument("--disable-observatory");
+
+    if (preference == InitializationPreference::kSnapshotsInitialize ||
+        preference == InitializationPreference::kMultiAOTInitialize) {
+      SetSnapshots();
+    }
+    if (preference == InitializationPreference::kAOTDataInitialize ||
+        preference == InitializationPreference::kMultiAOTInitialize) {
+      SetAOTDataElf();
+    }
   }
 }
 
@@ -132,6 +141,10 @@
   }
 }
 
+void EmbedderConfigBuilder::SetAOTDataElf() {
+  project_args_.aot_data = context_.GetAOTData();
+}
+
 void EmbedderConfigBuilder::SetIsolateCreateCallbackHook() {
   project_args_.root_isolate_create_callback =
       EmbedderTestContext::GetIsolateCreateCallbackHook();
diff --git a/shell/platform/embedder/tests/embedder_config_builder.h b/shell/platform/embedder/tests/embedder_config_builder.h
index 1dfdd7f..6512f15 100644
--- a/shell/platform/embedder/tests/embedder_config_builder.h
+++ b/shell/platform/embedder/tests/embedder_config_builder.h
@@ -31,13 +31,15 @@
 class EmbedderConfigBuilder {
  public:
   enum class InitializationPreference {
-    kInitialize,
+    kSnapshotsInitialize,
+    kAOTDataInitialize,
+    kMultiAOTInitialize,
     kNoInitialize,
   };
 
   EmbedderConfigBuilder(EmbedderTestContext& context,
                         InitializationPreference preference =
-                            InitializationPreference::kInitialize);
+                            InitializationPreference::kSnapshotsInitialize);
 
   ~EmbedderConfigBuilder();
 
@@ -51,6 +53,8 @@
 
   void SetSnapshots();
 
+  void SetAOTDataElf();
+
   void SetIsolateCreateCallbackHook();
 
   void SetSemanticsCallbackHooks();
diff --git a/shell/platform/embedder/tests/embedder_test_context.cc b/shell/platform/embedder/tests/embedder_test_context.cc
index f0cdd88..bb2c717 100644
--- a/shell/platform/embedder/tests/embedder_test_context.cc
+++ b/shell/platform/embedder/tests/embedder_test_context.cc
@@ -8,6 +8,7 @@
 #include "flutter/fml/paths.h"
 #include "flutter/runtime/dart_vm.h"
 #include "flutter/shell/platform/embedder/tests/embedder_assertions.h"
+#include "flutter/testing/testing.h"
 #include "third_party/dart/runtime/bin/elf_loader.h"
 #include "third_party/skia/include/core/SkSurface.h"
 
@@ -19,6 +20,7 @@
       aot_symbols_(LoadELFSymbolFromFixturesIfNeccessary()),
       native_resolver_(std::make_shared<TestDartNativeResolver>()) {
   SetupAOTMappingsIfNecessary();
+  SetupAOTDataIfNecessary();
   isolate_create_callbacks_.push_back(
       [weak_resolver =
            std::weak_ptr<TestDartNativeResolver>{native_resolver_}]() {
@@ -44,6 +46,24 @@
       aot_symbols_.vm_isolate_instrs, 0u);
 }
 
+void EmbedderTestContext::SetupAOTDataIfNecessary() {
+  if (!DartVM::IsRunningPrecompiledCode()) {
+    return;
+  }
+  FlutterEngineAOTDataSource data_in = {};
+  FlutterEngineAOTData data_out = nullptr;
+
+  const auto elf_path =
+      fml::paths::JoinPaths({GetFixturesPath(), kAOTAppELFFileName});
+
+  data_in.type = kFlutterEngineAOTDataSourceTypeElfPath;
+  data_in.elf_path = elf_path.c_str();
+
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kSuccess);
+
+  aot_data_.reset(data_out);
+}
+
 const std::string& EmbedderTestContext::GetAssetsPath() const {
   return assets_path_;
 }
@@ -65,6 +85,10 @@
   return isolate_snapshot_instructions_.get();
 }
 
+FlutterEngineAOTData EmbedderTestContext::GetAOTData() const {
+  return aot_data_.get();
+}
+
 void EmbedderTestContext::SetRootSurfaceTransformation(SkMatrix matrix) {
   root_surface_transformation_ = matrix;
 }
diff --git a/shell/platform/embedder/tests/embedder_test_context.h b/shell/platform/embedder/tests/embedder_test_context.h
index fc7bf59..182edcd 100644
--- a/shell/platform/embedder/tests/embedder_test_context.h
+++ b/shell/platform/embedder/tests/embedder_test_context.h
@@ -28,6 +28,16 @@
 using SemanticsActionCallback =
     std::function<void(const FlutterSemanticsCustomAction*)>;
 
+struct AOTDataDeleter {
+  void operator()(FlutterEngineAOTData aot_data) {
+    if (aot_data) {
+      FlutterEngineCollectAOTData(aot_data);
+    }
+  }
+};
+
+using UniqueAOTData = std::unique_ptr<_FlutterEngineAOTData, AOTDataDeleter>;
+
 class EmbedderTestContext {
  public:
   EmbedderTestContext(std::string assets_path = "");
@@ -44,6 +54,8 @@
 
   const fml::Mapping* GetIsolateSnapshotInstructions() const;
 
+  FlutterEngineAOTData GetAOTData() const;
+
   void SetRootSurfaceTransformation(SkMatrix matrix);
 
   void AddIsolateCreateCallback(fml::closure closure);
@@ -79,6 +91,7 @@
   std::unique_ptr<fml::Mapping> vm_snapshot_instructions_;
   std::unique_ptr<fml::Mapping> isolate_snapshot_data_;
   std::unique_ptr<fml::Mapping> isolate_snapshot_instructions_;
+  UniqueAOTData aot_data_;
   std::vector<fml::closure> isolate_create_callbacks_;
   std::shared_ptr<TestDartNativeResolver> native_resolver_;
   SemanticsNodeCallback update_semantics_node_callback_;
@@ -101,6 +114,8 @@
 
   void SetupAOTMappingsIfNecessary();
 
+  void SetupAOTDataIfNecessary();
+
   void SetupCompositor();
 
   void FireIsolateCreateCallbacks();
diff --git a/shell/platform/embedder/tests/embedder_unittests.cc b/shell/platform/embedder/tests/embedder_unittests.cc
index 78286fc..82e392b 100644
--- a/shell/platform/embedder/tests/embedder_unittests.cc
+++ b/shell/platform/embedder/tests/embedder_unittests.cc
@@ -4207,5 +4207,114 @@
   latch.Wait();
 }
 
+TEST_F(EmbedderTest, InvalidAOTDataSourcesMustReturnError) {
+  if (!DartVM::IsRunningPrecompiledCode()) {
+    GTEST_SKIP();
+    return;
+  }
+  FlutterEngineAOTDataSource data_in = {};
+  FlutterEngineAOTData data_out = nullptr;
+
+  // Null source specified.
+  ASSERT_EQ(FlutterEngineCreateAOTData(nullptr, &data_out), kInvalidArguments);
+  ASSERT_EQ(data_out, nullptr);
+
+  // Null data_out specified.
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, nullptr), kInvalidArguments);
+
+  // Invalid FlutterEngineAOTDataSourceType type specified.
+  data_in.type = FlutterEngineAOTDataSourceType(-1);
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kInvalidArguments);
+  ASSERT_EQ(data_out, nullptr);
+
+  // Invalid ELF path specified.
+  data_in.type = kFlutterEngineAOTDataSourceTypeElfPath;
+  data_in.elf_path = nullptr;
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kInvalidArguments);
+  ASSERT_EQ(data_in.type, kFlutterEngineAOTDataSourceTypeElfPath);
+  ASSERT_EQ(data_in.elf_path, nullptr);
+  ASSERT_EQ(data_out, nullptr);
+
+  // Invalid ELF path specified.
+  data_in.elf_path = "";
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kInvalidArguments);
+  ASSERT_EQ(data_in.type, kFlutterEngineAOTDataSourceTypeElfPath);
+  ASSERT_EQ(data_in.elf_path, "");
+  ASSERT_EQ(data_out, nullptr);
+
+  // Could not find VM snapshot data.
+  data_in.elf_path = "/bin/true";
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kInvalidArguments);
+  ASSERT_EQ(data_in.type, kFlutterEngineAOTDataSourceTypeElfPath);
+  ASSERT_EQ(data_in.elf_path, "/bin/true");
+  ASSERT_EQ(data_out, nullptr);
+}
+
+TEST_F(EmbedderTest, MustNotRunWithMultipleAOTSources) {
+  if (!DartVM::IsRunningPrecompiledCode()) {
+    GTEST_SKIP();
+    return;
+  }
+  auto& context = GetEmbedderContext();
+
+  EmbedderConfigBuilder builder(
+      context,
+      EmbedderConfigBuilder::InitializationPreference::kMultiAOTInitialize);
+
+  builder.SetSoftwareRendererConfig();
+
+  auto engine = builder.LaunchEngine();
+  ASSERT_FALSE(engine.is_valid());
+}
+
+TEST_F(EmbedderTest, CanCreateAndCollectAValidElfSource) {
+  if (!DartVM::IsRunningPrecompiledCode()) {
+    GTEST_SKIP();
+    return;
+  }
+  FlutterEngineAOTDataSource data_in = {};
+  FlutterEngineAOTData data_out = nullptr;
+
+  // Collecting a null object should be allowed
+  ASSERT_EQ(FlutterEngineCollectAOTData(data_out), kSuccess);
+
+  const auto elf_path =
+      fml::paths::JoinPaths({GetFixturesPath(), kAOTAppELFFileName});
+
+  data_in.type = kFlutterEngineAOTDataSourceTypeElfPath;
+  data_in.elf_path = elf_path.c_str();
+
+  ASSERT_EQ(FlutterEngineCreateAOTData(&data_in, &data_out), kSuccess);
+  ASSERT_EQ(data_in.type, kFlutterEngineAOTDataSourceTypeElfPath);
+  ASSERT_EQ(data_in.elf_path, elf_path.c_str());
+  ASSERT_NE(data_out, nullptr);
+
+  ASSERT_EQ(FlutterEngineCollectAOTData(data_out), kSuccess);
+}
+
+TEST_F(EmbedderTest, CanLaunchAndShutdownWithAValidElfSource) {
+  if (!DartVM::IsRunningPrecompiledCode()) {
+    GTEST_SKIP();
+    return;
+  }
+  auto& context = GetEmbedderContext();
+
+  fml::AutoResetWaitableEvent latch;
+  context.AddIsolateCreateCallback([&latch]() { latch.Signal(); });
+
+  EmbedderConfigBuilder builder(
+      context,
+      EmbedderConfigBuilder::InitializationPreference::kAOTDataInitialize);
+
+  builder.SetSoftwareRendererConfig();
+
+  auto engine = builder.LaunchEngine();
+  ASSERT_TRUE(engine.is_valid());
+
+  // Wait for the root isolate to launch.
+  latch.Wait();
+  engine.reset();
+}
+
 }  // namespace testing
 }  // namespace flutter
diff --git a/testing/elf_loader.cc b/testing/elf_loader.cc
index 9b1eab1..618859d 100644
--- a/testing/elf_loader.cc
+++ b/testing/elf_loader.cc
@@ -12,8 +12,6 @@
 namespace flutter {
 namespace testing {
 
-static constexpr const char* kAOTAppELFFileName = "app_elf_snapshot.so";
-
 ELFAOTSymbols LoadELFSymbolFromFixturesIfNeccessary() {
   if (!DartVM::IsRunningPrecompiledCode()) {
     return {};
diff --git a/testing/elf_loader.h b/testing/elf_loader.h
index 1d9a93d..a8c37c9 100644
--- a/testing/elf_loader.h
+++ b/testing/elf_loader.h
@@ -14,6 +14,8 @@
 namespace flutter {
 namespace testing {
 
+inline constexpr const char* kAOTAppELFFileName = "app_elf_snapshot.so";
+
 struct LoadedELFDeleter {
   void operator()(Dart_LoadedElf* elf) { ::Dart_UnloadELF(elf); }
 };