1/*	$NetBSD$	*/
2
3/*
4** Id: ltablib.c,v 1.38.1.3 2008/02/14 16:46:58 roberto 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
24static 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
41static 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
58static 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
74static int getn (lua_State *L) {
75  lua_pushinteger(L, aux_getn(L, 1));
76  return 1;
77}
78
79
80static 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
92static 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
120static 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
137static 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
146static 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
175static void set2 (lua_State *L, int i, int j) {
176  lua_rawseti(L, 1, i);
177  lua_rawseti(L, 1, j);
178}
179
180static 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
195static 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
258static 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
271static 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
285LUALIB_API int luaopen_table (lua_State *L) {
286  luaL_register(L, LUA_TABLIBNAME, tab_funcs);
287  return 1;
288}
289
290