From 52b92621d4c1a334cf96f942a2d18990f09e2009 Mon Sep 17 00:00:00 2001 From: luisleo526 Date: Sun, 19 Jul 2026 07:23:22 +0800 Subject: [PATCH] Fix terminal primitive array get returns --- pineforge_codegen/analyzer/base.py | 15 ++ pineforge_codegen/analyzer/contracts.py | 13 +- pineforge_codegen/analyzer/types.py | 108 ++++++++++ pineforge_codegen/codegen/emit_top.py | 4 +- tests/test_array_checked_access.py | 2 +- tests/test_array_terminal_returns.py | 258 ++++++++++++++++++++++++ 6 files changed, 391 insertions(+), 9 deletions(-) create mode 100644 tests/test_array_terminal_returns.py diff --git a/pineforge_codegen/analyzer/base.py b/pineforge_codegen/analyzer/base.py index 8cc9690..f3d8083 100644 --- a/pineforge_codegen/analyzer/base.py +++ b/pineforge_codegen/analyzer/base.py @@ -1989,6 +1989,13 @@ def _visit_FuncDef(self, node: FuncDef) -> PineType: for name, spec in zip(node.params, inferred_param_specs) }, ) + terminal_array_get_return = self._terminal_array_get_return( + terminal_ret_expr, + { + name: spec + for name, spec in zip(node.params, inferred_param_specs) + }, + ) self._symbols.exit_scope() @@ -2049,6 +2056,14 @@ def _visit_FuncDef(self, node: FuncDef) -> PineType: if terminal_spec is not None: self._func_return_type_specs[node.name] = terminal_spec + if terminal_array_get_return is not None: + body_type, terminal_spec = terminal_array_get_return + # Preserve the exact primitive TypeSpec as well as the coarse + # PineType. Downstream collection construction consults this cache + # directly; without it, asking for the UDF's element type can + # re-visit the call and duplicate stateful call sites. + self._func_return_type_specs[node.name] = terminal_spec + # Store return type self._func_return_types[node.name] = body_type diff --git a/pineforge_codegen/analyzer/contracts.py b/pineforge_codegen/analyzer/contracts.py index 04b3d20..6aed714 100644 --- a/pineforge_codegen/analyzer/contracts.py +++ b/pineforge_codegen/analyzer/contracts.py @@ -88,11 +88,12 @@ class FuncInfo: # ``pivot hi`` parameter emits as ``pivot hi`` (not ``double hi``) and an # untyped ``s`` used as a string emits as ``std::string s``. param_type_specs: list = field(default_factory=list) - # ``TypeSpec`` of the function's return value when it is a collection the - # coarse ``return_type`` (PineType) cannot represent — arrays and maps - # (e.g. ``buildPDLevels() => array.from(...)`` -> ``std::vector``). - # UDT / drawing-handle returns use - # ``udt_return_type``; tuple returns use ``returns_tuple``. + # Exact ``TypeSpec`` of the function's return value when downstream + # expression inference needs more than the coarse ``return_type`` slot. + # This primarily carries arrays/maps and also direct-terminal primitive + # array-element reads, whose cached spec prevents inference from re-visiting + # a stateful UDF call. UDT / drawing-handle returns use ``udt_return_type``; + # tuple returns use ``returns_tuple``. return_type_spec: Any = None @@ -249,7 +250,7 @@ class AnalyzerContext: # Per-function var_members + series_vars (used when emitting per-function call-site variants): func_var_members: dict = field(default_factory=dict) func_series_vars: dict = field(default_factory=dict) - # Per-function collection-return TypeSpec (see FuncInfo.return_type_spec). + # Per-function exact return TypeSpec (see FuncInfo.return_type_spec). func_return_type_specs: dict = field(default_factory=dict) # var_name -> UDT type name for variables instantiated via TypeName.new(...) udt_var_types: dict[str, str] = field(default_factory=dict) diff --git a/pineforge_codegen/analyzer/types.py b/pineforge_codegen/analyzer/types.py index 8677362..8cb85e7 100644 --- a/pineforge_codegen/analyzer/types.py +++ b/pineforge_codegen/analyzer/types.py @@ -666,6 +666,114 @@ def _terminal_map_call_return( return PineType.STRING, None return None + def _terminal_array_get_return( + self, + value: ASTNode | None, + parameter_specs: dict[str, TypeSpec | None] | None = None, + ) -> tuple[PineType, TypeSpec] | None: + """Return exact metadata for a direct terminal ``array.get`` call. + + The regular expression ``TypeSpec`` pass already knows the element + type of both ``array.get(values, index)`` and ``values.get(index)``. + The coarse visitor, however, deliberately returns ``VOID`` for most + array methods. When such a call is the final expression of a UDF, + that ``VOID`` used to make codegen emit a ``double`` return type even + for primitive elements such as strings, booleans, and integers. + + Keep the refinement intentionally narrow: only the established + positional/keyword shapes of ``get`` participate, and the receiver + must resolve to an exact array ``TypeSpec``. Other array accessors, + mutations, range/view semantics, reference/ID-like elements, and + unresolved receivers remain on their existing paths. Returning UDTs + or nested collections by value would lose Pine reference identity even + when the generated C++ happens to compile, so this helper must not + expose those specs as an apparent fix. + """ + if not isinstance(value, FuncCall) or not isinstance( + value.callee, MemberAccess + ): + return None + + callee = value.callee + if callee.member != "get": + return None + + receiver: ASTNode | None = None + object_spec = None + object_is_parameter = False + object_is_visible_binding = False + if isinstance(callee.object, Identifier): + object_is_visible_binding = ( + self._symbols.resolve(callee.object.name) is not None + ) + object_is_parameter = ( + parameter_specs is not None + and callee.object.name in parameter_specs + ) + if object_is_parameter: + object_spec = parameter_specs.get(callee.object.name) + else: + object_spec = self._type_spec_from_expr(callee.object) + object_is_array_value = ( + object_spec is not None and object_spec.kind == "array" + ) + if ( + isinstance(callee.object, Identifier) + and callee.object.name == "array" + and not object_is_array_value + and not object_is_visible_binding + ): + # Functional forms supported by the checked-array emitter: + # array.get(values, index) + # array.get(values, index=index) + # array.get(id=values, index=index) + if len(value.args) == 2 and not value.kwargs: + receiver = value.args[0] + elif len(value.args) == 1 and set(value.kwargs) == {"index"}: + receiver = value.args[0] + elif not value.args and set(value.kwargs) == {"id", "index"}: + receiver = value.kwargs["id"] + else: + # Method forms supported by the checked-array emitter: + # values.get(index) + # values.get(index=index) + if len(value.args) == 1 and not value.kwargs: + receiver = callee.object + elif not value.args and set(value.kwargs) == {"index"}: + receiver = callee.object + + if receiver is None: + return None + # Temporary/arbitrary receivers have their own stateful-call and + # reference-identity questions. Re-inferring them here would revisit + # nested calls and mint extra call-site state. The registered residual + # is limited to exact identifier receivers. + if not isinstance(receiver, Identifier): + return None + + recv_spec = None + if parameter_specs is not None and receiver.name in parameter_specs: + recv_spec = parameter_specs.get(receiver.name) + # An unresolved parameter shadows any same-named global. + if recv_spec is None: + return None + else: + recv_spec = self._type_spec_from_expr(receiver) + if ( + recv_spec is None + or recv_spec.kind != "array" + or recv_spec.element is None + ): + return None + + element_spec = recv_spec.element + if element_spec.kind != "primitive": + return None + return_type = self._element_pine_type(element_spec) + if return_type == PineType.VOID: + return None + return return_type, element_spec + @staticmethod def _pine_type_to_spec(pine_type: PineType) -> TypeSpec: mapping = { diff --git a/pineforge_codegen/codegen/emit_top.py b/pineforge_codegen/codegen/emit_top.py index bb0d443..db1f45b 100644 --- a/pineforge_codegen/codegen/emit_top.py +++ b/pineforge_codegen/codegen/emit_top.py @@ -1303,8 +1303,8 @@ def _emit_func_def(self, fi: FuncInfo, lines: list[str], call_site_idx: int | No # struct (Line/Box/Label/Linefill), not the unknown lowercase name. ret_type = DRAWING_TYPE_TO_CPP.get(fi.udt_return_type, fi.udt_return_type) elif getattr(fi, "return_type_spec", None) is not None: - # Collection-returning function (array or map terminal) — emit the - # concrete C++ collection type from its inferred TypeSpec. + # Exact return TypeSpec (collections plus narrowly cached primitive + # terminal reads) takes precedence over the coarse PineType slot. ret_type = self._type_spec_to_cpp(fi.return_type_spec) else: ret_type = PINE_TYPE_TO_CPP.get(fi.return_type, "double") diff --git a/tests/test_array_checked_access.py b/tests/test_array_checked_access.py index bde01be..d105502 100644 --- a/tests/test_array_checked_access.py +++ b/tests/test_array_checked_access.py @@ -880,7 +880,7 @@ def test_array_typed_udf_parameter_methods_route_by_typespec(): # user-defined element type). assert "double rank_pos(std::vector& source, int index)" in cpp assert "double rank_kw(std::vector& source, int index)" in cpp - assert "double mutate_int(std::vector& source)" in cpp + assert "int mutate_int(std::vector& source)" in cpp assert "bool probe_bool(std::vector& source)" in cpp assert "double mutate_pivot(std::vector& source)" in cpp assert ( diff --git a/tests/test_array_terminal_returns.py b/tests/test_array_terminal_returns.py new file mode 100644 index 0000000..b3d2112 --- /dev/null +++ b/tests/test_array_terminal_returns.py @@ -0,0 +1,258 @@ +"""Exact primitive UDF return typing for direct terminal array ``get`` calls.""" + +from __future__ import annotations + +from itertools import product + +import pytest + +from pineforge_codegen import transpile +from tests._compile import compile_cpp + + +_GET_SHAPES_SOURCE = r'''//@version=6 +strategy("Terminal array get shapes") + +method_positional() => + array values = array.from("method-positional") + values.get(0) + +method_keyword(array values) => + values.get(index=0) + +functional_positional() => + array values = array.from("functional-positional") + array.get(values, 0) + +functional_mixed(array values) => + array.get(values, index=0) + +functional_keyword() => + array values = array.from("functional-keyword") + array.get(id=values, index=0) + +shadowed_namespace() => + array array = array.from("shadowed") + array.get(0) + +parameter_values = array.from("parameter") +a = method_positional() +b = method_keyword(parameter_values) +c = functional_positional() +d = functional_mixed(parameter_values) +e = functional_keyword() +f = shadowed_namespace() +if bar_index == 0 and a == "method-positional" and b == "parameter" and c == "functional-positional" and d == "parameter" and e == "functional-keyword" and f == "shadowed" + strategy.entry("L", strategy.long) +''' + + +_PRIMITIVE_TYPES_SOURCE = r'''//@version=6 +strategy("Terminal array get primitive types") + +get_int() => + array values = array.from(1) + values.get(0) + +get_float() => + array values = array.from(1.5) + array.get(values, 0) + +get_bool() => + array values = array.from(true) + values.get(index=0) + +get_string() => + array values = array.from("x") + array.get(id=values, index=0) + +get_color() => + array values = array.new() + values.push(color.red) + values.get(0) + +i = get_int() +f = get_float() +b = get_bool() +s = get_string() +c = get_color() +if bar_index == 0 and i == 1 and f == 1.5 and b and s == "x" and c == color.red + strategy.entry("L", strategy.long) +''' + + +def _masking_matrix_source( + *, + functional: bool, + parameter: bool, + reuse_receiver_name: bool, + string_first: bool, +) -> str: + """Build one frozen representation/order cell around the same defect.""" + + def function(kind: str) -> list[str]: + is_string = kind == "string" + func_name = f"read_{kind}" + type_name = "string" if is_string else "int" + literal = '"ok"' if is_string else "7" + receiver = "values" if reuse_receiver_name else f"{kind}_values" + if parameter: + lines = [f"{func_name}(array<{type_name}> {receiver}) =>"] + else: + lines = [ + f"{func_name}() =>", + f" array<{type_name}> {receiver} = array.from({literal})", + ] + terminal = ( + f"array.get({receiver}, 0)" + if functional + else f"{receiver}.get(0)" + ) + lines.append(f" {terminal}") + return lines + + ordered = ["string", "int"] if string_first else ["int", "string"] + lines = [ + "//@version=6", + 'strategy("Terminal array get masking matrix")', + ] + for kind in ordered: + lines.extend(function(kind)) + if parameter: + lines.extend( + [ + 'string_inputs = array.from("ok")', + "int_inputs = array.from(7)", + "observed_string = read_string(string_inputs)", + "observed_int = read_int(int_inputs)", + ] + ) + else: + lines.extend( + [ + "observed_string = read_string()", + "observed_int = read_int()", + ] + ) + lines.extend( + [ + 'if bar_index == 0 and observed_string == "ok" and observed_int == 7', + ' strategy.entry("L", strategy.long)', + ] + ) + return "\n".join(lines) + "\n" + + +def test_all_established_get_shapes_infer_string_and_compile(): + cpp = transpile(_GET_SHAPES_SOURCE) + for name in ( + "method_positional", + "method_keyword", + "functional_positional", + "functional_mixed", + "functional_keyword", + "shadowed_namespace", + ): + assert f"std::string {name}(" in cpp + compile_cpp(cpp) + + +def test_all_primitive_element_types_emit_exact_udf_returns_and_compile(): + cpp = transpile(_PRIMITIVE_TYPES_SOURCE) + for signature in ( + "int get_int(", + "double get_float(", + "bool get_bool(", + "std::string get_string(", + "int get_color(", + ): + assert signature in cpp + compile_cpp(cpp) + + +def test_exact_return_spec_prevents_stateful_call_revisit_in_array_from(): + source = r'''//@version=6 +strategy("Stateful terminal array get") +read_value() => + unused = ta.sma(close, 2) + values = array.from("x") + values.get(0) +wrapped = array.from(read_value()) +if bar_index == 0 and wrapped.get(0) == "x" + strategy.entry("L", strategy.long) +''' + cpp = transpile(source) + assert "std::string read_value_cs0(" in cpp + assert "read_value_cs1" not in cpp + assert "_ta_sma_1_cs" not in cpp + assert cpp.count("ta::SMA _ta_sma_1;") == 1 + compile_cpp(cpp) + + +@pytest.mark.parametrize( + ("functional", "parameter", "reuse_receiver_name", "string_first"), + product((False, True), repeat=4), +) +def test_terminal_get_representation_and_order_matrix( + functional: bool, + parameter: bool, + reuse_receiver_name: bool, + string_first: bool, +): + """Complete 2^4 input matrix; base/candidate is the fifth publish factor.""" + source = _masking_matrix_source( + functional=functional, + parameter=parameter, + reuse_receiver_name=reuse_receiver_name, + string_first=string_first, + ) + cpp = transpile(source) + assert "std::string read_string(" in cpp + assert "int read_int(" in cpp + compile_cpp(cpp) + + +def test_reference_like_array_elements_are_not_value_return_refined(): + source = r'''//@version=6 +strategy("Terminal UDT array get remains identity-gated") +type Item + int value +pick() => + array values = array.from(Item.new(1)) + values.get(0) +observed = pick() +''' + cpp = transpile(source) + assert "Item pick(" not in cpp + assert "double pick(" in cpp + + +def test_scalar_array_binding_is_not_misclassified_as_builtin_namespace(): + source = r'''//@version=6 +strategy("Array namespace shadow remains fail closed") +global_values = array.from("x") +blocked() => + int array = 1 + array.get(global_values, 0) +observed = blocked() +''' + cpp = transpile(source) + assert "std::string blocked(" not in cpp + assert "double blocked(" in cpp + + +def test_temporary_receiver_does_not_duplicate_nested_stateful_call_sites(): + source = r'''//@version=6 +strategy("Temporary receiver stays outside terminal get refinement") +producer() => + unused = ta.sma(close, 2) + "x" +outer() => + array.get(array.from(producer()), 0) +observed = outer() +''' + cpp = transpile(source) + assert "std::string outer_cs0(" not in cpp + assert "double outer_cs0(" in cpp + assert "producer_cs3" in cpp + assert "producer_cs4" not in cpp