Add support for SourceLocExpr during differentiation to enable code containing assert and other source-location builtins.

Author: Vedant2005goyal

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -99,6 +99,7 @@ class BaseForwardModeVisitor
   StmtDiff VisitWhileStmt(const clang::WhileStmt* WS);
   StmtDiff VisitDoStmt(const clang::DoStmt* DS);
   StmtDiff VisitContinueStmt(const clang::ContinueStmt* ContStmt);
+  StmtDiff VisitSourceLocExpr(const clang::SourceLocExpr* E);
 
   StmtDiff VisitSwitchStmt(const clang::SwitchStmt* SS);
   StmtDiff VisitBreakStmt(const clang::BreakStmt* BS);

include/clad/Differentiator/ReverseModeVisitor.h

@@ -485,6 +485,7 @@ namespace clad {
     StmtDiff VisitImplicitCastExpr(const clang::ImplicitCastExpr* ICE);
     StmtDiff VisitGNUNullExpr(const clang::GNUNullExpr* E);
     StmtDiff VisitPredefinedExpr(const clang::PredefinedExpr* E);
+    StmtDiff VisitSourceLocExpr(const clang::SourceLocExpr* SLE);
 
 #if CLANG_VERSION_MAJOR > 16
     StmtDiff VisitLambdaExpr(const clang::LambdaExpr* LE);

include/clad/Differentiator/StmtClone.h

@@ -79,6 +79,7 @@ namespace utils {
     DECLARE_CLONE_FN(CXXCatchStmt)
     DECLARE_CLONE_FN(CXXTryStmt)
     DECLARE_CLONE_FN(PredefinedExpr)
+    DECLARE_CLONE_FN(SourceLocExpr)
     DECLARE_CLONE_FN(CharacterLiteral)
     DECLARE_CLONE_FN(FloatingLiteral)
     DECLARE_CLONE_FN(ImaginaryLiteral)

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -581,7 +581,8 @@ StmtDiff BaseForwardModeVisitor::VisitConditionalOperator(
           .ActOnConditionalOp(noLoc, noLoc, cond, ifTrueDiff.getExpr(),
                               ifFalseDiff.getExpr())
           .get();
-
+  if (condExpr->getType()->isVoidType())
+    return StmtDiff(condExpr, nullptr);
   // cond is already used by the value conditional above; clone it for the
   // derivative conditional so the two do not share the stored condition.
   Expr* condExprDiff =
@@ -1874,6 +1875,13 @@ StmtDiff BaseForwardModeVisitor::VisitStringLiteral(const StringLiteral* SL) {
                                  SL->getType(), utils::GetValidSLoc(m_Sema)));
 }
 
+StmtDiff
+BaseForwardModeVisitor::VisitSourceLocExpr(const clang::SourceLocExpr* E) {
+  auto* Constant0 =
+      ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, /*val=*/0);
+  return StmtDiff(CloneNode(E), Constant0);
+}
+
 StmtDiff BaseForwardModeVisitor::VisitWhileStmt(const WhileStmt* WS) {
   // Scope for the whole while loop.
   ScopeRAII whileScope(*this, Scope::ContinueScope | Scope::BreakScope |

lib/Differentiator/ReverseModeVisitor.cpp

@@ -2942,6 +2942,10 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
       for (Stmt* S : revBlock)
         addToCurrentBlock(S, direction::reverse);
       return {cloneE, derivedE};
+    } else if (opCode == UnaryOperatorKind::UO_Extension) {
+      diff = Visit(E);
+      ResultRef = diff.getExpr_dx();
+      valueForRevPass = diff.getRevSweepAsExpr();
     } else {
       if (opCode != UO_LNot)
         // We should only output warnings on visiting boolean conditions
@@ -3904,6 +3908,13 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     return StmtDiff(CloneNode(E), Constant0);
   }
 
+  StmtDiff
+  ReverseModeVisitor::VisitSourceLocExpr(const clang::SourceLocExpr* E) {
+    auto* Constant0 = ConstantFolder::synthesizeLiteral(m_Context.IntTy,
+                                                        m_Context, /*val=*/0);
+    return StmtDiff(CloneNode(E), Constant0);
+  }
+
   StmtDiff ReverseModeVisitor::VisitCXXFunctionalCastExpr(
       const clang::CXXFunctionalCastExpr* FCE) {
     StmtDiff castExprDiff = Visit(FCE->getSubExpr(), dfdx());

lib/Differentiator/StmtClone.cpp

@@ -92,6 +92,16 @@ Stmt* StmtClone::VisitPredefinedExpr(PredefinedExpr* Node) {
   clad_compat::ExprSetDeps(result, Node);
   return result;
 }
+#if CLANG_VERSION_MAJOR >= 15
+DEFINE_CLONE_EXPR(SourceLocExpr,
+                  (Ctx, Node->getIdentKind(), CloneType(Node->getType()),
+                   Node->getBeginLoc(), Node->getEndLoc(),
+                   Node->getParentContext()))
+#else
+DEFINE_CLONE_EXPR(SourceLocExpr,
+                  (Ctx, Node->getIdentKind(), Node->getBeginLoc(),
+                   Node->getEndLoc(), Node->getParentContext()))
+#endif
 DEFINE_CLONE_EXPR(CharacterLiteral,
                   (Node->getValue(), Node->getKind(),
                    CloneType(Node->getType()), Node->getLocation()))

test/Features/DiffAssert.cpp

@@ -0,0 +1,22 @@
+// RUN: %cladclang -std=c++17 -I%S/../../include %s -o %t
+// RUN: %t | %filecheck_exec %s
+
+#include <cassert>
+#include "clad/Differentiator/Differentiator.h"
+
+double func(double x){
+    assert(x>0.0);
+    return x*x*x;
+}
+
+int main(){
+    auto d_func=clad::gradient(func);
+    double dx=0;
+    d_func.execute(3.0,&dx);
+    printf("Diff result: %.2f\n", dx);
+    //CHECK-EXEC: Diff result: 27.00
+    auto d_func1=clad::differentiate(func);
+    auto res=d_func1.execute(3.0);
+    printf("Diff result: %.2f\n", res);
+    //CHECK-EXEC: Diff result: 27.00
+}