Implement more descriptive pretty stack printing for generated derivatives.

Author: vgvassilev

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -39,6 +39,16 @@ class BaseForwardModeVisitor
   ///
   DerivativeAndOverload Derive() override;
 
+  StmtDiff Visit(const clang::Stmt* S) {
+    m_CurVisitedStmt = S;
+#ifndef NDEBUG
+    // Enable testing of the pretty printing of the state when clad crashes.
+    if (const char* Env = std::getenv("CLAD_FORCE_CRASH"))
+      std::terminate();
+#endif // NDEBUG
+    return clang::ConstStmtVisitor<BaseForwardModeVisitor, StmtDiff>::Visit(S);
+  }
+
   virtual void ExecuteInsidePushforwardFunctionBlock() {}
 
   virtual StmtDiff

include/clad/Differentiator/ReverseModeVisitor.h

@@ -108,6 +108,13 @@ namespace clad {
       return m_Stack.top();
     }
     StmtDiff Visit(const clang::Stmt* stmt, clang::Expr* dfdS = nullptr) {
+      m_CurVisitedStmt = stmt;
+#ifndef NDEBUG
+      // Enable testing of the pretty printing of the state when clad crashes.
+      if (const char* Env = std::getenv("CLAD_FORCE_CRASH"))
+        std::terminate();
+#endif // NDEBUG
+
       // No need to push the same expr multiple times.
       bool push = !(!m_Stack.empty() && (dfdS == dfdx()));
       if (push)

include/clad/Differentiator/VisitorBase.h

@@ -20,6 +20,7 @@
 #include "clang/Sema/Sema.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/PrettyStackTrace.h"
 
 #include <array>
 #include <stack>
@@ -29,6 +30,10 @@ namespace clang {
 class NestedNameSpecifier;
 } // namespace clang
 
+namespace llvm {
+class raw_ostream;
+} // namespace llvm
+
 namespace clad {
   class MultiplexExternalRMVSource;
   /// A class that represents the result of Visit of ForwardModeVisitor.
@@ -136,6 +141,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 +679,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

@@ -13,6 +13,7 @@
 #include "clad/Differentiator/DiffPlanner.h"
 #include "clad/Differentiator/ErrorEstimator.h"
 #include "clad/Differentiator/ParseDiffArgsTypes.h"
+#include "clad/Differentiator/VisitorBase.h"
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTLambda.h"
@@ -64,6 +65,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

@@ -15,6 +15,7 @@
 #include "clad/Differentiator/ExternalRMVSource.h"
 #include "clad/Differentiator/MultiplexExternalRMVSource.h"
 #include "clad/Differentiator/StmtClone.h"
+#include "clad/Differentiator/VisitorBase.h"
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTLambda.h"
@@ -235,6 +236,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

@@ -22,6 +22,7 @@
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/TemplateBase.h"
 #include "clang/Basic/OperatorKinds.h"
+#include "clang/Basic/SourceManager.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Sema/Lookup.h"
 #include "clang/Sema/Overload.h"
@@ -1072,4 +1073,32 @@ 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();
+      clang::SourceLocation 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 << "\n--- Begin Stmt Dump ---\n";
+      (*m_Stmt)->printPretty(OS, /*Helper=*/nullptr, P);
+      OS << "\n--- End 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

test/Misc/CrashDiags.cpp

@@ -0,0 +1,23 @@
+// RUN: env CLAD_FORCE_CRASH= not %cladclang %s -I%S/../../include 2>&1 | FileCheck %s
+// RUN: env CLAD_FORCE_CRASH= not %cladclang -DREVERSE %s -I%S/../../include 2>&1 | FileCheck %s
+// REQUIRES: asserts
+
+#include "clad/Differentiator/Differentiator.h"
+
+double fn1(double x) {
+  return x * x + 3 * x + 5;
+}
+
+int main() {
+#ifdef REVERSE
+  auto grad = clad::gradient(fn1);
+#else
+  auto dx = clad::differentiate(fn1);
+#endif
+}
+
+// CHECK: Building code for '<double fn1(double x)>[name=fn1, order=1, mode={{.*}}, args='', tbr]'
+// CHECK-NEXT: While visiting <CompoundStmt> [ '
+// CHECK: --- Begin Stmt Dump ---
+// CHECK return x * x + 3 * x + 5;
+// CHECK: --- End Stmt Dump ---