[clang][Interp] Implement dynamic memory allocation handling

Author: tbaederr

clang/lib/AST/CMakeLists.txt

@@ -76,6 +76,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/ByteCodeExprGen.cpp

@@ -2471,6 +2471,85 @@ bool ByteCodeExprGen<Emitter>::VisitCXXInheritedCtorInitExpr(
   return this->emitCall(F, 0, E);
 }
 
+template <class Emitter>
+bool ByteCodeExprGen<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);
+
+  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;
+    assert(ArraySizeExpr);
+    PrimType SizeT = classifyPrim((*ArraySizeExpr)->getType());
+
+    if (!this->visit(*ArraySizeExpr))
+      return false;
+
+    if (ElemT) {
+      // N primitive elements.
+      if (!this->emitAllocN(SizeT, *ElemT, E, E))
+        return false;
+    } else {
+      // N Composite elements.
+      if (!this->emitAllocCN(SizeT, Desc, E))
+        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 ByteCodeExprGen<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 ByteCodeExprGen<Emitter>::VisitExpressionTraitExpr(
     const ExpressionTraitExpr *E) {

clang/lib/AST/Interp/ByteCodeExprGen.h

@@ -128,6 +128,8 @@ class ByteCodeExprGen : public ConstStmtVisitor<ByteCodeExprGen<Emitter>, bool>,
   bool VisitConvertVectorExpr(const ConvertVectorExpr *E);
   bool VisitShuffleVectorExpr(const ShuffleVectorExpr *E);
   bool VisitObjCBoxedExpr(const ObjCBoxedExpr *E);
+  bool VisitCXXNewExpr(const CXXNewExpr *E);
+  bool VisitCXXDeleteExpr(const CXXDeleteExpr *E);
 
 protected:
   bool visitExpr(const Expr *E) override;

clang/lib/AST/Interp/DynamicAllocator.cpp

@@ -0,0 +1,119 @@
+//==-------- 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) {
+  assert(NumElements > 0);
+  // 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);
+}
+
+Block *DynamicAllocator::allocate(const Descriptor *ElementDesc,
+                                  size_t NumElements) {
+  assert(NumElements > 0);
+  // 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);
+}
+
+Block *DynamicAllocator::allocate(const Descriptor *D) {
+  assert(D->asExpr());
+
+  auto Memory =
+      std::make_unique<std::byte[]>(sizeof(Block) + D->getAllocSize());
+  auto *B = new (Memory.get()) Block(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;
+  assert(ID->Desc);
+
+  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,95 @@
+//==--------- 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 allocat 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);
+  /// Allocate \p NumElements primitive elements of the given type.
+  Block *allocate(const Expr *Source, PrimType T, size_t NumElements);
+  /// Allocate \p NumElements elements of the given descriptor.
+  Block *allocate(const Descriptor *D, size_t NumElements);
+
+  /// 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;
+  }
+
+  // FIXME: Public because I'm not sure how to expose an iterator to it.
+  llvm::DenseMap<const Expr *, AllocationSite> AllocationSites;
+
+private:
+  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

clang/lib/AST/Interp/EvalEmitter.cpp

@@ -124,9 +124,17 @@ bool EvalEmitter::fallthrough(const LabelTy &Label) {
   return true;
 }
 
+static bool checkReturnState(InterpState &S) {
+  return S.maybeDiagnoseDanglingAllocations();
+}
+
 template <PrimType OpType> bool EvalEmitter::emitRet(const SourceInfo &Info) {
   if (!isActive())
     return true;
+
+  if (!checkReturnState(S))
+    return false;
+
   using T = typename PrimConv<OpType>::T;
   EvalResult.setValue(S.Stk.pop<T>().toAPValue());
   return true;
@@ -137,6 +145,12 @@ template <> bool EvalEmitter::emitRet<PT_Ptr>(const SourceInfo &Info) {
     return true;
 
   const Pointer &Ptr = S.Stk.pop<Pointer>();
+
+  if (!EvalResult.checkReturnValue(S, Ctx, Ptr, Info))
+    return false;
+  if (!checkReturnState(S))
+    return false;
+
   // Implicitly convert lvalue to rvalue, if requested.
   if (ConvertResultToRValue) {
     if (std::optional<APValue> V = Ptr.toRValue(Ctx)) {
@@ -163,18 +177,30 @@ template <> bool EvalEmitter::emitRet<PT_Ptr>(const SourceInfo &Info) {
 template <> bool EvalEmitter::emitRet<PT_FnPtr>(const SourceInfo &Info) {
   if (!isActive())
     return true;
+
+  if (!checkReturnState(S))
+    return false;
   // Function pointers cannot be converted to rvalues.
   EvalResult.setFunctionPointer(S.Stk.pop<FunctionPointer>());
   return true;
 }
 
 bool EvalEmitter::emitRetVoid(const SourceInfo &Info) {
+  if (!checkReturnState(S))
+    return false;
   EvalResult.setValid();
   return true;
 }
 
 bool EvalEmitter::emitRetValue(const SourceInfo &Info) {
   const auto &Ptr = S.Stk.pop<Pointer>();
+
+  if (!EvalResult.checkReturnValue(S, Ctx, Ptr, Info))
+    return false;
+
+  if (!checkReturnState(S))
+    return false;
+
   if (std::optional<APValue> APV = Ptr.toRValue(S.getCtx())) {
     EvalResult.setValue(*APV);
     return true;