Lambda support in the reverse mode

include/clad/Differentiator/CladUtils.h

@@ -158,6 +158,8 @@ namespace clad {
     clang::StringLiteral* CreateStringLiteral(clang::ASTContext& C,
                                               llvm::StringRef str);
 
+    bool isLambdaQType(clang::QualType QT);
+
     /// Returns true if `QT` is Array or Pointer Type, otherwise returns false.
     bool isArrayOrPointerType(clang::QualType QT);
 

include/clad/Differentiator/Compatibility.h

@@ -422,6 +422,16 @@ getConstantArrayType(const ASTContext& Ctx, QualType EltTy,
    #define CLAD_COMPAT_CLANG15_Declarator_DeclarationAttrs_ExtraParam clang::ParsedAttributesView::none(),
 #endif
 
+#if CLANG_VERSION_MAJOR > 12
+#define CLAD_COMPAT_CXXRecordDecl_CreateLambda_DependencyKind(                 \
+    LAMBDACXXRECORDDECL)                                                       \
+  LAMBDACXXRECORDDECL->getLambdaDependencyKind()
+#else
+#define CLAD_COMPAT_CXXRecordDecl_CreateLambda_DependencyKind(                 \
+    LAMBDACXXRECORDDECL)                                                       \
+  LAMBDACXXRECORDDECL->isDependentLambda()
+#endif
+
 // Clang 12 add one extra param (FPO) that we get from Node in Create method of:
 // ImplicitCastExpr, CStyleCastExpr, CXXStaticCastExpr and CXXFunctionalCastExpr
 

include/clad/Differentiator/DiffPlanner.h

@@ -79,6 +79,9 @@ struct DiffRequest {
   /// order derivatives.
   const clang::CXXRecordDecl* Functor = nullptr;
 
+  /// If we're creating a pullback for a lambda function, we need to put the diffed function into a call operator of another lambda class. Therefore, we need to used a pre-created function declaration to put our derived code into. Only applicable if working with lambdas in the reverse mode.
+  clang::FunctionDecl* LambdaPreCreatedDerivativeTarget = nullptr;
+
   /// Stores differentiation parameters information. Stored information
   /// includes info on indices range for array parameters, and nested data
   /// member information for record (class) type parameters.

include/clad/Differentiator/ReverseModeVisitor.h

@@ -357,6 +357,11 @@ namespace clad {
         llvm::SmallVectorImpl<clang::Expr*>& outputArgs,
         clang::Expr* CUDAExecConfig = nullptr);
 
+    clang::CXXRecordDecl*
+    diffLambdaCXXRecordDecl(const clang::CXXRecordDecl* Original);
+    clang::CXXMethodDecl*
+    DifferentiateCallOperatorIfLambda(const clang::CXXRecordDecl* RD, clang::FunctionDecl* preCreatedDerivativeDecl);
+
   public:
     ReverseModeVisitor(DerivativeBuilder& builder, const DiffRequest& request);
     virtual ~ReverseModeVisitor();
@@ -381,8 +386,10 @@ namespace clad {
     /// y" will give 'f_grad_0_1' and "x, z" will give 'f_grad_0_2'.
     DerivativeAndOverload Derive();
     DerivativeAndOverload DerivePullback();
+    DerivativeAndOverload DerivePullback(clang::FunctionDecl* preDerivative);
     StmtDiff VisitArraySubscriptExpr(const clang::ArraySubscriptExpr* ASE);
     StmtDiff VisitBinaryOperator(const clang::BinaryOperator* BinOp);
+    StmtDiff VisitLambdaExpr(const clang::LambdaExpr* LE);
     StmtDiff VisitCallExpr(const clang::CallExpr* CE);
     virtual StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS);
     StmtDiff VisitConditionalOperator(const clang::ConditionalOperator* CO);

include/clad/Differentiator/VisitorBase.h

@@ -293,7 +293,7 @@ namespace clad {
                  clang::Scope* scope, clang::Expr* Init = nullptr,
                  bool DirectInit = false, clang::TypeSourceInfo* TSI = nullptr,
                  clang::VarDecl::InitializationStyle IS =
-                     clang::VarDecl::InitializationStyle::CInit);
+                     clang::VarDecl::InitializationStyle::CInit, bool pushCodeSynthCtxt = false);
     /// Builds variable declaration to be used inside the derivative
     /// body.
     /// \param[in] Type The type of variable declaration to build.
@@ -310,7 +310,7 @@ namespace clad {
                  clang::Expr* Init = nullptr, bool DirectInit = false,
                  clang::TypeSourceInfo* TSI = nullptr,
                  clang::VarDecl::InitializationStyle IS =
-                     clang::VarDecl::InitializationStyle::CInit);
+                     clang::VarDecl::InitializationStyle::CInit, bool pushCodeSynthCtxt = false);
     /// Builds variable declaration to be used inside the derivative
     /// body.
     /// \param[in] Type The type of variable declaration to build.
@@ -326,7 +326,7 @@ namespace clad {
                  clang::Expr* Init = nullptr, bool DirectInit = false,
                  clang::TypeSourceInfo* TSI = nullptr,
                  clang::VarDecl::InitializationStyle IS =
-                     clang::VarDecl::InitializationStyle::CInit);
+                     clang::VarDecl::InitializationStyle::CInit, bool pushCodeSynthCtxt = false);
     /// Builds variable declaration to be used inside the derivative
     /// body in the derivative function global scope.
     clang::VarDecl*

lib/Differentiator/CladUtils.cpp

@@ -316,6 +316,13 @@ namespace clad {
       return false;
     }
 
+    bool isLambdaQType(QualType QT) {
+      if (const RecordType* RT = QT->getAs<RecordType>()) 
+        if (const CXXRecordDecl* RD = dyn_cast<CXXRecordDecl>(RT->getDecl()))
+          return RD->isLambda();
+      return false;
+    }
+
     bool IsReferenceOrPointerArg(const Expr* arg) {
       // The argument is passed by reference if it's passed as an L-value.
       // However, if arg is a MaterializeTemporaryExpr, then arg is a

lib/Differentiator/DerivativeBuilder.cpp

@@ -422,7 +422,7 @@ static void registerDerivative(FunctionDecl* derivedFD, Sema& semaRef) {
         InitErrorEstimation(m_ErrorEstHandler, m_EstModel, *this, request);
         V.AddExternalSource(*m_ErrorEstHandler.back());
       }
-      result = V.DerivePullback();
+      result = V.DerivePullback(request.LambdaPreCreatedDerivativeTarget);
       if (!m_ErrorEstHandler.empty())
         CleanupErrorEstimation(m_ErrorEstHandler, m_EstModel);
     } else if (request.Mode == DiffMode::reverse_mode_forward_pass) {