refactor!: Store function argument names in FuncInput

Author: mark-koch

guppylang-internals/src/guppylang_internals/checker/expr_checker.py

@@ -496,12 +496,7 @@ def visit_Attribute(self, node: ast.Attribute) -> tuple[ast.expr, Type]:
             )
             # Make a closure by partially applying the `self` argument
             # TODO: Try to infer some type args based on `self`
-            result_ty = FunctionType(
-                func.ty.inputs[1:],
-                func.ty.output,
-                func.ty.input_names[1:] if func.ty.input_names else None,
-                func.ty.params,
-            )
+            result_ty = FunctionType(func.ty.inputs[1:], func.ty.output, func.ty.params)
             return with_loc(node, PartialApply(func=name, args=[node.value])), result_ty
         raise GuppyTypeError(AttributeNotFoundError(attr_span, ty, node.attr))
 

guppylang-internals/src/guppylang_internals/checker/func_checker.py

@@ -134,12 +134,12 @@ def check_global_func_def(
     """Type checks a top-level function definition."""
     args = func_def.args.args
     returns_none = isinstance(ty.output, NoneType)
-    assert ty.input_names is not None
+    assert all(inp.name is not None for inp in ty.inputs)
 
     cfg = CFGBuilder().build(func_def.body, returns_none, globals)
     inputs = [
-        Variable(x, inp.ty, loc, inp.flags, is_func_input=True)
-        for x, inp, loc in zip(ty.input_names, ty.inputs, args, strict=True)
+        Variable(cast(str, inp.name), inp.ty, loc, inp.flags, is_func_input=True)
+        for inp, loc in zip(ty.inputs, args, strict=True)
         # Comptime inputs are turned into generic args, so are not included here
         if InputFlags.Comptime not in inp.flags
     ]
@@ -194,10 +194,14 @@ def check_nested_func_def(
 
     # Construct inputs for checking the body CFG
     inputs = [v for v, _ in captured.values()] + [
-        Variable(x, inp.ty, func_def.args.args[i], inp.flags, is_func_input=True)
-        for i, (x, inp) in enumerate(
-            zip(func_ty.input_names, func_ty.inputs, strict=True)
+        Variable(
+            cast(str, inp.name),
+            inp.ty,
+            func_def.args.args[i],
+            inp.flags,
+            is_func_input=True,
         )
+        for i, inp in enumerate(func_ty.inputs)
         # Comptime inputs are turned into generic args, so are not included here
         if InputFlags.Comptime not in inp.flags
     ]
@@ -305,7 +309,6 @@ def check_signature(
     return FunctionType(
         inputs,
         output,
-        input_names,
         sorted(param_var_mapping.values(), key=lambda v: v.idx),
     )
 

guppylang-internals/src/guppylang_internals/definition/struct.py

@@ -273,13 +273,16 @@ def compile(self, args: list[Wire]) -> list[Wire]:
 
         constructor_sig = FunctionType(
             inputs=[
-                FuncInput(f.ty, InputFlags.Owned if f.ty.linear else InputFlags.NoFlags)
+                FuncInput(
+                    f.ty,
+                    InputFlags.Owned if f.ty.linear else InputFlags.NoFlags,
+                    f.name,
+                )
                 for f in self.fields
             ],
             output=StructType(
                 defn=self, args=[p.to_bound(i) for i, p in enumerate(self.params)]
             ),
-            input_names=[f.name for f in self.fields],
             params=self.params,
         )
         constructor_def = CustomFunctionDef(

guppylang-internals/src/guppylang_internals/tys/parsing.py

@@ -289,7 +289,7 @@ def check_function_arg(
         ctx.param_var_mapping[name] = ConstParam(
             len(ctx.param_var_mapping), name, ty, from_comptime_arg=True
         )
-    return FuncInput(ty, flags)
+    return FuncInput(ty, flags, name)
 
 
 if sys.version_info >= (3, 12):

guppylang-internals/src/guppylang_internals/tys/printing.py

@@ -165,10 +165,10 @@ def _wrap(s: str, inside_row: bool) -> str:
 
 def signature_to_str(name: str, sig: FunctionType) -> str:
     """Displays a function signature in Python syntax including the function name."""
-    assert sig.input_names is not None
+    assert all(inp.name is not None for inp in sig.inputs)
     s = f"def {name}("
     s += ", ".join(
-        f"{name}: {inp.ty}{TypePrinter._print_flags(inp.flags)}"
-        for name, inp in zip(sig.input_names, sig.inputs, strict=True)
+        f"{inp.name}: {inp.ty}{TypePrinter._print_flags(inp.flags)}"
+        for inp in sig.inputs
     )
     return s + ") -> " + str(sig.output)

guppylang-internals/src/guppylang_internals/tys/ty.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from enum import Enum, Flag, auto
 from functools import cached_property, total_ordering
 from typing import TYPE_CHECKING, ClassVar, TypeAlias, cast
@@ -389,6 +389,10 @@ class FuncInput:
     ty: "Type"
     flags: InputFlags
 
+    #: Name of this input, or `None` if it is an unnamed argument (e.g. inside a
+    #: higher-order `Callable` type)
+    name: str | None = field(default=None, compare=False)
+
 
 @dataclass(frozen=True, init=False)
 class FunctionType(ParametrizedTypeBase):
@@ -397,7 +401,6 @@ class FunctionType(ParametrizedTypeBase):
     inputs: Sequence[FuncInput]
     output: "Type"
     params: Sequence[Parameter]
-    input_names: Sequence[str] | None
     comptime_args: Sequence[ConstArg]
 
     args: Sequence[Argument] = field(init=False)
@@ -411,7 +414,6 @@ def __init__(
         self,
         inputs: Sequence[FuncInput],
         output: "Type",
-        input_names: Sequence[str] | None = None,
         params: Sequence[Parameter] | None = None,
         comptime_args: Sequence[ConstArg] | None = None,
     ) -> None:
@@ -433,7 +435,6 @@ def __init__(
         object.__setattr__(self, "comptime_args", comptime_args)
         object.__setattr__(self, "inputs", inputs)
         object.__setattr__(self, "output", output)
-        object.__setattr__(self, "input_names", input_names or [])
         object.__setattr__(self, "params", params)
 
     @property
@@ -449,6 +450,16 @@ def bound_vars(self) -> set[BoundVar]:
             return set()
         return super().bound_vars
 
+    @cached_property
+    def input_names(self) -> Sequence[str] | None:
+        """Names of all inputs or `None` if there are unnamed inputs."""
+        names: list[str] = []
+        for inp in self.inputs:
+            if inp.name is None:
+                return None
+            names.append(inp.name)
+        return names
+
     def cast(self) -> "Type":
         """Casts an implementor of `TypeBase` into a `Type`."""
         return self
@@ -507,12 +518,8 @@ def visit(self, visitor: Visitor) -> None:
     def transform(self, transformer: Transformer) -> "Type":
         """Accepts a transformer on this type."""
         return transformer.transform(self) or FunctionType(
-            [
-                FuncInput(inp.ty.transform(transformer), inp.flags)
-                for inp in self.inputs
-            ],
+            [replace(inp, ty=inp.ty.transform(transformer)) for inp in self.inputs],
             self.output.transform(transformer),
-            self.input_names,
             self.params,
         )
 
@@ -542,9 +549,8 @@ def instantiate_partial(self, args: "PartialInst") -> "FunctionType":
 
         inst = Instantiator(full_inst)
         return FunctionType(
-            [FuncInput(inp.ty.transform(inst), inp.flags) for inp in self.inputs],
+            [replace(inp, ty=inp.ty.transform(inst)) for inp in self.inputs],
             self.output.transform(inst),
-            self.input_names,
             remaining_params,
             # Comptime type arguments also need to be instantiated
             comptime_args=[