Add table.slice

PlutoLang · Sep 23, 2024 · fb20a00 · fb20a00
1 parent 3a53bee
commit fb20a00
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 0 deletions.
diff --git a/src/ltablib.cpp b/src/ltablib.cpp
@@ -842,10 +842,50 @@ static int tcountvalues (lua_State *L) {
 }
 
 
+static int tslice (lua_State *L) {
+  luaL_checktype(L, 1, LUA_TTABLE);
+
+  lua_Integer idx_start = luaL_checkinteger(L, 2);
+  lua_Integer idx_end = luaL_checkinteger(L, 3);
+  const lua_Integer l = luaL_len(L, 1);
+
+  if (idx_start < 0) {
+    idx_start = l + idx_start + 1;
+  }
+
+  if (idx_end < 0) {
+    idx_end = l + idx_end + 1;
+  }
+  else if (idx_end > l) {
+    idx_end = l;
+  }
+
+  lua_newtable(L);
+  lua_Integer idx_result = 1;
+  for (lua_Integer i = idx_start; i <= idx_end; ++i) {
+    lua_pushinteger(L, i);
+    lua_gettable(L, 1);
+    if (!lua_isnoneornil(L, -1)) {
+      lua_pushinteger(L, idx_result++);
+      lua_pushvalue(L, -2);
+      lua_settable(L, 4);
+    }
+    else {
+      lua_pop(L, 1);
+    }
+  }
+
+  lua_settop(L, 4);
+
+  return 1;
+}
+
+
 /* }====================================================== */
 
 
 static const luaL_Reg tab_funcs[] = {
+  {"slice", tslice},
   {"countvalues", tcountvalues},
   {"keys", tkeys},
   {"modget", modget},

diff --git a/testes/pluto/basic.pluto b/testes/pluto/basic.pluto
@@ -2006,6 +2006,23 @@ do
     assert(test_table_7[4] == 1)
     assert(test_table_7[test_table_6] == 2)
 end
+do
+    local t = {1, 2, 3, 4, 5}
+
+    assert(table.slice(t, 1, 3)[1] == 1)
+    assert(table.slice(t, 1, 3)[2] == 2)
+    assert(table.slice(t, 1, 3)[3] == 3)
+    assert(#table.slice(t, 1, 3) == 3)
+
+    assert(table.slice(t, -3, -1)[1] == 3)
+    assert(table.slice(t, -3, -1)[2] == 4)
+    assert(table.slice(t, -3, -1)[3] == 5)
+    assert(#table.slice(t, -3, -1) == 3)
+
+    assert(#table.slice(t, 6, 10) == 0)
+    assert(#table.slice(t, 2, 10) == 4)
+    assert(#table.slice(t, 0, 10) == #t)
+end
 do
     assert(compareversions("0.1.0", "0.1.0") == 0)
     assert(compareversions("0.1.0", "0.2.0") ~= 0)