refactor(tvix/eval): statically resolve select from constant attrs r/5831

When resolving a select expression (`attrs.name` or `attrs.name or
default`), if the set compiles to a constant attribute set (as is most
notably the case with `builtins`) we can backtrack and replace that
attribute set directly with the compiled value.

For something like `builtins.length`, this will directly emit an
`OpConstant` that leaves the `length` builtin on the stack.

Change-Id: I639654e065a06e8cfcbcacb528c6da7ec9e513ee
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7957
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

3 files changed, 94 insertions, 50 deletions

diff --git a/tvix/eval/docs/known-optimisation-potential.md b/tvix/eval/docs/known-optimisation-potential.md
index df13577de3ba..f45f1ee6c48a 100644
--- a/tvix/eval/docs/known-optimisation-potential.md
+++ b/tvix/eval/docs/known-optimisation-potential.md
@@ -50,17 +50,11 @@ optimisations, but note the most important ones here.
   can directly use the `value::function::Lambda` representation where
   possible.
 
-* Optimise inner builtin access [medium]
+* Apply `compiler::optimise_select` to other set operations [medium]
 
-  When accessing identifiers like `builtins.foo`, the compiler should
-  not go through the trouble of setting up the attribute set on the
-  stack and accessing `foo` from it if it knows that the scope for
-  `builtins` is unshadowed. The same optimisation can also be done
-  for the other set operations like `builtins ? foo` and
-  `builtins.foo or alternative-implementation`.
-
-  The same thing goes for resolving `with builtins;`, which should
-  definitely resolve statically.
+  In addition to selects, statically known attribute resolution could
+  also be used for things like `?` or `with`. The latter might be a
+  little more complicated but is worth investigating.
 
 * Inline fully applied builtins with equivalent operators [medium]
 
diff --git a/tvix/eval/src/compiler/bindings.rs b/tvix/eval/src/compiler/bindings.rs
index 59ca222b7219..9ce50d36db3e 100644
--- a/tvix/eval/src/compiler/bindings.rs
+++ b/tvix/eval/src/compiler/bindings.rs
@@ -226,7 +226,7 @@ impl TrackedBindings {
 
         // If the first element of the path is not statically known, the entry
         // can not be merged.
-        let name = match c.expr_static_attr_str(name) {
+        let name = match expr_static_attr_str(name) {
             Some(name) => name,
             None => return false,
         };
@@ -336,7 +336,7 @@ impl Compiler<'_> {
 
                 None => {
                     for attr in inherit.attrs() {
-                        let name = match self.expr_static_attr_str(&attr) {
+                        let name = match expr_static_attr_str(&attr) {
                             Some(name) => name,
                             None => {
                                 self.emit_error(&attr, ErrorKind::DynamicKeyInScope("inherit"));
@@ -387,7 +387,7 @@ impl Compiler<'_> {
 
                 Some(from) => {
                     for attr in inherit.attrs() {
-                        let name = match self.expr_static_attr_str(&attr) {
+                        let name = match expr_static_attr_str(&attr) {
                             Some(name) => name,
                             None => {
                                 self.emit_error(&attr, ErrorKind::DynamicKeyInScope("inherit"));
@@ -476,7 +476,7 @@ impl Compiler<'_> {
             *count += 1;
 
             let key_span = self.span_for(&key);
-            let key_slot = match self.expr_static_attr_str(&key) {
+            let key_slot = match expr_static_attr_str(&key) {
                 Some(name) if kind.is_attrs() => KeySlot::Static {
                     name,
                     slot: self.scope_mut().declare_phantom(key_span, false),
@@ -819,37 +819,4 @@ impl Compiler<'_> {
         self.contexts[ctx_idx].lambda.upvalue_count += 1;
         idx
     }
-
-    /// Convert a non-dynamic string expression to a string if possible.
-    fn expr_static_str(&self, node: &ast::Str) -> Option<SmolStr> {
-        let mut parts = node.normalized_parts();
-
-        if parts.len() != 1 {
-            return None;
-        }
-
-        if let Some(ast::InterpolPart::Literal(lit)) = parts.pop() {
-            return Some(SmolStr::new(lit));
-        }
-
-        None
-    }
-
-    /// Convert the provided `ast::Attr` into a statically known string if
-    /// possible.
-    fn expr_static_attr_str(&self, node: &ast::Attr) -> Option<SmolStr> {
-        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,
-            },
-        }
-    }
 }
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 9e648fb99873..82f1dd5f285d 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -29,7 +29,7 @@ use std::sync::Arc;
 use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
 use crate::observer::CompilerObserver;
-use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, UpvalueIdx};
+use crate::opcode::{CodeIdx, ConstantIdx, Count, JumpOffset, OpCode, UpvalueIdx};
 use crate::spans::LightSpan;
 use crate::spans::ToSpan;
 use crate::value::{Closure, Formals, Lambda, NixAttrs, Thunk, Value};
@@ -631,17 +631,63 @@ impl Compiler<'_> {
         self.push_op(OpCode::OpHasAttr, node);
     }
 
+    /// When compiling select or select_or expressions, an optimisation is
+    /// possible of compiling the set emitted a constant attribute set by
+    /// immediately replacing it with the actual value.
+    ///
+    /// We take care not to emit an error here, as that would interfere with
+    /// thunking behaviour (there can be perfectly valid Nix code that accesses
+    /// a statically known attribute set that is lacking a key, because that
+    /// thunk is never evaluated). If anything is missing, just inform the
+    /// caller that the optimisation did not take place and move on. We may want
+    /// to emit warnings here in the future.
+    fn optimise_select(&mut self, path: &ast::Attrpath) -> bool {
+        // If compiling the set emitted a constant attribute set, the
+        // associated constant can immediately be replaced with the
+        // actual value.
+        //
+        // We take care not to emit an error here, as that would
+        // interfere with thunking behaviour (there can be perfectly
+        // valid Nix code that accesses a statically known attribute
+        // set that is lacking a key, because that thunk is never
+        // evaluated). If anything is missing, just move on. We may
+        // want to emit warnings here in the future.
+        if let Some(OpCode::OpConstant(ConstantIdx(idx))) = self.chunk().code.last().cloned() {
+            let constant = &mut self.chunk().constants[idx];
+            if let Value::Attrs(attrs) = constant {
+                let mut path_iter = path.attrs();
+
+                // Only do this optimisation if there is a *single*
+                // element in the attribute path. It is extremely
+                // unlikely that we'd have a static nested set.
+                if let (Some(attr), None) = (path_iter.next(), path_iter.next()) {
+                    // Only do this optimisation for statically known attrs.
+                    if let Some(ident) = expr_static_attr_str(&attr) {
+                        if let Some(selected_value) = attrs.select(ident.as_str()) {
+                            *constant = selected_value.clone();
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+
+        false
+    }
+
     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;
+            return self.compile_select_or(slot, set, path, node.default_expr().unwrap());
         }
 
         // Push the set onto the stack
         self.compile(slot, set);
+        if self.optimise_select(&path) {
+            return;
+        }
 
         // Compile each key fragment and emit access instructions.
         //
@@ -693,6 +739,10 @@ impl Compiler<'_> {
         default: ast::Expr,
     ) {
         self.compile(slot, set);
+        if self.optimise_select(&path) {
+            return;
+        }
+
         let mut jumps = vec![];
 
         for fragment in path.attrs() {
@@ -1211,6 +1261,39 @@ impl Compiler<'_> {
     }
 }
 
+/// Convert a non-dynamic string expression to a string if possible.
+fn expr_static_str(node: &ast::Str) -> Option<SmolStr> {
+    let mut parts = node.normalized_parts();
+
+    if parts.len() != 1 {
+        return None;
+    }
+
+    if let Some(ast::InterpolPart::Literal(lit)) = parts.pop() {
+        return Some(SmolStr::new(lit));
+    }
+
+    None
+}
+
+/// Convert the provided `ast::Attr` into a statically known string if
+/// possible.
+fn expr_static_attr_str(node: &ast::Attr) -> Option<SmolStr> {
+    match node {
+        ast::Attr::Ident(ident) => Some(ident.ident_token().unwrap().text().into()),
+        ast::Attr::Str(s) => 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) => expr_static_str(&s),
+            _ => None,
+        },
+    }
+}
+
 /// Perform tail-call optimisation if the last call within a
 /// compiled chunk is another call.
 fn optimise_tail_call(chunk: &mut Chunk) {