From a612bb771baed0a3444926f8faa7b9e599b3eec5 Mon Sep 17 00:00:00 2001
From: Roman Godmaire <godmaire@twilit.io>
Date: Sat, 4 May 2024 17:26:26 -0400
Subject: [PATCH] feat: count function for collections

---
 src/env/core.rs  | 13 +++++++++++++
 src/evaluator.rs |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/src/env/core.rs b/src/env/core.rs
index 6e6e433..b7c4ff0 100644
--- a/src/env/core.rs
+++ b/src/env/core.rs
@@ -142,6 +142,19 @@ pub(super) fn core() -> HashMap<String, Node> {
                 }
             }),
         ),
+        (
+            "count",
+            Node::NativeFunc(|_env, args| {
+                arg_count!(1, args.len());
+
+                match args[0].borrow() {
+                    Node::List(list) => Ok(Node::Int(list.len() as i64)),
+                    Node::Vector(vec) => Ok(Node::Int(vec.len() as i64)),
+                    Node::Map(map) => Ok(Node::Int(map.len() as i64)),
+                    _ => Err(Error::ExpectedCollection)?,
+                }
+            }),
+        ),
         // Ordering
         (
             "eq?",
diff --git a/src/evaluator.rs b/src/evaluator.rs
index 1deb090..b79a418 100644
--- a/src/evaluator.rs
+++ b/src/evaluator.rs
@@ -114,18 +114,24 @@ mod test {
     #[case("(list? 23)", "false")]
     #[case("(empty? (list 1 2 3))", "false")]
     #[case("(empty? (list))", "true")]
+    #[case("(count (list))", "0")]
+    #[case("(count (list 1 2 3))", "3")]
     // Vectors
     #[case("(vector 1 2 3)", "[1 2 3]")]
     #[case("(vector? [])", "true")]
     #[case("(vector? 23)", "false")]
     #[case("(empty? [1 2 3])", "false")]
     #[case("(empty? [])", "true")]
+    #[case("(count [])", "0")]
+    #[case("(count [1 2 3])", "3")]
     // Hashmaps
     #[case("(hashmap :a 1)", "{:a: 1}")]
     #[case("(hashmap? {})", "true")]
     #[case("(hashmap? 23)", "false")]
     #[case("(empty? {:a 1})", "false")]
     #[case("(empty? {})", "true")]
+    #[case("(count {})", "0")]
+    #[case("(count {:a 1})", "1")]
     // Environment manipulation
     #[case("(define! asdf (+ 2 2)) (+ asdf 2)", "4\n6")]
     #[case("(let* (a 2) (+ a 2))", "4")]