Keep track of memory allocated by alloca so that it is freed appropriately

llvm-svn: 1776

2 files changed, 37 insertions, 5 deletions

diff --git a/llvm/lib/ExecutionEngine/Interpreter/Execution.cpp b/llvm/lib/ExecutionEngine/Interpreter/Execution.cpp
index 9f6317fd393..8b910f68af9 100644
--- a/llvm/lib/ExecutionEngine/Interpreter/Execution.cpp
+++ b/llvm/lib/ExecutionEngine/Interpreter/Execution.cpp
@@ -741,15 +741,16 @@ void Interpreter::executeAllocInst(AllocationInst *I, ExecutionContext &SF) {
   }
 
   // Allocate enough memory to hold the type...
-  GenericValue Result;
   // FIXME: Don't use CALLOC, use a tainted malloc.
-  Result.PointerVal = (PointerTy)calloc(NumElements, TD.getTypeSize(Ty));
+  void *Memory = calloc(NumElements, TD.getTypeSize(Ty));
+
+  GenericValue Result;
+  Result.PointerVal = (PointerTy)Memory;
   assert(Result.PointerVal != 0 && "Null pointer returned by malloc!");
   SetValue(I, Result, SF);
 
-  if (I->getOpcode() == Instruction::Alloca) {
-    // TODO: FIXME: alloca should keep track of memory to free it later...
-  }
+  if (I->getOpcode() == Instruction::Alloca)
+    ECStack.back().Allocas.add(Memory);
 }
 
 static void executeFreeInst(FreeInst *I, ExecutionContext &SF) {
diff --git a/llvm/lib/ExecutionEngine/Interpreter/Interpreter.h b/llvm/lib/ExecutionEngine/Interpreter/Interpreter.h
index 050b7d332d2..8576f9ce225 100644
--- a/llvm/lib/ExecutionEngine/Interpreter/Interpreter.h
+++ b/llvm/lib/ExecutionEngine/Interpreter/Interpreter.h
@@ -42,6 +42,36 @@ union GenericValue {
   PointerTy       PointerVal;
 };
 
+// AllocaHolder - Object to track all of the blocks of memory allocated by
+// alloca.  When the function returns, this object is poped off the execution
+// stack, which causes the dtor to be run, which frees all the alloca'd memory.
+//
+class AllocaHolder {
+  friend class AllocaHolderHandle;
+  std::vector<void*> Allocations;
+  unsigned RefCnt;
+public:
+  AllocaHolder() : RefCnt(0) {}
+  void add(void *mem) { Allocations.push_back(mem); }
+  ~AllocaHolder() {
+    for (unsigned i = 0; i < Allocations.size(); ++i)
+      free(Allocations[i]);
+  }
+};
+
+// AllocaHolderHandle gives AllocaHolder value semantics so we can stick it into
+// a vector...
+//
+class AllocaHolderHandle {
+  AllocaHolder *H;
+public:
+  AllocaHolderHandle() : H(new AllocaHolder()) { H->RefCnt++; }
+  AllocaHolderHandle(const AllocaHolderHandle &AH) : H(AH.H) { H->RefCnt++; }
+  ~AllocaHolderHandle() { if (--H->RefCnt == 0) delete H; }
+
+  void add(void *mem) { H->add(mem); }
+};
+
 typedef std::vector<GenericValue> ValuePlaneTy;
 
 // ExecutionContext struct - This struct represents one stack frame currently
@@ -57,6 +87,7 @@ struct ExecutionContext {
   BasicBlock           *PrevBB;     // The previous BB or null if in first BB
   CallInst             *Caller;     // Holds the call that called subframes.
                                     // NULL if main func or debugger invoked fn
+  AllocaHolderHandle    Allocas;    // Track memory allocated by alloca
 };
 
 // Interpreter - This class represents the entirety of the interpreter.