feat: Extend comptime arguments to arbitrary non-linear types

guppylang/tys/parsing.py

@@ -36,6 +36,7 @@
     NonLinearOwnedError,
 )
 from guppylang.tys.param import ConstParam, Parameter, TypeParam
+from guppylang.tys.subst import BoundVarFinder
 from guppylang.tys.ty import (
     FuncInput,
     FunctionType,
@@ -365,16 +366,15 @@ def type_with_flags_from_ast(
                 flags |= InputFlags.Comptime
                 if not ty.copyable or not ty.droppable:
                     raise GuppyError(LinearComptimeError(node.right, ty))
-
-                # TODO: For now we can only do `nat` comptime args since they lower to
-                #  Hugr bounded nats. Extend to arbitrary types via monomorphization.
-                #  See https://github.com/CQCL/guppylang/issues/1008
-                if (
-                    not isinstance(ty, NumericType)
-                    or not ty.kind == NumericType.Kind.Nat
-                ):
+                # For now, we don't allow comptime annotations on generic inputs
+                # TODO: In the future we might want to allow stuff like
+                #  `def foo[T: (Copy, Discard](x: T @comptime)`.
+                #   Also see the todo in `parse_parameter`.
+                var_finder = BoundVarFinder()
+                ty.visit(var_finder)
+                if var_finder.bound_vars:
                     raise GuppyError(
-                        UnsupportedError(node.right, f"`{ty}` comptime arguments")
+                        UnsupportedError(node.left, "Generic comptime arguments")
                     )
             case _:
                 raise GuppyError(InvalidFlagError(node.right))

guppylang/tys/subst.py

@@ -4,7 +4,7 @@
 
 from guppylang.error import InternalGuppyError
 from guppylang.tys.arg import Argument, ConstArg, TypeArg
-from guppylang.tys.common import Transformer
+from guppylang.tys.common import Transformer, Visitor
 from guppylang.tys.const import BoundConstVar, Const, ConstBase, ExistentialConstVar
 from guppylang.tys.ty import (
     BoundTypeVar,
@@ -13,7 +13,7 @@
     Type,
     TypeBase,
 )
-from guppylang.tys.var import ExistentialVar
+from guppylang.tys.var import BoundVar, ExistentialVar
 
 Subst = dict[ExistentialVar, Type | Const]
 Inst = Sequence[Argument]
@@ -82,3 +82,26 @@ def _transform_FunctionType(self, ty: FunctionType) -> Type | None:
         if ty.parametrized:
             raise InternalGuppyError("Tried to instantiate under binder")
         return None
+
+
+class BoundVarFinder(Visitor):
+    """Type visitor that looks for occurrences of bound variables."""
+
+    bound_vars: set[BoundVar]
+
+    def __init__(self) -> None:
+        self.bound_vars = set()
+
+    @functools.singledispatchmethod
+    def visit(self, ty: Any) -> bool:  # type: ignore[override]
+        return False
+
+    @visit.register
+    def _transform_BoundTypeVar(self, ty: BoundTypeVar) -> bool:
+        self.bound_vars.add(ty)
+        return False
+
+    @visit.register
+    def _transform_BoundConstVar(self, c: BoundConstVar) -> bool:
+        self.bound_vars.add(c)
+        return False

tests/error/comptime_arg_errors/generic.err

@@ -1,8 +1,8 @@
-Error: Unsupported (at $FILE:8:15)
+Error: Unsupported (at $FILE:8:11)
   | 
 6 | 
 7 | @guppy
-8 | def main(q: T @comptime) -> None:
-  |                ^^^^^^^^ `T` comptime arguments are not supported
+8 | def foo(q: T @comptime) -> T:
+  |            ^ Generic comptime arguments are not supported
 
 Guppy compilation failed due to 1 previous error

tests/error/comptime_arg_errors/generic.py

@@ -1,12 +1,17 @@
 from guppylang import guppy
 from guppylang.std.builtins import comptime
 
-T = guppy.type_var("T")
+T = guppy.type_var("T", copyable=True, droppable=True)
 
 
 @guppy
-def main(q: T @comptime) -> None:
-    pass
+def foo(q: T @comptime) -> T:
+    return T
+
+
+@guppy
+def main() -> int:
+    return foo(42)
 
 
 guppy.compile(main)

tests/integration/test_comptime_args.py

@@ -2,7 +2,7 @@
 from guppylang.std.builtins import nat, comptime, array
 
 
-def test_basic(validate):
+def test_basic_nat(validate):
     @guppy
     def foo(n: nat @ comptime) -> nat:
         return nat(n + 1)
@@ -14,6 +14,42 @@ def main() -> nat:
     validate(guppy.compile(main))
 
 
+def test_basic_int(validate):
+    @guppy
+    def foo(n: int @ comptime) -> int:
+        return n + 1
+
+    @guppy
+    def main() -> int:
+        return foo(42)
+
+    validate(guppy.compile(main))
+
+
+def test_basic_float(validate):
+    @guppy
+    def foo(f: float @ comptime) -> float:
+        return f + 1.5
+
+    @guppy
+    def main() -> float:
+        return foo(42.0)
+
+    validate(guppy.compile(main))
+
+
+def test_basic_bool(validate):
+    @guppy
+    def foo(b: bool @ comptime) -> bool:
+        return not b
+
+    @guppy
+    def main() -> bool:
+        return foo(True)
+
+    validate(guppy.compile(main))
+
+
 def test_multiple(validate):
     @guppy
     def foo(