[arcmt] When a NSData's 'bytes' family of methods are used on a local var,

add __attribute__((objc_precise_lifetime)) to make sure that the object
(and its data) will not get released before the var goes out-of-scope.

rdar://9206226

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

1 files changed, 34 insertions, 0 deletions

diff --git a/lib/ARCMigrate/TransAPIUses.cpp b/lib/ARCMigrate/TransAPIUses.cpp
index 58fd3d07d0..296aa07000 100644
--- a/lib/ARCMigrate/TransAPIUses.cpp
+++ b/lib/ARCMigrate/TransAPIUses.cpp
@@ -13,6 +13,8 @@
 //
 // - NSInvocation's [get/set]ReturnValue and [get/set]Argument are only safe
 //   with __unsafe_unretained objects.
+// - When a NSData's 'bytes' family of methods are used on a local var,
+//   add __attribute__((objc_precise_lifetime)) to make it safer.
 //
 //===----------------------------------------------------------------------===//
 
@@ -28,9 +30,13 @@ namespace {
 
 class APIChecker : public RecursiveASTVisitor<APIChecker> {
   MigrationPass &Pass;
+
   Selector getReturnValueSel, setReturnValueSel;
   Selector getArgumentSel, setArgumentSel;
 
+  Selector bytesSel, getBytesSel, getBytesLengthSel, getBytesRangeSel;
+
+  llvm::DenseSet<VarDecl *> ChangedNSDataVars;
 public:
   APIChecker(MigrationPass &pass) : Pass(pass) {
     SelectorTable &sels = Pass.Ctx.Selectors;
@@ -44,6 +50,14 @@ public:
     getArgumentSel = sels.getSelector(2, selIds);
     selIds[0] = &ids.get("setArgument");
     setArgumentSel = sels.getSelector(2, selIds);
+
+    bytesSel = sels.getNullarySelector(&ids.get("bytes"));
+    getBytesSel = sels.getUnarySelector(&ids.get("getBytes"));
+    selIds[0] = &ids.get("getBytes");
+    selIds[1] = &ids.get("length");
+    getBytesLengthSel = sels.getSelector(2, selIds);
+    selIds[1] = &ids.get("range");
+    getBytesRangeSel = sels.getSelector(2, selIds);
   }
 
   bool VisitObjCMessageExpr(ObjCMessageExpr *E) {
@@ -78,6 +92,26 @@ public:
       return true;
     }
 
+    if (E->isInstanceMessage() &&
+        E->getReceiverInterface() &&
+        E->getReceiverInterface()->getName() == "NSData" &&
+        E->getInstanceReceiver() &&
+        (E->getSelector() == bytesSel ||
+         E->getSelector() == getBytesSel ||
+         E->getSelector() == getBytesLengthSel ||
+         E->getSelector() == getBytesRangeSel)) {
+      Expr *rec = E->getInstanceReceiver();
+      rec = rec->IgnoreParenCasts();
+      if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(rec))
+        if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
+          if (VD->hasLocalStorage() && !ChangedNSDataVars.count(VD)) {
+            Transaction Trans(Pass.TA);
+            Pass.TA.insertAfterToken(VD->getLocation(),
+                                     " __attribute__((objc_precise_lifetime))");
+            ChangedNSDataVars.insert(VD);
+          }
+    }
+
     return true;
   }
 };