Add a cache to remember previously realized classes

Author: DataCorrupted

clang/lib/CodeGen/CGObjCMac.cpp

@@ -2190,6 +2190,20 @@ CodeGen::RValue CGObjCCommonMac::EmitMessageSend(
     CallSite->setDoesNotReturn();
   }
 
+  // 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())
+      // 1. Class methods have forced class realization (regardless direct or
+      // not)
+      // 2. Direct methods whose receiver is not null means the class is
+      // previously realized.
+      if (Method->isClassMethod() ||
+          (Method->isInstanceMethod() && !ReceiverCanBeNull)) {
+        CGF.ObjCRealizedClasses[CurrentBB].insert(ClassReceiver);
+      }
+  }
+
   return nullReturn.complete(CGF, Return, rvalue, ResultType, CallArgs,
                              RequiresNullCheck ? Method : nullptr);
 }

clang/lib/CodeGen/CGObjCRuntime.cpp

@@ -437,16 +437,30 @@ bool CGObjCRuntime::canClassObjectBeUnrealized(
       return false;
   }
 
-  // TODO: Heuristic 3: previously realized
-  // Walk through previous instructions can be inefficient, since
-  // `canClassObjectBeUnrealized` is called everytime we emit a class method.
-  // Besides, a realized subclass means parent class is realized. Therefore,
-  // a code like below also requires some special handling.
+  // Heuristic 3: previously realized classes
+  // If we've already emitted a class method call for this class (or a subclass)
+  // earlier, then the class must be realized.
   //
-  // ```
-  // +[Child foo];
-  // +[Parent foo];
-  // ```
+  // 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;
+        }
+      }
+    }
+  }
 
   // Otherwise, assume it can be unrealized.
   return true;

clang/lib/CodeGen/CodeGenFunction.h

@@ -870,6 +870,15 @@ class CodeGenFunction : public CodeGenTypeCache {
   /// 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.
+  llvm::DenseMap<llvm::BasicBlock *,
+                 llvm::SmallPtrSet<const ObjCInterfaceDecl *, 4>>
+      ObjCRealizedClasses;
+
   /// A class controlling the emission of a finally block.
   class FinallyInfo {
     /// Where the catchall's edge through the cleanup should go.

clang/test/CodeGenObjC/expose-direct-method-opt-class-realization.m

@@ -128,3 +128,40 @@ - (int)callerInstanceMethod {
 }
 
 @end
+
+// ============================================================================
+// HEURISTIC 3: Previously realized classes in the same basic block skip thunk.
+// If we've already called a class method (which realizes the class),
+// subsequent calls to the same class or its superclasses can skip the thunk.
+// ============================================================================
+
+// CHECK-LABEL: define{{.*}} i32 @testPreviouslyRealizedParentClass
+int testPreviouslyRealizedParentClass(int flag) {
+  if (flag) {
+    // First call to ClassWithoutLoad - needs thunk (class might not be realized)
+    // CHECK: call i32 @"+[ClassWithoutLoad classDirectMethod]_thunk"(ptr noundef
+    int a = [ClassWithoutLoad classDirectMethod];
+
+    // Second call to same class - should skip thunk (class was just realized)
+    // CHECK: call i32 @"+[ClassWithoutLoad classDirectMethod]"(ptr noundef
+    // CHECK-NOT: call i32 @"+[ClassWithoutLoad classDirectMethod]_thunk"
+    int b = [ClassWithoutLoad classDirectMethod];
+
+    // Call to Root (parent of ClassWithoutLoad) - should skip thunk
+    // because realizing ClassWithoutLoad also realizes its superclass Root.
+    // CHECK: call i32 @"+[Root rootDirectMethod]"(ptr noundef
+    // CHECK-NOT: call i32 @"+[Root rootDirectMethod]_thunk"
+    int c = [Root rootDirectMethod];
+    return a + b + c;
+
+  }
+  // New block, we are not sure if prev block is executed, so we have to conservatively realize again.
+  // CHECK: call i32 @"+[ClassWithoutLoad classDirectMethod]_thunk"
+  // CHECK-NOT: call i32 @"+[ClassWithoutLoad classDirectMethod]"(ptr noundef
+  int b = [ClassWithoutLoad classDirectMethod];
+  // CHECK: call i32 @"+[Root rootDirectMethod]"(ptr noundef
+  // CHECK-NOT: call i32 @"+[Root rootDirectMethod]_thunk"
+  int c = [Root rootDirectMethod];
+
+  return b + c;
+}

clang/test/CodeGenObjC/expose-direct-method.m

@@ -266,12 +266,11 @@ int useSRet(Root *r) {
     // TODO: we should know that this instance is non nil.
     // CHECK: call void @"-[Root getAggregate]_thunk"
     [r getAggregate].a +
-    // TODO: The compiler is not smart enough to know the class object must be realized yet.
     // CHECK-NOT: call i64 @"+[Root classGetComplex]"(ptr noundef
     // CHECK: call i64 @"+[Root classGetComplex]_thunk"(ptr noundef
     [Root classGetComplex].a +
-    // CHECK-NOT: call void @"+[Root classGetAggregate]"(ptr {{.*}}sret
-    // CHECK: call void @"+[Root classGetAggregate]_thunk"(ptr {{.*}}sret
+    // CHECK-NOT: call void @"+[Root classGetAggregate]_thunk"(ptr {{.*}}sret
+    // CHECK: call void @"+[Root classGetAggregate]"(ptr {{.*}}sret
     [Root classGetAggregate].a
   );
 }
@@ -291,4 +290,4 @@ int useSRet(Root *r) {
 // CHECK:   ret void
 
 // CHECK: define {{.*}} @"+[Root classGetComplex]_thunk"
-// CHECK: define {{.*}} @"+[Root classGetAggregate]_thunk"
+// CHECK-NOT: define {{.*}} @"+[Root classGetAggregate]_thunk"