From 9c43d4fc86d060cf8ca7dc97ff4fc2ce84cc7a2f Mon Sep 17 00:00:00 2001
From: "petro.zarytskyi" <petro.zarytskyi@gmail.com>
Date: Tue, 21 Nov 2023 22:06:04 +0200
Subject: [PATCH] Do not change predecessors in VarsData merge.

---
 lib/Differentiator/TBRAnalyzer.cpp | 57 +++++++++++++++++++++---------
 1 file changed, 40 insertions(+), 17 deletions(-)

diff --git a/lib/Differentiator/TBRAnalyzer.cpp b/lib/Differentiator/TBRAnalyzer.cpp
index d9c226540..185e30725 100644
--- a/lib/Differentiator/TBRAnalyzer.cpp
+++ b/lib/Differentiator/TBRAnalyzer.cpp
@@ -396,19 +396,19 @@ void TBRAnalyzer::VisitCFGBlock(const CFGBlock& block) {
     /// a loop condition and the loop body has already been visited 2 times.
     /// This means we should not visit the loop body anymore.
     if (notLastPass) {
-        /// Add the successor to the queue.
-        m_CFGQueue.insert(succ->getBlockID());
-
-        /// This part is necessary for loops. For other cases, this is not
-        /// supposed to do anything.
-        if (succ->getBlockID() < block.getBlockID()) {
-            /// If there is another loop condition present inside a loop,
-            /// We have to set it's loop pass counter to 0 (it might be 3
-            /// from the previous outer loop pass).
-            m_BlockPassCounter[succ->getBlockID()] = 0;
-            /// Remove VarsData left after the previous pass.
-            varsData->clear();
-        }
+      /// Add the successor to the queue.
+      m_CFGQueue.insert(succ->getBlockID());
+
+      /// This part is necessary for loops. For other cases, this is not
+      /// supposed to do anything.
+      if (succ->getBlockID() < block.getBlockID()) {
+        /// If there is another loop condition present inside a loop,
+        /// We have to set it's loop pass counter to 0 (it might be 3
+        /// from the previous outer loop pass).
+        m_BlockPassCounter[succ->getBlockID()] = 0;
+        /// Remove VarsData left after the previous pass.
+        varsData->clear();
+      }
     }
 
     /// If the successor's previous block is not this one,
@@ -495,7 +495,6 @@ void TBRAnalyzer::merge(VarsData* targetData, VarsData* mergeData) {
   auto* LCA = findLowestCommonAncestor(targetData, mergeData);
   auto collectedMergeData =
       collectDataFromPredecessors(mergeData, /*limit=*/LCA);
-  auto collectedTargetData = collectDataFromPredecessors(targetData, /*limit=*/LCA);
 
   /// For every variable in 'collectedMergeData', search it in targetData
   /// and all its predecessors (if found in a predecessor, make a copy to
@@ -526,19 +525,43 @@ void TBRAnalyzer::merge(VarsData* targetData, VarsData* mergeData) {
       (*targetData)[pair.first] = copy(*pair.second);
   }
 
-  /// For every variable in collectedTargetData, search it inside
+  /// For every variable in collected targetData predecessors,
+  /// search it inside collectedMergeData. If it's not found,
+  /// that means it was not used anywhere between LCA and mergeData.
+  /// To correctly merge, we have to take it from LCA's
+  /// predecessors and merge it to targetData.
+  /// If targetData is LCA, LCA will come after targetData->prev and
+  /// collectDataFromPredecessors will not reach the limit.
+  if (targetData != LCA) {
+    for (auto& pair : collectDataFromPredecessors(targetData->prev, /*limit=*/LCA)) {
+      auto elemSearch = collectedMergeData.find(pair.first);
+      if (elemSearch == collectedMergeData.end()) {
+        auto* branch = LCA;
+        while (branch) {
+          auto it = branch->find(pair.first);
+          if (it != branch->end()) {
+            (*targetData)[pair.first] = copy(*pair.second);
+            merge((*targetData)[pair.first], it->second);
+            break;
+          }
+          branch = branch->prev;
+        }
+      }
+    }
+  }
+  /// For every variable in targetData, search it inside
   /// collectedMergeData. If it's not found, that means it
   /// was not used anywhere between LCA and mergeData.
   /// To correctly merge, we have to take it from LCA's
   /// predecessors and merge it to targetData.
-  for (auto& pair : collectedTargetData) {
+  for (auto& pair : *targetData) {
     auto elemSearch = collectedMergeData.find(pair.first);
     if (elemSearch == collectedMergeData.end()) {
       auto* branch = LCA;
       while (branch) {
         auto it = branch->find(pair.first);
         if (it != branch->end()) {
-          merge(*pair.second, it->second);
+          merge(pair.second, it->second);
           break;
         }
         branch = branch->prev;