jbe@0: #include <lua.h>
jbe@0: #include <lauxlib.h>
jbe@8: 
jbe@14: static int seqlua_aux_call(lua_State *L) {  // TODO: avoid
jbe@8:   lua_pushvalue(L, lua_upvalueindex(1));
jbe@8:   lua_pushvalue(L, lua_upvalueindex(2));
jbe@8:   lua_call(L, 1, 1);
jbe@8:   return 1;
jbe@8: }
jbe@0: 
jbe@0: static int seqlua_ipairsaux_raw(lua_State *L) {
jbe@0:   lua_Integer i;
jbe@0:   luaL_checktype(L, 1, LUA_TTABLE);
jbe@0:   i = luaL_checkinteger(L, 2) + 1;
jbe@0:   lua_pushinteger(L, i);
jbe@0:   lua_rawgeti(L, 1, i);  // TODO: Lua 5.3 returns type
jbe@0:   return lua_isnil(L, -1) ? 1 : 2;
jbe@0: }
jbe@0: 
jbe@0: static int seqlua_ipairsaux_meta(lua_State *L) {
jbe@0:   lua_Integer i;
jbe@0:   i = luaL_checkinteger(L, 2) + 1;
jbe@0:   lua_pushinteger(L, i);
jbe@0:   lua_pushinteger(L, i);
jbe@0:   lua_gettable(L, 1);  // TODO: Lua 5.3 returns type
jbe@0:   return lua_isnil(L, -1) ? 1 : 2;
jbe@0: }
jbe@0: 
jbe@0: static int seqlua_ipairsaux_func(lua_State *L) {
jbe@0:   luaL_checktype(L, 1, LUA_TFUNCTION);
jbe@0:   lua_pushinteger(L, luaL_checkinteger(L, 2) + 1);
jbe@0:   lua_insert(L, 1);
jbe@0:   lua_settop(L, 2);
jbe@0:   lua_call(L, 0, LUA_MULTRET);
jbe@0:   if (lua_isnoneornil(L, 2)) {
jbe@0:     lua_settop(L, 0);
jbe@0:     lua_pushnil(L);
jbe@0:     return 1;
jbe@0:   } else {
jbe@0:     return lua_gettop(L);
jbe@0:   }
jbe@0: }
jbe@0: 
jbe@8: static int seqlua_ipairsaux_triplet(lua_State *L) {
jbe@0:   lua_Integer i = lua_tointeger(L, lua_upvalueindex(4)) + 1;
jbe@0:   lua_settop(L, 0);
jbe@0:   lua_pushinteger(L, i);
jbe@0:   lua_replace(L, lua_upvalueindex(4));
jbe@0:   lua_pushinteger(L, i);
jbe@0:   lua_pushvalue(L, lua_upvalueindex(1));
jbe@0:   lua_pushvalue(L, lua_upvalueindex(2));
jbe@0:   lua_pushvalue(L, lua_upvalueindex(3));
jbe@0:   lua_call(L, 2, LUA_MULTRET);
jbe@0:   if (lua_isnoneornil(L, 2)) {
jbe@0:     lua_settop(L, 0);
jbe@0:     lua_pushnil(L);
jbe@0:     return 1;
jbe@0:   } else {
jbe@0:     lua_pushvalue(L, 2);
jbe@0:     lua_replace(L, lua_upvalueindex(3));
jbe@0:     return lua_gettop(L);
jbe@0:   }
jbe@0: }
jbe@0: 
jbe@0: static int seqlua_ipairs(lua_State *L) {
jbe@8:   luaL_checkany(L, 1);  // provides better error message
jbe@0:   if (lua_type(L, 1) == LUA_TFUNCTION) {
jbe@8:     seqlua_ipairs_function:
jbe@0:     if (!lua_isnoneornil(L, 2) || !lua_isnoneornil(L, 3)) {
jbe@0:       lua_settop(L, 3);
jbe@0:       lua_pushinteger(L, 0);
jbe@8:       lua_pushcclosure(L, seqlua_ipairsaux_triplet, 4);
jbe@0:       return 1;
jbe@0:     }
jbe@0:     lua_pushcfunction(L, seqlua_ipairsaux_func);
jbe@8:   } else if (luaL_getmetafield(L, 1, "__call")) {
jbe@8:     if (lua_type(L, -1) != LUA_TFUNCTION) luaL_error(L, "__call is not a function");
jbe@8:     lua_pushvalue(L, 1);
jbe@8:     lua_pushcclosure(L, seqlua_aux_call, 2);
jbe@8:     lua_replace(L, 1);
jbe@8:     goto seqlua_ipairs_function;
jbe@0:   } else if (luaL_getmetafield(L, 1, "__index")) {
jbe@0:     lua_pushcfunction(L, seqlua_ipairsaux_meta);
jbe@0:   } else {
jbe@0:     luaL_checktype(L, 1, LUA_TTABLE);
jbe@0:     lua_pushcfunction(L, seqlua_ipairsaux_raw);
jbe@0:   }
jbe@0:   lua_pushvalue(L, 1);
jbe@0:   lua_pushinteger(L, 0);
jbe@0:   return 3;
jbe@0: }
jbe@0: 
jbe@8: static int seqlua_iteratoraux_raw(lua_State *L) {
jbe@8:   lua_Integer i = lua_tointeger(L, lua_upvalueindex(2)) + 1;
jbe@8:   lua_rawgeti(L, lua_upvalueindex(1), i);
jbe@8:   if (lua_isnil(L, -1)) return 1;
jbe@8:   lua_pushinteger(L, i);
jbe@8:   lua_replace(L, lua_upvalueindex(2));
jbe@8:   return 1;
jbe@8: }
jbe@8: 
jbe@8: static int seqlua_iteratoraux_meta(lua_State *L) {
jbe@8:   lua_Integer i = lua_tointeger(L, lua_upvalueindex(2)) + 1;
jbe@8:   lua_pushinteger(L, i);
jbe@8:   lua_gettable(L, lua_upvalueindex(1));
jbe@8:   if (lua_isnil(L, -1)) return 1;
jbe@8:   lua_pushinteger(L, i);
jbe@8:   lua_replace(L, lua_upvalueindex(2));
jbe@8:   return 1;
jbe@8: }
jbe@8: 
jbe@8: static int seqlua_iteratoraux_triplet(lua_State *L) {
jbe@0:   lua_settop(L, 0);
jbe@0:   lua_pushvalue(L, lua_upvalueindex(1));
jbe@0:   lua_pushvalue(L, lua_upvalueindex(2));
jbe@0:   lua_pushvalue(L, lua_upvalueindex(3));
jbe@0:   lua_call(L, 2, LUA_MULTRET);
jbe@0:   if (lua_isnoneornil(L, 1)) {
jbe@0:     lua_settop(L, 1);
jbe@0:     return 1;
jbe@0:   }
jbe@0:   lua_pushvalue(L, 1);
jbe@0:   lua_replace(L, lua_upvalueindex(3));
jbe@0:   return lua_gettop(L);
jbe@0: }
jbe@0: 
jbe@8: int seqlua_iterator(lua_State *L) {
jbe@8:   luaL_checkany(L, 1);  // provides better error message
jbe@8:   lua_settop(L, 3);
jbe@0:   if (lua_type(L, 1) == LUA_TFUNCTION) {
jbe@8:     seqlua_iterator_function:
jbe@8:     if (lua_isnil(L, 2) && lua_isnil(L, 3)) {
jbe@0:       lua_settop(L, 1);
jbe@0:     } else {
jbe@8:       lua_pushcclosure(L, seqlua_iteratoraux_triplet, 3);
jbe@0:     }
jbe@8:   } else if (luaL_getmetafield(L, 1, "__call")) {
jbe@8:     if (lua_type(L, -1) != LUA_TFUNCTION) luaL_error(L, "__call is not a function");
jbe@8:     lua_pushvalue(L, 1);
jbe@8:     lua_pushcclosure(L, seqlua_aux_call, 2);
jbe@8:     lua_replace(L, 1);
jbe@8:     goto seqlua_iterator_function;
jbe@8:   } else if (luaL_getmetafield(L, 1, "__index")) {
jbe@8:     lua_pushvalue(L, 1);
jbe@8:     lua_pushinteger(L, 0);
jbe@8:     lua_pushcclosure(L, seqlua_iteratoraux_raw, 2);
jbe@0:   } else {
jbe@8:     luaL_checktype(L, 1, LUA_TTABLE);
jbe@8:     lua_pushvalue(L, 1);
jbe@8:     lua_pushinteger(L, 0);
jbe@8:     lua_pushcclosure(L, seqlua_iteratoraux_meta, 2);
jbe@0:   }
jbe@0:   return 1;
jbe@0: }
jbe@0: 
jbe@0: int luaopen_seqlua(lua_State *L) {
jbe@0:   lua_pushcfunction(L, seqlua_ipairs);
jbe@0:   lua_setglobal(L, "ipairs");
jbe@0:   lua_pushcfunction(L, seqlua_iterator);
jbe@0:   lua_setglobal(L, "iterator");
jbe@0:   return 1;
jbe@0: }