[clang][Serialization] Don't duplicate the body of LambdaExpr during deserialization

[riccibruno]

Summary

05843dc6ab97a00cbde7aa4f08bf3692eb83109d changed the serialization of the body
of LambdaExpr to avoid a mutation in LambdaExpr::getBody and to avoid a missing
body in LambdaExpr::children.

Unfortunately this replaced one bug by another: we are now duplicating the body during
deserialization; that is after deserialization the identity:

E->getBody() == E->getCallOperator()->getBody() does not hold.

Fix that by instead lazily loading the body from the call operator when needed.

[vabridgers]

I cherry picked this change and retried the case that was failing. This change addresses the issue I was seeing. Thanks.

[martong]

In D83009#2127106, @vabridgers wrote:

I cherry picked this change and retried the case that was failing. This change addresses the issue I was seeing. Thanks.

Yes, I can confirm. Thanks!
Looks good to me!

[aaron.ballman]

LGTM

clang/include/clang/AST/ExprCXX.h

LambdaExpr(QualType T, SourceRange IntroducerRange,
              SourceLocation ClosingBrace, bool ContainsUnexpandedParameterPack);
 
   /// Construct an empty lambda expression.
   LambdaExpr(EmptyShell Empty, unsigned NumCaptures);
 
   Stmt **getStoredStmts() { return getTrailingObjects<Stmt *>(); }
   Stmt *const *getStoredStmts() const { return getTrailingObjects<Stmt *>(); }
 
+  void initBodyIfNeeded() const;
+
 public:
   friend class ASTStmtReader;
   friend class ASTStmtWriter;
   friend TrailingObjects;
 
   /// Construct a new lambda expression.
   static LambdaExpr *
   Create(const ASTContext &C, CXXRecordDecl *Class, SourceRange IntroducerRange,
public:
 
   /// Whether this is a generic lambda.
   bool isGenericLambda() const { return getTemplateParameterList(); }
 
   /// Retrieve the body of the lambda. This will be most of the time
   /// a \p CompoundStmt, but can also be \p CoroutineBodyStmt wrapping
   /// a \p CompoundStmt. Note that unlike functions, lambda-expressions
   /// cannot have a function-try-block.
-  Stmt *getBody() const { return getStoredStmts()[capture_size()]; }
+  Stmt *getBody() const;
 
   /// Retrieve the \p CompoundStmt representing the body of the lambda.
   /// This is a convenience function for callers who do not need
   /// to handle node(s) which may wrap a \p CompoundStmt.
-  const CompoundStmt *getCompoundStmtBody() const {
-    Stmt *Body = getBody();
-    if (const auto *CoroBody = dyn_cast<CoroutineBodyStmt>(Body))
-      return cast<CompoundStmt>(CoroBody->getBody());
-    return cast<CompoundStmt>(Body);
-  }
+  const CompoundStmt *getCompoundStmtBody() const;
   CompoundStmt *getCompoundStmtBody() {
     const auto *ConstThis = this;
     return const_cast<CompoundStmt *>(ConstThis->getCompoundStmtBody());
   }
 
   /// Determine whether the lambda is mutable, meaning that any
   /// captures values can be modified.
   bool isMutable() const;
public:
   }
 
   SourceLocation getBeginLoc() const LLVM_READONLY {
     return IntroducerRange.getBegin();
   }
 
   SourceLocation getEndLoc() const LLVM_READONLY { return ClosingBrace; }
 
-  child_range children() {
-    // Includes initialization exprs plus body stmt
-    return child_range(getStoredStmts(), getStoredStmts() + capture_size() + 1);
-  }
-
-  const_child_range children() const {
-    return const_child_range(getStoredStmts(),
-                             getStoredStmts() + capture_size() + 1);
-  }
+  /// Includes the captures and the body of the lambda.
+  child_range children();
+  const_child_range children() const;
 };
 
 /// An expression "T()" which creates a value-initialized rvalue of type
 /// T, which is a non-class type.  See (C++98 [5.2.3p2]).
 class CXXScalarValueInitExpr : public Expr {
   friend class ASTStmtReader;
 
   TypeSourceInfo *TypeInfo;

clang/lib/AST/ExprCXX.cpp

LambdaExpr::LambdaExpr(QualType T, SourceRange IntroducerRange,
   *Stored++ = getCallOperator()->getBody();
 
   setDependence(computeDependence(this, ContainsUnexpandedParameterPack));
 }
 
 LambdaExpr::LambdaExpr(EmptyShell Empty, unsigned NumCaptures)
     : Expr(LambdaExprClass, Empty) {
   LambdaExprBits.NumCaptures = NumCaptures;
+
+  // Initially don't initialize the body of the LambdaExpr. The body will
+  // be lazily deserialized when needed.
+  getStoredStmts()[NumCaptures] = nullptr; // Not one past the end.
 }
 
 LambdaExpr *LambdaExpr::Create(const ASTContext &Context, CXXRecordDecl *Class,
                                SourceRange IntroducerRange,
                                LambdaCaptureDefault CaptureDefault,
                                SourceLocation CaptureDefaultLoc,
                                bool ExplicitParams, bool ExplicitResultType,
                                ArrayRef<Expr *> CaptureInits,

 
 LambdaExpr *LambdaExpr::CreateDeserialized(const ASTContext &C,
                                            unsigned NumCaptures) {
   unsigned Size = totalSizeToAlloc<Stmt *>(NumCaptures + 1);
   void *Mem = C.Allocate(Size);
   return new (Mem) LambdaExpr(EmptyShell(), NumCaptures);
 }
 
+void LambdaExpr::initBodyIfNeeded() const {
+  if (!getStoredStmts()[capture_size()]) {
+    auto *This = const_cast<LambdaExpr *>(this);
+    This->getStoredStmts()[capture_size()] = getCallOperator()->getBody();
+  }
+}
+
+Stmt *LambdaExpr::getBody() const {
+  initBodyIfNeeded();
+  return getStoredStmts()[capture_size()];
+}
+
+const CompoundStmt *LambdaExpr::getCompoundStmtBody() const {
+  Stmt *Body = getBody();
+  if (const auto *CoroBody = dyn_cast<CoroutineBodyStmt>(Body))
+    return cast<CompoundStmt>(CoroBody->getBody());
+  return cast<CompoundStmt>(Body);
+}
+
 bool LambdaExpr::isInitCapture(const LambdaCapture *C) const {
   return (C->capturesVariable() && C->getCapturedVar()->isInitCapture() &&
           (getCallOperator() == C->getCapturedVar()->getDeclContext()));
 }
 
 LambdaExpr::capture_iterator LambdaExpr::capture_begin() const {
   return getLambdaClass()->getLambdaData().Captures;
 }

 
 ArrayRef<NamedDecl *> LambdaExpr::getExplicitTemplateParameters() const {
   const CXXRecordDecl *Record = getLambdaClass();
   return Record->getLambdaExplicitTemplateParameters();
 }
 
 bool LambdaExpr::isMutable() const { return !getCallOperator()->isConst(); }
 
+LambdaExpr::child_range LambdaExpr::children() {
+  initBodyIfNeeded();
+  return child_range(getStoredStmts(), getStoredStmts() + capture_size() + 1);
+}
+
+LambdaExpr::const_child_range LambdaExpr::children() const {
+  initBodyIfNeeded();
+  return const_child_range(getStoredStmts(),
+                           getStoredStmts() + capture_size() + 1);
+}
+
 ExprWithCleanups::ExprWithCleanups(Expr *subexpr,
                                    bool CleanupsHaveSideEffects,
                                    ArrayRef<CleanupObject> objects)
     : FullExpr(ExprWithCleanupsClass, subexpr) {
   ExprWithCleanupsBits.CleanupsHaveSideEffects = CleanupsHaveSideEffects;
   ExprWithCleanupsBits.NumObjects = objects.size();
   for (unsigned i = 0, e = objects.size(); i != e; ++i)
     getTrailingObjects<CleanupObject>()[i] = objects[i];

clang/lib/Serialization/ASTReaderStmt.cpp

void ASTStmtReader::VisitLambdaExpr(LambdaExpr *E) {
   E->LambdaExprBits.CaptureDefault = Record.readInt();
   E->CaptureDefaultLoc = readSourceLocation();
   E->LambdaExprBits.ExplicitParams = Record.readInt();
   E->LambdaExprBits.ExplicitResultType = Record.readInt();
   E->ClosingBrace = readSourceLocation();
 
   // Read capture initializers.
   for (LambdaExpr::capture_init_iterator C = E->capture_init_begin(),
                                          CEnd = E->capture_init_end();
        C != CEnd; ++C)
     *C = Record.readSubExpr();
 
-  // Ok, not one past the end.
-  E->getStoredStmts()[NumCaptures] = Record.readSubStmt();
+  // The body will be lazily deserialized when needed from the call operator
+  // declaration.
 }
 
 void
 ASTStmtReader::VisitCXXStdInitializerListExpr(CXXStdInitializerListExpr *E) {
   VisitExpr(E);
   E->SubExpr = Record.readSubExpr();
 }
 

clang/lib/Serialization/ASTWriterStmt.cpp

void ASTStmtWriter::VisitLambdaExpr(LambdaExpr *E) {
 
   // Add capture initializers.
   for (LambdaExpr::capture_init_iterator C = E->capture_init_begin(),
                                       CEnd = E->capture_init_end();
        C != CEnd; ++C) {
     Record.AddStmt(*C);
   }
 
-  Record.AddStmt(E->getBody());
+  // Don't serialize the body. It belongs to the call operator declaration.
+  // LambdaExpr only stores a copy of the Stmt *.
 
   Code = serialization::EXPR_LAMBDA;
 }
 
 void ASTStmtWriter::VisitCXXStdInitializerListExpr(CXXStdInitializerListExpr *E) {
   VisitExpr(E);
   Record.AddStmt(E->getSubExpr());
   Code = serialization::EXPR_CXX_STD_INITIALIZER_LIST;

clang/test/AST/ast-dump-lambda-body-not-duplicated.cpp

+// Test without serialization:
+// RUN: %clang_cc1 -triple x86_64-unknown-unknown -Wno-unused-value -ast-dump %s \
+// RUN: | FileCheck %s
+//
+// Test with serialization:
+// RUN: %clang_cc1 -triple x86_64-unknown-unknown -Wno-unused-value -emit-pch -o %t %s
+// RUN: %clang_cc1 -x c++ -triple x86_64-unknown-unknown -Wno-unused-value \
+// RUN: -include-pch %t -ast-dump-all /dev/null \
+// RUN: | FileCheck %s
+
+// Make sure that the Stmt * for the body of the LambdaExpr is
+// equal to the Stmt * for the body of the call operator.
+void Test0() {
+  []() {
+    return 42;
+  };
+}
+
+// CHECK: FunctionDecl {{.*}} Test0
+//
+// CHECK: CXXMethodDecl {{.*}} operator() 'int () const' inline
+// CHECK-NEXT: CompoundStmt 0x[[TMP0:.*]]
+// CHECK: IntegerLiteral {{.*}} 'int' 42
+//
+// CHECK: CompoundStmt 0x[[TMP0]]
+// Check: IntegerLiteral {{.*}} 'int' 42
+
+void Test1() {
+  [](auto x) { return x; };
+}
+
+// CHECK: FunctionDecl {{.*}} Test1
+//
+// CHECK: CXXMethodDecl {{.*}} operator() 'auto (auto) const' inline
+// CHECK-NEXT: ParmVarDecl {{.*}} referenced x 'auto'
+// CHECK-NEXT: CompoundStmt 0x[[TMP1:.*]]
+// CHECK: DeclRefExpr {{.*}} 'x' 'auto'
+//
+// CHECK: CompoundStmt 0x[[TMP1]]
+// CHECK: DeclRefExpr {{.*}} 'x' 'auto'