diff --git a/bolt/include/bolt/Profile/DataAggregator.h b/bolt/include/bolt/Profile/DataAggregator.h
--- a/bolt/include/bolt/Profile/DataAggregator.h
+++ b/bolt/include/bolt/Profile/DataAggregator.h
@@ -199,10 +199,10 @@
   /// execution order.
   ///
   /// Return true if the trace is valid, false otherwise.
-  bool recordTrace(
-      BinaryFunction &BF, const LBREntry &First, const LBREntry &Second,
-      uint64_t Count = 1,
-      SmallVector<std::pair<uint64_t, uint64_t>, 16> *Branches = nullptr) const;
+  bool
+  recordTrace(BinaryFunction &BF, const LBREntry &First, const LBREntry &Second,
+              uint64_t Count,
+              SmallVector<std::pair<uint64_t, uint64_t>, 16> &Branches) const;
 
   /// Return a vector of offsets corresponding to a trace in a function
   /// (see recordTrace() above).
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -838,11 +838,9 @@
 }
 
 bool DataAggregator::recordTrace(
-    BinaryFunction &BF,
-    const LBREntry &FirstLBR,
-    const LBREntry &SecondLBR,
+    BinaryFunction &BF, const LBREntry &FirstLBR, const LBREntry &SecondLBR,
     uint64_t Count,
-    SmallVector<std::pair<uint64_t, uint64_t>, 16> *Branches) const {
+    SmallVector<std::pair<uint64_t, uint64_t>, 16> &Branches) const {
   BinaryContext &BC = BF.getBinaryContext();
 
   if (!BF.isSimple())
@@ -902,24 +900,27 @@
       return false;
     }
 
-    // Record fall-through jumps
-    BinaryBasicBlock::BinaryBranchInfo &BI = BB->getBranchInfo(*NextBB);
-    BI.Count += Count;
-
-    if (Branches) {
-      const MCInst *Instr = BB->getLastNonPseudoInstr();
-      uint64_t Offset = 0;
-      if (Instr)
-        Offset = BC.MIB->getOffsetWithDefault(*Instr, 0);
-      else
-        Offset = BB->getOffset();
+    const MCInst *Instr = BB->getLastNonPseudoInstr();
+    uint64_t Offset = 0;
+    if (Instr)
+      Offset = BC.MIB->getOffsetWithDefault(*Instr, 0);
+    else
+      Offset = BB->getOffset();
 
-      Branches->emplace_back(Offset, NextBB->getOffset());
-    }
+    Branches.emplace_back(Offset, NextBB->getOffset());
 
     BB = NextBB;
   }
 
+  // Record fall-through jumps
+  for (const auto &[FromOffset, ToOffset] : Branches) {
+    BinaryBasicBlock *FromBB = BF.getBasicBlockContainingOffset(FromOffset);
+    BinaryBasicBlock *ToBB = BF.getBasicBlockAtOffset(ToOffset);
+    assert(FromBB && ToBB);
+    BinaryBasicBlock::BinaryBranchInfo &BI = FromBB->getBranchInfo(*ToBB);
+    BI.Count += Count;
+  }
+
   return true;
 }
 
@@ -930,7 +931,7 @@
                                        uint64_t Count) const {
   SmallVector<std::pair<uint64_t, uint64_t>, 16> Res;
 
-  if (!recordTrace(BF, FirstLBR, SecondLBR, Count, &Res))
+  if (!recordTrace(BF, FirstLBR, SecondLBR, Count, Res))
     return std::nullopt;
 
   return Res;