[clang][bytecode] Use std::allocator calls for Descriptor source (llvm#123900)

... for the dynamic blocks created for operator new calls. This way we
get the type of memory allocated right. As a side-effect, the
diagnostics now point to the std::allocator calls, which is an
improvement.

Author: tbaederr

clang/lib/AST/ByteCode/Interp.cpp

@@ -873,13 +873,17 @@ bool CheckNewDeleteForms(InterpState &S, CodePtr OpPC,
 
 bool CheckDeleteSource(InterpState &S, CodePtr OpPC, const Expr *Source,
                        const Pointer &Ptr) {
-  // The two sources we currently allow are new expressions and
-  // __builtin_operator_new calls.
+  // Regular new type(...) call.
   if (isa_and_nonnull<CXXNewExpr>(Source))
     return true;
-  if (const CallExpr *CE = dyn_cast_if_present<CallExpr>(Source);
+  // operator new.
+  if (const auto *CE = dyn_cast_if_present<CallExpr>(Source);
       CE && CE->getBuiltinCallee() == Builtin::BI__builtin_operator_new)
     return true;
+  // std::allocator.allocate() call
+  if (const auto *MCE = dyn_cast_if_present<CXXMemberCallExpr>(Source);
+      MCE && MCE->getMethodDecl()->getIdentifier()->isStr("allocate"))
+    return true;
 
   // Whatever this is, we didn't heap allocate it.
   const SourceInfo &Loc = S.Current->getSource(OpPC);
@@ -1489,7 +1493,8 @@ bool CheckNewTypeMismatch(InterpState &S, CodePtr OpPC, const Expr *E,
   const auto *NewExpr = cast<CXXNewExpr>(E);
   QualType StorageType = Ptr.getType();
 
-  if (isa_and_nonnull<CXXNewExpr>(Ptr.getFieldDesc()->asExpr()) &&
+  if ((isa_and_nonnull<CXXNewExpr>(Ptr.getFieldDesc()->asExpr()) ||
+       isa_and_nonnull<CXXMemberCallExpr>(Ptr.getFieldDesc()->asExpr())) &&
       StorageType->isPointerType()) {
     // FIXME: Are there other cases where this is a problem?
     StorageType = StorageType->getPointeeType();

clang/lib/AST/ByteCode/InterpBuiltin.cpp

@@ -1584,6 +1584,7 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
   // Walk up the call stack to find the appropriate caller and get the
   // element type from it.
   QualType ElemType;
+  const CallExpr *NewCall = nullptr;
 
   for (const InterpFrame *F = Frame; F; F = F->Caller) {
     const Function *Func = F->getFunction();
@@ -1606,6 +1607,7 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
     if (CTSD->isInStdNamespace() && ClassII && ClassII->isStr("allocator") &&
         TAL.size() >= 1 && TAL[0].getKind() == TemplateArgument::Type) {
       ElemType = TAL[0].getAsType();
+      NewCall = cast<CallExpr>(F->Caller->getExpr(F->getRetPC()));
       break;
     }
   }
@@ -1616,6 +1618,7 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
                        : diag::note_constexpr_new);
     return false;
   }
+  assert(NewCall);
 
   if (ElemType->isIncompleteType() || ElemType->isFunctionType()) {
     S.FFDiag(Call, diag::note_constexpr_new_not_complete_object_type)
@@ -1654,7 +1657,7 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
   if (ElemT) {
     if (NumElems.ule(1)) {
       const Descriptor *Desc =
-          S.P.createDescriptor(Call, *ElemT, Descriptor::InlineDescMD,
+          S.P.createDescriptor(NewCall, *ElemT, Descriptor::InlineDescMD,
                                /*IsConst=*/false, /*IsTemporary=*/false,
                                /*IsMutable=*/false);
       Block *B = Allocator.allocate(Desc, S.getContext().getEvalID(),
@@ -1667,7 +1670,7 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
     assert(NumElems.ugt(1));
 
     Block *B =
-        Allocator.allocate(Call, *ElemT, NumElems.getZExtValue(),
+        Allocator.allocate(NewCall, *ElemT, NumElems.getZExtValue(),
                            S.Ctx.getEvalID(), DynamicAllocator::Form::Operator);
     assert(B);
     S.Stk.push<Pointer>(B);

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

@@ -602,8 +602,7 @@ namespace std {
   using size_t = decltype(sizeof(0));
   template<typename T> struct allocator {
     constexpr T *allocate(size_t N) {
-      return (T*)__builtin_operator_new(sizeof(T) * N); // expected-note 2{{allocation performed here}} \
-                                                        // #alloc
+      return (T*)__builtin_operator_new(sizeof(T) * N); // #alloc
     }
     constexpr void deallocate(void *p) {
       __builtin_operator_delete(p); // both-note 2{{std::allocator<...>::deallocate' used to delete pointer to object allocated with 'new'}} \
@@ -641,7 +640,7 @@ namespace OperatorNewDelete {
       p = new int[1]; // both-note {{heap allocation performed here}}
       break;
     case 2:
-      p = std::allocator<int>().allocate(1); // ref-note 2{{heap allocation performed here}}
+      p = std::allocator<int>().allocate(1); // both-note 2{{heap allocation performed here}}
       break;
     }
     switch (dealloc_kind) {
@@ -838,6 +837,26 @@ namespace ToplevelScopeInTemplateArg {
   }
 }
 
+template <typename T>
+struct SS {
+    constexpr SS(unsigned long long N)
+    : data(nullptr){
+        data = alloc.allocate(N);  // #call
+        for(std::size_t i = 0; i < N; i ++)
+            std::construct_at<T>(data + i, i); // #construct_call
+    }
+    constexpr T operator[](std::size_t i) const {
+      return data[i];
+    }
+
+    constexpr ~SS() {
+        alloc.deallocate(data);
+    }
+    std::allocator<T> alloc;
+    T* data;
+};
+constexpr unsigned short ssmall = SS<unsigned short>(100)[42];
+
 #else
 /// Make sure we reject this prior to C++20
 constexpr int a() { // both-error {{never produces a constant expression}}