[cxx-interop][IRGen] Do not try to retain/release a null pointer

egorzhdan · egorzhdan · commit bd6da5dc645b · 2024-05-07T20:20:45.000+01:00
This teaches IRGen to only emit a lifetime operation (retain or release) for a C++ foreign reference type if the pointer is not `nullptr`.

Previously the compiler would in some cases emit a release call for `nullptr`, which breaks the assumption that the argument to a custom release function is `_Nonnull`. For instance:
```
var globalOptional: MyRefType? = nil
func foo() { globalOptional = MyRefType.create() }
```
When emitting IR for the assignment operation to `globalOptional`, the compiler would emit code to first retrieve the existing value of `globalOptional` and release it. If the value is `nil`, it does not need to be released.

rdar://97532642
diff --git a/lib/IRGen/ClassTypeInfo.h b/lib/IRGen/ClassTypeInfo.h
@@ -99,38 +99,56 @@ class ClassTypeInfo : public HeapTypeInfo<ClassTypeInfo> {
     HeapTypeInfo::emitScalarRetain(IGF, value, atomicity);
   }
 
+  void strongCustomRetain(IRGenFunction &IGF, Explosion &e,
+                          bool needsNullCheck) const {
+    assert(getReferenceCounting() == ReferenceCounting::Custom &&
+           "only supported for custom ref-counting");
+
+    llvm::Value *value = e.claimNext();
+    auto retainFn =
+        evaluateOrDefault(
+            getClass()->getASTContext().evaluator,
+            CustomRefCountingOperation(
+                {getClass(), CustomRefCountingOperationKind::retain}),
+            {})
+            .operation;
+    IGF.emitForeignReferenceTypeLifetimeOperation(retainFn, value,
+                                                  needsNullCheck);
+  }
+
   // Implement the primary retain/release operations of ReferenceTypeInfo
   // using basic reference counting.
   void strongRetain(IRGenFunction &IGF, Explosion &e,
                     Atomicity atomicity) const override {
     if (getReferenceCounting() == ReferenceCounting::Custom) {
-      llvm::Value *value = e.claimNext();
-      auto retainFn =
-          evaluateOrDefault(
-              getClass()->getASTContext().evaluator,
-              CustomRefCountingOperation(
-                  {getClass(), CustomRefCountingOperationKind::retain}),
-              {})
-              .operation;
-      IGF.emitForeignReferenceTypeLifetimeOperation(retainFn, value);
+      strongCustomRetain(IGF, e, /*needsNullCheck*/ false);
       return;
     }
 
     HeapTypeInfo::strongRetain(IGF, e, atomicity);
   }
 
+  void strongCustomRelease(IRGenFunction &IGF, Explosion &e,
+                           bool needsNullCheck) const {
+    assert(getReferenceCounting() == ReferenceCounting::Custom &&
+           "only supported for custom ref-counting");
+
+    llvm::Value *value = e.claimNext();
+    auto releaseFn =
+        evaluateOrDefault(
+            getClass()->getASTContext().evaluator,
+            CustomRefCountingOperation(
+                {getClass(), CustomRefCountingOperationKind::release}),
+            {})
+            .operation;
+    IGF.emitForeignReferenceTypeLifetimeOperation(releaseFn, value,
+                                                  needsNullCheck);
+  }
+
   void strongRelease(IRGenFunction &IGF, Explosion &e,
                      Atomicity atomicity) const override {
     if (getReferenceCounting() == ReferenceCounting::Custom) {
-      llvm::Value *value = e.claimNext();
-      auto releaseFn =
-          evaluateOrDefault(
-              getClass()->getASTContext().evaluator,
-              CustomRefCountingOperation(
-                  {getClass(), CustomRefCountingOperationKind::release}),
-              {})
-              .operation;
-      IGF.emitForeignReferenceTypeLifetimeOperation(releaseFn, value);
+      strongCustomRelease(IGF, e, /*needsNullCheck*/ false);
       return;
     }
 
diff --git a/lib/IRGen/GenEnum.cpp b/lib/IRGen/GenEnum.cpp
@@ -2668,7 +2668,8 @@ namespace {
         if (Refcounting == ReferenceCounting::Custom) {
           Explosion e;
           e.add(ptr);
-          getPayloadTypeInfo().as<ClassTypeInfo>().strongRetain(IGF, e, IGF.getDefaultAtomicity());
+          getPayloadTypeInfo().as<ClassTypeInfo>().strongCustomRetain(
+              IGF, e, /*needsNullCheck*/ true);
           return;
         }
 
@@ -2704,7 +2705,8 @@ namespace {
         if (Refcounting == ReferenceCounting::Custom) {
           Explosion e;
           e.add(ptr);
-          getPayloadTypeInfo().as<ClassTypeInfo>().strongRelease(IGF, e, IGF.getDefaultAtomicity());
+          getPayloadTypeInfo().as<ClassTypeInfo>().strongCustomRelease(
+              IGF, e, /*needsNullCheck*/ true);
           return;
         }
 
diff --git a/lib/IRGen/GenObjC.cpp b/lib/IRGen/GenObjC.cpp
@@ -1712,7 +1712,7 @@ void IRGenFunction::emitBlockRelease(llvm::Value *value) {
 }
 
 void IRGenFunction::emitForeignReferenceTypeLifetimeOperation(
-    ValueDecl *fn, llvm::Value *value) {
+    ValueDecl *fn, llvm::Value *value, bool needsNullCheck) {
   assert(fn->getClangDecl() && isa<clang::FunctionDecl>(fn->getClangDecl()));
 
   auto clangFn = cast<clang::FunctionDecl>(fn->getClangDecl());
@@ -1723,6 +1723,28 @@ void IRGenFunction::emitForeignReferenceTypeLifetimeOperation(
       cast<llvm::FunctionType>(llvmFn->getFunctionType())->getParamType(0);
   value = Builder.CreateBitCast(value, argType);
 
-  auto call = Builder.CreateCall(llvmFn->getFunctionType(), llvmFn, value);
+  llvm::CallInst *call = nullptr;
+  if (needsNullCheck) {
+    // Check if the pointer is null.
+    auto nullValue = llvm::Constant::getNullValue(argType);
+    auto hasValue = Builder.CreateICmpNE(value, nullValue);
+
+    auto nonNullValueBB = createBasicBlock("lifetime.nonnull-value");
+    auto contBB = createBasicBlock("lifetime.cont");
+
+    // If null, just continue.
+    Builder.CreateCondBr(hasValue, nonNullValueBB, contBB);
+
+    // If non-null, emit a call to release/retain function.
+    Builder.emitBlock(nonNullValueBB);
+    call = Builder.CreateCall(llvmFn->getFunctionType(), llvmFn, value);
+
+    Builder.CreateBr(contBB);
+
+    Builder.emitBlock(contBB);
+  } else {
+    call = Builder.CreateCall(llvmFn->getFunctionType(), llvmFn, value);
+  }
+
   call->setDoesNotThrow();
 }
diff --git a/lib/IRGen/IRGenFunction.h b/lib/IRGen/IRGenFunction.h
@@ -575,7 +575,8 @@ class IRGenFunction {
   void emitBlockRelease(llvm::Value *value);
 
   void emitForeignReferenceTypeLifetimeOperation(ValueDecl *fn,
-                                                 llvm::Value *value);
+                                                 llvm::Value *value,
+                                                 bool needsNullCheck = false);
 
   // Routines for an unknown reference-counting style (meaning,
   // dynamically something compatible with either the ObjC or Swift styles).
diff --git a/test/Interop/Cxx/foreign-reference/Inputs/reference-counted.h b/test/Interop/Cxx/foreign-reference/Inputs/reference-counted.h
@@ -46,6 +46,21 @@ __attribute__((swift_attr("release:GCRelease"))) GlobalCount {
 inline void GCRetain(GlobalCount *x) { globalCount++; }
 inline void GCRelease(GlobalCount *x) { globalCount--; }
 
+struct __attribute__((swift_attr("import_as_ref")))
+__attribute__((swift_attr("retain:GCRetainNullableInit")))
+__attribute__((swift_attr("release:GCReleaseNullableInit")))
+GlobalCountNullableInit {
+  static GlobalCountNullableInit *_Nullable create(bool wantNullptr) {
+    if (wantNullptr)
+      return nullptr;
+    return new (malloc(sizeof(GlobalCountNullableInit)))
+        GlobalCountNullableInit();
+  }
+};
+
+inline void GCRetainNullableInit(GlobalCountNullableInit *x) { globalCount++; }
+inline void GCReleaseNullableInit(GlobalCountNullableInit *x) { globalCount--; }
+
 SWIFT_END_NULLABILITY_ANNOTATIONS
 
 #endif // TEST_INTEROP_CXX_FOREIGN_REFERENCE_INPUTS_REFERENCE_COUNTED_H
diff --git a/test/Interop/Cxx/foreign-reference/reference-counted-irgen.swift b/test/Interop/Cxx/foreign-reference/reference-counted-irgen.swift
@@ -0,0 +1,53 @@
+// RUN: %target-swift-emit-irgen %s -I %S/Inputs -cxx-interoperability-mode=default -Xcc -fignore-exceptions -disable-availability-checking | %FileCheck %s
+// XFAIL: OS=linux-android, OS=linux-androideabi
+
+import ReferenceCounted
+
+
+public func getLocalCount() -> NS.LocalCount {
+    let result = NS.LocalCount.create()
+    return result
+}
+
+// CHECK:      define {{.*}}swiftcc ptr @"$s4main13getLocalCountSo2NSO0cD0VyF"()
+// CHECK-NEXT: entry:
+// CHECK:        %0 = call ptr @{{_ZN2NS10LocalCount6createEv|"\?create\@LocalCount\@NS\@\@SAPEAU12\@XZ"}}()
+// CHECK-NEXT:   call void @{{_Z8LCRetainPN2NS10LocalCountE|"\?LCRetain\@\@YAXPEAULocalCount\@NS\@\@\@Z"}}(ptr %0)
+// CHECK:        ret ptr %0
+// CHECK-NEXT: }
+
+
+public func get42() -> Int32 {
+    let result = NS.LocalCount.create()
+    return result.returns42()
+}
+
+// CHECK:      define {{.*}}swiftcc i32 @"$s4main5get42s5Int32VyF"()
+// CHECK-NEXT: entry:
+// CHECK:        %0 = call ptr @{{_ZN2NS10LocalCount6createEv|"\?create\@LocalCount\@NS\@\@SAPEAU12\@XZ"}}()
+// CHECK-NEXT:   call void @{{_Z8LCRetainPN2NS10LocalCountE|"\?LCRetain\@\@YAXPEAULocalCount\@NS\@\@\@Z"}}(ptr %0)
+// CHECK:        %1 = call i32 @{{_ZN2NS10LocalCount9returns42Ev|"\?returns42\@LocalCount\@NS\@\@QEAAHXZ"}}
+// CHECK:        ret i32 %1
+// CHECK-NEXT: }
+
+
+public func getNullable(wantNullptr: Bool) -> GlobalCountNullableInit? {
+    let result = GlobalCountNullableInit.create(wantNullptr)
+    return result
+}
+
+// CHECK:      define {{.*}}swiftcc i64 @"$s4main11getNullable11wantNullptrSo011GlobalCountC4InitVSgSb_tF"(i1 %0)
+// CHECK-NEXT: entry:
+// CHECK:        %1 = call ptr @{{_ZN23GlobalCountNullableInit6createEb|"\?create\@GlobalCountNullableInit\@\@SAPEAU1\@_N\@Z"}}
+// CHECK-NEXT:   %2 = ptrtoint ptr %1 to i64
+// CHECK-NEXT:   %3 = inttoptr i64 %2 to ptr
+// CHECK-NEXT:   %4 = icmp ne ptr %3, null
+// CHECK-NEXT:   br i1 %4, label %lifetime.nonnull-value, label %lifetime.cont
+
+// CHECK:      lifetime.nonnull-value:
+// CHECK-NEXT:   call void @{{_Z20GCRetainNullableInitP23GlobalCountNullableInit|"\?GCRetainNullableInit\@\@YAXPEAUGlobalCountNullableInit\@\@\@Z"}}(ptr %3)
+// CHECK-NEXT:   br label %lifetime.cont
+
+// CHECK:      lifetime.cont:
+// CHECK:          ret i64 %2
+// CHECK-NEXT: }
