[vm/ffi] AOT kernel concatenation with native assets

gen_snapshot can now consume a concatenated dill file.

Restrictions:
* The first dill file must be the main program.
* The subsequent dill files must only contain dummy libraries with
  objects which are recognized classes in the VM and not treeshaken
  in the main program.

This will enable compiling Dart AOT first, then using the tree-shaking
information to tree-shake native assets, and then embedding the native
assets mapping in the kernel file before running gen_snapshot.

TEST=tests/ffi/native_assets/asset_relative_test.dart

Closes: https://github.com/dart-lang/sdk/issues/50152
Bug: https://github.com/dart-lang/sdk/issues/55377
Change-Id: Id562aa39840d5eb467198efaa4cf3152d860f8b5
q-Include-Trybots: dart/try:vm-aot-linux-debug-x64-try,vm-aot-linux-debug-x64c-try,vm-aot-mac-release-arm64-try,vm-aot-mac-release-x64-try,vm-aot-obfuscate-linux-release-x64-try,vm-aot-optimization-level-linux-release-x64-try,vm-aot-win-debug-arm64-try,vm-aot-win-debug-x64-try,vm-aot-win-debug-x64c-try
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/362381
Reviewed-by: Alexander Markov <alexmarkov@google.com>
Commit-Queue: Daco Harkes <dacoharkes@google.com>
diff --git a/runtime/vm/bootstrap.cc b/runtime/vm/bootstrap.cc
index 4ad7fc3..401b710 100644
--- a/runtime/vm/bootstrap.cc
+++ b/runtime/vm/bootstrap.cc
@@ -5,6 +5,7 @@
 #include "vm/bootstrap.h"
 
 #include <memory>
+#include <utility>
 
 #include "include/dart_api.h"
 
@@ -108,22 +109,10 @@
   cls.EnsureIsFinalized(thread);
 }
 
-static ErrorPtr BootstrapFromKernel(Thread* thread,
-                                    const uint8_t* kernel_buffer,
-                                    intptr_t kernel_buffer_size) {
+static ErrorPtr BootstrapFromKernelSingleProgram(
+    Thread* thread,
+    std::unique_ptr<kernel::Program> program) {
   Zone* zone = thread->zone();
-  const char* error = nullptr;
-  std::unique_ptr<kernel::Program> program = kernel::Program::ReadFromBuffer(
-      kernel_buffer, kernel_buffer_size, &error);
-  if (program == nullptr) {
-    const intptr_t kMessageBufferSize = 512;
-    char message_buffer[kMessageBufferSize];
-    Utils::SNPrint(message_buffer, kMessageBufferSize,
-                   "Can't load Kernel binary: %s.", error);
-    const String& msg = String::Handle(String::New(message_buffer, Heap::kOld));
-    return ApiError::New(msg, Heap::kOld);
-  }
-
   LongJumpScope jump;
   if (setjmp(*jump.Set()) == 0) {
     kernel::KernelLoader loader(program.get(), /*uri_to_source_table=*/nullptr);
@@ -167,6 +156,66 @@
   return Thread::Current()->StealStickyError();
 }
 
+static ErrorPtr BootstrapFromKernel(Thread* thread,
+                                    const uint8_t* kernel_buffer,
+                                    intptr_t kernel_buffer_size) {
+  Zone* zone = thread->zone();
+  const char* error = nullptr;
+  std::unique_ptr<kernel::Program> program = kernel::Program::ReadFromBuffer(
+      kernel_buffer, kernel_buffer_size, &error);
+  if (program == nullptr) {
+    const intptr_t kMessageBufferSize = 512;
+    char message_buffer[kMessageBufferSize];
+    Utils::SNPrint(message_buffer, kMessageBufferSize,
+                   "Can't load Kernel binary: %s.", error);
+    const String& msg = String::Handle(String::New(message_buffer, Heap::kOld));
+    return ApiError::New(msg, Heap::kOld);
+  }
+
+  if (program->is_single_program()) {
+    return BootstrapFromKernelSingleProgram(thread, std::move(program));
+  }
+
+  GrowableArray<intptr_t> subprogram_file_starts;
+  {
+    kernel::Reader reader(program->binary());
+    kernel::KernelLoader::index_programs(&reader, &subprogram_file_starts);
+  }
+  intptr_t subprogram_count = subprogram_file_starts.length() - 1;
+
+  // Create "fake programs" for each sub-program.
+  auto& load_result = Error::Handle(zone);
+  for (intptr_t i = 0; i < subprogram_count; i++) {
+    intptr_t subprogram_start = subprogram_file_starts.At(i);
+    intptr_t subprogram_end = subprogram_file_starts.At(i + 1);
+    const auto& component = TypedDataBase::Handle(
+        program->binary().ViewFromTo(subprogram_start, subprogram_end));
+    kernel::Reader reader(component);
+    const char* error = nullptr;
+    std::unique_ptr<kernel::Program> subprogram =
+        kernel::Program::ReadFrom(&reader, &error);
+    if (subprogram == nullptr) {
+      FATAL("Failed to load kernel file: %s", error);
+    }
+    ASSERT(subprogram->is_single_program());
+    if (i == 0) {
+      // The first subprogram must be the main Dart program.
+      load_result ^=
+          BootstrapFromKernelSingleProgram(thread, std::move(subprogram));
+    } else {
+      // Restrictions on the subsequent programs: Must contain only
+      // contain dummy libraries with VM recognized classes (or classes kept
+      // fully intact by tree-shaking).
+      // Currently only used for concatenating native assets mappings.
+      kernel::KernelLoader loader(subprogram.get(),
+                                  /*uri_to_source_table=*/nullptr);
+      load_result ^= loader.LoadProgram(false);
+    }
+    if (load_result.IsError()) return load_result.ptr();
+  }
+  return Error::null();
+}
+
 ErrorPtr Bootstrap::DoBootstrapping(const uint8_t* kernel_buffer,
                                     intptr_t kernel_buffer_size) {
   Thread* thread = Thread::Current();
diff --git a/runtime/vm/kernel_loader.h b/runtime/vm/kernel_loader.h
index f7e46e0..d817594 100644
--- a/runtime/vm/kernel_loader.h
+++ b/runtime/vm/kernel_loader.h
@@ -220,6 +220,9 @@
                                         const Function& parent_function,
                                         const Object& closure_owner);
 
+  static void index_programs(kernel::Reader* reader,
+                             GrowableArray<intptr_t>* subprogram_file_starts);
+
  private:
   // Pragma bits
   using HasPragma = BitField<uint32_t, bool, 0, 1>;
@@ -279,8 +282,6 @@
 
   uint8_t CharacterAt(StringIndex string_index, intptr_t index);
 
-  static void index_programs(kernel::Reader* reader,
-                             GrowableArray<intptr_t>* subprogram_file_starts);
   void walk_incremental_kernel(BitVector* modified_libs,
                                bool* is_empty_program,
                                intptr_t* p_num_classes,
diff --git a/tests/ffi/native_assets/asset_relative_test.dart b/tests/ffi/native_assets/asset_relative_test.dart
index b083c05..e75793a 100644
--- a/tests/ffi/native_assets/asset_relative_test.dart
+++ b/tests/ffi/native_assets/asset_relative_test.dart
@@ -60,8 +60,7 @@
   await invokeSelf(
     selfSourceUri: selfSourceUri,
     runtime: Runtime.aot,
-    // TODO(https://dartbug.com/55377): Support concatenation in AOT.
-    // kernelCombine: KernelCombine.concatenation,
+    kernelCombine: KernelCombine.concatenation,
     relativePath: RelativePath.up,
     arguments: [runTestsArg],
   );