[ObjCDirect] Move nil check to a thunk function

Signed-off-by: Peter Rong <[email protected]>

Author: DataCorrupted

clang/include/clang/AST/DeclObjC.h

@@ -482,6 +482,12 @@ class ObjCMethodDecl : public NamedDecl, public DeclContext {
   /// True if the method is tagged as objc_direct
   bool isDirectMethod() const;
 
+  // Only direct instance method that have a fixed number of arguments can have
+  // nil check thunk functions.
+  bool canHaveNilCheckThunk() const {
+    return isDirectMethod() && isInstanceMethod() && !isVariadic();
+  }
+
   /// True if the method has a parameter that's destroyed in the callee.
   bool hasParamDestroyedInCallee() const;
 

clang/include/clang/AST/Mangle.h

@@ -40,6 +40,13 @@ namespace clang {
   struct ThunkInfo;
   class VarDecl;
 
+/// Extract mangling function name from MangleContext such that swift can call
+/// it to prepare for ObjCDirect in swift.
+void mangleObjCMethodName(raw_ostream &OS, bool includePrefixByte,
+                          bool isInstanceMethod, StringRef ClassName,
+                          std::optional<StringRef> CategoryName,
+                          StringRef MethodName, bool isThunk);
+
 /// MangleContext - Context for tracking state which persists across multiple
 /// calls to the C++ name mangler.
 class MangleContext {
@@ -157,7 +164,8 @@ class MangleContext {
 
   void mangleObjCMethodName(const ObjCMethodDecl *MD, raw_ostream &OS,
                             bool includePrefixByte = true,
-                            bool includeCategoryNamespace = true);
+                            bool includeCategoryNamespace = true,
+                            bool isThunk = true);
   void mangleObjCMethodNameAsSourceName(const ObjCMethodDecl *MD,
                                         raw_ostream &);
 

clang/include/clang/Basic/CodeGenOptions.def

@@ -356,6 +356,7 @@ CODEGENOPT(EmitLLVMUseLists, 1, 0) ///< Control whether to serialize use-lists.
 
 CODEGENOPT(WholeProgramVTables, 1, 0) ///< Whether to apply whole-program
                                       ///  vtable optimization.
+CODEGENOPT(ObjCEmitNilCheckThunk , 1, 0) ///< Whether objc_direct methods should emit a nil check thunk.
 
 CODEGENOPT(VirtualFunctionElimination, 1, 0) ///< Whether to apply the dead
                                              /// virtual function elimination

clang/include/clang/Basic/LangOptions.def

@@ -371,6 +371,7 @@ LANGOPT(IncludeDefaultHeader, 1, 0, "Include default header file for OpenCL")
 LANGOPT(DeclareOpenCLBuiltins, 1, 0, "Declare OpenCL builtin functions")
 BENIGN_LANGOPT(DelayedTemplateParsing , 1, 0, "delayed template parsing")
 LANGOPT(BlocksRuntimeOptional , 1, 0, "optional blocks runtime")
+LANGOPT(ObjCEmitNilCheckThunk, 1, 0, "Emit a thunk to do nil check for objc direct methods")
 LANGOPT(
     CompleteMemberPointers, 1, 0,
     "Require member pointer base types to be complete at the point where the "

clang/include/clang/Driver/Options.td

@@ -2990,6 +2990,8 @@ def fthin_link_bitcode_EQ : Joined<["-"], "fthin-link-bitcode=">,
   Visibility<[ClangOption, CLOption, CC1Option]>, Group<f_Group>,
   HelpText<"Write minimized bitcode to <file> for the ThinLTO thin link only">,
   MarshallingInfoString<CodeGenOpts<"ThinLinkBitcodeFile">>;
+def fobjc_emit_nil_check_thunk:
+  Flag<["-"], "fobjc-emit-nil-check-thunk">, Visibility<[ClangOption, CC1Option]>, Group<f_Group>;
 defm fat_lto_objects : BoolFOption<"fat-lto-objects",
   CodeGenOpts<"FatLTO">, DefaultFalse,
   PosFlag<SetTrue, [], [ClangOption, CC1Option], "Enable">,

clang/lib/AST/Mangle.cpp

@@ -29,6 +29,24 @@
 
 using namespace clang;
 
+void clang::mangleObjCMethodName(raw_ostream &OS, bool includePrefixByte,
+                                 bool isInstanceMethod, StringRef ClassName,
+                                 std::optional<StringRef> CategoryName,
+                                 StringRef MethodName, bool isThunk) {
+  // \01+[ContainerName(CategoryName) SelectorName]
+  if (includePrefixByte)
+    OS << "\01";
+  OS << (isInstanceMethod ? '-' : '+');
+  OS << '[';
+  OS << ClassName;
+  if (CategoryName)
+    OS << "(" << CategoryName.value() << ")";
+  OS << " ";
+  OS << MethodName;
+  OS << ']';
+  if (!isThunk)
+    OS << "_inner";
+}
 // FIXME: For blocks we currently mimic GCC's mangling scheme, which leaves
 // much to be desired. Come up with a better mangling scheme.
 
@@ -327,7 +345,8 @@ void MangleContext::mangleBlock(const DeclContext *DC, const BlockDecl *BD,
 void MangleContext::mangleObjCMethodName(const ObjCMethodDecl *MD,
                                          raw_ostream &OS,
                                          bool includePrefixByte,
-                                         bool includeCategoryNamespace) {
+                                         bool includeCategoryNamespace,
+                                         bool isThunk) {
   if (getASTContext().getLangOpts().ObjCRuntime.isGNUFamily()) {
     // This is the mangling we've always used on the GNU runtimes, but it
     // has obvious collisions in the face of underscores within class
@@ -361,24 +380,24 @@ void MangleContext::mangleObjCMethodName(const ObjCMethodDecl *MD,
   }
 
   // \01+[ContainerName(CategoryName) SelectorName]
-  if (includePrefixByte) {
-    OS << '\01';
-  }
-  OS << (MD->isInstanceMethod() ? '-' : '+') << '[';
+  auto CategoryName = std::optional<StringRef>();
+  StringRef ClassName;
   if (const auto *CID = MD->getCategory()) {
-    OS << CID->getClassInterface()->getName();
-    if (includeCategoryNamespace) {
-      OS << '(' << *CID << ')';
-    }
+    if (includeCategoryNamespace)
+      CategoryName = CID->getName();
+    ClassName = CID->getClassInterface()->getName();
+
   } else if (const auto *CD =
                  dyn_cast<ObjCContainerDecl>(MD->getDeclContext())) {
-    OS << CD->getName();
+    ClassName = CD->getName();
   } else {
     llvm_unreachable("Unexpected ObjC method decl context");
   }
-  OS << ' ';
-  MD->getSelector().print(OS);
-  OS << ']';
+  std::string MethodName;
+  llvm::raw_string_ostream MethodNameOS(MethodName);
+  MD->getSelector().print(MethodNameOS);
+  clang::mangleObjCMethodName(OS, includePrefixByte, MD->isInstanceMethod(),
+                              ClassName, CategoryName, MethodName, isThunk);
 }
 
 void MangleContext::mangleObjCMethodNameAsSourceName(const ObjCMethodDecl *MD,

clang/lib/CodeGen/CGCall.cpp

@@ -2704,6 +2704,18 @@ void CodeGenModule::ConstructAttributeList(StringRef Name,
 
     ArgAttrs[IRArgs.first] = llvm::AttributeSet::get(getLLVMContext(), Attrs);
   }
+  // Direct method prologue should not contain nil check anymore.
+  // As a result, we can set `self` to be NonNull to prepare for further
+  // optimizations.
+  if (getLangOpts().ObjCEmitNilCheckThunk && TargetDecl) {
+    auto OMD = dyn_cast<ObjCMethodDecl>(TargetDecl);
+    bool isDirect = OMD && OMD->isDirectMethod();
+    if (isDirect && !IsThunk) {
+      auto IRArgs = IRFunctionArgs.getIRArgs(0);
+      ArgAttrs[IRArgs.first] = ArgAttrs[IRArgs.first].addAttribute(
+          getLLVMContext(), llvm::Attribute::NonNull);
+    }
+  }
 
   unsigned ArgNo = 0;
   for (CGFunctionInfo::const_arg_iterator I = FI.arg_begin(),
@@ -3388,7 +3400,13 @@ void CodeGenFunction::EmitFunctionProlog(const CGFunctionInfo &FI,
     }
   }
 
-  if (getTarget().getCXXABI().areArgsDestroyedLeftToRightInCallee()) {
+  if (InnerFn) {
+    // Don't emit any arg except for `self` if we are in a thunk function.
+    // We still need self for nil check, other arguments aren't used in this
+    // function and thus is not needed. Avoid emitting them also prevents
+    // accidental release/retain.
+    EmitParmDecl(*Args[0], ArgVals[0], 1);
+  } else if (getTarget().getCXXABI().areArgsDestroyedLeftToRightInCallee()) {
     for (int I = Args.size() - 1; I >= 0; --I)
       EmitParmDecl(*Args[I], ArgVals[I], I + 1);
   } else {

clang/lib/CodeGen/CGObjC.cpp

@@ -756,12 +756,13 @@ void CodeGenFunction::StartObjCMethod(const ObjCMethodDecl *OMD,
   if (OMD->hasAttr<NoDebugAttr>())
     DebugInfo = nullptr; // disable debug info indefinitely for this function
 
-  llvm::Function *Fn = CGM.getObjCRuntime().GenerateMethod(OMD, CD);
+  bool isInner = CGM.shouldHaveNilCheckThunk(OMD) && !InnerFn;
+  llvm::Function *Fn = CGM.getObjCRuntime().GenerateMethod(OMD, CD, !isInner);
 
   const CGFunctionInfo &FI = CGM.getTypes().arrangeObjCMethodDeclaration(OMD);
   if (OMD->isDirectMethod()) {
     Fn->setVisibility(llvm::Function::HiddenVisibility);
-    CGM.SetLLVMFunctionAttributes(OMD, FI, Fn, /*IsThunk=*/false);
+    CGM.SetLLVMFunctionAttributes(OMD, FI, Fn, /*IsThunk=*/InnerFn);
     CGM.SetLLVMFunctionAttributesForDefinition(OMD, Fn);
   } else {
     CGM.SetInternalFunctionAttributes(OMD, Fn, FI);
@@ -780,11 +781,21 @@ void CodeGenFunction::StartObjCMethod(const ObjCMethodDecl *OMD,
                 OMD->getLocation(), StartLoc);
 
   if (OMD->isDirectMethod()) {
-    // This function is a direct call, it has to implement a nil check
-    // on entry.
-    //
-    // TODO: possibly have several entry points to elide the check
-    CGM.getObjCRuntime().GenerateDirectMethodPrologue(*this, Fn, OMD, CD);
+    if (CGM.getLangOpts().ObjCRuntime.isNeXTFamily()) {
+      // Having `InnerFn` indicates that we are generating a nil check thunk.
+      // In that case our job is done here.
+      if (InnerFn)
+        return;
+      if (CGM.shouldHaveNilCheckThunk(OMD))
+        // Go generate  a nil check thunk around `Fn`
+        CodeGenFunction(CGM, /*InnerFn=*/Fn).GenerateObjCDirectThunk(OMD, CD);
+      else
+        CGM.getObjCRuntime().GenerateObjCDirectNilCheck(*this, OMD, CD);
+      CGM.getObjCRuntime().GenerateCmdIfNecessary(*this, OMD);
+    } else {
+      // For GNU family, since GNU Step2 also supports direct methods now.
+      CGM.getObjCRuntime().GenerateDirectMethodPrologue(*this, Fn, OMD, CD);
+    }
   }
 
   // In ARC, certain methods get an extra cleanup.
@@ -1637,6 +1648,44 @@ void CodeGenFunction::GenerateObjCSetter(ObjCImplementationDecl *IMP,
   FinishFunction(OMD->getEndLoc());
 }
 
+void CodeGenFunction::GenerateObjCDirectThunk(const ObjCMethodDecl *OMD,
+                                              const ObjCContainerDecl *CD) {
+  assert(InnerFn && CGM.shouldHaveNilCheckThunk(OMD) &&
+         "Should only generate wrapper when the flag is set.");
+  StartObjCMethod(OMD, CD);
+
+  // Manually pop all the clean up that doesn't need to happen in the outer
+  // function. InnerFn will do this for us.
+  while (EHStack.stable_begin() != PrologueCleanupDepth)
+    EHStack.popCleanup();
+
+  // Generate a nil check.
+  CGM.getObjCRuntime().GenerateObjCDirectNilCheck(*this, OMD, CD);
+  // Call the InnerFn and pass the return value
+  SmallVector<llvm::Value *> Args(CurFn->arg_size());
+  std::transform(CurFn->arg_begin(), CurFn->arg_end(), Args.begin(),
+                 [](llvm::Argument &arg) { return &arg; });
+
+  // This will be optimized into a tail call.
+  auto *CallInst = EmitCallOrInvoke(InnerFn, Args);
+  // Preserve the inner function's attributes to the call instruction.
+  CallInst->setAttributes(InnerFn->getAttributes());
+  llvm::Value *RetVal = CallInst;
+
+  // If `AutoreleaseResult` is set, the return value is not void.
+  if (AutoreleaseResult)
+    RetVal = EmitARCRetainAutoreleasedReturnValue(RetVal);
+
+  // This excessive store is totally unnecessary.
+  // But `FinishFunction` really wants us to store the result so it can
+  // clean up the function properly.
+  // The unnecessary store-load of the ret value will be optimized out anyway.
+  if (!CurFn->getReturnType()->isVoidTy())
+    Builder.CreateStore(RetVal, ReturnValue);
+
+  // Nil check's end location is the function's start location.
+  FinishFunction(OMD->getBeginLoc());
+}
 namespace {
   struct DestroyIvar final : EHScopeStack::Cleanup {
   private:

clang/lib/CodeGen/CGObjCGNU.cpp

@@ -594,8 +594,23 @@ class CGObjCGNU : public CGObjCRuntime {
   }
   llvm::Constant *GetEHType(QualType T) override;
 
+  void GenerateObjCDirectNilCheck(CodeGenFunction &CGF,
+                                  const ObjCMethodDecl *OMD,
+                                  const ObjCContainerDecl *CD) override {
+    // GNU runtime doesn't support nil check thunks at this time
+  };
+  void GenerateCmdIfNecessary(CodeGenFunction &CGF,
+                              const ObjCMethodDecl *OMD) override {
+    // GNU runtime doesn't support nil check thunks at this time
+  }
+  llvm::Function *GenerateMethod(const ObjCMethodDecl *OMD,
+                                 const ObjCContainerDecl *CD,
+                                 bool isThunk) override {
+    // isThunk is irrelevent for GNU.
+    return GenerateMethod(OMD, CD);
+  };
   llvm::Function *GenerateMethod(const ObjCMethodDecl *OMD,
-                                 const ObjCContainerDecl *CD) override;
+                                 const ObjCContainerDecl *CD);
 
   // Map to unify direct method definitions.
   llvm::DenseMap<const ObjCMethodDecl *, llvm::Function *>