5 files changed, 91 insertions, 11 deletions

diff --git a/lib/Core/ExecutionState.cpp b/lib/Core/ExecutionState.cpp
index c6a2cad4..db685639 100644
--- a/lib/Core/ExecutionState.cpp
+++ b/lib/Core/ExecutionState.cpp
@@ -86,9 +86,46 @@ ExecutionState::ExecutionState(const std::vector<ref<Expr> > &assumptions)
 }
 
 ExecutionState::~ExecutionState() {
+  for (unsigned int i=0; i<symbolics.size(); i++)
+  {
+    const MemoryObject *mo = symbolics[i].first;
+    assert(mo->refCount > 0);
+    mo->refCount--;
+    if (mo->refCount == 0)
+      delete mo;
+  }
+
   while (!stack.empty()) popFrame();
 }
 
+ExecutionState::ExecutionState(const ExecutionState& state)
+  : fnAliases(state.fnAliases),
+    fakeState(state.fakeState),
+    underConstrained(state.underConstrained),
+    depth(state.depth),
+    pc(state.pc),
+    prevPC(state.prevPC),
+    stack(state.stack),
+    constraints(state.constraints),
+    queryCost(state.queryCost),
+    weight(state.weight),
+    addressSpace(state.addressSpace),
+    pathOS(state.pathOS),
+    symPathOS(state.symPathOS),
+    instsSinceCovNew(state.instsSinceCovNew),
+    coveredNew(state.coveredNew),
+    forkDisabled(state.forkDisabled),
+    coveredLines(state.coveredLines),
+    ptreeNode(state.ptreeNode),
+    symbolics(state.symbolics),
+    arrayNames(state.arrayNames),
+    shadowObjects(state.shadowObjects),
+    incomingBBIndex(state.incomingBBIndex)
+{
+  for (unsigned int i=0; i<symbolics.size(); i++)
+    symbolics[i].first->refCount++;
+}
+
 ExecutionState *ExecutionState::branch() {
   depth++;
 
@@ -114,6 +151,10 @@ void ExecutionState::popFrame() {
   stack.pop_back();
 }
 
+void ExecutionState::addSymbolic(const MemoryObject *mo, const Array *array) { 
+  mo->refCount++;
+  symbolics.push_back(std::make_pair(mo, array));
+}
 ///
 
 std::string ExecutionState::getFnAlias(std::string fn) {
diff --git a/lib/Core/Memory.cpp b/lib/Core/Memory.cpp
index d54264a3..7b3655f8 100644
--- a/lib/Core/Memory.cpp
+++ b/lib/Core/Memory.cpp
@@ -17,6 +17,7 @@
 #include "klee/util/BitArray.h"
 
 #include "ObjectHolder.h"
+#include "MemoryManager.h"
 
 #include <llvm/Function.h>
 #include <llvm/Instruction.h>
@@ -63,6 +64,8 @@ ObjectHolder &ObjectHolder::operator=(const ObjectHolder &b) {
 int MemoryObject::counter = 0;
 
 MemoryObject::~MemoryObject() {
+  if (parent)
+    parent->markFreed(this);
 }
 
 void MemoryObject::getAllocInfo(std::string &result) const {
@@ -100,6 +103,7 @@ ObjectState::ObjectState(const MemoryObject *mo)
     updates(0, 0),
     size(mo->size),
     readOnly(false) {
+  mo->refCount++;
   if (!UseConstantArrays) {
     // FIXME: Leaked.
     static unsigned id = 0;
@@ -120,6 +124,7 @@ ObjectState::ObjectState(const MemoryObject *mo, const Array *array)
     updates(array, 0),
     size(mo->size),
     readOnly(false) {
+  mo->refCount++;
   makeSymbolic();
 }
 
@@ -135,6 +140,8 @@ ObjectState::ObjectState(const ObjectState &os)
     size(os.size),
     readOnly(false) {
   assert(!os.readOnly && "no need to copy read only object?");
+  if (object)
+    object->refCount++;
 
   if (os.knownSymbolics) {
     knownSymbolics = new ref<Expr>[size];
@@ -150,6 +157,16 @@ ObjectState::~ObjectState() {
   if (flushMask) delete flushMask;
   if (knownSymbolics) delete[] knownSymbolics;
   delete[] concreteStore;
+
+  if (object)
+  {
+    assert(object->refCount > 0);
+    object->refCount--;
+    if (object->refCount == 0)
+    {
+      delete object;
+    }
+  }
 }
 
 /***/
diff --git a/lib/Core/Memory.h b/lib/Core/Memory.h
index 9ce8e09b..637f8772 100644
--- a/lib/Core/Memory.h
+++ b/lib/Core/Memory.h
@@ -30,9 +30,12 @@ class Solver;
 
 class MemoryObject {
   friend class STPBuilder;
+  friend class ObjectState;
+  friend class ExecutionState;
 
 private:
   static int counter;
+  mutable unsigned refCount;
 
 public:
   unsigned id;
@@ -50,6 +53,8 @@ public:
   bool fake_object;
   bool isUserSpecified;
 
+  MemoryManager *parent;
+
   /// "Location" for which this memory object was allocated. This
   /// should be either the allocating instruction or the global object
   /// it was allocated for (or whatever else makes sense).
@@ -69,17 +74,21 @@ public:
   // XXX this is just a temp hack, should be removed
   explicit
   MemoryObject(uint64_t _address) 
-    : id(counter++),
+    : refCount(0),
+      id(counter++), 
       address(_address),
       size(0),
       isFixed(true),
+      parent(NULL),
       allocSite(0) {
   }
 
   MemoryObject(uint64_t _address, unsigned _size, 
                bool _isLocal, bool _isGlobal, bool _isFixed,
-               const llvm::Value *_allocSite) 
-    : id(counter++),
+               const llvm::Value *_allocSite,
+               MemoryManager *_parent)
+    : refCount(0), 
+      id(counter++),
       address(_address),
       size(_size),
       name("unnamed"),
@@ -88,6 +97,7 @@ public:
       isFixed(_isFixed),
       fake_object(false),
       isUserSpecified(false),
+      parent(_parent), 
       allocSite(_allocSite) {
   }
 
diff --git a/lib/Core/MemoryManager.cpp b/lib/Core/MemoryManager.cpp
index 79fbcecf..06c234a2 100644
--- a/lib/Core/MemoryManager.cpp
+++ b/lib/Core/MemoryManager.cpp
@@ -25,8 +25,10 @@ using namespace klee;
 
 MemoryManager::~MemoryManager() { 
   while (!objects.empty()) {
-    MemoryObject *mo = objects.back();
-    objects.pop_back();
+    MemoryObject *mo = *objects.begin();
+    if (!mo->isFixed)
+      free((void *)mo->address);
+    objects.erase(mo);
     delete mo;
   }
 }
@@ -44,8 +46,8 @@ MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
   
   ++stats::allocations;
   MemoryObject *res = new MemoryObject(address, size, isLocal, isGlobal, false,
-                                       allocSite);
-  objects.push_back(res);
+                                       allocSite, this);
+  objects.insert(res);
   return res;
 }
 
@@ -62,11 +64,20 @@ MemoryObject *MemoryManager::allocateFixed(uint64_t address, uint64_t size,
 
   ++stats::allocations;
   MemoryObject *res = new MemoryObject(address, size, false, true, true,
-                                       allocSite);
-  objects.push_back(res);
+                                       allocSite, this);
+  objects.insert(res);
   return res;
 }
 
 void MemoryManager::deallocate(const MemoryObject *mo) {
   assert(0);
 }
+
+void MemoryManager::markFreed(MemoryObject *mo) {
+  if (objects.find(mo) != objects.end())
+  {
+    if (!mo->isFixed)
+      free((void *)mo->address);
+    objects.erase(mo);
+  }
+}
diff --git a/lib/Core/MemoryManager.h b/lib/Core/MemoryManager.h
index adb2ba22..f398db62 100644
--- a/lib/Core/MemoryManager.h
+++ b/lib/Core/MemoryManager.h
@@ -10,7 +10,7 @@
 #ifndef KLEE_MEMORYMANAGER_H
 #define KLEE_MEMORYMANAGER_H
 
-#include <vector>
+#include <set>
 #include <stdint.h>
 
 namespace llvm {
@@ -22,7 +22,7 @@ namespace klee {
 
   class MemoryManager {
   private:
-    typedef std::vector<MemoryObject*> objects_ty;
+    typedef std::set<MemoryObject*> objects_ty;
     objects_ty objects;
 
   public:
@@ -34,6 +34,7 @@ namespace klee {
     MemoryObject *allocateFixed(uint64_t address, uint64_t size,
                                 const llvm::Value *allocSite);
     void deallocate(const MemoryObject *mo);
+    void markFreed(MemoryObject *mo);
   };
 
 } // End klee namespace