feat: lazily import TLAs

CertainLach · Nov 1, 2023 · 2ceb920 · 2ceb920
1 parent 3baf492
commit 2ceb920
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cmds/jrsonnet/src/main.rs b/cmds/jrsonnet/src/main.rs
@@ -9,7 +9,7 @@ use jrsonnet_cli::{GcOpts, ManifestOpts, MiscOpts, OutputOpts, StdOpts, TlaOpts,
 use jrsonnet_evaluator::{
 	apply_tla, bail,
 	error::{Error as JrError, ErrorKind},
-	ResultExt, State, Val,
+	ImportResolver, ResultExt, State, Val,
 };
 
 #[cfg(feature = "mimalloc")]
@@ -186,13 +186,7 @@ fn main_real(s: &State, opts: Opts) -> Result<(), Error> {
 		s.import(&input)?
 	};
 
-	let (tla, tla_str_paths, tla_code_paths) = opts.tla.tla_opts()?;
-	for path in tla_str_paths {
-		s.import_str(path)?;
-	}
-	for path in tla_code_paths {
-		s.import(path)?;
-	}
+	let tla = opts.tla.into_args_in(s)?;
 	#[allow(unused_mut)]
 	let mut val = apply_tla(s.clone(), &tla, val)?;
 

diff --git a/crates/jrsonnet-cli/src/tla.rs b/crates/jrsonnet-cli/src/tla.rs
@@ -1,12 +1,15 @@
+use std::path::PathBuf;
+
 use clap::Parser;
 use jrsonnet_evaluator::{
 	error::{ErrorKind, Result},
 	function::TlaArg,
 	gc::GcHashMap,
-	IStr,
+	val::ThunkValue,
+	IStr, State, Thunk, Val,
 };
+use jrsonnet_gcmodule::Trace;
 use jrsonnet_parser::{ParserSettings, Source};
-use std::path::PathBuf;
 
 use crate::ExtStr;
 
@@ -33,15 +36,24 @@ pub struct TlaOpts {
 	tla_code_file: Vec<PathBuf>,
 }
 impl TlaOpts {
-	pub fn tla_opts(&self) -> Result<(GcHashMap<IStr, TlaArg>, &Vec<PathBuf>, &Vec<PathBuf>)> {
+	pub fn into_args_in(self, state: &State) -> Result<GcHashMap<IStr, TlaArg>> {
 		let mut out = GcHashMap::new();
 		for (name, value) in self.tla_str.iter().map(|c| (&c.name, &c.value)) {
 			out.insert(name.into(), TlaArg::String(value.into()));
 		}
+		for path in self.tla_str_file {
+			out.insert(
+				path.to_string_lossy().into(),
+				TlaArg::Lazy(Thunk::new(ImportStrThunk {
+					state: state.clone(),
+					path,
+				})),
+			);
+		}
 		for (name, code) in self.tla_code.iter().map(|c| (&c.name, &c.value)) {
 			let source = Source::new_virtual(format!("<top-level-arg:{name}>").into(), code.into());
 			out.insert(
-				(name as &str).into(),
+				name.as_str().into(),
 				TlaArg::Code(
 					jrsonnet_parser::parse(
 						code,
@@ -56,6 +68,41 @@ impl TlaOpts {
 				),
 			);
 		}
-		Ok((out, &self.tla_str_file, &self.tla_code_file))
+		for path in self.tla_code_file {
+			out.insert(
+				path.to_string_lossy().into(),
+				TlaArg::Lazy(Thunk::new(ImportCodeThunk {
+					state: state.clone(),
+					path,
+				})),
+			);
+		}
+		Ok(out)
+	}
+}
+
+#[derive(Trace)]
+struct ImportStrThunk {
+	path: PathBuf,
+	state: State,
+}
+impl ThunkValue for ImportStrThunk {
+	type Output = Val;
+
+	fn get(self: Box<Self>) -> Result<Self::Output> {
+		self.state.import_str(self.path).map(|s| Val::Str(s.into()))
+	}
+}
+
+#[derive(Trace)]
+struct ImportCodeThunk {
+	path: PathBuf,
+	state: State,
+}
+impl ThunkValue for ImportCodeThunk {
+	type Output = Val;
+
+	fn get(self: Box<Self>) -> Result<Self::Output> {
+		self.state.import(self.path)
 	}
 }
diff --git a/crates/jrsonnet-evaluator/src/evaluate/mod.rs b/crates/jrsonnet-evaluator/src/evaluate/mod.rs
@@ -667,7 +667,7 @@ pub fn evaluate(ctx: Context, expr: &LocExpr) -> Result<Val> {
 			IndexableVal::into_untyped(indexable.into_indexable()?.slice(start, end, step)?)?
 		}
 		i @ (Import(path) | ImportStr(path) | ImportBin(path)) => {
-			let Expr::Str(path) = &*path.0 else {
+			let Str(path) = &*path.0 else {
 				bail!("computed imports are not supported")
 			};
 			let tmp = loc.clone().0;

diff --git a/crates/jrsonnet-evaluator/src/lib.rs b/crates/jrsonnet-evaluator/src/lib.rs
@@ -346,7 +346,7 @@ impl State {
 		let file_name = Source::new(path.clone(), code.clone());
 		if file.parsed.is_none() {
 			file.parsed = Some(
-				jrsonnet_parser::parse(
+				parse(
 					&code,
 					&ParserSettings {
 						source: file_name.clone(),
@@ -522,7 +522,7 @@ impl State {
 	pub fn evaluate_snippet(&self, name: impl Into<IStr>, code: impl Into<IStr>) -> Result<Val> {
 		let code = code.into();
 		let source = Source::new_virtual(name.into(), code.clone());
-		let parsed = jrsonnet_parser::parse(
+		let parsed = parse(
 			&code,
 			&ParserSettings {
 				source: source.clone(),
@@ -543,7 +543,7 @@ impl State {
 	) -> Result<Val> {
 		let code = code.into();
 		let source = Source::new_virtual(name.into(), code.clone());
-		let parsed = jrsonnet_parser::parse(
+		let parsed = parse(
 			&code,
 			&ParserSettings {
 				source: source.clone(),