1 /* $NetBSD: ltablib.c,v 1.4 2015/10/08 13:21:00 mbalmer Exp $ */ 2 3 /* 4 ** Id: ltablib.c,v 1.80 2015/01/13 16:27:29 roberto Exp 5 ** Library for Table Manipulation 6 ** See Copyright Notice in lua.h 7 */ 8 9 #define ltablib_c 10 #define LUA_LIB 11 12 #include "lprefix.h" 13 14 15 #ifndef _KERNEL 16 #include <limits.h> 17 #include <stddef.h> 18 #endif 19 20 #include "lua.h" 21 22 #include "lauxlib.h" 23 #include "lualib.h" 24 25 26 27 /* 28 ** Structure with table-access functions 29 */ 30 typedef struct { 31 int (*geti) (lua_State *L, int idx, lua_Integer n); 32 void (*seti) (lua_State *L, int idx, lua_Integer n); 33 } TabA; 34 35 36 /* 37 ** Check that 'arg' has a table and set access functions in 'ta' to raw 38 ** or non-raw according to the presence of corresponding metamethods. 39 */ 40 static void checktab (lua_State *L, int arg, TabA *ta) { 41 ta->geti = NULL; ta->seti = NULL; 42 if (lua_getmetatable(L, arg)) { 43 lua_pushliteral(L, "__index"); /* 'index' metamethod */ 44 if (lua_rawget(L, -2) != LUA_TNIL) 45 ta->geti = lua_geti; 46 lua_pushliteral(L, "__newindex"); /* 'newindex' metamethod */ 47 if (lua_rawget(L, -3) != LUA_TNIL) 48 ta->seti = lua_seti; 49 lua_pop(L, 3); /* pop metatable plus both metamethods */ 50 } 51 if (ta->geti == NULL || ta->seti == NULL) { 52 luaL_checktype(L, arg, LUA_TTABLE); /* must be table for raw methods */ 53 if (ta->geti == NULL) ta->geti = lua_rawgeti; 54 if (ta->seti == NULL) ta->seti = lua_rawseti; 55 } 56 } 57 58 59 #define aux_getn(L,n,ta) (checktab(L, n, ta), luaL_len(L, n)) 60 61 62 #if defined(LUA_COMPAT_MAXN) 63 static int maxn (lua_State *L) { 64 lua_Number max = 0; 65 luaL_checktype(L, 1, LUA_TTABLE); 66 lua_pushnil(L); /* first key */ 67 while (lua_next(L, 1)) { 68 lua_pop(L, 1); /* remove value */ 69 if (lua_type(L, -1) == LUA_TNUMBER) { 70 lua_Number v = lua_tonumber(L, -1); 71 if (v > max) max = v; 72 } 73 } 74 lua_pushnumber(L, max); 75 return 1; 76 } 77 #endif 78 79 80 static int tinsert (lua_State *L) { 81 TabA ta; 82 lua_Integer e = aux_getn(L, 1, &ta) + 1; /* first empty element */ 83 lua_Integer pos; /* where to insert new element */ 84 switch (lua_gettop(L)) { 85 case 2: { /* called with only 2 arguments */ 86 pos = e; /* insert new element at the end */ 87 break; 88 } 89 case 3: { 90 lua_Integer i; 91 pos = luaL_checkinteger(L, 2); /* 2nd argument is the position */ 92 luaL_argcheck(L, 1 <= pos && pos <= e, 2, "position out of bounds"); 93 for (i = e; i > pos; i--) { /* move up elements */ 94 (*ta.geti)(L, 1, i - 1); 95 (*ta.seti)(L, 1, i); /* t[i] = t[i - 1] */ 96 } 97 break; 98 } 99 default: { 100 return luaL_error(L, "wrong number of arguments to 'insert'"); 101 } 102 } 103 (*ta.seti)(L, 1, pos); /* t[pos] = v */ 104 return 0; 105 } 106 107 108 static int tremove (lua_State *L) { 109 TabA ta; 110 lua_Integer size = aux_getn(L, 1, &ta); 111 lua_Integer pos = luaL_optinteger(L, 2, size); 112 if (pos != size) /* validate 'pos' if given */ 113 luaL_argcheck(L, 1 <= pos && pos <= size + 1, 1, "position out of bounds"); 114 (*ta.geti)(L, 1, pos); /* result = t[pos] */ 115 for ( ; pos < size; pos++) { 116 (*ta.geti)(L, 1, pos + 1); 117 (*ta.seti)(L, 1, pos); /* t[pos] = t[pos + 1] */ 118 } 119 lua_pushnil(L); 120 (*ta.seti)(L, 1, pos); /* t[pos] = nil */ 121 return 1; 122 } 123 124 125 static int tmove (lua_State *L) { 126 TabA ta; 127 lua_Integer f = luaL_checkinteger(L, 2); 128 lua_Integer e = luaL_checkinteger(L, 3); 129 lua_Integer t = luaL_checkinteger(L, 4); 130 int tt = !lua_isnoneornil(L, 5) ? 5 : 1; /* destination table */ 131 if (e >= f) { /* otherwise, nothing to move */ 132 lua_Integer n, i; 133 ta.geti = (luaL_getmetafield(L, 1, "__index") == LUA_TNIL) 134 ? (luaL_checktype(L, 1, LUA_TTABLE), lua_rawgeti) 135 : lua_geti; 136 ta.seti = (luaL_getmetafield(L, tt, "__newindex") == LUA_TNIL) 137 ? (luaL_checktype(L, tt, LUA_TTABLE), lua_rawseti) 138 : lua_seti; 139 luaL_argcheck(L, f > 0 || e < LUA_MAXINTEGER + f, 3, 140 "too many elements to move"); 141 n = e - f + 1; /* number of elements to move */ 142 luaL_argcheck(L, t <= LUA_MAXINTEGER - n + 1, 4, 143 "destination wrap around"); 144 if (t > f) { 145 for (i = n - 1; i >= 0; i--) { 146 (*ta.geti)(L, 1, f + i); 147 (*ta.seti)(L, tt, t + i); 148 } 149 } 150 else { 151 for (i = 0; i < n; i++) { 152 (*ta.geti)(L, 1, f + i); 153 (*ta.seti)(L, tt, t + i); 154 } 155 } 156 } 157 lua_pushvalue(L, tt); /* return "to table" */ 158 return 1; 159 } 160 161 162 static void addfield (lua_State *L, luaL_Buffer *b, TabA *ta, lua_Integer i) { 163 (*ta->geti)(L, 1, i); 164 if (!lua_isstring(L, -1)) 165 luaL_error(L, "invalid value (%s) at index %d in table for 'concat'", 166 luaL_typename(L, -1), i); 167 luaL_addvalue(b); 168 } 169 170 171 static int tconcat (lua_State *L) { 172 TabA ta; 173 luaL_Buffer b; 174 size_t lsep; 175 lua_Integer i, last; 176 const char *sep = luaL_optlstring(L, 2, "", &lsep); 177 checktab(L, 1, &ta); 178 i = luaL_optinteger(L, 3, 1); 179 last = luaL_opt(L, luaL_checkinteger, 4, luaL_len(L, 1)); 180 luaL_buffinit(L, &b); 181 for (; i < last; i++) { 182 addfield(L, &b, &ta, i); 183 luaL_addlstring(&b, sep, lsep); 184 } 185 if (i == last) /* add last value (if interval was not empty) */ 186 addfield(L, &b, &ta, i); 187 luaL_pushresult(&b); 188 return 1; 189 } 190 191 192 /* 193 ** {====================================================== 194 ** Pack/unpack 195 ** ======================================================= 196 */ 197 198 static int pack (lua_State *L) { 199 int i; 200 int n = lua_gettop(L); /* number of elements to pack */ 201 lua_createtable(L, n, 1); /* create result table */ 202 lua_insert(L, 1); /* put it at index 1 */ 203 for (i = n; i >= 1; i--) /* assign elements */ 204 lua_rawseti(L, 1, i); 205 lua_pushinteger(L, n); 206 lua_setfield(L, 1, "n"); /* t.n = number of elements */ 207 return 1; /* return table */ 208 } 209 210 211 static int unpack (lua_State *L) { 212 TabA ta; 213 lua_Integer i, e; 214 lua_Unsigned n; 215 checktab(L, 1, &ta); 216 i = luaL_optinteger(L, 2, 1); 217 e = luaL_opt(L, luaL_checkinteger, 3, luaL_len(L, 1)); 218 if (i > e) return 0; /* empty range */ 219 n = (lua_Unsigned)e - i; /* number of elements minus 1 (avoid overflows) */ 220 if (n >= (unsigned int)INT_MAX || !lua_checkstack(L, (int)(++n))) 221 return luaL_error(L, "too many results to unpack"); 222 do { /* must have at least one element */ 223 (*ta.geti)(L, 1, i); /* push arg[i..e] */ 224 } while (i++ < e); 225 226 return (int)n; 227 } 228 229 /* }====================================================== */ 230 231 232 233 /* 234 ** {====================================================== 235 ** Quicksort 236 ** (based on 'Algorithms in MODULA-3', Robert Sedgewick; 237 ** Addison-Wesley, 1993.) 238 ** ======================================================= 239 */ 240 241 242 static void set2 (lua_State *L, TabA *ta, int i, int j) { 243 (*ta->seti)(L, 1, i); 244 (*ta->seti)(L, 1, j); 245 } 246 247 static int sort_comp (lua_State *L, int a, int b) { 248 if (!lua_isnil(L, 2)) { /* function? */ 249 int res; 250 lua_pushvalue(L, 2); 251 lua_pushvalue(L, a-1); /* -1 to compensate function */ 252 lua_pushvalue(L, b-2); /* -2 to compensate function and 'a' */ 253 lua_call(L, 2, 1); 254 res = lua_toboolean(L, -1); 255 lua_pop(L, 1); 256 return res; 257 } 258 else /* a < b? */ 259 return lua_compare(L, a, b, LUA_OPLT); 260 } 261 262 static void auxsort (lua_State *L, TabA *ta, int l, int u) { 263 while (l < u) { /* for tail recursion */ 264 int i, j; 265 /* sort elements a[l], a[(l+u)/2] and a[u] */ 266 (*ta->geti)(L, 1, l); 267 (*ta->geti)(L, 1, u); 268 if (sort_comp(L, -1, -2)) /* a[u] < a[l]? */ 269 set2(L, ta, l, u); /* swap a[l] - a[u] */ 270 else 271 lua_pop(L, 2); 272 if (u-l == 1) break; /* only 2 elements */ 273 i = (l+u)/2; 274 (*ta->geti)(L, 1, i); 275 (*ta->geti)(L, 1, l); 276 if (sort_comp(L, -2, -1)) /* a[i]<a[l]? */ 277 set2(L, ta, i, l); 278 else { 279 lua_pop(L, 1); /* remove a[l] */ 280 (*ta->geti)(L, 1, u); 281 if (sort_comp(L, -1, -2)) /* a[u]<a[i]? */ 282 set2(L, ta, i, u); 283 else 284 lua_pop(L, 2); 285 } 286 if (u-l == 2) break; /* only 3 elements */ 287 (*ta->geti)(L, 1, i); /* Pivot */ 288 lua_pushvalue(L, -1); 289 (*ta->geti)(L, 1, u-1); 290 set2(L, ta, i, u-1); 291 /* a[l] <= P == a[u-1] <= a[u], only need to sort from l+1 to u-2 */ 292 i = l; j = u-1; 293 for (;;) { /* invariant: a[l..i] <= P <= a[j..u] */ 294 /* repeat ++i until a[i] >= P */ 295 while ((*ta->geti)(L, 1, ++i), sort_comp(L, -1, -2)) { 296 if (i>=u) luaL_error(L, "invalid order function for sorting"); 297 lua_pop(L, 1); /* remove a[i] */ 298 } 299 /* repeat --j until a[j] <= P */ 300 while ((*ta->geti)(L, 1, --j), sort_comp(L, -3, -1)) { 301 if (j<=l) luaL_error(L, "invalid order function for sorting"); 302 lua_pop(L, 1); /* remove a[j] */ 303 } 304 if (j<i) { 305 lua_pop(L, 3); /* pop pivot, a[i], a[j] */ 306 break; 307 } 308 set2(L, ta, i, j); 309 } 310 (*ta->geti)(L, 1, u-1); 311 (*ta->geti)(L, 1, i); 312 set2(L, ta, u-1, i); /* swap pivot (a[u-1]) with a[i] */ 313 /* a[l..i-1] <= a[i] == P <= a[i+1..u] */ 314 /* adjust so that smaller half is in [j..i] and larger one in [l..u] */ 315 if (i-l < u-i) { 316 j=l; i=i-1; l=i+2; 317 } 318 else { 319 j=i+1; i=u; u=j-2; 320 } 321 auxsort(L, ta, j, i); /* call recursively the smaller one */ 322 } /* repeat the routine for the larger one */ 323 } 324 325 static int sort (lua_State *L) { 326 TabA ta; 327 int n = (int)aux_getn(L, 1, &ta); 328 luaL_checkstack(L, 50, ""); /* assume array is smaller than 2^50 */ 329 if (!lua_isnoneornil(L, 2)) /* is there a 2nd argument? */ 330 luaL_checktype(L, 2, LUA_TFUNCTION); 331 lua_settop(L, 2); /* make sure there are two arguments */ 332 auxsort(L, &ta, 1, n); 333 return 0; 334 } 335 336 /* }====================================================== */ 337 338 339 static const luaL_Reg tab_funcs[] = { 340 {"concat", tconcat}, 341 #if defined(LUA_COMPAT_MAXN) 342 {"maxn", maxn}, 343 #endif 344 {"insert", tinsert}, 345 {"pack", pack}, 346 {"unpack", unpack}, 347 {"remove", tremove}, 348 {"move", tmove}, 349 {"sort", sort}, 350 {NULL, NULL} 351 }; 352 353 354 LUAMOD_API int luaopen_table (lua_State *L) { 355 luaL_newlib(L, tab_funcs); 356 #if defined(LUA_COMPAT_UNPACK) 357 /* _G.unpack = table.unpack */ 358 lua_getfield(L, -1, "unpack"); 359 lua_setglobal(L, "unpack"); 360 #endif 361 return 1; 362 } 363 364