1 files changed, 138 insertions, 0 deletions
diff --git a/core/snapshot_util.rs b/core/snapshot_util.rs
index 38ab18b4f..74bd7d8a3 100644
--- a/core/snapshot_util.rs
+++ b/core/snapshot_util.rs
@@ -107,3 +107,141 @@ pub fn get_js_files(
   js_files.sort();
   js_files
 }
+
+fn data_error_to_panic(err: v8::DataError) -> ! {
+  match err {
+    v8::DataError::BadType { actual, expected } => {
+      panic!(
+        "Invalid type for snapshot data: expected {expected}, got {actual}"
+      );
+    }
+    v8::DataError::NoData { expected } => {
+      panic!("No data for snapshot data: expected {expected}");
+    }
+  }
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub(crate) enum SnapshotOptions {
+  Load,
+  CreateFromExisting,
+  Create,
+  None,
+}
+
+impl SnapshotOptions {
+  pub fn loaded(&self) -> bool {
+    matches!(self, Self::Load | Self::CreateFromExisting)
+  }
+
+  pub fn will_snapshot(&self) -> bool {
+    matches!(self, Self::Create | Self::CreateFromExisting)
+  }
+
+  pub fn from_bools(snapshot_loaded: bool, will_snapshot: bool) -> Self {
+    match (snapshot_loaded, will_snapshot) {
+      (true, true) => Self::CreateFromExisting,
+      (false, true) => Self::Create,
+      (true, false) => Self::Load,
+      (false, false) => Self::None,
+    }
+  }
+}
+
+pub(crate) struct SnapshottedData {
+  pub module_map_data: v8::Global<v8::Array>,
+  pub module_handles: Vec<v8::Global<v8::Module>>,
+}
+
+static MODULE_MAP_CONTEXT_DATA_INDEX: usize = 0;
+
+pub(crate) fn get_snapshotted_data(
+  scope: &mut v8::HandleScope<()>,
+  context: v8::Local<v8::Context>,
+) -> SnapshottedData {
+  let mut scope = v8::ContextScope::new(scope, context);
+
+  // The 0th element is the module map itself, followed by X number of module
+  // handles. We need to deserialize the "next_module_id" field from the
+  // map to see how many module handles we expect.
+  let result = scope.get_context_data_from_snapshot_once::<v8::Array>(
+    MODULE_MAP_CONTEXT_DATA_INDEX,
+  );
+
+  let val = match result {
+    Ok(v) => v,
+    Err(err) => data_error_to_panic(err),
+  };
+
+  let next_module_id = {
+    let info_data: v8::Local<v8::Array> =
+      val.get_index(&mut scope, 1).unwrap().try_into().unwrap();
+    info_data.length()
+  };
+
+  // Over allocate so executing a few scripts doesn't have to resize this vec.
+  let mut module_handles = Vec::with_capacity(next_module_id as usize + 16);
+  for i in 1..=next_module_id {
+    match scope.get_context_data_from_snapshot_once::<v8::Module>(i as usize) {
+      Ok(val) => {
+        let module_global = v8::Global::new(&mut scope, val);
+        module_handles.push(module_global);
+      }
+      Err(err) => data_error_to_panic(err),
+    }
+  }
+
+  SnapshottedData {
+    module_map_data: v8::Global::new(&mut scope, val),
+    module_handles,
+  }
+}
+
+pub(crate) fn set_snapshotted_data(
+  scope: &mut v8::HandleScope<()>,
+  context: v8::Global<v8::Context>,
+  snapshotted_data: SnapshottedData,
+) {
+  let local_context = v8::Local::new(scope, context);
+  let local_data = v8::Local::new(scope, snapshotted_data.module_map_data);
+  let offset = scope.add_context_data(local_context, local_data);
+  assert_eq!(offset, MODULE_MAP_CONTEXT_DATA_INDEX);
+
+  for (index, handle) in snapshotted_data.module_handles.into_iter().enumerate()
+  {
+    let module_handle = v8::Local::new(scope, handle);
+    let offset = scope.add_context_data(local_context, module_handle);
+    assert_eq!(offset, index + 1);
+  }
+}
+
+/// Returns an isolate set up for snapshotting.
+pub(crate) fn create_snapshot_creator(
+  external_refs: &'static v8::ExternalReferences,
+  maybe_startup_snapshot: Option<Snapshot>,
+) -> v8::OwnedIsolate {
+  if let Some(snapshot) = maybe_startup_snapshot {
+    match snapshot {
+      Snapshot::Static(data) => {
+        v8::Isolate::snapshot_creator_from_existing_snapshot(
+          data,
+          Some(external_refs),
+        )
+      }
+      Snapshot::JustCreated(data) => {
+        v8::Isolate::snapshot_creator_from_existing_snapshot(
+          data,
+          Some(external_refs),
+        )
+      }
+      Snapshot::Boxed(data) => {
+        v8::Isolate::snapshot_creator_from_existing_snapshot(
+          data,
+          Some(external_refs),
+        )
+      }
+    }
+  } else {
+    v8::Isolate::snapshot_creator(Some(external_refs))
+  }
+}