refactor(tvix/eval): bye compiler::attrs, hello compiler::bindings r/4972

Changes the module structure of the compiler to have a module
dedicated to the logic of setting up bindings. This logic is in the
process of being merged between attribute sets and `let`-expressions,
and the structure of the modules makes more sense when ecapsulating
that specifically.

(Other bits of code related to e.g. attribute sets are pretty
straightforward and can just live in the main compiler module).

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

1 files changed, 121 insertions, 438 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 21f131e4543a..de614be4dad6 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -13,7 +13,7 @@
 //! by the code in this module, `debug_assert!` has been used to catch
 //! mistakes early during development.
 
-mod attrs;
+mod bindings;
 mod scope;
 mod spans;
 
@@ -507,6 +507,126 @@ impl Compiler<'_> {
         self.scope_mut().end_scope();
     }
 
+    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),
+        }
+    }
+
+    fn compile_has_attr(&mut self, slot: LocalIdx, node: ast::HasAttr) {
+        // Put the attribute set on the stack.
+        self.compile(slot, node.expr().unwrap());
+        self.emit_force(&node);
+
+        // 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.emit_force(&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::OpHasAttr, &node);
+    }
+
+    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());
+
+        // 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() {
+            // Force the current set value.
+            self.emit_force(&fragment);
+
+            self.compile_attr(slot, fragment.clone());
+            self.push_op(OpCode::OpAttrsSelect, &fragment);
+        }
+    }
+
+    /// 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());
+        let mut jumps = vec![];
+
+        for fragment in path.attrs() {
+            self.emit_force(&fragment);
+            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());
@@ -552,311 +672,6 @@ impl Compiler<'_> {
         self.patch_jump(else_idx); // patch jump *over* else body
     }
 
-    fn compile_recursive_scope<N>(&mut self, slot: LocalIdx, rec_attrs: bool, node: &N) -> usize
-    where
-        N: ToSpan + ast::HasEntry,
-    {
-        let mut count = 0;
-        self.scope_mut().begin_scope();
-
-        // First pass to find all plain inherits (if they are not useless).
-        // Since they always resolve to a higher scope, we can just compile and
-        // declare them immediately. This needs to happen *before* we declare
-        // any other locals in the scope or the stack order gets messed up.
-        // While we are looping through the inherits, already note all inherit
-        // (from) expressions, that may very well resolve recursively and need
-        // to be compiled like normal let in bindings.
-        let mut inherit_froms: Vec<(ast::Expr, String, Span)> = vec![];
-        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
-                // statically resolved.
-                None if !rec_attrs && !self.scope().has_with() => {
-                    self.emit_warning(&inherit, WarningKind::UselessInherit);
-                    continue;
-                }
-
-                None => {
-                    for attr in inherit.attrs() {
-                        let name = match self.expr_static_attr_str(&attr) {
-                            Some(name) => name,
-                            None => {
-                                // TODO(tazjin): error variant for dynamic
-                                // key in *inherit* (or generalise it)
-                                self.emit_error(&attr, ErrorKind::DynamicKeyInLet);
-                                continue;
-                            }
-                        };
-
-                        count += 1;
-
-                        // If the identifier resolves statically in a
-                        // `let`, it has precedence over dynamic
-                        // bindings, and the inherit is useless.
-                        if !rec_attrs
-                            && matches!(
-                                self.scope_mut().resolve_local(&name),
-                                LocalPosition::Known(_)
-                            )
-                        {
-                            self.emit_warning(&attr, WarningKind::UselessInherit);
-                            continue;
-                        }
-
-                        if rec_attrs {
-                            self.emit_constant(Value::String(SmolStr::new(&name).into()), &attr);
-                            let span = self.span_for(&attr);
-                            self.scope_mut().declare_phantom(span, true);
-                        }
-
-                        self.compile_identifier_access(slot, &name, &attr);
-                        let idx = self.declare_local(&attr, &name);
-                        self.scope_mut().mark_initialised(idx);
-                    }
-                }
-
-                Some(from) => {
-                    for attr in inherit.attrs() {
-                        let name = match self.expr_static_attr_str(&attr) {
-                            Some(name) => name,
-                            None => {
-                                // TODO(tazjin): error variant for dynamic
-                                // key in *inherit* (or generalise it)
-                                self.emit_error(&attr, ErrorKind::DynamicKeyInLet);
-                                continue;
-                            }
-                        };
-
-                        count += 1;
-                        inherit_froms.push((from.expr().unwrap(), name, self.span_for(&attr)));
-                    }
-                }
-            }
-        }
-
-        // Data structures to track the bindings observed in the
-        // second pass, and forward the information needed to compile
-        // their value.
-        enum BindingKind {
-            InheritFrom {
-                namespace: ast::Expr,
-                name: SmolStr,
-                span: Span,
-            },
-
-            Plain {
-                expr: ast::Expr,
-            },
-        }
-
-        struct KeySlot {
-            slot: LocalIdx,
-            name: SmolStr,
-        }
-
-        struct TrackedBinding {
-            key_slot: Option<KeySlot>,
-            value_slot: LocalIdx,
-            kind: BindingKind,
-        }
-
-        // Vector to track these observed bindings.
-        let mut bindings: Vec<TrackedBinding> = vec![];
-
-        // Begin second pass to ensure that all remaining identifiers
-        // (that may resolve recursively) are known.
-
-        // Begin with the inherit (from)s since they all become a thunk anyway
-        for (from, name, span) in inherit_froms {
-            let key_slot = if rec_attrs {
-                Some(KeySlot {
-                    slot: self.scope_mut().declare_phantom(span, false),
-                    name: SmolStr::new(&name),
-                })
-            } else {
-                None
-            };
-
-            let value_slot = self.declare_local(&span, &name);
-
-            bindings.push(TrackedBinding {
-                key_slot,
-                value_slot,
-                kind: BindingKind::InheritFrom {
-                    namespace: from,
-                    name: SmolStr::new(&name),
-                    span,
-                },
-            });
-        }
-
-        // Declare all regular bindings
-        for entry in node.attrpath_values() {
-            count += 1;
-
-            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 {
-                self.emit_error(
-                    &entry,
-                    ErrorKind::NotImplemented("nested bindings in recursive scope :("),
-                );
-                continue;
-            }
-
-            let key_slot = if rec_attrs {
-                let span = self.span_for(&entry.attrpath().unwrap());
-                Some(KeySlot {
-                    slot: self.scope_mut().declare_phantom(span, false),
-                    name: SmolStr::new(&path[0]),
-                })
-            } else {
-                None
-            };
-
-            let value_slot = self.declare_local(&entry.attrpath().unwrap(), path.pop().unwrap());
-
-            bindings.push(TrackedBinding {
-                key_slot,
-                value_slot,
-                kind: BindingKind::Plain {
-                    expr: entry.value().unwrap(),
-                },
-            });
-        }
-
-        // Third pass to place the values in the correct stack slots.
-        let mut value_indices: Vec<LocalIdx> = vec![];
-        for binding in bindings.into_iter() {
-            value_indices.push(binding.value_slot);
-
-            if let Some(key_slot) = binding.key_slot {
-                let span = self.scope()[key_slot.slot].span;
-                self.emit_constant(Value::String(key_slot.name.into()), &span);
-                self.scope_mut().mark_initialised(key_slot.slot);
-            }
-
-            match binding.kind {
-                // This entry is an inherit (from) expr. The value is
-                // placed on the stack by selecting an attribute.
-                BindingKind::InheritFrom {
-                    namespace,
-                    name,
-                    span,
-                } => {
-                    // Create a thunk wrapping value (which may be one as well) to
-                    // avoid forcing the from expr too early.
-                    self.thunk(binding.value_slot, &namespace, move |c, n, s| {
-                        c.compile(s, n.clone());
-                        c.emit_force(n);
-
-                        c.emit_constant(Value::String(name.into()), &span);
-                        c.push_op(OpCode::OpAttrsSelect, &span);
-                    })
-                }
-
-                // Binding is "just" a plain expression that needs to
-                // be compiled.
-                BindingKind::Plain { expr } => self.compile(binding.value_slot, expr),
-            }
-
-            // Any code after this point will observe the value in the
-            // right stack slot, so mark it as initialised.
-            self.scope_mut().mark_initialised(binding.value_slot);
-        }
-
-        // Fourth pass to emit finaliser instructions if necessary.
-        for idx in value_indices {
-            if self.scope()[idx].needs_finaliser {
-                let stack_idx = self.scope().stack_index(idx);
-                self.push_op(OpCode::OpFinalise(stack_idx), node);
-            }
-        }
-
-        count
-    }
-
-    /// Compile a standard `let ...; in ...` expression.
-    ///
-    /// 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.compile_recursive_scope(slot, false, &node);
-
-        // Deal with the body, then clean up the locals afterwards.
-        self.compile(slot, node.body().unwrap());
-        self.cleanup_scope(&node);
-    }
-
-    /// Resolve and compile access to an identifier in the scope.
-    fn compile_identifier_access<N: ToSpan + Clone>(
-        &mut self,
-        slot: LocalIdx,
-        ident: &str,
-        node: &N,
-    ) {
-        // If the identifier is a global, and it is not poisoned, emit
-        // the global directly.
-        if let Some(global) = self.globals.get(ident) {
-            if !self.scope().is_poisoned(ident) {
-                global.clone()(self, self.span_for(node));
-                return;
-            }
-        }
-
-        match self.scope_mut().resolve_local(ident) {
-            LocalPosition::Unknown => {
-                // Are we possibly dealing with an upvalue?
-                if let Some(idx) = self.resolve_upvalue(self.contexts.len() - 1, ident, node) {
-                    self.push_op(OpCode::OpGetUpvalue(idx), node);
-                    return;
-                }
-
-                // If there is a non-empty `with`-stack (or a parent
-                // context with one), emit a runtime dynamic
-                // resolution instruction.
-                if self.has_dynamic_ancestor() {
-                    self.emit_constant(Value::String(SmolStr::new(ident).into()), node);
-                    self.push_op(OpCode::OpResolveWith, node);
-                    return;
-                }
-
-                // Otherwise, this variable is missing.
-                self.emit_error(node, ErrorKind::UnknownStaticVariable);
-            }
-
-            LocalPosition::Known(idx) => {
-                let stack_idx = self.scope().stack_index(idx);
-                self.push_op(OpCode::OpGetLocal(stack_idx), node);
-            }
-
-            // This identifier is referring to a value from the same
-            // scope which is not yet defined. This identifier access
-            // must be thunked.
-            LocalPosition::Recursive(idx) => self.thunk(slot, node, move |compiler, node, _| {
-                let upvalue_idx = compiler.add_upvalue(
-                    compiler.contexts.len() - 1,
-                    node,
-                    UpvalueKind::Local(idx),
-                );
-                compiler.push_op(OpCode::OpGetUpvalue(upvalue_idx), node);
-            }),
-        };
-    }
-
-    fn compile_ident(&mut self, slot: LocalIdx, node: ast::Ident) {
-        let ident = node.ident_token().unwrap();
-        self.compile_identifier_access(slot, ident.text(), &node);
-    }
-
     /// Compile `with` expressions by emitting instructions that
     /// pop/remove the indices of attribute sets that are implicitly
     /// in scope through `with` on the "with-stack".
@@ -1061,15 +876,6 @@ impl Compiler<'_> {
         self.push_op(OpCode::OpCall, &node);
     }
 
-    fn compile_legacy_let(&mut self, slot: LocalIdx, node: ast::LegacyLet) {
-        self.emit_warning(&node, WarningKind::DeprecatedLegacyLet);
-        self.scope_mut().begin_scope();
-        self.compile_recursive_scope(slot, true, &node);
-        self.push_op(OpCode::OpAttrs(Count(node.entries().count())), &node);
-        self.emit_constant(Value::String(SmolStr::new_inline("body").into()), &node);
-        self.push_op(OpCode::OpAttrsSelect, &node);
-    }
-
     /// Compile an expression into a runtime thunk which should be
     /// lazily evaluated when accessed.
     // TODO: almost the same as Compiler::compile_lambda; unify?
@@ -1246,40 +1052,6 @@ impl Compiler<'_> {
         self.scope_mut().declare_local(name, span)
     }
 
-    fn resolve_upvalue<N: ToSpan>(
-        &mut self,
-        ctx_idx: usize,
-        name: &str,
-        node: &N,
-    ) -> Option<UpvalueIdx> {
-        if ctx_idx == 0 {
-            // There can not be any upvalue at the outermost context.
-            return None;
-        }
-
-        // Determine whether the upvalue is a local in the enclosing context.
-        match self.contexts[ctx_idx - 1].scope.resolve_local(name) {
-            // recursive upvalues are dealt with the same way as
-            // standard known ones, as thunks and closures are
-            // guaranteed to be placed on the stack (i.e. in the right
-            // position) *during* their runtime construction
-            LocalPosition::Known(idx) | LocalPosition::Recursive(idx) => {
-                return Some(self.add_upvalue(ctx_idx, node, UpvalueKind::Local(idx)))
-            }
-
-            LocalPosition::Unknown => { /* continue below */ }
-        };
-
-        // If the upvalue comes from even further up, we need to
-        // recurse to make sure that the upvalues are created at each
-        // level.
-        if let Some(idx) = self.resolve_upvalue(ctx_idx - 1, name, node) {
-            return Some(self.add_upvalue(ctx_idx, node, UpvalueKind::Upvalue(idx)));
-        }
-
-        None
-    }
-
     /// Determine whether the current lambda context has any ancestors
     /// that use dynamic scope resolution, and mark contexts as
     /// needing to capture their enclosing `with`-stack in their
@@ -1301,31 +1073,6 @@ impl Compiler<'_> {
         ancestor_has_with
     }
 
-    fn add_upvalue<N: ToSpan>(
-        &mut self,
-        ctx_idx: usize,
-        node: &N,
-        kind: UpvalueKind,
-    ) -> UpvalueIdx {
-        // If there is already an upvalue closing over the specified
-        // index, retrieve that instead.
-        for (idx, existing) in self.contexts[ctx_idx].scope.upvalues.iter().enumerate() {
-            if existing.kind == kind {
-                return UpvalueIdx(idx);
-            }
-        }
-
-        let span = self.span_for(node);
-        self.contexts[ctx_idx]
-            .scope
-            .upvalues
-            .push(Upvalue { kind, span });
-
-        let idx = UpvalueIdx(self.contexts[ctx_idx].lambda.upvalue_count);
-        self.contexts[ctx_idx].lambda.upvalue_count += 1;
-        idx
-    }
-
     fn emit_force<N: ToSpan>(&mut self, node: &N) {
         self.push_op(OpCode::OpForce, node);
     }
@@ -1339,70 +1086,6 @@ impl Compiler<'_> {
         let span = self.span_for(node);
         self.errors.push(Error { kind, span })
     }
-
-    /// Convert a non-dynamic string expression to a string if possible.
-    fn expr_static_str(&self, node: &ast::Str) -> Option<String> {
-        let mut parts = node.normalized_parts();
-
-        if parts.len() != 1 {
-            return None;
-        }
-
-        if let Some(ast::InterpolPart::Literal(lit)) = parts.pop() {
-            return Some(lit);
-        }
-
-        None
-    }
-
-    /// Convert the provided `ast::Attr` into a statically known
-    /// string if possible.
-    // TODO(tazjin): these should probably be SmolStr
-    fn expr_static_attr_str(&self, node: &ast::Attr) -> Option<String> {
-        match node {
-            ast::Attr::Ident(ident) => Some(ident.ident_token().unwrap().text().into()),
-            ast::Attr::Str(s) => self.expr_static_str(s),
-
-            // The dynamic node type is just a wrapper. C++ Nix does not
-            // care about the dynamic wrapper when determining whether the
-            // node itself is dynamic, it depends solely on the expression
-            // inside (i.e. `let ${"a"} = 1; in a` is valid).
-            ast::Attr::Dynamic(ref dynamic) => match dynamic.expr().unwrap() {
-                ast::Expr::Str(s) => self.expr_static_str(&s),
-                _ => None,
-            },
-        }
-    }
-
-    /// Construct the error returned when a dynamic attribute is found
-    /// in a `let`-expression.
-    fn dynamic_key_error<N>(&self, node: &N) -> Error
-    where
-        N: ToSpan,
-    {
-        Error {
-            kind: ErrorKind::DynamicKeyInLet,
-            span: self.span_for(node),
-        }
-    }
-
-    /// Convert a single identifier path fragment of a let binding to
-    /// a string if possible, or raise an error about the node being
-    /// dynamic.
-    fn binding_name(&self, node: ast::Attr) -> EvalResult<String> {
-        self.expr_static_attr_str(&node)
-            .ok_or_else(|| self.dynamic_key_error(&node))
-    }
-
-    /// Normalises identifier fragments into a single string vector
-    /// for `let`-expressions; fails if fragments requiring dynamic
-    /// computation are encountered.
-    fn normalise_ident_path<I: Iterator<Item = ast::Attr>>(
-        &self,
-        path: I,
-    ) -> EvalResult<Vec<String>> {
-        path.map(|node| self.binding_name(node)).collect()
-    }
 }
 
 /// Perform tail-call optimisation if the last call within a