2 files changed, 8 insertions, 10 deletions

diff --git a/clang/lib/CodeGen/CGBlocks.cpp b/clang/lib/CodeGen/CGBlocks.cpp
index 344ac60d98a..4e4698273fc 100644
--- a/clang/lib/CodeGen/CGBlocks.cpp
+++ b/clang/lib/CodeGen/CGBlocks.cpp
@@ -254,8 +254,8 @@ llvm::Value *CodeGenFunction::BuildBlockLiteralTmp(const BlockExpr *BE) {
           if (BDRE->isByRef()) {
             // FIXME: For only local, or all byrefs?
             NoteForHelper[helpersize].flag = BLOCK_FIELD_IS_BYREF |
-              (0?BLOCK_FIELD_IS_WEAK : 0);
-            // FIXME: Add weak support
+              // FIXME: Someone double check this.
+              (VD->getType().isObjCGCWeak() ? BLOCK_FIELD_IS_WEAK : 0);
             const llvm::Type *Ty = Types[i+5];
             llvm::Value *Loc = LocalDeclMap[VD];
             Loc = Builder.CreateStructGEP(Loc, 1, "forwarding");
@@ -484,7 +484,7 @@ llvm::Value *CodeGenFunction::GetAddrOfBlockDecl(const BlockDeclRefExpr *E) {
     ErrorUnsupported(E, "__block variable in block literal");
   else if (!Enable__block && E->getType()->isBlockPointerType())
     ErrorUnsupported(E, "block pointer in block literal");
-  else if (E->getDecl()->getAttr<ObjCNSObjectAttr>() ||
+  else if (!Enable__block && E->getDecl()->getAttr<ObjCNSObjectAttr>() ||
            getContext().isObjCNSObjectType(E->getType()))
     ErrorUnsupported(E, "__attribute__((NSObject)) variable in block "
                      "literal");
@@ -946,11 +946,6 @@ BlockFunction::GeneratebyrefDestroyHelperFunction(const llvm::Type *T,
   V = Builder.CreateStructGEP(V, 6, "x");
   V = Builder.CreateBitCast(V, PtrToInt8Ty);
 
-  // FIXME: Move to other one.
-  // int flag = BLOCK_FIELD_IS_BYREF;
-  // FIXME: Add weak support
-  if (0)
-    flag |= BLOCK_FIELD_IS_WEAK;
   flag |= BLOCK_BYREF_CALLER;
   BuildBlockRelease(V, flag);
   CGF.FinishFunction();
diff --git a/clang/lib/CodeGen/CGDecl.cpp b/clang/lib/CodeGen/CGDecl.cpp
index 1322ca223ac..abb13376c7b 100644
--- a/clang/lib/CodeGen/CGDecl.cpp
+++ b/clang/lib/CodeGen/CGDecl.cpp
@@ -340,6 +340,7 @@ void CodeGenFunction::EmitLocalBlockVarDecl(const VarDecl &D) {
 
     // The block literal will need a copy/destroy helper.
     BlockHasCopyDispose = true;
+    needsDispose = true;
 
     if (Ty->isBlockPointerType()) {
       flag |= BLOCK_FIELD_IS_BLOCK;
@@ -348,7 +349,10 @@ void CodeGenFunction::EmitLocalBlockVarDecl(const VarDecl &D) {
       flag |= BLOCK_FIELD_IS_OBJECT;
       flags |= BLOCK_HAS_COPY_DISPOSE;
     }
-    // FIXME: Need to set BLOCK_FIELD_IS_WEAK as appropriate.
+
+    // FIXME: Someone double check this.
+    if (Ty.isObjCGCWeak())
+      flag |= BLOCK_FIELD_IS_WEAK;
 
     int isa = 0;
     if (flag&BLOCK_FIELD_IS_WEAK)
@@ -380,7 +384,6 @@ void CodeGenFunction::EmitLocalBlockVarDecl(const VarDecl &D) {
       Builder.CreateStore(BuildbyrefDestroyHelper(DeclPtr->getType(), flag),
                           destroy_helper);
     }
-    needsDispose = true;
   }
 
   // Handle the cleanup attribute