Preserve nullable matrix types

pineforge_codegen/analyzer/types.py

@@ -42,7 +42,7 @@
 from ..ast_nodes import (
     ASTNode, BinOp, BoolLiteral, ExprStmt, FuncCall, Identifier, IfStmt,
     MemberAccess, NaLiteral, NumberLiteral, StringLiteral, Subscript, Ternary,
-    TupleLiteral, UnaryOp,
+    SwitchStmt, TupleLiteral, UnaryOp,
 )
 from ..symbols import PineType, TypeSpec
 
@@ -61,6 +61,15 @@
     "copy",
 })
 
+# Keep this analyzer-owned mirror in sync with
+# codegen.tables.MATRIX_RETURNING_METHODS.  The analyzer cannot import from
+# codegen (codegen already imports analyzer), but nullable selections need the
+# exact matrix result type before codegen registers global aggregate members.
+_MATRIX_RETURNING_METHODS = frozenset({
+    "copy", "submatrix", "transpose", "concat", "diff", "mult", "pow",
+    "inv", "pinv", "eigenvectors", "kron",
+})
+
 
 class TypeHelper:
     """Pine type-hint / expression inference.
@@ -188,12 +197,59 @@ def _array_from_element_spec(self, value: ASTNode | None) -> TypeSpec | None:
                 return self._pine_type_to_spec(sym.pine_type)
         return None
 
+    @staticmethod
+    def _selection_terminal_expr(
+        body: list[ASTNode] | None,
+    ) -> ASTNode | None:
+        """Return one if/switch branch's value expression, if present."""
+        if not body:
+            return None
+        terminal = body[-1]
+        return terminal.expr if isinstance(terminal, ExprStmt) else terminal
+
+    @staticmethod
+    def _selection_node_is_na(node: ASTNode | None) -> bool:
+        """Whether a selection arm is explicit or implicit Pine ``na``."""
+        return (
+            node is None
+            or isinstance(node, NaLiteral)
+            or (isinstance(node, Identifier) and node.name == "na")
+        )
+
+    def _nullable_collection_selection_spec(
+        self,
+        branches: list[tuple[ASTNode | None, TypeSpec | None]],
+    ) -> TypeSpec | None:
+        """Unify compatible map/matrix selection arms around typed ``na``.
+
+        A missing ``if``/``switch`` fallback is Pine's implicit ``na`` arm.
+        Every concrete arm must carry the same nullable collection TypeSpec;
+        an unknown or incompatible concrete arm fails closed.
+        """
+        concrete: list[TypeSpec] = []
+        for node, spec in branches:
+            if self._selection_node_is_na(node):
+                continue
+            if spec is None or spec.kind not in {"map", "matrix"}:
+                return None
+            concrete.append(spec)
+        if not concrete:
+            return None
+        first = concrete[0]
+        return first if all(spec == first for spec in concrete[1:]) else None
+
     def _type_spec_from_expr(self, value: ASTNode | None) -> TypeSpec | None:
         if value is None:
             return None
         if isinstance(value, Ternary):
             true_spec = self._type_spec_from_expr(value.true_val)
             false_spec = self._type_spec_from_expr(value.false_val)
+            collection_spec = self._nullable_collection_selection_spec([
+                (value.true_val, true_spec),
+                (value.false_val, false_spec),
+            ])
+            if collection_spec is not None:
+                return collection_spec
 
             def direct_user_udt_ctor_name(node: ASTNode) -> str | None:
                 if not isinstance(node, FuncCall):
@@ -216,21 +272,6 @@ def direct_user_udt_ctor_name(node: ASTNode) -> str | None:
                     and true_spec.kind == "udt"
                     and true_spec == false_spec):
                 return true_spec
-            if (true_spec is not None
-                    and true_spec.kind == "map"
-                    and true_spec == false_spec):
-                return true_spec
-            # A Pine ``na`` arm acquires the other arm's map type.  Keep this
-            # narrow to maps so unrelated scalar/array inference and generated
-            # output retain their established behavior.
-            if (true_spec is not None
-                    and true_spec.kind == "map"
-                    and isinstance(value.false_val, NaLiteral)):
-                return true_spec
-            if (false_spec is not None
-                    and false_spec.kind == "map"
-                    and isinstance(value.true_val, NaLiteral)):
-                return false_spec
             # A direct user-UDT constructor selected against bare ``na`` has
             # one unambiguous value type.  Require the constructor AST itself,
             # not merely an inferred UDT expression, so temporary array-element
@@ -277,18 +318,16 @@ def direct_user_udt_ctor_name(node: ASTNode) -> str | None:
                 return receiver_spec
             return None
         if isinstance(value, IfStmt):
-            def terminal_expr(body):
-                if not body:
-                    return None
-                terminal = body[-1]
-                return terminal.expr if isinstance(terminal, ExprStmt) else terminal
-
-            true_node = terminal_expr(value.body)
-            false_node = terminal_expr(value.else_body)
-            if true_node is None or false_node is None:
-                return None
+            true_node = self._selection_terminal_expr(value.body)
+            false_node = self._selection_terminal_expr(value.else_body)
             true_spec = self._type_spec_from_expr(true_node)
             false_spec = self._type_spec_from_expr(false_node)
+            collection_spec = self._nullable_collection_selection_spec([
+                (true_node, true_spec),
+                (false_node, false_spec),
+            ])
+            if collection_spec is not None:
+                return collection_spec
             true_is_na = (
                 isinstance(true_node, NaLiteral)
                 or (isinstance(true_node, Identifier)
@@ -299,18 +338,6 @@ def terminal_expr(body):
                 or (isinstance(false_node, Identifier)
                     and false_node.name == "na")
             )
-            if (true_spec is not None
-                    and true_spec.kind == "map"
-                    and true_spec == false_spec):
-                return true_spec
-            if (true_spec is not None
-                    and true_spec.kind == "map"
-                    and isinstance(false_node, NaLiteral)):
-                return true_spec
-            if (false_spec is not None
-                    and false_spec.kind == "map"
-                    and isinstance(true_node, NaLiteral)):
-                return false_spec
             if (true_spec is not None
                     and true_spec.kind == "udt"
                     and true_spec.name in _DRAWING_TYPE_NAMES
@@ -327,6 +354,22 @@ def terminal_expr(body):
                     and true_is_na):
                 return false_spec
             return None
+        if isinstance(value, SwitchStmt):
+            branches: list[tuple[ASTNode | None, TypeSpec | None]] = []
+            for _case_expr, case_body in value.cases:
+                terminal = self._selection_terminal_expr(case_body)
+                branches.append((
+                    terminal,
+                    self._type_spec_from_expr(terminal),
+                ))
+            default_terminal = self._selection_terminal_expr(
+                value.default_body
+            )
+            branches.append((
+                default_terminal,
+                self._type_spec_from_expr(default_terminal),
+            ))
+            return self._nullable_collection_selection_spec(branches)
         if isinstance(value, FuncCall):
             cal = value.callee
             func = cal.member if isinstance(cal, MemberAccess) else None
@@ -387,6 +430,11 @@ def terminal_expr(body):
                 else:
                     elem = TypeSpec.primitive("float")
                 return TypeSpec.matrix(elem)
+            if ns == "matrix" and func in _MATRIX_RETURNING_METHODS:
+                receiver = value.args[0] if value.args else value.kwargs.get("id")
+                receiver_spec = self._type_spec_from_expr(receiver)
+                if receiver_spec is not None and receiver_spec.kind == "matrix":
+                    return receiver_spec
             if ns == "map" and func == "new":
                 key = self._type_spec_from_hint(targs[0]) if len(targs) > 0 else TypeSpec.primitive("string")
                 val = self._type_spec_from_hint(targs[1]) if len(targs) > 1 else TypeSpec.primitive("float")
@@ -452,7 +500,7 @@ def terminal_expr(body):
                     if func == "size":
                         return TypeSpec.primitive("int")
                 if recv_spec is not None and recv_spec.kind == "matrix":
-                    if func in ("copy", "submatrix", "transpose", "concat"):
+                    if func in _MATRIX_RETURNING_METHODS:
                         return recv_spec
                     if func in ("row", "col"):
                         return TypeSpec.array(recv_spec.element)

pineforge_codegen/codegen/base.py

@@ -1213,11 +1213,19 @@ def allocate_flag(base: str) -> str:
                     and self._is_na_expr(stmt.value)
                 )
             )
+            nullable_collection_selection_decl_site = (
+                stmt_spec is not None
+                and stmt_spec.kind in {"map", "matrix"}
+                and isinstance(stmt.value, (IfStmt, SwitchStmt))
+            )
             established_decl_site = (
                 (not is_callable_scoped or drawing_cpp is not None)
-                and self._is_runtime_scalar_var_initializer(
-                    member_name, ptype, init_str, stmt.value, drawing_cpp,
-                    is_series
+                and (
+                    nullable_collection_selection_decl_site
+                    or self._is_runtime_scalar_var_initializer(
+                        member_name, ptype, init_str, stmt.value, drawing_cpp,
+                        is_series
+                    )
                 )
             )
             if not callable_decl_site and not established_decl_site:
@@ -1857,6 +1865,29 @@ def _register_global_aggregate_member_types(self) -> None:
         ``global_expr_map``; without registering it here, ``m`` was emitted as a scalar
         while ``on_bar`` still assigned ``PineMatrix``.
         """
+        # A typed ``matrix<T> m = na`` or an inferred nullable selection has no
+        # constructor call for the legacy RHS scan below to recognize. Resolve
+        # declarations in source order from the authored hint/RHS, not from
+        # ``ctx.collection_types``: that raw-name table describes the analyzer's
+        # final state and can already contain an invalid later reassignment's
+        # element type. The original declaration must remain the compatibility
+        # baseline used to reject matrix<int> := matrix<float>.
+        for stmt, _in_loop in self._walk_global_scope_with_loopflag(
+            getattr(self.ctx.ast, "body", []), False
+        ):
+            if not isinstance(stmt, VarDecl):
+                continue
+            if stmt.type_hint:
+                declared_spec = self._type_spec_from_hint_name(stmt.type_hint)
+            elif isinstance(stmt.value, (Ternary, IfStmt, SwitchStmt)):
+                declared_spec = self._type_spec_from_expr(stmt.value)
+            else:
+                continue
+            if declared_spec is None or declared_spec.kind != "matrix":
+                continue
+            self._matrix_specs.setdefault(stmt.name, declared_spec)
+            self._collection_types[stmt.name] = self._matrix_specs[stmt.name]
+
         gem = getattr(self.ctx, "global_expr_map", {}) or {}
         for name, _ptype in self.ctx.global_var_decls:
             expr = gem.get(name)
@@ -2055,8 +2086,70 @@ def _check_matrix_method_allowed(self, meth_name, recv_spec, node) -> None:
             else:
                 self._codegen_error(node, "matrix.sort requires int, bool, string, or float element type; UDT matrices cannot be sorted")
 
+    def _iter_declared_matrix_specs(self):
+        """Yield matrix TypeSpecs reachable from declared source types.
+
+        A typed ``matrix<T> x = na`` has no matrix call for a syntactic scan
+        to find. Matrix types can also be nested in UDT fields or appear only
+        at a callable boundary, so include selection must start from analyzer
+        metadata and recurse through aggregate types.
+        """
+        roots = list(self._collection_types.values())
+        roots.extend(
+            spec
+            for specs in self._func_collection_types.values()
+            for spec in specs.values()
+        )
+        roots.extend(
+            spec
+            for specs in self._block_collection_types.values()
+            for spec in specs.values()
+            if spec is not None
+        )
+        roots.extend(
+            spec
+            for fields in self._udt_field_type_specs.values()
+            for spec in fields.values()
+        )
+        for fi in self.ctx.func_infos:
+            roots.extend(
+                spec
+                for spec in (getattr(fi, "param_type_specs", []) or [])
+                if spec is not None
+            )
+            return_spec = getattr(fi, "return_type_spec", None)
+            if return_spec is not None:
+                roots.append(return_spec)
+
+        def walk(spec: TypeSpec | None, visiting_udts: frozenset[str]):
+            if spec is None:
+                return
+            if spec.kind == "matrix":
+                yield spec
+                return
+            if spec.kind == "array":
+                yield from walk(spec.element, visiting_udts)
+                return
+            if spec.kind == "map":
+                yield from walk(spec.key, visiting_udts)
+                yield from walk(spec.value, visiting_udts)
+                return
+            if spec.kind == "udt" and spec.name:
+                if spec.name in visiting_udts:
+                    return
+                nested_visiting = visiting_udts | {spec.name}
+                for field_spec in self._udt_field_type_specs.get(
+                    spec.name, {}
+                ).values():
+                    yield from walk(field_spec, nested_visiting)
+
+        for root in roots:
+            yield from walk(root, frozenset())
+
     def _detect_matrix_usage(self) -> bool:
-        """True if emitted C++ will need runtime/matrix.hpp (PineMatrix)."""
+        """True if emitted C++ will need a matrix runtime header."""
+        if next(self._iter_declared_matrix_specs(), None) is not None:
+            return True
         for _, _, init_str in self.ctx.var_members:
             if init_str and "matrix.new" in str(init_str):
                 return True
@@ -2067,6 +2160,32 @@ def _detect_matrix_usage(self) -> bool:
                     return True
         return False
 
+    def _detect_generic_matrix_usage(self) -> bool:
+        """True if emitted C++ mentions a non-float matrix specialization."""
+        float_spec = TypeSpec.primitive("float")
+        if any(
+            spec.element != float_spec
+            for spec in self._iter_declared_matrix_specs()
+        ):
+            return True
+
+        # A standalone matrix.new<T>() expression may have no declaration
+        # TypeSpec. Read its explicit template argument directly so the
+        # generated call still receives the generic runtime declaration.
+        for node in self._walk_ast(self.ctx.ast):
+            if not isinstance(node, FuncCall):
+                continue
+            fn, ns = self._resolve_callee(node.callee)
+            if ns != "matrix" or fn != "new":
+                continue
+            targs = self._template_args_from_call(node)
+            if not targs:
+                continue
+            elem_spec = self._type_spec_from_hint_name(targs[0])
+            if elem_spec is not None and elem_spec != float_spec:
+                return True
+        return False
+
     def _detect_map_usage(self) -> bool:
         """True when emitted C++ needs the PineMap handle/runtime helpers.
 
@@ -3453,7 +3572,7 @@ def generate(self) -> str:
                         f.default,
                         target_cpp_type=(
                             cpp_type
-                            if (cpp_type.startswith("PineMap<")
+                            if (self._is_nullable_collection_cpp_type(cpp_type)
                                 or cpp_type in self._udt_defs
                                 or cpp_type in DRAWING_TYPE_TO_CPP.values())
                             else None