refactor(tvix/eval): optimise initialisation of locals r/4663

Instead of looking up the local to be initialised by its name again,
we can simply track the index at which it was declared from the point
where the declaration was made.

This reduces some string cloning and removes unnecessary logic. It
also theoretically makes the *current* index available during locals
compilation, which can be used to optimise some recursion cases.

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

1 files changed, 19 insertions, 23 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index fb08662da84b..24199a34ea3d 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -620,11 +620,11 @@ impl Compiler {
                         }
 
                         self.compile_ident(ident.clone());
-                        self.declare_local(
+                        let idx = self.declare_local(
                             ident.syntax().clone(),
                             ident.ident_token().unwrap().text(),
                         );
-                        self.mark_initialised(ident.ident_token().unwrap().text());
+                        self.mark_initialised(idx);
                     }
                 }
 
@@ -633,11 +633,11 @@ impl Compiler {
                         self.compile(from.expr().unwrap());
                         self.emit_literal_ident(&ident);
                         self.chunk().push_op(OpCode::OpAttrsSelect);
-                        self.declare_local(
+                        let idx = self.declare_local(
                             ident.syntax().clone(),
                             ident.ident_token().unwrap().text(),
                         );
-                        self.mark_initialised(ident.ident_token().unwrap().text());
+                        self.mark_initialised(idx);
                     }
                 }
             }
@@ -656,7 +656,7 @@ impl Compiler {
 
         // First pass to ensure that all identifiers are known;
         // required for resolving recursion.
-        let mut entries: Vec<(String, rnix::ast::Expr)> = vec![];
+        let mut entries: Vec<(usize, rnix::ast::Expr)> = vec![];
         for entry in node.attrpath_values() {
             let mut path = match normalise_ident_path(entry.attrpath().unwrap().attrs()) {
                 Ok(p) => p,
@@ -670,18 +670,20 @@ impl Compiler {
                 todo!("nested bindings in let expressions :(")
             }
 
-            let name = path.pop().unwrap();
-            self.declare_local(entry.attrpath().unwrap().syntax().clone(), &name);
-            entries.push((name, entry.value().unwrap()));
+            let idx = self.declare_local(
+                entry.attrpath().unwrap().syntax().clone(),
+                path.pop().unwrap(),
+            );
+            entries.push((idx, entry.value().unwrap()));
         }
 
         // Second pass to place the values in the correct stack slots.
-        for (name, value) in entries.into_iter() {
+        for (idx, value) in entries.into_iter() {
             self.compile(value);
 
             // Any code after this point will observe the value in the
             // right stack slot, so mark it as initialised.
-            self.mark_initialised(&name);
+            self.mark_initialised(idx);
         }
 
         // Deal with the body, then clean up the locals afterwards.
@@ -794,8 +796,8 @@ impl Compiler {
                     .text()
                     .to_string();
 
-                self.declare_local(param.syntax().clone(), &name);
-                self.mark_initialised(&name);
+                let idx = self.declare_local(param.syntax().clone(), &name);
+                self.mark_initialised(idx);
             }
         }
 
@@ -936,7 +938,7 @@ impl Compiler {
     /// Declare a local variable known in the scope that is being
     /// compiled by pushing it to the locals. This is used to
     /// determine the stack offset of variables.
-    fn declare_local<S: Into<String>>(&mut self, node: rnix::SyntaxNode, name: S) {
+    fn declare_local<S: Into<String>>(&mut self, node: rnix::SyntaxNode, name: S) -> usize {
         let name = name.into();
         let depth = self.scope().scope_depth;
 
@@ -967,6 +969,7 @@ impl Compiler {
             );
         }
 
+        let idx = self.scope().locals.len();
         self.scope_mut().locals.push(Local {
             name,
             depth,
@@ -975,6 +978,7 @@ impl Compiler {
             phantom: false,
             used: false,
         });
+        idx
     }
 
     /// Declare a local variable that occupies a stack slot and should
@@ -993,16 +997,8 @@ impl Compiler {
     }
 
     /// Mark local as initialised after compiling its expression.
-    fn mark_initialised(&mut self, name: &str) {
-        let depth = self.scope().scope_depth;
-        for local in self.scope_mut().locals.iter_mut().rev() {
-            if !local.initialised && local.depth == depth && local.name == name {
-                local.initialised = true;
-                return;
-            }
-        }
-
-        panic!("critical compiler error: unbalanced locals stack");
+    fn mark_initialised(&mut self, local_idx: usize) {
+        self.scope_mut().locals[local_idx].initialised = true;
     }
 
     fn resolve_upvalue(&mut self, ctx_idx: usize, name: &str) -> Option<UpvalueIdx> {