From b6f8f15bae323c817625f476726ec84afaf70718 Mon Sep 17 00:00:00 2001
From: Vaibhav Thakkar <vaibhav.thakkar.22.12.99@gmail.com>
Date: Mon, 20 May 2024 00:26:58 +0200
Subject: [PATCH] Fix recursive processing of DiffRequests

---
 include/clad/Differentiator/DynamicGraph.h |  4 ++++
 tools/ClangPlugin.cpp                      | 17 +++++++++++------
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/include/clad/Differentiator/DynamicGraph.h b/include/clad/Differentiator/DynamicGraph.h
index ae84e0b73..55e5ad350 100644
--- a/include/clad/Differentiator/DynamicGraph.h
+++ b/include/clad/Differentiator/DynamicGraph.h
@@ -98,6 +98,10 @@ template <typename T> class DynamicGraph {
     m_currentId = -1;
   }
 
+  /// Check if currently processing a node.
+  /// \returns True if currently processing a node, false otherwise.
+  bool isProcessingNode() { return m_currentId != -1; }
+
   /// Get the nodes in the graph.
   std::vector<T> getNodes() { return m_nodes; }
 
diff --git a/tools/ClangPlugin.cpp b/tools/ClangPlugin.cpp
index c667ddd2c..062afcc70 100644
--- a/tools/ClangPlugin.cpp
+++ b/tools/ClangPlugin.cpp
@@ -407,12 +407,17 @@ namespace clad {
       Sema::GlobalEagerInstantiationScope GlobalInstantiations(S, Enabled);
       Sema::LocalEagerInstantiationScope LocalInstantiations(S);
 
-      DiffRequest request = m_DiffRequestGraph.getNextToProcessNode();
-      while (request.Function != nullptr) {
-        m_DiffRequestGraph.setCurrentProcessingNode(request);
-        ProcessDiffRequest(request);
-        m_DiffRequestGraph.markCurrentNodeProcessed();
-        request = m_DiffRequestGraph.getNextToProcessNode();
+      if (!m_DiffRequestGraph.isProcessingNode()) {
+        // This check is to avoid recursive processing of the graph, as
+        // HandleTopLevelDecl can be called recursively in non-standard
+        // setup for code generation.
+        DiffRequest request = m_DiffRequestGraph.getNextToProcessNode();
+        while (request.Function != nullptr) {
+          m_DiffRequestGraph.setCurrentProcessingNode(request);
+          ProcessDiffRequest(request);
+          m_DiffRequestGraph.markCurrentNodeProcessed();
+          request = m_DiffRequestGraph.getNextToProcessNode();
+        }
       }
 
       // Put the TUScope in a consistent state after clad is done.