diff --git a/pineforge_codegen/analyzer/types.py b/pineforge_codegen/analyzer/types.py index f4857bc..f1e400d 100644 --- a/pineforge_codegen/analyzer/types.py +++ b/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) diff --git a/pineforge_codegen/codegen/base.py b/pineforge_codegen/codegen/base.py index 521cc9b..1f886c9 100644 --- a/pineforge_codegen/codegen/base.py +++ b/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 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 := matrix. + 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 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() 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 diff --git a/pineforge_codegen/codegen/emit_top.py b/pineforge_codegen/codegen/emit_top.py index 621df0c..f5f89f9 100644 --- a/pineforge_codegen/codegen/emit_top.py +++ b/pineforge_codegen/codegen/emit_top.py @@ -86,7 +86,7 @@ ExprStmt, FuncCall, Identifier, IfStmt, SwitchStmt, VarDecl, ) from ..analyzer import FuncInfo -from ..symbols import PineType, TypeSpec +from ..symbols import PineType from .tables import ( BAR_SERIES_PUSH, DRAWING_TYPE_TO_CPP, @@ -137,27 +137,7 @@ def _emit_includes(self, lines: list[str]) -> None: # is present in the TU. Float matrices route through PineMatrix in # matrix.hpp; pulling in the generic header otherwise is a wasted # include. - float_spec = TypeSpec.primitive("float") - scoped_matrix_specs = ( - spec - for specs in self._func_collection_types.values() - for spec in specs.values() - if spec.kind == "matrix" - ) - block_matrix_specs = ( - spec - for specs in self._block_collection_types.values() - for spec in specs.values() - if spec is not None and spec.kind == "matrix" - ) - if any( - spec.element != float_spec - for spec in ( - *self._matrix_specs.values(), - *scoped_matrix_specs, - *block_matrix_specs, - ) - ): + if self._detect_generic_matrix_usage(): lines.append('#include ') # Drawing-objects-as-data runtime (line/box/label/linefill arenas + # ChartPoint). Gated on _uses_drawing so non-drawing strategies stay @@ -575,7 +555,14 @@ def _emit_constructor(self, lines: list[str]) -> None: continue seen_ctor_vars.add(name) safe = self._safe_name(name) - if name in self._array_vars or name in self._map_vars: + if (name in self._array_vars + or name in self._map_vars + or name in self._matrix_specs): + # Collection members are initialized in the first-bar block. + # A matrix declared as ``var matrix x = na`` has no + # constructor call to replay there and its default constructor + # is already the correctly typed null ID. Letting the generic + # scalar path add ``x(na())`` is ill-typed. continue # Callable-scoped ``var`` members initialize at their exact source # declarations; exclude them from the constructor so source order, @@ -881,6 +868,20 @@ def _emit_on_bar(self, lines: list[str]) -> None: f"{runtime_info['drawing_cpp']}{{}}", ) continue + # Declaration-site initialization owns nullable collection + # members carrying runtime metadata. A persistent map/matrix + # if/switch initializer must not also enter the legacy + # first-bar aggregate preamble (which cannot lower a block + # expression and formerly emitted ``/* unknown */`` for maps). + # Keep primitive and ternary paths on their established order. + runtime_spec = ( + runtime_info.get("type_spec") + if runtime_info is not None + else None + ) + if (runtime_spec is not None + and runtime_spec.kind in {"map", "matrix"}): + continue if name in self._array_vars: for stmt in self.ctx.ast.body: if isinstance(stmt, VarDecl) and stmt.name == name: @@ -1393,7 +1394,7 @@ def _emit_func_def(self, fi: FuncInfo, lines: list[str], call_site_idx: int | No ret_type = PINE_TYPE_TO_CPP.get(fi.return_type, "double") rhs_return_cpp_type = ( ret_type - if (ret_type.startswith("PineMap<") + if (self._is_nullable_collection_cpp_type(ret_type) or ret_type in DRAWING_TYPE_TO_CPP.values() or ret_type in self._udt_defs) else None diff --git a/pineforge_codegen/codegen/types.py b/pineforge_codegen/codegen/types.py index 6e9c017..44e4dfc 100644 --- a/pineforge_codegen/codegen/types.py +++ b/pineforge_codegen/codegen/types.py @@ -34,7 +34,7 @@ from __future__ import annotations from ..ast_nodes import ( - Assignment, BinOp, BoolLiteral, ExprStmt, FuncCall, FuncDef, Identifier, IfStmt, + ASTNode, Assignment, BinOp, BoolLiteral, ExprStmt, FuncCall, FuncDef, Identifier, IfStmt, MemberAccess, NaLiteral, NumberLiteral, StringLiteral, SwitchStmt, Subscript, Ternary, TupleLiteral, UnaryOp, VarDecl, ) @@ -47,6 +47,7 @@ BAR_FIELDS, DRAWING_NS, DRAWING_TYPE_TO_CPP, + MATRIX_RETURNING_METHODS, PINE_TYPE_TO_CPP, TA_RETURNS_BOOL, ) @@ -154,7 +155,23 @@ def _type_spec_to_cpp(self, spec: TypeSpec | None) -> str: return "double" @staticmethod - def _default_for_type(cpp_type: str) -> str: + def _is_nullable_collection_cpp_type(cpp_type: str | None) -> bool: + """Whether ``cpp_type`` has a default-constructed Pine ``na`` ID. + + Arrays still lower to ``std::vector`` and cannot distinguish a null ID + from a valid empty array. Keep them outside this predicate until their + runtime representation carries that distinction. + """ + return bool( + cpp_type + and ( + cpp_type.startswith("PineMap<") + or cpp_type == "PineMatrix" + or cpp_type.startswith("PineGenericMatrix<") + ) + ) + + def _default_for_type(self, cpp_type: str) -> str: """Default initialiser for a primitive C++ type (matches Pine ``na``).""" if cpp_type == "std::string": return 'std::string("")' @@ -166,6 +183,8 @@ def _default_for_type(cpp_type: str) -> str: return f"{cpp_type}{{}}" if cpp_type.startswith("std::vector") or cpp_type.startswith("PineMap"): return f"{cpp_type}()" + if self._is_nullable_collection_cpp_type(cpp_type): + return f"{cpp_type}{{}}" return "0.0" def _default_for_spec(self, spec: TypeSpec | None) -> str: @@ -184,6 +203,8 @@ def _default_for_spec(self, spec: TypeSpec | None) -> str: cpp_type = self._type_spec_to_cpp(spec) if cpp_type.startswith("std::vector") or cpp_type.startswith("PineMap"): return f"{cpp_type}()" + if self._is_nullable_collection_cpp_type(cpp_type): + return f"{cpp_type}{{}}" return self._default_for_type(cpp_type) def _collection_spec_for_name(self, name: str) -> TypeSpec | None: @@ -357,6 +378,40 @@ def _array_from_element_spec(self, node) -> TypeSpec | None: return TypeSpec.primitive(primitive_name) 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``.""" + 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, node) -> TypeSpec | None: """Best-effort TypeSpec inference for an expression node. @@ -420,6 +475,12 @@ def _type_spec_from_expr(self, node) -> TypeSpec | None: if isinstance(node, Ternary): true_spec = self._type_spec_from_expr(node.true_val) false_spec = self._type_spec_from_expr(node.false_val) + collection_spec = self._nullable_collection_selection_spec([ + (node.true_val, true_spec), + (node.false_val, false_spec), + ]) + if collection_spec is not None: + return collection_spec if true_spec is not None and true_spec == false_spec: return true_spec if (true_spec is not None @@ -434,22 +495,16 @@ def _type_spec_from_expr(self, node) -> TypeSpec | None: return false_spec return None if isinstance(node, 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(node.body) - false_node = terminal_expr(node.else_body) - if true_node is None or false_node is None: - return None + true_node = self._selection_terminal_expr(node.body) + false_node = self._selection_terminal_expr(node.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) @@ -476,6 +531,22 @@ def terminal_expr(body): and true_is_na): return false_spec return None + if isinstance(node, SwitchStmt): + branches: list[tuple[ASTNode | None, TypeSpec | None]] = [] + for _case_expr, case_body in node.cases: + terminal = self._selection_terminal_expr(case_body) + branches.append(( + terminal, + self._type_spec_from_expr(terminal), + )) + default_terminal = self._selection_terminal_expr( + node.default_body + ) + branches.append(( + default_terminal, + self._type_spec_from_expr(default_terminal), + )) + return self._nullable_collection_selection_spec(branches) if isinstance(node, MemberAccess): owner = self._type_spec_from_expr(node.object) if owner is not None and owner.kind == "udt" and owner.name: @@ -510,6 +581,13 @@ def terminal_expr(body): return TypeSpec.udt("chart.point") if namespace == "str" and func_name == "split": return TypeSpec.array(TypeSpec.primitive("string")) + if namespace == "matrix" and func_name == "new": + elem = ( + self._type_spec_from_hint_name(targs[0]) + if targs + else TypeSpec.primitive("float") + ) + return TypeSpec.matrix(elem or TypeSpec.primitive("float")) if namespace == "array" and func_name in ( "new", "new_float", "new_int", "new_bool", "new_string", "from", ) or (namespace == "array" and func_name in ARRAY_DRAWING_NEW_CTORS): @@ -545,6 +623,11 @@ def terminal_expr(body): if func_name in ("copy", "slice"): return arg_spec return arg_spec.element + if namespace == "matrix" and func_name in MATRIX_RETURNING_METHODS: + receiver_node = node.args[0] if node.args else node.kwargs.get("id") + receiver_spec = self._type_spec_from_expr(receiver_node) + if receiver_spec is not None and receiver_spec.kind == "matrix": + return receiver_spec if namespace == "map" and func_name == "new": key = self._type_spec_from_hint_name(targs[0]) if len(targs) > 0 else TypeSpec.primitive("string") val = self._type_spec_from_hint_name(targs[1]) if len(targs) > 1 else TypeSpec.primitive("float") @@ -608,7 +691,7 @@ def terminal_expr(body): if return_spec is not None: return return_spec if recv_spec is not None and recv_spec.kind == "matrix": - if func_name in ("copy", "submatrix", "transpose", "concat"): + if func_name in MATRIX_RETURNING_METHODS: return recv_spec if func_name in ("row", "col"): return TypeSpec.array(recv_spec.element) @@ -886,10 +969,10 @@ def promote_wide_int(cpp_type: str) -> str: # its lexical scopes are popped. In particular, an inferred local # such as ``selected = cond ? na : global_map`` has a map TypeSpec even # though generic C++ expression inference sees ``na`` as ``double``. - # Consume only the map form here; arrays/matrices and every non-map - # declaration retain their established inference/output paths. + # Consume nullable map/matrix forms here. Arrays still lack a runtime + # null-ID representation and retain their established inference path. captured = self._callable_collection_bindings.get(id(node)) - if captured is not None and captured.kind == "map": + if captured is not None and captured.kind in {"map", "matrix"}: return self._type_spec_to_cpp(captured) # Drawing handle local (L-N6): a hintless local whose RHS resolves to a # drawing udt must declare as the handle struct, not the analyzer's @@ -1355,13 +1438,14 @@ def _na_reassign_cpp_type(self, name: str) -> str | None: for ``double`` (already the default lowering), so those paths are unchanged. """ - # Maps use their default-constructed null ID for Pine ``na``. Arrays, - # matrices, UDTs and drawing handles retain their established paths. + # Maps and matrices use their default-constructed null ID for Pine + # ``na``. Arrays still lack a nullable runtime representation. collection_spec = self._collection_spec_for_name(name) - if collection_spec is not None and collection_spec.kind == "map": + if (collection_spec is not None + and collection_spec.kind in {"map", "matrix"}): return self._type_spec_to_cpp(collection_spec) if ((collection_spec is not None - and collection_spec.kind in {"array", "matrix"}) + and collection_spec.kind == "array") or self._identifier_udt_type(name) is not None): return None cpp_type: str | None = None diff --git a/pineforge_codegen/codegen/visit_call.py b/pineforge_codegen/codegen/visit_call.py index 1f7763b..2969285 100644 --- a/pineforge_codegen/codegen/visit_call.py +++ b/pineforge_codegen/codegen/visit_call.py @@ -805,6 +805,23 @@ def _visit_udt_method_series_arg( and param_index < len(func_info.node.params) else None ) + param_specs = list( + getattr(func_info, "param_type_specs", ()) or () + ) + param_spec = ( + param_specs[param_index] + if param_index < len(param_specs) + else None + ) + if param_spec is not None and param_spec.kind in {"map", "matrix"}: + # Both UDT-method dispatch paths funnel through this helper after + # positional/keyword/default merging. Lower a bare ``na`` (and + # nullable selections containing it) against the declared exact + # parameter type instead of the legacy numeric ``na()``. + return self._visit_rhs_value( + arg_node, + target_cpp_type=self._type_spec_to_cpp(param_spec), + ) method_series = self.ctx.func_series_vars.get(func_info.name, set()) if param_name not in method_series: return self._visit_expr(arg_node) @@ -1900,7 +1917,7 @@ def _visit_func_call(self, node: FuncCall) -> str: value_node, target_cpp_type=( f_cpp_type - if (f_cpp_type.startswith("PineMap<") + if (self._is_nullable_collection_cpp_type(f_cpp_type) or f_cpp_type in self._udt_defs or f_cpp_type in DRAWING_TYPE_TO_CPP.values()) else None @@ -1911,7 +1928,7 @@ def _visit_func_call(self, node: FuncCall) -> str: val = val.replace("na()", "na()") else: val = f"(int64_t)({val})" - elif (f_cpp_type.startswith("PineMap<") + elif (self._is_nullable_collection_cpp_type(f_cpp_type) and val == "na()"): val = f"{f_cpp_type}{{}}" field_inits.append(f".{f.name} = {val}") @@ -2034,17 +2051,14 @@ def _visit_arg_for_series(arg_node, arg_idx): f"else {member}.update(_sv); " f"return {member}; }}())" ) - # A concrete map specialization learned through an untyped - # wrapper chain also target-types its call arguments. In - # particular, ``choose(cond, value, na)`` must pass a typed null - # PineMap handle rather than the generic ``na()``. Every - # non-map destination remains on the byte-identical expression - # path below. + # A concrete nullable collection specialization learned through an + # untyped wrapper chain also target-types its call arguments. if fi_lookup is not None: param_specs = getattr(fi_lookup, "param_type_specs", []) or [] if arg_idx < len(param_specs): param_spec = param_specs[arg_idx] - if param_spec is not None and param_spec.kind == "map": + if (param_spec is not None + and param_spec.kind in {"map", "matrix"}): return self._visit_rhs_value( arg_node, target_cpp_type=self._type_spec_to_cpp(param_spec), diff --git a/pineforge_codegen/codegen/visit_expr.py b/pineforge_codegen/codegen/visit_expr.py index 4872420..a40946c 100644 --- a/pineforge_codegen/codegen/visit_expr.py +++ b/pineforge_codegen/codegen/visit_expr.py @@ -342,23 +342,23 @@ def _visit_rhs_value(self, value_node, target_name: str | None = None, return f"{drawing_target}{{}}" if target_cpp_type in self._udt_defs: return f"{target_cpp_type}{{}}" - if target_cpp_type and target_cpp_type.startswith("PineMap<"): - # PineMap's default constructor is the typed ``na`` ID. A - # map.new call is the only operation that allocates storage. + if self._is_nullable_collection_cpp_type(target_cpp_type): + # Maps and matrices use default construction for a typed + # ``na`` ID. Their ``*.new`` factories create a valid ID, + # including a valid empty collection. return f"{target_cpp_type}{{}}" if target_cpp_type in ("std::string", "int", "int64_t", "bool"): return f"na<{target_cpp_type}>()" if (isinstance(value_node, Ternary) - and ((target_cpp_type - and target_cpp_type.startswith("PineMap<")) + and (self._is_nullable_collection_cpp_type(target_cpp_type) or drawing_target is not None or target_cpp_type in self._udt_defs)): # C++ cannot deduce a common type for ``na()`` and a - # PineMap/drawing handle. Pine's ternary is target typed, so + # collection/drawing handle. Pine's ternary is target typed, so # propagate the exact declared/reassignment target into both arms. - # Arrays, matrices, and scalar ternaries retain the established - # generic expression path; collection IDs require a nullable - # runtime representation before target typing alone can be safe. + # Arrays and scalar ternaries retain the established generic path; + # target typing is safe only for collections with nullable runtime + # representations. branch_target = drawing_target or target_cpp_type condition = self._visit_expr(value_node.condition) true_value = self._visit_rhs_value( diff --git a/pineforge_codegen/codegen/visit_stmt.py b/pineforge_codegen/codegen/visit_stmt.py index b49161d..4ca77b3 100644 --- a/pineforge_codegen/codegen/visit_stmt.py +++ b/pineforge_codegen/codegen/visit_stmt.py @@ -91,6 +91,7 @@ MethodDef, StrategyDecl, SwitchStmt, + Ternary, TupleAssign, TupleLiteral, TypeDecl, @@ -428,27 +429,27 @@ def _concat_receiver_name(self, expr) -> str | None: return recv return None - def _map_target_cpp_type( + def _nullable_collection_target_cpp_type( self, *, name: str | None = None, target_node=None, type_hint: str | None = None, ) -> str | None: - """Return the exact C++ map type for a contextual RHS target. + """Return the exact nullable collection type for an RHS target. Declarations need the explicit hint before their lexical collection binding is activated; reassignments can use the active name registry, while UDT fields are resolved from the target expression itself. - Returning ``None`` for every non-map shape deliberately preserves the - established generic expression lowering byte-for-byte. + Arrays are deliberately excluded until their runtime representation + can distinguish ``na`` from a valid empty ID. """ spec = self._type_spec_from_hint_name(type_hint) if type_hint else None if spec is None and target_node is not None: spec = self._type_spec_from_expr(target_node) if spec is None and name is not None: spec = self._collection_spec_for_name(name) - if spec is None or spec.kind != "map": + if spec is None or spec.kind not in {"map", "matrix"}: return None return self._type_spec_to_cpp(spec) @@ -503,6 +504,27 @@ def _visit_var_decl(self, node: VarDecl, lines: list[str], pad: str) -> None: and type_spec.kind in {"map", "matrix"} ): target_cpp_type = self._type_spec_to_cpp(type_spec) + if ( + target_cpp_type is not None + and type_spec is not None + and type_spec.kind in {"map", "matrix"} + and isinstance(node.value, (IfStmt, SwitchStmt)) + ): + lines.append(f"{pad}if (!{flag_expr}) {{") + self._visit_if_switch_expr( + node.value, + target_expr, + lines, + len(pad) // 4 + 1, + target_cpp_type=target_cpp_type, + ) + # The flag belongs after the complete selection. If it + # were emitted in an individual branch, unmatched paths + # or later arms could retry a Pine ``var`` initializer, + # violating first-reach / one-time semantics. + lines.append(f"{pad} {flag_expr} = true;") + lines.append(f"{pad}}}") + return if target_cpp_type is not None: init_cpp = self._visit_rhs_value( node.value, @@ -739,10 +761,10 @@ def remember_local_type(cpp_type: str | None) -> None: selection_cpp_type = ( cpp_type if (cpp_type is not None - and (cpp_type.startswith("PineMap<") + and (self._is_nullable_collection_cpp_type(cpp_type) or cpp_type in DRAWING_TYPE_TO_CPP.values() or cpp_type in self._udt_defs)) - else self._map_target_cpp_type( + else self._nullable_collection_target_cpp_type( name=node.name, type_hint=node.type_hint, ) @@ -844,7 +866,7 @@ def remember_local_type(cpp_type: str | None) -> None: # General declaration cpp_type = self._type_for_decl(node) if not is_global_member else None - target_cpp_type = cpp_type or self._map_target_cpp_type( + target_cpp_type = cpp_type or self._nullable_collection_target_cpp_type( name=node.name, type_hint=node.type_hint, ) @@ -883,15 +905,96 @@ def _compound_assign_rhs(target_read: str, op: str, val_cpp: str) -> str | None: return f"std::fmod((double)({target_read}), (double)({val_cpp}))" return None + def _matrix_rhs_specs(self, value: ASTNode | None) -> list[TypeSpec]: + """Return every known concrete matrix type in an RHS selection. + + ``_type_spec_from_expr`` intentionally returns ``None`` for a + selection whose concrete arms disagree. Reassignment validation must + still inspect those arms individually: otherwise + ``matrix := cond ? matrix : matrix`` evades the type + check precisely because the complete selection has no unified type. + Explicit and implicit ``na`` arms contribute no concrete type. + """ + if self._selection_node_is_na(value): + return [] + if isinstance(value, Ternary): + return ( + self._matrix_rhs_specs(value.true_val) + + self._matrix_rhs_specs(value.false_val) + ) + if isinstance(value, IfStmt): + return ( + self._matrix_rhs_specs( + self._selection_terminal_expr(value.body) + ) + + self._matrix_rhs_specs( + self._selection_terminal_expr(value.else_body) + ) + ) + if isinstance(value, SwitchStmt): + specs: list[TypeSpec] = [] + for _case_expr, case_body in value.cases: + specs.extend( + self._matrix_rhs_specs( + self._selection_terminal_expr(case_body) + ) + ) + specs.extend( + self._matrix_rhs_specs( + self._selection_terminal_expr(value.default_body) + ) + ) + return specs + spec = self._type_spec_from_expr(value) + return [spec] if spec is not None and spec.kind == "matrix" else [] + + def _validate_matrix_reassignment( + self, + node: Assignment, + target_name: str | None, + ) -> None: + """Reject any known matrix RHS whose element type changes the LHS. + + This check belongs to the lexical target rather than any particular + storage representation. A matrix may be a persistent ``var``, a plain + global member, a callable local, or a UDT field, but Pine forbids + changing its declared element type in every case. + """ + if node.op != ":=": + return + target_spec = self._type_spec_from_expr(node.target) + if target_spec is None and target_name is not None: + target_spec = self._collection_spec_for_name(target_name) + if target_spec is None or target_spec.kind != "matrix": + return + target_label = target_name + if target_label is None and isinstance(node.target, MemberAccess): + target_label = node.target.member + target_label = target_label or "" + for rhs_spec in self._matrix_rhs_specs(node.value): + if rhs_spec.element == target_spec.element: + continue + self._codegen_error( + node, + f"matrix '{target_label}' element type mismatch on reassignment: " + f"expected {self._type_spec_to_cpp(target_spec)}, " + f"got {self._type_spec_to_cpp(rhs_spec)}", + ) + def _visit_assignment(self, node: Assignment, lines: list[str], pad: str) -> None: if isinstance(node.value, FuncCall) and self._is_skip_expr(node.value): return + # Validate against the active lexical target before any syntax-specific + # early return. In particular, if/switch expressions and UDT-field + # targets otherwise bypass the ordinary identifier assignment path. + target_name = self._get_target_name(node.target) + self._validate_matrix_reassignment(node, target_name) + # If/switch expression in assignment: x := if cond ... if isinstance(node.value, (IfStmt, SwitchStmt)): - target_name = self._get_target_name(node.target) safe = self._safe_name(target_name) if target_name else self._visit_expr(node.target) - selection_cpp_type = self._map_target_cpp_type( + selection_cpp_type = self._nullable_collection_target_cpp_type( name=target_name, target_node=node.target if target_name is None else None, ) @@ -916,7 +1019,6 @@ def _visit_assignment(self, node: Assignment, lines: list[str], pad: str) -> Non return # Get target name - target_name = self._get_target_name(node.target) if target_name is None: # Assignment to a UDT field that was dropped from the emitted # struct because it had a drawing-only type (label/line/box/ @@ -935,7 +1037,7 @@ def _visit_assignment(self, node: Assignment, lines: list[str], pad: str) -> Non return # General expression target (e.g., member access) target_cpp = self._visit_expr(node.target) - target_cpp_type = self._map_target_cpp_type( + target_cpp_type = self._nullable_collection_target_cpp_type( target_node=node.target, ) if target_cpp_type is None: @@ -987,39 +1089,12 @@ def _visit_assignment(self, node: Assignment, lines: list[str], pad: str) -> Non op_char = node.op[0] # e.g., "+" from "+=" lines.append(f"{pad}{safe}.update({safe}[0] {op_char} {val_cpp});") elif target_name in self._var_names: - target_spec = self._collection_spec_for_name(target_name) - if (node.op == ":=" - and target_spec is not None - and target_spec.kind == "matrix" - and isinstance(node.value, FuncCall)): - rhs_fn, rhs_ns = self._resolve_callee(node.value.callee) - rhs_spec = None - if rhs_ns == "matrix" and rhs_fn == "new": - targs = self._template_args_from_call(node.value) if hasattr(node.value, "annotations") else [] - elem = self._type_spec_from_hint_name(targs[0]) if targs else TypeSpec.primitive("float") - rhs_spec = TypeSpec.matrix(elem) - elif rhs_ns == "matrix" and rhs_fn in MATRIX_RETURNING_METHODS: - rcv = self._extract_receiver_name(node.value) - rhs_spec = ( - self._collection_spec_for_name(rcv) - if rcv is not None - else None - ) - if rhs_spec is not None: - lhs_spec = target_spec - if rhs_spec.element != lhs_spec.element: - self._codegen_error( - node, - f"matrix '{target_name}' element type mismatch on reassignment: " - f"expected {self._type_spec_to_cpp(lhs_spec)}, " - f"got {self._type_spec_to_cpp(rhs_spec)}", - ) # A bare-``na`` reassignment must adopt the target's declared scalar # type (``x := na`` -> ``na()`` not ``na()``); otherwise # a double NaN is stored into an int/int64_t/bool member (UB, defeats # is_na()). Only computed for bare na — every other RHS is # unaffected. - tct = self._map_target_cpp_type(name=target_name) + tct = self._nullable_collection_target_cpp_type(name=target_name) if tct is None: tct = self._udt_target_cpp_type(target_name=target_name) if tct is None and self._is_na_expr(node.value): @@ -1034,7 +1109,7 @@ def _visit_assignment(self, node: Assignment, lines: list[str], pad: str) -> Non else: lines.append(f"{pad}{safe} {node.op} {val_cpp};") else: - tct = self._map_target_cpp_type(name=target_name) + tct = self._nullable_collection_target_cpp_type(name=target_name) if tct is None: tct = self._udt_target_cpp_type(target_name=target_name) if tct is None and self._is_na_expr(node.value): @@ -1791,6 +1866,18 @@ def _visit_if_switch_expr( ) -> None: """Emit an if/switch used as an expression, assigning to target.""" pad = " " * indent + nullable_collection_target = self._is_nullable_collection_cpp_type( + target_cpp_type + ) + + def emit_implicit_na_fallback() -> None: + """Reset a nullable ID when an expression has no matching arm.""" + lines.append(f"{pad}else {{") + lines.append( + f"{pad} {target} = {target_cpp_type}{{}};" + ) + lines.append(f"{pad}}}") + if isinstance(node, IfStmt): cond = self._visit_expr(node.condition) lines.append(f"{pad}if ({cond}) {{") @@ -1822,6 +1909,12 @@ def _visit_if_switch_expr( target_cpp_type=target_cpp_type, ) lines.append(f"{pad}}}") + elif nullable_collection_target: + # A Pine if-expression without an else evaluates to ``na`` on + # the unmatched path. This assignment is essential for + # non-var globals and reassignments: retaining the prior bar's + # map/matrix ID would turn the expression into implicit state. + emit_implicit_na_fallback() elif isinstance(node, SwitchStmt): if node.expr: expr_var = f"__switch_val_{self._switch_counter}" @@ -1853,12 +1946,28 @@ def _visit_if_switch_expr( ) lines.append(f"{pad}}}") if node.default_body: - lines.append(f"{pad}else {{") - self._emit_block_with_assign( - node.default_body, - target, - lines, - indent + 1, - target_cpp_type=target_cpp_type, - ) - lines.append(f"{pad}}}") + if node.cases: + lines.append(f"{pad}else {{") + self._emit_block_with_assign( + node.default_body, + target, + lines, + indent + 1, + target_cpp_type=target_cpp_type, + ) + lines.append(f"{pad}}}") + else: + self._emit_block_with_assign( + node.default_body, + target, + lines, + indent, + target_cpp_type=target_cpp_type, + ) + elif nullable_collection_target: + if node.cases: + emit_implicit_na_fallback() + else: + lines.append( + f"{pad}{target} = {target_cpp_type}{{}};" + ) diff --git a/tests/test_matrix_na_lifecycle.py b/tests/test_matrix_na_lifecycle.py new file mode 100644 index 0000000..01e811b --- /dev/null +++ b/tests/test_matrix_na_lifecycle.py @@ -0,0 +1,741 @@ +"""Pine v6 target-typed ``na`` lifecycle for matrix IDs.""" + +from __future__ import annotations + +import re + +import pytest + +from pineforge_codegen import transpile +from pineforge_codegen.errors import CompileError +from tests._compile import compile_cpp +from tests.test_runtime_var_initialization import _compile_and_run + + +_DIRECT_SOURCE = r'''//@version=6 +strategy("matrix na direct lifecycle") +probe(bool gate, float seed) => + if gate + var matrix state = na + bool startedNull = na(state) + if startedNull + state := matrix.new(1, 1, seed) + state.set(0, 0, state.get(0, 0) + 1.0) + if bar_index > 10 + state := na + (startedNull ? 1000.0 : 0.0) + state.get(0, 0) + else + -1.0 +never = probe(false, 1.0) +early = probe(bar_index >= 1, 10.0) +late = probe(bar_index >= 2, 100.0) +''' + + +_METHOD_SOURCE = r'''//@version=6 +strategy("matrix na method lifecycle") +type Carrier + float seed +method probe(Carrier self, bool gate) => + if gate + var matrix state = na + bool startedNull = na(state) + if startedNull + state := matrix.new(1, 1, self.seed) + state.set(0, 0, state.get(0, 0) + 1.0) + if bar_index > 10 + state := na + (startedNull ? 1000.0 : 0.0) + state.get(0, 0) + else + -1.0 +var Carrier firstCarrier = Carrier.new(1.0) +var Carrier secondCarrier = Carrier.new(10.0) +never = firstCarrier.probe(false) +early = secondCarrier.probe(bar_index >= 1) +late = firstCarrier.probe(bar_index >= 2) +''' + + +_DRIVER = r''' +#include +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + Bar{3.0, 4.0, 2.0, 3.0, 1.0, 120000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 3); + if (!strategy.last_error().empty()) return 2; + std::cout << std::fixed << std::setprecision(1) + << strategy.never << " " + << strategy.early << " " + << strategy.late << "\n"; +} +''' + + +def _assert_cloned_state_is_target_typed(cpp: str) -> None: + assert not re.search(r"state(?:_cs\d+)? = na\(\);", cpp) + for target in ("state", "state_cs1", "state_cs2"): + assert re.search( + rf"^\s+(?:this->)?{target} = PineMatrix\{{\}};$", cpp, re.M + ) + + +def test_callable_matrix_na_lifecycle_keeps_written_callsites_independent() -> None: + cpp = transpile(_DIRECT_SOURCE) + _assert_cloned_state_is_target_typed(cpp) + compile_cpp(cpp, label="matrix-na-direct-lifecycle") + assert _compile_and_run(cpp + _DRIVER) == "-1.0 12.0 1101.0\n" + + +def test_method_matrix_na_lifecycle_keeps_written_callsites_independent() -> None: + cpp = transpile(_METHOD_SOURCE) + _assert_cloned_state_is_target_typed(cpp) + compile_cpp(cpp, label="matrix-na-method-lifecycle") + assert _compile_and_run(cpp + _DRIVER) == "-1.0 12.0 1002.0\n" + + +def test_matrix_na_reset_is_reachable_and_reinitializes_on_next_bar() -> None: + source = r'''//@version=6 +strategy("matrix na reachable reset") +probe(bool gate, bool clear) => + var matrix state = na + bool startedNull = na(state) + float result = -1.0 + if gate + if startedNull + state := matrix.new(1, 1, 0.0) + state.set(0, 0, state.get(0, 0) + 1.0) + result := (startedNull ? 1000.0 : 0.0) + state.get(0, 0) + if clear + state := na + result +value = probe(bar_index > 0, bar_index == 2) +var float step0 = na +var float step1 = na +var float step2 = na +var float step3 = na +var float step4 = na +if bar_index == 0 + step0 := value +if bar_index == 1 + step1 := value +if bar_index == 2 + step2 := value +if bar_index == 3 + step3 := value +if bar_index == 4 + step4 := value +''' + cpp = transpile(source) + assert "state = PineMatrix{};" in cpp + driver = r''' +#include +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + Bar{3.0, 4.0, 2.0, 3.0, 1.0, 120000}, + Bar{4.0, 5.0, 3.0, 4.0, 1.0, 180000}, + Bar{5.0, 6.0, 4.0, 5.0, 1.0, 240000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 5); + if (!strategy.last_error().empty()) return 2; + std::cout << std::fixed << std::setprecision(1) + << strategy.step0 << " " << strategy.step1 << " " + << strategy.step2 << " " << strategy.step3 << " " + << strategy.step4 << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "-1.0 1001.0 2.0 1001.0 2.0\n" + + +def test_all_supported_matrix_element_families_use_typed_na_ids() -> None: + source = r'''//@version=6 +strategy("typed matrix na families") +matrix floats = na +matrix ints = na +matrix strings = na +matrix bools = na +bool startedNull = na(floats) and na(ints) and na(strings) and na(bools) +floats := matrix.new(0, 0) +ints := matrix.new(0, 0) +strings := matrix.new(0, 0) +bools := matrix.new(0, 0) +bool emptyIdsAreValid = not na(floats) and not na(ints) and not na(strings) and not na(bools) +floats := na +ints := na +strings := na +bools := na +bool resetNull = na(floats) and na(ints) and na(strings) and na(bools) +''' + cpp = transpile(source) + expected = { + "floats": "PineMatrix", + "ints": "PineGenericMatrix", + "strings": "PineGenericMatrix", + "bools": "PineGenericMatrix", + } + for name, cpp_type in expected.items(): + assert cpp.count(f"{name} = {cpp_type}{{}};") == 2 + assert f"{name} = na();" not in cpp + compile_cpp(cpp, label="matrix-na-element-families") + driver = r''' +#include +int main() { + Bar bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}; + GeneratedStrategy strategy; + strategy.run(&bar, 1); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.startedNull << " " + << strategy.emptyIdsAreValid << " " + << strategy.resetNull << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "1 1 1\n" + + +def test_matrix_na_selection_arms_inherit_the_declared_target_type() -> None: + source = r'''//@version=6 +strategy("matrix na selections") +matrix ternaryValue = bar_index > 0 ? matrix.new(0, 0) : na +matrix ifValue = if bar_index > 0 + matrix.new(0, 0) +else + na +matrix switchValue = switch bar_index + 0 => na + => matrix.new(0, 0) +ternaryValue := bar_index > 1 ? na : matrix.new(0, 0) +ifValue := if bar_index > 1 + na +else + matrix.new(0, 0) +switchValue := switch bar_index + 2 => na + => matrix.new(0, 0) +''' + cpp = transpile(source) + assert "PineGenericMatrix{}" in cpp + assert "PineGenericMatrix{}" in cpp + assert "PineGenericMatrix{}" in cpp + assert not re.search( + r"(?:ternaryValue|ifValue|switchValue) = na\(\);", cpp + ) + compile_cpp(cpp, label="matrix-na-selections") + + +def test_type_only_matrix_contexts_emit_headers_and_target_typed_na() -> None: + sources = { + "float-global": r'''//@version=6 +strategy("typed na float global") +matrix state = na +bool startedNull = na(state) +''', + "generic-parameter": r'''//@version=6 +strategy("typed na generic parameter") +isNull(matrix state) => na(state) +bool startedNull = isNull(na) +''', + "generic-udt-field": r'''//@version=6 +strategy("typed na generic UDT field") +type Holder + matrix state = na +var Holder defaulted = Holder.new() +var Holder provided = Holder.new(na) +provided.state := na +bool startedNull = na(defaulted.state) and na(provided.state) +''', + } + for label, source in sources.items(): + cpp = transpile(source) + assert "#include " in cpp + if label.startswith("generic"): + assert "#include " in cpp + assert "PineGenericMatrix{}" in cpp + assert not re.search(r"(?:state|\.state) = na\(\)", cpp) + compile_cpp(cpp, label=f"matrix-na-{label}") + + +def test_explicit_typed_na_matrix_rejects_element_type_change() -> None: + source = r'''//@version=6 +strategy("typed na matrix mismatch") +matrix state = na +state := matrix.new(1, 1, 0.0) +''' + with pytest.raises( + CompileError, + match=r"element type mismatch.*expected PineGenericMatrix.*got PineMatrix", + ): + transpile(source) + + +@pytest.mark.parametrize( + ("label", "assignment"), + [ + ("identifier", "state := other"), + ("ternary", "state := bar_index == 0 ? other : na"), + ( + "if", + """state := if bar_index == 0 + other +else + na""", + ), + ( + "switch", + """state := switch bar_index + 0 => other + => na""", + ), + ("method-return", "state := other.copy()"), + ("mixed-arms", "state := bar_index == 0 ? same : other"), + ], +) +def test_matrix_reassignment_rejects_complete_incompatible_rhs( + label: str, + assignment: str, +) -> None: + source = f'''//@version=6 +strategy("matrix reassignment {label}") +matrix state = na +matrix same = matrix.new(1, 1, 0) +matrix other = matrix.new(1, 1, 0.0) +{assignment} +''' + with pytest.raises( + CompileError, + match=r"element type mismatch.*expected PineGenericMatrix.*got PineMatrix", + ): + transpile(source) + + +def test_callable_local_matrix_reassignment_rejects_identifier_type_change() -> None: + source = r'''//@version=6 +strategy("callable matrix reassignment mismatch") +probe() => + matrix state = na + matrix other = matrix.new(1, 1, 0.0) + state := other + 0 +value = probe() +''' + with pytest.raises( + CompileError, + match=r"element type mismatch.*expected PineGenericMatrix.*got PineMatrix", + ): + transpile(source) + + +@pytest.mark.parametrize( + "assignment", + [ + "holder.state := other", + """holder.state := if bar_index == 0 + other +else + na""", + ], +) +def test_udt_field_matrix_reassignment_rejects_element_type_change( + assignment: str, +) -> None: + source = f'''//@version=6 +strategy("UDT field matrix reassignment mismatch") +type Holder + matrix state = na +var Holder holder = Holder.new() +matrix other = matrix.new(1, 1, 0.0) +{assignment} +''' + with pytest.raises( + CompileError, + match=r"element type mismatch.*expected PineGenericMatrix.*got PineMatrix", + ): + transpile(source) + + +def test_matrix_reassignment_complete_rhs_accepts_same_element_type() -> None: + source = r'''//@version=6 +strategy("same-type matrix reassignments") +matrix state = na +matrix other = matrix.new(1, 1, 0) +state := other +state := bar_index == 0 ? other : na +state := if bar_index == 0 + other +else + na +state := switch bar_index + 0 => other.copy() + => na +''' + cpp = transpile(source) + compile_cpp(cpp, label="matrix-na-same-type-reassignments") + + +def test_global_inferred_matrix_na_selection_keeps_matrix_type() -> None: + source = r'''//@version=6 +strategy("inferred matrix na global") +var base = matrix.new(0, 0) +selected = bar_index > 0 ? na : base +selectedNull = na(selected) +''' + cpp = transpile(source) + assert "PineGenericMatrix selected;" in cpp + assert "selected = na();" not in cpp + compile_cpp(cpp, label="matrix-na-inferred-global") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 2); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.selectedNull << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "1\n" + + +def test_callable_inferred_matrix_na_selections_keep_matrix_type() -> None: + source = r'''//@version=6 +strategy("inferred matrix na callable") +probe(int mode) => + base = matrix.new(0, 0) + ternaryValue = mode == 0 ? na : base + ifValue = if mode == 1 + na + else + base + switchValue = switch mode + 2 => na + => base + (na(ternaryValue) ? 100 : 0) + (na(ifValue) ? 10 : 0) + (na(switchValue) ? 1 : 0) +mode0 = probe(0) +mode1 = probe(1) +mode2 = probe(2) +''' + cpp = transpile(source) + for name in ("ternaryValue", "ifValue", "switchValue"): + assert re.search( + rf"^\s+PineGenericMatrix {name}(?: =|;)", cpp, re.M + ) + assert f"{name} = na();" not in cpp + compile_cpp(cpp, label="matrix-na-inferred-callable") + driver = r''' +#include +int main() { + Bar bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}; + GeneratedStrategy strategy; + strategy.run(&bar, 1); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.mode0 << " " << strategy.mode1 << " " + << strategy.mode2 << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "100 10 1\n" + + +def test_global_direct_matrix_constructors_infer_nullable_selection_types() -> None: + source = r'''//@version=6 +strategy("direct matrix constructor selections") +ternaryValue = bar_index == 0 ? matrix.new(0, 0) : na +ifValue = if bar_index == 0 + matrix.new(0, 0) +switchValue = switch bar_index + 0 => matrix.new(0, 0) +ternaryNull = na(ternaryValue) +ifNull = na(ifValue) +switchNull = na(switchValue) +''' + cpp = transpile(source) + expected = { + "ternaryValue": "PineGenericMatrix", + "ifValue": "PineGenericMatrix", + "switchValue": "PineGenericMatrix", + } + for name, cpp_type in expected.items(): + assert f"{cpp_type} {name};" in cpp + assert f"{name} = na();" not in cpp + assert "ifValue = PineGenericMatrix{};" in cpp + assert "switchValue = PineGenericMatrix{};" in cpp + compile_cpp(cpp, label="matrix-na-direct-constructor-selections") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 2); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.ternaryNull << " " + << strategy.ifNull << " " + << strategy.switchNull << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "1 1 1\n" + + +def test_missing_if_switch_arms_reset_reassigned_matrix_and_map_ids() -> None: + source = r'''//@version=6 +strategy("nullable collection unmatched reset") +var matrix matrixState = na +var map mapState = na +if bar_index == 0 + matrixState := matrix.new(0, 0) + mapState := map.new() +matrixState := switch bar_index + 0 => matrixState +mapState := if bar_index == 0 + mapState +var bool matrixWasValid = false +var bool mapWasValid = false +if bar_index == 0 + matrixWasValid := not na(matrixState) + mapWasValid := not na(mapState) +matrixIsNull = na(matrixState) +mapIsNull = na(mapState) +''' + cpp = transpile(source) + assert "matrixState = PineGenericMatrix{};" in cpp + assert "mapState = PineMap{};" in cpp + compile_cpp(cpp, label="nullable-collection-unmatched-reset") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 2); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.matrixWasValid << " " + << strategy.mapWasValid << " " + << strategy.matrixIsNull << " " + << strategy.mapIsNull << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "1 1 1 1\n" + + +def test_udt_method_matrix_na_arguments_use_the_declared_parameter_type() -> None: + source = r'''//@version=6 +strategy("typed matrix method arguments") +type Carrier + float seed +method isNull(Carrier self, matrix state = na) => + na(state) +var Carrier carrier = Carrier.new(1.0) +identifierPositional = carrier.isNull(na) +identifierKeyword = carrier.isNull(state = na) +identifierOmitted = carrier.isNull() +temporaryPositional = Carrier.new(2.0).isNull(na) +temporaryKeyword = Carrier.new(3.0).isNull(state = na) +temporaryOmitted = Carrier.new(4.0).isNull() +''' + cpp = transpile(source) + assert "PineGenericMatrix state" in cpp + assert "na()" not in cpp + assert cpp.count("PineGenericMatrix{}") >= 6 + compile_cpp(cpp, label="matrix-na-udt-method-arguments") + driver = r''' +#include +int main() { + Bar bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}; + GeneratedStrategy strategy; + strategy.run(&bar, 1); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.identifierPositional << " " + << strategy.identifierKeyword << " " + << strategy.identifierOmitted << " " + << strategy.temporaryPositional << " " + << strategy.temporaryKeyword << " " + << strategy.temporaryOmitted << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "1 1 1 1 1 1\n" + + +def test_nullable_matrix_builtin_returns_keep_receiver_element_type() -> None: + source = r'''//@version=6 +strategy("nullable matrix builtin returns") +var matrix base = matrix.new(1, 1, 2.0) +methodInv = bar_index == 0 ? na : base.inv() +methodPinv = if bar_index == 0 + base.pinv() +else + na +functionalCopy = switch bar_index + 0 => matrix.copy(base) + => na +functionalTranspose = bar_index == 0 ? na : matrix.transpose(base) +functionalInv = if bar_index > 0 + matrix.inv(base) +methodInvNull = na(methodInv) +methodPinvNull = na(methodPinv) +functionalCopyNull = na(functionalCopy) +functionalTransposeNull = na(functionalTranspose) +functionalInvNull = na(functionalInv) +''' + cpp = transpile(source) + for name in ( + "methodInv", + "methodPinv", + "functionalCopy", + "functionalTranspose", + "functionalInv", + ): + assert f"PineMatrix {name};" in cpp + assert f"{name} = na();" not in cpp + assert "functionalInv = PineMatrix{};" in cpp + compile_cpp(cpp, label="matrix-na-builtin-return-selections") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 2); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.methodInvNull << " " + << strategy.methodPinvNull << " " + << strategy.functionalCopyNull << " " + << strategy.functionalTransposeNull << " " + << strategy.functionalInvNull << "\n"; +} +''' + assert _compile_and_run(cpp + driver) == "0 1 1 0 0\n" + + +def test_global_persistent_collection_selections_initialize_once() -> None: + source = r'''//@version=6 +strategy("global persistent collection selections") +var matrix explicitMatrix = if bar_index == 0 + matrix.new(1, 1, 7) +else + na +var map explicitMap = switch bar_index + 0 => map.new() + => na +var matrix missingMatrix = if bar_index > 0 + matrix.new(1, 1, 9) +var map missingMap = switch bar_index + 1 => map.new() +explicitMatrixNull = na(explicitMatrix) +explicitMapNull = na(explicitMap) +missingMatrixNull = na(missingMatrix) +missingMapNull = na(missingMap) +''' + cpp = transpile(source) + assert "/* unknown */" not in cpp + assert "missingMatrix = PineGenericMatrix{};" in cpp + assert "missingMap = PineMap{};" in cpp + for name in ( + "explicitMatrix", + "explicitMap", + "missingMatrix", + "missingMap", + ): + assert f"if (!_pf_var_init_{name})" in cpp + assert f"_pf_var_init_{name} = true;" in cpp + compile_cpp(cpp, label="global-persistent-collection-selections") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + Bar{3.0, 4.0, 2.0, 3.0, 1.0, 120000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 3); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.explicitMatrixNull << " " + << strategy.explicitMapNull << " " + << strategy.missingMatrixNull << " " + << strategy.missingMapNull << "\n"; +} +''' + # The explicit constructor arms run on bar 0 and persist. The missing + # fallbacks become typed null on bar 0 and must not retry on bars 1/2. + assert _compile_and_run(cpp + driver) == "0 0 1 1\n" + + +def test_callable_persistent_collection_selections_initialize_on_first_reach() -> None: + source = r'''//@version=6 +strategy("callable persistent collection selections") +probe(bool gate) => + int result = -1 + if gate + var matrix explicitMatrix = if bar_index == 1 + matrix.new(1, 1, 7) + else + na + var map explicitMap = switch bar_index + 1 => map.new() + => na + var matrix missingMatrix = if bar_index > 1 + matrix.new(1, 1, 9) + var map missingMap = switch bar_index + 2 => map.new() + result := (not na(explicitMatrix) ? 1000 : 0) + + (not na(explicitMap) ? 100 : 0) + + (na(missingMatrix) ? 10 : 0) + + (na(missingMap) ? 1 : 0) + result +value = probe(bar_index >= 1) +var int step0 = -99 +var int step1 = -99 +var int step2 = -99 +if bar_index == 0 + step0 := value +if bar_index == 1 + step1 := value +if bar_index == 2 + step2 := value +''' + cpp = transpile(source) + assert "/* unknown */" not in cpp + assert "this->missingMatrix = PineGenericMatrix{};" in cpp + assert "this->missingMap = PineMap{};" in cpp + for name in ( + "explicitMatrix", + "explicitMap", + "missingMatrix", + "missingMap", + ): + assert re.search(rf"if \(!this->_pf_var_init_.*{name}", cpp) + compile_cpp(cpp, label="callable-persistent-collection-selections") + driver = r''' +#include +int main() { + Bar bars[] = { + Bar{1.0, 2.0, 0.0, 1.0, 1.0, 0}, + Bar{2.0, 3.0, 1.0, 2.0, 1.0, 60000}, + Bar{3.0, 4.0, 2.0, 3.0, 1.0, 120000}, + }; + GeneratedStrategy strategy; + strategy.run(bars, 3); + if (!strategy.last_error().empty()) return 2; + std::cout << strategy.step0 << " " + << strategy.step1 << " " + << strategy.step2 << "\n"; +} +''' + # The declaration is first reached on bar 1. Its constructor/null choices + # persist on bar 2 even though every condition has flipped by then. + assert _compile_and_run(cpp + driver) == "-1 1111 1111\n"