Support casting ffi.C functions to cdata numbers.

For example, ffi.cast('intptr_t', ffi.C.malloc) gives the
address of the malloc function.

7 files changed, 76 insertions, 8 deletions

diff --git a/call_x64.h b/call_x64.h
index e0acdee..34ef3eb 100644
--- a/call_x64.h
+++ b/call_x64.h
@@ -1193,7 +1193,7 @@ void compile_function(lua_State* L, cfunction func, int ct_usr, const struct cty
         cfunction f = compile(Dst, L, func, LUA_NOREF);
         /* add a callback as an upval so that the jitted code gets cleaned up when
          * the function gets gc'd */
-        push_callback(L, f);
+        push_callback(L, f, func);
         lua_pushcclosure(L, (lua_CFunction) f, num_upvals+1);
     }
 }
diff --git a/call_x64win.h b/call_x64win.h
index 99306b8..9e12a97 100644
--- a/call_x64win.h
+++ b/call_x64win.h
@@ -1179,7 +1179,7 @@ void compile_function(lua_State* L, cfunction func, int ct_usr, const struct cty
         cfunction f = compile(Dst, L, func, LUA_NOREF);
         /* add a callback as an upval so that the jitted code gets cleaned up when
          * the function gets gc'd */
-        push_callback(L, f);
+        push_callback(L, f, func);
         lua_pushcclosure(L, (lua_CFunction) f, num_upvals+1);
     }
 }
diff --git a/call_x86.dasc b/call_x86.dasc
index d7a7c4c..3d1e83f 100644
--- a/call_x86.dasc
+++ b/call_x86.dasc
@@ -1590,7 +1590,7 @@ void compile_function(lua_State* L, cfunction func, int ct_usr, const struct cty
         cfunction f = compile(Dst, L, func, LUA_NOREF);
         /* add a callback as an upval so that the jitted code gets cleaned up when
          * the function gets gc'd */
-        push_callback(L, f);
+        push_callback(L, f, func);
         lua_pushcclosure(L, (lua_CFunction) f, num_upvals+1);
     }
 }
diff --git a/call_x86.h b/call_x86.h
index 3142aba..0a288c8 100644
--- a/call_x86.h
+++ b/call_x86.h
@@ -1162,7 +1162,7 @@ void compile_function(lua_State* L, cfunction func, int ct_usr, const struct cty
         cfunction f = compile(Dst, L, func, LUA_NOREF);
         /* add a callback as an upval so that the jitted code gets cleaned up when
          * the function gets gc'd */
-        push_callback(L, f);
+        push_callback(L, f, func);
         lua_pushcclosure(L, (lua_CFunction) f, num_upvals+1);
     }
 }
diff --git a/ctype.c b/ctype.c
index 037d86c..1e12fb7 100644
--- a/ctype.c
+++ b/ctype.c
@@ -177,10 +177,11 @@ void* push_cdata(lua_State* L, int ct_usr, const struct ctype* ct)
     return cd+1;
 }
 
-void push_callback(lua_State* L, cfunction f)
+void push_callback(lua_State* L, cfunction luafunc, cfunction cfunc)
 {
-    cfunction* pf = (cfunction*) lua_newuserdata(L, sizeof(cfunction));
-    *pf = f;
+    cfunction* pf = (cfunction*) lua_newuserdata(L, 2 * sizeof(cfunction));
+    pf[0] = luafunc;
+    pf[1] = cfunc;
 
     push_upval(L, &callback_mt_key);
     lua_setmetatable(L, -2);
diff --git a/ffi.c b/ffi.c
index 86a22cb..fa989e9 100644
--- a/ffi.c
+++ b/ffi.c
@@ -160,10 +160,13 @@ static int64_t check_intptr(lua_State* L, int idx, void* p, struct ctype* ct)
     }
 }
 
+static int get_cfunction_address(lua_State* L, int idx, cfunction* addr);
+
 #define TO_NUMBER(TYPE, ALLOW_POINTERS)                                     \
     TYPE ret = 0;                                                           \
     void* p;                                                                \
     struct ctype ct;                                                        \
+    cfunction f;                                                            \
                                                                             \
     switch (lua_type(L, idx)) {                                             \
     case LUA_TBOOLEAN:                                                      \
@@ -188,6 +191,16 @@ static int64_t check_intptr(lua_State* L, int idx, void* p, struct ctype* ct)
         ret = (TYPE) (intptr_t) lua_topointer(L, idx);                      \
         break;                                                              \
                                                                             \
+    case LUA_TFUNCTION:                                                     \
+        if (!ALLOW_POINTERS) {                                              \
+            type_error(L, idx, #TYPE, 0, NULL);                             \
+        }                                                                   \
+        if (!get_cfunction_address(L, idx, &f)) {                           \
+            type_error(L, idx, #TYPE, 0, NULL);                             \
+        }                                                                   \
+        ret = (TYPE) (intptr_t) f;                                          \
+        break;                                                              \
+                                                                            \
     case LUA_TUSERDATA:                                                     \
         p = to_cdata(L, idx, &ct);                                          \
                                                                             \
@@ -622,6 +635,48 @@ err:
     return NULL;
 }
 
+/**
+ * gets the address of the wrapped C function for the lua function value at idx
+ * and returns 1 if it exists; otherwise returns 0 and nothing is pushed */
+static int get_cfunction_address(lua_State* L, int idx, cfunction* addr)
+{
+    if (!lua_isfunction(L, idx)) return 0;
+
+    int top = lua_gettop(L);
+
+    // Get the last upvalue
+    int n = 2;
+    while (lua_getupvalue(L, idx, n)) {
+        lua_pop(L, 1);
+        n++;
+    }
+
+    if (!lua_getupvalue(L, idx, n - 1))
+        return 0;
+
+    if (!lua_isuserdata(L, -1) || !lua_getmetatable(L, -1)) {
+        lua_pop(L, 1);
+        return 0;
+    }
+
+    push_upval(L, &callback_mt_key);
+    if (!lua_rawequal(L, -1, -2)) {
+        lua_pop(L, 3);
+        return 0;
+    }
+
+    /* stack is:
+     * userdata upval
+     * metatable
+     * callback_mt
+     */
+
+    cfunction* f = lua_touserdata(L, -3);
+    *addr = f[1];
+    lua_pop(L, 3);
+    return 1;
+}
+
 /* to_cfunction converts a value at idx with usr table at to_usr and type tt
  * into a function. Leaves the stack unchanged. */
 static cfunction check_cfunction(lua_State* L, int idx, int to_usr, const struct ctype* tt, int check_pointers)
@@ -636,6 +691,10 @@ static cfunction check_cfunction(lua_State* L, int idx, int to_usr, const struct
 
     switch (lua_type(L, idx)) {
     case LUA_TFUNCTION:
+        if (get_cfunction_address(L, idx, &f)) {
+            return f;
+        }
+
         /* Function cdatas are pinned and must be manually cleaned up by
          * calling func:free(). */
         push_upval(L, &callbacks_key);
@@ -1133,6 +1192,14 @@ static int do_new(lua_State* L, int is_cast)
     /* don't push a callback when we have a c function, as cb:set needs a
      * compiled callback from a lua function to work */
     if (!ct.pointers && ct.type == FUNCTION_PTR_TYPE && (lua_isnil(L, 2) || lua_isfunction(L, 2))) {
+        // Get the bound C function if this is a ffi lua function
+        cfunction func;
+        if (get_cfunction_address(L, 2, &func)) {
+            p = push_cdata(L, -1, &ct);
+            *(cfunction*) p = func;
+            return 1;
+        }
+
         /* Function cdatas are pinned and must be manually cleaned up by
          * calling func:free(). */
         compile_callback(L, 2, -1, &ct);
diff --git a/ffi.h b/ffi.h
index 4991a3a..2c27966 100644
--- a/ffi.h
+++ b/ffi.h
@@ -415,7 +415,7 @@ static float cimagf(complex_float c) {
 void set_defined(lua_State* L, int ct_usr, struct ctype* ct);
 struct ctype* push_ctype(lua_State* L, int ct_usr, const struct ctype* ct);
 void* push_cdata(lua_State* L, int ct_usr, const struct ctype* ct); /* called from asm */
-void push_callback(lua_State* L, cfunction f);
+void push_callback(lua_State* L, cfunction luafunc, cfunction cfunc);
 void check_ctype(lua_State* L, int idx, struct ctype* ct);
 void* to_cdata(lua_State* L, int idx, struct ctype* ct);
 void* check_cdata(lua_State* L, int idx, struct ctype* ct);