[arcmt]

-Make sure we don't change to '__weak' a __block variable used as output.
-Make sure we don't apply __weak twice.

Fixes rdar://10520757&10521362

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

3 files changed, 39 insertions, 27 deletions

diff --git a/lib/ARCMigrate/TransBlockObjCVariable.cpp b/lib/ARCMigrate/TransBlockObjCVariable.cpp
index 48c0ca9cef..5f9592739d 100644
--- a/lib/ARCMigrate/TransBlockObjCVariable.cpp
+++ b/lib/ARCMigrate/TransBlockObjCVariable.cpp
@@ -38,7 +38,7 @@ namespace {
 class RootBlockObjCVarRewriter :
                           public RecursiveASTVisitor<RootBlockObjCVarRewriter> {
   MigrationPass &Pass;
-  llvm::DenseSet<VarDecl *> CheckedVars;
+  llvm::DenseSet<VarDecl *> &VarsToChange;
 
   class BlockVarChecker : public RecursiveASTVisitor<BlockVarChecker> {
     VarDecl *Var;
@@ -71,7 +71,9 @@ class RootBlockObjCVarRewriter :
   };
 
 public:
-  RootBlockObjCVarRewriter(MigrationPass &pass) : Pass(pass) { }
+  RootBlockObjCVarRewriter(MigrationPass &pass,
+                           llvm::DenseSet<VarDecl *> &VarsToChange)
+    : Pass(pass), VarsToChange(VarsToChange) { }
 
   bool VisitBlockDecl(BlockDecl *block) {
     SmallVector<VarDecl *, 4> BlockVars;
@@ -80,7 +82,6 @@ public:
            I = block->capture_begin(), E = block->capture_end(); I != E; ++I) {
       VarDecl *var = I->getVariable();
       if (I->isByRef() &&
-          !isAlreadyChecked(var) &&
           var->getType()->isObjCObjectPointerType() &&
           isImplicitStrong(var->getType())) {
         BlockVars.push_back(var);
@@ -89,32 +90,19 @@ public:
 
     for (unsigned i = 0, e = BlockVars.size(); i != e; ++i) {
       VarDecl *var = BlockVars[i];
-      CheckedVars.insert(var);
 
       BlockVarChecker checker(var);
       bool onlyValueOfVarIsNeeded = checker.TraverseStmt(block->getBody());
-      if (onlyValueOfVarIsNeeded) {
-        BlocksAttr *attr = var->getAttr<BlocksAttr>();
-        if(!attr)
-          continue;
-        bool useWeak = canApplyWeak(Pass.Ctx, var->getType());
-        SourceManager &SM = Pass.Ctx.getSourceManager();
-        Transaction Trans(Pass.TA);
-        Pass.TA.replaceText(SM.getExpansionLoc(attr->getLocation()),
-                            "__block",
-                            useWeak ? "__weak" : "__unsafe_unretained");
-      }
-
+      if (onlyValueOfVarIsNeeded)
+        VarsToChange.insert(var);
+      else
+        VarsToChange.erase(var);
     }
 
     return true;
   }
 
 private:
-  bool isAlreadyChecked(VarDecl *VD) {
-    return CheckedVars.count(VD);
-  }
-
   bool isImplicitStrong(QualType ty) {
     if (isa<AttributedType>(ty.getTypePtr()))
       return false;
@@ -124,19 +112,39 @@ private:
 
 class BlockObjCVarRewriter : public RecursiveASTVisitor<BlockObjCVarRewriter> {
   MigrationPass &Pass;
+  llvm::DenseSet<VarDecl *> &VarsToChange;
 
 public:
-  BlockObjCVarRewriter(MigrationPass &pass) : Pass(pass) { }
+  BlockObjCVarRewriter(MigrationPass &pass,
+                       llvm::DenseSet<VarDecl *> &VarsToChange)
+    : Pass(pass), VarsToChange(VarsToChange) { }
 
   bool TraverseBlockDecl(BlockDecl *block) {
-    RootBlockObjCVarRewriter(Pass).TraverseDecl(block);
+    RootBlockObjCVarRewriter(Pass, VarsToChange).TraverseDecl(block);
     return true;
   }
 };
 
 } // anonymous namespace
 
-void trans::rewriteBlockObjCVariable(MigrationPass &pass) {
-  BlockObjCVarRewriter trans(pass);
-  trans.TraverseDecl(pass.Ctx.getTranslationUnitDecl());
+void BlockObjCVariableTraverser::traverseBody(BodyContext &BodyCtx) {
+  MigrationPass &Pass = BodyCtx.getMigrationContext().Pass;
+  llvm::DenseSet<VarDecl *> VarsToChange;
+
+  BlockObjCVarRewriter trans(Pass, VarsToChange);
+  trans.TraverseStmt(BodyCtx.getTopStmt());
+
+  for (llvm::DenseSet<VarDecl *>::iterator
+         I = VarsToChange.begin(), E = VarsToChange.end(); I != E; ++I) {
+    VarDecl *var = *I;
+    BlocksAttr *attr = var->getAttr<BlocksAttr>();
+    if(!attr)
+      continue;
+    bool useWeak = canApplyWeak(Pass.Ctx, var->getType());
+    SourceManager &SM = Pass.Ctx.getSourceManager();
+    Transaction Trans(Pass.TA);
+    Pass.TA.replaceText(SM.getExpansionLoc(attr->getLocation()),
+                        "__block",
+                        useWeak ? "__weak" : "__unsafe_unretained");
+  }
 }
diff --git a/lib/ARCMigrate/Transforms.cpp b/lib/ARCMigrate/Transforms.cpp
index 91121f0245..6d08d1e7d2 100644
--- a/lib/ARCMigrate/Transforms.cpp
+++ b/lib/ARCMigrate/Transforms.cpp
@@ -553,6 +553,7 @@ static void traverseAST(MigrationPass &pass) {
     MigrateCtx.addTraverser(new GCAttrsTraverser());
   }
   MigrateCtx.addTraverser(new PropertyRewriteTraverser());
+  MigrateCtx.addTraverser(new BlockObjCVariableTraverser());
 
   MigrateCtx.traverse(pass.Ctx.getTranslationUnitDecl());
 }
@@ -564,7 +565,6 @@ static void independentTransforms(MigrationPass &pass) {
   removeZeroOutPropsInDeallocFinalize(pass);
   makeAssignARCSafe(pass);
   rewriteUnbridgedCasts(pass);
-  rewriteBlockObjCVariable(pass);
   checkAPIUses(pass);
   traverseAST(pass);
 }
diff --git a/lib/ARCMigrate/Transforms.h b/lib/ARCMigrate/Transforms.h
index 84912bcd26..445c3e599d 100644
--- a/lib/ARCMigrate/Transforms.h
+++ b/lib/ARCMigrate/Transforms.h
@@ -37,7 +37,6 @@ void rewriteUnbridgedCasts(MigrationPass &pass);
 void makeAssignARCSafe(MigrationPass &pass);
 void removeRetainReleaseDeallocFinalize(MigrationPass &pass);
 void removeZeroOutPropsInDeallocFinalize(MigrationPass &pass);
-void rewriteBlockObjCVariable(MigrationPass &pass);
 void rewriteUnusedInitDelegate(MigrationPass &pass);
 void checkAPIUses(MigrationPass &pass);
 
@@ -130,6 +129,11 @@ public:
   virtual void traverseObjCImplementation(ObjCImplementationContext &ImplCtx);
 };
 
+class BlockObjCVariableTraverser : public ASTTraverser {
+public:
+  virtual void traverseBody(BodyContext &BodyCtx);
+};
+
 // GC transformations
 
 class GCAttrsTraverser : public ASTTraverser {