1 files changed, 27 insertions, 7 deletions
diff --git a/lib/Core/Memory.h b/lib/Core/Memory.h
index e4260c55..2ab7cd22 100644
--- a/lib/Core/Memory.h
+++ b/lib/Core/Memory.h
@@ -35,10 +35,13 @@ class MemoryObject {
   friend class STPBuilder;
   friend class ObjectState;
   friend class ExecutionState;
+  friend class ref<MemoryObject>;
+  friend class ref<const MemoryObject>;
 
 private:
   static int counter;
-  mutable unsigned refCount;
+  /// @brief Required by klee::ref-managed objects
+  mutable class ReferenceCounter _refCount;
 
 public:
   unsigned id;
@@ -75,8 +78,7 @@ public:
   // XXX this is just a temp hack, should be removed
   explicit
   MemoryObject(uint64_t _address) 
-    : refCount(0),
-      id(counter++), 
+    : id(counter++),
       address(_address),
       size(0),
       isFixed(true),
@@ -88,8 +90,7 @@ public:
                bool _isLocal, bool _isGlobal, bool _isFixed,
                const llvm::Value *_allocSite,
                MemoryManager *_parent)
-    : refCount(0), 
-      id(counter++),
+    : id(counter++),
       address(_address),
       size(_size),
       name("unnamed"),
@@ -143,6 +144,25 @@ public:
       return ConstantExpr::alloc(0, Expr::Bool);
     }
   }
+
+  /// Compare this object with memory object b.
+  /// \param b memory object to compare with
+  /// \return <0 if this is smaller, 0 if both are equal, >0 if b is smaller
+  int compare(const MemoryObject &b) const {
+    // Short-cut with id
+    if (id == b.id)
+      return 0;
+    if (address != b.address)
+      return (address < b.address ? -1 : 1);
+
+    if (size != b.size)
+      return (size < b.size ? -1 : 1);
+
+    if (allocSite != b.allocSite)
+      return (allocSite < b.allocSite ? -1 : 1);
+
+    return 0;
+  }
 };
 
 class ObjectState {
@@ -155,7 +175,7 @@ private:
   /// @brief Required by klee::ref-managed objects
   class ReferenceCounter _refCount;
 
-  const MemoryObject *object;
+  ref<const MemoryObject> object;
 
   uint8_t *concreteStore;
 
@@ -188,7 +208,7 @@ public:
   ObjectState(const ObjectState &os);
   ~ObjectState();
 
-  const MemoryObject *getObject() const { return object; }
+  const MemoryObject *getObject() const { return object.get(); }
 
   void setReadOnly(bool ro) { readOnly = ro; }