use anyhow::Error;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;
use std::rc::Rc;
use std::time::Instant;
use crate::modules::ModuleMapSnapshotData;
use crate::Extension;
use crate::JsRuntimeForSnapshot;
use crate::RuntimeOptions;
use super::ExtensionTranspiler;
pub type WithRuntimeCb = dyn Fn(&mut JsRuntimeForSnapshot);
pub type SnapshotDataId = u32;
const ULEN: usize = std::mem::size_of::<usize>();
pub(crate) struct V8Snapshot(pub(crate) &'static [u8]);
pub(crate) fn deconstruct(
slice: &'static [u8],
) -> (V8Snapshot, SerializableSnapshotSidecarData) {
let len =
usize::from_le_bytes(slice[slice.len() - ULEN..].try_into().unwrap());
let data = SerializableSnapshotSidecarData::from_slice(
&slice[len..slice.len() - ULEN],
);
(V8Snapshot(&slice[0..len]), data)
}
pub(crate) fn serialize(
v8_data: v8::StartupData,
sidecar_data: SerializableSnapshotSidecarData,
) -> Box<[u8]> {
let v8_data_len = v8_data.len();
let sidecar_data = sidecar_data.into_bytes();
let mut data = Vec::with_capacity(v8_data_len + sidecar_data.len() + ULEN);
data.extend_from_slice(&v8_data);
data.extend_from_slice(&sidecar_data);
data.extend_from_slice(&v8_data_len.to_le_bytes());
data.into_boxed_slice()
}
#[derive(Default)]
pub struct SnapshotLoadDataStore {
data: Vec<Option<v8::Global<v8::Data>>>,
}
impl SnapshotLoadDataStore {
pub fn get<'s, T>(
&mut self,
scope: &mut v8::HandleScope<'s>,
id: SnapshotDataId,
) -> v8::Global<T>
where
v8::Local<'s, T>: TryFrom<v8::Local<'s, v8::Data>>,
{
let Some(data) = self.data.get_mut(id as usize) else {
panic!(
"Attempted to read snapshot data out of range: {id} (of {})",
self.data.len()
);
};
let Some(data) = data.take() else {
panic!("Attempted to read the snapshot data at index {id} twice");
};
let local = v8::Local::new(scope, data);
let local = v8::Local::<T>::try_from(local).unwrap_or_else(|_| {
panic!(
"Invalid data type at index {id}, expected '{}'",
std::any::type_name::<T>()
)
});
v8::Global::new(scope, local)
}
}
#[derive(Default)]
pub struct SnapshotStoreDataStore {
data: Vec<v8::Global<v8::Data>>,
}
impl SnapshotStoreDataStore {
pub fn register<T>(&mut self, global: v8::Global<T>) -> SnapshotDataId
where
for<'s> v8::Local<'s, v8::Data>: From<v8::Local<'s, T>>,
{
let id = self.data.len();
unsafe {
self.data.push(std::mem::transmute(global));
}
id as _
}
}
pub struct CreateSnapshotOptions {
pub cargo_manifest_dir: &'static str,
pub startup_snapshot: Option<&'static [u8]>,
pub skip_op_registration: bool,
pub extensions: Vec<Extension>,
pub extension_transpiler: Option<Rc<ExtensionTranspiler>>,
pub with_runtime_cb: Option<Box<WithRuntimeCb>>,
}
pub struct CreateSnapshotOutput {
pub files_loaded_during_snapshot: Vec<PathBuf>,
pub output: Box<[u8]>,
}
#[must_use = "The files listed by create_snapshot should be printed as 'cargo:rerun-if-changed' lines"]
pub fn create_snapshot(
create_snapshot_options: CreateSnapshotOptions,
warmup_script: Option<&'static str>,
) -> Result<CreateSnapshotOutput, Error> {
let mut mark = Instant::now();
println!("Creating a snapshot...",);
let warmup_exts = warmup_script.map(|_| {
create_snapshot_options
.extensions
.iter()
.map(|e| e.for_warmup())
.collect::<Vec<_>>()
});
let mut js_runtime = JsRuntimeForSnapshot::new(RuntimeOptions {
startup_snapshot: create_snapshot_options.startup_snapshot,
extensions: create_snapshot_options.extensions,
extension_transpiler: create_snapshot_options.extension_transpiler,
skip_op_registration: create_snapshot_options.skip_op_registration,
..Default::default()
});
println!("JsRuntimeForSnapshot prepared, took {:#?}", mark.elapsed(),);
mark = Instant::now();
let files_loaded_during_snapshot = js_runtime
.files_loaded_from_fs_during_snapshot()
.iter()
.map(PathBuf::from)
.collect::<Vec<_>>();
if let Some(ref with_runtime_cb) = create_snapshot_options.with_runtime_cb {
with_runtime_cb(&mut js_runtime);
}
let mut snapshot = js_runtime.snapshot();
if let Some(warmup_script) = warmup_script {
let leaked_snapshot = Box::leak(snapshot);
let warmup_exts = warmup_exts.unwrap();
let mut js_runtime = JsRuntimeForSnapshot::new(RuntimeOptions {
startup_snapshot: Some(leaked_snapshot),
extensions: warmup_exts,
skip_op_registration: true,
..Default::default()
});
if let Some(with_runtime_cb) = create_snapshot_options.with_runtime_cb {
with_runtime_cb(&mut js_runtime);
}
js_runtime.execute_script("warmup", warmup_script)?;
snapshot = js_runtime.snapshot();
}
println!(
"Snapshot size: {}, took {:#?}",
snapshot.len(),
mark.elapsed(),
);
mark = Instant::now();
println!(
"Snapshot written, took: {:#?}",
Instant::now().saturating_duration_since(mark),
);
Ok(CreateSnapshotOutput {
files_loaded_during_snapshot,
output: snapshot,
})
}
pub type FilterFn = Box<dyn Fn(&PathBuf) -> bool>;
pub fn get_js_files(
cargo_manifest_dir: &'static str,
directory: &str,
filter: Option<FilterFn>,
) -> Vec<PathBuf> {
let manifest_dir = Path::new(cargo_manifest_dir);
let mut js_files = std::fs::read_dir(directory)
.unwrap()
.map(|dir_entry| {
let file = dir_entry.unwrap();
manifest_dir.join(file.path())
})
.filter(|path| {
path.extension().unwrap_or_default() == "js"
&& filter.as_ref().map(|filter| filter(path)).unwrap_or(true)
})
.collect::<Vec<PathBuf>>();
js_files.sort();
js_files
}
#[derive(Serialize, Deserialize)]
pub(crate) struct SnapshottedData<'snapshot> {
pub js_handled_promise_rejection_cb: Option<u32>,
pub module_map_data: ModuleMapSnapshotData,
pub externals_count: u32,
pub extension_count: usize,
pub op_count: usize,
pub source_count: usize,
pub addl_refs_count: usize,
#[serde(borrow)]
pub ext_source_maps: HashMap<String, &'snapshot [u8]>,
#[serde(borrow)]
pub external_strings: Vec<&'snapshot [u8]>,
}
#[derive(Serialize, Deserialize)]
pub(crate) struct SerializableSnapshotSidecarData<'snapshot> {
data_count: u32,
#[serde(borrow)]
pub snapshot_data: SnapshottedData<'snapshot>,
}
impl<'snapshot> SerializableSnapshotSidecarData<'snapshot> {
fn from_slice(slice: &'snapshot [u8]) -> Self {
bincode::deserialize(slice).expect("Failed to deserialize snapshot data")
}
fn into_bytes(self) -> Box<[u8]> {
bincode::serialize(&self).unwrap().into_boxed_slice()
}
}
pub(crate) fn load_snapshotted_data_from_snapshot<'snapshot>(
scope: &mut v8::HandleScope<()>,
context: v8::Local<v8::Context>,
raw_data: SerializableSnapshotSidecarData<'snapshot>,
) -> (SnapshottedData<'snapshot>, SnapshotLoadDataStore) {
let scope = &mut v8::ContextScope::new(scope, context);
let mut data = SnapshotLoadDataStore::default();
for i in 0..raw_data.data_count {
let item = scope
.get_context_data_from_snapshot_once::<v8::Data>(i as usize)
.unwrap();
let item = v8::Global::new(scope, item);
data.data.push(Some(item));
}
(raw_data.snapshot_data, data)
}
pub(crate) fn store_snapshotted_data_for_snapshot<'snapshot>(
scope: &mut v8::HandleScope,
context: v8::Global<v8::Context>,
snapshotted_data: SnapshottedData<'snapshot>,
data_store: SnapshotStoreDataStore,
) -> SerializableSnapshotSidecarData<'snapshot> {
let context = v8::Local::new(scope, context);
let raw_snapshot_data = SerializableSnapshotSidecarData {
data_count: data_store.data.len() as _,
snapshot_data: snapshotted_data,
};
for data in data_store.data {
let data = v8::Local::new(scope, data);
scope.add_context_data(context, data);
}
raw_snapshot_data
}
pub(crate) fn create_snapshot_creator(
external_refs: &'static v8::ExternalReferences,
maybe_startup_snapshot: Option<V8Snapshot>,
) -> v8::OwnedIsolate {
if let Some(snapshot) = maybe_startup_snapshot {
v8::Isolate::snapshot_creator_from_existing_snapshot(
snapshot.0,
Some(external_refs),
)
} else {
v8::Isolate::snapshot_creator(Some(external_refs))
}
}