Fix local consts data ref

dace_fortran/inliner/ast_desugaring/analysis.py

@@ -1302,8 +1302,13 @@ def _const_eval_int(expr: Base, alias_map: types.SPEC_TABLE) -> int | None:
     return None
 
 
-#: Key of a tracked local constant: a variable's spec, or a (root spec, component spec) pair for ``root % comp``.
-ConstKey = Union[types.SPEC, tuple[types.SPEC, types.SPEC]]
+#: Key of a tracked local constant: a variable's spec, a (root spec, component
+#: spec) pair for ``root % comp``, or a (root spec, component_path) pair where
+#: component_path is a tuple of component names and integer subscript tuples for
+#: indexed components (e.g. ``root % baz(3) % scalar`` or
+#: ``root % baz(2) % arr(4)``).
+IndexedComponent = tuple[Union[str, tuple[int, ...]], ...]
+ConstKey = Union[types.SPEC, tuple[types.SPEC, types.SPEC], tuple[types.SPEC, tuple[IndexedComponent, ...]]]
 
 
 class RootComponent(NamedTuple):
@@ -1325,6 +1330,69 @@ def track_local_consts(
     plus: dict[ConstKey, types.LITERAL_TYPES] = copy(plus_init) if plus_init else {}
     minus: set[ConstKey] = copy(minus_init) if minus_init else set()
 
+    def _indexed_key(ref: Base) -> ConstKey | None:
+        """Build a constant key for a scalar array element or component reference.
+
+        For ``arr(i)`` returns ``(arr_spec, ((i,),))``.  For ``root % comp(i)``
+        returns ``(root_spec, (('comp', (i,)),))``.  For
+        ``root % baz(3) % scalar`` returns ``(root_spec, (('baz', (3,)), 'scalar'))``.
+        Returns ``None`` for non-scalar references (slices), references to
+        non-local roots, or non-constant subscripts.
+        """
+        scope_spec = search_scope_spec(ref)
+        if not scope_spec:
+            return None
+
+        components: list[IndexedComponent] = []
+
+        def _append_part(pr: f03.Part_Ref) -> bool:
+            name, subs = pr.children
+            if not isinstance(name, f03.Name):
+                return False
+            if subs is None:
+                return False
+            subscripts: list[int] = []
+            for s in subs.children:
+                if isinstance(s, f03.Subscript_Triplet):
+                    return False
+                sv = const_eval_basic_type(s, alias_map)
+                if sv is None or not types.is_numpy_int(sv):
+                    return False
+                subscripts.append(int(sv))
+            components.append((name.string.lower(), tuple(subscripts)))
+            return True
+
+        if isinstance(ref, f03.Part_Ref):
+            if not _append_part(ref):
+                return None
+            name, _ = ref.children
+            loc = search_real_local_alias_spec(name, alias_map)
+            if not loc:
+                return None
+            base_spec = ident_spec(alias_map[loc])
+        elif isinstance(ref, f03.Data_Ref):
+            root, _root_type, rest = dataref_root(ref, scope_spec, alias_map)
+            if not isinstance(root, f03.Name):
+                return None
+            loc = search_real_local_alias_spec(root, alias_map)
+            if not loc:
+                return None
+            base_spec = ident_spec(alias_map[loc])
+            for comp in rest:
+                if isinstance(comp, f03.Part_Ref):
+                    if not _append_part(comp):
+                        return None
+                elif isinstance(comp, f03.Name):
+                    components.append(comp.string.lower())
+                else:
+                    return None
+        else:
+            return None
+
+        if not components:
+            return None
+        return (base_spec, tuple(components))
+
     def _root_comp(dref: Union[f03.Data_Ref, f03.Data_Pointer_Object]) -> RootComponent | None:
         scope_spec = search_scope_spec(dref)
         assert scope_spec
@@ -1410,6 +1478,14 @@ def _inject_knowns(x: Base, value: bool = True, pointer: bool = True) -> None:
             if (pointer and xtyp and xtyp.pointer) or value:
                 _subst_and_renorm(x, plus[spec])
         elif isinstance(x, f03.Data_Ref):
+            idx_key = _indexed_key(x)
+            if idx_key is not None and idx_key in plus:
+                assert idx_key not in minus
+                scope_spec = find_scope_spec(x)
+                xtyp = find_type_dataref(x, scope_spec, alias_map)
+                if (pointer and xtyp.pointer) or value:
+                    _subst_and_renorm(x, plus[idx_key])
+                return
             comp_key = _root_comp(x)
             if comp_key is None or comp_key not in plus:
                 for pr in x.children[1:]:
@@ -1426,6 +1502,12 @@ def _inject_knowns(x: Base, value: bool = True, pointer: bool = True) -> None:
             if (pointer and xtyp.pointer) or value:
                 _subst_and_renorm(x, plus[comp_key])
         elif isinstance(x, f03.Part_Ref):
+            idx_key = _indexed_key(x)
+            if idx_key is not None and idx_key in plus:
+                assert idx_key not in minus
+                if value:
+                    _subst_and_renorm(x, plus[idx_key])
+                return
             par, subsc = x.children
             _inject_knowns(par, value=False, pointer=True)
             assert isinstance(subsc, f03.Section_Subscript_List)
@@ -1503,15 +1585,30 @@ def _inject_knowns(x: Base, value: bool = True, pointer: bool = True) -> None:
                 lspec = ident_spec(alias_map[loc])
                 if isinstance(alias_map[lspec], f03.Entity_Decl):
                     ltyp = find_type_of_entity(alias_map[lspec], alias_map)
+        elif isinstance(lv, f03.Part_Ref):
+            lspec = _indexed_key(lv)
+            if lspec is not None:
+                name, _ = lv.children
+                name_spec = search_real_local_alias_spec(name, alias_map)
+                assert name_spec
+                ltyp = find_type_of_entity(alias_map[name_spec], alias_map)
         elif isinstance(lv, f03.Data_Ref):
-            lspec = _root_comp(lv)
-            scope_spec = find_scope_spec(lv)
-            ltyp = find_type_dataref(lv, scope_spec, alias_map)
+            idx_key = _indexed_key(lv)
+            if idx_key is not None:
+                lspec = idx_key
+                scope_spec = find_scope_spec(lv)
+                ltyp = find_type_dataref(lv, scope_spec, alias_map)
+            else:
+                lspec = _root_comp(lv)
+                scope_spec = find_scope_spec(lv)
+                ltyp = find_type_dataref(lv, scope_spec, alias_map)
         if lspec and ltyp:
             rval = const_eval_basic_type(rv, alias_map)
             if rval is None:
                 _integrate_subresults({}, {lspec})
-            elif not ltyp.shape:
+            elif not ltyp.shape or (isinstance(lspec, tuple) and len(lspec) >= 2):
+                # Scalar variable, or a scalar array element / component element
+                # represented by an indexed key.
                 plus[lspec] = types.numpy_type_to_literal(rval)
                 if lspec in minus:
                     minus.remove(lspec)
@@ -1535,11 +1632,26 @@ def _inject_knowns(x: Base, value: bool = True, pointer: bool = True) -> None:
             lspec = _root_comp(lv)
             scope_spec = find_scope_spec(lv)
             ltyp = find_type_dataref(lv, scope_spec, alias_map)
+        remove: set[ConstKey] = set()
         if lspec and ltyp and ltyp.pointer:
-            plus[lspec] = rv
-            if lspec in minus:
-                minus.remove(lspec)
-        tp, tm = track_local_consts(rv, alias_map)
+            # Only treat the pointer as a known alias when the RHS is a plain
+            # target name/data-ref or a constant-indexed element (e.g.
+            # ``ptr => arr(2)``).  If the RHS indexes with a variable (e.g.
+            # ``ptr => arr(i)``) the target is not a compile-time constant and
+            # substituting it would be wrong.
+            rv_scope_spec = search_scope_spec(rv)
+            if rv_scope_spec is None:
+                remove = {lspec}
+            else:
+                root_name = dataref_root(rv, rv_scope_spec, alias_map)[0]
+                bad = list(walk(rv.children, (f03.Name, f03.Data_Ref, f03.Part_Ref)))
+                if not bad or all(n is root_name for n in bad):
+                    plus[lspec] = rv
+                    if lspec in minus:
+                        minus.remove(lspec)
+                else:
+                    remove = {lspec}
+        tp, tm = track_local_consts(rv, alias_map, {}, remove)
         _integrate_subresults(tp, tm)
     elif isinstance(node, f03.If_Stmt):
         cond, body = node.children

tests/inliner/optimizations_test.py

@@ -1,4 +1,6 @@
 # Copyright 2025-2026 ETH Zurich and the dace-fortran authors. All rights reserved.
+import pytest
+
 from dace_fortran.inliner.ast_desugaring import optimizations, types
 from tests.inliner.fortran_test_helper import SourceCodeBuilder, parse_and_improve
 
@@ -905,6 +907,7 @@ def test_exploit_locally_constant_pointers():
   end type cfg
   type(cfg) :: c
   real, target :: data = 0.
+  real, target :: arr(3)
   real, pointer :: ptr => null()
   integer, target :: i
   integer, pointer :: iptr => null()
@@ -924,6 +927,13 @@ def test_exploit_locally_constant_pointers():
       iarr(iptr) = iarr(iptr-1) + 1
     end if
   end do
+
+  ptr => arr(2)
+  data = ptr
+  ptr => arr(i)
+  data = ptr  ! Must not do constant replacement
+  ptr => arr(iarr(1))
+  data = ptr  ! Must not do constant replacement
 end subroutine main
 """)
         .check_with_gfortran()
@@ -941,6 +951,7 @@ def test_exploit_locally_constant_pointers():
   END TYPE cfg
   TYPE(cfg) :: c
   REAL, TARGET :: data = 0.
+  REAL, TARGET :: arr(3)
   REAL, POINTER :: ptr => NULL()
   INTEGER, TARGET :: i
   INTEGER, POINTER :: iptr => NULL()
@@ -959,6 +970,170 @@ def test_exploit_locally_constant_pointers():
       iarr(i) = iarr(i - 1) + 1
     END IF
   END DO
+  ptr => arr(2)
+  data = arr(2)
+  ptr => arr(i)
+  data = ptr
+  ptr => arr(iarr(1))
+  data = ptr
+END SUBROUTINE main
+""".strip()
+    assert got == want
+    SourceCodeBuilder().add_file(got).check_with_gfortran()
+
+
+def test_exploit_locally_constant_part_refs():
+    """Indexed ``Part_Ref`` array elements are tracked and propagated."""
+    sources, _ = (
+        SourceCodeBuilder()
+        .add_file("""
+subroutine main
+  implicit none
+  integer :: iarr(4) = 0
+  real :: rarr(4,4) = 0.0
+
+  iarr(2) = 7
+  iarr(3) = iarr(2)
+
+  rarr(1,1) = 3.2
+  rarr(2:3,1) = rarr(1,1)
+end subroutine main
+""")
+        .check_with_gfortran()
+        .get()
+    )
+    ast = parse_and_improve(sources)
+    ast = optimizations.exploit_locally_constant_variables(ast)
+
+    got = ast.tofortran()
+    want = """
+SUBROUTINE main
+  IMPLICIT NONE
+  INTEGER :: iarr(4) = 0
+  REAL :: rarr(4, 4) = 0.0
+  iarr(2) = 7
+  iarr(3) = 7
+  rarr(1, 1) = 3.2
+  rarr(2 : 3, 1) = 3.2
+END SUBROUTINE main
+""".strip()
+    assert got == want
+    SourceCodeBuilder().add_file(got).check_with_gfortran()
+
+
+def test_exploit_locally_constant_data_refs():
+    """Indexed ``Data_Ref`` component accesses are tracked and propagated."""
+    sources, _ = (
+        SourceCodeBuilder()
+        .add_file("""
+subroutine main
+  implicit none
+  type bar
+    integer :: iarr(4) = 0
+    real :: rarr(4, 4) = 0.0
+    real :: scalar
+  end type bar
+  type foo
+    integer :: iarr(6) = 0
+    type(bar) :: baz(3)
+  end type foo
+  type(foo) :: quux
+
+  quux % iarr(2) = 7
+  quux % iarr(1) = quux % iarr(2)
+
+  quux % baz(3) % scalar = 6.8
+  quux % baz(2) % iarr(4) = 9
+  quux % baz(2) % rarr(3,2) = 3.2
+  quux % baz(1) % rarr(1, 1) = quux % baz(3) % scalar
+  quux % iarr(3) = quux % baz(2) % iarr(4)
+  quux % baz(1) % rarr(2, 1) = quux % baz(2) % rarr(3,2)
+end subroutine main
+""")
+        .check_with_gfortran()
+        .get()
+    )
+    ast = parse_and_improve(sources)
+    ast = optimizations.exploit_locally_constant_variables(ast)
+
+    got = ast.tofortran()
+    want = """
+SUBROUTINE main
+  IMPLICIT NONE
+  TYPE :: bar
+    INTEGER :: iarr(4) = 0
+    REAL :: rarr(4, 4) = 0.0
+    REAL :: scalar
+  END TYPE bar
+  TYPE :: foo
+    INTEGER :: iarr(6) = 0
+    TYPE(bar) :: baz(3)
+  END TYPE foo
+  TYPE(foo) :: quux
+  quux % iarr(2) = 7
+  quux % iarr(1) = 7
+  quux % baz(3) % scalar = 6.8
+  quux % baz(2) % iarr(4) = 9
+  quux % baz(2) % rarr(3, 2) = 3.2
+  quux % baz(1) % rarr(1, 1) = 6.8
+  quux % iarr(3) = 9
+  quux % baz(1) % rarr(2, 1) = 3.2
+END SUBROUTINE main
+""".strip()
+    assert got == want
+    SourceCodeBuilder().add_file(got).check_with_gfortran()
+
+
+@pytest.mark.skip(reason="Data-ref pointer target substitution is not yet supported in the dace-fortran port.")
+def test_exploit_locally_constant_data_ref_pointers():
+    """Pointer target substitution works for indexed data-ref components."""
+    sources, _ = (
+        SourceCodeBuilder()
+        .add_file("""
+subroutine main
+  implicit none
+  type foo
+    integer, pointer :: ptr(:, :)
+    integer, pointer :: qtr
+  end type foo
+  type(foo) :: bar(3)
+  integer, target :: x(2,3)
+  integer :: idx
+
+  bar(1) % ptr => x(:,:)
+  bar(1) % ptr(idx,1) = 5
+  bar(1) % ptr(2,2) = 2
+
+  bar(2) % qtr => x(idx, 2)
+  bar(2) % ptr(1,1) = bar(2) % qtr
+  bar(3) % ptr(2,2) = bar(1) % ptr(idx,1)
+  bar(3) % ptr(2,3) = bar(1) % ptr(2,2)
+end subroutine main
+""")
+        .check_with_gfortran()
+        .get()
+    )
+    ast = parse_and_improve(sources)
+    ast = optimizations.exploit_locally_constant_variables(ast)
+
+    got = ast.tofortran()
+    want = """
+SUBROUTINE main
+  IMPLICIT NONE
+  TYPE :: foo
+    INTEGER, POINTER :: ptr(:, :)
+    INTEGER, POINTER :: qtr
+  END TYPE foo
+  TYPE(foo) :: bar(3)
+  INTEGER, TARGET :: x(2, 3)
+  INTEGER :: idx
+  bar(1) % ptr => x(:, :)
+  x(idx, 1) = 5
+  x(2, 2) = 2
+  bar(2) % qtr => x(idx, 2)
+  bar(2) % ptr(1, 1) = x(idx, 2)
+  bar(3) % ptr(2, 2) = x(idx, 1)
+  bar(3) % ptr(2, 3) = x(2, 2)
 END SUBROUTINE main
 """.strip()
     assert got == want