feat(lua): add for_each to expr Lua API

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/bindings/lua/expr.cpp

@@ -14,6 +14,7 @@ Author: Leonardo de Moura
 #include "kernel/expr.h"
 #include "kernel/abstract.h"
 #include "kernel/formatter.h"
+#include "kernel/for_each.h"
 #include "library/expr_lt.h"
 #include "bindings/lua/util.h"
 #include "bindings/lua/name.h"
@@ -286,6 +287,26 @@ static int expr_pred(lua_State * L) {
     return 1;
 }
 
+static int expr_for_each(lua_State * L) {
+    expr & e = to_nonnull_expr(L, 1);
+    if (!lua_isfunction(L, 2))
+        throw exception("arg #2 must be a function");
+    auto f = [&](expr const & a, unsigned offset) {
+        lua_pushvalue(L, 2); // push user function
+        push_expr(L, a);
+        lua_pushinteger(L, offset);
+        pcall(L, 2, 1, 0);
+        bool r = true;
+        if (lua_isboolean(L, -1))
+            r = lua_toboolean(L, -1);
+        lua_pop(L, 1);
+        return r;
+    };
+    for_each_fn<decltype(f)> proc(f);
+    proc(e);
+    return 0;
+}
+
 static const struct luaL_Reg expr_m[] = {
     {"__gc",           expr_gc}, // never throws
     {"__tostring",     safe_function<expr_tostring>},
@@ -308,6 +329,7 @@ static const struct luaL_Reg expr_m[] = {
     {"args",           safe_function<expr_args>},
     {"num_args",       safe_function<expr_num_args>},
     {"arg",            safe_function<expr_arg>},
+    {"for_each",       safe_function<expr_for_each>},
     {0, 0}
 };
 

src/bindings/lua/leanlua_state.cpp

@@ -285,9 +285,7 @@ int state_dostring(lua_State * L) {
 
     copy_values(L, first, last, S->m_state);
 
-    result = lua_pcall(S->m_state, first > last ? 0 : last - first + 1, LUA_MULTRET, 0);
-    if (result)
-        throw lua_exception(lua_tostring(S->m_state, -1));
+    pcall(S->m_state, first > last ? 0 : last - first + 1, LUA_MULTRET, 0);
 
     int sz_after = lua_gettop(S->m_state);
 

src/bindings/lua/util.cpp

@@ -59,9 +59,7 @@ size_t objlen(lua_State * L, int idx) {
 }
 
 static void exec(lua_State * L) {
-    int result = lua_pcall(L, 0, LUA_MULTRET, 0);
-    if (result)
-        throw lua_exception(lua_tostring(L, -1));
+    pcall(L, 0, LUA_MULTRET, 0);
 }
 
 void dofile(lua_State * L, char const * fname) {
@@ -78,6 +76,12 @@ void dostring(lua_State * L, char const * str) {
     exec(L);
 }
 
+void pcall(lua_State * L, int nargs, int nresults, int errorfun) {
+    int result = lua_pcall(L, nargs, nresults, errorfun);
+    if (result)
+        throw lua_exception(lua_tostring(L, -1));
+}
+
 int safe_function_wrapper(lua_State * L, lua_CFunction f){
     static thread_local std::string _error_msg;
     char const * error_msg;

src/bindings/lua/util.h

@@ -14,6 +14,7 @@ int load(lua_State * L, lua_Reader reader, void * data, char const * source);
 size_t objlen(lua_State * L, int idx);
 void dofile(lua_State * L, char const * fname);
 void dostring(lua_State * L, char const * str);
+void pcall(lua_State * L, int nargs, int nresults, int errorfun);
 /**
    \brief Wrapper for invoking function f, and catching Lean exceptions.
 */

tests/lean/lua10.lean

@@ -0,0 +1,19 @@
+
+Variables x1 x2 x3 : Bool
+Definition F : Bool := x1 /\ (x2 \/ x3)
+
+(**
+ local env    = get_environment()
+ local F      = env:find_object("F"):get_value()
+ print(F)
+
+ function expr_size(e)
+    local r = 0
+    e:for_each(function(e, offset) r = r + 1 end)
+    return r
+ end
+
+ print(expr_size(F))
+
+**)
+

tests/lean/lua10.lean.expected.out

@@ -0,0 +1,8 @@
+  Set: pp::colors
+  Set: pp::unicode
+  Assumed: x1
+  Assumed: x2
+  Assumed: x3
+  Defined: F
+x1 ∧ (x2 ∨ x3)
+7