test_scalar_module_variable_lowering.py

"""Scalar module-variable planning and lowering tests."""

from __future__ import annotations

from dataclasses import replace
from unittest.mock import Mock

import pytest

from tests.fortran._support.ownership_policy import parse_pyi_text
from prik.parsers.fortran.parser import parse_fortran_project
from prik.pipeline.build import _apply_source_python_exports, _merge_wrapper_modules
from prik.semantics.fortran2ir import fortran_project_to_semantic_modules
from prik.policy.ownership import AssignmentMode, SetterAction
from prik.policy.completion import complete_semantic_policies
from prik.policy.models import ModuleGetterAction
from prik.pipeline.wrapper import WrapperGenerator
from prik.planning import WrapperPlanner
from prik.codegen.c.binding import CBindingGenerator
from prik.codegen.fortran.bridge import FortranBridgeGenerator


SCALAR_MODULE_CONTRACT = """
limit: Final[Int32] = 12
counter: Int32 = 3
target_scale: Annotated[Float64, Aliased]
optional_scale: Allocatable[Float64]
selected_scale: Pointer[Float64]

def summarize() -> Int32: ...
"""


def _plan():
    module = parse_pyi_text(SCALAR_MODULE_CONTRACT, module_name="scalar_state")
    complete_semantic_policies(module)
    return WrapperPlanner().build(module)


def _computed_constant_plan():
    parsed = parse_fortran_project(
        {
            "computed_constants.f90": """
module computed_constants
  integer, parameter :: computed = kind(1.0) * 2
  character*1, parameter :: prefix = 'D'
end module computed_constants
"""
        }
    )
    modules = fortran_project_to_semantic_modules(parsed)
    _apply_source_python_exports(modules)
    module = _merge_wrapper_modules(modules, name="computed_constants_wrapper")
    complete_semantic_policies(module)
    return WrapperPlanner().build(module)


def _parameter_array_plan():
    parsed = parse_fortran_project(
        {
            "parameter_array.f90": """
module parameter_array
  use iso_fortran_env, only: real64
  real(real64), parameter :: dpmpar(3) = [epsilon(1.0_real64), tiny(1.0_real64), huge(1.0_real64)]
end module parameter_array
"""
        }
    )
    modules = fortran_project_to_semantic_modules(parsed)
    _apply_source_python_exports(modules)
    module = _merge_wrapper_modules(modules, name="parameter_array_wrapper")
    complete_semantic_policies(module)
    return WrapperPlanner().build(module)


def _source(artifacts, suffix: str) -> str:
    return next(item.text for item in artifacts.sources if item.path.name.endswith(suffix))


def _replace_variable(plan, python_name: str, edit):
    root = plan.namespaces[0]
    variables = tuple(
        edit(variable) if variable.binding.python_names == (python_name,) else variable for variable in root.variables
    )
    return replace(plan, namespaces=(replace(root, variables=variables), *plan.namespaces[1:]))


def test_module_variable_plan_contains_only_completed_dispatch_facts():
    plan = _plan()
    variables = {variable.binding.python_names[0]: variable for variable in plan.namespaces[0].variables}

    assert variables["limit"].binding.getter_action is ModuleGetterAction.CONSTANT_VALUE
    assert variables["limit"].binding.setter_action is SetterAction.OMIT
    assert variables["limit"].bridge.native_assignment is AssignmentMode.NONE
    assert variables["limit"].binding.constant_value == 12
    assert variables["counter"].binding.getter_action is ModuleGetterAction.DIRECT_VALUE
    assert variables["counter"].binding.setter_action is SetterAction.WRITE_THROUGH
    assert variables["counter"].bridge.native_assignment is AssignmentMode.VALUE_COPY
    assert variables["counter"].binding.initializer == 3
    assert variables["target_scale"].bridge.native_assignment is AssignmentMode.VALUE_COPY
    assert variables["optional_scale"].binding.getter_action is ModuleGetterAction.NULLABLE_SNAPSHOT
    assert variables["optional_scale"].entrypoint.descriptor_kind == "allocatable"
    assert variables["optional_scale"].binding.setter_action is SetterAction.REJECT_REPLACEMENT
    assert variables["optional_scale"].bridge.native_assignment is AssignmentMode.NONE
    assert variables["selected_scale"].entrypoint.descriptor_kind == "pointer"
    assert variables["selected_scale"].bridge.native_assignment is AssignmentMode.NONE


def test_symbolic_source_parameter_reuses_scalar_bridge_getter_for_module_initialization():
    plan = _computed_constant_plan()
    variables = {variable.binding.python_names[0]: variable for variable in plan.namespaces[1].variables}
    computed = variables["computed"]
    assert computed.binding.getter_action is ModuleGetterAction.NATIVE_CONSTANT_VALUE
    assert computed.binding.constant_value is None
    assert computed.entrypoint.getter_role == "computed_constants.computed:getter"
    assert computed.binding.setter_action is SetterAction.OMIT

    artifacts = WrapperGenerator().generate(plan)
    c_source = _source(artifacts, ".c")
    fortran_source = _source(artifacts, ".f90")
    assert "int32_t bind_c_get_computed(void);" in c_source
    assert "int32_t constant_computed_value_0 = bind_c_get_computed();" in c_source
    assert 'PyUnicode_FromString("D")' in c_source
    assert "native_computed => computed" in fortran_source
    assert "function bind_c_get_computed()" in fortran_source
    assert "result = native_computed" in fortran_source
    assert "bind_c_set_computed" not in c_source
    assert "bind_c_set_computed" not in fortran_source


def test_parameter_array_uses_one_immutable_python_owned_import_snapshot():
    plan = _parameter_array_plan()
    variable = next(
        variable
        for namespace in plan.namespaces
        for variable in namespace.variables
        if variable.binding.python_names == ("dpmpar",)
    )
    assert variable.binding.getter_action is ModuleGetterAction.NATIVE_CONSTANT_ARRAY_VALUE
    assert variable.binding.setter_action is SetterAction.OMIT
    assert variable.binding.constant_value is None
    assert variable.array is not None
    assert variable.array.shape == ("3",)

    artifacts = WrapperGenerator().generate(plan)
    c_source = _source(artifacts, ".c")
    fortran_source = _source(artifacts, ".f90")
    assert "void * bind_c_get_dpmpar(int64_t * extent_0);" in c_source
    assert "PyArray_EMPTY(1, constant_dpmpar_value_0_dimensions, NPY_FLOAT64, 1)" in c_source
    assert "memcpy(PyArray_DATA((PyArrayObject *)constant_dpmpar_object_0)" in c_source
    assert "PyArray_CLEARFLAGS((PyArrayObject *)constant_dpmpar_object_0, NPY_ARRAY_WRITEABLE)" in c_source
    assert 'PyModule_AddObject(namespace_parameter_array, "dpmpar", constant_dpmpar_object_0)' in c_source
    assert "real(c_double), allocatable, target, save, dimension(:) :: parameter_snapshot" in fortran_source
    assert "parameter_snapshot = native_dpmpar" in fortran_source
    assert "result = c_loc(parameter_snapshot)" in fortran_source


def test_module_variable_visitors_consume_their_backend_owned_actions():
    plan = _plan()
    counter = next(
        variable for variable in plan.namespaces[0].variables if variable.binding.python_names == ("counter",)
    )
    split_actions = replace(
        counter,
        binding=replace(
            counter.binding,
            getter_action=ModuleGetterAction.CONSTANT_VALUE,
            setter_action=SetterAction.OMIT,
        ),
        bridge=replace(
            counter.bridge,
            native_getter_action=ModuleGetterAction.DIRECT_VALUE,
            native_assignment=AssignmentMode.VALUE_COPY,
        ),
    )

    assert CBindingGenerator().visit(split_actions) == ()
    bridge = FortranBridgeGenerator()
    bridge.visit(plan)
    assert [procedure.name for procedure in bridge.visit(split_actions)] == [
        "bind_c_get_counter",
        "bind_c_set_counter",
    ]


def test_fortran_module_setter_rejects_unsupported_bridge_assignment():
    plan = _plan()
    counter = next(
        variable for variable in plan.namespaces[0].variables if variable.binding.python_names == ("counter",)
    )
    invalid = replace(counter, bridge=replace(counter.bridge, native_assignment=AssignmentMode.ALIAS))

    bridge = FortranBridgeGenerator()
    bridge.visit(plan)
    with pytest.raises(ValueError, match="Unsupported Fortran module setter assignment"):
        bridge.visit(invalid)


@pytest.mark.parametrize(
    ("python_name", "assignment"),
    [
        ("counter", AssignmentMode.NONE),
        ("counter", AssignmentMode.ALIAS),
        ("optional_scale", AssignmentMode.VALUE_COPY),
        ("optional_scale", AssignmentMode.ALIAS),
        ("limit", AssignmentMode.VALUE_COPY),
    ],
)
def test_module_setter_assignment_mismatch_fails_before_backend_preflight_or_lowering(
    python_name,
    assignment,
):
    invalid = _replace_variable(
        _plan(),
        python_name,
        lambda variable: replace(
            variable,
            bridge=replace(variable.bridge, native_assignment=assignment),
        ),
    )
    c_generator = Mock(spec=CBindingGenerator)
    fortran_generator = Mock(spec=FortranBridgeGenerator)
    c_printer = Mock()
    fortran_printer = Mock()
    generator = WrapperGenerator(
        c_generator=c_generator,
        fortran_generator=fortran_generator,
        c_printer=c_printer,
        fortran_printer=fortran_printer,
    )

    with pytest.raises(ValueError, match="invalid-module-native-assignment") as error:
        generator.generate(invalid)

    assert "Unsupported Fortran module setter assignment" not in str(error.value)
    c_generator.require_supported.assert_not_called()
    fortran_generator.require_supported.assert_not_called()
    c_generator.visit.assert_not_called()
    fortran_generator.visit.assert_not_called()
    c_generator.requires_native_support.assert_not_called()
    c_printer.doprint.assert_not_called()
    fortran_printer.doprint.assert_not_called()


def test_module_variable_generators_dispatch_get_set_and_rejection_from_plan():
    artifacts = WrapperGenerator().generate(_plan())
    c_source = _source(artifacts, ".c")
    fortran_source = _source(artifacts, ".f90")

    assert "scalar_state_root_module_property_setup_getattro" in c_source
    assert "scalar_state_root_module_property_setup_setattro" in c_source
    assert 'PyModule_AddObject(mod, "limit"' in c_source
    assert "bind_c_set_counter(3);" in c_source
    assert "return bind_c_get_counter();" not in c_source
    assert "bind_c_get_counter()" in c_source
    assert "bind_c_set_counter(value)" in c_source
    assert "module variable optional_scale is read-only" in c_source
    assert "module variable selected_scale is read-only" in c_source
    assert 'getenv("PRIK_WRAPPER_FAIL_ALLOC")' in c_source
    assert "result = native_counter" in fortran_source
    assert "native_counter = value" in fortran_source
    assert "allocated(native_optional_scale)" in fortran_source
    assert "associated(native_selected_scale)" in fortran_source
    assert "optional_scale = value" not in fortran_source
    assert "selected_scale = value" not in fortran_source


def test_generated_support_procedure_symbol_is_shared_by_both_boundary_lowerers():
    plan = _plan()
    procedure = next(
        item
        for item in plan.entrypoint.support_procedures
        if item.owner_path == "scalar_state.counter" and item.role == "module:set"
    )
    renamed = replace(procedure, symbol_name="planned_counter_assignment")
    edited = replace(
        plan,
        entrypoint=replace(
            plan.entrypoint,
            support_procedures=tuple(
                renamed if item is procedure else item for item in plan.entrypoint.support_procedures
            ),
        ),
    )

    artifacts = WrapperGenerator().generate(edited)
    c_source = _source(artifacts, ".c")
    fortran_source = _source(artifacts, ".f90")

    assert "void planned_counter_assignment(int32_t value);" in c_source
    assert "planned_counter_assignment(value);" in c_source
    assert "subroutine planned_counter_assignment(value)" in fortran_source
    assert 'bind(c, name="planned_counter_assignment")' in fortran_source


def test_missing_generated_support_procedure_fails_before_lowering():
    plan = _plan()
    edited = replace(
        plan,
        entrypoint=replace(
            plan.entrypoint,
            support_procedures=tuple(
                item
                for item in plan.entrypoint.support_procedures
                if not (item.owner_path == "scalar_state.counter" and item.role == "module:set")
            ),
        ),
    )

    with pytest.raises(ValueError, match="incomplete-auxiliary-entrypoint-inventory"):
        WrapperGenerator().generate(edited)


def test_bridge_local_module_target_edit_does_not_change_the_c_boundary():
    plan = _plan()
    baseline = _source(WrapperGenerator().generate(plan), ".c")
    counter = next(
        variable for variable in plan.namespaces[0].variables if variable.binding.python_names == ("counter",)
    )
    edited_counter = replace(
        counter,
        bridge=replace(counter.bridge, native_name="counter_alternate"),
    )
    root = replace(
        plan.namespaces[0],
        variables=tuple(
            edited_counter if variable is counter else variable for variable in plan.namespaces[0].variables
        ),
    )
    edited = replace(plan, namespaces=(root, *plan.namespaces[1:]))

    artifacts = WrapperGenerator().generate(edited)

    assert _source(artifacts, ".c") == baseline
    assert "native_counter => counter_alternate" in _source(artifacts, ".f90")


def test_generator_rejects_python_module_setter_without_bridge_handoff():
    plan = _plan()
    counter = next(
        variable for variable in plan.namespaces[0].variables if variable.binding.python_names == ("counter",)
    )
    invalid_counter = replace(counter, entrypoint=replace(counter.entrypoint, setter_role=None))
    invalid = replace(
        plan,
        namespaces=(
            replace(
                plan.namespaces[0],
                variables=tuple(
                    invalid_counter if variable is counter else variable for variable in plan.namespaces[0].variables
                ),
            ),
        ),
    )

    with pytest.raises(ValueError, match="missing-module-setter-role"):
        WrapperGenerator().generate(invalid)