Add table.map

PlutoLang · Aug 25, 2023 · 688ee37 · 688ee37
1 parent c99b009
commit 688ee37
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/src/ltablib.cpp b/src/ltablib.cpp
@@ -513,10 +513,42 @@ static int tfilter (lua_State *L) {
 }
 
 
+static int tmap (lua_State *L) {
+  luaL_checktype(L, 1, LUA_TTABLE);
+  luaL_checktype(L, 2, LUA_TFUNCTION);
+
+  lua_newtable(L);
+  lua_Integer idx = 1;
+  lua_pushvalue(L, 1);
+  lua_pushnil(L);
+  /* stack now: out, in, key */
+  while (lua_next(L, -2)) {
+    /* stack now: out, in, key, value */
+    lua_pushvalue(L, 2);
+    lua_pushvalue(L, -2);
+    /* stack now: out, in, key, value, function, value */
+    lua_call(L, 1, 1);
+    /* stack now: out, in, key, value, mapped_value */
+    lua_pushinteger(L, idx++);
+    lua_pushvalue(L, -2);
+    /* stack now: out, in, key, value, mapped_value, idx, mapped_value */
+    lua_settable(L, -7);
+    /* stack now: out, in, key, value, mapped_value */
+    lua_pop(L, 2);
+    /* stack now: out, in, key */
+  }
+  /* stack now: out, in */
+  lua_pop(L, 1);
+  /* stack now: out */
+  return 1;
+}
+
+
 /* }====================================================== */
 
 
 static const luaL_Reg tab_funcs[] = {
+  {"map", tmap},
   {"filter", tfilter},
   {"foreach", foreach},
   {"contains", tcontains},

diff --git a/tests/basic.pluto b/tests/basic.pluto
@@ -762,6 +762,14 @@ do
     assert(t[1] == 5)
     assert(t[2] == 3)
     assert(t[3] == 1)
+
+    local menu = {
+        { id = 1, name = "Pizza" },
+        { id = 2, name = "Doener" }
+    }
+    local betterDoener = { id = 2, name = "Better Doener" }
+    menu = table.map(menu, |item| -> item.id == betterDoener.id ? betterDoener : item)
+    assert(menu[2].name == "Better Doener")
 end
 
 print "Testing default table metatable."