diff --git a/pineforge_codegen/codegen/tables.py b/pineforge_codegen/codegen/tables.py index de4c8e2..483b792 100644 --- a/pineforge_codegen/codegen/tables.py +++ b/pineforge_codegen/codegen/tables.py @@ -686,6 +686,22 @@ def _checked_array_percentrank(a: str, args: list[str]) -> str: "put_all": lambda m, args: f"{m}.insert({args[0]}.begin(), {args[0]}.end())", } +# Declared signature order after removing the method receiver ``id``. This is +# consumed by the typed-map UDF-parameter lane; established global/local and +# namespace map lowerings intentionally retain their existing output. +MAP_METHOD_KWARGS: dict[str, list[str]] = { + "put": ["key", "value"], + "get": ["key"], + "remove": ["key"], + "contains": ["key"], + "size": [], + "clear": [], + "keys": [], + "values": [], + "copy": [], + "put_all": ["id2"], +} + def _matrix_add_row(m: str, args: list) -> str: """Pine ``matrix.add_row`` codegen. diff --git a/pineforge_codegen/codegen/visit_call.py b/pineforge_codegen/codegen/visit_call.py index c58f5e3..65969f7 100644 --- a/pineforge_codegen/codegen/visit_call.py +++ b/pineforge_codegen/codegen/visit_call.py @@ -152,6 +152,7 @@ DRAWING_NS, DRAWING_TYPE_TO_CPP, MAP_METHODS, + MAP_METHOD_KWARGS, MATH_FUNC_MAP, MATRIX_METHODS, MATRIX_METHOD_KWARGS, @@ -298,6 +299,47 @@ def _array_function_arg_nodes(self, method: str, node: FuncCall) -> list: node.args, node.kwargs, ["id", *param_names], lambda arg: arg ) + def _map_param_method_arg_nodes(self, method: str, node: FuncCall) -> list: + """Merge typed-map parameter method kwargs into signature order.""" + param_names = MAP_METHOD_KWARGS.get(method) + if param_names is None: + return list(node.args) + return _merge_kwargs(node.args, node.kwargs, param_names, lambda arg: arg) + + def _map_param_method_expr( + self, map_expr: str, method: str, arg_nodes: list, spec: TypeSpec, + ) -> str: + """Lower one typed-map parameter method with ordered, single-use args. + + Existing map templates may reference a key more than once (``get`` / + ``contains`` / ``remove``), while C++ assignment evaluates the RHS of + ``put`` before its subscript LHS. Bind supplied expressions through + nested lambdas so this lane deterministically evaluates key then value + exactly once without changing established global/local map output. + """ + args = [self._visit_expr(arg) for arg in arg_nodes] + occupied = "\n".join((map_expr, *args)) + counter = getattr(self, "_map_param_arg_counter", 0) + bindings: list[tuple[str, str]] = [] + for arg in args: + while True: + token = f"__pf_map_param_arg_{counter}" + counter += 1 + if token not in occupied: + break + bindings.append((token, arg)) + self._map_param_arg_counter = counter + + lowered = self._map_method_expr( + map_expr, method, [token for token, _arg in bindings], spec + ) + for token, arg in reversed(bindings): + lowered = ( + f"[&](auto&& {token})->decltype(auto){{ " + f"return {lowered}; }}(({arg}))" + ) + return lowered + def _visit_func_call(self, node: FuncCall) -> str: callee = node.callee if isinstance(callee, MemberAccess): @@ -418,6 +460,18 @@ def _visit_func_call(self, node: FuncCall) -> str: return self._array_method_expr( arr, meth_raw, margs, param_spec ) + if ( + param_spec is not None + and param_spec.kind == "map" + and meth_raw in MAP_METHODS + ): + m = self._safe_name(oname) + arg_nodes = self._map_param_method_arg_nodes( + meth_raw, node + ) + return self._map_param_method_expr( + m, meth_raw, arg_nodes, param_spec + ) # map.put / map.get / … must lower to unordered_map ops, not `.put` on C++ if oname in self._map_vars and meth_raw in MAP_METHODS: m = self._safe_name(oname) diff --git a/tests/test_map_param_methods.py b/tests/test_map_param_methods.py new file mode 100644 index 0000000..9bee017 --- /dev/null +++ b/tests/test_map_param_methods.py @@ -0,0 +1,342 @@ +"""Typed-map UDF parameter method dispatch regression coverage.""" + +from __future__ import annotations + +import os +import subprocess +import tempfile +from pathlib import Path + +import pytest + +from pineforge_codegen.errors import CompileError +from pineforge_codegen import transpile +from tests import _compile as compile_env + + +_SOURCE = """//@version=6 +strategy("Typed map parameter methods") +var order = array.new() +var array shadow_array = array.from(99.0) +var map shadow_map = map.new() +var matrix shadow_matrix = matrix.new(1, 1, 99.0) + +next_key() => + order.push(1) + "ordered" + +next_value() => + order.push(2) + 7 + +read_key() => + order.push(3) + "ordered" + +mutate_int(map target) => + target.put(value=next_value(), key=next_key()) + target.get(key=read_key()) + +mutate_float(map target) => + target.put("float", 2.5) + target.get("float") + +probe_bool(map target) => + target.put(key="bool", value=true) + target.contains(key="bool") and target.get(key="bool") + +probe_string(map target) => + target.put(value="ok", key="string") + target.get(key="string") == "ok" + +missing_float(map target) => + target.get(key="missing") == 0.0 + +missing_int(map target) => + target.get(key="missing") == 0 + +missing_bool(map target) => + not target.get(key="missing") + +missing_string(map target) => + target.get(key="missing") == "" + +inspect_int(map target) => + array key_list = target.keys() + array value_list = target.values() + map copied = target.copy() + key_list.size() + value_list.size() + copied.size() + target.size() + +remove_int(map target) => + target.remove(key="ordered") + +clear_int(map target) => + target.clear() + target.size() + +merge_int(map target, map source) => + target.put_all(id2=source) + target.get(key="other") + +shadow_array_probe(map shadow_array) => + shadow_array.put(key="a", value=11) + shadow_array.get(key="a") + +shadow_map_probe(map shadow_map) => + shadow_map.put(key="m", value="map") + shadow_map.get(key="m") == "map" + +shadow_matrix_probe(map shadow_matrix) => + shadow_matrix.put(key="mx", value=true) + shadow_matrix.get(key="mx") + +int_values = map.new() +float_values = map.new() +bool_values = map.new() +string_values = map.new() +empty_float_values = map.new() +empty_int_values = map.new() +empty_bool_values = map.new() +empty_string_values = map.new() +clear_values = map.new() +merge_target = map.new() +merge_source = map.new() +array_shadow_values = map.new() +map_shadow_values = map.new() +matrix_shadow_values = map.new() +map.put(clear_values, "clear", 1) +map.put(merge_source, "other", 8) +int_result = mutate_int(int_values) +int_size_after_put = map.size(int_values) +inspect_result = inspect_int(int_values) +float_result = mutate_float(float_values) +bool_result = probe_bool(bool_values) +string_result = probe_string(string_values) +missing_float_result = missing_float(empty_float_values) +missing_int_result = missing_int(empty_int_values) +missing_bool_result = missing_bool(empty_bool_values) +missing_string_result = missing_string(empty_string_values) +merge_result = merge_int(merge_target, merge_source) +merge_caller_size = map.size(merge_target) +shadow_array_result = shadow_array_probe(array_shadow_values) +shadow_map_result = shadow_map_probe(map_shadow_values) +shadow_matrix_result = shadow_matrix_probe(matrix_shadow_values) +removed_result = remove_int(int_values) +int_size_after_remove = map.size(int_values) +clear_result = clear_int(clear_values) +clear_caller_size = map.size(clear_values) +order_code = order.get(0) * 100 + order.get(1) * 10 + order.get(2) +order_size = order.size() +""" + + +def _function_body(cpp: str, signature: str) -> str: + start = cpp.index(signature) + end = cpp.index("\n }", start) + len("\n }") + return cpp[start:end] + + +def test_typed_map_parameter_methods_route_by_typespec_and_keywords(): + cpp = transpile(_SOURCE) + + assert "double mutate_int(std::unordered_map& target)" in cpp + assert "double mutate_float(std::unordered_map& target)" in cpp + assert "bool probe_bool(std::unordered_map& target)" in cpp + assert ( + "bool probe_string(std::unordered_map& target)" + in cpp + ) + + for raw_call in ( + "target.put(", + "target.get(", + "target.contains(", + "target.remove(", + "target.put_all(", + "shadow_array.put(", + "shadow_array.get(", + "shadow_map.put(", + "shadow_map.get(", + "shadow_matrix.put(", + "shadow_matrix.get(", + ): + assert raw_call not in cpp + + # Reversed keyword spelling still binds key before value. Both expressions + # execute once, and get's duplicated template key is also bound once. + put_line = next( + line for line in cpp.splitlines() + if "next_key()" in line and "next_value()" in line + ) + assert put_line.count("next_key()") == 1 + assert put_line.count("next_value()") == 1 + assert put_line.index("__pf_map_param_arg_0") < put_line.index( + "__pf_map_param_arg_1" + ) + assert put_line.index("}((next_value()))") < put_line.index( + "}((next_key()))" + ) + get_line = next( + line for line in cpp.splitlines() + if "read_key()" in line and "__pf_map_param_arg" in line + ) + assert get_line.count("read_key()") == 1 + + # Primitive missing-key defaults retain the parameter's value TypeSpec. + assert ": 0.0)" in _function_body(cpp, "bool missing_float(") + assert ": 0)" in _function_body(cpp, "bool missing_int(") + assert ": false)" in _function_body(cpp, "bool missing_bool(") + assert 'std::string("")' in _function_body(cpp, "bool missing_string(") + + # Zero-argument methods and typed collection results route through the + # helper without synthetic arguments. + inspect = _function_body(cpp, "double inspect_int(") + assert "std::vector v" in inspect + assert "std::vector v" in inspect + assert "std::unordered_map(target)" in inspect + assert "(double)target.size()" in inspect + assert "target.clear();" in _function_body(cpp, "double clear_int(") + + # Parameter TypeSpecs win over same-named global array/map/matrix entries. + assert ": 0)" in _function_body(cpp, "double shadow_array_probe(") + assert 'std::string("")' in _function_body(cpp, "bool shadow_map_probe(") + assert ": false)" in _function_body(cpp, "double shadow_matrix_probe(") + + merge = _function_body(cpp, "double merge_int(") + assert "target.insert(" in merge + assert ".begin(), __pf_map_param_arg_" in merge + assert ".end()); }((source));" in merge + + +def test_udt_valued_map_parameter_remains_outside_supported_subset(): + source = '''//@version=6 +strategy("T") +type Pivot + float level +probe(map values) => + values.get("x").level +values = map.new() +observed = probe(values) +''' + with pytest.raises(CompileError, match="map values must be primitive"): + transpile(source) + + +def test_typed_map_parameter_methods_compile(): + compile_env.compile_cpp(transpile(_SOURCE), label="typed_map_parameter_methods") + + +def _find_engine_library() -> Path | None: + explicit = os.environ.get("PINEFORGE_ENGINE_LIB") + if explicit: + path = Path(explicit).expanduser().resolve() + return path if path.is_file() else None + if compile_env._ENGINE_INC is None: + return None + candidates: list[Path] = [] + for pattern in ("build*/lib/libpineforge.a", "build*/lib/libpineforge.dylib"): + candidates.extend(sorted(compile_env._ENGINE_INC.parent.glob(pattern))) + return candidates[0].resolve() if candidates else None + + +def _compile_and_run(cpp_source: str) -> str: + compile_env.skip_if_no_compile_env() + engine_lib = _find_engine_library() + if engine_lib is None: + pytest.skip("built libpineforge not found; set PINEFORGE_ENGINE_LIB") + compiler = compile_env._COMPILER + engine_inc = compile_env._ENGINE_INC + eigen_inc = compile_env._EIGEN_INC + assert compiler is not None and engine_inc is not None and eigen_inc is not None + + with tempfile.TemporaryDirectory(prefix="pineforge-map-param-") as tmp: + cpp_path = Path(tmp) / "probe.cpp" + exe_path = Path(tmp) / "probe" + cpp_path.write_text(cpp_source) + command = [ + compiler, + "-std=c++17", + "-O0", + "-I", + str(engine_inc), + "-I", + str(eigen_inc), + ] + if compile_env._GENERATED_INC is not None: + command += ["-I", str(compile_env._GENERATED_INC)] + command += [str(cpp_path), str(engine_lib), "-pthread", "-o", str(exe_path)] + built = subprocess.run(command, capture_output=True, text=True, timeout=120) + if built.returncode != 0: + raise AssertionError( + "typed-map runtime probe failed to link\n" + + "\n".join((built.stderr or built.stdout).splitlines()[:100]) + ) + ran = subprocess.run( + [str(exe_path)], capture_output=True, text=True, timeout=30 + ) + if ran.returncode != 0: + raise AssertionError( + f"typed-map runtime probe exited {ran.returncode}\n" + f"stdout:\n{ran.stdout}\nstderr:\n{ran.stderr}" + ) + return ran.stdout + + +def test_typed_map_parameter_methods_preserve_runtime_aliases_and_order(): + driver = r""" +#include +int main() { + GeneratedStrategy strategy; + Bar bar{1.0, 1.0, 1.0, 1.0, 1.0, 0}; + strategy.run(&bar, 1); + if (!strategy.last_error().empty()) { + std::cerr << strategy.last_error() << "\n"; + return 2; + } + std::cout << strategy.int_result << " " + << strategy.int_size_after_put << " " + << strategy.inspect_result << " " + << strategy.float_result << " " + << strategy.bool_result << " " + << strategy.string_result << " " + << strategy.missing_float_result << " " + << strategy.missing_int_result << " " + << strategy.missing_bool_result << " " + << strategy.missing_string_result << " " + << strategy.merge_result << " " + << strategy.merge_caller_size << " " + << strategy.shadow_array_result << " " + << strategy.shadow_map_result << " " + << strategy.shadow_matrix_result << " " + << strategy.removed_result << " " + << strategy.int_size_after_remove << " " + << strategy.clear_result << " " + << strategy.clear_caller_size << " " + << strategy.order_code << " " + << strategy.order_size << "\n"; +} +""" + output = _compile_and_run(transpile(_SOURCE) + driver) + assert tuple(float(value) for value in output.split()) == ( + 7.0, + 1.0, + 4.0, + 2.5, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 8.0, + 1.0, + 11.0, + 1.0, + 1.0, + 7.0, + 0.0, + 0.0, + 0.0, + 123.0, + 3.0, + )