[clang][bytecode] Check for memory leaks after destroying global scope

clang/lib/AST/ByteCode/Compiler.cpp

@@ -4132,10 +4132,16 @@ template <class Emitter>
 bool Compiler<Emitter>::visitExpr(const Expr *E, bool DestroyToplevelScope) {
   LocalScope<Emitter> RootScope(this);
 
+  // If we won't destroy the toplevel scope, check for memory leaks first.
+  if (!DestroyToplevelScope) {
+    if (!this->emitCheckAllocations(E))
+      return false;
+  }
+
   auto maybeDestroyLocals = [&]() -> bool {
     if (DestroyToplevelScope)
-      return RootScope.destroyLocals();
-    return true;
+      return RootScope.destroyLocals() && this->emitCheckAllocations(E);
+    return this->emitCheckAllocations(E);
   };
 
   // Void expressions.
@@ -4171,8 +4177,7 @@ bool Compiler<Emitter>::visitExpr(const Expr *E, bool DestroyToplevelScope) {
     return this->emitRetValue(E) && maybeDestroyLocals();
   }
 
-  (void)maybeDestroyLocals();
-  return false;
+  return maybeDestroyLocals() && this->emitCheckAllocations(E) && false;
 }
 
 template <class Emitter>
@@ -4214,7 +4219,8 @@ bool Compiler<Emitter>::visitDeclAndReturn(const VarDecl *VD,
     DeclScope<Emitter> LS(this, VD);
     if (!this->visit(VD->getAnyInitializer()))
       return false;
-    return this->emitRet(VarT.value_or(PT_Ptr), VD) && LS.destroyLocals();
+    return this->emitRet(VarT.value_or(PT_Ptr), VD) && LS.destroyLocals() &&
+           this->emitCheckAllocations(VD);
   }
 
   LocalScope<Emitter> VDScope(this, VD);
@@ -4260,7 +4266,7 @@ bool Compiler<Emitter>::visitDeclAndReturn(const VarDecl *VD,
     return false;
   }
 
-  return VDScope.destroyLocals();
+  return VDScope.destroyLocals() && this->emitCheckAllocations(VD);
 }
 
 template <class Emitter>

clang/lib/AST/ByteCode/Context.cpp

@@ -78,8 +78,7 @@ bool Context::evaluate(State &Parent, const Expr *E, APValue &Result,
   Compiler<EvalEmitter> C(*this, *P, Parent, Stk);
 
   auto Res = C.interpretExpr(E, /*ConvertResultToRValue=*/false,
-                             /*DestroyToplevelScope=*/Kind ==
-                                 ConstantExprKind::ClassTemplateArgument);
+                             /*DestroyToplevelScope=*/true);
   if (Res.isInvalid()) {
     C.cleanup();
     Stk.clearTo(StackSizeBefore);

clang/lib/AST/ByteCode/EvalEmitter.cpp

@@ -132,17 +132,10 @@ 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(Ctx.getASTContext()));
   return true;
@@ -159,9 +152,6 @@ template <> bool EvalEmitter::emitRet<PT_Ptr>(const SourceInfo &Info) {
   if (CheckFullyInitialized && !EvalResult.checkFullyInitialized(S, Ptr))
     return false;
 
-  if (!checkReturnState(S))
-    return false;
-
   // Implicitly convert lvalue to rvalue, if requested.
   if (ConvertResultToRValue) {
     if (!Ptr.isZero() && !Ptr.isDereferencable())
@@ -194,16 +184,12 @@ 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;
 }
@@ -216,9 +202,6 @@ bool EvalEmitter::emitRetValue(const SourceInfo &Info) {
   if (CheckFullyInitialized && !EvalResult.checkFullyInitialized(S, Ptr))
     return false;
 
-  if (!checkReturnState(S))
-    return false;
-
   if (std::optional<APValue> APV =
           Ptr.toRValue(S.getASTContext(), EvalResult.getSourceType())) {
     EvalResult.setValue(*APV);

clang/lib/AST/ByteCode/Interp.h

@@ -3007,6 +3007,10 @@ static inline bool IsConstantContext(InterpState &S, CodePtr OpPC) {
   return true;
 }
 
+static inline bool CheckAllocations(InterpState &S, CodePtr OpPC) {
+  return S.maybeDiagnoseDanglingAllocations();
+}
+
 /// Check if the initializer and storage types of a placement-new expression
 /// match.
 bool CheckNewTypeMismatch(InterpState &S, CodePtr OpPC, const Expr *E,

clang/lib/AST/ByteCode/Opcodes.td

@@ -836,3 +836,4 @@ def CheckNewTypeMismatchArray : Opcode {
 }
 
 def IsConstantContext: Opcode;
+def CheckAllocations : Opcode;

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

@@ -796,6 +796,28 @@ static_assert(virt_delete(false)); // both-error {{not an integral constant expr
                                    // both-note {{in call to}}
 
 
+namespace ToplevelScopeInTemplateArg {
+  class string {
+  public:
+    char *mem;
+    constexpr string() {
+      this->mem = new char(1);
+    }
+    constexpr ~string() {
+      delete this->mem;
+    }
+    constexpr unsigned size() const { return 4; }
+  };
+
+
+  template <unsigned N>
+  void test() {};
+
+  void f() {
+      test<string().size()>();
+      static_assert(string().size() == 4);
+  }
+}
 
 #else
 /// Make sure we reject this prior to C++20