From caa88359da12dc6b87a99d246f1c8dd7db55f408 Mon Sep 17 00:00:00 2001
From: perost <perost86@gmail.com>
Date: Thu, 17 Feb 2022 17:38:04 +0100
Subject: [PATCH] Improve flattening performance (#8573)

- Change the flattening of subscripted expressions with split indices so
  that the subscripting is done before flattening the subscripted
  expression, to avoid flattening large expressions over and over.
---
 OMCompiler/Compiler/NFFrontEnd/NFCall.mo    |  9 +++++---
 OMCompiler/Compiler/NFFrontEnd/NFFlatten.mo | 25 +++++++++++----------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/OMCompiler/Compiler/NFFrontEnd/NFCall.mo b/OMCompiler/Compiler/NFFrontEnd/NFCall.mo
index f0a3d0902f6..94ff5169050 100644
--- a/OMCompiler/Compiler/NFFrontEnd/NFCall.mo
+++ b/OMCompiler/Compiler/NFFrontEnd/NFCall.mo
@@ -1608,20 +1608,23 @@ public
       case UNTYPED_ARRAY_CONSTRUCTOR()
         algorithm
           e := func(call.exp);
+          iters := mapIteratorsExpShallow(call.iters, func);
         then
-          UNTYPED_ARRAY_CONSTRUCTOR(e, call.iters);
+          UNTYPED_ARRAY_CONSTRUCTOR(e, iters);
 
       case TYPED_ARRAY_CONSTRUCTOR()
         algorithm
           e := func(call.exp);
+          iters := mapIteratorsExpShallow(call.iters, func);
         then
-          TYPED_ARRAY_CONSTRUCTOR(call.ty, call.var, call.purity, e, call.iters);
+          TYPED_ARRAY_CONSTRUCTOR(call.ty, call.var, call.purity, e, iters);
 
       case UNTYPED_REDUCTION()
         algorithm
           e := func(call.exp);
+          iters := mapIteratorsExpShallow(call.iters, func);
         then
-          UNTYPED_REDUCTION(call.ref, e, call.iters);
+          UNTYPED_REDUCTION(call.ref, e, iters);
 
       case TYPED_REDUCTION()
         algorithm
diff --git a/OMCompiler/Compiler/NFFrontEnd/NFFlatten.mo b/OMCompiler/Compiler/NFFrontEnd/NFFlatten.mo
index 36f90c24505..7be058de1bc 100644
--- a/OMCompiler/Compiler/NFFrontEnd/NFFlatten.mo
+++ b/OMCompiler/Compiler/NFFrontEnd/NFFlatten.mo
@@ -1034,31 +1034,31 @@ end flattenBinding;
 public function flattenExp
   input output Expression exp;
   input ComponentRef prefix;
-algorithm
-  exp := Expression.map(exp, function flattenExp_traverse(prefix = prefix));
-end flattenExp;
-
-function flattenExp_traverse
-  input output Expression exp;
-  input ComponentRef prefix;
 algorithm
   exp := match exp
     case Expression.CREF(cref = ComponentRef.CREF())
       algorithm
+        exp.cref := ComponentRef.mapExpShallow(exp.cref, function flattenExp(prefix = prefix));
         exp.cref := flattenCref(exp.cref, prefix);
         exp.ty := flattenType(exp.ty, prefix);
       then
         exp;
 
-    case Expression.SUBSCRIPTED_EXP()
-      then replaceSplitIndices(exp.exp, exp.subscripts, prefix);
+    case Expression.SUBSCRIPTED_EXP(split = true)
+      then Expression.mapShallow(
+        replaceSplitIndices(exp.exp, exp.subscripts, prefix),
+        function flattenExp(prefix = prefix));
 
-    case Expression.IF(ty = Type.CONDITIONAL_ARRAY()) then flattenConditionalArrayIfExp(exp);
-    else exp;
+    case Expression.IF(ty = Type.CONDITIONAL_ARRAY())
+      then Expression.mapShallow(
+          flattenConditionalArrayIfExp(exp),
+          function flattenExp(prefix = prefix));
+
+    else Expression.mapShallow(exp, function flattenExp(prefix = prefix));
   end match;
 
   exp := flattenExpType(exp, prefix);
-end flattenExp_traverse;
+end flattenExp;
 
 function replaceSplitIndices
   input output Expression exp;
@@ -1086,6 +1086,7 @@ algorithm
 
   subs := Subscript.expandSplitIndices(subs);
   exp := Expression.applySubscripts(subs, exp);
+  exp := flattenExp(exp, prefix);
 end replaceSplitIndices;
 
 function replaceSplitIndices2