allow red functions to be passed as arguments and return values

spy/backend/c/cbackend.py

@@ -13,7 +13,7 @@
 from spy.fqn import FQN
 from spy.highlight import highlight_src
 from spy.vm.cell import W_Cell
-from spy.vm.function import W_ASTFunc
+from spy.vm.function import W_ASTFunc, W_FuncType
 from spy.vm.modules.unsafe.ptr import W_MemLocType
 from spy.vm.object import W_Object, W_Type
 from spy.vm.primitive import W_I32
@@ -145,7 +145,7 @@ def foo() -> i32:
             modname = fqn.modname
             w_mod = self.vm.modules_w[modname]
             if w_mod.filepath is None and not isinstance(
-                w_obj, (W_MemLocType, W_StructType)
+                w_obj, (W_MemLocType, W_StructType, W_FuncType)
             ):
                 continue
 
@@ -274,19 +274,31 @@ def get_type_deps(self, fqn: FQN) -> list[FQN]:
             # the forward-decl section: it needs T's own typedef to appear
             # first.
             return [w_type.w_itemT.fqn]
+        if isinstance(w_type, W_FuncType):
+            # A function-pointer typedef references its return and param types
+            # by name, so any by-value struct types must be defined first.
+            deps: list[FQN] = []
+            seen: set[FQN] = set()
+            for w_dep in [w_type.w_restype] + [p.w_T for p in w_type.params]:
+                if isinstance(w_dep, W_StructType):
+                    d = w_dep.fqn
+                    if d != fqn and d not in seen:
+                        seen.add(d)
+                        deps.append(d)
+            return deps
         if not isinstance(w_type, W_StructType) or not w_type.is_defined():
             return []
-        deps: list[FQN] = []
-        seen: set[FQN] = set()
+        struct_deps: list[FQN] = []
+        struct_seen: set[FQN] = set()
         for field in w_type.iterfields_w():
             w_fieldT = field.w_T
             if not isinstance(w_fieldT, W_StructType):
                 continue
             dep_fqn = w_fieldT.fqn
-            if dep_fqn != fqn and dep_fqn not in seen:
-                seen.add(dep_fqn)
-                deps.append(dep_fqn)
-        return deps
+            if dep_fqn != fqn and dep_fqn not in struct_seen:
+                struct_seen.add(dep_fqn)
+                struct_deps.append(dep_fqn)
+        return struct_deps
 
     def topo_sort_structdefs(
         self, content: list[tuple[FQN, W_Type]]

spy/backend/c/cstructwriter.py

@@ -7,6 +7,7 @@
 from spy.backend.c.context import C_Type, Context
 from spy.fqn import FQN
 from spy.textbuilder import TextBuilder
+from spy.vm.function import W_FuncType
 from spy.vm.modules.unsafe.ptr import W_PtrType, W_RefType
 from spy.vm.object import W_Type
 from spy.vm.struct import W_StructType
@@ -98,6 +99,8 @@ def emit_content(self) -> None:
             assert fqn == w_type.fqn  # sanity check
             if isinstance(w_type, W_StructType):
                 self.emit_StructType(fqn, w_type)
+            elif isinstance(w_type, W_FuncType):
+                self.emit_FuncType(fqn, w_type)
             elif isinstance(w_type, W_PtrType):
                 self.emit_PtrType(fqn, w_type)
             elif isinstance(w_type, W_RefType):
@@ -155,6 +158,15 @@ def emit_StructType(self, fqn: FQN, w_st: W_StructType) -> None:
         tb.wl("};")
         tb.wl("")
 
+    def emit_FuncType(self, fqn: FQN, w_ft: W_FuncType) -> None:
+        c_name = w_ft.fqn.c_name
+        c_ret = self.ctx.w2c(w_ft.w_restype)
+        if w_ft.params:
+            c_params = ", ".join(str(self.ctx.w2c(p.w_T)) for p in w_ft.params)
+        else:
+            c_params = "void"
+        self.tbh_fwdecl.wl(f"typedef {c_ret} (*{c_name})({c_params});")
+
     def emit_PtrType(self, fqn: FQN, w_ptrtype: W_PtrType) -> None:
         c_ptrtype = C_Type(w_ptrtype.fqn.c_name)
         w_itemT = w_ptrtype.w_itemT

spy/tests/compiler/out_of_tree/mymod/__init__.py

@@ -1,7 +1,10 @@
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Annotated
 
 from spy.build.build_info import BuildInfo, BuildTarget, BuildType
+from spy.vm.b import B
+from spy.vm.function import FuncParam, W_Func, W_FuncType
+from spy.vm.primitive import W_I32
 from spy.vm.registry import ModuleRegistry
 from spy.vm.str import W_Str
 
@@ -12,6 +15,10 @@
 
 MODULE = ModuleRegistry("mymod")
 
+# red def(i32) -> i32: a callback that takes one i32 and returns i32
+_w_cb_type = W_FuncType.new([FuncParam(B.w_i32, "simple")], B.w_i32)
+CB = Annotated[W_Func, _w_cb_type]
+
 
 def build_info(target: BuildTarget, build_type: BuildType) -> BuildInfo:
     return BuildInfo(
@@ -24,3 +31,10 @@ def build_info(target: BuildTarget, build_type: BuildType) -> BuildInfo:
 @MODULE.builtin_func
 def w_get_name(vm: "SPyVM") -> W_Str:
     return vm.wrap("hello from mymod")
+
+
+@MODULE.builtin_func
+def w_run_callback(vm: "SPyVM", w_cb: CB, w_x: W_I32) -> W_I32:
+    # At interp level, w_cb is the W_ASTFunc itself. Call it directly.
+    assert isinstance(w_cb, W_Func)
+    return vm.fast_call(w_cb, [w_x])  # type: ignore[return-value]

spy/tests/compiler/out_of_tree/mymod/mymod.c

@@ -10,3 +10,7 @@ spy_StrObject *spy_mymod$get_name(void) {
     memcpy(spy_StrObject_UTF8(s), NAME, n);
     return s;
 }
+
+int32_t spy_mymod$run_callback(int32_t (*cb)(int32_t), int32_t x) {
+    return cb(x);
+}

spy/tests/compiler/out_of_tree/mymod/mymod.h

@@ -2,7 +2,13 @@
 #define MYMOD_H
 
 #include <spy.h>
+#include <stdint.h>
 
 spy_StrObject *spy_mymod$get_name(void);
 
+// Accepts a function-pointer callback and calls it with x. The functype
+// typedef emitted in spy_structdefs.h expands to int32_t (*)(int32_t), so we
+// declare the parameter with the raw function pointer type.
+int32_t spy_mymod$run_callback(int32_t (*cb)(int32_t), int32_t x);
+
 #endif /* MYMOD_H */

spy/tests/compiler/out_of_tree/test_out_of_tree.py

@@ -37,3 +37,43 @@ def foo() -> str:
             return get_name()
         """)
         assert mod.foo() == "hello from mymod"
+
+    def test_c_callback(self):
+        self.vm = SPyVM(extra_vm_modules=[str(MYMOD_PATH)])
+        self.vm.path.append(str(self.tmpdir))
+        mod = self.compile("""
+        from mymod import run_callback
+
+        def double(x: i32) -> i32:
+            return x * 2
+
+        def run() -> i32:
+            return run_callback(double, 21)
+        """)
+        assert mod.run() == 42
+
+    def test_c_callback_blue_factory(self):
+        # A @blue function generates a specialised red callback per compile-time
+        # constant; each becomes a distinct C symbol that C can call back through.
+        self.vm = SPyVM(extra_vm_modules=[str(MYMOD_PATH)])
+        self.vm.path.append(str(self.tmpdir))
+        mod = self.compile("""
+        from mymod import run_callback
+
+        @functype
+        def CB(x: i32) -> i32:
+            pass
+
+        @blue
+        def make_adder(n: i32) -> CB:
+            def adder(x: i32) -> i32:
+                return x + n
+            return adder
+
+        add5 = make_adder(5)
+        add10 = make_adder(10)
+
+        def run() -> i32:
+            return run_callback(add5, 1) + run_callback(add10, 1)
+        """)
+        assert mod.run() == 17  # (1+5) + (1+10)

spy/tests/compiler/test_functype.py

@@ -0,0 +1,154 @@
+from spy.tests.support import CompilerTest, expect_errors, only_interp
+from spy.vm.b import TYPES, B
+from spy.vm.function import FuncParam, W_FuncType
+
+
+class TestFuncType(CompilerTest):
+    @only_interp
+    def test_type_construction(self):
+        # After compiling a @functype-decorated def, the named variable holds
+        # the W_FuncType for that signature.
+        mod = self.compile("""
+        @functype
+        def CB(a: i32, b: i32) -> i32:
+            pass
+        """)
+        w_CB = mod.w_mod.getattr("CB")
+        assert isinstance(w_CB, W_FuncType)
+        assert w_CB.w_restype is B.w_i32
+        assert [p.w_T for p in w_CB.params] == [B.w_i32, B.w_i32]
+        assert w_CB.color == "red"
+        assert w_CB.kind == "plain"
+        # FQN puts params first, restype last
+        assert str(w_CB.fqn) == "builtins::def[i32, i32, i32]"
+
+    @only_interp
+    def test_type_identical_to_red_def(self):
+        # The W_FuncType from @functype must be the SAME interned object as a
+        # directly constructed W_FuncType with the same signature — so passing
+        # a matching red function where that type is expected works by identity.
+        mod = self.compile("""
+        @functype
+        def CB(x: i32) -> i32:
+            pass
+        """)
+        w_CB = mod.w_mod.getattr("CB")
+        w_T_direct = W_FuncType.new(
+            [FuncParam(B.w_i32, "simple")], B.w_i32, color="red", kind="plain"
+        )
+        assert w_CB is w_T_direct
+
+    def test_functype_decorator(self):
+        self.compile("""
+        @functype
+        def CB(x: i32, y: i32) -> i32:
+            pass
+
+        def apply(cb: CB, x: i32, y: i32) -> i32:
+            return x + y
+
+        def my_add(a: i32, b: i32) -> i32:
+            return a + b
+
+        def run() -> i32:
+            return apply(my_add, 3, 4)
+        """)
+
+    def test_functype_at_call_site(self):
+        mod = self.compile("""
+        @functype
+        def CB(x: i32, y: i32) -> i32:
+            pass
+
+        def apply(cb: CB, x: i32, y: i32) -> i32:
+            return x + y
+
+        def my_add(a: i32, b: i32) -> i32:
+            return a + b
+
+        def run() -> i32:
+            return apply(my_add, 3, 4)
+        """)
+        assert mod.run() == 7
+
+    def test_signature_mismatch_wrong_ret(self):
+        src = """
+        @functype
+        def CB(x: i32) -> bool:
+            pass
+
+        def my_fn(x: i32) -> i32:
+            return x
+
+        def bad() -> CB:
+            return my_fn
+        """
+        errors = expect_errors("mismatched types")
+        self.compile_raises(src, "bad", errors)
+
+    def test_signature_mismatch_wrong_argcount(self):
+        src = """
+        @functype
+        def CB(x: i32, y: i32) -> i32:
+            pass
+
+        def my_fn(x: i32) -> i32:
+            return x
+
+        def bad() -> CB:
+            return my_fn
+        """
+        errors = expect_errors("mismatched types")
+        self.compile_raises(src, "bad", errors)
+
+    def test_signature_mismatch_wrong_arg_type(self):
+        src = """
+        @functype
+        def CB(x: f64) -> i32:
+            pass
+
+        def my_fn(x: i32) -> i32:
+            return x
+
+        def bad() -> CB:
+            return my_fn
+        """
+        errors = expect_errors("mismatched types")
+        self.compile_raises(src, "bad", errors)
+
+    def test_blue_func_rejected(self):
+        # A @blue function's functype has a different FQN than a red functype,
+        # so passing it where a @functype type is expected fails with a type mismatch.
+        src = """
+        @functype
+        def CB(x: i32) -> i32:
+            pass
+
+        @blue
+        def my_fn(x: i32) -> i32:
+            return x
+
+        def get_cb() -> CB:
+            return my_fn
+        """
+        errors = expect_errors("mismatched types")
+        self.compile_raises(src, "get_cb", errors)
+
+    def test_blue_factory(self):
+        # A @blue function can generate a red callback that captures compile-time
+        # constants. After redshifting, captures are inlined so each becomes a
+        # standalone C symbol.
+        self.compile("""
+        @functype
+        def CB(x: i32) -> i32:
+            pass
+
+        @blue
+        def make_adder(n: i32) -> CB:
+            def adder(x: i32) -> i32:
+                return x + n
+            return adder
+
+        add5: CB = make_adder(5)
+        add10: CB = make_adder(10)
+        """)

spy/vm/modules/builtins.py

@@ -34,6 +34,31 @@ def w_STATIC_TYPE(vm: "SPyVM", wam_obj: W_MetaArg) -> W_OpSpec:
     return W_OpSpec.const(wam_obj.w_static_T)
 
 
+@BUILTINS.builtin_func(color="blue", kind="metafunc")
+def w_functype(vm: "SPyVM", wam_func: W_MetaArg) -> W_OpSpec:
+    """
+    Decorator that gives a name to a function type.
+
+    Used like:
+        @functype
+        def CB(x: i32, y: i32) -> i32:
+            pass
+
+    CB then holds the W_FuncType for `def(i32, i32) -> i32`, which is
+    identical (by identity) to the functype of any matching red function.
+    """
+    w_T = wam_func.w_static_T
+    if not isinstance(w_T, W_FuncType):
+        t = w_T.fqn.human_name(vm)
+        raise SPyError.simple(
+            "W_TypeError",
+            f"functype expects a function, got `{t}`",
+            f"this is `{t}`",
+            wam_func.loc,
+        )
+    return W_OpSpec.const(w_T)
+
+
 @BUILTINS.builtin_func(color="blue", kind="metafunc")
 def w_print(vm: "SPyVM", *args_wam: W_MetaArg) -> W_OpSpec:
     vm.import_("_print")

spy/vm/vm.py

@@ -514,10 +514,9 @@ def make_fqn_const(self, w_val: W_Object) -> FQN:
             assert w_val.fqn not in self.globals_w
 
         elif isinstance(w_val, W_Type):
-            # for now types are only builtin so they must have an unique fqn,
-            # we might need to change this when we introduce custom types
             fqn = w_val.fqn
-            assert w_val.fqn not in self.globals_w
+            if fqn in self.globals_w:
+                return fqn
         else:
             w_T = self.dynamic_type(w_val)
             T = w_T.fqn.human_name(self)