feat(tvix/repl): Allow binding variables at the top-level r/8346

Allow binding variables at the REPL's toplevel in the same way the Nix
REPL does, using the syntax <ident> = <expr>. This fully, strictly
evaluates the value and sets it in the repl's "env", which gets passed
in at the toplevel when evaluating expressions.

The laziness behavior differs from Nix's, but I think this is good:

    ❯ nix repl
    Welcome to Nix version 2.3.18. Type :? for help.

    nix-repl> x = builtins.trace "x" 1

    nix-repl> x
    trace: x
    1

    nix-repl> x
    1

vs tvix:

    tvix-repl> x = builtins.trace "x" 1
    trace: "x" :: string
    tvix-repl> x
    => 1 :: int
    tvix-repl> x
    => 1 :: int

Bug: https://b.tvl.fyi/issues/371
Change-Id: Ieb2d626b7195fa87be638c9a4dae2eee45eb9ab1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11954
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Autosubmit: aspen <root@gws.fyi>

6 files changed, 157 insertions, 16 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 731731e72a1b..bbdeaeefb262 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -4367,7 +4367,9 @@ dependencies = [
  "dirs",
  "nix-compat",
  "rnix",
+ "rowan",
  "rustyline",
+ "smol_str",
  "thiserror",
  "tikv-jemallocator",
  "tokio",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 4d6e1d94628f..bf0958407907 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -13846,10 +13846,18 @@ rec {
             packageId = "rnix";
           }
           {
+            name = "rowan";
+            packageId = "rowan";
+          }
+          {
             name = "rustyline";
             packageId = "rustyline";
           }
           {
+            name = "smol_str";
+            packageId = "smol_str";
+          }
+          {
             name = "thiserror";
             packageId = "thiserror";
           }
diff --git a/tvix/cli/Cargo.toml b/tvix/cli/Cargo.toml
index 644393a7c4ed..f6d03ce5ab8f 100644
--- a/tvix/cli/Cargo.toml
+++ b/tvix/cli/Cargo.toml
@@ -20,6 +20,8 @@ clap = { version = "4.0", features = ["derive", "env"] }
 dirs = "4.0.0"
 rustyline = "10.0.0"
 rnix = "0.11.0"
+rowan = "*"
+smol_str = "0.2.0"
 thiserror = "1.0.38"
 tokio = "1.28.0"
 tracing = "0.1.40"
diff --git a/tvix/cli/src/assignment.rs b/tvix/cli/src/assignment.rs
new file mode 100644
index 000000000000..6fd9725d2956
--- /dev/null
+++ b/tvix/cli/src/assignment.rs
@@ -0,0 +1,74 @@
+use rnix::{Root, SyntaxKind, SyntaxNode};
+use rowan::ast::AstNode;
+
+/// An assignment of an identifier to a value in the context of a REPL.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct Assignment<'a> {
+    pub(crate) ident: &'a str,
+    pub(crate) value: rnix::ast::Expr,
+}
+
+impl<'a> Assignment<'a> {
+    /// Try to parse an [`Assignment`] from the given input string.
+    ///
+    /// Returns [`None`] if the parsing fails for any reason, since the intent is for us to
+    /// fall-back to trying to parse the input as a regular expression or other REPL commands for
+    /// any reason, since the intent is for us to fall-back to trying to parse the input as a
+    /// regular expression or other REPL command.
+    pub fn parse(input: &'a str) -> Option<Self> {
+        let mut tt = rnix::tokenizer::Tokenizer::new(input);
+        macro_rules! next {
+            ($kind:ident) => {{
+                loop {
+                    let (kind, tok) = tt.next()?;
+                    if kind == SyntaxKind::TOKEN_WHITESPACE {
+                        continue;
+                    }
+                    if kind != SyntaxKind::$kind {
+                        return None;
+                    }
+                    break tok;
+                }
+            }};
+        }
+
+        let ident = next!(TOKEN_IDENT);
+        let _equal = next!(TOKEN_ASSIGN);
+        let (green, errs) = rnix::parser::parse(tt);
+        let value = Root::cast(SyntaxNode::new_root(green))?.expr()?;
+
+        if !errs.is_empty() {
+            return None;
+        }
+
+        Some(Self { ident, value })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn simple_assignments() {
+        for input in ["x = 4", "x     =       \t\t\n\t4", "x=4"] {
+            let res = Assignment::parse(input).unwrap();
+            assert_eq!(res.ident, "x");
+            assert_eq!(res.value.to_string(), "4");
+        }
+    }
+
+    #[test]
+    fn complex_exprs() {
+        let input = "x = { y = 4; z = let q = 7; in [ q (y // { z = 9; }) ]; }";
+        let res = Assignment::parse(input).unwrap();
+        assert_eq!(res.ident, "x");
+    }
+
+    #[test]
+    fn not_an_assignment() {
+        let input = "{ x = 4; }";
+        let res = Assignment::parse(input);
+        assert!(res.is_none(), "{input:?}");
+    }
+}
diff --git a/tvix/cli/src/main.rs b/tvix/cli/src/main.rs
index 686513b77cc4..c4c70d2a1777 100644
--- a/tvix/cli/src/main.rs
+++ b/tvix/cli/src/main.rs
@@ -1,7 +1,10 @@
+mod assignment;
 mod repl;
 
 use clap::Parser;
 use repl::Repl;
+use smol_str::SmolStr;
+use std::collections::HashMap;
 use std::rc::Rc;
 use std::{fs, path::PathBuf};
 use tracing::{instrument, Level, Span};
@@ -150,18 +153,15 @@ impl AllowIncomplete {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 struct IncompleteInput;
 
-/// Interprets the given code snippet, printing out warnings, errors
-/// and the result itself. The return value indicates whether
-/// evaluation succeeded.
-#[instrument(skip_all, fields(indicatif.pb_show=1))]
-fn interpret(
+/// Interprets the given code snippet, printing out warnings and errors and returning the result
+fn evaluate(
     tvix_store_io: Rc<TvixStoreIO>,
     code: &str,
     path: Option<PathBuf>,
     args: &Args,
-    explain: bool,
     allow_incomplete: AllowIncomplete,
-) -> Result<bool, IncompleteInput> {
+    env: Option<&HashMap<SmolStr, Value>>,
+) -> Result<Option<Value>, IncompleteInput> {
     let span = Span::current();
     span.pb_start();
     span.pb_set_style(&tvix_tracing::PB_SPINNER_STYLE);
@@ -173,6 +173,9 @@ fn interpret(
     );
     eval.strict = args.strict;
     eval.builtins.extend(impure_builtins());
+    if let Some(env) = env {
+        eval.env = Some(env);
+    }
     add_derivation_builtins(&mut eval, Rc::clone(&tvix_store_io));
     add_fetcher_builtins(&mut eval, Rc::clone(&tvix_store_io));
     add_import_builtins(&mut eval, tvix_store_io);
@@ -226,7 +229,25 @@ fn interpret(
         }
     }
 
-    if let Some(value) = result.value.as_ref() {
+    Ok(result.value)
+}
+
+/// Interprets the given code snippet, printing out warnings, errors
+/// and the result itself. The return value indicates whether
+/// evaluation succeeded.
+#[instrument(skip_all, fields(indicatif.pb_show=1))]
+fn interpret(
+    tvix_store_io: Rc<TvixStoreIO>,
+    code: &str,
+    path: Option<PathBuf>,
+    args: &Args,
+    explain: bool,
+    allow_incomplete: AllowIncomplete,
+    env: Option<&HashMap<SmolStr, Value>>,
+) -> Result<bool, IncompleteInput> {
+    let result = evaluate(tvix_store_io, code, path, args, allow_incomplete, env)?;
+
+    if let Some(value) = result.as_ref() {
         if explain {
             println!("=> {}", value.explain());
         } else {
@@ -235,7 +256,7 @@ fn interpret(
     }
 
     // inform the caller about any errors
-    Ok(result.errors.is_empty())
+    Ok(result.is_some())
 }
 
 /// Interpret the given code snippet, but only run the Tvix compiler
@@ -298,6 +319,7 @@ fn main() {
             &args,
             false,
             AllowIncomplete::RequireComplete,
+            None, // TODO(aspen): Pass in --arg/--argstr here
         )
         .unwrap()
         {
@@ -325,6 +347,7 @@ fn run_file(io_handle: Rc<TvixStoreIO>, mut path: PathBuf, args: &Args) {
             args,
             false,
             AllowIncomplete::RequireComplete,
+            None,
         )
         .unwrap()
     };
diff --git a/tvix/cli/src/repl.rs b/tvix/cli/src/repl.rs
index 5a4830a027bc..758874016326 100644
--- a/tvix/cli/src/repl.rs
+++ b/tvix/cli/src/repl.rs
@@ -1,10 +1,13 @@
-use std::path::PathBuf;
 use std::rc::Rc;
+use std::{collections::HashMap, path::PathBuf};
 
 use rustyline::{error::ReadlineError, Editor};
+use smol_str::SmolStr;
+use tvix_eval::Value;
 use tvix_glue::tvix_store_io::TvixStoreIO;
 
-use crate::{interpret, AllowIncomplete, Args, IncompleteInput};
+use crate::evaluate;
+use crate::{assignment::Assignment, interpret, AllowIncomplete, Args, IncompleteInput};
 
 fn state_dir() -> Option<PathBuf> {
     let mut path = dirs::data_dir();
@@ -17,6 +20,7 @@ fn state_dir() -> Option<PathBuf> {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ReplCommand<'a> {
     Expr(&'a str),
+    Assign(Assignment<'a>),
     Explain(&'a str),
     Print(&'a str),
     Quit,
@@ -29,11 +33,12 @@ Welcome to the Tvix REPL!
 
 The following commands are supported:
 
-  <expr>    Evaluate a Nix language expression and print the result, along with its inferred type
-  :d <expr> Evaluate a Nix language expression and print a detailed description of the result
-  :p <expr> Evaluate a Nix language expression and print the result recursively
-  :q        Exit the REPL
-  :?, :h    Display this help text
+  <expr>       Evaluate a Nix language expression and print the result, along with its inferred type
+  <x> = <expr> Bind the result of an expression to a variable
+  :d <expr>    Evaluate a Nix language expression and print a detailed description of the result
+  :p <expr>    Evaluate a Nix language expression and print the result recursively
+  :q           Exit the REPL
+  :?, :h       Display this help text
 ";
 
     pub fn parse(input: &'a str) -> Self {
@@ -52,6 +57,10 @@ The following commands are supported:
             }
         }
 
+        if let Some(assignment) = Assignment::parse(input) {
+            return Self::Assign(assignment);
+        }
+
         Self::Expr(input)
     }
 }
@@ -61,6 +70,8 @@ pub struct Repl {
     /// In-progress multiline input, when the input so far doesn't parse as a complete expression
     multiline_input: Option<String>,
     rl: Editor<()>,
+    /// Local variables defined at the top-level in the repl
+    env: HashMap<SmolStr, Value>,
 }
 
 impl Repl {
@@ -69,6 +80,7 @@ impl Repl {
         Self {
             multiline_input: None,
             rl,
+            env: HashMap::new(),
         }
     }
 
@@ -125,7 +137,25 @@ impl Repl {
                             args,
                             false,
                             AllowIncomplete::Allow,
+                            Some(&self.env),
                         ),
+                        ReplCommand::Assign(Assignment { ident, value }) => {
+                            match evaluate(
+                                Rc::clone(&io_handle),
+                                &value.to_string(), /* FIXME: don't re-parse */
+                                None,
+                                args,
+                                AllowIncomplete::Allow,
+                                Some(&self.env),
+                            ) {
+                                Ok(Some(value)) => {
+                                    self.env.insert(ident.into(), value);
+                                    Ok(true)
+                                }
+                                Ok(None) => Ok(true),
+                                Err(incomplete) => Err(incomplete),
+                            }
+                        }
                         ReplCommand::Explain(input) => interpret(
                             Rc::clone(&io_handle),
                             input,
@@ -133,6 +163,7 @@ impl Repl {
                             args,
                             true,
                             AllowIncomplete::Allow,
+                            Some(&self.env),
                         ),
                         ReplCommand::Print(input) => interpret(
                             Rc::clone(&io_handle),
@@ -144,6 +175,7 @@ impl Repl {
                             },
                             false,
                             AllowIncomplete::Allow,
+                            Some(&self.env),
                         ),
                     };