1 /* $NetBSD: ltablib.c,v 1.1.1.2 2012/03/15 00:08:07 alnsn Exp $ */ 2 3 /* 4 ** $Id: ltablib.c,v 1.1.1.2 2012/03/15 00:08:07 alnsn Exp $ 5 ** Library for Table Manipulation 6 ** See Copyright Notice in lua.h 7 */ 8 9 10 #include <stddef.h> 11 12 #define ltablib_c 13 #define LUA_LIB 14 15 #include "lua.h" 16 17 #include "lauxlib.h" 18 #include "lualib.h" 19 20 21 #define aux_getn(L,n) (luaL_checktype(L, n, LUA_TTABLE), luaL_getn(L, n)) 22 23 24 static int foreachi (lua_State *L) { 25 int i; 26 int n = aux_getn(L, 1); 27 luaL_checktype(L, 2, LUA_TFUNCTION); 28 for (i=1; i <= n; i++) { 29 lua_pushvalue(L, 2); /* function */ 30 lua_pushinteger(L, i); /* 1st argument */ 31 lua_rawgeti(L, 1, i); /* 2nd argument */ 32 lua_call(L, 2, 1); 33 if (!lua_isnil(L, -1)) 34 return 1; 35 lua_pop(L, 1); /* remove nil result */ 36 } 37 return 0; 38 } 39 40 41 static int foreach (lua_State *L) { 42 luaL_checktype(L, 1, LUA_TTABLE); 43 luaL_checktype(L, 2, LUA_TFUNCTION); 44 lua_pushnil(L); /* first key */ 45 while (lua_next(L, 1)) { 46 lua_pushvalue(L, 2); /* function */ 47 lua_pushvalue(L, -3); /* key */ 48 lua_pushvalue(L, -3); /* value */ 49 lua_call(L, 2, 1); 50 if (!lua_isnil(L, -1)) 51 return 1; 52 lua_pop(L, 2); /* remove value and result */ 53 } 54 return 0; 55 } 56 57 58 static int maxn (lua_State *L) { 59 lua_Number max = 0; 60 luaL_checktype(L, 1, LUA_TTABLE); 61 lua_pushnil(L); /* first key */ 62 while (lua_next(L, 1)) { 63 lua_pop(L, 1); /* remove value */ 64 if (lua_type(L, -1) == LUA_TNUMBER) { 65 lua_Number v = lua_tonumber(L, -1); 66 if (v > max) max = v; 67 } 68 } 69 lua_pushnumber(L, max); 70 return 1; 71 } 72 73 74 static int getn (lua_State *L) { 75 lua_pushinteger(L, aux_getn(L, 1)); 76 return 1; 77 } 78 79 80 static int setn (lua_State *L) { 81 luaL_checktype(L, 1, LUA_TTABLE); 82 #ifndef luaL_setn 83 luaL_setn(L, 1, luaL_checkint(L, 2)); 84 #else 85 luaL_error(L, LUA_QL("setn") " is obsolete"); 86 #endif 87 lua_pushvalue(L, 1); 88 return 1; 89 } 90 91 92 static int tinsert (lua_State *L) { 93 int e = aux_getn(L, 1) + 1; /* first empty element */ 94 int pos; /* where to insert new element */ 95 switch (lua_gettop(L)) { 96 case 2: { /* called with only 2 arguments */ 97 pos = e; /* insert new element at the end */ 98 break; 99 } 100 case 3: { 101 int i; 102 pos = luaL_checkint(L, 2); /* 2nd argument is the position */ 103 if (pos > e) e = pos; /* `grow' array if necessary */ 104 for (i = e; i > pos; i--) { /* move up elements */ 105 lua_rawgeti(L, 1, i-1); 106 lua_rawseti(L, 1, i); /* t[i] = t[i-1] */ 107 } 108 break; 109 } 110 default: { 111 return luaL_error(L, "wrong number of arguments to " LUA_QL("insert")); 112 } 113 } 114 luaL_setn(L, 1, e); /* new size */ 115 lua_rawseti(L, 1, pos); /* t[pos] = v */ 116 return 0; 117 } 118 119 120 static int tremove (lua_State *L) { 121 int e = aux_getn(L, 1); 122 int pos = luaL_optint(L, 2, e); 123 if (!(1 <= pos && pos <= e)) /* position is outside bounds? */ 124 return 0; /* nothing to remove */ 125 luaL_setn(L, 1, e - 1); /* t.n = n-1 */ 126 lua_rawgeti(L, 1, pos); /* result = t[pos] */ 127 for ( ;pos<e; pos++) { 128 lua_rawgeti(L, 1, pos+1); 129 lua_rawseti(L, 1, pos); /* t[pos] = t[pos+1] */ 130 } 131 lua_pushnil(L); 132 lua_rawseti(L, 1, e); /* t[e] = nil */ 133 return 1; 134 } 135 136 137 static void addfield (lua_State *L, luaL_Buffer *b, int i) { 138 lua_rawgeti(L, 1, i); 139 if (!lua_isstring(L, -1)) 140 luaL_error(L, "invalid value (%s) at index %d in table for " 141 LUA_QL("concat"), luaL_typename(L, -1), i); 142 luaL_addvalue(b); 143 } 144 145 146 static int tconcat (lua_State *L) { 147 luaL_Buffer b; 148 size_t lsep; 149 int i, last; 150 const char *sep = luaL_optlstring(L, 2, "", &lsep); 151 luaL_checktype(L, 1, LUA_TTABLE); 152 i = luaL_optint(L, 3, 1); 153 last = luaL_opt(L, luaL_checkint, 4, luaL_getn(L, 1)); 154 luaL_buffinit(L, &b); 155 for (; i < last; i++) { 156 addfield(L, &b, i); 157 luaL_addlstring(&b, sep, lsep); 158 } 159 if (i == last) /* add last value (if interval was not empty) */ 160 addfield(L, &b, i); 161 luaL_pushresult(&b); 162 return 1; 163 } 164 165 166 167 /* 168 ** {====================================================== 169 ** Quicksort 170 ** (based on `Algorithms in MODULA-3', Robert Sedgewick; 171 ** Addison-Wesley, 1993.) 172 */ 173 174 175 static void set2 (lua_State *L, int i, int j) { 176 lua_rawseti(L, 1, i); 177 lua_rawseti(L, 1, j); 178 } 179 180 static int sort_comp (lua_State *L, int a, int b) { 181 if (!lua_isnil(L, 2)) { /* function? */ 182 int res; 183 lua_pushvalue(L, 2); 184 lua_pushvalue(L, a-1); /* -1 to compensate function */ 185 lua_pushvalue(L, b-2); /* -2 to compensate function and `a' */ 186 lua_call(L, 2, 1); 187 res = lua_toboolean(L, -1); 188 lua_pop(L, 1); 189 return res; 190 } 191 else /* a < b? */ 192 return lua_lessthan(L, a, b); 193 } 194 195 static void auxsort (lua_State *L, int l, int u) { 196 while (l < u) { /* for tail recursion */ 197 int i, j; 198 /* sort elements a[l], a[(l+u)/2] and a[u] */ 199 lua_rawgeti(L, 1, l); 200 lua_rawgeti(L, 1, u); 201 if (sort_comp(L, -1, -2)) /* a[u] < a[l]? */ 202 set2(L, l, u); /* swap a[l] - a[u] */ 203 else 204 lua_pop(L, 2); 205 if (u-l == 1) break; /* only 2 elements */ 206 i = (l+u)/2; 207 lua_rawgeti(L, 1, i); 208 lua_rawgeti(L, 1, l); 209 if (sort_comp(L, -2, -1)) /* a[i]<a[l]? */ 210 set2(L, i, l); 211 else { 212 lua_pop(L, 1); /* remove a[l] */ 213 lua_rawgeti(L, 1, u); 214 if (sort_comp(L, -1, -2)) /* a[u]<a[i]? */ 215 set2(L, i, u); 216 else 217 lua_pop(L, 2); 218 } 219 if (u-l == 2) break; /* only 3 elements */ 220 lua_rawgeti(L, 1, i); /* Pivot */ 221 lua_pushvalue(L, -1); 222 lua_rawgeti(L, 1, u-1); 223 set2(L, i, u-1); 224 /* a[l] <= P == a[u-1] <= a[u], only need to sort from l+1 to u-2 */ 225 i = l; j = u-1; 226 for (;;) { /* invariant: a[l..i] <= P <= a[j..u] */ 227 /* repeat ++i until a[i] >= P */ 228 while (lua_rawgeti(L, 1, ++i), sort_comp(L, -1, -2)) { 229 if (i>u) luaL_error(L, "invalid order function for sorting"); 230 lua_pop(L, 1); /* remove a[i] */ 231 } 232 /* repeat --j until a[j] <= P */ 233 while (lua_rawgeti(L, 1, --j), sort_comp(L, -3, -1)) { 234 if (j<l) luaL_error(L, "invalid order function for sorting"); 235 lua_pop(L, 1); /* remove a[j] */ 236 } 237 if (j<i) { 238 lua_pop(L, 3); /* pop pivot, a[i], a[j] */ 239 break; 240 } 241 set2(L, i, j); 242 } 243 lua_rawgeti(L, 1, u-1); 244 lua_rawgeti(L, 1, i); 245 set2(L, u-1, i); /* swap pivot (a[u-1]) with a[i] */ 246 /* a[l..i-1] <= a[i] == P <= a[i+1..u] */ 247 /* adjust so that smaller half is in [j..i] and larger one in [l..u] */ 248 if (i-l < u-i) { 249 j=l; i=i-1; l=i+2; 250 } 251 else { 252 j=i+1; i=u; u=j-2; 253 } 254 auxsort(L, j, i); /* call recursively the smaller one */ 255 } /* repeat the routine for the larger one */ 256 } 257 258 static int sort (lua_State *L) { 259 int n = aux_getn(L, 1); 260 luaL_checkstack(L, 40, ""); /* assume array is smaller than 2^40 */ 261 if (!lua_isnoneornil(L, 2)) /* is there a 2nd argument? */ 262 luaL_checktype(L, 2, LUA_TFUNCTION); 263 lua_settop(L, 2); /* make sure there is two arguments */ 264 auxsort(L, 1, n); 265 return 0; 266 } 267 268 /* }====================================================== */ 269 270 271 static const luaL_Reg tab_funcs[] = { 272 {"concat", tconcat}, 273 {"foreach", foreach}, 274 {"foreachi", foreachi}, 275 {"getn", getn}, 276 {"maxn", maxn}, 277 {"insert", tinsert}, 278 {"remove", tremove}, 279 {"setn", setn}, 280 {"sort", sort}, 281 {NULL, NULL} 282 }; 283 284 285 LUALIB_API int luaopen_table (lua_State *L) { 286 luaL_register(L, LUA_TABLIBNAME, tab_funcs); 287 return 1; 288 } 289 290