- initial import of revision 374 from cnc
[apt.git] / lua / lib / ltablib.c
1 /*
2 ** $Id: ltablib.c,v 1.21 2003/04/03 13:35:34 roberto Exp $
3 ** Library for Table Manipulation
4 ** See Copyright Notice in lua.h
5 */
6
7
8 #include <stddef.h>
9
10 #define ltablib_c
11
12 #include "lua.h"
13
14 #include "lauxlib.h"
15 #include "lualib.h"
16
17
18 #define aux_getn(L,n)   (luaL_checktype(L, n, LUA_TTABLE), luaL_getn(L, n))
19
20
21 static int luaB_foreachi (lua_State *L) {
22   int i;
23   int n = aux_getn(L, 1);
24   luaL_checktype(L, 2, LUA_TFUNCTION);
25   for (i=1; i<=n; i++) {
26     lua_pushvalue(L, 2);  /* function */
27     lua_pushnumber(L, (lua_Number)i);  /* 1st argument */
28     lua_rawgeti(L, 1, i);  /* 2nd argument */
29     lua_call(L, 2, 1);
30     if (!lua_isnil(L, -1))
31       return 1;
32     lua_pop(L, 1);  /* remove nil result */
33   }
34   return 0;
35 }
36
37
38 static int luaB_foreach (lua_State *L) {
39   luaL_checktype(L, 1, LUA_TTABLE);
40   luaL_checktype(L, 2, LUA_TFUNCTION);
41   lua_pushnil(L);  /* first key */
42   for (;;) {
43     if (lua_next(L, 1) == 0)
44       return 0;
45     lua_pushvalue(L, 2);  /* function */
46     lua_pushvalue(L, -3);  /* key */
47     lua_pushvalue(L, -3);  /* value */
48     lua_call(L, 2, 1);
49     if (!lua_isnil(L, -1))
50       return 1;
51     lua_pop(L, 2);  /* remove value and result */
52   }
53 }
54
55
56 static int luaB_getn (lua_State *L) {
57   lua_pushnumber(L, (lua_Number)aux_getn(L, 1));
58   return 1;
59 }
60
61
62 static int luaB_setn (lua_State *L) {
63   luaL_checktype(L, 1, LUA_TTABLE);
64   luaL_setn(L, 1, luaL_checkint(L, 2));
65   return 0;
66 }
67
68
69 static int luaB_tinsert (lua_State *L) {
70   int v = lua_gettop(L);  /* number of arguments */
71   int n = aux_getn(L, 1) + 1;
72   int pos;  /* where to insert new element */
73   if (v == 2)  /* called with only 2 arguments */
74     pos = n;  /* insert new element at the end */
75   else {
76     pos = luaL_checkint(L, 2);  /* 2nd argument is the position */
77     if (pos > n) n = pos;  /* `grow' array if necessary */
78     v = 3;  /* function may be called with more than 3 args */
79   }
80   luaL_setn(L, 1, n);  /* new size */
81   while (--n >= pos) {  /* move up elements */
82     lua_rawgeti(L, 1, n);
83     lua_rawseti(L, 1, n+1);  /* t[n+1] = t[n] */
84   }
85   lua_pushvalue(L, v);
86   lua_rawseti(L, 1, pos);  /* t[pos] = v */
87   return 0;
88 }
89
90
91 static int luaB_tremove (lua_State *L) {
92   int n = aux_getn(L, 1);
93   int pos = luaL_optint(L, 2, n);
94   if (n <= 0) return 0;  /* table is `empty' */
95   luaL_setn(L, 1, n-1);  /* t.n = n-1 */
96   lua_rawgeti(L, 1, pos);  /* result = t[pos] */
97   for ( ;pos<n; pos++) {
98     lua_rawgeti(L, 1, pos+1);
99     lua_rawseti(L, 1, pos);  /* t[pos] = t[pos+1] */
100   }
101   lua_pushnil(L);
102   lua_rawseti(L, 1, n);  /* t[n] = nil */
103   return 1;
104 }
105
106
107 static int str_concat (lua_State *L) {
108   luaL_Buffer b;
109   size_t lsep;
110   const char *sep = luaL_optlstring(L, 2, "", &lsep);
111   int i = luaL_optint(L, 3, 1);
112   int n = luaL_optint(L, 4, 0);
113   luaL_checktype(L, 1, LUA_TTABLE);
114   if (n == 0) n = luaL_getn(L, 1);
115   luaL_buffinit(L, &b);
116   for (; i <= n; i++) {
117     lua_rawgeti(L, 1, i);
118     luaL_argcheck(L, lua_isstring(L, -1), 1, "table contains non-strings");
119     luaL_addvalue(&b);
120     if (i != n)
121       luaL_addlstring(&b, sep, lsep);
122   }
123   luaL_pushresult(&b);
124   return 1;
125 }
126
127
128
129 /*
130 ** {======================================================
131 ** Quicksort
132 ** (based on `Algorithms in MODULA-3', Robert Sedgewick;
133 **  Addison-Wesley, 1993.)
134 */
135
136
137 static void set2 (lua_State *L, int i, int j) {
138   lua_rawseti(L, 1, i);
139   lua_rawseti(L, 1, j);
140 }
141
142 static int sort_comp (lua_State *L, int a, int b) {
143   if (!lua_isnil(L, 2)) {  /* function? */
144     int res;
145     lua_pushvalue(L, 2);
146     lua_pushvalue(L, a-1);  /* -1 to compensate function */
147     lua_pushvalue(L, b-2);  /* -2 to compensate function and `a' */
148     lua_call(L, 2, 1);
149     res = lua_toboolean(L, -1);
150     lua_pop(L, 1);
151     return res;
152   }
153   else  /* a < b? */
154     return lua_lessthan(L, a, b);
155 }
156
157 static void auxsort (lua_State *L, int l, int u) {
158   while (l < u) {  /* for tail recursion */
159     int i, j;
160     /* sort elements a[l], a[(l+u)/2] and a[u] */
161     lua_rawgeti(L, 1, l);
162     lua_rawgeti(L, 1, u);
163     if (sort_comp(L, -1, -2))  /* a[u] < a[l]? */
164       set2(L, l, u);  /* swap a[l] - a[u] */
165     else
166       lua_pop(L, 2);
167     if (u-l == 1) break;  /* only 2 elements */
168     i = (l+u)/2;
169     lua_rawgeti(L, 1, i);
170     lua_rawgeti(L, 1, l);
171     if (sort_comp(L, -2, -1))  /* a[i]<a[l]? */
172       set2(L, i, l);
173     else {
174       lua_pop(L, 1);  /* remove a[l] */
175       lua_rawgeti(L, 1, u);
176       if (sort_comp(L, -1, -2))  /* a[u]<a[i]? */
177         set2(L, i, u);
178       else
179         lua_pop(L, 2);
180     }
181     if (u-l == 2) break;  /* only 3 elements */
182     lua_rawgeti(L, 1, i);  /* Pivot */
183     lua_pushvalue(L, -1);
184     lua_rawgeti(L, 1, u-1);
185     set2(L, i, u-1);
186     /* a[l] <= P == a[u-1] <= a[u], only need to sort from l+1 to u-2 */
187     i = l; j = u-1;
188     for (;;) {  /* invariant: a[l..i] <= P <= a[j..u] */
189       /* repeat ++i until a[i] >= P */
190       while (lua_rawgeti(L, 1, ++i), sort_comp(L, -1, -2)) {
191         if (i>u) luaL_error(L, "invalid order function for sorting");
192         lua_pop(L, 1);  /* remove a[i] */
193       }
194       /* repeat --j until a[j] <= P */
195       while (lua_rawgeti(L, 1, --j), sort_comp(L, -3, -1)) {
196         if (j<l) luaL_error(L, "invalid order function for sorting");
197         lua_pop(L, 1);  /* remove a[j] */
198       }
199       if (j<i) {
200         lua_pop(L, 3);  /* pop pivot, a[i], a[j] */
201         break;
202       }
203       set2(L, i, j);
204     }
205     lua_rawgeti(L, 1, u-1);
206     lua_rawgeti(L, 1, i);
207     set2(L, u-1, i);  /* swap pivot (a[u-1]) with a[i] */
208     /* a[l..i-1] <= a[i] == P <= a[i+1..u] */
209     /* adjust so that smaller half is in [j..i] and larger one in [l..u] */
210     if (i-l < u-i) {
211       j=l; i=i-1; l=i+2;
212     }
213     else {
214       j=i+1; i=u; u=j-2;
215     }
216     auxsort(L, j, i);  /* call recursively the smaller one */
217   }  /* repeat the routine for the larger one */
218 }
219
220 static int luaB_sort (lua_State *L) {
221   int n = aux_getn(L, 1);
222   luaL_checkstack(L, 40, "");  /* assume array is smaller than 2^40 */
223   if (!lua_isnoneornil(L, 2))  /* is there a 2nd argument? */
224     luaL_checktype(L, 2, LUA_TFUNCTION);
225   lua_settop(L, 2);  /* make sure there is two arguments */
226   auxsort(L, 1, n);
227   return 0;
228 }
229
230 /* }====================================================== */
231
232
233 static const luaL_reg tab_funcs[] = {
234   {"concat", str_concat},
235   {"foreach", luaB_foreach},
236   {"foreachi", luaB_foreachi},
237   {"getn", luaB_getn},
238   {"setn", luaB_setn},
239   {"sort", luaB_sort},
240   {"insert", luaB_tinsert},
241   {"remove", luaB_tremove},
242   {NULL, NULL}
243 };
244
245
246 LUALIB_API int luaopen_table (lua_State *L) {
247   luaL_openlib(L, LUA_TABLIBNAME, tab_funcs, 0);
248   return 1;
249 }
250