allowing type[T] to be passed to functions

Author: shiinamiyuki

luisa_lang/_builtin_decor.py

@@ -69,7 +69,7 @@ class _ObjKind(Enum):
 
 
 def _make_func_template(f: Callable[..., Any], func_name: str, func_sig: Optional[MethodType],
-                        func_globals: Dict[str, Any], foreign_type_var_ns: Dict[TypeVar, hir.Type | hir.ComptimeValue],
+                        func_globals: Dict[str, Any], foreign_type_var_ns: Dict[TypeVar, hir.Type],
                         props: hir.FuncProperties, self_type: Optional[hir.Type] = None):
     # parsing_ctx = _parse.ParsingContext(func_name, func_globals)
     # func_sig_parser = _parse.FuncParser(func_name, f, parsing_ctx, self_type)
@@ -88,8 +88,7 @@ def _make_func_template(f: Callable[..., Any], func_name: str, func_sig: Optiona
         implicit_generic_params.add(p.param)
 
     def parsing_func(args: hir.FunctionTemplateResolvingArgs) -> hir.Function:
-        type_var_ns: Dict[TypeVar, hir.Type |
-                          hir.ComptimeValue] = foreign_type_var_ns.copy()
+        type_var_ns: Dict[TypeVar, hir.Type] = foreign_type_var_ns.copy()
         mapped_implicit_type_params: Dict[str,
                                           hir.Type] = dict()
         assert func_sig is not None
@@ -166,7 +165,7 @@ def _dsl_func_impl(f: _TT, kind: _ObjKind, attrs: Dict[str, Any]) -> _TT:
 
 
 _MakeTemplateFn = Callable[[List[hir.GenericParameter]], hir.Type]
-_InstantiateFn = Callable[[List[Any]], hir.Type]
+_InstantiateFn = Callable[[List[hir.Type]], hir.Type]
 
 
 def _dsl_struct_impl(cls: type[_TT], attrs: Dict[str, Any], ir_ty_override: hir.Type | Tuple[_MakeTemplateFn, _InstantiateFn] | None = None, opqaue_override: str | None = None) -> type[_TT]:
@@ -202,6 +201,8 @@ def parse_fields(tp: parse.TypeParser, self_ty: hir.Type):
 
     def parse_methods(type_var_ns: Dict[TypeVar, hir.Type | Any], self_ty: hir.Type,):
         for name in cls_info.methods:
+            if name == '__setitem__': # __setitem__ is ignored deliberately
+                continue
             method_object = getattr(cls, name)
             props: hir.FuncProperties
             if hasattr(method_object, '__luisa_func_props__'):
@@ -214,7 +215,7 @@ def parse_methods(type_var_ns: Dict[TypeVar, hir.Type | Any], self_ty: hir.Type,
                 method_object, get_full_name(method_object), cls_info.methods[name], globalns, type_var_ns, props, self_type=self_ty)
             if isinstance(self_ty, hir.BoundType):
                 assert isinstance(self_ty.instantiated,
-                                  (hir.StructType, hir.OpaqueType))
+                                  (hir.ArrayType, hir.StructType, hir.OpaqueType))
                 self_ty.instantiated.methods[name] = template
             else:
                 self_ty.methods[name] = template
@@ -235,7 +236,7 @@ def parse_methods(type_var_ns: Dict[TypeVar, hir.Type | Any], self_ty: hir.Type,
         parse_fields(type_parser, ir_ty)
     is_generic = len(cls_info.type_vars) > 0
     if is_generic:
-        def monomorphization_func(args: List[hir.Type | Any]) -> hir.Type:
+        def monomorphization_func(args: List[hir.Type]) -> hir.Type:
             assert isinstance(ir_ty, hir.ParametricType)
             type_var_ns = {}
             if len(args) != len(cls_info.type_vars):

luisa_lang/classinfo.py

@@ -56,9 +56,21 @@ def __repr__(self):
 
     def __eq__(self, other):
         return isinstance(other, SelfType)
+    
+class LiteralType:
+    value: Any
 
+    def __init__(self, value: Any):
+        self.value = value
 
-VarType = Union[TypeVar, Type, GenericInstance, UnionType, SelfType, AnyType]
+    def __repr__(self):
+        return f"Literal[{self.value}]"
+
+    def __eq__(self, other):
+        return isinstance(other, LiteralType) and self.value == other.value
+
+
+VarType = Union[TypeVar, Type, GenericInstance, UnionType, SelfType, AnyType, LiteralType]
 
 
 def subst_type(ty: VarType, env: Dict[TypeVar, VarType]) -> VarType:
@@ -204,6 +216,8 @@ def parse_type_hint(hint: Any) -> VarType:
             return GenericInstance(origin, [parse_type_hint(arg) for arg in args])
         elif origin is Union:
             return UnionType([parse_type_hint(arg) for arg in typing.get_args(hint)])
+        elif origin is Literal:
+            return LiteralType(typing.get_args(hint)[0])
         else:
             raise RuntimeError(f"Unsupported origin type: {origin}")
 

luisa_lang/codegen/cpp.py

@@ -56,8 +56,12 @@ def gen_impl(self, ty: hir.Type) -> str:
                         raise RuntimeError("invalid float type")
             case hir.BoolType():
                 return "lc_bool"
+            case hir.PointerType(element=element):
+                return f"lc_ptr<{self.gen(element)}>"
             case hir.VectorType(element=element, count=count):
                 return f"{self.gen(element)}{count}>"
+            case hir.ArrayType(element=element, count=count):
+                return f"lc_array<{self.gen(element)}, {count}>"
             case hir.StructType(name=name, fields=fields):
                 self.impl.writeln(f'struct {name} {{')
                 for field in fields:
@@ -80,9 +84,13 @@ def do():
                 assert ty.instantiated
                 return self.gen(ty.instantiated)
             case hir.FunctionType():
-                return ''
+                name = f'func_{unique_hash(ty.func_like.name)}_t'
+                self.impl.writeln(f'struct {name} {{}}; // function type of {ty.func_like.name}')
+                return name
             case hir.TypeConstructorType():
-                return ''
+                name = f'type_{unique_hash(self.gen(ty.inner))}_t'
+                self.impl.writeln(f'struct {name} {{}}; // type constructor of {ty.inner}')
+                return name
             case hir.OpaqueType():
                 def do():
                     match ty.name:
@@ -171,8 +179,8 @@ def mangle_impl(self, obj: Union[hir.Type, hir.Function]) -> str:
             case hir.Function(name=name, params=params, return_type=ret):
                 assert ret
                 name = mangle_name(name)
-                params = list(filter(lambda p: not isinstance(
-                    p.type, (hir.FunctionType)), params))
+                # params = list(filter(lambda p: not isinstance(
+                #     p.type, (hir.FunctionType)), params))
                 return f'{name}_' + unique_hash(f"F{name}_{self.mangle(ret)}{''.join(self.mangle(unwrap(p.type)) for p in params)}")
             case hir.StructType(name=name):
                 return name
@@ -184,6 +192,10 @@ def mangle_impl(self, obj: Union[hir.Type, hir.Function]) -> str:
                 return self.mangle(obj.instantiated)
             case hir.OpaqueType():
                 return obj.name
+            case hir.TypeConstructorType():
+                return self.mangle(obj.inner)
+            case hir.FunctionType():
+                return f'func_{unique_hash(obj.func_like.name)}'
             case _:
                 raise NotImplementedError(f"unsupported object: {obj}")
 
@@ -246,7 +258,7 @@ def __init__(self, base: CppCodeGen, func: hir.Function) -> None:
         self.name = base.mangling.mangle(func)
         self.func = func
         params = ",".join(self.gen_var(
-            p) for p in func.params if not isinstance(p.type, hir.FunctionType))
+            p) for p in func.params)
         assert func.return_type
 
         self.signature = f'auto {self.name}({params}) -> {base.type_cache.gen(func.return_type)}'
@@ -285,10 +297,13 @@ def do():
                     intrin_name = intrin.name
                     gened_args = [self.gen_value_or_ref(
                             arg) for arg in intrin.args]
-                    if intrin_name == 'buffer_ref':
-                        return f"{gened_args[0]}[{gened_args[1]}]"
-                    else:
-                        raise RuntimeError(f"unsupported intrinsic reference: {intrin_name}")
+                    match intrin_name:
+                        case 'buffer.ref' | 'array.ref':
+                            return f"{gened_args[0]}[{gened_args[1]}]"
+                        case 'buffer.size' | 'array.size':
+                            return f"{gened_args[0]}.size"
+                        case _:
+                            raise RuntimeError(f"unsupported intrinsic reference: {intrin_name}")
                 return do()
             case _:
                 raise NotImplementedError(f"unsupported reference: {ref}")
@@ -312,17 +327,40 @@ def gen_value_or_ref(self, value: hir.Value | hir.Ref) -> str:
     def gen_node_checked(self, node: hir.Node) -> str:
         if isinstance(node, hir.Constant):
             return self.gen_expr(node)
+        if isinstance(node, hir.TypedNode) and isinstance(node.type, (hir.TypeConstructorType, hir.FunctionType)):
+            assert node.type
+            return f'{self.base.type_cache.gen(node.type)}{{}}'
+            
         return self.node_map[node]
 
     def gen_expr(self, expr: hir.Value) -> str:
-        if expr.type and isinstance(expr.type, hir.FunctionType):
-            return ''
+        # if expr.type and isinstance(expr.type, hir.FunctionType):
+        #     return ''
+        if isinstance(expr, hir.Constant):
+            value = expr.value
+            if isinstance(value, int):
+                return f"{value}"
+            elif isinstance(value, float):
+                return f"{value}f"
+            elif isinstance(value, bool):
+                return "true" if value else "false"
+            elif isinstance(value, str):
+                return f"\"{value}\""
+            elif isinstance(value, hir.Function):
+                return self.gen_func(value)
+            else:
+                raise NotImplementedError(
+                    f"unsupported constant: {expr}")
         if expr in self.node_map:
             return self.node_map[expr]
         vid = self.new_vid()
 
         def impl() -> None:
             match expr:
+                case hir.TypeValue() as type_value:
+                    assert type_value.type
+                    self.base.type_cache.gen(type_value.type)
+                    return
                 case hir.Load() as load:
                     self.body.writeln(
                         f"const auto &v{vid} = {self.gen_ref(load.ref)}; // load")
@@ -337,36 +375,17 @@ def impl() -> None:
                 case hir.Call() as call:
                     op = self.gen_func(call.op)
                     args_s = ','.join(self.gen_value_or_ref(
-                        arg) for arg in call.args if not isinstance(arg.type, hir.FunctionType))
+                        arg) for arg in call.args)
                     if call.type != hir.UnitType():
                         self.body.writeln(
                             f"auto v{vid} = {op}({args_s});")
                     else:
                         self.body.writeln(f"{op}({args_s});")
-                case hir.Constant() as constant:
-                    value = constant.value
-                    if isinstance(value, int):
-                        self.body.writeln(f"const auto v{vid} = {value};")
-                    elif isinstance(value, float):
-                        self.body.writeln(f"const auto v{vid} = {value};")
-                    elif isinstance(value, bool):
-                        s = "true" if value else "false"
-                        self.body.writeln(f"const auto v{vid} = {s};")
-                    elif isinstance(value, str):
-                        self.body.writeln(f"const auto v{vid} = \"{value}\";")
-                    elif isinstance(value, hir.Function):
-                        name = self.gen_func(value)
-                        self.body.writeln(f"auto&& v{vid} = {name};")
-                    else:
-                        raise NotImplementedError(
-                            f"unsupported constant: {constant}")
                 case hir.AggregateInit():
                     assert expr.type
                     ty = self.base.type_cache.gen(expr.type)
                     self.body.writeln(
                         f"{ty} v{vid}{{ {','.join(self.gen_expr(e) for e in expr.args)} }};")
-                case hir.TypeValue():
-                    pass
                 case hir.Intrinsic() as intrin:
                     def do():
                         intrin_name = intrin.name
@@ -544,7 +563,7 @@ def gen_node(self, node: hir.Node) -> Optional[hir.BasicBlock]:
                 ty = self.base.type_cache.gen(alloca.type)
                 self.body.writeln(f"{ty} v{vid}{{}};")
                 self.node_map[alloca] = f"v{vid}"
-            case hir.AggregateInit() | hir.Intrinsic() | hir.Call() | hir.Constant() | hir.Load() | hir.Index() | hir.Member():
+            case hir.AggregateInit() | hir.Intrinsic() | hir.Call() | hir.Constant() | hir.Load() | hir.Index() | hir.Member() | hir.TypeValue() | hir.FunctionValue():
                 self.gen_expr(node)
             case hir.Member() | hir.Index():
                 pass
@@ -570,8 +589,8 @@ def gen_locals(self):
         for local in self.func.locals:
             if local.name in self.params:
                 continue
-            if isinstance(local.type, (hir.FunctionType, hir.TypeConstructorType)):
-                continue
+            # if isinstance(local.type, (hir.FunctionType, hir.TypeConstructorType)):
+            #     continue
             assert (
                 local.type
             ), f"Local variable `{local.name}` contains unresolved type"