[clang][bytecode] Save a per-block dynamic allocation ID (#154094)

This fixes an old todo item about wrong allocation counting and some
diagnostic differences.

Author: tbaederr

clang/lib/AST/ByteCode/Disasm.cpp

@@ -545,7 +545,7 @@ LLVM_DUMP_METHOD void Block::dump(llvm::raw_ostream &OS) const {
   OS << "  Initialized: " << IsInitialized << "\n";
   OS << "  Weak: " << isWeak() << "\n";
   OS << "  Dummy: " << isDummy() << '\n';
-  OS << "  Dynamic: " << IsDynamic << "\n";
+  OS << "  Dynamic: " << isDynamic() << "\n";
 }
 
 LLVM_DUMP_METHOD void EvaluationResult::dump() const {

clang/lib/AST/ByteCode/DynamicAllocator.cpp

@@ -101,13 +101,17 @@ Block *DynamicAllocator::allocate(const Descriptor *D, unsigned EvalID,
     ID->LifeState =
         AllocForm == Form::Operator ? Lifetime::Ended : Lifetime::Started;
 
-  B->IsDynamic = true;
-
-  if (auto It = AllocationSites.find(D->asExpr()); It != AllocationSites.end())
+  if (auto It = AllocationSites.find(D->asExpr());
+      It != AllocationSites.end()) {
     It->second.Allocations.emplace_back(std::move(Memory));
-  else
+    B->setDynAllocId(It->second.NumAllocs);
+    ++It->second.NumAllocs;
+  } else {
     AllocationSites.insert(
         {D->asExpr(), AllocationSite(std::move(Memory), AllocForm)});
+    B->setDynAllocId(0);
+  }
+  assert(B->isDynamic());
   return B;
 }
 

clang/lib/AST/ByteCode/DynamicAllocator.h

@@ -48,11 +48,13 @@ class DynamicAllocator final {
 
   struct AllocationSite {
     llvm::SmallVector<Allocation> Allocations;
+    unsigned NumAllocs = 0;
     Form AllocForm;
 
     AllocationSite(std::unique_ptr<std::byte[]> Memory, Form AllocForm)
         : AllocForm(AllocForm) {
       Allocations.push_back({std::move(Memory)});
+      ++NumAllocs;
     }
 
     size_t size() const { return Allocations.size(); }

clang/lib/AST/ByteCode/InterpBlock.cpp

@@ -56,7 +56,7 @@ void Block::removePointer(Pointer *P) {
 }
 
 void Block::cleanup() {
-  if (Pointers == nullptr && !IsDynamic && isDead())
+  if (Pointers == nullptr && !isDynamic() && isDead())
     (reinterpret_cast<DeadBlock *>(this + 1) - 1)->free();
 }
 
@@ -111,7 +111,7 @@ DeadBlock::DeadBlock(DeadBlock *&Root, Block *Blk)
   Prev = nullptr;
   Root = this;
 
-  B.IsDynamic = Blk->IsDynamic;
+  B.DynAllocId = Blk->DynAllocId;
 
   // Transfer pointers.
   B.Pointers = Blk->Pointers;

clang/lib/AST/ByteCode/InterpBlock.h

@@ -62,7 +62,7 @@ class Block final {
 
   Block(unsigned EvalID, const Descriptor *Desc, bool IsStatic = false,
         bool IsExtern = false, bool IsWeak = false, bool IsDummy = false)
-      : Desc(Desc), EvalID(EvalID), IsStatic(IsStatic), IsDynamic(false) {
+      : Desc(Desc), EvalID(EvalID), IsStatic(IsStatic) {
     assert(Desc);
     AccessFlags |= (ExternFlag * IsExtern);
     AccessFlags |= (WeakFlag * IsWeak);
@@ -80,7 +80,7 @@ class Block final {
   /// Checks if the block is temporary.
   bool isTemporary() const { return Desc->IsTemporary; }
   bool isWeak() const { return AccessFlags & WeakFlag; }
-  bool isDynamic() const { return IsDynamic; }
+  bool isDynamic() const { return (DynAllocId != std::nullopt); }
   bool isDummy() const { return AccessFlags & DummyFlag; }
   bool isDead() const { return AccessFlags & DeadFlag; }
   /// Returns the size of the block.
@@ -160,6 +160,9 @@ class Block final {
     AccessFlags |= (DummyFlag * IsDummy);
   }
 
+  /// To be called by DynamicAllocator.
+  void setDynAllocId(unsigned ID) { DynAllocId = ID; }
+
   /// Deletes a dead block at the end of its lifetime.
   void cleanup();
 
@@ -183,9 +186,8 @@ class Block final {
   /// Flag indicating if the block contents have been initialized
   /// via invokeCtor.
   bool IsInitialized = false;
-  /// Flag indicating if this block has been allocated via dynamic
-  /// memory allocation (e.g. malloc).
-  bool IsDynamic = false;
+  /// Allocation ID for this dynamic allocation, if it is one.
+  UnsignedOrNone DynAllocId = std::nullopt;
   /// AccessFlags containing IsExtern, IsDead, IsWeak, and IsDummy bits.
   uint8_t AccessFlags = 0;
 };

clang/lib/AST/ByteCode/Pointer.cpp

@@ -179,10 +179,7 @@ APValue Pointer::toAPValue(const ASTContext &ASTCtx) const {
   else if (const auto *E = Desc->asExpr()) {
     if (block()->isDynamic()) {
       QualType AllocatedType = getDeclPtr().getFieldDesc()->getDataType(ASTCtx);
-      // FIXME: Suboptimal counting of dynamic allocations. Move this to Context
-      // or InterpState?
-      static int ReportedDynamicAllocs = 0;
-      DynamicAllocLValue DA(ReportedDynamicAllocs++);
+      DynamicAllocLValue DA(*block()->DynAllocId);
       Base = APValue::LValueBase::getDynamicAlloc(DA, AllocatedType);
     } else {
       Base = E;

clang/test/AST/ByteCode/new-delete.cpp

@@ -82,8 +82,7 @@ static_assert(noInit() == 0, "");
 /// Try to delete a pointer that hasn't been heap allocated.
 constexpr int notHeapAllocated() { // both-error {{never produces a constant expression}}
   int A = 0; // both-note 2{{declared here}}
-  delete &A; // ref-note 2{{delete of pointer '&A' that does not point to a heap-allocated object}} \
-             // expected-note 2{{delete of pointer '&A' that does not point to a heap-allocated object}}
+  delete &A; // both-note 2{{delete of pointer '&A' that does not point to a heap-allocated object}}
 
   return 1;
 }
@@ -374,8 +373,7 @@ namespace delete_random_things {
   static_assert((delete &(new A)->n, true)); // both-error {{}} \
                                              // both-note {{delete of pointer to subobject }}
   static_assert((delete (new int + 1), true)); // both-error {{}} \
-                                               // ref-note {{delete of pointer '&{*new int#0} + 1' that does not point to complete object}} \
-                                               // expected-note {{delete of pointer '&{*new int#1} + 1' that does not point to complete object}}
+                                               // both-note {{delete of pointer '&{*new int#0} + 1' that does not point to complete object}}
   static_assert((delete[] (new int[3] + 1), true)); // both-error {{}} \
                                                     // both-note {{delete of pointer to subobject}}
   static_assert((delete &(int&)(int&&)0, true)); // both-error {{}} \