Fix Kokkos-typed declarations enter infinite loop

Author: Max Andriychuk

lib/Differentiator/TBRAnalyzer.cpp

@@ -219,10 +219,18 @@ TBRAnalyzer::VarData::VarData(QualType QT, const ASTContext& C,
     const auto* recordDecl = recordType->getDecl();
     auto& newArrMap = m_Val.m_ArrData;
     newArrMap = std::unique_ptr<ArrMap>(new ArrMap());
-    for (const auto* field : recordDecl->fields()) {
-      const auto varType = field->getType();
-      (*newArrMap)[getProfileID(field)] = VarData(varType, C);
-    }
+
+    // FIXME: For some reason if a variable is of kokkos-type we enter an infinite loop.
+    bool isInKokkosNS = false;
+    if (const auto *ns = llvm::dyn_cast<clang::NamespaceDecl>(recordDecl->getDeclContext())) 
+        if (ns->getName() == "Kokkos") 
+            isInKokkosNS = true;
+
+    if(!isInKokkosNS)
+       for (const auto* field : recordDecl->fields()) {
+        const auto varType = field->getType();
+        (*newArrMap)[getProfileID(field)] = VarData(varType, C);
+      }
   }
 }
 
@@ -280,19 +288,28 @@ void TBRAnalyzer::copyVarToCurBlock(const clang::VarDecl* VD) {
 }
 
 void TBRAnalyzer::addVar(const clang::VarDecl* VD, bool forceNonRefType) {
+  llvm::errs() << "\n==============21\n";
   auto& curBranch = getCurBlockVarsData();
 
   QualType varType;
   if (const auto* arrayParam = dyn_cast<ParmVarDecl>(VD))
     varType = arrayParam->getOriginalType();
   else
     varType = VD->getType();
+
   // If varType represents auto or auto*, get the type of init.
   if (utils::IsAutoOrAutoPtrType(varType))
     varType = VD->getInit()->getType();
 
+  llvm::errs() << "\n==============type\n";
+
+varType->dump();
+
   curBranch[VD] =
       VarData(varType, m_ADContext->getASTContext(), forceNonRefType);
+          
+          
+  llvm::errs() << "\n==============22\n";
 }
 
 void TBRAnalyzer::markLocation(const clang::Expr* E) {
@@ -314,6 +331,7 @@ void TBRAnalyzer::setIsRequired(const clang::Expr* E, bool isReq) {
 }
 
 void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
+  llvm::errs() << "\n===========start\n";
   m_BlockData.resize(m_ADContext->getCFG()->size());
   m_BlockPassCounter.resize(m_ADContext->getCFG()->size(), 0);
 
@@ -339,6 +357,7 @@ void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
     addVar(paramsRef[i], /*forceNonRefType=*/true);
   // Add the entry block to the queue.
   m_CFGQueue.insert(m_CurBlockID);
+  llvm::errs() << "\n===========1\n";
 
   // Visit CFG blocks in the queue until it's empty.
   while (!m_CFGQueue.empty()) {
@@ -349,6 +368,8 @@ void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
     CFGBlock& nextBlock = *getCFGBlockByID(m_CurBlockID);
     VisitCFGBlock(nextBlock);
   }
+    llvm::errs() << "\n===========end\n";
+
 #ifndef NDEBUG
   for (int id = m_CurBlockID; id >= 0; --id) {
     LLVM_DEBUG(llvm::dbgs() << "\n-----BLOCK" << id << "-----\n\n");
@@ -584,13 +605,21 @@ bool TBRAnalyzer::VisitDeclRefExpr(DeclRefExpr* DRE) {
 }
 
 bool TBRAnalyzer::VisitDeclStmt(DeclStmt* DS) {
+
   for (auto* D : DS->decls()) {
+          llvm::errs() << "\n===========2\n";
+
     if (auto* VD = dyn_cast<VarDecl>(D)) {
+VD->dump();
       addVar(VD);
+                                  llvm::errs() << "\n===========3\n";
+
       if (clang::Expr* init = VD->getInit()) {
+
         setMode(Mode::kMarkingMode);
         TraverseStmt(init);
         resetMode();
+
         auto& VDExpr = getCurBlockVarsData()[VD];
         // if the declared variable is ref type attach its VarData to the
         // VarData of the RHS variable.
@@ -601,6 +630,7 @@ bool TBRAnalyzer::VisitDeclStmt(DeclStmt* DS) {
           VDExpr.m_Val.m_RefData = ExprsToStore[0];
       }
     }
+
   }
   return true;
 }