From 2f1bd11400a195a752b7596455b21918d58ab6cf Mon Sep 17 00:00:00 2001
From: luisleo526 <luisleo52655@gmail.com>
Date: Mon, 20 Jul 2026 05:07:21 +0800
Subject: [PATCH] fix: preserve forward local temporary-array return types

---
 pineforge_codegen/analyzer/base.py   |  88 ++++++-
 tests/test_array_terminal_returns.py | 373 +++++++++++++++++++++++++++
 2 files changed, 459 insertions(+), 2 deletions(-)

diff --git a/pineforge_codegen/analyzer/base.py b/pineforge_codegen/analyzer/base.py
index b161c62..9a8897c 100644
--- a/pineforge_codegen/analyzer/base.py
+++ b/pineforge_codegen/analyzer/base.py
@@ -603,6 +603,85 @@ def _refresh_direct_terminal_array_temporary_returns(self) -> None:
             if not changed:
                 break
 
+    def _direct_terminal_array_temporary_return_expr(
+        self,
+        func_def: FuncDef,
+        terminal: ASTNode | None,
+    ) -> ASTNode | None:
+        """Return the exact temporary read that determines a UDF return.
+
+        The established path accepts a direct terminal ``array.get``.  A
+        single ordinary local may also carry that same value to a bare
+        terminal identity return::
+
+            reader() =>
+                value = array.from(later()).get(0)
+                value
+
+        Capture only that adjacent one-hop lexical shape while the function
+        scope is still live.  Deferred reconciliation can then reuse the
+        initializer AST without re-visiting it after ``later`` is known.
+        Persistent/typed locals, alias chains, intervening statements, block
+        bindings, argument-bearing element calls, and every non-direct
+        producer stay outside this path.
+        """
+        if self._terminal_array_get_uses_direct_temporary(terminal):
+            return terminal
+        if not isinstance(terminal, Identifier):
+            return None
+
+        declarations = [
+            (index, stmt)
+            for index, stmt in enumerate(func_def.body[:-1])
+            if isinstance(stmt, VarDecl) and stmt.name == terminal.name
+        ]
+        if len(declarations) != 1:
+            return None
+        declaration_index, declaration = declarations[0]
+        if (
+            declaration.is_var
+            or declaration.is_varip
+            or declaration.type_hint is not None
+            or declaration_index != len(func_def.body) - 2
+        ):
+            return None
+
+        # Raw-name lookup is safe only because the exact declaration identity
+        # was attached to its live lexical Symbol by _visit_VarDecl.
+        symbol = self._symbols.resolve(terminal.name)
+        if (
+            symbol is None
+            or getattr(symbol, "_pf_decl_node_id", None) != id(declaration)
+        ):
+            return None
+
+        initializer = declaration.value
+        if not self._terminal_array_get_uses_direct_temporary(initializer):
+            return None
+        element_call = self._direct_terminal_array_temporary_element_call(
+            initializer
+        )
+        if (
+            element_call is None
+            or element_call.args
+            or element_call.kwargs
+        ):
+            return None
+        known_definition = self._func_defs.get(element_call.callee.name)
+        element_symbol = self._symbols.resolve(element_call.callee.name)
+        if element_symbol is not None and (
+            known_definition is None or element_symbol.scope != "global"
+        ):
+            self._error(
+                "Direct temporary-array element call "
+                f"'{element_call.callee.name}()' resolves to a local or "
+                "parameter, not a user-defined function.",
+                element_call.loc,
+            )
+        if known_definition is not None and known_definition.params:
+            return None
+        return initializer
+
     @classmethod
     def _direct_terminal_array_temporary_element_call(
         cls,
@@ -3186,9 +3265,14 @@ def _visit_FuncDef(self, node: FuncDef) -> PineType:
         # this narrow to map terminals so general/nonterminal inference and
         # generated output remain unchanged.
         terminal_ret_expr = self._direct_terminal_return_expr(node)
-        if self._terminal_array_get_uses_direct_temporary(terminal_ret_expr):
+        temporary_return_expr = (
+            self._direct_terminal_array_temporary_return_expr(
+                node, terminal_ret_expr
+            )
+        )
+        if temporary_return_expr is not None:
             self._direct_terminal_array_temporary_exprs[node.name] = (
-                terminal_ret_expr
+                temporary_return_expr
             )
         terminal_direct_return_spec = self._type_spec_from_expr(
             terminal_ret_expr
diff --git a/tests/test_array_terminal_returns.py b/tests/test_array_terminal_returns.py
index b277a9b..067f124 100644
--- a/tests/test_array_terminal_returns.py
+++ b/tests/test_array_terminal_returns.py
@@ -761,3 +761,376 @@ def test_temporary_builtin_array_producer_matrix(
     vector_type = "std::vector<std::string>" if string else "std::vector<int>"
     assert vector_type in cpp
     compile_cpp(cpp)
+
+
+@pytest.mark.parametrize(
+    ("literal", "expected_cpp"),
+    (
+        ("7", "int"),
+        ("1.5", "double"),
+        ("true", "bool"),
+        ('"x"', "std::string"),
+        ("color.red", "int"),
+    ),
+)
+def test_forward_local_identity_return_preserves_every_primitive_signature(
+    literal: str,
+    expected_cpp: str,
+):
+    source = f'''//@version=6
+strategy("Forward local identity primitive")
+read_value() =>
+    local_value = array.from(produce()).copy().get(0)
+    local_value
+produce() =>
+    {literal}
+observed = read_value()
+'''
+
+    cpp = transpile(source)
+    assert f"{expected_cpp} read_value(" in cpp
+    assert f"{expected_cpp} local_value =" in cpp
+    compile_cpp(cpp)
+
+
+@pytest.mark.parametrize(
+    ("functional", "keyword", "copy_layer", "producer_first"),
+    product((False, True), repeat=4),
+)
+def test_forward_local_identity_stateful_representation_matrix(
+    functional: bool,
+    keyword: bool,
+    copy_layer: bool,
+    producer_first: bool,
+):
+    """All 2^4 source-order/get/copy cells preserve one stateful call site."""
+    producer = '''produce() =>
+    unused = ta.sma(close, 2)
+    "x"
+'''
+    receiver = "array.from(produce())"
+    if copy_layer:
+        receiver = (
+            f"array.copy(id={receiver})"
+            if functional
+            else f"{receiver}.copy()"
+        )
+    if functional:
+        initializer = (
+            f"array.get(id={receiver}, index=0)"
+            if keyword
+            else f"array.get({receiver}, 0)"
+        )
+    else:
+        initializer = (
+            f"{receiver}.get(index=0)"
+            if keyword
+            else f"{receiver}.get(0)"
+        )
+    reader = (
+        "read_value() =>\n"
+        f"    local_value = {initializer}\n"
+        "    local_value\n"
+    )
+    definitions = producer + reader if producer_first else reader + producer
+    source = (
+        "//@version=6\n"
+        'strategy("Forward local identity representation matrix")\n'
+        f"{definitions}"
+        "observed = read_value()\n"
+    )
+
+    cpp = transpile(source)
+    assert "std::string read_value_cs0(" in cpp
+    assert "double read_value_cs0(" not in cpp
+    assert "std::string local_value =" in cpp
+    assert "produce_cs0()" in cpp
+    assert "produce_cs1" not in cpp
+    assert cpp.count("ta::SMA _ta_sma_1;") == 1
+    compile_cpp(cpp)
+
+
+@pytest.mark.parametrize(
+    "definitions",
+    (
+        '''reader() =>
+    local_value = array.from(reader()).copy().get(0)
+    local_value
+''',
+        '''first_reader() =>
+    local_value = array.from(second_reader()).copy().get(0)
+    local_value
+second_reader() =>
+    first_reader()
+''',
+    ),
+)
+def test_forward_local_identity_temporary_reader_recursion_is_rejected(
+    definitions: str,
+):
+    source = (
+        "//@version=6\n"
+        'strategy("Forward local identity recursion")\n'
+        f"{definitions}"
+        + (
+            "observed = reader()\n"
+            if definitions.startswith("reader")
+            else "observed = first_reader()\n"
+        )
+    )
+
+    with pytest.raises(
+        CompileError,
+        match="Recursive direct temporary-array reader cycle",
+    ):
+        transpile(source)
+
+
+@pytest.mark.parametrize("storage", ("var", "varip"))
+def test_forward_local_identity_persistent_local_remains_fail_closed(
+    storage: str,
+):
+    source = f'''//@version=6
+strategy("Forward persistent local boundary")
+blocked() =>
+    {storage} local_value = array.from(produce()).copy().get(0)
+    local_value
+produce() =>
+    "x"
+observed = blocked()
+'''
+
+    if storage == "varip":
+        with pytest.raises(CompileError, match="varip is not supported"):
+            transpile(source)
+        return
+    cpp = transpile(source)
+    assert "std::string blocked(" not in cpp
+
+
+@pytest.mark.parametrize(
+    ("type_hint", "literal", "expected_cpp", "compiles"),
+    (
+        ("float", "7", "double", True),
+        ("string", "7", "std::string", False),
+    ),
+)
+def test_forward_local_identity_explicit_type_hint_remains_authoritative(
+    type_hint: str,
+    literal: str,
+    expected_cpp: str,
+    compiles: bool,
+):
+    source = f'''//@version=6
+strategy("Forward typed local boundary")
+read_value() =>
+    {type_hint} local_value = array.from(produce()).copy().get(0)
+    local_value
+produce() =>
+    {literal}
+observed = read_value()
+'''
+
+    cpp = transpile(source)
+    assert f"{expected_cpp} read_value(" in cpp
+    assert f"{expected_cpp} local_value =" in cpp
+    assert "int read_value(" not in cpp
+    if compiles:
+        compile_cpp(cpp)
+    else:
+        # Pine rejects the incompatible declaration. Keep its authoritative
+        # hint instead of making the invalid source appear valid by refining
+        # the UDF return to the initializer's element type.
+        assert "std::vector<int>" in cpp
+
+
+@pytest.mark.parametrize(
+    "earlier_declaration",
+    ("float local_value = 0.0", "var local_value = 0.0"),
+)
+def test_forward_local_identity_requires_one_unique_same_named_declaration(
+    earlier_declaration: str,
+):
+    source = f'''//@version=6
+strategy("Forward duplicate local boundary")
+blocked() =>
+    {earlier_declaration}
+    local_value = array.from(produce()).copy().get(0)
+    local_value
+produce() =>
+    "x"
+observed = blocked()
+'''
+
+    cpp = transpile(source)
+    assert "std::string blocked(" not in cpp
+
+
+@pytest.mark.parametrize(
+    ("signature", "shadow_declaration", "call"),
+    (
+        ("blocked()", "    produce = 1\n", "blocked()"),
+        ("blocked(float produce)", "", "blocked(1.0)"),
+    ),
+)
+@pytest.mark.parametrize("producer_first", (False, True))
+def test_forward_local_identity_element_callee_shadow_remains_fail_closed(
+    signature: str,
+    shadow_declaration: str,
+    call: str,
+    producer_first: bool,
+):
+    producer = '''produce() =>
+    "x"
+'''
+    reader = f'''{signature} =>
+{shadow_declaration}    local_value = array.from(produce()).copy().get(0)
+    local_value
+'''
+    definitions = producer + reader if producer_first else reader + producer
+    source = (
+        "//@version=6\n"
+        'strategy("Forward element callee shadow boundary")\n'
+        f"{definitions}"
+        f"observed = {call}\n"
+    )
+
+    with pytest.raises(
+        CompileError,
+        match="resolves to a local or parameter",
+    ):
+        transpile(source)
+
+
+@pytest.mark.parametrize(
+    "body",
+    (
+        '''    first_value = array.from(produce()).copy().get(0)
+    local_value = first_value
+    local_value
+''',
+        '''    local_value = array.from(produce()).copy().get(0)
+    local_value := "replacement"
+    local_value
+''',
+        '''    [local_value, other] = [array.from(produce()).copy().get(0), 1]
+    local_value
+''',
+        '''    if close > open
+        local_value = array.from(produce()).copy().get(0)
+    local_value
+''',
+        '''    local_value = array.from(produce()).copy().get(0)
+    extras = array.from(local_value)
+    local_value
+''',
+    ),
+)
+def test_forward_local_identity_non_direct_bindings_remain_fail_closed(
+    body: str,
+):
+    source = (
+        "//@version=6\n"
+        'strategy("Forward non-direct local boundary")\n'
+        "blocked() =>\n"
+        f"{body}"
+        "produce() =>\n"
+        '    "x"\n'
+        "observed = blocked()\n"
+    )
+
+    cpp = transpile(source)
+    assert "std::string blocked(" not in cpp
+
+
+@pytest.mark.parametrize(
+    "initializer",
+    (
+        'array.from(produce("x")).copy().get(0)',
+        "make_values().get(0)",
+        "array.slice(array.from(produce()), 0, 1).get(0)",
+    ),
+)
+def test_forward_local_identity_non_direct_producers_remain_fail_closed(
+    initializer: str,
+):
+    producer = (
+        '''produce(string value) =>
+    value
+'''
+        if 'produce("x")' in initializer
+        else '''produce() =>
+    "x"
+make_values() =>
+    array.from("x")
+'''
+    )
+    source = (
+        "//@version=6\n"
+        'strategy("Forward non-direct producer boundary")\n'
+        "blocked() =>\n"
+        f"    local_value = {initializer}\n"
+        "    local_value\n"
+        f"{producer}"
+        "observed = blocked()\n"
+    )
+
+    if 'produce("x")' in initializer:
+        with pytest.raises(CompileError, match="Unknown function 'produce"):
+            transpile(source)
+        return
+    cpp = transpile(source)
+    assert "std::string blocked(" not in cpp
+
+
+def test_forward_local_identity_defaulted_parameter_still_fails_closed():
+    source = r'''//@version=6
+strategy("Forward local default parameter boundary")
+blocked() =>
+    local_value = array.from(produce()).copy().get(0)
+    local_value
+produce(string value = "x") =>
+    value
+observed = blocked()
+'''
+
+    with pytest.raises(CompileError, match="Unknown function 'produce"):
+        transpile(source)
+
+
+@pytest.mark.parametrize(
+    "preamble_and_initializer",
+    (
+        '''type Item
+    int value
+blocked() =>
+    local_value = array.from(Item.new(1)).copy().get(0)
+''',
+        '''blocked() =>
+    local_value = array.from(array.from("x")).copy().get(0)
+''',
+        '''blocked() =>
+    int array = 1
+    local_value = array.get(array.from(produce()), 0)
+''',
+    ),
+)
+def test_forward_local_identity_reference_or_shadowed_boundaries_stay_closed(
+    preamble_and_initializer: str,
+):
+    source = (
+        "//@version=6\n"
+        'strategy("Forward reference boundary")\n'
+        f"{preamble_and_initializer}"
+        "    local_value\n"
+        "produce() =>\n"
+        '    "x"\n'
+        "observed = blocked()\n"
+    )
+
+    if "int array = 1" in preamble_and_initializer:
+        with pytest.raises(CompileError, match="Unknown function 'produce"):
+            transpile(source)
+        return
+    cpp = transpile(source)
+    assert "std::string blocked(" not in cpp