about summary refs log tree commit diff
path: root/tvix/eval/src/value
diff options
context:
space:
mode:
Diffstat (limited to 'tvix/eval/src/value')
-rw-r--r--tvix/eval/src/value/attrs.rs58
-rw-r--r--tvix/eval/src/value/function.rs20
2 files changed, 73 insertions, 5 deletions
diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index e9d5a239a3cf..6ee3efee679b 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -202,7 +202,7 @@ impl NixAttrs {
         self.0.contains(key)
     }
 
-    /// Provide an iterator over all values of the attribute set.
+    /// Construct an iterator over all the key-value pairs in the attribute set.
     #[allow(clippy::needless_lifetimes)]
     pub fn iter<'a>(&'a self) -> Iter<KeyValue<'a>> {
         Iter(match &self.0 {
@@ -215,11 +215,20 @@ impl NixAttrs {
             } => KeyValue::KV {
                 name,
                 value,
-                at: IterKV::Name,
+                at: IterKV::default(),
             },
         })
     }
 
+    /// Construct an iterator over all the keys of the attribute set
+    pub fn keys(&self) -> Keys {
+        Keys(match &self.0 {
+            AttrsRep::Empty => KeysInner::Empty,
+            AttrsRep::Map(m) => KeysInner::Map(m.keys()),
+            AttrsRep::KV { .. } => KeysInner::KV(IterKV::default()),
+        })
+    }
+
     /// 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>) -> Result<Self, ErrorKind> {
@@ -405,13 +414,24 @@ fn set_attr(attrs: &mut NixAttrs, key: NixString, value: Value) -> Result<(), Er
 
 /// Internal helper type to track the iteration status of an iterator
 /// over the name/value representation.
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub enum IterKV {
+    #[default]
     Name,
     Value,
     Done,
 }
 
+impl IterKV {
+    fn next(&mut self) {
+        match *self {
+            Self::Name => *self = Self::Value,
+            Self::Value => *self = Self::Done,
+            Self::Done => {}
+        }
+    }
+}
+
 /// Iterator representation over the keys *and* values of an attribute
 /// set.
 #[derive(Debug)]
@@ -443,12 +463,12 @@ impl<'a> Iterator for Iter<KeyValue<'a>> {
 
             KeyValue::KV { name, value, at } => match at {
                 IterKV::Name => {
-                    *at = IterKV::Value;
+                    at.next();
                     Some((NixString::NAME_REF, name))
                 }
 
                 IterKV::Value => {
-                    *at = IterKV::Done;
+                    at.next();
                     Some((NixString::VALUE_REF, value))
                 }
 
@@ -457,3 +477,31 @@ impl<'a> Iterator for Iter<KeyValue<'a>> {
         }
     }
 }
+
+enum KeysInner<'a> {
+    Empty,
+    KV(IterKV),
+    Map(btree_map::Keys<'a, NixString, Value>),
+}
+
+pub struct Keys<'a>(KeysInner<'a>);
+
+impl<'a> Iterator for Keys<'a> {
+    type Item = &'a NixString;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.0 {
+            KeysInner::Empty => None,
+            KeysInner::KV(at @ IterKV::Name) => {
+                at.next();
+                Some(NixString::NAME_REF)
+            }
+            KeysInner::KV(at @ IterKV::Value) => {
+                at.next();
+                Some(NixString::VALUE_REF)
+            }
+            KeysInner::KV(IterKV::Done) => None,
+            KeysInner::Map(m) => m.next(),
+        }
+    }
+}
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index e58aab19c0e7..a3aa9325f598 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -2,9 +2,12 @@
 use std::{
     cell::{Ref, RefCell, RefMut},
     collections::HashMap,
+    hash::Hash,
     rc::Rc,
 };
 
+use codemap::Span;
+
 use crate::{
     chunk::Chunk,
     upvalues::{UpvalueCarrier, Upvalues},
@@ -19,6 +22,23 @@ pub(crate) struct Formals {
 
     /// Do the formals of this function accept extra arguments
     pub(crate) ellipsis: bool,
+
+    /// The span of the formals themselves, to use to emit errors
+    pub(crate) span: Span,
+}
+
+impl Formals {
+    /// Returns true if the given arg is a valid argument to these formals.
+    ///
+    /// This is true if it is either listed in the list of arguments, or the formals have an
+    /// ellipsis
+    pub(crate) fn contains<Q>(&self, arg: &Q) -> bool
+    where
+        Q: ?Sized + Hash + Eq,
+        NixString: std::borrow::Borrow<Q>,
+    {
+        self.ellipsis || self.arguments.contains_key(&arg)
+    }
 }
 
 /// The opcodes for a thunk or closure, plus the number of