refactor(tvix/eval): move attrset-related code to compiler::attrs r/4810

There's about to be a lot more code for attrsets (hopefully
temporarily as part of an expand&contract cycle), while nested
attribute logic is being refactored in preparation for recursive
attribute sets.

This does not change any functionality.

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

2 files changed, 209 insertions, 202 deletions

diff --git a/tvix/eval/src/compiler/attrs.rs b/tvix/eval/src/compiler/attrs.rs
new file mode 100644
index 0000000000..6ad31ab850
--- /dev/null
+++ b/tvix/eval/src/compiler/attrs.rs
@@ -0,0 +1,208 @@
+//! This module implements compiler logic related to attribute sets
+//! (construction, access operators, ...).
+
+use super::*;
+
+impl Compiler<'_, '_> {
+    pub(super) fn compile_attr(&mut self, slot: LocalIdx, node: ast::Attr) {
+        match node {
+            ast::Attr::Dynamic(dynamic) => {
+                self.compile(slot, dynamic.expr().unwrap());
+                self.emit_force(&dynamic.expr().unwrap());
+            }
+
+            ast::Attr::Str(s) => {
+                self.compile_str(slot, s.clone());
+                self.emit_force(&s);
+            }
+
+            ast::Attr::Ident(ident) => self.emit_literal_ident(&ident),
+        }
+    }
+
+    /// Compile attribute set literals into equivalent bytecode.
+    ///
+    /// This is complicated by a number of features specific to Nix
+    /// attribute sets, most importantly:
+    ///
+    /// 1. Keys can be dynamically constructed through interpolation.
+    /// 2. Keys can refer to nested attribute sets.
+    /// 3. Attribute sets can (optionally) be recursive.
+    pub(super) fn compile_attr_set(&mut self, slot: LocalIdx, node: ast::AttrSet) {
+        if node.rec_token().is_some() {
+            todo!("recursive attribute sets are not yet implemented")
+        }
+
+        // Open a scope to track the positions of the temporaries used
+        // by the `OpAttrs` instruction.
+        self.begin_scope();
+
+        let mut count = self.compile_inherit_attrs(slot, node.inherits());
+
+        for kv in node.attrpath_values() {
+            count += 1;
+
+            // Because attribute set literals can contain nested keys,
+            // there is potentially more than one key fragment. If
+            // this is the case, a special operation to construct a
+            // runtime value representing the attribute path is
+            // emitted.
+            let mut key_count = 0;
+            let key_span = self.span_for(&kv.attrpath().unwrap());
+            let key_idx = self.scope_mut().declare_phantom(key_span, false);
+
+            for fragment in kv.attrpath().unwrap().attrs() {
+                // Key fragments can contain dynamic expressions,
+                // which makes accounting for their stack slots very
+                // tricky.
+                //
+                // In order to ensure the locals are correctly cleaned
+                // up, we essentially treat any key fragment after the
+                // first one (which has a locals index that will
+                // become that of the final key itself) as being part
+                // of a separate scope, which results in the somewhat
+                // annoying setup logic below.
+                let fragment_slot = match key_count {
+                    0 => key_idx,
+                    1 => {
+                        self.begin_scope();
+                        self.scope_mut().declare_phantom(key_span, false)
+                    }
+                    _ => self.scope_mut().declare_phantom(key_span, false),
+                };
+
+                key_count += 1;
+                self.compile_attr(fragment_slot, fragment);
+                self.scope_mut().mark_initialised(fragment_slot);
+            }
+
+            // We're done with the key if there was only one fragment,
+            // otherwise we need to emit an instruction to construct
+            // the attribute path.
+            if key_count > 1 {
+                self.push_op(
+                    OpCode::OpAttrPath(Count(key_count)),
+                    &kv.attrpath().unwrap(),
+                );
+
+                // Close the temporary scope that was set up for the
+                // key fragments.
+                self.scope_mut().end_scope();
+            }
+
+            // The value is just compiled as normal so that its
+            // resulting value is on the stack when the attribute set
+            // is constructed at runtime.
+            let value_span = self.span_for(&kv.value().unwrap());
+            let value_slot = self.scope_mut().declare_phantom(value_span, false);
+            self.compile(slot, kv.value().unwrap());
+            self.scope_mut().mark_initialised(value_slot);
+        }
+
+        self.push_op(OpCode::OpAttrs(Count(count)), &node);
+
+        // Remove the temporary scope, but do not emit any additional
+        // cleanup (OpAttrs consumes all of these locals).
+        self.scope_mut().end_scope();
+    }
+
+    pub(super) fn compile_has_attr(&mut self, slot: LocalIdx, node: ast::HasAttr) {
+        // Put the attribute set on the stack.
+        self.compile(slot, node.expr().unwrap());
+
+        // Push all path fragments with an operation for fetching the
+        // next nested element, for all fragments except the last one.
+        for (count, fragment) in node.attrpath().unwrap().attrs().enumerate() {
+            if count > 0 {
+                self.push_op(OpCode::OpAttrsTrySelect, &fragment);
+            }
+
+            self.compile_attr(slot, fragment);
+        }
+
+        // After the last fragment, emit the actual instruction that
+        // leaves a boolean on the stack.
+        self.push_op(OpCode::OpAttrsIsSet, &node);
+    }
+
+    pub(super) fn compile_select(&mut self, slot: LocalIdx, node: ast::Select) {
+        let set = node.expr().unwrap();
+        let path = node.attrpath().unwrap();
+
+        if node.or_token().is_some() {
+            self.compile_select_or(slot, set, path, node.default_expr().unwrap());
+            return;
+        }
+
+        // Push the set onto the stack
+        self.compile(slot, set.clone());
+        self.emit_force(&set);
+
+        // Compile each key fragment and emit access instructions.
+        //
+        // TODO: multi-select instruction to avoid re-pushing attrs on
+        // nested selects.
+        for fragment in path.attrs() {
+            self.compile_attr(slot, fragment);
+            self.push_op(OpCode::OpAttrsSelect, &node);
+        }
+    }
+
+    /// Compile an `or` expression into a chunk of conditional jumps.
+    ///
+    /// If at any point during attribute set traversal a key is
+    /// missing, the `OpAttrOrNotFound` instruction will leave a
+    /// special sentinel value on the stack.
+    ///
+    /// After each access, a conditional jump evaluates the top of the
+    /// stack and short-circuits to the default value if it sees the
+    /// sentinel.
+    ///
+    /// Code like `{ a.b = 1; }.a.c or 42` yields this bytecode and
+    /// runtime stack:
+    ///
+    /// ```notrust
+    ///            Bytecode                     Runtime stack
+    ///  ┌────────────────────────────┐   ┌─────────────────────────┐
+    ///  │    ...                     │   │ ...                     │
+    ///  │ 5  OP_ATTRS(1)             │ → │ 5  [ { a.b = 1; }     ] │
+    ///  │ 6  OP_CONSTANT("a")        │ → │ 6  [ { a.b = 1; } "a" ] │
+    ///  │ 7  OP_ATTR_OR_NOT_FOUND    │ → │ 7  [ { b = 1; }       ] │
+    ///  │ 8  JUMP_IF_NOT_FOUND(13)   │ → │ 8  [ { b = 1; }       ] │
+    ///  │ 9  OP_CONSTANT("C")        │ → │ 9  [ { b = 1; } "c"   ] │
+    ///  │ 10 OP_ATTR_OR_NOT_FOUND    │ → │ 10 [ NOT_FOUND        ] │
+    ///  │ 11 JUMP_IF_NOT_FOUND(13)   │ → │ 11 [                  ] │
+    ///  │ 12 JUMP(14)                │   │ ..     jumped over      │
+    ///  │ 13 CONSTANT(42)            │ → │ 12 [ 42 ]               │
+    ///  │ 14 ...                     │   │ ..   ....               │
+    ///  └────────────────────────────┘   └─────────────────────────┘
+    /// ```
+    fn compile_select_or(
+        &mut self,
+        slot: LocalIdx,
+        set: ast::Expr,
+        path: ast::Attrpath,
+        default: ast::Expr,
+    ) {
+        self.compile(slot, set.clone());
+        self.emit_force(&set);
+        let mut jumps = vec![];
+
+        for fragment in path.attrs() {
+            self.compile_attr(slot, fragment.clone());
+            self.push_op(OpCode::OpAttrsTrySelect, &fragment);
+            jumps.push(self.push_op(OpCode::OpJumpIfNotFound(JumpOffset(0)), &fragment));
+        }
+
+        let final_jump = self.push_op(OpCode::OpJump(JumpOffset(0)), &path);
+
+        for jump in jumps {
+            self.patch_jump(jump);
+        }
+
+        // Compile the default value expression and patch the final
+        // jump to point *beyond* it.
+        self.compile(slot, default);
+        self.patch_jump(final_jump);
+    }
+}
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 79754e4de0..06170c1839 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -13,6 +13,7 @@
 //! by the code in this module, `debug_assert!` has been used to catch
 //! mistakes early during development.
 
+mod attrs;
 mod scope;
 
 use path_clean::PathClean;
@@ -404,41 +405,6 @@ impl Compiler<'_, '_> {
         self.push_op(OpCode::OpAssertBool, &node);
     }
 
-    fn compile_has_attr(&mut self, slot: LocalIdx, node: ast::HasAttr) {
-        // Put the attribute set on the stack.
-        self.compile(slot, node.expr().unwrap());
-
-        // Push all path fragments with an operation for fetching the
-        // next nested element, for all fragments except the last one.
-        for (count, fragment) in node.attrpath().unwrap().attrs().enumerate() {
-            if count > 0 {
-                self.push_op(OpCode::OpAttrsTrySelect, &fragment);
-            }
-
-            self.compile_attr(slot, fragment);
-        }
-
-        // After the last fragment, emit the actual instruction that
-        // leaves a boolean on the stack.
-        self.push_op(OpCode::OpAttrsIsSet, &node);
-    }
-
-    fn compile_attr(&mut self, slot: LocalIdx, node: ast::Attr) {
-        match node {
-            ast::Attr::Dynamic(dynamic) => {
-                self.compile(slot, dynamic.expr().unwrap());
-                self.emit_force(&dynamic.expr().unwrap());
-            }
-
-            ast::Attr::Str(s) => {
-                self.compile_str(slot, s.clone());
-                self.emit_force(&s);
-            }
-
-            ast::Attr::Ident(ident) => self.emit_literal_ident(&ident),
-        }
-    }
-
     /// Compile list literals into equivalent bytecode. List
     /// construction is fairly simple, consisting of pushing code for
     /// each literal element and an instruction with the element
@@ -536,173 +502,6 @@ impl Compiler<'_, '_> {
         count
     }
 
-    /// Compile attribute set literals into equivalent bytecode.
-    ///
-    /// This is complicated by a number of features specific to Nix
-    /// attribute sets, most importantly:
-    ///
-    /// 1. Keys can be dynamically constructed through interpolation.
-    /// 2. Keys can refer to nested attribute sets.
-    /// 3. Attribute sets can (optionally) be recursive.
-    fn compile_attr_set(&mut self, slot: LocalIdx, node: ast::AttrSet) {
-        if node.rec_token().is_some() {
-            todo!("recursive attribute sets are not yet implemented")
-        }
-
-        // Open a scope to track the positions of the temporaries used
-        // by the `OpAttrs` instruction.
-        self.begin_scope();
-
-        let mut count = self.compile_inherit_attrs(slot, node.inherits());
-
-        for kv in node.attrpath_values() {
-            count += 1;
-
-            // Because attribute set literals can contain nested keys,
-            // there is potentially more than one key fragment. If
-            // this is the case, a special operation to construct a
-            // runtime value representing the attribute path is
-            // emitted.
-            let mut key_count = 0;
-            let key_span = self.span_for(&kv.attrpath().unwrap());
-            let key_idx = self.scope_mut().declare_phantom(key_span, false);
-
-            for fragment in kv.attrpath().unwrap().attrs() {
-                // Key fragments can contain dynamic expressions,
-                // which makes accounting for their stack slots very
-                // tricky.
-                //
-                // In order to ensure the locals are correctly cleaned
-                // up, we essentially treat any key fragment after the
-                // first one (which has a locals index that will
-                // become that of the final key itself) as being part
-                // of a separate scope, which results in the somewhat
-                // annoying setup logic below.
-                let fragment_slot = match key_count {
-                    0 => key_idx,
-                    1 => {
-                        self.begin_scope();
-                        self.scope_mut().declare_phantom(key_span, false)
-                    }
-                    _ => self.scope_mut().declare_phantom(key_span, false),
-                };
-
-                key_count += 1;
-                self.compile_attr(fragment_slot, fragment);
-                self.scope_mut().mark_initialised(fragment_slot);
-            }
-
-            // We're done with the key if there was only one fragment,
-            // otherwise we need to emit an instruction to construct
-            // the attribute path.
-            if key_count > 1 {
-                self.push_op(
-                    OpCode::OpAttrPath(Count(key_count)),
-                    &kv.attrpath().unwrap(),
-                );
-
-                // Close the temporary scope that was set up for the
-                // key fragments.
-                self.scope_mut().end_scope();
-            }
-
-            // The value is just compiled as normal so that its
-            // resulting value is on the stack when the attribute set
-            // is constructed at runtime.
-            let value_span = self.span_for(&kv.value().unwrap());
-            let value_slot = self.scope_mut().declare_phantom(value_span, false);
-            self.compile(slot, kv.value().unwrap());
-            self.scope_mut().mark_initialised(value_slot);
-        }
-
-        self.push_op(OpCode::OpAttrs(Count(count)), &node);
-
-        // Remove the temporary scope, but do not emit any additional
-        // cleanup (OpAttrs consumes all of these locals).
-        self.scope_mut().end_scope();
-    }
-
-    fn compile_select(&mut self, slot: LocalIdx, node: ast::Select) {
-        let set = node.expr().unwrap();
-        let path = node.attrpath().unwrap();
-
-        if node.or_token().is_some() {
-            self.compile_select_or(slot, set, path, node.default_expr().unwrap());
-            return;
-        }
-
-        // Push the set onto the stack
-        self.compile(slot, set.clone());
-        self.emit_force(&set);
-
-        // Compile each key fragment and emit access instructions.
-        //
-        // TODO: multi-select instruction to avoid re-pushing attrs on
-        // nested selects.
-        for fragment in path.attrs() {
-            self.compile_attr(slot, fragment);
-            self.push_op(OpCode::OpAttrsSelect, &node);
-        }
-    }
-
-    /// Compile an `or` expression into a chunk of conditional jumps.
-    ///
-    /// If at any point during attribute set traversal a key is
-    /// missing, the `OpAttrOrNotFound` instruction will leave a
-    /// special sentinel value on the stack.
-    ///
-    /// After each access, a conditional jump evaluates the top of the
-    /// stack and short-circuits to the default value if it sees the
-    /// sentinel.
-    ///
-    /// Code like `{ a.b = 1; }.a.c or 42` yields this bytecode and
-    /// runtime stack:
-    ///
-    /// ```notrust
-    ///            Bytecode                     Runtime stack
-    ///  ┌────────────────────────────┐   ┌─────────────────────────┐
-    ///  │    ...                     │   │ ...                     │
-    ///  │ 5  OP_ATTRS(1)             │ → │ 5  [ { a.b = 1; }     ] │
-    ///  │ 6  OP_CONSTANT("a")        │ → │ 6  [ { a.b = 1; } "a" ] │
-    ///  │ 7  OP_ATTR_OR_NOT_FOUND    │ → │ 7  [ { b = 1; }       ] │
-    ///  │ 8  JUMP_IF_NOT_FOUND(13)   │ → │ 8  [ { b = 1; }       ] │
-    ///  │ 9  OP_CONSTANT("C")        │ → │ 9  [ { b = 1; } "c"   ] │
-    ///  │ 10 OP_ATTR_OR_NOT_FOUND    │ → │ 10 [ NOT_FOUND        ] │
-    ///  │ 11 JUMP_IF_NOT_FOUND(13)   │ → │ 11 [                  ] │
-    ///  │ 12 JUMP(14)                │   │ ..     jumped over      │
-    ///  │ 13 CONSTANT(42)            │ → │ 12 [ 42 ]               │
-    ///  │ 14 ...                     │   │ ..   ....               │
-    ///  └────────────────────────────┘   └─────────────────────────┘
-    /// ```
-    fn compile_select_or(
-        &mut self,
-        slot: LocalIdx,
-        set: ast::Expr,
-        path: ast::Attrpath,
-        default: ast::Expr,
-    ) {
-        self.compile(slot, set.clone());
-        self.emit_force(&set);
-        let mut jumps = vec![];
-
-        for fragment in path.attrs() {
-            self.compile_attr(slot, fragment.clone());
-            self.push_op(OpCode::OpAttrsTrySelect, &fragment);
-            jumps.push(self.push_op(OpCode::OpJumpIfNotFound(JumpOffset(0)), &fragment));
-        }
-
-        let final_jump = self.push_op(OpCode::OpJump(JumpOffset(0)), &path);
-
-        for jump in jumps {
-            self.patch_jump(jump);
-        }
-
-        // Compile the default value expression and patch the final
-        // jump to point *beyond* it.
-        self.compile(slot, default);
-        self.patch_jump(final_jump);
-    }
-
     fn compile_assert(&mut self, slot: LocalIdx, node: ast::Assert) {
         // Compile the assertion condition to leave its value on the stack.
         self.compile(slot, node.condition().unwrap());