3 files changed, 42 insertions, 14 deletions

diff --git a/lib/CodeGen/CGBlocks.cpp b/lib/CodeGen/CGBlocks.cpp
index 43b098e021..b8827793d6 100644
--- a/lib/CodeGen/CGBlocks.cpp
+++ b/lib/CodeGen/CGBlocks.cpp
@@ -442,10 +442,10 @@ RValue CodeGenFunction::EmitBlockCallExpr(const CallExpr* E) {
 }
 
 llvm::Value *CodeGenFunction::GetAddrOfBlockDecl(const BlockDeclRefExpr *E) {
-  uint64_t &offset = BlockDecls[E->getDecl()];
+  const ValueDecl *VD = E->getDecl();
+  
+  uint64_t &offset = BlockDecls[VD];
 
-  const llvm::Type *Ty;
-  Ty = CGM.getTypes().ConvertType(E->getDecl()->getType());
 
   // See if we have already allocated an offset for this variable.
   if (offset == 0) {
@@ -462,20 +462,23 @@ llvm::Value *CodeGenFunction::GetAddrOfBlockDecl(const BlockDeclRefExpr *E) {
                                                          offset),
                                      "block.literal");
   if (E->isByRef()) {
-    bool needsCopyDispose = BlockRequiresCopying(E->getType());
     const llvm::Type *PtrStructTy
-      = llvm::PointerType::get(BuildByRefType(E->getDecl()), 0);
+      = llvm::PointerType::get(BuildByRefType(VD), 0);
     // The block literal will need a copy/destroy helper.
     BlockHasCopyDispose = true;
-    Ty = PtrStructTy;
+    
+    const llvm::Type *Ty = PtrStructTy;
     Ty = llvm::PointerType::get(Ty, 0);
     V = Builder.CreateBitCast(V, Ty);
     V = Builder.CreateLoad(V, false);
     V = Builder.CreateStructGEP(V, 1, "forwarding");
     V = Builder.CreateLoad(V, false);
     V = Builder.CreateBitCast(V, PtrStructTy);
-    V = Builder.CreateStructGEP(V, needsCopyDispose*2 + 4, "x");
+    V = Builder.CreateStructGEP(V, getByRefValueLLVMField(VD), 
+                                VD->getNameAsString());
   } else {
+    const llvm::Type *Ty = CGM.getTypes().ConvertType(VD->getType());
+
     Ty = llvm::PointerType::get(Ty, 0);
     V = Builder.CreateBitCast(V, Ty);
   }
diff --git a/lib/CodeGen/CGDecl.cpp b/lib/CodeGen/CGDecl.cpp
index 82c1c02bf8..eae5a0d113 100644
--- a/lib/CodeGen/CGDecl.cpp
+++ b/lib/CodeGen/CGDecl.cpp
@@ -199,6 +199,12 @@ void CodeGenFunction::EmitStaticBlockVarDecl(const VarDecl &D) {
   }
 }
 
+unsigned CodeGenFunction::getByRefValueLLVMField(const ValueDecl *VD) const {
+  assert(ByRefValueInfo.count(VD) && "Did not find value!");
+  
+  return ByRefValueInfo.find(VD)->second.second;
+}
+
 /// BuildByRefType - This routine changes a __block variable declared as T x
 ///   into:
 ///
@@ -212,8 +218,11 @@ void CodeGenFunction::EmitStaticBlockVarDecl(const VarDecl &D) {
 ///        T x;
 ///      } x
 ///
-/// Align is the alignment needed in bytes for x.
 const llvm::Type *CodeGenFunction::BuildByRefType(const ValueDecl *D) {
+  std::pair<const llvm::Type *, unsigned> &Info = ByRefValueInfo[D];
+  if (Info.first)
+    return Info.first;
+  
   QualType Ty = D->getType();
   uint64_t Align = getContext().getDeclAlignInBytes(D);
   (void) Align;
@@ -237,7 +246,10 @@ const llvm::Type *CodeGenFunction::BuildByRefType(const ValueDecl *D) {
   // FIXME: Align this on at least an Align boundary, assert if we can't.
   assert((Align <= unsigned(Target.getPointerAlign(0))/8)
          && "Can't align more than pointer yet");
-  Types[needsCopyDispose*2 + 4] = LTy;
+  
+  unsigned FieldNumber = needsCopyDispose*2 + 4;
+  
+  Types[FieldNumber] = LTy;
   
   const llvm::Type *T = llvm::StructType::get(VMContext, Types, false);
   
@@ -245,7 +257,10 @@ const llvm::Type *CodeGenFunction::BuildByRefType(const ValueDecl *D) {
   CGM.getModule().addTypeName("struct.__block_byref_" + D->getNameAsString(), 
                               ByRefTypeHolder.get());
   
-  return ByRefTypeHolder.get();
+  Info.first = ByRefTypeHolder.get();
+  Info.second = FieldNumber;
+  
+  return Info.first;
 }
 
 /// EmitLocalBlockVarDecl - Emit code and set up an entry in LocalDeclMap for a
@@ -371,10 +386,10 @@ void CodeGenFunction::EmitLocalBlockVarDecl(const VarDecl &D) {
 
   if (Init) {
     llvm::Value *Loc = DeclPtr;
-    if (isByRef) {
-      bool needsCopyDispose = BlockRequiresCopying(Ty);
-      Loc = Builder.CreateStructGEP(DeclPtr, needsCopyDispose*2+4, "x");
-    }
+    if (isByRef)
+      Loc = Builder.CreateStructGEP(DeclPtr, getByRefValueLLVMField(&D), 
+                                    D.getNameAsString());
+
     if (Ty->isReferenceType()) {
       RValue RV = EmitReferenceBindingToExpr(Init, Ty, /*IsInitializer=*/true);
       EmitStoreOfScalar(RV.getScalarVal(), Loc, false, Ty);
diff --git a/lib/CodeGen/CodeGenFunction.h b/lib/CodeGen/CodeGenFunction.h
index 5886035f29..a81f67bda3 100644
--- a/lib/CodeGen/CodeGenFunction.h
+++ b/lib/CodeGen/CodeGenFunction.h
@@ -292,6 +292,16 @@ private:
   /// we know how many temporaries were created by a certain expression.
   llvm::SmallVector<size_t, 4> ConditionalTempDestructionStack;
 
+
+  /// ByrefValueInfoMap - For each __block variable, contains a pair of the LLVM
+  /// type as well as the field number that contains the actual data.
+  llvm::DenseMap<const ValueDecl *, std::pair<const llvm::Type *, 
+                                              unsigned> > ByRefValueInfo;
+  
+  /// getByrefValueFieldNumber - Given a declaration, returns the LLVM field
+  /// number that holds the value.
+  unsigned getByRefValueLLVMField(const ValueDecl *VD) const;
+  
 public:
   CodeGenFunction(CodeGenModule &cgm);