fix(tvix/eval): declare let inherit (from) locals before compiling r/4805

The recent change that split declaration of let based locals and the
compilation of their values did not touch locals bound by inherit in
let. These were previously declared and compiled immediately before
starting to work on the other locals introduced in a let.

In the case of plain inherits, this behavior is kept in this change,
because there's nothing wrong with it: The value of a plain inherit will
always resolve to a higher scope, either statically or dynamically.

Since inherit (from) expression might refer to other locals bound in the
same let, we need to handle them in the same three steps as ordinary let
based locals:

1. We need to declare the (uninitialised) locals.

2. We need to compile the expression that obtains their value. For this,
   we create a new thunk, since the from expression may very well return
   a thunk which we need to force before selecting the value we are
   interested in.

3. Thunks need to be finalised.

For 1., we create an extra pass over the inherits that already declares
and initialises plain inherits and notes inherit (from) expressions in
the entries vector after declaring them. 2. only needs a bit of adapting
to create the thunks for selecting if appropriate, the rest of the
existing code can be reused.

Change-Id: Ie4ac1c0f9ffcbf7c07c452036aa8e577443af773
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6490
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>
Reviewed-by: tazjin <tazjin@tvl.su>

5 files changed, 69 insertions, 49 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 1ac5ae7fec..c64492d546 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -745,13 +745,38 @@ impl Compiler<'_, '_> {
         self.patch_jump(else_idx); // patch jump *over* else body
     }
 
-    /// Compile an `inherit` node of a `let`-expression.
-    fn compile_let_inherit<I: Iterator<Item = ast::Inherit>>(
-        &mut self,
-        slot: LocalIdx,
-        inherits: I,
-    ) {
-        for inherit in inherits {
+    /// Compile a standard `let ...; in ...` statement.
+    ///
+    /// Unless in a non-standard scope, the encountered values are
+    /// simply pushed on the stack and their indices noted in the
+    /// entries vector.
+    fn compile_let_in(&mut self, slot: LocalIdx, node: ast::LetIn) {
+        self.begin_scope();
+
+        // First pass to ensure that all identifiers are known;
+        // required for resolving recursion.
+        let mut entries: Vec<(LocalIdx, ast::Expr, Option<ast::Ident>)> = vec![];
+        for entry in node.attrpath_values() {
+            let mut path = match self.normalise_ident_path(entry.attrpath().unwrap().attrs()) {
+                Ok(p) => p,
+                Err(err) => {
+                    self.errors.push(err);
+                    continue;
+                }
+            };
+
+            if path.len() != 1 {
+                todo!("nested bindings in let expressions :(")
+            }
+
+            let idx = self.declare_local(&entry.attrpath().unwrap(), path.pop().unwrap());
+
+            entries.push((idx, entry.value().unwrap(), None));
+        }
+
+        // We also need to do such a pass for all inherits, we can even populate
+        // plain inherits directly, since they can't be (self) recursive.
+        for inherit in node.inherits() {
             match inherit.from() {
                 // Within a `let` binding, inheriting from the outer
                 // scope is a no-op *if* the identifier can be
@@ -783,55 +808,32 @@ impl Compiler<'_, '_> {
 
                 Some(from) => {
                     for ident in inherit.idents() {
-                        self.compile(slot, from.expr().unwrap());
-                        self.emit_force(&from.expr().unwrap());
-
-                        self.emit_literal_ident(&ident);
-                        self.push_op(OpCode::OpAttrsSelect, &ident);
                         let idx = self.declare_local(&ident, ident.ident_token().unwrap().text());
-                        self.scope_mut().mark_initialised(idx);
+                        entries.push((idx, from.expr().unwrap(), Some(ident)))
                     }
                 }
             }
         }
-    }
-
-    /// Compile a standard `let ...; in ...` statement.
-    ///
-    /// Unless in a non-standard scope, the encountered values are
-    /// simply pushed on the stack and their indices noted in the
-    /// entries vector.
-    fn compile_let_in(&mut self, slot: LocalIdx, node: ast::LetIn) {
-        self.begin_scope();
-
-        self.compile_let_inherit(slot, node.inherits());
-
-        // First pass to ensure that all identifiers are known;
-        // required for resolving recursion.
-        let mut entries: Vec<(LocalIdx, ast::Expr)> = vec![];
-        for entry in node.attrpath_values() {
-            let mut path = match self.normalise_ident_path(entry.attrpath().unwrap().attrs()) {
-                Ok(p) => p,
-                Err(err) => {
-                    self.errors.push(err);
-                    continue;
-                }
-            };
-
-            if path.len() != 1 {
-                todo!("nested bindings in let expressions :(")
-            }
-
-            let idx = self.declare_local(&entry.attrpath().unwrap(), path.pop().unwrap());
-
-            entries.push((idx, entry.value().unwrap()));
-        }
 
         // Second pass to place the values in the correct stack slots.
         let mut indices: Vec<LocalIdx> = vec![];
-        for (idx, value) in entries.into_iter() {
+        for (idx, value, path) in entries.into_iter() {
             indices.push(idx);
-            self.compile(idx, value);
+
+            // This entry is an inherit (from) expr, we need to select an attr
+            if let Some(ident) = path {
+                // Create a thunk wrapping value (which may be one as well) to
+                // avoid forcing the from expr too early.
+                self.thunk(idx, &value, move |c, n, s| {
+                    c.compile(s, n.clone());
+                    c.emit_force(n);
+
+                    c.emit_literal_ident(&ident);
+                    c.push_op(OpCode::OpAttrsSelect, &ident);
+                })
+            } else {
+                self.compile(idx, value);
+            }
 
             // Any code after this point will observe the value in the
             // right stack slot, so mark it as initialised.
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.exp
new file mode 100644
index 0000000000..409940768f
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.exp
@@ -0,0 +1 @@
+23
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.nix
new file mode 100644
index 0000000000..21196f48bc
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit-from-later-bound.nix
@@ -0,0 +1,13 @@
+let
+  inherit (c) d;
+  inherit (a) b c;
+
+  a = {
+    b = 20;
+    c = {
+      d = 3;
+    };
+  };
+in
+
+b + d
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.exp
index d00491fd7e..0cfbf08886 100644
--- a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.exp
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.exp
@@ -1 +1 @@
-1
+2
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.nix
index 12eed10e13..4ec270e3bf 100644
--- a/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.nix
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-let-inherit.nix
@@ -4,5 +4,9 @@ let
   };
 in
   let
+    set2 = {
+      b = 1;
+    };
     inherit (set) a;
-  in a
+    inherit (set2) b;
+  in a + b