Implement more descriptive pretty stack printing for generated derivatives.

Author: vgvassilev

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -39,6 +39,11 @@ class BaseForwardModeVisitor
   ///
   DerivativeAndOverload Derive() override;
 
+  virtual StmtDiff Visit(const clang::Stmt* S) {
+    m_CurVisitedStmt = S;
+    return clang::ConstStmtVisitor<BaseForwardModeVisitor, StmtDiff>::Visit(S);
+  }
+
   virtual void ExecuteInsidePushforwardFunctionBlock() {}
 
   virtual StmtDiff

include/clad/Differentiator/VisitorBase.h

@@ -136,6 +136,9 @@ namespace clad {
     // expression to be of object type in the reverse mode as well.
     clang::Expr* m_ThisExprDerivative = nullptr;
 
+    /// The currently visited statement. Useful for crash pretty-printing.
+    const clang::Stmt* m_CurVisitedStmt = nullptr;
+
     /// A function used to wrap result of visiting E in a lambda. Returns a call
     /// to the built lambda. Func is a functor that will be invoked inside
     /// lambda scope and block. Statements inside lambda are expected to be
@@ -671,6 +674,23 @@ namespace clad {
     clang::TemplateDecl* m_CladConstructorPushforwardTag = nullptr;
     clang::TemplateDecl* m_CladConstructorReverseForwTag = nullptr;
   };
+
+  /// A class that generates prettier stack traces when we crash on generating
+  /// a derivative.
+  class PrettyStackTraceDerivative : public llvm::PrettyStackTraceEntry {
+    const DiffRequest& m_DiffReq;
+    using Blocks = std::vector<llvm::SmallVector<clang::Stmt*, 16>>;
+    const Blocks& m_Blocks;
+    const clang::Sema& m_Sema;
+    const clang::Stmt** m_Stmt = nullptr;
+
+  public:
+    PrettyStackTraceDerivative(const DiffRequest& DiffReq, const Blocks& B,
+                               const clang::Sema& Sema, const clang::Stmt** S)
+        : m_DiffReq(DiffReq), m_Blocks(B), m_Sema(Sema), m_Stmt(S) {}
+    void print(llvm::raw_ostream& OS) const override;
+  };
+
 } // end namespace clad
 
 #endif // CLAD_VISITOR_BASE_H

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -63,6 +63,9 @@ DerivativeAndOverload BaseForwardModeVisitor::Derive() {
   assert(!m_DerivativeInFlight &&
          "Doesn't support recursive diff. Use DiffPlan.");
 
+  PrettyStackTraceDerivative CrashInfo(m_DiffReq, m_Blocks, m_Sema,
+                                       &m_CurVisitedStmt);
+
   llvm::SaveAndRestore<bool> saveInFlight(m_DerivativeInFlight,
                                           /*NewValue=*/true);
 

lib/Differentiator/ReverseModeVisitor.cpp

@@ -234,6 +234,10 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
 
   DerivativeAndOverload ReverseModeVisitor::Derive() {
     assert(m_DiffReq.Function && "Must not be null.");
+
+    PrettyStackTraceDerivative CrashInfo(m_DiffReq, m_Blocks, m_Sema,
+                                         &m_CurVisitedStmt);
+
     if (m_ExternalSource)
       m_ExternalSource->ActOnStartOfDerive();
     if (m_DiffReq.Mode == DiffMode::error_estimation)

lib/Differentiator/VisitorBase.cpp

@@ -1072,4 +1072,31 @@ namespace clad {
     }
     return nullptr;
   }
+
+  void PrettyStackTraceDerivative::print(llvm::raw_ostream& OS) const {
+    OS << "Building code for '" << (std::string)m_DiffReq << "'\n";
+    clang::PrintingPolicy P(m_Sema.getASTContext().getLangOpts());
+    if (m_Stmt) {
+      const Stmt* S = *m_Stmt;
+      clang::SourceLocation B = S->getBeginLoc(), E = S->getEndLoc();
+      clang::SourceManager& SM = m_Sema.getSourceManager();
+
+      OS << "While visiting <" << S->getStmtClassName() << ">"
+         << " [ '";
+      B.print(OS, SM);
+      OS << "', '";
+      E.print(OS, SM);
+      OS << "']\n";
+
+      OS << "\nBegin Stmt Dump ---\n";
+      (*m_Stmt)->printPretty(OS, /*Helper=*/nullptr, P);
+      OS << "\nEnd Stmt Dump ---\n";
+    }
+
+    if (!m_Blocks.empty() && !m_Blocks.back().empty()) {
+      OS << "Last forward statement ";
+      m_Blocks.back().back()->printPretty(OS, /*Helper=*/nullptr, P);
+    }
+  }
+
 } // end namespace clad