Allow in-place compound statment body replacement

[v.g.vassilev]

Summary

CompoundStmt::setStmts was reimplemented to save a pointer (r321429/llvm-mirror/clang@d0ee47a).

The patch assumes allocation of a new CompoundStmt if its statements are modified. Unfortunately this needs the newly created CompoundStmt swapped in its users making the AST visitation harder. This break external tools such as cling which alters CompoundStmts to implement its interactive extensions.

This patch patch brings the old facilities only when the new interface CompoundStmt::replaceStmts is called.

[v.g.vassilev]

Thanks to @Hahnfeld, that's not needed anymore.

include/clang/AST/Stmt.h

protected:
   };
 
   class CompoundStmtBitfields {
     friend class ASTStmtReader;
     friend class CompoundStmt;
 
     unsigned : NumStmtBits;
 
-    unsigned NumStmts : 32 - NumStmtBits;
+    unsigned WasReplaced : 1;
+    unsigned NumStmts : 32 - (NumStmtBits + 1);
 
     /// The location of the opening "{".
     SourceLocation LBraceLoc;
   };
 
   class LabelStmtBitfields {
     friend class LabelStmt;
 

 public:
   static CompoundStmt *Create(const ASTContext &C, ArrayRef<Stmt *> Stmts,
                               SourceLocation LB, SourceLocation RB);
 
   // Build an empty compound statement with a location.
   explicit CompoundStmt(SourceLocation Loc)
       : Stmt(CompoundStmtClass), RBraceLoc(Loc) {
     CompoundStmtBits.NumStmts = 0;
+    CompoundStmtBits.WasReplaced = 0;
     CompoundStmtBits.LBraceLoc = Loc;
   }
 
   // Build an empty compound statement.
   static CompoundStmt *CreateEmpty(const ASTContext &C, unsigned NumStmts);
 
   bool body_empty() const { return CompoundStmtBits.NumStmts == 0; }
   unsigned size() const { return CompoundStmtBits.NumStmts; }
 
   using body_iterator = Stmt **;
   using body_range = llvm::iterator_range<body_iterator>;
 
   body_range body() { return body_range(body_begin(), body_end()); }
-  body_iterator body_begin() { return getTrailingObjects<Stmt *>(); }
+  body_iterator body_begin() {
+    Stmt** begin = getTrailingObjects<Stmt *>();
+    return !CompoundStmtBits.WasReplaced ? begin : (body_iterator)begin[0];
+  }
   body_iterator body_end() { return body_begin() + size(); }
   Stmt *body_front() { return !body_empty() ? body_begin()[0] : nullptr; }
 
   Stmt *body_back() {
     return !body_empty() ? body_begin()[size() - 1] : nullptr;
   }
 
   using const_body_iterator = Stmt *const *;
   using body_const_range = llvm::iterator_range<const_body_iterator>;
 
   body_const_range body() const {
     return body_const_range(body_begin(), body_end());
   }
 
   const_body_iterator body_begin() const {
-    return getTrailingObjects<Stmt *>();
+    return const_cast<CompoundStmt*>(this)->body_begin();
   }
 
   const_body_iterator body_end() const { return body_begin() + size(); }
 
   const Stmt *body_front() const {
     return !body_empty() ? body_begin()[0] : nullptr;
   }
 
public:
   const_reverse_body_iterator body_rbegin() const {
     return const_reverse_body_iterator(body_end());
   }
 
   const_reverse_body_iterator body_rend() const {
     return const_reverse_body_iterator(body_begin());
   }
 
+  void replaceStmts(const ASTContext &C, llvm::ArrayRef<Stmt*> Stmts);
+
   // Get the Stmt that StmtExpr would consider to be the result of this
   // compound statement. This is used by StmtExpr to properly emulate the GCC
   // compound expression extension, which ignores trailing NullStmts when
   // getting the result of the expression.
   // i.e. ({ 5;;; })
   //           ^^ ignored
   // If we don't find something that isn't a NullStmt, just return the last
   // Stmt.

lib/AST/Stmt.cpp

 int64_t Stmt::getID(const ASTContext &Context) const {
   return Context.getAllocator().identifyKnownAlignedObject<Stmt>(this);
 }
 
 CompoundStmt::CompoundStmt(ArrayRef<Stmt *> Stmts, SourceLocation LB,
                            SourceLocation RB)
     : Stmt(CompoundStmtClass), RBraceLoc(RB) {
   CompoundStmtBits.NumStmts = Stmts.size();
+  CompoundStmtBits.WasReplaced = 0;
   setStmts(Stmts);
   CompoundStmtBits.LBraceLoc = LB;
 }
 
 void CompoundStmt::setStmts(ArrayRef<Stmt *> Stmts) {
   assert(CompoundStmtBits.NumStmts == Stmts.size() &&
          "NumStmts doesn't fit in bits of CompoundStmtBits.NumStmts!");
-
+  assert(!CompoundStmtBits.WasReplaced && "Call replaceStmts!");
   std::copy(Stmts.begin(), Stmts.end(), body_begin());
 }
 
 CompoundStmt *CompoundStmt::Create(const ASTContext &C, ArrayRef<Stmt *> Stmts,
                                    SourceLocation LB, SourceLocation RB) {
   void *Mem =
       C.Allocate(totalSizeToAlloc<Stmt *>(Stmts.size()), alignof(CompoundStmt));
   return new (Mem) CompoundStmt(Stmts, LB, RB);
 }
 
 CompoundStmt *CompoundStmt::CreateEmpty(const ASTContext &C,
                                         unsigned NumStmts) {
   void *Mem =
       C.Allocate(totalSizeToAlloc<Stmt *>(NumStmts), alignof(CompoundStmt));
   CompoundStmt *New = new (Mem) CompoundStmt(EmptyShell());
+  New->CompoundStmtBits.WasReplaced = 0;
   New->CompoundStmtBits.NumStmts = NumStmts;
   return New;
 }
 
+void CompoundStmt::replaceStmts(const ASTContext &C,
+                                llvm::ArrayRef<Stmt*> Stmts) {
+  Stmt** Body = body_begin();
+
+  if (CompoundStmtBits.WasReplaced)
+    C.Deallocate(Body);
+  else
+    memset(body_begin(), 0, size());
+
+  CompoundStmtBits.NumStmts = Stmts.size();
+  assert(CompoundStmtBits.NumStmts == Stmts.size() &&
+         "NumStmts doesn't fit in bits of CompoundStmtBits.NumStmts!");
+
+  Body = new (C) Stmt*[Stmts.size()];
+  std::copy(Stmts.begin(), Stmts.end(), Body);
+
+  getTrailingObjects<Stmt *>()[0] = reinterpret_cast<Stmt*>(Body);
+
+  CompoundStmtBits.WasReplaced = 1;
+}
+
 const Expr *ValueStmt::getExprStmt() const {
   const Stmt *S = this;
   do {
     if (const auto *E = dyn_cast<Expr>(S))
       return E;
 
     if (const auto *LS = dyn_cast<LabelStmt>(S))
       S = LS->getSubStmt();

unittests/AST/CMakeLists.txt

add_clang_unittest(ASTTests
   DeclTest.cpp
   EvaluateAsRValueTest.cpp
   ExternalASTSourceTest.cpp
   Language.cpp
   NamedDeclPrinterTest.cpp
   OMPStructuredBlockTest.cpp
   SourceLocationTest.cpp
   StmtPrinterTest.cpp
+  StmtTest.cpp
   StructuralEquivalenceTest.cpp
   )
 
 clang_target_link_libraries(ASTTests
   PRIVATE
   clangAST
   clangASTMatchers
   clangBasic
   clangFrontend
   clangSerialization
   clangTooling
   )

unittests/AST/StmtTest.cpp

+//===- unittests/AST/StmtTest.cpp --- Declaration tests -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Unit tests for Stmt nodes in the AST.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/Frontend/ASTUnit.h"
+#include "clang/Tooling/Tooling.h"
+
+#include "DeclMatcher.h"
+
+#include "gtest/gtest.h"
+
+#include <string>
+#include <vector>
+
+using namespace clang;
+using namespace clang::ast_matchers;
+using namespace clang::tooling;
+
+
+TEST(Stmt, CompoundStmtReplaceStmts) {
+  const char *const InputFileName = "input.cc";
+  std::string Code = "void f() { int i = 1; } void g() { float f = 2.; }";
+  std::vector<std::string> Args;
+  std::unique_ptr<ASTUnit> AST =
+    tooling::buildASTFromCodeWithArgs(Code, Args, InputFileName);
+
+  ASTContext &C = AST->getASTContext();
+  TranslationUnitDecl *TU = C.getTranslationUnitDecl();
+  auto MakeMatcher = [](const char* Id) { return functionDecl(hasName(Id)); };
+
+  auto FDf = FirstDeclMatcher<FunctionDecl>().match(TU, MakeMatcher("f"));
+  auto FDg = FirstDeclMatcher<FunctionDecl>().match(TU, MakeMatcher("g"));
+
+  auto CSf = cast<CompoundStmt>(FDf->getBody());
+  auto CSg = cast<CompoundStmt>(FDg->getBody());
+
+  llvm::SmallVector<Stmt*, 2> Stmts;
+  Stmts.push_back(CSf->body_front());
+  // FIXME: Clone the CSg content to keep the AST invariants.
+  Stmts.push_back(CSg->body_front());
+  CSf->replaceStmts(C, Stmts);
+
+  ASSERT_TRUE(CSf->size() == 2);
+}