Reapply "[clang][Interp] Implement dynamic memory allocation handling (#70306)"

This reverts commit 48d703e.

Author: tbaederr

clang/lib/AST/CMakeLists.txt

@@ -75,6 +75,7 @@ add_clang_library(clangAST
   Interp/InterpBuiltin.cpp
   Interp/Floating.cpp
   Interp/EvaluationResult.cpp
+  Interp/DynamicAllocator.cpp
   Interp/Interp.cpp
   Interp/InterpBlock.cpp
   Interp/InterpFrame.cpp

clang/lib/AST/Interp/Compiler.cpp

@@ -2771,6 +2771,117 @@ bool Compiler<Emitter>::VisitCXXInheritedCtorInitExpr(
   return this->emitCall(F, 0, E);
 }
 
+template <class Emitter>
+bool Compiler<Emitter>::VisitCXXNewExpr(const CXXNewExpr *E) {
+  assert(classifyPrim(E->getType()) == PT_Ptr);
+  const Expr *Init = E->getInitializer();
+  QualType ElementType = E->getAllocatedType();
+  std::optional<PrimType> ElemT = classify(ElementType);
+  unsigned PlacementArgs = E->getNumPlacementArgs();
+  bool IsNoThrow = false;
+
+  // FIXME: Better diagnostic. diag::note_constexpr_new_placement
+  if (PlacementArgs != 0) {
+    // The only new-placement list we support is of the form (std::nothrow).
+    //
+    // FIXME: There is no restriction on this, but it's not clear that any
+    // other form makes any sense. We get here for cases such as:
+    //
+    //   new (std::align_val_t{N}) X(int)
+    //
+    // (which should presumably be valid only if N is a multiple of
+    // alignof(int), and in any case can't be deallocated unless N is
+    // alignof(X) and X has new-extended alignment).
+    if (PlacementArgs != 1 || !E->getPlacementArg(0)->getType()->isNothrowT())
+      return this->emitInvalid(E);
+
+    if (!this->discard(E->getPlacementArg(0)))
+      return false;
+    IsNoThrow = true;
+  }
+
+  const Descriptor *Desc;
+  if (ElemT) {
+    if (E->isArray())
+      Desc = nullptr; // We're not going to use it in this case.
+    else
+      Desc = P.createDescriptor(E, *ElemT, Descriptor::InlineDescMD,
+                                /*IsConst=*/false, /*IsTemporary=*/false,
+                                /*IsMutable=*/false);
+  } else {
+    Desc = P.createDescriptor(
+        E, ElementType.getTypePtr(),
+        E->isArray() ? std::nullopt : Descriptor::InlineDescMD,
+        /*IsConst=*/false, /*IsTemporary=*/false, /*IsMutable=*/false, Init);
+  }
+
+  if (E->isArray()) {
+    std::optional<const Expr *> ArraySizeExpr = E->getArraySize();
+    if (!ArraySizeExpr)
+      return false;
+
+    const Expr *Stripped = *ArraySizeExpr;
+    for (; auto *ICE = dyn_cast<ImplicitCastExpr>(Stripped);
+         Stripped = ICE->getSubExpr())
+      if (ICE->getCastKind() != CK_NoOp &&
+          ICE->getCastKind() != CK_IntegralCast)
+        break;
+
+    PrimType SizeT = classifyPrim(Stripped->getType());
+
+    if (!this->visit(Stripped))
+      return false;
+
+    if (ElemT) {
+      // N primitive elements.
+      if (!this->emitAllocN(SizeT, *ElemT, E, IsNoThrow, E))
+        return false;
+    } else {
+      // N Composite elements.
+      if (!this->emitAllocCN(SizeT, Desc, IsNoThrow, E))
+        return false;
+    }
+
+    if (Init && !this->visitInitializer(Init))
+      return false;
+
+  } else {
+    // Allocate just one element.
+    if (!this->emitAlloc(Desc, E))
+      return false;
+
+    if (Init) {
+      if (ElemT) {
+        if (!this->visit(Init))
+          return false;
+
+        if (!this->emitInit(*ElemT, E))
+          return false;
+      } else {
+        // Composite.
+        if (!this->visitInitializer(Init))
+          return false;
+      }
+    }
+  }
+
+  if (DiscardResult)
+    return this->emitPopPtr(E);
+
+  return true;
+}
+
+template <class Emitter>
+bool Compiler<Emitter>::VisitCXXDeleteExpr(const CXXDeleteExpr *E) {
+  const Expr *Arg = E->getArgument();
+
+  // Arg must be an lvalue.
+  if (!this->visit(Arg))
+    return false;
+
+  return this->emitFree(E->isArrayForm(), E);
+}
+
 template <class Emitter>
 bool Compiler<Emitter>::VisitExpressionTraitExpr(const ExpressionTraitExpr *E) {
   assert(Ctx.getLangOpts().CPlusPlus);

clang/lib/AST/Interp/Compiler.h

@@ -190,6 +190,8 @@ class Compiler : public ConstStmtVisitor<Compiler<Emitter>, bool>,
   bool VisitObjCBoxedExpr(const ObjCBoxedExpr *E);
   bool VisitCXXStdInitializerListExpr(const CXXStdInitializerListExpr *E);
   bool VisitStmtExpr(const StmtExpr *E);
+  bool VisitCXXNewExpr(const CXXNewExpr *E);
+  bool VisitCXXDeleteExpr(const CXXDeleteExpr *E);
 
   // Statements.
   bool visitCompoundStmt(const CompoundStmt *S);

clang/lib/AST/Interp/DynamicAllocator.cpp

@@ -0,0 +1,118 @@
+//==-------- DynamicAllocator.cpp - Dynamic allocations ----------*- C++ -*-==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "DynamicAllocator.h"
+#include "InterpBlock.h"
+#include "InterpState.h"
+
+using namespace clang;
+using namespace clang::interp;
+
+DynamicAllocator::~DynamicAllocator() { cleanup(); }
+
+void DynamicAllocator::cleanup() {
+  // Invoke destructors of all the blocks and as a last restort,
+  // reset all the pointers pointing to them to null pointees.
+  // This should never show up in diagnostics, but it's necessary
+  // for us to not cause use-after-free problems.
+  for (auto &Iter : AllocationSites) {
+    auto &AllocSite = Iter.second;
+    for (auto &Alloc : AllocSite.Allocations) {
+      Block *B = reinterpret_cast<Block *>(Alloc.Memory.get());
+      B->invokeDtor();
+      if (B->hasPointers()) {
+        while (B->Pointers) {
+          Pointer *Next = B->Pointers->Next;
+          B->Pointers->PointeeStorage.BS.Pointee = nullptr;
+          B->Pointers = Next;
+        }
+        B->Pointers = nullptr;
+      }
+    }
+  }
+
+  AllocationSites.clear();
+}
+
+Block *DynamicAllocator::allocate(const Expr *Source, PrimType T,
+                                  size_t NumElements, unsigned EvalID) {
+  // Create a new descriptor for an array of the specified size and
+  // element type.
+  const Descriptor *D = allocateDescriptor(
+      Source, T, Descriptor::InlineDescMD, NumElements, /*IsConst=*/false,
+      /*IsTemporary=*/false, /*IsMutable=*/false);
+
+  return allocate(D, EvalID);
+}
+
+Block *DynamicAllocator::allocate(const Descriptor *ElementDesc,
+                                  size_t NumElements, unsigned EvalID) {
+  // Create a new descriptor for an array of the specified size and
+  // element type.
+  const Descriptor *D = allocateDescriptor(
+      ElementDesc->asExpr(), ElementDesc, Descriptor::InlineDescMD, NumElements,
+      /*IsConst=*/false, /*IsTemporary=*/false, /*IsMutable=*/false);
+  return allocate(D, EvalID);
+}
+
+Block *DynamicAllocator::allocate(const Descriptor *D, unsigned EvalID) {
+  assert(D);
+  assert(D->asExpr());
+
+  auto Memory =
+      std::make_unique<std::byte[]>(sizeof(Block) + D->getAllocSize());
+  auto *B = new (Memory.get()) Block(EvalID, D, /*isStatic=*/false);
+  B->invokeCtor();
+
+  InlineDescriptor *ID = reinterpret_cast<InlineDescriptor *>(B->rawData());
+  ID->Desc = D;
+  ID->IsActive = true;
+  ID->Offset = sizeof(InlineDescriptor);
+  ID->IsBase = false;
+  ID->IsFieldMutable = false;
+  ID->IsConst = false;
+  ID->IsInitialized = false;
+
+  B->IsDynamic = true;
+
+  if (auto It = AllocationSites.find(D->asExpr()); It != AllocationSites.end())
+    It->second.Allocations.emplace_back(std::move(Memory));
+  else
+    AllocationSites.insert(
+        {D->asExpr(), AllocationSite(std::move(Memory), D->isArray())});
+  return B;
+}
+
+bool DynamicAllocator::deallocate(const Expr *Source,
+                                  const Block *BlockToDelete, InterpState &S) {
+  auto It = AllocationSites.find(Source);
+  if (It == AllocationSites.end())
+    return false;
+
+  auto &Site = It->second;
+  assert(Site.size() > 0);
+
+  // Find the Block to delete.
+  auto AllocIt = llvm::find_if(Site.Allocations, [&](const Allocation &A) {
+    const Block *B = reinterpret_cast<const Block *>(A.Memory.get());
+    return BlockToDelete == B;
+  });
+
+  assert(AllocIt != Site.Allocations.end());
+
+  Block *B = reinterpret_cast<Block *>(AllocIt->Memory.get());
+  B->invokeDtor();
+
+  S.deallocate(B);
+  Site.Allocations.erase(AllocIt);
+
+  if (Site.size() == 0)
+    AllocationSites.erase(It);
+
+  return true;
+}

clang/lib/AST/Interp/DynamicAllocator.h

@@ -0,0 +1,102 @@
+//==--------- DynamicAllocator.h - Dynamic allocations ------------*- C++ -*-=//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_AST_INTERP_DYNAMIC_ALLOCATOR_H
+#define LLVM_CLANG_AST_INTERP_DYNAMIC_ALLOCATOR_H
+
+#include "Descriptor.h"
+#include "InterpBlock.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/Support/Allocator.h"
+
+namespace clang {
+class Expr;
+namespace interp {
+class Block;
+class InterpState;
+
+/// Manages dynamic memory allocations done during bytecode interpretation.
+///
+/// We manage allocations as a map from their new-expression to a list
+/// of allocations. This is called an AllocationSite. For each site, we
+/// record whether it was allocated using new or new[], the
+/// IsArrayAllocation flag.
+///
+/// For all array allocations, we need to allocate new Descriptor instances,
+/// so the DynamicAllocator has a llvm::BumpPtrAllocator similar to Program.
+class DynamicAllocator final {
+  struct Allocation {
+    std::unique_ptr<std::byte[]> Memory;
+    Allocation(std::unique_ptr<std::byte[]> Memory)
+        : Memory(std::move(Memory)) {}
+  };
+
+  struct AllocationSite {
+    llvm::SmallVector<Allocation> Allocations;
+    bool IsArrayAllocation = false;
+
+    AllocationSite(std::unique_ptr<std::byte[]> Memory, bool Array)
+        : IsArrayAllocation(Array) {
+      Allocations.push_back({std::move(Memory)});
+    }
+
+    size_t size() const { return Allocations.size(); }
+  };
+
+public:
+  DynamicAllocator() = default;
+  ~DynamicAllocator();
+
+  void cleanup();
+
+  unsigned getNumAllocations() const { return AllocationSites.size(); }
+
+  /// Allocate ONE element of the given descriptor.
+  Block *allocate(const Descriptor *D, unsigned EvalID);
+  /// Allocate \p NumElements primitive elements of the given type.
+  Block *allocate(const Expr *Source, PrimType T, size_t NumElements,
+                  unsigned EvalID);
+  /// Allocate \p NumElements elements of the given descriptor.
+  Block *allocate(const Descriptor *D, size_t NumElements, unsigned EvalID);
+
+  /// Deallocate the given source+block combination.
+  /// Returns \c true if anything has been deallocatd, \c false otherwise.
+  bool deallocate(const Expr *Source, const Block *BlockToDelete,
+                  InterpState &S);
+
+  /// Checks whether the allocation done at the given source is an array
+  /// allocation.
+  bool isArrayAllocation(const Expr *Source) const {
+    if (auto It = AllocationSites.find(Source); It != AllocationSites.end())
+      return It->second.IsArrayAllocation;
+    return false;
+  }
+
+  /// Allocation site iterator.
+  using const_virtual_iter =
+      llvm::DenseMap<const Expr *, AllocationSite>::const_iterator;
+  llvm::iterator_range<const_virtual_iter> allocation_sites() const {
+    return llvm::make_range(AllocationSites.begin(), AllocationSites.end());
+  }
+
+private:
+  llvm::DenseMap<const Expr *, AllocationSite> AllocationSites;
+
+  using PoolAllocTy = llvm::BumpPtrAllocatorImpl<llvm::MallocAllocator>;
+  PoolAllocTy DescAllocator;
+
+  /// Allocates a new descriptor.
+  template <typename... Ts> Descriptor *allocateDescriptor(Ts &&...Args) {
+    return new (DescAllocator) Descriptor(std::forward<Ts>(Args)...);
+  }
+};
+
+} // namespace interp
+} // namespace clang
+#endif