From cfdfeb87e5009b0939a16053c050c9cdf8c7fe0b Mon Sep 17 00:00:00 2001
From: kabdelhak <38032125+kabdelhak@users.noreply.github.com>
Date: Sat, 27 Jan 2024 16:32:52 +0100
Subject: [PATCH] [NB] update event collection (#11895)

- do not collect inside when equation bodies because they cannot trigger events
---
 .../Compiler/NBackEnd/Classes/NBEquation.mo   | 37 +++++++++++++++-
 .../NBackEnd/Modules/2_Pre/NBEvents.mo        | 42 ++++++++++++++-----
 2 files changed, 67 insertions(+), 12 deletions(-)

diff --git a/OMCompiler/Compiler/NBackEnd/Classes/NBEquation.mo b/OMCompiler/Compiler/NBackEnd/Classes/NBEquation.mo
index af718bc4f13..4944629d575 100644
--- a/OMCompiler/Compiler/NBackEnd/Classes/NBEquation.mo
+++ b/OMCompiler/Compiler/NBackEnd/Classes/NBEquation.mo
@@ -2191,6 +2191,21 @@ public
       input MapFuncExp funcExp;
       input Option<MapFuncCref> funcCrefOpt;
       input MapFuncExpWrapper mapFunc;
+    algorithm
+      ifBody := mapEqnExpCref(
+        ifBody      = ifBody,
+        func        = function Pointer.apply(func = function Equation.map(funcExp = funcExp, funcCrefOpt = funcCrefOpt, mapFunc = mapFunc)),
+        funcExp     = funcExp,
+        funcCrefOpt = funcCrefOpt,
+        mapFunc     = mapFunc);
+    end map;
+
+    function mapEqnExpCref
+      input output IfEquationBody ifBody;
+      input MapFuncEqnPtr func;
+      input MapFuncExp funcExp;
+      input Option<MapFuncCref> funcCrefOpt;
+      input MapFuncExpWrapper mapFunc;
     protected
       Expression condition;
       IfEquationBody else_if, old_else_if;
@@ -2201,7 +2216,7 @@ public
       end if;
 
       // referenceEq for lists?
-      ifBody.then_eqns := List.map(ifBody.then_eqns, function Pointer.apply(func = function Equation.map(funcExp = funcExp, funcCrefOpt = funcCrefOpt, mapFunc = mapFunc)));
+      ifBody.then_eqns := List.map(ifBody.then_eqns, func);
 
       if Util.isSome(ifBody.else_if) then
         old_else_if := Util.getOption(ifBody.else_if);
@@ -2210,7 +2225,7 @@ public
           ifBody.else_if := SOME(else_if);
         end if;
       end if;
-    end map;
+    end mapEqnExpCref;
 
     function size
       "only considers first branch"
@@ -2429,6 +2444,24 @@ public
       whenBody.else_when := Util.applyOption(whenBody.else_when, function map(funcExp = funcExp, funcCrefOpt = funcCrefOpt, mapFunc = mapFunc));
     end map;
 
+    function mapCondition
+      "only maps the conditions and not the body"
+      input output WhenEquationBody whenBody;
+      input MapFuncExp funcExp;
+      input Option<MapFuncCref> funcCrefOpt;
+      input MapFuncExpWrapper mapFunc;
+    protected
+      Expression condition;
+    algorithm
+      condition := mapFunc(whenBody.condition, funcExp);
+      if not referenceEq(condition, whenBody.condition) then
+        whenBody.condition := condition;
+      end if;
+
+      // map else when
+      whenBody.else_when := Util.applyOption(whenBody.else_when, function mapCondition(funcExp = funcExp, funcCrefOpt = funcCrefOpt, mapFunc = mapFunc));
+    end mapCondition;
+
     function split
       "this function splits up when equations while respecting to keep
       correct branches for assigned discrete states and reinitialized states.
diff --git a/OMCompiler/Compiler/NBackEnd/Modules/2_Pre/NBEvents.mo b/OMCompiler/Compiler/NBackEnd/Modules/2_Pre/NBEvents.mo
index 89d7a48ea5d..3815f04ddcd 100644
--- a/OMCompiler/Compiler/NBackEnd/Modules/2_Pre/NBEvents.mo
+++ b/OMCompiler/Compiler/NBackEnd/Modules/2_Pre/NBEvents.mo
@@ -946,10 +946,20 @@ protected
     input Pointer<Bucket> bucket_ptr;
     input FunctionTree funcTree;
   protected
-    Equation eqn = Pointer.access(eqn_ptr);
+    Equation eqn = Pointer.access(eqn_ptr), body_eqn;
     Iterator iter;
     Boolean createAux = not Equation.isAlgorithm(eqn_ptr);
+    BEquation.MapFuncExp collector;
   algorithm
+    // create the traverser function
+    iter := Equation.getForIterator(eqn);
+    collector := function collectEventsTraverse(
+          bucket_ptr  = bucket_ptr,
+          iter        = iter,
+          eqn         = eqn_ptr,
+          funcTree    = funcTree,
+          createAux   = createAux);
+
     eqn := match eqn
       case Equation.ALGORITHM() algorithm
         for stmt in eqn.alg.statements loop
@@ -957,15 +967,27 @@ protected
         end for;
       then eqn;
 
-      else algorithm
-        iter := Equation.getForIterator(eqn);
-      then Equation.map(eqn, function collectEventsTraverse(
-          bucket_ptr  = bucket_ptr,
-          iter        = iter,
-          eqn         = eqn_ptr,
-          funcTree    = funcTree,
-          createAux   = createAux),
-        NONE(), Expression.mapReverse);
+      // For when equations only map the condition and not the body
+      case Equation.WHEN_EQUATION() algorithm
+        eqn.body := WhenEquationBody.mapCondition(eqn.body, collector, NONE(), Expression.mapReverse);
+      then eqn;
+
+      // Also don't do it for when equations in for-equations
+      case Equation.FOR_EQUATION(body = {body_eqn as Equation.WHEN_EQUATION()}) algorithm
+        body_eqn.body := WhenEquationBody.mapCondition(body_eqn.body, collector, NONE(), Expression.mapReverse);
+        eqn.body := {body_eqn};
+      then eqn;
+
+      // Map if equation body with this function to ensure that when equation bodies are not traversed
+      case Equation.IF_EQUATION() algorithm
+        eqn.body := IfEquationBody.mapEqnExpCref(eqn.body,
+          func        = function collectEvents(bucket_ptr = bucket_ptr, funcTree = funcTree),
+          funcExp     = collector,
+          funcCrefOpt = NONE(),
+          mapFunc     = Expression.mapReverse);
+      then eqn;
+
+      else Equation.map(eqn, collector, NONE(), Expression.mapReverse);
     end match;
 
     if not referenceEq(eqn, Pointer.access(eqn_ptr)) then