diff --git a/op-finder/OperationFinder.cpp b/op-finder/OperationFinder.cpp
index bc4eef3..3fcc54a 100644
--- a/op-finder/OperationFinder.cpp
+++ b/op-finder/OperationFinder.cpp
@@ -115,26 +115,19 @@ void OperationFinder::processArraySubscript(const clang::ArraySubscriptExpr* sub
   llvm::outs() << "\n";
 }
 
-void OperationFinder::fallthroughBranchEntered()
+void OperationFinder::branchEntered()
 {
-  _in_fallthrough = true;
+  _current_branch++;
+  llvm::outs() << "Branch entered: " << _current_branch << "\n";
 }
 
-void OperationFinder::fallthroughBranchExited()
+void OperationFinder::branchExited()
 {
-  _in_fallthrough = false;
-}
+  llvm::outs() << "Branch exited: " << _current_branch << "\n";
 
-void OperationFinder::forLoopEntered()
-{
-  _for_loop_stack.push_back(_next_for_loop_id);
+  _current_branch--;
 
-  _next_for_loop_id++;
-}
-
-void OperationFinder::forLoopExited()
-{
-  _for_loop_stack.pop_back();
+  assert(_current_branch > -1);
 }
 
 std::unique_ptr<OperationLog::BasicOperation>
@@ -167,8 +160,7 @@ std::pair<std::string, OperationLog> OperationFinder::_createBaseOperationLog(co
 
   OperationLog log;
   log.line = line_number;
-  log.current_for_loops = _for_loop_stack;
-  log.is_fallthrough = _in_fallthrough;
+  log.branch_number = _current_branch;
 
   llvm::outs()  << file_name << ":"
                 << line_number << ":"
diff --git a/op-finder/OperationFinder.hpp b/op-finder/OperationFinder.hpp
index 780771a..9ec60bd 100644
--- a/op-finder/OperationFinder.hpp
+++ b/op-finder/OperationFinder.hpp
@@ -19,11 +19,8 @@ public:
   void processFunctionCall(const clang::CallExpr* call, const clang::SourceManager& source_manager);
   void processArraySubscript(const clang::ArraySubscriptExpr* subscript, const clang::SourceManager& source_manager);
 
-  void fallthroughBranchEntered();
-  void fallthroughBranchExited();
-
-  void forLoopEntered();
-  void forLoopExited();
+  void branchEntered();
+  void branchExited();
 
 private:
   std::unique_ptr<OperationLog::BasicOperation>
@@ -31,9 +28,7 @@ private:
 
   std::pair<std::string, OperationLog> _createBaseOperationLog(const clang::Stmt* stmt, const clang::SourceManager& source_manager);
 
-  int _next_for_loop_id = 0;
-  bool _in_fallthrough = false;
-  std::vector<int> _for_loop_stack;
+  int _current_branch = 0;
   IOperationOutput* _storage;
 };
 
diff --git a/op-finder/OperationFinderAstVisitor.cpp b/op-finder/OperationFinderAstVisitor.cpp
index ce9178d..39833c8 100644
--- a/op-finder/OperationFinderAstVisitor.cpp
+++ b/op-finder/OperationFinderAstVisitor.cpp
@@ -65,17 +65,10 @@ bool OperationFinderAstVisitor::dataTraverseStmtPre(clang::Stmt* stmt)
 {
   assert(_context);
 
-  if (auto* loop = clang::dyn_cast<clang::ForStmt>(stmt))
+  if (clang::Stmt* branch_entry = _isBranchEntry(stmt))
   {
-    if (loop->getInit())
-    {
-      _loop_init = loop->getInit();
-      _op_finder->fallthroughBranchEntered();
-    }
-    else
-    {
-      _op_finder->forLoopEntered();
-    }
+    _branch_stack.push_back(branch_entry);
+    _op_finder->branchEntered();
   }
 
   return true;
@@ -85,16 +78,77 @@ bool OperationFinderAstVisitor::dataTraverseStmtPost(clang::Stmt* stmt)
 {
   assert(_context);
 
-  if (_loop_init && _loop_init == stmt)
+  if (_loop_header.in_loop_header && (_loop_header.init == stmt || _loop_header.header_end == stmt))
   {
-    _op_finder->forLoopEntered();
-    _op_finder->fallthroughBranchExited();
-    _loop_init = nullptr;
+    if (_loop_header.header_end == stmt)
+    {
+      assert(_branch_stack.back() == stmt);
+      _branch_stack.pop_back();
+      _op_finder->branchExited();
+
+      if (_loop_header.init)
+      {
+        assert(_branch_stack.back() == _loop_header.init);
+        _branch_stack.pop_back();
+        _op_finder->branchExited();
+      }
+
+      _loop_header = {};
+    }
+    else if (_loop_header.init == stmt && !_loop_header.header_end)
+    {
+      assert(_branch_stack.back() == stmt);
+      _branch_stack.pop_back();
+      _op_finder->branchExited();
+
+      _loop_header = {};
+    }
   }
-  else if (clang::dyn_cast<clang::ForStmt>(stmt))
+  else if (!_branch_stack.empty() && _branch_stack.back() == stmt)
   {
-    _op_finder->forLoopExited();
+    _branch_stack.pop_back();
+    _op_finder->branchExited();
+  }
+  else
+  {
+    assert(std::find(_branch_stack.cbegin(), _branch_stack.cend(), stmt)
+           == _branch_stack.cend());
   }
 
   return true;
 }
+
+clang::Stmt* OperationFinderAstVisitor::_isBranchEntry(clang::Stmt* stmt)
+{
+  if (auto* loop = clang::dyn_cast<clang::ForStmt>(stmt))
+  {
+    _loop_header.init = loop->getInit();
+
+    if (loop->getCond())
+    {
+      _loop_header.header_start = loop->getCond();
+      _loop_header.header_end = loop->getCond();
+    }
+
+    if (loop->getInc())
+    {
+      if (!_loop_header.header_start)
+        _loop_header.header_start = loop->getInc();
+
+      _loop_header.header_end = loop->getInc();
+    }
+
+    _loop_header.in_loop_header = _loop_header.init || _loop_header.header_end;
+
+    if (_loop_header.init)
+      return _loop_header.init;
+    else if (_loop_header.header_start)
+      return _loop_header.header_end;
+  }
+  else if (_loop_header.in_loop_header && _loop_header.header_start == stmt)
+  {
+    return _loop_header.header_end;
+  }
+
+  return nullptr;
+}
diff --git a/op-finder/OperationFinderAstVisitor.hpp b/op-finder/OperationFinderAstVisitor.hpp
index d82fe46..93e4a45 100644
--- a/op-finder/OperationFinderAstVisitor.hpp
+++ b/op-finder/OperationFinderAstVisitor.hpp
@@ -30,7 +30,19 @@ private:
   clang::ASTContext* _context;
   OperationFinder* _op_finder;
 
-  clang::Stmt* _loop_init = nullptr;
+  struct _LoopHeaderStateMachine
+  {
+    bool in_loop_header = false;
+    clang::Stmt* init = nullptr;
+    clang::Stmt* header_start = nullptr;
+    clang::Stmt* header_end = nullptr;
+  };
+
+  _LoopHeaderStateMachine _loop_header;
+
+  std::vector<clang::Stmt*> _branch_stack;
+
+  clang::Stmt* _isBranchEntry(clang::Stmt* stmt);
 };
 
 #endif //C_ANALYZER_OPERATIONFINDERASTVISITOR_HPP
diff --git a/op-finder/OperationLog.hpp b/op-finder/OperationLog.hpp
index d55c8f7..e8fdc65 100644
--- a/op-finder/OperationLog.hpp
+++ b/op-finder/OperationLog.hpp
@@ -39,13 +39,12 @@ struct OperationLog
     [[nodiscard]] nlohmann::json toJson() const override;
   };
 
-  unsigned int line;
+  unsigned int line = 0;
+  int branch_number = 0;
+
   std::string entry_type;
   std::unique_ptr<IEntry> entry;
 
-  bool is_fallthrough;
-  std::vector<int> current_for_loops;
-
   OperationLog() = default;
   OperationLog(const OperationLog&) = delete;
   OperationLog(OperationLog&&) = default;
@@ -67,8 +66,7 @@ inline void to_json(nlohmann::json& j, const OperationLog& l)
           {"line", l.line},
           {"entry_type", l.entry_type},
           {"entry", l.entry->toJson()},
-          {"is_fallthrough", l.is_fallthrough},
-          {"current_for_loops", l.current_for_loops}
+          {"branch_number", l.branch_number}
   };
 }
 
@@ -78,17 +76,16 @@ inline void from_json(const nlohmann::json& j, OperationLog& l)
   j.at("entry_type").get_to(l.entry_type);
   l.DecodeEntry(j["entry"]);
 
-  j.at("is_fallthrough").get_to(l.is_fallthrough);
-  j.at("current_for_loops").get_to(l.current_for_loops);
+  j.at("branch_number").get_to(l.branch_number);
 }
 
 inline void to_json(nlohmann::json& j, const OperationLog::BasicOperation& bo)
 {
   j = nlohmann::json{
-        {"operation_name", bo.operation_name},
-        {"type_lhs", bo.type_lhs},
-        {"type_rhs", bo.type_rhs},
-        {"type_result", bo.type_result}
+          {"operation_name", bo.operation_name},
+          {"type_lhs", bo.type_lhs},
+          {"type_rhs", bo.type_rhs},
+          {"type_result", bo.type_result}
   };
 }