fix(tvix/repl): Share globals and sourcemap across evaluations r/8355

Now that we can bind (potentially lazy, potentially lambda-containing)
values in the REPL and then reference them in subsequent evaluations,
it's important that the values to which we construct shared references
are shared across those subsequent evaluations - otherwise, we get
panics due to unknown source map locations, or dropped weak references
to globals.

This change assigns both the globals and the source map as fields on the
Repl after the first evaluation, and then passes those in (to the
EvaluationBuilder) on subsequent evaluations.

On the EvaluationBuilder side, there's some panicking introduced - this
is intentional, as my intent is for the builder to be configured
statically enough that panicking is the best way to report errors
here (it's always a bug to misconfigure an Evaluation, and we'd never
want to handle it dynamically).

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

3 files changed, 84 insertions, 21 deletions

diff --git a/tvix/cli/src/lib.rs b/tvix/cli/src/lib.rs
index 008593c5d84f..800ffb4e0e75 100644
--- a/tvix/cli/src/lib.rs
+++ b/tvix/cli/src/lib.rs
@@ -8,7 +8,7 @@ use tvix_build::buildservice;
 use tvix_eval::{
     builtins::impure_builtins,
     observer::{DisassemblingObserver, TracingObserver},
-    ErrorKind, EvalIO, Value,
+    ErrorKind, EvalIO, GlobalsMap, SourceCode, Value,
 };
 use tvix_glue::{
     builtins::{add_derivation_builtins, add_fetcher_builtins, add_import_builtins},
@@ -83,7 +83,13 @@ impl AllowIncomplete {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct IncompleteInput;
 
+pub struct EvalResult {
+    value: Option<Value>,
+    globals: Rc<GlobalsMap>,
+}
+
 /// Interprets the given code snippet, printing out warnings and errors and returning the result
+#[allow(clippy::too_many_arguments)]
 pub fn evaluate(
     tvix_store_io: Rc<TvixStoreIO>,
     code: &str,
@@ -91,7 +97,9 @@ pub fn evaluate(
     args: &Args,
     allow_incomplete: AllowIncomplete,
     env: Option<&HashMap<SmolStr, Value>>,
-) -> Result<Option<Value>, IncompleteInput> {
+    globals: Option<Rc<GlobalsMap>>,
+    source_map: Option<SourceCode>,
+) -> Result<EvalResult, IncompleteInput> {
     let span = Span::current();
     span.pb_start();
     span.pb_set_style(&tvix_tracing::PB_SPINNER_STYLE);
@@ -102,16 +110,27 @@ pub fn evaluate(
     )) as Box<dyn EvalIO>)
     .enable_import()
     .with_strict(args.strict)
-    .add_builtins(impure_builtins())
     .env(env);
 
-    eval_builder = add_derivation_builtins(eval_builder, Rc::clone(&tvix_store_io));
-    eval_builder = add_fetcher_builtins(eval_builder, Rc::clone(&tvix_store_io));
-    eval_builder = add_import_builtins(eval_builder, tvix_store_io);
-    eval_builder = configure_nix_path(eval_builder, &args.nix_search_path);
+    match globals {
+        Some(globals) => {
+            eval_builder = eval_builder.with_globals(globals);
+        }
+        None => {
+            eval_builder = eval_builder.add_builtins(impure_builtins());
+            eval_builder = add_derivation_builtins(eval_builder, Rc::clone(&tvix_store_io));
+            eval_builder = add_fetcher_builtins(eval_builder, Rc::clone(&tvix_store_io));
+            eval_builder = add_import_builtins(eval_builder, tvix_store_io);
+            eval_builder = configure_nix_path(eval_builder, &args.nix_search_path);
+        }
+    };
+
+    if let Some(source_map) = source_map {
+        eval_builder = eval_builder.with_source_map(source_map);
+    }
 
     let source_map = eval_builder.source_map().clone();
-    let result = {
+    let (result, globals) = {
         let mut compiler_observer =
             DisassemblingObserver::new(source_map.clone(), std::io::stderr());
         if args.dump_bytecode {
@@ -129,7 +148,9 @@ pub fn evaluate(
         span.pb_set_message("Evaluating…");
 
         let eval = eval_builder.build();
-        eval.evaluate(code, path)
+        let globals = eval.globals();
+        let result = eval.evaluate(code, path);
+        (result, globals)
     };
 
     if allow_incomplete.allow()
@@ -160,19 +181,24 @@ pub fn evaluate(
         }
     }
 
-    Ok(result.value)
+    Ok(EvalResult {
+        globals,
+        value: result.value,
+    })
 }
 
 pub struct InterpretResult {
     output: String,
     success: bool,
+    pub(crate) globals: Option<Rc<GlobalsMap>>,
 }
 
 impl InterpretResult {
-    pub fn empty_success() -> Self {
+    pub fn empty_success(globals: Option<Rc<GlobalsMap>>) -> Self {
         Self {
             output: String::new(),
             success: true,
+            globals,
         }
     }
 
@@ -194,6 +220,7 @@ impl InterpretResult {
 /// and the result itself. The return value indicates whether
 /// evaluation succeeded.
 #[instrument(skip_all, fields(indicatif.pb_show=1))]
+#[allow(clippy::too_many_arguments)]
 pub fn interpret(
     tvix_store_io: Rc<TvixStoreIO>,
     code: &str,
@@ -202,11 +229,22 @@ pub fn interpret(
     explain: bool,
     allow_incomplete: AllowIncomplete,
     env: Option<&HashMap<SmolStr, Value>>,
+    globals: Option<Rc<GlobalsMap>>,
+    source_map: Option<SourceCode>,
 ) -> Result<InterpretResult, IncompleteInput> {
     let mut output = String::new();
-    let result = evaluate(tvix_store_io, code, path, args, allow_incomplete, env)?;
+    let result = evaluate(
+        tvix_store_io,
+        code,
+        path,
+        args,
+        allow_incomplete,
+        env,
+        globals,
+        source_map,
+    )?;
 
-    if let Some(value) = result.as_ref() {
+    if let Some(value) = result.value.as_ref() {
         if explain {
             writeln!(&mut output, "=> {}", value.explain()).unwrap();
         } else if args.raw {
@@ -219,6 +257,7 @@ pub fn interpret(
     // inform the caller about any errors
     Ok(InterpretResult {
         output,
-        success: result.is_some(),
+        success: result.value.is_some(),
+        globals: Some(result.globals),
     })
 }
diff --git a/tvix/cli/src/main.rs b/tvix/cli/src/main.rs
index 0bd3be37eeb2..f927665aeb76 100644
--- a/tvix/cli/src/main.rs
+++ b/tvix/cli/src/main.rs
@@ -75,6 +75,8 @@ fn main() {
             false,
             AllowIncomplete::RequireComplete,
             None, // TODO(aspen): Pass in --arg/--argstr here
+            None,
+            None,
         )
         .unwrap()
         .finalize()
@@ -104,6 +106,8 @@ fn run_file(io_handle: Rc<TvixStoreIO>, mut path: PathBuf, args: &Args) {
             false,
             AllowIncomplete::RequireComplete,
             None,
+            None,
+            None,
         )
         .unwrap()
         .finalize()
diff --git a/tvix/cli/src/repl.rs b/tvix/cli/src/repl.rs
index 5098fbaeedc3..6b34b6552da0 100644
--- a/tvix/cli/src/repl.rs
+++ b/tvix/cli/src/repl.rs
@@ -3,7 +3,7 @@ use std::{collections::HashMap, path::PathBuf};
 
 use rustyline::{error::ReadlineError, Editor};
 use smol_str::SmolStr;
-use tvix_eval::Value;
+use tvix_eval::{GlobalsMap, SourceCode, Value};
 use tvix_glue::tvix_store_io::TvixStoreIO;
 
 use crate::{
@@ -92,6 +92,8 @@ pub struct Repl<'a> {
 
     io_handle: Rc<TvixStoreIO>,
     args: &'a Args,
+    source_map: SourceCode,
+    globals: Option<Rc<GlobalsMap>>,
 }
 
 impl<'a> Repl<'a> {
@@ -103,6 +105,8 @@ impl<'a> Repl<'a> {
             env: HashMap::new(),
             io_handle,
             args,
+            source_map: Default::default(),
+            globals: None,
         }
     }
 
@@ -179,7 +183,7 @@ impl<'a> Repl<'a> {
             }
             ReplCommand::Help => {
                 println!("{}", ReplCommand::HELP);
-                Ok(InterpretResult::empty_success())
+                Ok(InterpretResult::empty_success(None))
             }
             ReplCommand::Expr(input) => interpret(
                 Rc::clone(&self.io_handle),
@@ -189,6 +193,8 @@ impl<'a> Repl<'a> {
                 false,
                 AllowIncomplete::Allow,
                 Some(&self.env),
+                self.globals.clone(),
+                Some(self.source_map.clone()),
             ),
             ReplCommand::Assign(Assignment { ident, value }) => {
                 match evaluate(
@@ -198,12 +204,15 @@ impl<'a> Repl<'a> {
                     self.args,
                     AllowIncomplete::Allow,
                     Some(&self.env),
+                    self.globals.clone(),
+                    Some(self.source_map.clone()),
                 ) {
-                    Ok(Some(value)) => {
-                        self.env.insert(ident.into(), value);
-                        Ok(InterpretResult::empty_success())
+                    Ok(result) => {
+                        if let Some(value) = result.value {
+                            self.env.insert(ident.into(), value);
+                        }
+                        Ok(InterpretResult::empty_success(Some(result.globals)))
                     }
-                    Ok(None) => Ok(InterpretResult::empty_success()),
                     Err(incomplete) => Err(incomplete),
                 }
             }
@@ -215,6 +224,8 @@ impl<'a> Repl<'a> {
                 true,
                 AllowIncomplete::Allow,
                 Some(&self.env),
+                self.globals.clone(),
+                Some(self.source_map.clone()),
             ),
             ReplCommand::Print(input) => interpret(
                 Rc::clone(&self.io_handle),
@@ -227,13 +238,22 @@ impl<'a> Repl<'a> {
                 false,
                 AllowIncomplete::Allow,
                 Some(&self.env),
+                self.globals.clone(),
+                Some(self.source_map.clone()),
             ),
         };
 
         match res {
-            Ok(InterpretResult { output, .. }) => {
+            Ok(InterpretResult {
+                output,
+                globals,
+                success: _,
+            }) => {
                 self.rl.add_history_entry(input);
                 self.multiline_input = None;
+                if globals.is_some() {
+                    self.globals = globals;
+                }
                 CommandResult {
                     output,
                     continue_: true,