about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--tvix/eval/src/value/attrs.rs75
-rw-r--r--tvix/eval/src/value/attrs/tests.rs47
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()))],
+    );
+}