feat(tvix/eval): implement NixAttrs::iter() r/4702

Implementing iteration over NixAttrs requires a custom iterator type
in order to encapsulate the different representations. The BTreeMap
for example has its own iterator type which needs to be encapsulated.

This is mostly boilerplate code, but for a change some simple unit
tests have been added in.

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

2 files changed, 121 insertions, 1 deletions

diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index 62fb9b710c..d05cc177d4 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -206,7 +206,7 @@ impl NixAttrs {
         }
     }
 
-    // Select a value from an attribute set by key.
+    /// Select a value from an attribute set by key.
     pub fn select(&self, key: &str) -> Option<&Value> {
         self.0.select(key)
     }
@@ -215,6 +215,24 @@ impl NixAttrs {
         self.0.contains(key)
     }
 
+    pub fn iter<'a>(&'a self) -> Iter<KeyValue<'a>> {
+        Iter(match &self.0 {
+            AttrsRep::Map(map) => KeyValue::Map(map.iter()),
+            AttrsRep::Empty => KeyValue::Empty,
+
+            AttrsRep::KV {
+                ref name,
+                ref value,
+            } => KeyValue::KV {
+                name,
+                value,
+                at: IterKV::Name,
+            },
+        })
+    }
+
+    /// Provide an iterator over all values of the attribute set.
+
     /// 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> {
@@ -397,3 +415,58 @@ fn set_nested_attr(
 
     Ok(())
 }
+
+/// Internal helper type to track the iteration status of an iterator
+/// over the name/value representation.
+#[derive(Debug)]
+pub enum IterKV {
+    Name,
+    Value,
+    Done,
+}
+
+/// Iterator representation over the keys *and* values of an attribute
+/// set.
+#[derive(Debug)]
+pub enum KeyValue<'a> {
+    Empty,
+
+    KV {
+        name: &'a Value,
+        value: &'a Value,
+        at: IterKV,
+    },
+
+    Map(btree_map::Iter<'a, NixString, Value>),
+}
+
+/// Iterator over a Nix attribute set.
+// This wrapper type exists to make the inner "raw" iterator
+// inaccessible.
+#[repr(transparent)]
+pub struct Iter<T>(T);
+
+impl<'a> Iterator for Iter<KeyValue<'a>> {
+    type Item = (&'a NixString, &'a Value);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.0 {
+            KeyValue::Map(inner) => inner.next(),
+            KeyValue::Empty => None,
+
+            KeyValue::KV { name, value, at } => match at {
+                IterKV::Name => {
+                    *at = IterKV::Value;
+                    Some((NixString::NAME_REF, name))
+                }
+
+                IterKV::Value => {
+                    *at = IterKV::Done;
+                    Some((NixString::VALUE_REF, value))
+                }
+
+                IterKV::Done => None,
+            },
+        }
+    }
+}
diff --git a/tvix/eval/src/value/attrs/tests.rs b/tvix/eval/src/value/attrs/tests.rs
index 647a358655..c9e4022472 100644
--- a/tvix/eval/src/value/attrs/tests.rs
+++ b/tvix/eval/src/value/attrs/tests.rs
@@ -52,3 +52,50 @@ fn test_kv_attrs() {
         ),
     }
 }
+
+#[test]
+fn test_empty_attrs_iter() {
+    let attrs = NixAttrs::construct(0, vec![]).unwrap();
+    assert_eq!(attrs.iter().next(), None);
+}
+
+#[test]
+fn test_kv_attrs_iter() {
+    let name_val = Value::String("name".into());
+    let value_val = Value::String("value".into());
+    let meaning_val = Value::String("meaning".into());
+    let forty_two_val = Value::Integer(42);
+
+    let kv_attrs = NixAttrs::construct(
+        2,
+        vec![
+            value_val.clone(),
+            forty_two_val.clone(),
+            name_val.clone(),
+            meaning_val.clone(),
+        ],
+    )
+    .expect("constructing K/V pair attrs should succeed");
+
+    assert_eq!(
+        kv_attrs.iter().collect::<Vec<_>>(),
+        vec![
+            (NixString::NAME_REF, &meaning_val),
+            (NixString::VALUE_REF, &forty_two_val)
+        ]
+    );
+}
+
+#[test]
+fn test_map_attrs_iter() {
+    let attrs = NixAttrs::construct(
+        1,
+        vec![Value::String("key".into()), Value::String("value".into())],
+    )
+    .expect("simple attr construction should succeed");
+
+    assert_eq!(
+        attrs.iter().collect::<Vec<_>>(),
+        vec![(&NixString::from("key"), &Value::String("value".into()))],
+    );
+}