From 27460066e0299821c280e56594f3661f3b867826 Mon Sep 17 00:00:00 2001
From: Michael Staneker <michael.staneker@ecmwf.int>
Date: Thu, 14 Dec 2023 14:45:58 +0100
Subject: [PATCH 1/5] introduce 'flatten_arrays()' (to overcome pointer hack)

---
 loki/backend/cgen.py                       | 15 +---
 loki/transform/fortran_c_transform.py      |  4 +-
 loki/transform/transform_array_indexing.py | 46 +++++++++-
 tests/test_transform_array_indexing.py     | 97 +++++++++++++++++++++-
 4 files changed, 144 insertions(+), 18 deletions(-)

diff --git a/loki/backend/cgen.py b/loki/backend/cgen.py
index c7a476325..9e4880a4c 100644
--- a/loki/backend/cgen.py
+++ b/loki/backend/cgen.py
@@ -158,9 +158,8 @@ def visit_Subroutine(self, o, **kwargs):
         # Generate header with argument signature
         aptr = []
         for a in o.arguments:
-            # TODO: Oh dear, the pointer derivation is beyond hacky; clean up!
             if isinstance(a, Array) > 0:
-                aptr += ['* restrict v_']
+                aptr += ['* restrict ']
             elif isinstance(a.type.dtype, DerivedType):
                 aptr += ['*']
             elif a.type.pointer:
@@ -176,18 +175,6 @@ def visit_Subroutine(self, o, **kwargs):
         # ...and generate the spec without imports and argument declarations
         body = [self.visit(o.spec, skip_imports=True, skip_argument_declarations=True, **kwargs)]
 
-        # Generate the array casts for pointer arguments
-        if any(isinstance(a, Array) for a in o.arguments):
-            body += [self.format_line('/* Array casts for pointer arguments */')]
-            for a in o.arguments:
-                if isinstance(a, Array):
-                    dtype = self.visit(a.type, **kwargs)
-                    # str(d).lower() is a bad hack to ensure caps-alignment
-                    outer_dims = ''.join(f'[{self.visit(d, **kwargs).lower()}]'
-                                         for d in a.dimensions[1:])
-                    body += [self.format_line(dtype, ' (*', a.name.lower(), ')', outer_dims, ' = (',
-                                              dtype, ' (*)', outer_dims, ') v_', a.name.lower(), ';')]
-
         # Fill the body
         body += [self.visit(o.body, **kwargs)]
         body += [self.format_line('return 0;')]
diff --git a/loki/transform/fortran_c_transform.py b/loki/transform/fortran_c_transform.py
index 37eb8fe56..d261183cc 100644
--- a/loki/transform/fortran_c_transform.py
+++ b/loki/transform/fortran_c_transform.py
@@ -11,7 +11,8 @@
 from loki.transform.transformation import Transformation
 from loki.transform.transform_array_indexing import (
     shift_to_zero_indexing, invert_array_indices,
-    resolve_vector_notation, normalize_range_indexing
+    resolve_vector_notation, normalize_range_indexing,
+    flatten_arrays
 )
 from loki.transform.transform_associates import resolve_associates
 from loki.transform.transform_utilities import (
@@ -374,6 +375,7 @@ def generate_c_kernel(self, routine, **kwargs):
         # TODO: Resolve reductions (eg. SUM(myvar(:)))
         invert_array_indices(kernel)
         shift_to_zero_indexing(kernel)
+        flatten_arrays(kernel, order="C", start_index=0)
 
         # Inline all known parameters, since they can be used in declarations,
         # and thus need to be known before we can fetch them via getters.
diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py
index 946709ff1..34d571c15 100644
--- a/loki/transform/transform_array_indexing.py
+++ b/loki/transform/transform_array_indexing.py
@@ -28,7 +28,8 @@
     'shift_to_zero_indexing', 'invert_array_indices',
     'resolve_vector_notation', 'normalize_range_indexing',
     'promote_variables', 'promote_nonmatching_variables',
-    'promotion_dimensions_from_loop_nest', 'demote_variables'
+    'promotion_dimensions_from_loop_nest', 'demote_variables',
+    'flatten_arrays'
 ]
 
 
@@ -496,3 +497,46 @@ def demote_variables(routine, variable_names, dimensions):
     routine.spec = Transformer(decl_map).visit(routine.spec)
 
     info(f'[Loki::Transform] Demoted variables in {routine.name}: {", ".join(variable_names)}')
+
+
+def flatten_arrays(routine, order='F', start_index=1):
+    """
+    Flatten arrays, converting multi-dimensional arrays to
+    one-dimensional arrays.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The subroutine in which the variables should be promoted.
+    order : str
+        Assume Fortran (F) vs. C memory/array order.
+    start_index : int
+        Assume array indexing starts with `start_index`.
+    """
+    def new_dims(dim, shape):
+        if len(dim) > 1:
+            assert not isinstance(shape[-1], sym.RangeIndex)
+            _dim = [sym.Sum((dim[-2], sym.Product((shape[-2], dim[-1] - start_index))))]
+            new_dim = dim[:-2]
+            new_dim.extend(_dim)
+            return new_dims(new_dim, shape[:-1])
+        return as_tuple(dim)
+
+    assert order in ['F', 'C']
+    if order == 'C':
+        array_map = {
+            var: var.clone(dimensions=new_dims(list(var.dimensions)[::-1], list(var.shape)[::-1]))
+            for var in FindVariables().visit(routine.body)
+                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+        }
+    elif order == 'F':
+        array_map = {
+            var: var.clone(dimensions=new_dims(list(var.dimensions), list(var.shape)))
+            for var in FindVariables().visit(routine.body)
+                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+        }
+    routine.body = SubstituteExpressions(array_map).visit(routine.body)
+
+    routine.variables = [v.clone(dimensions=as_tuple(sym.Product(v.shape)),
+                                 type=v.type.clone(shape=as_tuple(sym.Product(v.shape))))
+                         if isinstance(v, sym.Array) else v for v in routine.variables]
diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py
index ab48695e5..0986470f4 100644
--- a/tests/test_transform_array_indexing.py
+++ b/tests/test_transform_array_indexing.py
@@ -10,9 +10,12 @@
 import numpy as np
 
 from conftest import jit_compile, clean_test, available_frontends
-from loki import Subroutine
+from loki import Subroutine, FindNodes, Assignment
 from loki.expression import symbols as sym
-from loki.transform import promote_variables, demote_variables, normalize_range_indexing
+from loki.transform import (
+        promote_variables, demote_variables, normalize_range_indexing,
+        invert_array_indices, flatten_arrays
+        )
 
 
 @pytest.fixture(scope='module', name='here')
@@ -315,3 +318,93 @@ def test_transform_demote_dimension_arguments(here, frontend):
     assert np.all(vec1 == 3) and np.sum(vec1) == 3
     assert np.all(vec2 == 2) and np.sum(vec2) == 6
     assert np.all(matrix == 16) and np.sum(matrix) == 32
+
+
+@pytest.mark.parametrize('frontend', available_frontends())
+def test_transform_flatten_arrays(here, frontend):
+    """
+    Test flattening or arrays, meaning converting multi-dimensional
+    arrays to one-dimensional arrays including corresponding
+    index arithmetic.
+    """
+    fcode = """
+    subroutine transform_flatten_arrays(x1, x2, x3, x4, l1, l2, l3, l4)
+        implicit none
+        integer :: i1, i2, i3, i4
+        integer, intent(in) :: l1, l2, l3, l4
+        integer, intent(inout) :: x1(l1), x2(l2, l1), x3(l3, l2, l1), x4(l4, l3, l2, l1)
+        
+        do i1=1,l1
+            x1(i1) = 2 * l1
+            do i2=1,l2
+                x2(i2, i1) = 2 * l2
+                do i3=1,l3
+                    x3(i3, i2, i1) = 2 * l3
+                    do i4=1,l4
+                        x4(i4, i3, i2, i1) = 2 * l4
+                    end do
+                end do
+            end do
+        end do
+        
+        
+    end subroutine transform_flatten_arrays
+    """
+
+    def init_arguments(l1, l2, l3, l4, flattened=False):
+        x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
+        x2 = np.zeros(shape=(l2*l1) if flattened else (l2,l1,), order='F', dtype=np.int32)
+        x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l3,l2,l1,), order='F', dtype=np.int32)
+        x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l4,l3,l2,l1,), order='F', dtype=np.int32)
+        return x1, x2, x3, x4
+
+    def validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4):
+        assert np.all(x1 == 2 * l1)
+        assert np.all(x2 == 2 * l2)
+        assert np.all(x3 == 2 * l3)
+        assert np.all(x4 == 2 * l4)
+
+    def validate_routine(routine):
+        assignments = FindNodes(Assignment).visit(routine.body)
+        lhs = [assignment.lhs for assignment in assignments]
+        dims = [_lhs.dimensions for _lhs in lhs]
+        assert all(len(dim) == 1 for dim in dims)
+        assert dims[1][0].children[0] == 'i2'
+        assert dims[2][0].children[0] == 'i3'
+        assert dims[3][0].children[0] == 'i4'
+
+    l1 = 2
+    l2 = 4
+    l3 = 6
+    l4 = 8
+    # Test the original implementation
+    routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_range_indexing(routine) # Fix OMNI nonsense
+    filepath = here/(f'{routine.name}_{frontend}.f90')
+    function = jit_compile(routine, filepath=filepath, objname=routine.name)
+    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4)
+    function(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+
+    f_routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_range_indexing(f_routine) # Fix OMNI nonsense
+    flatten_arrays(routine=f_routine, order='F', start_index=1)
+    filepath = here/(f'{f_routine.name}_flattened_F_{frontend}.f90')
+    function = jit_compile(f_routine, filepath=filepath, objname=routine.name)
+    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4, flattened=True)
+    function(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_routine(f_routine)
+
+    c_routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_range_indexing(c_routine) # Fix OMNI nonsense
+    invert_array_indices(c_routine)
+    flatten_arrays(routine=c_routine, order='C', start_index=1)
+    filepath = here/(f'{c_routine.name}_flattened_C_{frontend}.f90')
+    function = jit_compile(c_routine, filepath=filepath, objname=routine.name)
+    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4, flattened=True)
+    function(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_routine(c_routine)
+
+    assert f_routine.body == c_routine.body

From 9edf6a1414169e877c82a97adde282621177aaa2 Mon Sep 17 00:00:00 2001
From: Michael Staneker <michael.staneker@ecmwf.int>
Date: Mon, 18 Dec 2023 18:08:34 +0000
Subject: [PATCH 2/5] introduce 'normalize_array_access' for prior execution to
 'flatten_arrays' to allow for arrays start counting not with '1'

---
 loki/transform/fortran_c_transform.py      |   9 +-
 loki/transform/transform_array_indexing.py |  64 ++++--
 tests/test_transform_array_indexing.py     | 220 ++++++++++++++++-----
 tests/test_transpile.py                    |   3 +-
 4 files changed, 231 insertions(+), 65 deletions(-)

diff --git a/loki/transform/fortran_c_transform.py b/loki/transform/fortran_c_transform.py
index d261183cc..ad563c243 100644
--- a/loki/transform/fortran_c_transform.py
+++ b/loki/transform/fortran_c_transform.py
@@ -12,7 +12,7 @@
 from loki.transform.transform_array_indexing import (
     shift_to_zero_indexing, invert_array_indices,
     resolve_vector_notation, normalize_range_indexing,
-    flatten_arrays
+    normalize_array_access, flatten_arrays
 )
 from loki.transform.transform_associates import resolve_associates
 from loki.transform.transform_utilities import (
@@ -37,7 +37,6 @@
 from loki.tools import as_tuple, flatten
 from loki.types import BasicType, DerivedType, SymbolAttributes
 
-
 __all__ = ['FortranCTransformation']
 
 
@@ -51,8 +50,9 @@ class FortranCTransformation(Transformation):
     # Set of standard module names that have no C equivalent
     __fortran_intrinsic_modules = ['ISO_FORTRAN_ENV', 'ISO_C_BINDING']
 
-    def __init__(self, header_modules=None, inline_elementals=True):
+    def __init__(self, header_modules=None, inline_elementals=True, order='F'):
         self.inline_elementals = inline_elementals
+        self.order = order
 
         # Maps from original type name to ISO-C and C-struct types
         self.c_structs = OrderedDict()
@@ -369,13 +369,14 @@ def generate_c_kernel(self, routine, **kwargs):
 
         # Clean up Fortran vector notation
         resolve_vector_notation(kernel)
+        normalize_array_access(kernel)
         normalize_range_indexing(kernel)
 
         # Convert array indexing to C conventions
         # TODO: Resolve reductions (eg. SUM(myvar(:)))
         invert_array_indices(kernel)
         shift_to_zero_indexing(kernel)
-        flatten_arrays(kernel, order="C", start_index=0)
+        flatten_arrays(kernel, order=self.order, start_index=0)
 
         # Inline all known parameters, since they can be used in declarations,
         # and thus need to be known before we can fetch them via getters.
diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py
index 34d571c15..33587f3b3 100644
--- a/loki/transform/transform_array_indexing.py
+++ b/loki/transform/transform_array_indexing.py
@@ -29,7 +29,7 @@
     'resolve_vector_notation', 'normalize_range_indexing',
     'promote_variables', 'promote_nonmatching_variables',
     'promotion_dimensions_from_loop_nest', 'demote_variables',
-    'flatten_arrays'
+    'flatten_arrays', 'normalize_array_access'
 ]
 
 
@@ -498,6 +498,43 @@ def demote_variables(routine, variable_names, dimensions):
 
     info(f'[Loki::Transform] Demoted variables in {routine.name}: {", ".join(variable_names)}')
 
+def normalize_array_access(routine):
+    """
+    Shift all arrays to start counting at "1"
+    """
+    def is_range_index(dim):
+        return isinstance(dim, sym.RangeIndex) and not dim.lower == 1
+
+    vmap = {}
+    for v in FindVariables(unique=False).visit(routine.body):
+        if isinstance(v, sym.Array):
+            new_dims = []
+            for i, d in enumerate(v.shape):
+                if isinstance(d, sym.RangeIndex):
+                    if isinstance(v.dimensions[i], sym.RangeIndex):
+                        start = simplify(v.dimensions[i].start - d.start + 1) if d.start is not None else None
+                        stop = simplify(v.dimensions[i].stop - d.start + 1) if d.stop is not None else None
+                        new_dims += [sym.RangeIndex((start, stop, d.step))]
+                    else:
+                        start = simplify(v.dimensions[i] - d.start + 1) if d.start is not None else None
+                        new_dims += [start]
+                else:
+                    new_dims += [v.dimensions[i]]
+            vmap[v] = v.clone(dimensions=as_tuple(new_dims))
+    routine.body = SubstituteExpressions(vmap).visit(routine.body)
+
+    vmap = {}
+    for v in routine.variables:
+        if isinstance(v, sym.Array):
+            new_dims = [sym.RangeIndex((1, simplify(d.upper - d.lower + 1)))
+                if is_range_index(d) else d for d in v.dimensions]
+            new_shape = [sym.RangeIndex((1, simplify(d.upper - d.lower + 1)))
+                if is_range_index(d) else d for d in v.shape]
+            new_type = v.type.clone(shape=as_tuple(new_shape))
+            vmap[v] = v.clone(dimensions=as_tuple(new_dims), type=new_type)
+    routine.variables = [vmap.get(v, v) for v in routine.variables]
+    normalize_range_indexing(routine)
+
 
 def flatten_arrays(routine, order='F', start_index=1):
     """
@@ -515,26 +552,29 @@ def flatten_arrays(routine, order='F', start_index=1):
     """
     def new_dims(dim, shape):
         if len(dim) > 1:
-            assert not isinstance(shape[-1], sym.RangeIndex)
-            _dim = [sym.Sum((dim[-2], sym.Product((shape[-2], dim[-1] - start_index))))]
+            if isinstance(shape[-2], sym.RangeIndex):
+                raise TypeError(f'Resolve shapes being of type RangeIndex, e.g., "{shape[-2]}" before flattening!')
+            _dim = (sym.Sum((dim[-2], sym.Product((shape[-2], dim[-1] - start_index)))),)
             new_dim = dim[:-2]
-            new_dim.extend(_dim)
+            new_dim += _dim
             return new_dims(new_dim, shape[:-1])
-        return as_tuple(dim)
+        return dim
 
-    assert order in ['F', 'C']
-    if order == 'C':
+    if order == 'F':
         array_map = {
-            var: var.clone(dimensions=new_dims(list(var.dimensions)[::-1], list(var.shape)[::-1]))
+            var: var.clone(dimensions=new_dims(var.dimensions[::-1], var.shape[::-1]))
             for var in FindVariables().visit(routine.body)
-                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+            if isinstance(var, sym.Array) and var.shape and len(var.shape)
         }
-    elif order == 'F':
+    elif order == 'C':
         array_map = {
-            var: var.clone(dimensions=new_dims(list(var.dimensions), list(var.shape)))
+            var: var.clone(dimensions=new_dims(var.dimensions, var.shape))
             for var in FindVariables().visit(routine.body)
-                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+            if isinstance(var, sym.Array) and var.shape and len(var.shape)
         }
+    else:
+        raise ValueError(f'Unsupported array order "{order}"')
+
     routine.body = SubstituteExpressions(array_map).visit(routine.body)
 
     routine.variables = [v.clone(dimensions=as_tuple(sym.Product(v.shape)),
diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py
index 0986470f4..f023d7d70 100644
--- a/tests/test_transform_array_indexing.py
+++ b/tests/test_transform_array_indexing.py
@@ -9,19 +9,26 @@
 import pytest
 import numpy as np
 
-from conftest import jit_compile, clean_test, available_frontends
-from loki import Subroutine, FindNodes, Assignment
+from conftest import jit_compile, jit_compile_lib, clean_test, available_frontends
+from loki import Subroutine, FindVariables, Array
 from loki.expression import symbols as sym
 from loki.transform import (
         promote_variables, demote_variables, normalize_range_indexing,
-        invert_array_indices, flatten_arrays
+        invert_array_indices, flatten_arrays,
+        normalize_array_access
         )
-
+from loki.transform import (
+    FortranCTransformation
+    )
+from loki.build import Builder
 
 @pytest.fixture(scope='module', name='here')
 def fixture_here():
     return Path(__file__).parent
 
+@pytest.fixture(scope='module', name='builder')
+def fixture_builder(here):
+    return Builder(source_dirs=here, build_dir=here/'build')
 
 @pytest.mark.parametrize('frontend', available_frontends())
 def test_transform_promote_variable_scalar(here, frontend):
@@ -319,38 +326,133 @@ def test_transform_demote_dimension_arguments(here, frontend):
     assert np.all(vec2 == 2) and np.sum(vec2) == 6
     assert np.all(matrix == 16) and np.sum(matrix) == 32
 
+@pytest.mark.parametrize('frontend', available_frontends())
+@pytest.mark.parametrize('start_index', (0, 1, 5))
+def test_normalize_array_access(here, frontend, start_index):
+    """
+    Test flattening or arrays, meaning converting multi-dimensional
+    arrays to one-dimensional arrays including corresponding
+    index arithmetic.
+    """
+    fcode = f"""
+    subroutine transform_normalize_array_access(x1, x2, x3, x4, l1, l2, l3, l4)
+        implicit none
+        integer :: i1, i2, i3, i4, c1, c2, c3, c4
+        integer, intent(in) :: l1, l2, l3, l4
+        integer, intent(inout) :: x1({start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x2({start_index}:l2+{start_index}-1, &
+         & {start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x3({start_index}:l3+{start_index}-1, &
+         & {start_index}:l2+{start_index}-1, {start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x4({start_index}:l4+{start_index}-1, &
+         & {start_index}:l3+{start_index}-1, {start_index}:l2+{start_index}-1, &
+         & {start_index}:l1+{start_index}-1)
+        c1 = 1
+        c2 = 1
+        c3 = 1
+        c4 = 1
+        x1({start_index}:l4+{start_index}-1) = 0
+        do i1={start_index},l1+{start_index}-1
+            x1(i1) = c1
+            do i2={start_index},l2+{start_index}-1
+                x2(i2, i1) = c2*10 + c1
+                do i3={start_index},l3+{start_index}-1
+                    x3(i3, i2, i1) = c3*100 + c2*10 + c1
+                    do i4={start_index},l4+{start_index}-1
+                        x4(i4, i3, i2, i1) = c4*1000 + c3*100 + c2*10 + c1
+                        c4 = c4 + 1
+                    end do
+                    c3 = c3 + 1
+                end do
+                c2 = c2 + 1
+            end do
+            c1 = c1 + 1
+        end do
+
+    end subroutine transform_normalize_array_access
+    """
+    def init_arguments(l1, l2, l3, l4):
+        x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
+        x2 = np.zeros(shape=(l2,l1,), order='F', dtype=np.int32)
+        x3 = np.zeros(shape=(l3,l2,l1,), order='F', dtype=np.int32)
+        x4 = np.zeros(shape=(l4,l3,l2,l1,), order='F', dtype=np.int32)
+        return x1, x2, x3, x4
+
+    def validate_routine(routine):
+        arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, Array)]
+        for arr in arrays:
+            assert all(not isinstance(shape, sym.RangeIndex) for shape in arr.shape)
+
+    l1 = 2
+    l2 = 3
+    l3 = 4
+    l4 = 5
+    routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_range_indexing(routine) # Fix OMNI nonsense
+    filepath = here/(f'{routine.name}_{frontend}.f90')
+    function = jit_compile(routine, filepath=filepath, objname=routine.name)
+    orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4)
+    function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4)
+
+    routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_array_access(routine)
+    normalize_range_indexing(routine)
+    filepath = here/(f'{routine.name}_normalized_{frontend}.f90')
+    function = jit_compile(routine, filepath=filepath, objname=routine.name)
+    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4)
+    function(x1, x2, x3, x4, l1, l2, l3, l4)
+    validate_routine(routine)
+
+    assert (x1 == orig_x1).all()
+    assert (x2 == orig_x2).all()
+    assert (x3 == orig_x3).all()
+    assert (x4 == orig_x4).all()
+
 
 @pytest.mark.parametrize('frontend', available_frontends())
-def test_transform_flatten_arrays(here, frontend):
+@pytest.mark.parametrize('start_index', (0, 1, 5))
+def test_transform_flatten_arrays(here, frontend, builder, start_index):
     """
     Test flattening or arrays, meaning converting multi-dimensional
     arrays to one-dimensional arrays including corresponding
     index arithmetic.
     """
-    fcode = """
+    fcode = f"""
     subroutine transform_flatten_arrays(x1, x2, x3, x4, l1, l2, l3, l4)
         implicit none
-        integer :: i1, i2, i3, i4
+        integer :: i1, i2, i3, i4, c1, c2, c3, c4
         integer, intent(in) :: l1, l2, l3, l4
-        integer, intent(inout) :: x1(l1), x2(l2, l1), x3(l3, l2, l1), x4(l4, l3, l2, l1)
-        
-        do i1=1,l1
-            x1(i1) = 2 * l1
-            do i2=1,l2
-                x2(i2, i1) = 2 * l2
-                do i3=1,l3
-                    x3(i3, i2, i1) = 2 * l3
-                    do i4=1,l4
-                        x4(i4, i3, i2, i1) = 2 * l4
+        integer, intent(inout) :: x1({start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x2({start_index}:l2+{start_index}-1, &
+         & {start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x3({start_index}:l3+{start_index}-1, &
+         & {start_index}:l2+{start_index}-1, {start_index}:l1+{start_index}-1)
+        integer, intent(inout) :: x4({start_index}:l4+{start_index}-1, &
+         & {start_index}:l3+{start_index}-1, {start_index}:l2+{start_index}-1, &
+         & {start_index}:l1+{start_index}-1)
+        c1 = 1
+        c2 = 1
+        c3 = 1
+        c4 = 1
+        do i1={start_index},l1+{start_index}-1
+            x1(i1) = c1
+            do i2={start_index},l2+{start_index}-1
+                x2(i2, i1) = c2*10 + c1
+                do i3={start_index},l3+{start_index}-1
+                    x3(i3, i2, i1) = c3*100 + c2*10 + c1
+                    do i4={start_index},l4+{start_index}-1
+                        x4(i4, i3, i2, i1) = c4*1000 + c3*100 + c2*10 + c1
+                        c4 = c4 + 1
                     end do
+                    c3 = c3 + 1
                 end do
+                c2 = c2 + 1
             end do
+            c1 = c1 + 1
         end do
-        
-        
+
     end subroutine transform_flatten_arrays
     """
-
     def init_arguments(l1, l2, l3, l4, flattened=False):
         x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
         x2 = np.zeros(shape=(l2*l1) if flattened else (l2,l1,), order='F', dtype=np.int32)
@@ -358,53 +460,75 @@ def init_arguments(l1, l2, l3, l4, flattened=False):
         x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l4,l3,l2,l1,), order='F', dtype=np.int32)
         return x1, x2, x3, x4
 
-    def validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4):
-        assert np.all(x1 == 2 * l1)
-        assert np.all(x2 == 2 * l2)
-        assert np.all(x3 == 2 * l3)
-        assert np.all(x4 == 2 * l4)
-
     def validate_routine(routine):
-        assignments = FindNodes(Assignment).visit(routine.body)
-        lhs = [assignment.lhs for assignment in assignments]
-        dims = [_lhs.dimensions for _lhs in lhs]
-        assert all(len(dim) == 1 for dim in dims)
-        assert dims[1][0].children[0] == 'i2'
-        assert dims[2][0].children[0] == 'i3'
-        assert dims[3][0].children[0] == 'i4'
+        arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, Array)]
+        assert all(len(arr.dimensions) == 1 for arr in arrays)
+        assert all(len(arr.shape) == 1 for arr in arrays)
 
     l1 = 2
-    l2 = 4
-    l3 = 6
-    l4 = 8
+    l2 = 3
+    l3 = 4
+    l4 = 5
     # Test the original implementation
     routine = Subroutine.from_source(fcode, frontend=frontend)
     normalize_range_indexing(routine) # Fix OMNI nonsense
-    filepath = here/(f'{routine.name}_{frontend}.f90')
+    filepath = here/(f'{routine.name}_{start_index}_{frontend}.f90')
     function = jit_compile(routine, filepath=filepath, objname=routine.name)
-    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4)
-    function(x1, x2, x3, x4, l1, l2, l3, l4)
-    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+    orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4)
+    function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4)
 
+    # Test flattening order='F'
     f_routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_array_access(f_routine)
     normalize_range_indexing(f_routine) # Fix OMNI nonsense
     flatten_arrays(routine=f_routine, order='F', start_index=1)
-    filepath = here/(f'{f_routine.name}_flattened_F_{frontend}.f90')
+    filepath = here/(f'{f_routine.name}_{start_index}_flattened_F_{frontend}.f90')
     function = jit_compile(f_routine, filepath=filepath, objname=routine.name)
-    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4, flattened=True)
-    function(x1, x2, x3, x4, l1, l2, l3, l4)
-    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+    f_x1, f_x2, f_x3, f_x4 = init_arguments(l1, l2, l3, l4, flattened=True)
+    function(f_x1, f_x2, f_x3, f_x4, l1, l2, l3, l4)
     validate_routine(f_routine)
 
+    assert (f_x1 == orig_x1.flatten()).all()
+    assert (f_x2 == orig_x2.flatten()).all()
+    assert (f_x3 == orig_x3.flatten()).all()
+    assert (f_x4 == orig_x4.flatten()).all()
+
+    # Test flattening order='C'
     c_routine = Subroutine.from_source(fcode, frontend=frontend)
+    normalize_array_access(c_routine)
     normalize_range_indexing(c_routine) # Fix OMNI nonsense
     invert_array_indices(c_routine)
     flatten_arrays(routine=c_routine, order='C', start_index=1)
-    filepath = here/(f'{c_routine.name}_flattened_C_{frontend}.f90')
+    filepath = here/(f'{c_routine.name}_{start_index}_flattened_C_{frontend}.f90')
     function = jit_compile(c_routine, filepath=filepath, objname=routine.name)
-    x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4, flattened=True)
-    function(x1, x2, x3, x4, l1, l2, l3, l4)
-    validate_arguments(x1, x2, x3, x4, l1, l2, l3, l4)
+    c_x1, c_x2, c_x3, c_x4 = init_arguments(l1, l2, l3, l4, flattened=True)
+    function(c_x1, c_x2, c_x3, c_x4, l1, l2, l3, l4)
     validate_routine(c_routine)
 
     assert f_routine.body == c_routine.body
+
+    assert (c_x1 == orig_x1.flatten()).all()
+    assert (c_x2 == orig_x2.flatten()).all()
+    assert (c_x3 == orig_x3.flatten()).all()
+    assert (c_x4 == orig_x4.flatten()).all()
+
+    # Test C transpilation (which includes flattening)
+    f2c_routine = Subroutine.from_source(fcode, frontend=frontend)
+    f2c = FortranCTransformation(order="C")
+    f2c.apply(source=f2c_routine, path=here)
+    libname = f'fc_{f2c_routine.name}_{start_index}_{frontend}'
+    c_kernel = jit_compile_lib([f2c.wrapperpath, f2c.c_path], path=here, name=libname, builder=builder)
+    fc_function = c_kernel.transform_flatten_arrays_fc_mod.transform_flatten_arrays_fc
+    f2c_x1, f2c_x2, f2c_x3, f2c_x4 = init_arguments(l1, l2, l3, l4, flattened=True)
+    fc_function(f2c_x1, f2c_x2, f2c_x3, f2c_x4, l1, l2, l3, l4)
+    validate_routine(c_routine)
+
+    assert (f2c_x1 == orig_x1.flatten()).all()
+    assert (f2c_x2 == orig_x2.flatten()).all()
+    assert (f2c_x3 == orig_x3.flatten()).all()
+    assert (f2c_x4 == orig_x4.flatten()).all()
+
+    builder.clean()
+    clean_test(filepath)
+    f2c.wrapperpath.unlink()
+    f2c.c_path.unlink()
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
index c4a9a2b43..17235bad9 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -940,7 +940,8 @@ def test_transpile_expressions(here, builder, frontend):
 
     # Make sure minus signs are represented correctly in the C code
     ccode = cgen(routine)
-    assert 'vector[i - 1 - 1]' in ccode  # double minus due to index shift to 0
+    # double minus due to index shift to 0
+    assert 'vector[i - 1 - 1]' in ccode or 'vector[-1 + i - 1]' in ccode
     assert 'vector[i - 1]' in ccode
     assert '-scalar' in ccode  # scalar with negative sign
 

From 5bdfafeae56d774ff29b69f808c76703b631b8a0 Mon Sep 17 00:00:00 2001
From: Michael Staneker <michael.staneker@ecmwf.int>
Date: Tue, 19 Dec 2023 08:57:23 +0000
Subject: [PATCH 3/5] adding missing 'clean_tests' for flattening array tests

---
 tests/test_transform_array_indexing.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py
index f023d7d70..fda95708a 100644
--- a/tests/test_transform_array_indexing.py
+++ b/tests/test_transform_array_indexing.py
@@ -393,6 +393,7 @@ def validate_routine(routine):
     function = jit_compile(routine, filepath=filepath, objname=routine.name)
     orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4)
     function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4)
+    clean_test(filepath)
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
     normalize_array_access(routine)
@@ -402,6 +403,7 @@ def validate_routine(routine):
     x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4)
     function(x1, x2, x3, x4, l1, l2, l3, l4)
     validate_routine(routine)
+    clean_test(filepath)
 
     assert (x1 == orig_x1).all()
     assert (x2 == orig_x2).all()
@@ -476,6 +478,7 @@ def validate_routine(routine):
     function = jit_compile(routine, filepath=filepath, objname=routine.name)
     orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4)
     function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4)
+    clean_test(filepath)
 
     # Test flattening order='F'
     f_routine = Subroutine.from_source(fcode, frontend=frontend)
@@ -487,6 +490,7 @@ def validate_routine(routine):
     f_x1, f_x2, f_x3, f_x4 = init_arguments(l1, l2, l3, l4, flattened=True)
     function(f_x1, f_x2, f_x3, f_x4, l1, l2, l3, l4)
     validate_routine(f_routine)
+    clean_test(filepath)
 
     assert (f_x1 == orig_x1.flatten()).all()
     assert (f_x2 == orig_x2.flatten()).all()
@@ -504,6 +508,7 @@ def validate_routine(routine):
     c_x1, c_x2, c_x3, c_x4 = init_arguments(l1, l2, l3, l4, flattened=True)
     function(c_x1, c_x2, c_x3, c_x4, l1, l2, l3, l4)
     validate_routine(c_routine)
+    clean_test(filepath)
 
     assert f_routine.body == c_routine.body
 

From 1e88158d116b2f0fc12fba5d5553378eba487764 Mon Sep 17 00:00:00 2001
From: Michael Staneker <michael.staneker@ecmwf.int>
Date: Tue, 19 Dec 2023 09:22:36 +0000
Subject: [PATCH 4/5] Fix memory order (C vs Fortran) issue/logic

---
 loki/transform/fortran_c_transform.py      |  5 ++--
 loki/transform/transform_array_indexing.py |  4 +--
 tests/test_transform_array_indexing.py     | 32 ++++++++++++----------
 3 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/loki/transform/fortran_c_transform.py b/loki/transform/fortran_c_transform.py
index ad563c243..1295a90c2 100644
--- a/loki/transform/fortran_c_transform.py
+++ b/loki/transform/fortran_c_transform.py
@@ -50,9 +50,8 @@ class FortranCTransformation(Transformation):
     # Set of standard module names that have no C equivalent
     __fortran_intrinsic_modules = ['ISO_FORTRAN_ENV', 'ISO_C_BINDING']
 
-    def __init__(self, header_modules=None, inline_elementals=True, order='F'):
+    def __init__(self, header_modules=None, inline_elementals=True):
         self.inline_elementals = inline_elementals
-        self.order = order
 
         # Maps from original type name to ISO-C and C-struct types
         self.c_structs = OrderedDict()
@@ -376,7 +375,7 @@ def generate_c_kernel(self, routine, **kwargs):
         # TODO: Resolve reductions (eg. SUM(myvar(:)))
         invert_array_indices(kernel)
         shift_to_zero_indexing(kernel)
-        flatten_arrays(kernel, order=self.order, start_index=0)
+        flatten_arrays(kernel, order='C', start_index=0)
 
         # Inline all known parameters, since they can be used in declarations,
         # and thus need to be known before we can fetch them via getters.
diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py
index 33587f3b3..b711a6d3a 100644
--- a/loki/transform/transform_array_indexing.py
+++ b/loki/transform/transform_array_indexing.py
@@ -560,13 +560,13 @@ def new_dims(dim, shape):
             return new_dims(new_dim, shape[:-1])
         return dim
 
-    if order == 'F':
+    if order == 'C':
         array_map = {
             var: var.clone(dimensions=new_dims(var.dimensions[::-1], var.shape[::-1]))
             for var in FindVariables().visit(routine.body)
             if isinstance(var, sym.Array) and var.shape and len(var.shape)
         }
-    elif order == 'C':
+    elif order == 'F':
         array_map = {
             var: var.clone(dimensions=new_dims(var.dimensions, var.shape))
             for var in FindVariables().visit(routine.body)
diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py
index fda95708a..86d3f6f1e 100644
--- a/tests/test_transform_array_indexing.py
+++ b/tests/test_transform_array_indexing.py
@@ -328,7 +328,7 @@ def test_transform_demote_dimension_arguments(here, frontend):
 
 @pytest.mark.parametrize('frontend', available_frontends())
 @pytest.mark.parametrize('start_index', (0, 1, 5))
-def test_normalize_array_access(here, frontend, start_index):
+def test_transform_normalize_array_access(here, frontend, start_index):
     """
     Test flattening or arrays, meaning converting multi-dimensional
     arrays to one-dimensional arrays including corresponding
@@ -460,6 +460,10 @@ def init_arguments(l1, l2, l3, l4, flattened=False):
         x2 = np.zeros(shape=(l2*l1) if flattened else (l2,l1,), order='F', dtype=np.int32)
         x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l3,l2,l1,), order='F', dtype=np.int32)
         x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l4,l3,l2,l1,), order='F', dtype=np.int32)
+        #x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
+        #x2 = np.zeros(shape=(l2*l1) if flattened else (l1,l2), order='F', dtype=np.int32)
+        #x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l1,l2,l3), order='F', dtype=np.int32)
+        #x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l1,l2,l3,l4), order='F', dtype=np.int32)
         return x1, x2, x3, x4
 
     def validate_routine(routine):
@@ -492,10 +496,10 @@ def validate_routine(routine):
     validate_routine(f_routine)
     clean_test(filepath)
 
-    assert (f_x1 == orig_x1.flatten()).all()
-    assert (f_x2 == orig_x2.flatten()).all()
-    assert (f_x3 == orig_x3.flatten()).all()
-    assert (f_x4 == orig_x4.flatten()).all()
+    assert (f_x1 == orig_x1.flatten(order='F')).all()
+    assert (f_x2 == orig_x2.flatten(order='F')).all()
+    assert (f_x3 == orig_x3.flatten(order='F')).all()
+    assert (f_x4 == orig_x4.flatten(order='F')).all()
 
     # Test flattening order='C'
     c_routine = Subroutine.from_source(fcode, frontend=frontend)
@@ -512,14 +516,14 @@ def validate_routine(routine):
 
     assert f_routine.body == c_routine.body
 
-    assert (c_x1 == orig_x1.flatten()).all()
-    assert (c_x2 == orig_x2.flatten()).all()
-    assert (c_x3 == orig_x3.flatten()).all()
-    assert (c_x4 == orig_x4.flatten()).all()
+    assert (c_x1 == orig_x1.flatten(order='F')).all()
+    assert (c_x2 == orig_x2.flatten(order='F')).all()
+    assert (c_x3 == orig_x3.flatten(order='F')).all()
+    assert (c_x4 == orig_x4.flatten(order='F')).all()
 
     # Test C transpilation (which includes flattening)
     f2c_routine = Subroutine.from_source(fcode, frontend=frontend)
-    f2c = FortranCTransformation(order="C")
+    f2c = FortranCTransformation()
     f2c.apply(source=f2c_routine, path=here)
     libname = f'fc_{f2c_routine.name}_{start_index}_{frontend}'
     c_kernel = jit_compile_lib([f2c.wrapperpath, f2c.c_path], path=here, name=libname, builder=builder)
@@ -528,10 +532,10 @@ def validate_routine(routine):
     fc_function(f2c_x1, f2c_x2, f2c_x3, f2c_x4, l1, l2, l3, l4)
     validate_routine(c_routine)
 
-    assert (f2c_x1 == orig_x1.flatten()).all()
-    assert (f2c_x2 == orig_x2.flatten()).all()
-    assert (f2c_x3 == orig_x3.flatten()).all()
-    assert (f2c_x4 == orig_x4.flatten()).all()
+    assert (f2c_x1 == orig_x1.flatten(order='F')).all()
+    assert (f2c_x2 == orig_x2.flatten(order='F')).all()
+    assert (f2c_x3 == orig_x3.flatten(order='F')).all()
+    assert (f2c_x4 == orig_x4.flatten(order='F')).all()
 
     builder.clean()
     clean_test(filepath)

From b2fef9b528ff0e69b50175e2e720207f54552682 Mon Sep 17 00:00:00 2001
From: Michael Staneker <michael.staneker@ecmwf.int>
Date: Tue, 19 Dec 2023 16:32:58 +0000
Subject: [PATCH 5/5] rename new utility to 'normalize_array_shape_and_access'

---
 loki/transform/fortran_c_transform.py      |  7 +++----
 loki/transform/transform_array_indexing.py |  4 ++--
 tests/test_transform_array_indexing.py     | 18 +++++++-----------
 3 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/loki/transform/fortran_c_transform.py b/loki/transform/fortran_c_transform.py
index 1295a90c2..45445a65e 100644
--- a/loki/transform/fortran_c_transform.py
+++ b/loki/transform/fortran_c_transform.py
@@ -11,8 +11,8 @@
 from loki.transform.transformation import Transformation
 from loki.transform.transform_array_indexing import (
     shift_to_zero_indexing, invert_array_indices,
-    resolve_vector_notation, normalize_range_indexing,
-    normalize_array_access, flatten_arrays
+    resolve_vector_notation, normalize_array_shape_and_access,
+    flatten_arrays
 )
 from loki.transform.transform_associates import resolve_associates
 from loki.transform.transform_utilities import (
@@ -368,8 +368,7 @@ def generate_c_kernel(self, routine, **kwargs):
 
         # Clean up Fortran vector notation
         resolve_vector_notation(kernel)
-        normalize_array_access(kernel)
-        normalize_range_indexing(kernel)
+        normalize_array_shape_and_access(kernel)
 
         # Convert array indexing to C conventions
         # TODO: Resolve reductions (eg. SUM(myvar(:)))
diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py
index b711a6d3a..bce200115 100644
--- a/loki/transform/transform_array_indexing.py
+++ b/loki/transform/transform_array_indexing.py
@@ -29,7 +29,7 @@
     'resolve_vector_notation', 'normalize_range_indexing',
     'promote_variables', 'promote_nonmatching_variables',
     'promotion_dimensions_from_loop_nest', 'demote_variables',
-    'flatten_arrays', 'normalize_array_access'
+    'flatten_arrays', 'normalize_array_shape_and_access'
 ]
 
 
@@ -498,7 +498,7 @@ def demote_variables(routine, variable_names, dimensions):
 
     info(f'[Loki::Transform] Demoted variables in {routine.name}: {", ".join(variable_names)}')
 
-def normalize_array_access(routine):
+def normalize_array_shape_and_access(routine):
     """
     Shift all arrays to start counting at "1"
     """
diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py
index 86d3f6f1e..2ad912972 100644
--- a/tests/test_transform_array_indexing.py
+++ b/tests/test_transform_array_indexing.py
@@ -15,7 +15,7 @@
 from loki.transform import (
         promote_variables, demote_variables, normalize_range_indexing,
         invert_array_indices, flatten_arrays,
-        normalize_array_access
+        normalize_array_shape_and_access
         )
 from loki.transform import (
     FortranCTransformation
@@ -328,14 +328,14 @@ def test_transform_demote_dimension_arguments(here, frontend):
 
 @pytest.mark.parametrize('frontend', available_frontends())
 @pytest.mark.parametrize('start_index', (0, 1, 5))
-def test_transform_normalize_array_access(here, frontend, start_index):
+def test_transform_normalize_array_shape_and_access(here, frontend, start_index):
     """
     Test flattening or arrays, meaning converting multi-dimensional
     arrays to one-dimensional arrays including corresponding
     index arithmetic.
     """
     fcode = f"""
-    subroutine transform_normalize_array_access(x1, x2, x3, x4, l1, l2, l3, l4)
+    subroutine transform_normalize_array_shape_and_access(x1, x2, x3, x4, l1, l2, l3, l4)
         implicit none
         integer :: i1, i2, i3, i4, c1, c2, c3, c4
         integer, intent(in) :: l1, l2, l3, l4
@@ -369,7 +369,7 @@ def test_transform_normalize_array_access(here, frontend, start_index):
             c1 = c1 + 1
         end do
 
-    end subroutine transform_normalize_array_access
+    end subroutine transform_normalize_array_shape_and_access
     """
     def init_arguments(l1, l2, l3, l4):
         x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
@@ -396,7 +396,7 @@ def validate_routine(routine):
     clean_test(filepath)
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
-    normalize_array_access(routine)
+    normalize_array_shape_and_access(routine)
     normalize_range_indexing(routine)
     filepath = here/(f'{routine.name}_normalized_{frontend}.f90')
     function = jit_compile(routine, filepath=filepath, objname=routine.name)
@@ -460,10 +460,6 @@ def init_arguments(l1, l2, l3, l4, flattened=False):
         x2 = np.zeros(shape=(l2*l1) if flattened else (l2,l1,), order='F', dtype=np.int32)
         x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l3,l2,l1,), order='F', dtype=np.int32)
         x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l4,l3,l2,l1,), order='F', dtype=np.int32)
-        #x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32)
-        #x2 = np.zeros(shape=(l2*l1) if flattened else (l1,l2), order='F', dtype=np.int32)
-        #x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l1,l2,l3), order='F', dtype=np.int32)
-        #x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l1,l2,l3,l4), order='F', dtype=np.int32)
         return x1, x2, x3, x4
 
     def validate_routine(routine):
@@ -486,7 +482,7 @@ def validate_routine(routine):
 
     # Test flattening order='F'
     f_routine = Subroutine.from_source(fcode, frontend=frontend)
-    normalize_array_access(f_routine)
+    normalize_array_shape_and_access(f_routine)
     normalize_range_indexing(f_routine) # Fix OMNI nonsense
     flatten_arrays(routine=f_routine, order='F', start_index=1)
     filepath = here/(f'{f_routine.name}_{start_index}_flattened_F_{frontend}.f90')
@@ -503,7 +499,7 @@ def validate_routine(routine):
 
     # Test flattening order='C'
     c_routine = Subroutine.from_source(fcode, frontend=frontend)
-    normalize_array_access(c_routine)
+    normalize_array_shape_and_access(c_routine)
     normalize_range_indexing(c_routine) # Fix OMNI nonsense
     invert_array_indices(c_routine)
     flatten_arrays(routine=c_routine, order='C', start_index=1)