From 049aeef62fca28931633a24c58fd1b587913e92a Mon Sep 17 00:00:00 2001
From: Konstantin Nazarov <mail@knazarov.com>
Date: Sun, 15 Sep 2024 16:21:58 +0100
Subject: [PATCH] Implement a map function that works on lists

---
 src/stdlib.cpp | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/stdlib.hpp |  1 +
 src/vm.cpp     |  1 -
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/stdlib.cpp b/src/stdlib.cpp
index 8d882d2..0eceffc 100644
--- a/src/stdlib.cpp
+++ b/src/stdlib.cpp
@@ -161,6 +161,60 @@ Result<StackFrame> stdlib_size(const StackFrame& stack) {
   return res.ret(0);
 }
 
+Result<StackFrame> stdlib_map(const StackFrame& stack_param) {
+  auto stack = TRY(stack_param.copy());
+  auto params = TRY(stack.get(0));
+  auto stack_size = TRY(stack.size());
+
+  auto fun = TRY(params.get(1));
+
+  if (stack_size == 1) {
+    // This is the first call, need to set up arguments on the stack
+    auto collection = TRY(params.get(0));
+
+    auto num_params = TRY(params.size());
+    if (num_params != 2) return ERROR(ArgumentCountMismatch);
+
+    if (!collection.is<Pair>()) return ERROR(TypeMismatch);
+
+    // Result accumulator will contain nil value initially
+    auto nil = Value(TRY(Nil::create()));
+    stack = TRY(stack.set(1, nil));
+
+    // Remaining elements in the collection
+    stack = TRY(stack.set(2, TRY(collection.rest())));
+
+    stack = TRY(stack.set(3, fun));
+    stack = TRY(stack.set(4, TRY(collection.first())));
+
+    return TRY(stack.call(fun, 3, 5));
+  }
+
+  auto accumulator = TRY(stack.get(1));
+  auto remaining = TRY(stack.get(2));
+  auto fun_result = TRY(stack.get(3));
+
+  accumulator = Value(TRY(Pair::create(fun_result, accumulator)));
+
+  stack = TRY(stack.set(1, accumulator));
+
+  if (!remaining.is<Nil>()) {
+    stack = TRY(stack.set(2, TRY(remaining.rest())));
+    stack = TRY(stack.set(3, fun));
+    stack = TRY(stack.set(4, TRY(remaining.first())));
+
+    return TRY(stack.call(fun, 3, 5));
+  }
+
+  accumulator = TRY(reverse(accumulator));
+
+  auto res = TRY(stack.set(0, accumulator));
+  res = TRY(res.ret(0));
+  res = TRY(res.incpc());
+
+  return res;
+}
+
 #define STDLIB_FUNCTION(name, id)                                  \
   [(uint64_t)StdlibFunctionId::id] = {#name, StdlibFunctionId::id, \
                                       stdlib_##name}
@@ -177,6 +231,7 @@ static StdlibFunctionEntry function_entries[] = {
     STDLIB_FUNCTION(get, Get),
     STDLIB_FUNCTION(srcloc, SrcLoc),
     STDLIB_FUNCTION(size, Size),
+    STDLIB_FUNCTION(map, Map),
     [(uint64_t)StdlibFunctionId::Max] = {0, StdlibFunctionId::Max,
                                          stdlib_unknown},
 };
diff --git a/src/stdlib.hpp b/src/stdlib.hpp
index c04fd9a..b12e7d7 100644
--- a/src/stdlib.hpp
+++ b/src/stdlib.hpp
@@ -16,6 +16,7 @@ enum class StdlibFunctionId : uint64_t {
   Get,
   SrcLoc,
   Size,
+  Map,
   Max,
 };
 
diff --git a/src/vm.cpp b/src/vm.cpp
index 2a6535f..043e4ba 100644
--- a/src/vm.cpp
+++ b/src/vm.cpp
@@ -283,7 +283,6 @@ Result<void> VM::step_native() {
   auto fun_id = fun.to<StdlibFunction>()->fun_id();
 
   _stack = TRY(call_stdlib_function(fun_id, _stack));
-  _stack = TRY(_stack.incpc());
   return Result<void>();
 }