feat(tvix/eval): Allow passing in an env to evaluation r/8345

Allow passing in a top-level env, a map from name to value, to
evaluation. The intent is to support bound identifiers in the REPL just
like upstream nix does.

Getting this working involves mucking around a bit with internals - most
notably, locals now only optionally have a Span (since locals don't have
an easy span we can use) - and getting that working requires propagating
some minor hacks to places where we currently *need* a span (and which
would require too much changing now to make spans optional; my guess is
that that would essentially end up making spans optional throughout the
codebase).

Also, some extra care has to be taken to close out the scope in the case
that we do pass in an env, to avoid breaking our assumptions about the
size of the stack when we return from the toplevel

Change-Id: Ie475b2d3dfc72ccbf298d2a3ea28c63ac877d653
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11953
Tested-by: BuildkiteCI
Autosubmit: aspen <root@gws.fyi>
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 14 insertions, 4 deletions

diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index 398da4d6e22e..97690d13610e 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -36,6 +36,7 @@ mod test_utils;
 #[cfg(test)]
 mod tests;
 
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::rc::Rc;
 use std::str::FromStr;
@@ -56,6 +57,7 @@ pub use crate::value::{NixContext, NixContextElement};
 pub use crate::vm::generators;
 pub use crate::warnings::{EvalWarning, WarningKind};
 pub use builtin_macros;
+use smol_str::SmolStr;
 
 pub use crate::value::{Builtin, CoercionKind, NixAttrs, NixList, NixString, Value};
 
@@ -68,7 +70,7 @@ pub use crate::io::StdIO;
 ///
 /// Public fields are intended to be set by the caller. Setting all
 /// fields is optional.
-pub struct Evaluation<'co, 'ro, IO> {
+pub struct Evaluation<'co, 'ro, 'env, IO> {
     /// Source code map used for error reporting.
     source_map: SourceCode,
 
@@ -83,6 +85,9 @@ pub struct Evaluation<'co, 'ro, IO> {
     /// be compiled and inserted in the builtins set.
     pub src_builtins: Vec<(&'static str, &'static str)>,
 
+    /// Top-level variables to define in the evaluation
+    pub env: Option<&'env HashMap<SmolStr, Value>>,
+
     /// Implementation of file-IO to use during evaluation, e.g. for
     /// impure builtins.
     ///
@@ -131,7 +136,7 @@ pub struct EvaluationResult {
     pub expr: Option<rnix::ast::Expr>,
 }
 
-impl<'co, 'ro, IO> Evaluation<'co, 'ro, IO>
+impl<'co, 'ro, 'env, IO> Evaluation<'co, 'ro, 'env, IO>
 where
     IO: AsRef<dyn EvalIO> + 'static,
 {
@@ -146,6 +151,7 @@ where
             io_handle,
             builtins,
             src_builtins: vec![],
+            env: None,
             strict: false,
             nix_path: None,
             compiler_observer: None,
@@ -154,7 +160,7 @@ where
     }
 }
 
-impl<'co, 'ro> Evaluation<'co, 'ro, Box<dyn EvalIO>> {
+impl<'co, 'ro, 'env> Evaluation<'co, 'ro, 'env, Box<dyn EvalIO>> {
     /// Initialize an `Evaluation`, without the import statement available, and
     /// all IO operations stubbed out.
     pub fn new_pure() -> Self {
@@ -188,7 +194,7 @@ impl<'co, 'ro> Evaluation<'co, 'ro, Box<dyn EvalIO>> {
     }
 }
 
-impl<'co, 'ro, IO> Evaluation<'co, 'ro, IO>
+impl<'co, 'ro, 'env, IO> Evaluation<'co, 'ro, 'env, IO>
 where
     IO: AsRef<dyn EvalIO> + 'static,
 {
@@ -229,6 +235,7 @@ where
             source,
             self.builtins,
             self.src_builtins,
+            self.env,
             self.enable_import,
             compiler_observer,
         );
@@ -270,6 +277,7 @@ where
             source.clone(),
             self.builtins,
             self.src_builtins,
+            self.env,
             self.enable_import,
             compiler_observer,
         ) {
@@ -341,6 +349,7 @@ fn parse_compile_internal(
     source: SourceCode,
     builtins: Vec<(&'static str, Value)>,
     src_builtins: Vec<(&'static str, &'static str)>,
+    env: Option<&HashMap<SmolStr, Value>>,
     enable_import: bool,
     compiler_observer: &mut dyn CompilerObserver,
 ) -> Option<(Rc<Lambda>, Rc<GlobalsMap>)> {
@@ -368,6 +377,7 @@ fn parse_compile_internal(
         result.expr.as_ref().unwrap(),
         location,
         builtins,
+        env,
         &source,
         &file,
         compiler_observer,