From e834a2cbc47cf18d5f35bb258ccc754c54c4f4e1 Mon Sep 17 00:00:00 2001
From: William Carroll <wpcarro@gmail.com>
Date: Mon, 5 Sep 2022 11:43:17 -0700
Subject: feat(tvix/eval): Support builtins.attrNames

Define `.len()` method on `NixAttrs` to preallocate the capacity of the result
vector.

Also anchor an errant comment to its context (I think).

Change-Id: I268f15025d453d7b3ae1146558c80e51433dd2a8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6546
Reviewed-by: wpcarro <wpcarro@gmail.com>
Reviewed-by: sterni <sternenseemann@systemli.org>
Autosubmit: wpcarro <wpcarro@gmail.com>
Tested-by: BuildkiteCI
---
 tvix/eval/src/builtins/mod.rs                               | 12 ++++++++++++
 .../src/tests/tvix_tests/eval-okay-builtins-attrnames.exp   |  1 +
 .../src/tests/tvix_tests/eval-okay-builtins-attrnames.nix   | 13 +++++++++++++
 tvix/eval/src/value/attrs.rs                                | 12 ++++++++++--
 4 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.exp
 create mode 100644 tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.nix

(limited to 'tvix/eval')

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 3598ac71f7db..33abfe492d10 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -51,6 +51,18 @@ fn pure_builtins() -> Vec<Builtin> {
                 args.pop().unwrap().to_str()?.as_str().to_owned(),
             ));
         }),
+        Builtin::new("attrNames", 1, |args, vm| {
+            force!(vm, &args[0], value, {
+                let xs = value.to_attrs()?;
+                let mut output = Vec::with_capacity(xs.len());
+
+                for (key, _val) in xs.iter() {
+                    output.push(Value::String(key.clone()));
+                }
+
+                Ok(Value::List(NixList::construct(output.len(), output)))
+            })
+        }),
         Builtin::new("catAttrs", 2, |mut args, _| {
             let list = args.pop().unwrap().to_list()?;
             let key = args.pop().unwrap().to_str()?;
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.exp
new file mode 100644
index 000000000000..6521066a8ea5
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.exp
@@ -0,0 +1 @@
+[ [ ] [ "bar" "baz" "foo" ] [ "Baz" "Foo" "bar" ] [ "Eric Idle" "Graham Chapman" "John Cleese" "Michael Palin" "Terry Gilliam" "Terry Jones" ] ]
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.nix
new file mode 100644
index 000000000000..67f7dcee5672
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-builtins-attrnames.nix
@@ -0,0 +1,13 @@
+[
+  (builtins.attrNames {})
+  (builtins.attrNames { foo = 1; bar = 2; baz = 3; })
+  (builtins.attrNames { Foo = 1; bar = 2; Baz = 3; })
+  (builtins.attrNames {
+    "Graham Chapman" = true;
+    "John Cleese" = true;
+    "Terry Gilliam" = true;
+    "Eric Idle" = true;
+    "Terry Jones" = true;
+    "Michael Palin" = true;
+  })
+]
diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs
index 4c18ad2f55b2..86e52206b846 100644
--- a/tvix/eval/src/value/attrs.rs
+++ b/tvix/eval/src/value/attrs.rs
@@ -207,6 +207,15 @@ impl NixAttrs {
         }
     }
 
+    /// Return the number of key-value entries in an attrset.
+    pub fn len(&self) -> usize {
+        match &self.0 {
+            AttrsRep::Map(map) => map.len(),
+            AttrsRep::Empty => 0,
+            AttrsRep::KV { .. } => 2,
+        }
+    }
+
     /// Select a value from an attribute set by key.
     pub fn select(&self, key: &str) -> Option<&Value> {
         self.0.select(key)
@@ -216,6 +225,7 @@ impl NixAttrs {
         self.0.contains(key)
     }
 
+    /// Provide an iterator over all values of the attribute set.
     #[allow(clippy::needless_lifetimes)]
     pub fn iter<'a>(&'a self) -> Iter<KeyValue<'a>> {
         Iter(match &self.0 {
@@ -233,8 +243,6 @@ impl NixAttrs {
         })
     }
 
-    /// 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>) -> Result<Self, ErrorKind> {
-- 
cgit 1.4.1