From b680a0ca9e8a8b2fc64da425afaa753706bd5384 Mon Sep 17 00:00:00 2001
From: ovdiiuv <104850830+ovdiiuv@users.noreply.github.com>
Date: Thu, 22 Aug 2024 18:55:07 +0200
Subject: [PATCH] Fix collectDataFromPredecessors to collect data from the
 passed branch (#1051)

Before the fix, we didn't collect data from varsData at all, and it wasn't merge afterward either. Now, we start collecting data from varsData, ensuring the final data of the CFG block is correct.

Fixes:#855
---
 lib/Differentiator/TBRAnalyzer.cpp |  2 +-
 test/Analyses/TBR.cpp              | 12 +++++++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/lib/Differentiator/TBRAnalyzer.cpp b/lib/Differentiator/TBRAnalyzer.cpp
index ffd0f7ce6..b3e54af36 100644
--- a/lib/Differentiator/TBRAnalyzer.cpp
+++ b/lib/Differentiator/TBRAnalyzer.cpp
@@ -484,7 +484,7 @@ TBRAnalyzer::collectDataFromPredecessors(VarsData* varsData,
   std::unordered_map<const clang::VarDecl*, VarData*> result;
   if (varsData != limit) {
     // Copy data from every predecessor.
-    for (auto* pred = varsData->m_Prev; pred != limit; pred = pred->m_Prev) {
+    for (auto* pred = varsData; pred != limit; pred = pred->m_Prev) {
       // If a variable from 'pred' is not present in 'result', place it there.
       for (auto& pair : *pred)
         if (result.find(pair.first) == result.end())
diff --git a/test/Analyses/TBR.cpp b/test/Analyses/TBR.cpp
index dec504de3..87a29f971 100644
--- a/test/Analyses/TBR.cpp
+++ b/test/Analyses/TBR.cpp
@@ -11,6 +11,16 @@ double f1(double x) {
   return t;
 } // == x^3
 
+double f2(double val) {
+  double res = 0;
+  for (int i=1; i<5; ++i) {
+    if (i == 3)
+      continue;
+    res += i * val;
+  }
+  return res;
+}
+
 #define TEST(F, x) { \
   result[0] = 0; \
   auto F##grad = clad::gradient<clad::opts::enable_tbr>(F);\
@@ -21,5 +31,5 @@ double f1(double x) {
 int main() {
   double result[3] = {};
   TEST(f1, 3); // CHECK-EXEC: {27.00}
-
+  TEST(f2, 3); // CHECK-EXEC: {9.00}
 }