From 31e5e43f9d0f24888a8fb720a8fc8d2b16bcac14 Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 17 Oct 2023 09:38:09 +0000
Subject: [PATCH] Module: Ensure nested scopes (TypeDefs) are rebuild during
 clone

---
 loki/program_unit.py  |  7 +++---
 tests/test_modules.py | 51 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/loki/program_unit.py b/loki/program_unit.py
index 11877e8ce..798fa0480 100644
--- a/loki/program_unit.py
+++ b/loki/program_unit.py
@@ -290,12 +290,13 @@ def clone(self, **kwargs):
         kwargs.setdefault('incomplete', self._incomplete)
 
         # Rebuild IRs
+        rebuild = Transformer({}, rebuild_scopes=True)
         if 'docstring' in kwargs:
-            kwargs['docstring'] = Transformer({}).visit(kwargs['docstring'])
+            kwargs['docstring'] = rebuild.visit(kwargs['docstring'])
         if 'spec' in kwargs:
-            kwargs['spec'] = Transformer({}).visit(kwargs['spec'])
+            kwargs['spec'] = rebuild.visit(kwargs['spec'])
         if 'contains' in kwargs:
-            kwargs['contains'] = Transformer({}).visit(kwargs['contains'])
+            kwargs['contains'] = rebuild.visit(kwargs['contains'])
 
         # Rescope symbols if not explicitly disabled
         kwargs.setdefault('rescope_symbols', True)
diff --git a/tests/test_modules.py b/tests/test_modules.py
index 55a21f3e1..bbb52726d 100644
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -13,7 +13,7 @@
     OFP, OMNI, Module, Subroutine, VariableDeclaration, TypeDef, fexprgen,
     BasicType, Assignment, FindNodes, FindInlineCalls, FindTypedSymbols,
     Transformer, fgen, SymbolAttributes, Variable, Import, Section, Intrinsic,
-    Scalar, DeferredTypeSymbol
+    Scalar, DeferredTypeSymbol, FindVariables, SubstituteExpressions, Literal
 )
 
 
@@ -529,6 +529,55 @@ def test_module_rescope_clone(frontend):
     with pytest.raises(AttributeError):
         fgen(other_module_copy)
 
+@pytest.mark.parametrize('frontend', available_frontends(
+    xfail=[(OMNI, 'Parsing fails without dummy module provided')]
+))
+def test_module_deep_clone(frontend):
+    """
+    Test the rescoping of variables in clone with nested scopes.
+    """
+    fcode = """
+module test_module_rescope_clone
+  use parkind1, only : jpim, jprb
+  implicit none
+
+  integer :: n
+
+  real :: array(n)
+
+  type my_type
+    real :: vector(n)
+    real :: matrix(n, n)
+  end type
+
+end module test_module_rescope_clone
+"""
+    module = Module.from_source(fcode, frontend=frontend)
+
+    # Deep-copy/clone the module
+    new_module = module.clone()
+
+    n = [v for v in FindVariables().visit(new_module.spec) if v.name == 'n'][0]
+    n_decl = FindNodes(VariableDeclaration).visit(new_module.spec)[0]
+
+    # Remove the declaration of `n` and replace it with `3`
+    new_module.spec = Transformer({n_decl: None}).visit(new_module.spec)
+    new_module.spec = SubstituteExpressions({n: Literal(3)}).visit(new_module.spec)
+
+    # Check the new module has been changed
+    assert len(FindNodes(VariableDeclaration).visit(new_module.spec)) == 1
+    new_type_decls = FindNodes(VariableDeclaration).visit(new_module['my_type'].body)
+    assert len(new_type_decls) == 2
+    assert new_type_decls[0].symbols[0] == 'vector(3)'
+    assert new_type_decls[1].symbols[0] == 'matrix(3, 3)'
+
+    # Check the old one has not changed
+    assert len(FindNodes(VariableDeclaration).visit(module.spec)) == 2
+    type_decls = FindNodes(VariableDeclaration).visit(module['my_type'].body)
+    assert len(type_decls) == 2
+    assert type_decls[0].symbols[0] == 'vector(n)'
+    assert type_decls[1].symbols[0] == 'matrix(n, n)'
+
 
 @pytest.mark.parametrize('frontend', available_frontends())
 def test_module_access_spec_none(frontend):