use crate::compiler::Program;
use crate::host::{FuncDef, TryDef};
use fusevm::Chunk;
use rkyv::{Archive, Deserialize as RkyvDe, Serialize as RkyvSer};
use serde::{Deserialize, Serialize};
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
const SCHEMA: u64 = 12;
#[derive(Archive, RkyvSer, RkyvDe, Default)]
#[archive(check_bytes)]
struct Shard {
entries: Vec<Entry>,
}
#[derive(Archive, RkyvSer, RkyvDe)]
#[archive(check_bytes)]
struct Entry {
key: u64,
verify: u64,
build: u64,
blob: Vec<u8>,
}
#[derive(Serialize, Deserialize)]
struct CProg {
main: Chunk,
functions: Vec<(String, FuncDef)>,
tries: Vec<TryDef>,
#[serde(default)]
sites: crate::host::SiteTables,
#[serde(default)]
strict: bool,
}
const BUILD_VERSION: &str = env!("CARGO_PKG_VERSION");
fn build_id() -> u64 {
use std::sync::OnceLock;
static ID: OnceLock<u64> = OnceLock::new();
*ID.get_or_init(|| {
std::env::current_exe()
.and_then(|p| p.metadata())
.and_then(|m| m.modified())
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map(|d| d.as_nanos() as u64)
.unwrap_or(0)
})
}
pub fn key_for(src: &str) -> u64 {
let mut h = rustc_hash::FxHasher::default();
SCHEMA.hash(&mut h);
BUILD_VERSION.hash(&mut h);
build_id().hash(&mut h);
src.hash(&mut h);
h.finish()
}
fn verify_for(src: &str) -> u64 {
use std::collections::hash_map::DefaultHasher;
let mut h = DefaultHasher::new();
SCHEMA.hash(&mut h);
BUILD_VERSION.hash(&mut h);
build_id().hash(&mut h);
src.len().hash(&mut h);
src.hash(&mut h);
h.finish()
}
fn shard_path() -> Option<PathBuf> {
let dir = dirs::home_dir()?.join(".node-js");
let _ = std::fs::create_dir_all(&dir);
Some(dir.join("scripts.rkyv"))
}
type Stamp = Option<(std::time::SystemTime, u64)>;
fn shard_stamp() -> Stamp {
let path = shard_path()?;
let md = std::fs::metadata(&path).ok()?;
Some((md.modified().ok()?, md.len()))
}
fn load_shard() -> Shard {
let Some(path) = shard_path() else {
return Shard::default();
};
let Ok(bytes) = std::fs::read(&path) else {
return Shard::default();
};
rkyv::from_bytes::<Shard>(&bytes).unwrap_or_default()
}
fn write_shard(shard: &Shard) -> Result<(), String> {
let path = shard_path().ok_or("no home dir for cache")?;
let bytes = rkyv::to_bytes::<_, 4096>(shard).map_err(|e| format!("cache serialize: {e}"))?;
static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let n = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let tmp = path.with_extension(format!("rkyv.tmp.{}.{n}", std::process::id()));
std::fs::write(&tmp, &bytes).map_err(|e| format!("cache write: {e}"))?;
std::fs::rename(&tmp, &path).map_err(|e| {
let _ = std::fs::remove_file(&tmp);
format!("cache rename: {e}")
})
}
#[derive(Default)]
struct ShardMem {
backing: Vec<u8>,
disk: rustc_hash::FxHashMap<u64, (u64, std::ops::Range<usize>)>,
added: rustc_hash::FxHashMap<u64, (u64, Vec<u8>)>,
dirty: bool,
stamp: Stamp,
}
impl ShardMem {
fn get(&self, key: u64) -> Option<(u64, &[u8])> {
if let Some((v, b)) = self.added.get(&key) {
return Some((*v, b.as_slice()));
}
let (v, r) = self.disk.get(&key)?;
Some((*v, &self.backing[r.clone()]))
}
fn iter(&self) -> impl Iterator<Item = (u64, u64, &[u8])> {
self.added
.iter()
.map(|(k, (v, b))| (*k, *v, b.as_slice()))
.chain(
self.disk
.iter()
.filter(|(k, _)| !self.added.contains_key(k))
.map(|(k, (v, r))| (*k, *v, &self.backing[r.clone()])),
)
}
}
fn load_shard_indexed() -> ShardMem {
let mut mem = ShardMem {
stamp: shard_stamp(),
..ShardMem::default()
};
let Some(path) = shard_path() else {
return mem;
};
let Ok(bytes) = std::fs::read(&path) else {
return mem;
};
let build = build_id();
let Ok(shard) = rkyv::check_archived_root::<Shard>(&bytes) else {
return mem;
};
let base = bytes.as_ptr() as usize;
for e in shard.entries.iter() {
if u64::from(e.build) != build {
continue;
}
let blob: &[u8] = &e.blob;
let off = blob.as_ptr() as usize - base;
mem.disk
.insert(e.key.into(), (e.verify.into(), off..off + blob.len()));
}
mem.backing = bytes;
mem
}
thread_local! {
static SHARD: std::cell::RefCell<Option<ShardMem>> = const { std::cell::RefCell::new(None) };
}
fn with_shard<T>(f: impl FnOnce(&mut ShardMem) -> T) -> T {
SHARD.with(|c| {
let mut slot = c.borrow_mut();
let mem = slot.get_or_insert_with(load_shard_indexed);
f(mem)
})
}
pub fn load(src: &str) -> Option<Program> {
let key = key_for(src);
let verify = verify_for(src);
let cp: CProg = with_shard(|m| {
let (v, blob) = m.get(key)?;
if v != verify {
return None;
}
bincode::deserialize(blob).ok()
})?;
let mut prog = Program {
main: cp.main,
functions: cp.functions,
tries: cp.tries,
strict: cp.strict,
source: Some(src.into()),
};
rehash(&mut prog);
crate::host::restore_site_tables(&cp.sites);
Some(prog)
}
fn rehash(prog: &mut Program) {
fn one(c: &mut Chunk) {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut h = DefaultHasher::new();
c.ops.hash(&mut h);
c.constants.hash(&mut h);
c.op_hash = h.finish();
}
one(&mut prog.main);
for (_, f) in &mut prog.functions {
one(&mut f.chunk);
}
for t in &mut prog.tries {
one(&mut t.block);
if let Some((_, h)) = &mut t.handler {
one(h);
}
if let Some(f) = &mut t.finalizer {
one(f);
}
}
}
pub fn store(src: &str, prog: &Program) -> Result<(), String> {
let cp = CProg {
main: prog.main.clone(),
functions: prog.functions.clone(),
tries: prog.tries.clone(),
sites: crate::host::site_tables(),
strict: prog.strict,
};
let blob = bincode::serialize(&cp).map_err(|e| format!("cache encode: {e}"))?;
let key = key_for(src);
let verify = verify_for(src);
with_shard(|m| {
m.added.insert(key, (verify, blob));
m.dirty = true;
});
Ok(())
}
pub fn flush() {
let build = build_id();
let pending = SHARD.with(|c| {
let mut slot = c.borrow_mut();
match slot.as_mut() {
Some(m) if m.dirty => {
m.dirty = false;
Some(
m.iter()
.map(|(k, v, b)| (k, (v, b.to_vec())))
.collect::<rustc_hash::FxHashMap<u64, (u64, Vec<u8>)>>(),
)
}
_ => None,
}
});
let Some(mut merged) = pending else { return };
let stamp = SHARD.with(|c| c.borrow().as_ref().and_then(|m| m.stamp));
if shard_stamp() != stamp {
for e in load_shard().entries {
if e.build == build {
merged.entry(e.key).or_insert((e.verify, e.blob));
}
}
}
let shard = Shard {
entries: merged
.into_iter()
.map(|(key, (verify, blob))| Entry {
key,
verify,
build,
blob,
})
.collect(),
};
let _ = write_shard(&shard);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_keys_depend_on_the_build_not_just_the_schema() {
use std::collections::hash_map::DefaultHasher;
let src = "console.log(1)\n";
let mut bare = rustc_hash::FxHasher::default();
SCHEMA.hash(&mut bare);
src.hash(&mut bare);
assert_ne!(
key_for(src),
bare.finish(),
"key_for must hash the build identity, not just SCHEMA"
);
let mut version_only = rustc_hash::FxHasher::default();
SCHEMA.hash(&mut version_only);
BUILD_VERSION.hash(&mut version_only);
src.hash(&mut version_only);
assert_ne!(
key_for(src),
version_only.finish(),
"key_for must hash the per-build id, so two dev builds of one \
version cannot share cached bytecode"
);
let mut bare = DefaultHasher::new();
SCHEMA.hash(&mut bare);
src.len().hash(&mut bare);
src.hash(&mut bare);
assert_ne!(
verify_for(src),
bare.finish(),
"verify_for must hash the build identity, not just SCHEMA"
);
let mut version_only = DefaultHasher::new();
SCHEMA.hash(&mut version_only);
BUILD_VERSION.hash(&mut version_only);
src.len().hash(&mut version_only);
src.hash(&mut version_only);
assert_ne!(
verify_for(src),
version_only.finish(),
"verify_for must hash the per-build id"
);
assert_eq!(BUILD_VERSION, env!("CARGO_PKG_VERSION"));
assert_ne!(build_id(), 0, "build_id must read the running binary");
assert_ne!(key_for(src), key_for("console.log(2)\n"));
}
}