diff --git a/pineforge_codegen/codegen/ta.py b/pineforge_codegen/codegen/ta.py index 54dcca1..7fa6b2f 100644 --- a/pineforge_codegen/codegen/ta.py +++ b/pineforge_codegen/codegen/ta.py @@ -293,23 +293,27 @@ def _expr_safe_for_ta_precalc(self, expr) -> bool: return False return False - def _ta_call_nodes_under_and_rhs(self) -> set[int]: - """``id(FuncCall)`` for ``ta.*`` sites below a lazy ``and`` RHS. - - This is deliberately narrower than a general conditional-execution - classifier. The campaign oracle pins chart-context ``ta.sma`` under - ``and``; it does not justify changing precalc for ``or``, ``?:``, or - other TA families. ``under_and_rhs`` is propagated through nested - expression nodes so calls such as ``ta.change(ta.sma(...))`` retain the - enclosing ``and`` context. + def _ta_call_nodes_by_lazy_scope(self) -> tuple[set[int], set[int]]: + """Classify chart ``ta.*`` sites below Pine-v6 lazy-expression edges. + + The first set contains sites below an ``and`` RHS, preserving the + already accepted lazy-SMA scope. The second contains sites below any + Pine-v6 lazy edge: an ``and``/``or`` RHS or either ``?:`` branch. The + broader set is consumed only by the recursive-EMA route; + other TA families retain their current precalculation behavior. """ - cached = getattr(self, "_and_rhs_ta_call_nodes", None) + cached = getattr(self, "_lazy_scope_ta_call_nodes", None) if cached is not None: return cached and_rhs: set[int] = set() + lazy_rhs: set[int] = set() - def note_ta_calls(expr, under_and_rhs: bool) -> None: + def note_ta_calls( + expr, + under_and_rhs: bool, + under_lazy_rhs: bool, + ) -> None: if expr is None: return if isinstance(expr, FuncCall): @@ -320,17 +324,20 @@ def note_ta_calls(expr, under_and_rhs: bool) -> None: and callee.object.name == "request" and callee.member in ("security", "security_lower_tf") ) - if under_and_rhs and isinstance(expr.callee, MemberAccess): + if isinstance(expr.callee, MemberAccess): obj = expr.callee.object if isinstance(obj, Identifier) and obj.name == "ta": - and_rhs.add(id(expr)) + if under_and_rhs: + and_rhs.add(id(expr)) + if under_lazy_rhs: + lazy_rhs.add(id(expr)) # A call target may itself be an evaluated expression, as in # ``array.new_float(...).get(0)``. Walk the callee subtree so # TA nested in a chained receiver inherits the surrounding # lazy context. Plain identifiers and namespace receivers are # leaves, so ordinary ``ta.sma`` / ``request.security`` calls # remain unaffected here. - note_ta_calls(callee, under_and_rhs) + note_ta_calls(callee, under_and_rhs, under_lazy_rhs) for idx, arg in enumerate(getattr(expr, "args", ()) or ()): # The third request.security* argument is evaluated by its # own security evaluator. Symbol, timeframe, and remaining @@ -338,60 +345,70 @@ def note_ta_calls(expr, under_and_rhs: bool) -> None: # inspected for lazy chart TA. if is_security and idx == 2: continue - note_ta_calls(arg, under_and_rhs) + note_ta_calls(arg, under_and_rhs, under_lazy_rhs) for key, value in (getattr(expr, "kwargs", None) or {}).items(): if is_security and key == "expression": continue - note_ta_calls(value, under_and_rhs) + note_ta_calls(value, under_and_rhs, under_lazy_rhs) return if isinstance(expr, BinOp): if expr.op == "and": # LHS always runs first; RHS is short-circuit conditional. - note_ta_calls(expr.left, under_and_rhs) - note_ta_calls(expr.right, True) + note_ta_calls(expr.left, under_and_rhs, under_lazy_rhs) + note_ta_calls(expr.right, True, True) + elif expr.op == "or": + # ``or`` preserves an enclosing ``and`` scope, but its own + # RHS is a Pine-v6 lazy edge for the EMA classifier. + note_ta_calls(expr.left, under_and_rhs, under_lazy_rhs) + note_ta_calls(expr.right, under_and_rhs, True) else: - # ``or`` is not a new opt-out boundary, but preserve an - # enclosing ``and`` RHS while walking through it. - note_ta_calls(expr.left, under_and_rhs) - note_ta_calls(expr.right, under_and_rhs) + note_ta_calls(expr.left, under_and_rhs, under_lazy_rhs) + note_ta_calls(expr.right, under_and_rhs, under_lazy_rhs) return if isinstance(expr, Ternary): - # A ternary alone is not evidence for changing precalc. - note_ta_calls(expr.condition, under_and_rhs) - note_ta_calls(expr.true_val, under_and_rhs) - note_ta_calls(expr.false_val, under_and_rhs) + note_ta_calls(expr.condition, under_and_rhs, under_lazy_rhs) + note_ta_calls(expr.true_val, under_and_rhs, True) + note_ta_calls(expr.false_val, under_and_rhs, True) return if isinstance(expr, UnaryOp): - note_ta_calls(expr.operand, under_and_rhs) + note_ta_calls(expr.operand, under_and_rhs, under_lazy_rhs) return if isinstance(expr, (MemberAccess, Subscript)): - note_ta_calls(getattr(expr, "object", None), under_and_rhs) - note_ta_calls(getattr(expr, "index", None), under_and_rhs) + note_ta_calls( + getattr(expr, "object", None), under_and_rhs, under_lazy_rhs + ) + note_ta_calls( + getattr(expr, "index", None), under_and_rhs, under_lazy_rhs + ) return if isinstance(expr, TupleLiteral): for elem in expr.elements: - note_ta_calls(elem, under_and_rhs) + note_ta_calls(elem, under_and_rhs, under_lazy_rhs) return def walk_stmt(stmt) -> None: if stmt is None: return if isinstance(stmt, VarDecl): - note_ta_calls(stmt.value, False) + note_ta_calls(stmt.value, False, False) return if isinstance(stmt, ExprStmt): - note_ta_calls(getattr(stmt, "value", None) or getattr(stmt, "expr", None), False) + note_ta_calls( + getattr(stmt, "value", None) or getattr(stmt, "expr", None), + False, + False, + ) return if isinstance(stmt, Assignment): - note_ta_calls(getattr(stmt, "target", None), False) - note_ta_calls(getattr(stmt, "value", None), False) + note_ta_calls(getattr(stmt, "target", None), False, False) + note_ta_calls(getattr(stmt, "value", None), False, False) return if isinstance(stmt, TupleAssign): - note_ta_calls(getattr(stmt, "value", None), False) + note_ta_calls(getattr(stmt, "value", None), False, False) return if isinstance(stmt, TypeDecl): for field in getattr(stmt, "fields", ()) or (): - note_ta_calls(getattr(field, "default", None), False) + note_ta_calls(getattr(field, "default", None), False, False) return # If / for / while / assign-like — best-effort field walk for attr in ("condition", "body", "else_body", "else_ifs", "value", "target", "iterable"): @@ -406,9 +423,9 @@ def walk_stmt(stmt) -> None: elif hasattr(item, "body") or hasattr(item, "name") or hasattr(item, "value"): walk_stmt(item) else: - note_ta_calls(item, False) + note_ta_calls(item, False, False) elif hasattr(child, "op") or hasattr(child, "args") or hasattr(child, "left"): - note_ta_calls(child, False) + note_ta_calls(child, False, False) elif hasattr(child, "body") or hasattr(child, "value") or hasattr(child, "condition"): walk_stmt(child) @@ -424,8 +441,15 @@ def walk_stmt(stmt) -> None: for stmt in body: walk_stmt(stmt) - self._and_rhs_ta_call_nodes = and_rhs - return and_rhs + result = (and_rhs, lazy_rhs) + self._lazy_scope_ta_call_nodes = result + return result + + def _ta_call_nodes_under_and_rhs(self) -> set[int]: + return self._ta_call_nodes_by_lazy_scope()[0] + + def _ta_call_nodes_under_lazy_rhs(self) -> set[int]: + return self._ta_call_nodes_by_lazy_scope()[1] def _ta_site_uses_precalc(self, site: "TACallSite") -> bool: """Whether a static TA site can safely read from ``_precalc_*``. @@ -440,17 +464,25 @@ def _ta_site_uses_precalc(self, site: "TACallSite") -> bool: A chart-context ``ta.sma`` nested under an ``and`` RHS is also opted out: precalc advances it every bar, which is eager and TV-incorrect for - the pinned dual-volume-SMA case. Other TA families, ``or``/``?:``-only - sites, and request.security sites retain their existing behavior.""" + the pinned dual-volume-SMA case. Recursive chart ``ta.ema`` sites below + any Pine-v6 lazy edge follow the same call-clock rule. Other TA families + and request.security sites retain their existing behavior.""" if not getattr(site, "is_static", False): return False node = getattr(site, "node", None) + ta_name = self._ta_name_from_site(site) if ( - self._ta_name_from_site(site) == "sma" + ta_name == "sma" and node is not None and id(node) in self._ta_call_nodes_under_and_rhs() ): return False + if ( + ta_name == "ema" + and node is not None + and id(node) in self._ta_call_nodes_under_lazy_rhs() + ): + return False return all(self._expr_safe_for_ta_precalc(arg) for arg in site.compute_args) def _security_ta_compute_args_for_site( diff --git a/tests/test_codegen_validation_fixes.py b/tests/test_codegen_validation_fixes.py index 3bcf2d6..5246a62 100644 --- a/tests/test_codegen_validation_fixes.py +++ b/tests/test_codegen_validation_fixes.py @@ -531,8 +531,8 @@ def test_ta_precalc_skips_nested_sma_below_and_rhs(): assert "_ta_sma_" in cpp and ".compute(current_bar_.close)" in cpp -def test_ta_precalc_and_rhs_scope_keeps_other_ta_and_other_lazy_constructs(): - """Only SMA under an ``and`` RHS changes precalc eligibility.""" +def test_ta_precalc_lazy_scope_routes_recursive_ema_only(): + """EMA follows Pine-v6 lazy edges while unrelated TA stays precalculated.""" cpp = _cpp( "pred = close > open\n" "a = pred and ta.ema(close, 20) > close\n" @@ -541,15 +541,79 @@ def test_ta_precalc_and_rhs_scope_keeps_other_ta_and_other_lazy_constructs(): "d = pred and ta.highest(close, 3) < high\n" "e = pred or close > ta.sma(close, 5)\n" "f = pred ? ta.sma(close, 7) : close\n" - "plot((a or b or c or d or e) ? f : close)" + "g = pred or ta.ema(close, 21) > close\n" + "h = pred ? ta.ema(close, 22) : close\n" + "i = ta.ema(close, 23)\n" + "plot((a or b or c or d or e or g) ? f + h + i : close)" ) - for name in ("ema", "roc", "lowest", "highest"): + for name in ("roc", "lowest", "highest"): assert f"std::vector _precalc__ta_{name}_" in cpp assert f"_use_precalc ? _precalc__ta_{name}_" in cpp + assert len(re.findall(r"std::vector _precalc__ta_ema_", cpp)) == 1 assert len(re.findall(r"std::vector _precalc__ta_sma_", cpp)) == 2 assert len(re.findall(r"_use_precalc \? _precalc__ta_sma_", cpp)) >= 2 +def test_recursive_ema_lazy_edges_preserve_eager_operand_positions(): + """Only OR RHS / ternary arms opt out; eager positions still precalc.""" + cpp = _cpp( + "pred = close > open\n" + "orLeft = ta.ema(close, 10) > close or pred\n" + "orRight = pred or ta.ema(close, 11) > close\n" + "ternaryCond = ta.ema(close, 12) > close ? close : open\n" + "ternaryTrue = pred ? ta.ema(close, 13) : close\n" + "ternaryFalse = pred ? close : ta.ema(close, 14)\n" + "plot(orLeft ? ternaryCond + ternaryTrue + ternaryFalse " + ": orRight ? close : open)" + ) + for idx in (1, 3): + assert f"std::vector _precalc__ta_ema_{idx}" in cpp + assert f"_use_precalc ? _precalc__ta_ema_{idx}" in cpp + for idx in (2, 4, 5): + assert f"ta::EMA _ta_ema_{idx};" in cpp + assert f"std::vector _precalc__ta_ema_{idx}" not in cpp + assert f"_use_precalc ? _precalc__ta_ema_{idx}" not in cpp + assert f"_ta_ema_{idx}.compute(current_bar_.close)" in cpp + + +def test_lazy_chart_ema_does_not_reclassify_security_evaluator_ema(): + """Security payload EMA keeps evaluator-local state, not chart scope.""" + cpp = _cpp( + "pred = close > open\n" + "chart = pred or ta.ema(close, 3) > close\n" + "sec = request.security(syminfo.tickerid, \"60\", " + "close > open or ta.ema(close, 4) > close)\n" + "plot(chart ? sec : close)" + ) + assert "std::vector _precalc__ta_ema_1" not in cpp + assert "_ta_ema_1.compute(current_bar_.close)" in cpp + security_body = cpp.split("void _eval_security_0", 1)[1].split("\n }", 1)[0] + assert "ta::EMA _sec0__ta_ema_2;" in cpp + assert "_sec0__ta_ema_2.compute(bar.close)" in security_body + assert " || " in security_body + + +def test_lazy_udf_ema_uses_callsite_state_without_precalc(): + """A lazy UDF call keeps a distinct EMA clock from an eager sibling.""" + cpp = _cpp( + "f() =>\n" + " ta.ema(close, 3)\n" + "pred = close > open\n" + "lazy = pred and f() > close\n" + "eager = f()\n" + "plot(lazy ? eager : close)" + ) + assert "std::vector _precalc__ta_ema" not in cpp + assert "ta::EMA _ta_ema_1;" in cpp + assert "ta::EMA _ta_ema_1_cs1;" in cpp + lazy_body = cpp.split("double f_cs0()", 1)[1].split("\n }", 1)[0] + eager_body = cpp.split("double f_cs1()", 1)[1].split("\n }", 1)[0] + assert "_ta_ema_1.compute(current_bar_.close)" in lazy_body + assert "_ta_ema_1_cs1.compute(current_bar_.close)" in eager_body + assert "lazy = (pred &&" in cpp and "f_cs0()" in cpp + assert "eager = f_cs1();" in cpp + + def test_ta_precalc_keeps_security_context_and_rhs_sma(): """Security-local TA is distinct; a referenced chart expression is not.""" cpp = _cpp(