refactor(tvix/eval): explicitly construct attrs in phases r/4872

This makes the phases of attribute set construction that Nix has very
explicit (inherits, static keys, dynamic keys).

This change focuses on the split between dynamic/static keys by
collecting all dynamic ones while compiling the static ones, and then
phasing them in afterwards. It's possible we also need to do some
additional splitting inside of the inherits.

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

2 files changed, 101 insertions, 32 deletions

diff --git a/tvix/eval/src/compiler/attrs.rs b/tvix/eval/src/compiler/attrs.rs
index c9897daac2..db8410fcdb 100644
--- a/tvix/eval/src/compiler/attrs.rs
+++ b/tvix/eval/src/compiler/attrs.rs
@@ -20,42 +20,82 @@ impl Compiler<'_, '_> {
         }
     }
 
-    /// 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() {
-            let span = self.span_for(&node);
-            self.emit_warning(
-                span,
-                WarningKind::NotImplemented("recursive attribute sets"),
-            );
-        }
+    /// Compile the statically known entries of an attribute set. Which
+    /// keys are which is not known from the iterator, so discovered
+    /// dynamic keys are returned from here.
+    fn compile_static_attr_entries(
+        &mut self,
+        count: &mut usize,
+        entries: AstChildren<ast::AttrpathValue>,
+    ) -> Vec<ast::AttrpathValue> {
+        let mut dynamic_attrs: Vec<ast::AttrpathValue> = vec![];
 
-        // Open a scope to track the positions of the temporaries used
-        // by the `OpAttrs` instruction.
-        self.scope_mut().begin_scope();
+        'entries: for kv in entries {
+            // Attempt to turn the attrpath into a list of static
+            // strings, but abort this process if any dynamic
+            // fragments are encountered.
+            let static_attrpath: Option<Vec<String>> = kv
+                .attrpath()
+                .unwrap()
+                .attrs()
+                .map(|a| self.expr_static_attr_str(&a))
+                .collect();
 
-        let mut count = self.compile_inherit_attrs(slot, node.inherits());
+            let fragments = match static_attrpath {
+                Some(fragments) => fragments,
+                None => {
+                    dynamic_attrs.push(kv);
+                    continue 'entries;
+                }
+            };
+
+            // At this point we can increase the counter because we
+            // know that this particular attribute is static and can
+            // thus be processed here.
+            *count += 1;
+
+            let key_count = fragments.len();
+            for fragment in fragments.into_iter() {
+                self.emit_constant(Value::String(fragment.into()), &kv.attrpath().unwrap());
+            }
+
+            // 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(),
+                );
+            }
+
+            // 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(value_slot, kv.value().unwrap());
+            self.scope_mut().mark_initialised(value_slot);
+        }
 
-        for kv in node.attrpath_values() {
-            count += 1;
+        dynamic_attrs
+    }
+
+    /// Compile the dynamic entries of an attribute set, where keys
+    /// are only known at runtime.
+    fn compile_dynamic_attr_entries(
+        &mut self,
+        count: &mut usize,
+        entries: Vec<ast::AttrpathValue>,
+    ) {
+        for entry in entries.into_iter() {
+            *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_span = self.span_for(&entry.attrpath().unwrap());
             let key_idx = self.scope_mut().declare_phantom(key_span, false);
 
-            for fragment in kv.attrpath().unwrap().attrs() {
+            for fragment in entry.attrpath().unwrap().attrs() {
                 // Key fragments can contain dynamic expressions,
                 // which makes accounting for their stack slots very
                 // tricky.
@@ -86,7 +126,7 @@ impl Compiler<'_, '_> {
             if key_count > 1 {
                 self.push_op(
                     OpCode::OpAttrPath(Count(key_count)),
-                    &kv.attrpath().unwrap(),
+                    &entry.attrpath().unwrap(),
                 );
 
                 // Close the temporary scope that was set up for the
@@ -97,11 +137,39 @@ impl Compiler<'_, '_> {
             // 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_span = self.span_for(&entry.value().unwrap());
             let value_slot = self.scope_mut().declare_phantom(value_span, false);
-            self.compile(value_slot, kv.value().unwrap());
+            self.compile(value_slot, entry.value().unwrap());
             self.scope_mut().mark_initialised(value_slot);
         }
+    }
+
+    /// 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() {
+            let span = self.span_for(&node);
+            self.emit_warning(
+                span,
+                WarningKind::NotImplemented("recursive attribute sets"),
+            );
+        }
+
+        // Open a scope to track the positions of the temporaries used
+        // by the `OpAttrs` instruction.
+        self.scope_mut().begin_scope();
+
+        let mut count = self.compile_inherit_attrs(slot, node.inherits());
+
+        let dynamic_entries = self.compile_static_attr_entries(&mut count, node.attrpath_values());
+
+        self.compile_dynamic_attr_entries(&mut count, dynamic_entries);
 
         self.push_op(OpCode::OpAttrs(Count(count)), &node);
 
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 15ff0ea309..4a5966ae83 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -1257,6 +1257,7 @@ impl Compiler<'_, '_> {
 
     /// 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()),