fix: preserve target-typed UDT lifecycle semantics

pineforge_codegen/analyzer/base.py

@@ -2170,12 +2170,16 @@ def _resolved_user_call_name(call: FuncCall, owner: str | None) -> str | None:
                         spec = specs[param_idx] if param_idx < len(specs) else None
                         if spec is not None and spec.kind == "udt":
                             udt_name = spec.name
-                if udt_name is None:
-                    udt_name = self._udt_var_types.get(recv.name)
+            # Resolve the surviving exact global/lexical symbol before the
+            # flat raw-name registry. A later callable-local declaration can
+            # overwrite that registry and otherwise attach the wrong stateful
+            # method edge to wrappers and their written call sites.
             if udt_name is None:
                 spec = self._type_spec_from_expr(recv)
                 if spec is not None and spec.kind == "udt":
                     udt_name = spec.name
+            if udt_name is None and isinstance(recv, Identifier):
+                udt_name = self._udt_var_types.get(recv.name)
             key = f"{udt_name}.{method}" if udt_name else ""
             return key if key in func_defs else None
 
@@ -3201,6 +3205,63 @@ def _udt_name_from_ctor(self, value: ASTNode) -> str | None:
             return None
         return owner
 
+    def _udt_name_from_nullable_ctor_selection(
+        self, value: ASTNode | None
+    ) -> str | None:
+        """Exact user-UDT type for ctor-only nullable selections.
+
+        Keep this deliberately narrower than generic UDT expression
+        inference.  In particular, a terminal ``array.get(...UDT...)`` is a
+        reference-identity surface with its own fail-closed rules; treating
+        every UDT-valued expression selected against ``na`` as a by-value
+        return would accidentally bypass those rules.
+        """
+        nullable = object()
+
+        def terminal(body: list[ASTNode]) -> ASTNode | None:
+            if not body:
+                return None
+            node = body[-1]
+            return node.expr if isinstance(node, ExprStmt) else node
+
+        def resolve(node: ASTNode | None) -> str | None | object:
+            if node is None:
+                return nullable
+            if isinstance(node, ExprStmt):
+                return resolve(node.expr)
+            if isinstance(node, NaLiteral):
+                return nullable
+            direct = self._udt_name_from_ctor(node)
+            if direct in self._udt_fields:
+                return direct
+            if isinstance(node, Ternary):
+                return merge((resolve(node.true_val), resolve(node.false_val)))
+            if isinstance(node, IfStmt):
+                return merge((
+                    resolve(terminal(node.body)),
+                    resolve(terminal(node.else_body)),
+                ))
+            if isinstance(node, SwitchStmt):
+                results = [
+                    resolve(terminal(branch))
+                    for _case, branch in node.cases
+                ]
+                results.append(resolve(terminal(node.default_body)))
+                return merge(results)
+            return None
+
+        def merge(results) -> str | None | object:
+            resolved = list(results)
+            if any(item is None for item in resolved):
+                return None
+            concrete = {item for item in resolved if item is not nullable}
+            if not concrete:
+                return nullable
+            return next(iter(concrete)) if len(concrete) == 1 else None
+
+        result = resolve(value)
+        return result if isinstance(result, str) else None
+
     def _func_terminal_drawing_type(self, func_node: FuncDef) -> str | None:
         """Resolve the drawing-handle / UDT type of a function's terminal
         (return) expression for cases the direct ``_udt_name_from_ctor`` on the
@@ -3933,6 +3994,8 @@ def _visit_FuncDef(self, node: FuncDef) -> PineType:
         if node.body:
             ret_expr = terminal_ret_expr
             udt_ret = self._udt_name_from_ctor(ret_expr) if ret_expr is not None else None
+            if udt_ret is None:
+                udt_ret = self._udt_name_from_nullable_ctor_selection(ret_expr)
             if (udt_ret is None
                     and terminal_direct_return_spec is not None
                     and terminal_direct_return_spec.kind == "udt"):
@@ -4097,13 +4160,22 @@ def _visit_MethodDef(self, node) -> PineType:
         self._nested_ta_touched = set()
         terminal_ret_expr = self._direct_terminal_return_expr(node)
         return_type_spec = None
+        method_udt_return = None
         try:
             for stmt in node.body:
                 ret_type = self._visit(stmt)
             if terminal_ret_expr is not None:
                 terminal_spec = self._type_spec_from_expr(terminal_ret_expr)
                 if terminal_spec is not None and terminal_spec.kind == "map":
                     return_type_spec = terminal_spec
+                method_udt_return = (
+                    self._udt_name_from_ctor(terminal_ret_expr)
+                    or self._udt_name_from_nullable_ctor_selection(
+                        terminal_ret_expr
+                    )
+                )
+                if method_udt_return is not None:
+                    return_type_spec = TypeSpec.udt(method_udt_return)
         finally:
             self._global_scope = old_global
             self._collection_scope_stack.pop()
@@ -4123,6 +4195,8 @@ def _visit_MethodDef(self, node) -> PineType:
         if hi > lo:
             self._func_ta_ranges[method_key] = (lo, hi)
         self._symbols.exit_scope()
+        if method_udt_return is not None:
+            self._func_udt_return_types[method_key] = method_udt_return
 
         # Detect tuple return on UDT methods (mirrors the regular FuncDef logic
         # earlier in this file). Without this, codegen emits the method with a
@@ -4168,6 +4242,7 @@ def _visit_MethodDef(self, node) -> PineType:
             param_defaults=param_defaults,
             param_type_specs=param_specs,
             return_type_spec=return_type_spec,
+            udt_return_type=method_udt_return,
         )
         self._func_infos.append(fi)
         return PineType.VOID

pineforge_codegen/analyzer/types.py

@@ -194,6 +194,20 @@ def _type_spec_from_expr(self, value: ASTNode | None) -> TypeSpec | 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)
+
+            def direct_user_udt_ctor_name(node: ASTNode) -> str | None:
+                if not isinstance(node, FuncCall):
+                    return None
+                callee = node.callee
+                if not (
+                    isinstance(callee, MemberAccess)
+                    and isinstance(callee.object, Identifier)
+                    and callee.member == "new"
+                ):
+                    return None
+                name = callee.object.name
+                return name if name in self._udt_fields else None
+
             # Selecting between two values of the same user-defined type
             # preserves that receiver type.  Codegen already applies this
             # rule; the analyzer must agree so stateful method calls on a UDT
@@ -217,6 +231,22 @@ def _type_spec_from_expr(self, value: ASTNode | None) -> TypeSpec | 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
+            # identity returns continue to fail closed on their own surface.
+            true_ctor = direct_user_udt_ctor_name(value.true_val)
+            if (true_spec is not None
+                    and true_spec.kind == "udt"
+                    and true_spec.name == true_ctor
+                    and isinstance(value.false_val, NaLiteral)):
+                return true_spec
+            false_ctor = direct_user_udt_ctor_name(value.false_val)
+            if (false_spec is not None
+                    and false_spec.kind == "udt"
+                    and false_spec.name == false_ctor
+                    and isinstance(value.true_val, NaLiteral)):
+                return false_spec
             # Drawing handles are nullable reference-like values in Pine.  A
             # bare ``na`` arm therefore acquires the other arm's exact handle
             # type, just like the established PineMap path above.  Keep this