2 files changed, 38 insertions, 7 deletions
diff --git a/lib/Bitcode/Reader/Deserialize.cpp b/lib/Bitcode/Reader/Deserialize.cpp
index 99cb5d2e3a..5d0c724b8b 100644
--- a/lib/Bitcode/Reader/Deserialize.cpp
+++ b/lib/Bitcode/Reader/Deserialize.cpp
@@ -63,15 +63,16 @@ void Deserializer::ReadRecord() {
     Code = Stream.ReadCode();
   
     if (Code == bitc::ENTER_SUBBLOCK) {
-      // No known subblocks, always skip them.
+      BlockLocs.push_back(Stream.GetCurrentBitNo());
       unsigned id = Stream.ReadSubBlockID();
       Stream.EnterSubBlock(id);
       continue;
     }
 
-    if (Code == bitc::END_BLOCK) {
+    if (Code == bitc::END_BLOCK) {      
       bool x = Stream.ReadBlockEnd();
       assert (!x && "Error at block end.");
+      BlockLocs.pop_back();
       continue;
     }
     
@@ -88,6 +89,26 @@ void Deserializer::ReadRecord() {
   assert (Record.size() > 0 || Stream.AtEndOfStream());
 }
 
+Deserializer::Location Deserializer::GetCurrentBlockLocation() {
+  if (!inRecord())
+    ReadRecord();
+  
+  assert (!BlockLocs.empty());
+  return BlockLocs.back();
+}
+
+bool Deserializer::FinishedBlock(Location BlockLoc) {
+  if (!inRecord())
+    ReadRecord();
+  
+  for (llvm::SmallVector<Location,5>::reverse_iterator
+       I=BlockLocs.rbegin(), E=BlockLocs.rend(); I!=E; ++I)
+    if (*I == BlockLoc)
+      return false;
+  
+  return true;
+}
+
 bool Deserializer::AtEnd() {
   if (inRecord())
     return false;
@@ -159,7 +180,7 @@ void Deserializer::RegisterPtr(unsigned PtrId, const void* Ptr) {
 }
 
 void Deserializer::ReadUIntPtr(uintptr_t& PtrRef, bool AllowBackpatch) {
-  unsigned PtrId = ReadInt();
+  SerializedPtrID PtrId = ReadPtrID();
   
   if (PtrId == 0) {
     PtrRef = 0;
@@ -194,7 +215,7 @@ void Deserializer::ReadUIntPtr(uintptr_t& PtrRef, bool AllowBackpatch) {
 }
 
 uintptr_t Deserializer::ReadInternalRefPtr() {
-  unsigned PtrId = ReadInt();
+  SerializedPtrID PtrId = ReadPtrID();
   
   assert (PtrId != 0 && "References cannot refer the NULL address.");
 
diff --git a/lib/Bitcode/Writer/Serialize.cpp b/lib/Bitcode/Writer/Serialize.cpp
index b97462b629..3baf9ad59a 100644
--- a/lib/Bitcode/Writer/Serialize.cpp
+++ b/lib/Bitcode/Writer/Serialize.cpp
@@ -14,6 +14,10 @@
 #include "llvm/Bitcode/Serialize.h"
 #include "string.h"
 
+#ifdef DEBUG_BACKPATCH
+#include "llvm/Support/Streams.h"
+#endif
+
 using namespace llvm;
 
 Serializer::Serializer(BitstreamWriter& stream)
@@ -67,15 +71,13 @@ void Serializer::EmitCStr(const char* s, const char* end) {
     Record.push_back(*s);
     ++s;
   }
-
-  EmitRecord();
 }
 
 void Serializer::EmitCStr(const char* s) {
   EmitCStr(s,s+strlen(s));
 }
 
-unsigned Serializer::getPtrId(const void* ptr) {
+SerializedPtrID Serializer::getPtrId(const void* ptr) {
   if (!ptr)
     return 0;
   
@@ -83,12 +85,20 @@ unsigned Serializer::getPtrId(const void* ptr) {
   
   if (I == PtrMap.end()) {
     unsigned id = PtrMap.size()+1;
+#ifdef DEBUG_BACKPATCH
+    llvm::cerr << "Registered PTR: " << ptr << " => " << id << "\n";
+#endif
     PtrMap[ptr] = id;
     return id;
   }
   else return I->second;
 }
 
+bool Serializer::isRegistered(const void* ptr) const {
+  MapTy::const_iterator I = PtrMap.find(ptr);
+  return I != PtrMap.end();
+}
+
 
 #define INT_EMIT(TYPE)\
 void SerializeTrait<TYPE>::Emit(Serializer&S, TYPE X) { S.EmitInt(X); }