From 2bad87d0b0a7d987aef8e41ff653bc662fda5210 Mon Sep 17 00:00:00 2001
From: Konstantin Nazarov <mail@knazarov.com>
Date: Tue, 10 Sep 2024 03:18:00 +0100
Subject: [PATCH] Implement "size" function in stdlib

---
 src/common.cpp | 12 ++++++++++++
 src/common.hpp |  4 ++++
 src/stdlib.cpp | 10 ++++++++++
 src/stdlib.hpp |  1 +
 4 files changed, 27 insertions(+)

diff --git a/src/common.cpp b/src/common.cpp
index 5c36324..5e37c14 100644
--- a/src/common.cpp
+++ b/src/common.cpp
@@ -391,6 +391,18 @@ Result<Value> Pair::rest() const {
   return Value::create(val);
 }
 
+Result<uint64_t> Pair::size() const {
+  uint64_t res = 1;
+
+  Value cur = TRY(rest());
+  while (!cur.is<Nil>()) {
+    res += 1;
+    cur = TRY(cur.rest());
+  }
+
+  return res;
+}
+
 Result<Value> Pair::get(const Value& key) const {
   if (!key.is<Int64>()) return ERROR(TypeMismatch);
 
diff --git a/src/common.hpp b/src/common.hpp
index 8e1f49c..d6e54ff 100644
--- a/src/common.hpp
+++ b/src/common.hpp
@@ -161,6 +161,7 @@ class Nil : public Object {
   virtual Result<Value> copy_value() const final;
   Result<Nil> copy() const;
 
+  virtual Result<uint64_t> size() const { return 0; }
   virtual void move(Object* obj) final { new (obj) Nil(std::move(_value)); }
 
  private:
@@ -757,6 +758,8 @@ class Pair : public Object {
   Result<Value> third() const final;
   Result<Value> rest() const final;
 
+  virtual Result<uint64_t> size() const final;
+
   virtual Result<Value> copy_value() const final;
   Result<Pair> copy() const;
 
@@ -1273,6 +1276,7 @@ class Value {
   Result<Value> second() const { return ((Object*)buf)->second(); }
   Result<Value> third() const { return ((Object*)buf)->third(); }
   Result<Value> rest() const { return ((Object*)buf)->rest(); }
+  Result<uint64_t> size() const { return ((Object*)buf)->size(); }
 
   // TODO: cmp() probably doesn't need arena parameter
   // Result<bool> operator==(Value& rhs) { return TRY(cmp(rhs)) == 0; }
diff --git a/src/stdlib.cpp b/src/stdlib.cpp
index 237b790..e21e11e 100644
--- a/src/stdlib.cpp
+++ b/src/stdlib.cpp
@@ -128,6 +128,15 @@ Result<Value> stdlib_srcloc(const Array& params) {
   return Value(TRY(SrcLoc::create(sr)));
 }
 
+Result<Value> stdlib_size(const Array& params) {
+  auto size = TRY(params.size());
+  if (size != 1) return ERROR(ArgumentCountMismatch);
+  auto param = TRY(params.get(0));
+  auto psize = TRY(param.size());
+
+  return Value(TRY(Int64::create((int64_t)psize)));
+}
+
 #define STDLIB_FUNCTION(name, id)                                  \
   [(uint64_t)StdlibFunctionId::id] = {#name, StdlibFunctionId::id, \
                                       stdlib_##name}
@@ -143,6 +152,7 @@ static StdlibFunctionEntry function_entries[] = {
     STDLIB_FUNCTION(array, Array),
     STDLIB_FUNCTION(get, Get),
     STDLIB_FUNCTION(srcloc, SrcLoc),
+    STDLIB_FUNCTION(size, Size),
     [(uint64_t)StdlibFunctionId::Max] = {0, StdlibFunctionId::Max,
                                          stdlib_unknown},
 };
diff --git a/src/stdlib.hpp b/src/stdlib.hpp
index 0e963dc..4c71ec6 100644
--- a/src/stdlib.hpp
+++ b/src/stdlib.hpp
@@ -15,6 +15,7 @@ enum class StdlibFunctionId : uint64_t {
   Array,
   Get,
   SrcLoc,
+  Size,
   Max,
 };