feat(tvix/eval): crude caching builtins.import r/5366

Before this, tvix was spending most of its time furiously re-parsing
and re-compiling nixpkgs, each time hoping to get a different result...

Change-Id: I1c0cfbf9af622c276275b1f2fb8d4e976f1b5533
Signed-off-by: Adam Joseph <adam@westernsemico.com>
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7361
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

2 files changed, 20 insertions, 7 deletions

diff --git a/tvix/eval/src/builtins/impure.rs b/tvix/eval/src/builtins/impure.rs
index f28609fa763f..dc62e93ada40 100644
--- a/tvix/eval/src/builtins/impure.rs
+++ b/tvix/eval/src/builtins/impure.rs
@@ -122,6 +122,12 @@ pub fn builtins_import(globals: &Weak<GlobalsMap>, source: SourceCode) -> Builti
                 path.push("default.nix");
             }
 
+            let current_span = vm.current_span();
+            let entry = match vm.import_cache.entry(path.clone()) {
+                std::collections::btree_map::Entry::Occupied(oe) => return Ok(oe.get().clone()),
+                std::collections::btree_map::Entry::Vacant(ve) => ve,
+            };
+
             let contents =
                 std::fs::read_to_string(&path).map_err(|err| ErrorKind::ReadFileError {
                     path: path.clone(),
@@ -167,16 +173,20 @@ pub fn builtins_import(globals: &Weak<GlobalsMap>, source: SourceCode) -> Builti
                 });
             }
 
+            // Compilation succeeded, we can construct a thunk from whatever it spat
+            // out and return that.
+            let res = entry
+                .insert(Value::Thunk(Thunk::new_suspended(
+                    result.lambda,
+                    current_span,
+                )))
+                .clone();
+
             for warning in result.warnings {
                 vm.push_warning(warning);
             }
 
-            // Compilation succeeded, we can construct a thunk from whatever it spat
-            // out and return that.
-            Ok(Value::Thunk(Thunk::new_suspended(
-                result.lambda,
-                vm.current_span(),
-            )))
+            Ok(res)
         },
     )
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index f6605940a38c..b5f948a96a55 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -2,7 +2,7 @@
 //! Tvix bytecode.
 
 use serde_json::json;
-use std::{cmp::Ordering, ops::DerefMut, path::PathBuf, rc::Rc};
+use std::{cmp::Ordering, collections::BTreeMap, ops::DerefMut, path::PathBuf, rc::Rc};
 
 use crate::{
     chunk::Chunk,
@@ -60,6 +60,8 @@ pub struct VM<'o> {
     /// Runtime warnings collected during evaluation.
     warnings: Vec<EvalWarning>,
 
+    pub import_cache: Box<BTreeMap<PathBuf, Value>>,
+
     nix_search_path: NixSearchPath,
 
     observer: &'o mut dyn RuntimeObserver,
@@ -164,6 +166,7 @@ impl<'o> VM<'o> {
             stack: vec![],
             with_stack: vec![],
             warnings: vec![],
+            import_cache: Default::default(),
         }
     }