feat(tvix/eval): support builtins implemented in Nix itself r/5735

This makes it possible to inject builtins into the builtin set that
are written in Nix code, and which at runtime are represented by a
thunk that will compile them the first time they are used.

Change-Id: Ia632367328f66fb2f26cb64ae464f8f3dc9c6d30
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7891
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

4 files changed, 108 insertions, 17 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 99e1d4c53c6c..ed0daa0264bd 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -1224,6 +1224,57 @@ fn optimise_tail_call(chunk: &mut Chunk) {
     }
 }
 
+/// Create a delayed source-only builtin compilation, for a builtin
+/// which is written in Nix code.
+///
+/// **Important:** tvix *panics* if a builtin with invalid source code
+/// is supplied. This is because there is no user-friendly way to
+/// thread the errors out of this function right now.
+fn compile_src_builtin(
+    name: &'static str,
+    code: &str,
+    source: &SourceCode,
+    weak: &Weak<GlobalsMap>,
+) -> Value {
+    use std::fmt::Write;
+
+    let parsed = rnix::ast::Root::parse(code);
+
+    if !parsed.errors().is_empty() {
+        let mut out = format!("BUG: code for source-builtin '{}' had parser errors", name);
+        for error in parsed.errors() {
+            writeln!(out, "{}", error).unwrap();
+        }
+
+        panic!("{}", out);
+    }
+
+    let file = source.add_file(format!("<src-builtins/{}.nix>", name), code.to_string());
+    let weak = weak.clone();
+
+    Value::Thunk(Thunk::new_suspended_native(Rc::new(move |_| {
+        let result = compile(
+            &parsed.tree().expr().unwrap(),
+            None,
+            file.clone(),
+            weak.upgrade().unwrap(),
+            &mut crate::observer::NoOpObserver {},
+        )?;
+
+        if !result.errors.is_empty() {
+            return Err(ErrorKind::ImportCompilerError {
+                path: format!("src-builtins/{}.nix", name).into(),
+                errors: result.errors,
+            });
+        }
+
+        Ok(Value::Thunk(Thunk::new_suspended(
+            result.lambda,
+            LightSpan::Actual { span: file.span },
+        )))
+    })))
+}
+
 /// Prepare the full set of globals available in evaluated code. These
 /// are constructed from the set of builtins supplied by the caller,
 /// which are made available globally under the `builtins` identifier.
@@ -1234,6 +1285,7 @@ fn optimise_tail_call(chunk: &mut Chunk) {
 /// Optionally adds the `import` feature if desired by the caller.
 pub fn prepare_globals(
     builtins: Vec<(&'static str, Value)>,
+    src_builtins: Vec<(&'static str, &'static str)>,
     source: SourceCode,
     enable_import: bool,
 ) -> Rc<GlobalsMap> {
@@ -1251,17 +1303,10 @@ pub fn prepare_globals(
             builtins.insert("import", import);
         }
 
-        // Next, the actual map of globals is constructed and
-        // populated with (copies) of the values that should be
-        // available in the global scope (see [`GLOBAL_BUILTINS`]).
+        // Next, the actual map of globals which the compiler will use
+        // to resolve identifiers is constructed.
         let mut globals: GlobalsMap = HashMap::new();
 
-        for global in GLOBAL_BUILTINS {
-            if let Some(builtin) = builtins.get(global).cloned() {
-                globals.insert(global, builtin);
-            }
-        }
-
         // builtins contain themselves (`builtins.builtins`), which we
         // can resolve by manually constructing a suspended thunk that
         // dereferences the same weak pointer as above.
@@ -1278,17 +1323,32 @@ pub fn prepare_globals(
             }))),
         );
 
-        // This is followed by the actual `builtins` attribute set
-        // being constructed and inserted in the global scope.
+        // Insert top-level static value builtins.
+        globals.insert("true", Value::Bool(true));
+        globals.insert("false", Value::Bool(false));
+        globals.insert("null", Value::Null);
+
+        // If "source builtins" were supplied, compile them and insert
+        // them.
+        builtins.extend(src_builtins.into_iter().map(move |(name, code)| {
+            let compiled = compile_src_builtin(name, code, &source, &weak);
+            (name, compiled)
+        }));
+
+        // Construct the actual `builtins` attribute set and insert it
+        // in the global scope.
         globals.insert(
             "builtins",
-            Value::attrs(NixAttrs::from_iter(builtins.into_iter())),
+            Value::attrs(NixAttrs::from_iter(builtins.clone().into_iter())),
         );
 
-        // Finally insert the compiler-internal "magic" builtins for top-level values.
-        globals.insert("true", Value::Bool(true));
-        globals.insert("false", Value::Bool(false));
-        globals.insert("null", Value::Null);
+        // Finally, the builtins that should be globally available are
+        // "elevated" to the outer scope.
+        for global in GLOBAL_BUILTINS {
+            if let Some(builtin) = builtins.get(global).cloned() {
+                globals.insert(global, builtin);
+            }
+        }
 
         globals
     }))
diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index 0e0f12059188..f903a078ff3b 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -91,6 +91,10 @@ pub struct Evaluation<'code, 'co, 'ro> {
     /// the set of impure builtins, or other custom builtins.
     pub builtins: Vec<(&'static str, Value)>,
 
+    /// Set of builtins that are implemented in Nix itself and should
+    /// be compiled and inserted in the builtins set.
+    pub src_builtins: Vec<(&'static str, &'static str)>,
+
     /// Implementation of file-IO to use during evaluation, e.g. for
     /// impure builtins.
     ///
@@ -156,6 +160,7 @@ impl<'code, 'co, 'ro> Evaluation<'code, 'co, 'ro> {
             source_map,
             file,
             builtins,
+            src_builtins: vec![],
             io_handle: Box::new(DummyIO {}),
             enable_import: false,
             nix_path: None,
@@ -198,6 +203,7 @@ impl<'code, 'co, 'ro> Evaluation<'code, 'co, 'ro> {
             self.location,
             source,
             self.builtins,
+            self.src_builtins,
             self.enable_import,
             compiler_observer,
         );
@@ -220,6 +226,7 @@ impl<'code, 'co, 'ro> Evaluation<'code, 'co, 'ro> {
             self.location,
             source,
             self.builtins,
+            self.src_builtins,
             self.enable_import,
             compiler_observer,
         ) {
@@ -271,6 +278,7 @@ fn parse_compile_internal(
     location: Option<PathBuf>,
     source: SourceCode,
     builtins: Vec<(&'static str, Value)>,
+    src_builtins: Vec<(&'static str, &'static str)>,
     enable_import: bool,
     compiler_observer: &mut dyn CompilerObserver,
 ) -> Option<(Rc<Lambda>, Rc<GlobalsMap>)> {
@@ -290,7 +298,7 @@ fn parse_compile_internal(
     // the result, in case the caller needs it for something.
     result.expr = parsed.tree().expr();
 
-    let builtins = crate::compiler::prepare_globals(builtins, source, enable_import);
+    let builtins = crate::compiler::prepare_globals(builtins, src_builtins, source, enable_import);
 
     let compiler_result = match compiler::compile(
         result.expr.as_ref().unwrap(),
diff --git a/tvix/eval/src/tests/mod.rs b/tvix/eval/src/tests/mod.rs
index 98c93ba27084..aeec75b2ae9d 100644
--- a/tvix/eval/src/tests/mod.rs
+++ b/tvix/eval/src/tests/mod.rs
@@ -2,6 +2,10 @@ use builtin_macros::builtins;
 use pretty_assertions::assert_eq;
 use test_generator::test_resources;
 
+/// Module for one-off tests which do not follow the rest of the
+/// test layout.
+mod one_offs;
+
 #[builtins]
 mod mock_builtins {
     //! Builtins which are required by language tests, but should not
diff --git a/tvix/eval/src/tests/one_offs.rs b/tvix/eval/src/tests/one_offs.rs
new file mode 100644
index 000000000000..63bb8f7af378
--- /dev/null
+++ b/tvix/eval/src/tests/one_offs.rs
@@ -0,0 +1,19 @@
+use crate::*;
+
+#[test]
+fn test_source_builtin() {
+    // Test an evaluation with a source-only builtin. The test ensures
+    // that the artificially constructed thunking is correct.
+
+    let mut eval = Evaluation::new_impure("builtins.testSourceBuiltin", None);
+    eval.src_builtins.push(("testSourceBuiltin", "42"));
+
+    let result = eval.evaluate();
+    assert!(
+        result.errors.is_empty(),
+        "evaluation failed: {:?}",
+        result.errors
+    );
+
+    assert!(matches!(result.value.unwrap(), Value::Integer(42)));
+}