[arcmt] Use __bridge_retained when passing an objc object to a CF parameter

annotated with cf_consumed attribute.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@139709 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 40 insertions, 0 deletions

diff --git a/lib/ARCMigrate/TransUnbridgedCasts.cpp b/lib/ARCMigrate/TransUnbridgedCasts.cpp
index bc33c53b75..047dfb276a 100644
--- a/lib/ARCMigrate/TransUnbridgedCasts.cpp
+++ b/lib/ARCMigrate/TransUnbridgedCasts.cpp
@@ -242,6 +242,11 @@ private:
       if (implCE->getCastKind() == CK_ARCReclaimReturnedObject)
         return rewriteToBridgedCast(E, OBC_Bridge);
     }
+
+    bool isConsumed = false;
+    if (isPassedToCParamWithKnownOwnership(E, isConsumed))
+      return rewriteToBridgedCast(E, isConsumed ? OBC_BridgeRetained
+                                                : OBC_Bridge);
   }
 
   static ObjCMethodFamily getFamilyOfMessage(Expr *E) {
@@ -265,6 +270,29 @@ private:
     return false;
   }
 
+  bool isPassedToCParamWithKnownOwnership(Expr *E, bool &isConsumed) const {
+    if (CallExpr *callE = dyn_cast_or_null<CallExpr>(
+                                     StmtMap->getParentIgnoreParenImpCasts(E)))
+      if (FunctionDecl *
+            FD = dyn_cast_or_null<FunctionDecl>(callE->getCalleeDecl())) {
+        unsigned i = 0;
+        for (unsigned e = callE->getNumArgs(); i != e; ++i) {
+          Expr *arg = callE->getArg(i);
+          if (arg == E || arg->IgnoreParenImpCasts() == E)
+            break;
+        }
+        if (i < callE->getNumArgs()) {
+          ParmVarDecl *PD = FD->getParamDecl(i);
+          if (PD->getAttr<CFConsumedAttr>()) {
+            isConsumed = true;
+            return true;
+          }
+        }
+      }
+
+    return false;
+  }
+
   bool isSelf(Expr *E) const {
     E = E->IgnoreParenLValueCasts();
     if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(E))
diff --git a/test/ARCMT/Common.h b/test/ARCMT/Common.h
index f10ae0d2a7..2603730cad 100644
--- a/test/ARCMT/Common.h
+++ b/test/ARCMT/Common.h
@@ -4,6 +4,8 @@
 #define NS_AUTOMATED_REFCOUNT_UNAVAILABLE
 #endif
 
+#define CF_CONSUMED __attribute__((cf_consumed))
+
 #define nil ((void*) 0)
 
 typedef int BOOL;
diff --git a/test/ARCMT/nonobjc-to-objc-cast.m b/test/ARCMT/nonobjc-to-objc-cast.m
index b4e4080149..080399dfb1 100644
--- a/test/ARCMT/nonobjc-to-objc-cast.m
+++ b/test/ARCMT/nonobjc-to-objc-cast.m
@@ -40,6 +40,8 @@ void f(BOOL b, id p) {
 }
 @end
 
+extern void consumeParam(CFStringRef CF_CONSUMED p);
+
 void f2(NSString *s) {
   CFStringRef ref = [s string];
   ref = (CFStringRef)[s string];
@@ -53,4 +55,7 @@ void f2(NSString *s) {
   ref = CFRetain([s string]);
   ref = CFRetain(s);
   ref = [s retain];
+
+  consumeParam((CFStringRef)s);
+  consumeParam(s);
 }
diff --git a/test/ARCMT/nonobjc-to-objc-cast.m.result b/test/ARCMT/nonobjc-to-objc-cast.m.result
index 1fc2774c14..1e9dab979f 100644
--- a/test/ARCMT/nonobjc-to-objc-cast.m.result
+++ b/test/ARCMT/nonobjc-to-objc-cast.m.result
@@ -40,6 +40,8 @@ void f(BOOL b, id p) {
 }
 @end
 
+extern void consumeParam(CFStringRef CF_CONSUMED p);
+
 void f2(NSString *s) {
   CFStringRef ref = (__bridge CFStringRef)([s string]);
   ref = (__bridge CFStringRef)[s string];
@@ -53,4 +55,7 @@ void f2(NSString *s) {
   ref = (__bridge_retained CFTypeRef)([s string]);
   ref = (__bridge_retained CFTypeRef)(s);
   ref = (__bridge_retained CFStringRef)(s);
+
+  consumeParam((__bridge_retained CFStringRef)s);
+  consumeParam((__bridge_retained CFStringRef)(s));
 }