address reviewerusers/DataCorrupted/ExposeDirectMethod-opt

4 files changed, 67 insertions, 78 deletions

diff --git a/clang/lib/CodeGen/CGObjCMac.cpp b/clang/lib/CodeGen/CGObjCMac.cpp
index 6351a55ecabd..e67cbe1e9bb9 100644
--- a/clang/lib/CodeGen/CGObjCMac.cpp
+++ b/clang/lib/CodeGen/CGObjCMac.cpp
@@ -2193,7 +2193,7 @@ CodeGen::RValue CGObjCCommonMac::EmitMessageSend(
   // If this was a class method call on a non-weakly-linked class, record it
   // as realized for the "previously realized" heuristic.
   if (ClassReceiver && Method && !isWeakLinkedClass(ClassReceiver)) {
-    if (llvm::BasicBlock *CurrentBB = CGF.Builder.GetInsertBlock())
+    if (llvm::BasicBlock *CurrentBB = CGF.Builder.GetInsertBlock()) {
       // 1. Class methods have forced class realization (regardless direct or
       // not)
       // 2. Direct methods whose receiver is not null means the class is
@@ -2202,6 +2202,7 @@ CodeGen::RValue CGObjCCommonMac::EmitMessageSend(
           (Method->isInstanceMethod() && !ReceiverCanBeNull)) {
         CGF.ObjCRealizedClasses[CurrentBB].insert(ClassReceiver);
       }
+    }
   }
 
   return nullReturn.complete(CGF, Return, rvalue, ResultType, CallArgs,
diff --git a/clang/lib/CodeGen/CGObjCRuntime.cpp b/clang/lib/CodeGen/CGObjCRuntime.cpp
index 82780e3268f9..4e28acd76c8d 100644
--- a/clang/lib/CodeGen/CGObjCRuntime.cpp
+++ b/clang/lib/CodeGen/CGObjCRuntime.cpp
@@ -397,12 +397,10 @@ bool CGObjCRuntime::canMessageReceiverBeNull(
 
   // If we're emitting a method, and self is const (meaning just ARC, for now),
   // and the receiver is a load of self, then self is a valid object.
-  if (const auto *curMethod =
-          dyn_cast_or_null<ObjCMethodDecl>(CGF.CurCodeDecl)) {
-    const auto *self = curMethod->getSelfDecl();
+  if (auto *curMethod = dyn_cast_or_null<ObjCMethodDecl>(CGF.CurCodeDecl)) {
+    const ImplicitParamDecl *self = curMethod->getSelfDecl();
     if (self->getType().isConstQualified()) {
-      if (const auto *LI =
-              dyn_cast<llvm::LoadInst>(receiver->stripPointerCasts())) {
+      if (auto LI = dyn_cast<llvm::LoadInst>(receiver->stripPointerCasts())) {
         llvm::Value *selfAddr = CGF.GetAddrOfLocalVar(self).emitRawPointer(CGF);
         if (selfAddr == LI->getPointerOperand()) {
           return false;
@@ -420,20 +418,20 @@ bool CGObjCRuntime::canClassObjectBeUnrealized(
   if (!CalleeClassDecl || isWeakLinkedClass(CalleeClassDecl))
     return true;
 
+  const ObjCInterfaceDecl *CanonicalClassDecl =
+      CalleeClassDecl->getCanonicalDecl();
+
   // Heuristic 1: +load method on this class or any subclass
-  // If the class or any of its subclasses has a +load method, it's realized
-  // when the binary is loaded. We cache this information to avoid repeatedly
-  // scanning the translation unit.
-  if (getOrPopulateRealizedClasses().contains(CalleeClassDecl))
+  if (isClassRealizedByLoader(CanonicalClassDecl))
     return false;
 
   // Heuristic 2: using Self / Super
-  // If we're currently executing a method of ClassDecl (or a subclass),
-  // then ClassDecl must already be realized.
+  // If we're inside the body of method declared on CanonicalClassDecl or a
+  // subclass, CanonicalClassDecl must have been realized
   if (const auto *CurMethod =
           dyn_cast_or_null<ObjCMethodDecl>(CGF.CurCodeDecl)) {
-    const ObjCInterfaceDecl *CallerCalssDecl = CurMethod->getClassInterface();
-    if (CallerCalssDecl && CalleeClassDecl->isSuperClassOf(CallerCalssDecl))
+    const ObjCInterfaceDecl *CallerClassDecl = CurMethod->getClassInterface();
+    if (CallerClassDecl && CanonicalClassDecl->isSuperClassOf(CallerClassDecl))
       return false;
   }
 
@@ -444,21 +442,22 @@ bool CGObjCRuntime::canClassObjectBeUnrealized(
   // TODO: Iter over all dominating blocks instead of just looking at the
   // current block. While we can construct a DT using CFG.CurFn, it is expensive
   // to do so repeatly when CGF is still emitting blocks.
-  if (auto *CurBB = CGF.Builder.GetInsertBlock()) {
-    auto It = CGF.ObjCRealizedClasses.find(CurBB);
-    if (It != CGF.ObjCRealizedClasses.end()) {
-      // Check if CalleeClassDecl is the same as or a superclass of any
-      // realized class in the cache. A realized subclass implies the parent
-      // is realized.
-      for (const auto *RealizedClass : It->second) {
-        if (CalleeClassDecl == RealizedClass)
-          return false;
-        if (CalleeClassDecl->isSuperClassOf(RealizedClass)) {
-          // Also cache this class to reduce future `isSuperClassOf` calls
-          It->second.insert(CalleeClassDecl);
-          return false;
-        }
-      }
+  auto *CurBB = CGF.Builder.GetInsertBlock();
+  if (!CurBB)
+    return true;
+  auto It = CGF.ObjCRealizedClasses.find(CurBB);
+  if (It == CGF.ObjCRealizedClasses.end())
+    return true;
+
+  llvm::SmallPtrSet<const ObjCInterfaceDecl *, 4> &BlockCache = It->second;
+  if (BlockCache.contains(CanonicalClassDecl))
+    return false;
+  // Check if CanonicalClassDecl is a superclass of any realized class in
+  // the cache. A realized subclass implies the parent is realized.
+  for (const ObjCInterfaceDecl *RealizedClass : BlockCache) {
+    if (CanonicalClassDecl->isSuperClassOf(RealizedClass)) {
+      BlockCache.insert(CanonicalClassDecl);
+      return false;
     }
   }
 
@@ -466,36 +465,40 @@ bool CGObjCRuntime::canClassObjectBeUnrealized(
   return true;
 }
 
-const RealizedClassSet &CGObjCRuntime::getOrPopulateRealizedClasses() const {
-  if (RealizedClasses)
-    return *RealizedClasses;
-  RealizedClasses = llvm::DenseSet<const ObjCInterfaceDecl *>();
-
-  ASTContext &Ctx = CGM.getContext();
-  const IdentifierInfo *LoadII = &Ctx.Idents.get("load");
-  Selector LoadSel = Ctx.Selectors.getSelector(0, &LoadII);
-
-  TranslationUnitDecl *TUDecl = Ctx.getTranslationUnitDecl();
-  llvm::DenseSet<const ObjCInterfaceDecl *> VisitedClasses;
-  for (const auto *D : TUDecl->decls()) {
-    if (const auto *OID = dyn_cast<ObjCInterfaceDecl>(D)) {
-      if (VisitedClasses.contains(OID))
-        continue;
-      // Check if this class has a +load method
-      if (OID->lookupMethod(LoadSel, /*isInstance=*/false,
-                            /*shallowCategoryLookup=*/false,
-                            /*followSuper=*/false)) {
-        // Add this class and all its superclasses to the realized set
-        const ObjCInterfaceDecl *Cls = OID;
-        while (Cls) {
-          RealizedClasses->insert(Cls);
-          VisitedClasses.insert(Cls);
-          Cls = Cls->getSuperClass();
+bool CGObjCRuntime::isClassRealizedByLoader(
+    const ObjCInterfaceDecl *ClassDecl) const {
+  auto populateRealizedClasses = [&]() {
+    RealizedClasses = llvm::DenseSet<const ObjCInterfaceDecl *>();
+
+    ASTContext &Ctx = CGM.getContext();
+    const IdentifierInfo *LoadII = &Ctx.Idents.get("load");
+    Selector LoadSel = Ctx.Selectors.getSelector(0, &LoadII);
+
+    TranslationUnitDecl *TUDecl = Ctx.getTranslationUnitDecl();
+    for (const auto *D : TUDecl->decls()) {
+      if (const auto *OID = dyn_cast<ObjCInterfaceDecl>(D)) {
+        const ObjCInterfaceDecl *CanonicalOID = OID->getCanonicalDecl();
+        if (RealizedClasses->contains(CanonicalOID))
+          continue;
+        // Check if this class has a +load method
+        if (CanonicalOID->lookupMethod(LoadSel, /*isInstance=*/false,
+                                       /*shallowCategoryLookup=*/false,
+                                       /*followSuper=*/false)) {
+          // Add this class and all its superclasses to the realized set
+          const ObjCInterfaceDecl *Cls = CanonicalOID;
+          while (Cls) {
+            RealizedClasses->insert(Cls->getCanonicalDecl());
+            Cls = Cls->getSuperClass();
+          }
         }
       }
     }
-  }
-  return *RealizedClasses;
+  };
+
+  if (!RealizedClasses)
+    populateRealizedClasses();
+
+  return RealizedClasses->contains(ClassDecl);
 }
 
 bool CGObjCRuntime::isWeakLinkedClass(const ObjCInterfaceDecl *ID) {
diff --git a/clang/lib/CodeGen/CGObjCRuntime.h b/clang/lib/CodeGen/CGObjCRuntime.h
index 9f7243e2aa73..d38a3cf5ec43 100644
--- a/clang/lib/CodeGen/CGObjCRuntime.h
+++ b/clang/lib/CodeGen/CGObjCRuntime.h
@@ -69,13 +69,13 @@ protected:
   CodeGen::CodeGenModule &CGM;
   CGObjCRuntime(CodeGen::CodeGenModule &CGM) : CGM(CGM) {}
 
-  /// Cache of classes that are guaranteed to be realized because they or one
-  /// of their subclasses has a +load method. Lazily populated on first query.
+  /// Cache of classes known to be realized during loading.
   mutable std::optional<RealizedClassSet> RealizedClasses;
 
-  /// Populate the RealizedClasses cache by scanning all ObjCInterfaceDecls
-  /// in the translation unit for +load methods.
-  const RealizedClassSet &getOrPopulateRealizedClasses() const;
+  /// Query the cache to determine if a class is guaranteed to be realized
+  /// because of +load. Construct the cache if by scanning all
+  /// ObjCInterfaceDecls in the translation unit if we haven't done so.
+  bool isClassRealizedByLoader(const ObjCInterfaceDecl *CalleeClassDecl) const;
 
   // Utility functions for unified ivar access. These need to
   // eventually be folded into other places (the structure layout
@@ -340,18 +340,6 @@ public:
                                      QualType resultType,
                                      CallArgList &callArgs);
 
-  /// Check if the receiver of an ObjC message send can be null.
-  /// Returns true if the receiver may be null, false if provably non-null.
-  ///
-  /// This can be overridden by subclasses to add runtime-specific heuristics.
-  /// Base implementation checks:
-  /// - Super dispatch (always non-null)
-  /// - Self in const-qualified methods (ARC)
-  /// - Weak-linked classes
-  ///
-  /// Future enhancements in CGObjCCommonMac override:
-  /// - _Nonnull attributes
-  /// - Results of alloc, new, ObjC literals
   virtual bool canMessageReceiverBeNull(CodeGenFunction &CGF,
                                         const ObjCMethodDecl *method,
                                         bool isSuper,
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index 0cbccd32ce43..4e114d1fb3e5 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -870,11 +870,8 @@ public:
   /// rethrows.
   SmallVector<llvm::Value *, 8> ObjCEHValueStack;
 
-  /// Per-basic-block cache of ObjC classes that have been realized during
-  /// codegen. When a class method is emitted on a non-weakly-linked class,
-  /// we record it here. This supports the "previously realized" heuristic
-  /// in canClassObjectBeUnrealized. The structure supports future
-  /// dominator-based analysis where we can check dominating blocks.
+  /// A cache of objc classes that that are known to have been realized in each
+  /// basic block
   llvm::DenseMap<llvm::BasicBlock *,
                  llvm::SmallPtrSet<const ObjCInterfaceDecl *, 4>>
       ObjCRealizedClasses;