refactor(tvix/value): encapsulate attrset logic within value::attrs r/4454

The internal optimisations of the set representation were previously
leaking into the VM, which is highly undesirable.

Keeping it encapsulated allows us to do additional optimisations
within value::attrs without being concerned about its use in the VM.

Change-Id: I7e7020bb0983b9d355d3db747b049b2faa60131f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6108
Reviewed-by: eta <tvl@eta.st>
Tested-by: BuildkiteCI

3 files changed, 194 insertions, 191 deletions

diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index e2ebc7cb344a..c9c5e1780a3c 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -2,14 +2,21 @@
 /// backing implementations, as they are used in very versatile
 /// use-cases that are all exposed the same way in the language
 /// surface.
+///
+/// Due to this, construction and management of attribute sets has
+/// some peculiarities that are encapsulated within this module.
 use std::collections::BTreeMap;
 use std::fmt::Display;
+use std::rc::Rc;
+
+use crate::errors::{Error, EvalResult};
 
 use super::string::NixString;
 use super::Value;
 
 #[derive(Debug)]
 pub enum NixAttrs {
+    Empty,
     Map(BTreeMap<NixString, Value>),
     KV { name: Value, value: Value },
 }
@@ -30,6 +37,11 @@ impl Display for NixAttrs {
                     f.write_fmt(format_args!("{} = {}; ", name, value))?;
                 }
             }
+
+            NixAttrs::Empty => {
+                /* no values to print! */
+                f.write_str("/* optimised empty set! */")?;
+            }
         }
 
         f.write_str("}")
@@ -41,3 +53,179 @@ impl PartialEq for NixAttrs {
         todo!("attrset equality")
     }
 }
+
+impl NixAttrs {
+    /// Implement construction logic of an attribute set, to encapsulate
+    /// logic about attribute set optimisations inside of this module.
+    pub fn construct(count: usize, mut stack_slice: Vec<Value>) -> EvalResult<Self> {
+        debug_assert!(
+            stack_slice.len() == count * 2,
+            "construct_attrs called with count == {}, but slice.len() == {}",
+            count,
+            stack_slice.len(),
+        );
+
+        // Optimisation: Empty attribute set
+        if count == 0 {
+            return Ok(NixAttrs::Empty);
+        }
+
+        // Optimisation: KV pattern
+        if count == 2 {
+            if let Some(kv) = attempt_optimise_kv(&mut stack_slice) {
+                return Ok(kv);
+            }
+        }
+
+        // TODO(tazjin): extend_reserve(count) (rust#72631)
+        let mut attrs: BTreeMap<NixString, Value> = BTreeMap::new();
+
+        for _ in 0..count {
+            let value = stack_slice.pop().unwrap();
+
+            // It is at this point that nested attribute sets need to
+            // be constructed (if they exist).
+            //
+            let key = stack_slice.pop().unwrap();
+            match key {
+                Value::String(ks) => set_attr(&mut attrs, ks, value)?,
+
+                Value::AttrPath(mut path) => {
+                    set_nested_attr(
+                        &mut attrs,
+                        path.pop().expect("AttrPath is never empty"),
+                        path,
+                        value,
+                    )?;
+                }
+
+                other => {
+                    return Err(Error::InvalidKeyType {
+                        given: other.type_of(),
+                    })
+                }
+            }
+        }
+
+        Ok(NixAttrs::Map(attrs))
+    }
+}
+
+// In Nix, name/value attribute pairs are frequently constructed from
+// literals. This particular case should avoid allocation of a map,
+// additional heap values etc. and use the optimised `KV` variant
+// instead.
+//
+// `slice` is the top of the stack from which the attrset is being
+// constructed, e.g.
+//
+//   slice: [ "value" 5 "name" "foo" ]
+//   index:   0       1 2      3
+//   stack:   3       2 1      0
+fn attempt_optimise_kv(slice: &mut [Value]) -> Option<NixAttrs> {
+    let (name_idx, value_idx) = {
+        match (&slice[2], &slice[0]) {
+            (Value::String(NixString(s1)), Value::String(NixString(s2)))
+                if (s1 == "name" && s2 == "value") =>
+            {
+                (3, 1)
+            }
+
+            (Value::String(NixString(s1)), Value::String(NixString(s2)))
+                if (s1 == "value" && s2 == "name") =>
+            {
+                (1, 3)
+            }
+
+            // Technically this branch lets type errors pass,
+            // but they will be caught during normal attribute
+            // set construction instead.
+            _ => return None,
+        }
+    };
+
+    Some(NixAttrs::KV {
+        name: std::mem::replace(&mut slice[name_idx], Value::Blackhole),
+        value: std::mem::replace(&mut slice[value_idx], Value::Blackhole),
+    })
+}
+
+// Set an attribute on an in-construction attribute set, while
+// checking against duplicate key.s
+fn set_attr(
+    attrs: &mut BTreeMap<NixString, Value>,
+    key: NixString,
+    value: Value,
+) -> EvalResult<()> {
+    let entry = attrs.entry(key);
+
+    match entry {
+        std::collections::btree_map::Entry::Occupied(entry) => {
+            return Err(Error::DuplicateAttrsKey {
+                key: entry.key().0.clone(),
+            })
+        }
+
+        std::collections::btree_map::Entry::Vacant(entry) => {
+            entry.insert(value);
+            return Ok(());
+        }
+    };
+}
+
+// Set a nested attribute inside of an attribute set, throwing a
+// duplicate key error if a non-hashmap entry already exists on the
+// path.
+//
+// There is some optimisation potential for this simple implementation
+// if it becomes a problem.
+fn set_nested_attr(
+    attrs: &mut BTreeMap<NixString, Value>,
+    key: NixString,
+    mut path: Vec<NixString>,
+    value: Value,
+) -> EvalResult<()> {
+    // If there is no next key we are at the point where we
+    // should insert the value itself.
+    if path.is_empty() {
+        return set_attr(attrs, key, value);
+    }
+
+    let entry = attrs.entry(key);
+
+    // If there is not we go one step further down, in which case we
+    // need to ensure that there either is no entry, or the existing
+    // entry is a hashmap into which to insert the next value.
+    //
+    // If a value of a different type exists, the user specified a
+    // duplicate key.
+    match entry {
+        // Vacant entry -> new attribute set is needed.
+        std::collections::btree_map::Entry::Vacant(entry) => {
+            let mut map = BTreeMap::new();
+
+            // TODO(tazjin): technically recursing further is not
+            // required, we can create the whole hierarchy here, but
+            // it's noisy.
+            set_nested_attr(&mut map, path.pop().expect("next key exists"), path, value)?;
+
+            entry.insert(Value::Attrs(Rc::new(NixAttrs::Map(map))));
+        }
+
+        // Occupied entry: Either error out if there is something
+        // other than attrs, or insert the next value.
+        std::collections::btree_map::Entry::Occupied(mut entry) => match entry.get_mut() {
+            Value::Attrs(_attrs) => {
+                todo!("implement mutable attrsets")
+            }
+
+            _ => {
+                return Err(Error::DuplicateAttrsKey {
+                    key: entry.key().0.clone(),
+                })
+            }
+        },
+    }
+
+    Ok(())
+}
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 55d44048b6bc..99c7ee8647c3 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -25,6 +25,7 @@ pub enum Value {
     // Internal values that, while they technically exist at runtime,
     // are never returned to or created directly by users.
     AttrPath(Vec<NixString>),
+    Blackhole,
 }
 
 impl Value {
@@ -47,7 +48,7 @@ impl Value {
             Value::List(_) => "list",
 
             // Internal types
-            Value::AttrPath(_) => "internal",
+            Value::AttrPath(_) | Value::Blackhole => "internal",
         }
     }
 
@@ -85,7 +86,7 @@ impl Display for Value {
             Value::List(list) => list.fmt(f),
 
             // internal types
-            Value::AttrPath(_) => f.write_str("internal"),
+            Value::AttrPath(_) | Value::Blackhole => f.write_str("internal"),
         }
     }
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 90f5f45f614d..a58d77cc2720 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -1,7 +1,7 @@
 //! This module implements the virtual (or abstract) machine that runs
 //! Tvix bytecode.
 
-use std::{collections::BTreeMap, rc::Rc};
+use std::rc::Rc;
 
 use crate::{
     chunk::Chunk,
@@ -23,10 +23,6 @@ impl VM {
         op
     }
 
-    fn peek(&self, at: usize) -> &Value {
-        &self.stack[self.stack.len() - 1 - at]
-    }
-
     fn pop(&mut self) -> Value {
         self.stack.pop().expect("TODO")
     }
@@ -151,110 +147,8 @@ impl VM {
     }
 
     fn run_attrset(&mut self, count: usize) -> EvalResult<()> {
-        // If the attribute count happens to be 2, we might be able to
-        // create the optimised name/value struct instead.
-        if count == 2 {
-            // When determining whether we are dealing with a
-            // name/value pair, we return the stack locations of name
-            // and value, using `0` as a sentinel value (i.e. if
-            // either is 0, we are dealing with some other attrset).
-            let is_pair = {
-                // The keys are located 1 & 3 values back in the
-                // stack.
-                let k1 = self.peek(1);
-                let k2 = self.peek(3);
-
-                match (k1, k2) {
-                    (Value::String(NixString(s1)), Value::String(NixString(s2)))
-                        if (s1 == "name" && s2 == "value") =>
-                    {
-                        (1, 2)
-                    }
-
-                    (Value::String(NixString(s1)), Value::String(NixString(s2)))
-                        if (s1 == "value" && s2 == "name") =>
-                    {
-                        (2, 1)
-                    }
-
-                    // Technically this branch lets type errors pass,
-                    // but they will be caught during normal attribute
-                    // set construction instead.
-                    _ => (0, 0),
-                }
-            };
-
-            match is_pair {
-                (1, 2) => {
-                    // The value of 'name' is at stack slot 0, the
-                    // value of 'value' is at stack slot 2.
-                    let pair = Value::Attrs(Rc::new(NixAttrs::KV {
-                        name: self.pop(),
-                        value: {
-                            self.pop(); // ignore the key
-                            self.pop()
-                        },
-                    }));
-
-                    // Clean up the last key fragment.
-                    self.pop();
-
-                    self.push(pair);
-                    return Ok(());
-                }
-
-                (2, 1) => {
-                    // The value of 'name' is at stack slot 2, the
-                    // value of 'value' is at stack slot 0.
-                    let pair = Value::Attrs(Rc::new(NixAttrs::KV {
-                        value: self.pop(),
-                        name: {
-                            self.pop(); // ignore the key
-                            self.pop()
-                        },
-                    }));
-
-                    // Clean up the last key fragment.
-                    self.pop();
-
-                    self.push(pair);
-                    return Ok(());
-                }
-                _ => {}
-            }
-        }
-
-        let mut attrs: BTreeMap<NixString, Value> = BTreeMap::new();
-
-        for _ in 0..count {
-            let value = self.pop();
-
-            // It is at this point that nested attribute sets need to
-            // be constructed (if they exist).
-            //
-            let key = self.pop();
-            match key {
-                Value::String(ks) => set_attr(&mut attrs, ks, value)?,
-
-                Value::AttrPath(mut path) => {
-                    set_nested_attr(
-                        &mut attrs,
-                        path.pop().expect("AttrPath is never empty"),
-                        path,
-                        value,
-                    )?;
-                }
-
-                other => {
-                    return Err(Error::InvalidKeyType {
-                        given: other.type_of(),
-                    })
-                }
-            }
-        }
-
-        // TODO(tazjin): extend_reserve(count) (rust#72631)
-        self.push(Value::Attrs(Rc::new(NixAttrs::Map(attrs))));
+        let attrs = NixAttrs::construct(count, self.stack.split_off(self.stack.len() - count * 2))?;
+        self.push(Value::Attrs(Rc::new(attrs)));
         Ok(())
     }
 
@@ -294,86 +188,6 @@ pub enum NumberPair {
     Integer(i64, i64),
 }
 
-// Set an attribute on an in-construction attribute set, while
-// checking against duplicate key.s
-fn set_attr(
-    attrs: &mut BTreeMap<NixString, Value>,
-    key: NixString,
-    value: Value,
-) -> EvalResult<()> {
-    let entry = attrs.entry(key);
-
-    match entry {
-        std::collections::btree_map::Entry::Occupied(entry) => {
-            return Err(Error::DuplicateAttrsKey {
-                key: entry.key().0.clone(),
-            })
-        }
-
-        std::collections::btree_map::Entry::Vacant(entry) => {
-            entry.insert(value);
-            return Ok(());
-        }
-    };
-}
-
-// Set a nested attribute inside of an attribute set, throwing a
-// duplicate key error if a non-hashmap entry already exists on the
-// path.
-//
-// There is some optimisation potential for this simple implementation
-// if it becomes a problem.
-fn set_nested_attr(
-    attrs: &mut BTreeMap<NixString, Value>,
-    key: NixString,
-    mut path: Vec<NixString>,
-    value: Value,
-) -> EvalResult<()> {
-    // If there is no next key we are at the point where we
-    // should insert the value itself.
-    if path.is_empty() {
-        return set_attr(attrs, key, value);
-    }
-
-    let entry = attrs.entry(key);
-
-    // If there is not we go one step further down, in which case we
-    // need to ensure that there either is no entry, or the existing
-    // entry is a hashmap into which to insert the next value.
-    //
-    // If a value of a different type exists, the user specified a
-    // duplicate key.
-    match entry {
-        // Vacant entry -> new attribute set is needed.
-        std::collections::btree_map::Entry::Vacant(entry) => {
-            let mut map = BTreeMap::new();
-
-            // TODO(tazjin): technically recursing further is not
-            // required, we can create the whole hierarchy here, but
-            // it's noisy.
-            set_nested_attr(&mut map, path.pop().expect("next key exists"), path, value)?;
-
-            entry.insert(Value::Attrs(Rc::new(NixAttrs::Map(map))));
-        }
-
-        // Occupied entry: Either error out if there is something
-        // other than attrs, or insert the next value.
-        std::collections::btree_map::Entry::Occupied(mut entry) => match entry.get_mut() {
-            Value::Attrs(_attrs) => {
-                todo!("implement mutable attrsets")
-            }
-
-            _ => {
-                return Err(Error::DuplicateAttrsKey {
-                    key: entry.key().0.clone(),
-                })
-            }
-        },
-    }
-
-    Ok(())
-}
-
 pub fn run_chunk(chunk: Chunk) -> EvalResult<Value> {
     let mut vm = VM {
         chunk,