refactor(tvix/eval): more efficiently intersect attributes r/6930

builtins.intersectAttrs is used a _lot_ in nixpkgs eval, for whatever
reason. We previously had a very inefficient implementation that would
allocate for each comparison. It stuck out like a sore thumb in perf
analysis.

This moves to a custom algorithm with two iterators, one for the left
and one for the right side, advancing them along the (borrowed) map
keys until a match is found and allocation is required.

I've not made any effort to reduce the verbosity of this code, I don't
think it's worth it.

On my machine this reduces the mean runtime of evaluating
`nixpkgs.emacs.outPath` by ~8%.

Change-Id: Ie506d82cb8d5f45909628f771a6b73e0eca16b27
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9898
Autosubmit: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 70 insertions, 9 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 13355fa1c5d2..c14ad13d0f1c 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -5,6 +5,7 @@
 
 use builtin_macros::builtins;
 use genawaiter::rc::Gen;
+use imbl::OrdMap;
 use regex::Regex;
 use std::cmp::{self, Ordering};
 use std::collections::VecDeque;
@@ -484,16 +485,76 @@ mod pure_builtins {
 
     #[builtin("intersectAttrs")]
     async fn builtin_intersect_attrs(co: GenCo, x: Value, y: Value) -> Result<Value, ErrorKind> {
-        let attrs1 = x.to_attrs()?;
-        let attrs2 = y.to_attrs()?;
-        let res = attrs2.iter().filter_map(|(k, v)| {
-            if attrs1.contains(k) {
-                Some((k.clone(), v.clone()))
-            } else {
-                None
+        let left_set = x.to_attrs()?;
+        if left_set.is_empty() {
+            return Ok(Value::attrs(NixAttrs::empty()));
+        }
+        let mut left_keys = left_set.keys();
+
+        let right_set = y.to_attrs()?;
+        if right_set.is_empty() {
+            return Ok(Value::attrs(NixAttrs::empty()));
+        }
+        let mut right_keys = right_set.keys();
+
+        let mut out: OrdMap<NixString, Value> = OrdMap::new();
+
+        // Both iterators have at least one entry
+        let mut left = left_keys.next().unwrap();
+        let mut right = right_keys.next().unwrap();
+
+        // Calculate the intersection of the attribute sets by simultaneously
+        // advancing two key iterators, and inserting into the result set from
+        // the right side when the keys match. Iteration over Nix attribute sets
+        // is in sorted lexicographical order, so we can advance either iterator
+        // until it "catches up" with its counterpart.
+        //
+        // Only when keys match are the key and value clones actually allocated.
+        //
+        // We opted for this implementation over simpler ones because of the
+        // heavy use of this function in nixpkgs.
+        loop {
+            if left == right {
+                // We know that the key exists in the set, and can
+                // skip the check instructions.
+                unsafe {
+                    out.insert(
+                        right.clone(),
+                        right_set.select(right).unwrap_unchecked().clone(),
+                    );
+                }
+
+                left = match left_keys.next() {
+                    Some(x) => x,
+                    None => break,
+                };
+
+                right = match right_keys.next() {
+                    Some(x) => x,
+                    None => break,
+                };
+
+                continue;
             }
-        });
-        Ok(Value::attrs(NixAttrs::from_iter(res)))
+
+            if left < right {
+                left = match left_keys.next() {
+                    Some(x) => x,
+                    None => break,
+                };
+                continue;
+            }
+
+            if right < left {
+                right = match right_keys.next() {
+                    Some(x) => x,
+                    None => break,
+                };
+                continue;
+            }
+        }
+
+        Ok(Value::attrs(out.into()))
     }
 
     #[builtin("isAttrs")]