3 files changed, 16 insertions, 9 deletions

diff --git a/lib/Bitcode/Writer/BitcodeWriter.cpp b/lib/Bitcode/Writer/BitcodeWriter.cpp
index 860b7e9..98d567e 100644
--- a/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -577,10 +577,9 @@ static void WriteFunctionLocalMetadata(const Function &F,
                                        BitstreamWriter &Stream) {
   bool StartedMetadataBlock = false;
   SmallVector<uint64_t, 64> Record;
-  const ValueEnumerator::ValueList &Vals = VE.getMDValues();
-  
+  const SmallVector<const MDNode *, 8> &Vals = VE.getFunctionLocalMDValues();
   for (unsigned i = 0, e = Vals.size(); i != e; ++i)
-    if (const MDNode *N = dyn_cast<MDNode>(Vals[i].first))
+    if (const MDNode *N = Vals[i])
       if (N->isFunctionLocal() && N->getFunction() == &F) {
         if (!StartedMetadataBlock) {
           Stream.EnterSubblock(bitc::METADATA_BLOCK_ID, 3);
@@ -588,7 +587,7 @@ static void WriteFunctionLocalMetadata(const Function &F,
         }
         WriteMDNode(N, VE, Stream, Record);
       }
-
+      
   if (StartedMetadataBlock)
     Stream.ExitBlock();
 }
diff --git a/lib/Bitcode/Writer/ValueEnumerator.cpp b/lib/Bitcode/Writer/ValueEnumerator.cpp
index d2baec7..a303fb4 100644
--- a/lib/Bitcode/Writer/ValueEnumerator.cpp
+++ b/lib/Bitcode/Writer/ValueEnumerator.cpp
@@ -257,6 +257,8 @@ void ValueEnumerator::EnumerateMetadata(const Value *MD) {
       else
         EnumerateType(Type::getVoidTy(MD->getContext()));
     }
+    if (N->isFunctionLocal() && N->getFunction())
+      FunctionLocalMDs.push_back(N);
     return;
   }
   
@@ -414,7 +416,8 @@ void ValueEnumerator::incorporateFunction(const Function &F) {
 
   FirstInstID = Values.size();
 
-  SmallVector<MDNode *, 8> FunctionLocalMDs;
+  FunctionLocalMDs.clear();
+  SmallVector<MDNode *, 8> FnLocalMDVector;
   // Add all of the instructions.
   for (Function::const_iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
     for (BasicBlock::const_iterator I = BB->begin(), E = BB->end(); I!=E; ++I) {
@@ -423,7 +426,7 @@ void ValueEnumerator::incorporateFunction(const Function &F) {
         if (MDNode *MD = dyn_cast<MDNode>(*OI))
           if (MD->isFunctionLocal() && MD->getFunction())
             // Enumerate metadata after the instructions they might refer to.
-            FunctionLocalMDs.push_back(MD);
+            FnLocalMDVector.push_back(MD);
       }
       if (!I->getType()->isVoidTy())
         EnumerateValue(I);
@@ -431,8 +434,8 @@ void ValueEnumerator::incorporateFunction(const Function &F) {
   }
 
   // Add all of the function-local metadata.
-  for (unsigned i = 0, e = FunctionLocalMDs.size(); i != e; ++i)
-    EnumerateOperandType(FunctionLocalMDs[i]);
+  for (unsigned i = 0, e = FnLocalMDVector.size(); i != e; ++i)
+    EnumerateOperandType(FnLocalMDVector[i]);
 }
 
 void ValueEnumerator::purgeFunction() {
diff --git a/lib/Bitcode/Writer/ValueEnumerator.h b/lib/Bitcode/Writer/ValueEnumerator.h
index 4f8ebf5..2b9b15f 100644
--- a/lib/Bitcode/Writer/ValueEnumerator.h
+++ b/lib/Bitcode/Writer/ValueEnumerator.h
@@ -15,6 +15,7 @@
 #define VALUE_ENUMERATOR_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Attributes.h"
 #include <vector>
 
@@ -26,7 +27,7 @@ class Instruction;
 class BasicBlock;
 class Function;
 class Module;
-class MetadataBase;
+class MDNode;
 class NamedMDNode;
 class AttrListPtr;
 class TypeSymbolTable;
@@ -49,6 +50,7 @@ private:
   ValueMapType ValueMap;
   ValueList Values;
   ValueList MDValues;
+  SmallVector<const MDNode *, 8> FunctionLocalMDs;
   ValueMapType MDValueMap;
   
   typedef DenseMap<void*, unsigned> AttributeMapType;
@@ -105,6 +107,9 @@ public:
   
   const ValueList &getValues() const { return Values; }
   const ValueList &getMDValues() const { return MDValues; }
+  const SmallVector<const MDNode *, 8> &getFunctionLocalMDValues() const { 
+    return FunctionLocalMDs;
+  }
   const TypeList &getTypes() const { return Types; }
   const std::vector<const BasicBlock*> &getBasicBlocks() const {
     return BasicBlocks;