[clang][bytecode] Optimize InterpStack (#159400)

Replace `StackChunk::End` with `StackChunk::Size`, mark the allocating
code paths as unlikely and move `grow()` into the header, which allows
us to template this for the `Size` parameter. Since we only push our
primitive types on the stack and all the sizes are aligned to pointer
size multiples, this only results in a few instantiations.

Author: tbaederr

clang/lib/AST/ByteCode/InterpStack.cpp

@@ -24,9 +24,6 @@ InterpStack::~InterpStack() {
     std::free(Chunk->Next);
   if (Chunk)
     std::free(Chunk);
-  Chunk = nullptr;
-  StackSize = 0;
-  ItemTypes.clear();
 }
 
 // We keep the last chunk around to reuse.
@@ -56,54 +53,44 @@ void InterpStack::clearTo(size_t NewSize) {
   assert(size() == NewSize);
 }
 
-void *InterpStack::grow(size_t Size) {
-  assert(Size < ChunkSize - sizeof(StackChunk) && "Object too large");
-
-  if (!Chunk || sizeof(StackChunk) + Chunk->size() + Size > ChunkSize) {
-    if (Chunk && Chunk->Next) {
-      Chunk = Chunk->Next;
-    } else {
-      StackChunk *Next = new (std::malloc(ChunkSize)) StackChunk(Chunk);
-      if (Chunk)
-        Chunk->Next = Next;
-      Chunk = Next;
-    }
-  }
-
-  auto *Object = reinterpret_cast<void *>(Chunk->End);
-  Chunk->End += Size;
-  StackSize += Size;
-  return Object;
-}
-
 void *InterpStack::peekData(size_t Size) const {
   assert(Chunk && "Stack is empty!");
 
+  if (LLVM_LIKELY(Size <= Chunk->size()))
+    return reinterpret_cast<void *>(Chunk->start() + Chunk->Size - Size);
+
   StackChunk *Ptr = Chunk;
   while (Size > Ptr->size()) {
     Size -= Ptr->size();
     Ptr = Ptr->Prev;
     assert(Ptr && "Offset too large");
   }
 
-  return reinterpret_cast<void *>(Ptr->End - Size);
+  return reinterpret_cast<void *>(Ptr->start() + Ptr->Size - Size);
 }
 
 void InterpStack::shrink(size_t Size) {
   assert(Chunk && "Chunk is empty!");
 
+  // Likely case is that we simply remove something from the current chunk.
+  if (LLVM_LIKELY(Size <= Chunk->size())) {
+    Chunk->Size -= Size;
+    StackSize -= Size;
+    return;
+  }
+
   while (Size > Chunk->size()) {
     Size -= Chunk->size();
     if (Chunk->Next) {
       std::free(Chunk->Next);
       Chunk->Next = nullptr;
     }
-    Chunk->End = Chunk->start();
+    Chunk->Size = 0;
     Chunk = Chunk->Prev;
     assert(Chunk && "Offset too large");
   }
 
-  Chunk->End -= Size;
+  Chunk->Size -= Size;
   StackSize -= Size;
 }
 

clang/lib/AST/ByteCode/InterpStack.h

@@ -24,14 +24,14 @@ namespace interp {
 /// Stack frame storing temporaries and parameters.
 class InterpStack final {
 public:
-  InterpStack() {}
+  InterpStack() = default;
 
   /// Destroys the stack, freeing up storage.
   ~InterpStack();
 
   /// Constructs a value in place on the top of the stack.
   template <typename T, typename... Tys> void push(Tys &&...Args) {
-    new (grow(aligned_size<T>())) T(std::forward<Tys>(Args)...);
+    new (grow<aligned_size<T>()>()) T(std::forward<Tys>(Args)...);
     ItemTypes.push_back(toPrimType<T>());
   }
 
@@ -89,7 +89,7 @@ class InterpStack final {
 private:
   /// All stack slots are aligned to the native pointer alignment for storage.
   /// The size of an object is rounded up to a pointer alignment multiple.
-  template <typename T> constexpr size_t aligned_size() const {
+  template <typename T> static constexpr size_t aligned_size() {
     constexpr size_t PtrAlign = alignof(void *);
     return ((sizeof(T) + PtrAlign - 1) / PtrAlign) * PtrAlign;
   }
@@ -100,7 +100,30 @@ class InterpStack final {
   }
 
   /// Grows the stack to accommodate a value and returns a pointer to it.
-  void *grow(size_t Size);
+  template <size_t Size> void *grow() {
+    assert(Size < ChunkSize - sizeof(StackChunk) && "Object too large");
+    static_assert(aligned(Size));
+
+    // Allocate a new stack chunk if necessary.
+    if (LLVM_UNLIKELY(!Chunk)) {
+      Chunk = new (std::malloc(ChunkSize)) StackChunk(Chunk);
+    } else if (LLVM_UNLIKELY(Chunk->size() >
+                             ChunkSize - sizeof(StackChunk) - Size)) {
+      if (Chunk->Next) {
+        Chunk = Chunk->Next;
+      } else {
+        StackChunk *Next = new (std::malloc(ChunkSize)) StackChunk(Chunk);
+        Chunk->Next = Next;
+        Chunk = Next;
+      }
+    }
+
+    auto *Object = reinterpret_cast<void *>(Chunk->start() + Chunk->Size);
+    Chunk->Size += Size;
+    StackSize += Size;
+    return Object;
+  }
+
   /// Returns a pointer from the top of the stack.
   void *peekData(size_t Size) const;
   /// Shrinks the stack.
@@ -118,13 +141,13 @@ class InterpStack final {
   struct StackChunk {
     StackChunk *Next;
     StackChunk *Prev;
-    char *End;
+    uint32_t Size;
 
     StackChunk(StackChunk *Prev = nullptr)
-        : Next(nullptr), Prev(Prev), End(reinterpret_cast<char *>(this + 1)) {}
+        : Next(nullptr), Prev(Prev), Size(0) {}
 
     /// Returns the size of the chunk, minus the header.
-    size_t size() const { return End - start(); }
+    size_t size() const { return Size; }
 
     /// Returns a pointer to the start of the data region.
     char *start() { return reinterpret_cast<char *>(this + 1); }