diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs index 62fb9b710..d05cc177d 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> { + 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) -> EvalResult { @@ -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); + +impl<'a> Iterator for Iter> { + type Item = (&'a NixString, &'a Value); + + fn next(&mut self) -> Option { + 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 647a35865..c9e402247 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![ + (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![(&NixString::from("key"), &Value::String("value".into()))], + ); +}