From 7c2008b664cb20da231b3ddd78c47208a630ec96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Bylica?= <pawel.bylica@imapp.pl>
Date: Mon, 23 Feb 2015 14:22:31 +0100
Subject: [PATCH] Free memory allocated for dynamic stack

---
 libevmjit/Compiler.cpp       | 23 ++++++-------
 libevmjit/Compiler.h         |  4 +--
 libevmjit/RuntimeManager.cpp | 10 ++++++
 libevmjit/RuntimeManager.h   |  7 ++++
 libevmjit/Stack.cpp          | 62 ++++++++++++++++++++++++++++++++++--
 libevmjit/Stack.h            |  4 +++
 6 files changed, 94 insertions(+), 16 deletions(-)

diff --git a/libevmjit/Compiler.cpp b/libevmjit/Compiler.cpp
index a0b634bdb..22305e366 100644
--- a/libevmjit/Compiler.cpp
+++ b/libevmjit/Compiler.cpp
@@ -90,7 +90,7 @@ void Compiler::createBasicBlocks(code_iterator _codeBegin, code_iterator _codeEn
 	}
 }
 
-llvm::BasicBlock* Compiler::getJumpTableBlock()
+llvm::BasicBlock* Compiler::getJumpTableBlock(RuntimeManager& _runtimeManager)
 {
 	if (!m_jumpTableBlock)
 	{
@@ -98,7 +98,7 @@ llvm::BasicBlock* Compiler::getJumpTableBlock()
 		InsertPointGuard g{m_builder};
 		m_builder.SetInsertPoint(m_jumpTableBlock->llvm());
 		auto dest = m_builder.CreatePHI(Type::Word, 8, "target");
-		auto switchInstr = m_builder.CreateSwitch(dest, getBadJumpBlock());
+		auto switchInstr = m_builder.CreateSwitch(dest, getBadJumpBlock(_runtimeManager));
 		for (auto&& p : m_basicBlocks)
 		{
 			if (p.second.isJumpDest())
@@ -108,14 +108,14 @@ llvm::BasicBlock* Compiler::getJumpTableBlock()
 	return m_jumpTableBlock->llvm();
 }
 
-llvm::BasicBlock* Compiler::getBadJumpBlock()
+llvm::BasicBlock* Compiler::getBadJumpBlock(RuntimeManager& _runtimeManager)
 {
 	if (!m_badJumpBlock)
 	{
 		m_badJumpBlock.reset(new BasicBlock("BadJump", m_mainFunc, m_builder, true));
 		InsertPointGuard g{m_builder};
 		m_builder.SetInsertPoint(m_badJumpBlock->llvm());
-		m_builder.CreateRet(Constant::get(ReturnCode::BadJumpDestination));
+		_runtimeManager.exit(ReturnCode::BadJumpDestination);
 	}
 	return m_badJumpBlock->llvm();
 }
@@ -155,6 +155,7 @@ std::unique_ptr<llvm::Module> Compiler::compile(code_iterator _begin, code_itera
 	Memory memory(runtimeManager, gasMeter);
 	Ext ext(runtimeManager, memory);
 	Stack stack(m_builder, runtimeManager);
+	runtimeManager.setStack(stack); // Runtime Manager will free stack memory
 	Arith256 arith(m_builder);
 
 	// TODO: Create Stop basic block on demand
@@ -177,10 +178,10 @@ std::unique_ptr<llvm::Module> Compiler::compile(code_iterator _begin, code_itera
 	// Code for special blocks:
 	// TODO: move to separate function.
 	m_builder.SetInsertPoint(m_stopBB);
-	m_builder.CreateRet(Constant::get(ReturnCode::Stop));
+	runtimeManager.exit(ReturnCode::Stop);
 
 	m_builder.SetInsertPoint(abortBB);
-	m_builder.CreateRet(Constant::get(ReturnCode::OutOfGas));
+	runtimeManager.exit(ReturnCode::OutOfGas);
 
 	removeDeadBlocks();
 
@@ -589,7 +590,7 @@ void Compiler::compileBasicBlock(BasicBlock& _basicBlock, RuntimeManager& _runti
 				auto&& c = constant->getValue();
 				auto targetIdx = c.getActiveBits() <= 64 ? c.getZExtValue() : -1;
 				auto it = m_basicBlocks.find(targetIdx);
-				targetBlock = (it != m_basicBlocks.end() && it->second.isJumpDest()) ? it->second.llvm() : getBadJumpBlock();
+				targetBlock = (it != m_basicBlocks.end() && it->second.isJumpDest()) ? it->second.llvm() : getBadJumpBlock(_runtimeManager);
 			}
 
 			// TODO: Improve; check for constants
@@ -602,7 +603,7 @@ void Compiler::compileBasicBlock(BasicBlock& _basicBlock, RuntimeManager& _runti
 				else
 				{
 					_basicBlock.setJumpTarget(target);
-					m_builder.CreateBr(getJumpTableBlock());
+					m_builder.CreateBr(getJumpTableBlock(_runtimeManager));
 				}
 			}
 			else // JUMPI
@@ -618,7 +619,7 @@ void Compiler::compileBasicBlock(BasicBlock& _basicBlock, RuntimeManager& _runti
 				else
 				{
 					_basicBlock.setJumpTarget(target);
-					m_builder.CreateCondBr(cond, getJumpTableBlock(), _nextBasicBlock);
+					m_builder.CreateCondBr(cond, getJumpTableBlock(_runtimeManager), _nextBasicBlock);
 				}
 			}
 			break;
@@ -795,14 +796,14 @@ void Compiler::compileBasicBlock(BasicBlock& _basicBlock, RuntimeManager& _runti
 			_memory.require(index, size);
 			_runtimeManager.registerReturnData(index, size);
 
-			m_builder.CreateRet(Constant::get(ReturnCode::Return));
+			_runtimeManager.exit(ReturnCode::Return);
 			break;
 		}
 
 		case Instruction::SUICIDE:
 		{
 			_runtimeManager.registerSuicide(stack.pop());
-			m_builder.CreateRet(Constant::get(ReturnCode::Suicide));
+			_runtimeManager.exit(ReturnCode::Suicide);
 			break;
 		}
 
diff --git a/libevmjit/Compiler.h b/libevmjit/Compiler.h
index c9795fb99..4469389bb 100644
--- a/libevmjit/Compiler.h
+++ b/libevmjit/Compiler.h
@@ -38,9 +38,9 @@ private:
 
 	void compileBasicBlock(BasicBlock& _basicBlock, class RuntimeManager& _runtimeManager, class Arith256& _arith, class Memory& _memory, class Ext& _ext, class GasMeter& _gasMeter, llvm::BasicBlock* _nextBasicBlock);
 
-	llvm::BasicBlock* getJumpTableBlock();
+	llvm::BasicBlock* getJumpTableBlock(RuntimeManager& _runtimeManager);
 
-	llvm::BasicBlock* getBadJumpBlock();
+	llvm::BasicBlock* getBadJumpBlock(RuntimeManager& _runtimeManager);
 
 	void removeDeadBlocks();
 
diff --git a/libevmjit/RuntimeManager.cpp b/libevmjit/RuntimeManager.cpp
index dd053fd90..62c24326e 100644
--- a/libevmjit/RuntimeManager.cpp
+++ b/libevmjit/RuntimeManager.cpp
@@ -4,6 +4,8 @@
 #include <llvm/IR/IntrinsicInst.h>
 #include "preprocessor/llvm_includes_end.h"
 
+#include "Stack.h"
+
 namespace dev
 {
 namespace eth
@@ -165,6 +167,14 @@ void RuntimeManager::registerSuicide(llvm::Value* _balanceAddress)
 	set(RuntimeData::SuicideDestAddress, _balanceAddress);
 }
 
+void RuntimeManager::exit(ReturnCode _returnCode)
+{
+	if (m_stack)
+		m_stack->free();
+
+	m_builder.CreateRet(Constant::get(_returnCode));
+}
+
 void RuntimeManager::abort(llvm::Value* _jmpBuf)
 {
 	auto longjmp = llvm::Intrinsic::getDeclaration(getModule(), llvm::Intrinsic::eh_sjlj_longjmp);
diff --git a/libevmjit/RuntimeManager.h b/libevmjit/RuntimeManager.h
index 30c69ec88..2f1fb390a 100644
--- a/libevmjit/RuntimeManager.h
+++ b/libevmjit/RuntimeManager.h
@@ -11,6 +11,7 @@ namespace eth
 {
 namespace jit
 {
+class Stack;
 
 class RuntimeManager: public CompilerHelper
 {
@@ -35,9 +36,13 @@ public:
 	void registerReturnData(llvm::Value* _index, llvm::Value* _size);
 	void registerSuicide(llvm::Value* _balanceAddress);
 
+	void exit(ReturnCode _returnCode);
+
 	void abort(llvm::Value* _jmpBuf);
 	void abort() { abort(getJmpBufExt()); }
 
+	void setStack(Stack& _stack) { m_stack = &_stack; }
+
 	static llvm::StructType* getRuntimeType();
 	static llvm::StructType* getRuntimeDataType();
 
@@ -53,6 +58,8 @@ private:
 
 	code_iterator m_codeBegin = {};
 	code_iterator m_codeEnd = {};
+
+	Stack* m_stack = nullptr;
 };
 
 }
diff --git a/libevmjit/Stack.cpp b/libevmjit/Stack.cpp
index fc3b301bb..e2df49c70 100644
--- a/libevmjit/Stack.cpp
+++ b/libevmjit/Stack.cpp
@@ -8,6 +8,7 @@
 #include "Runtime.h"
 
 #include <iostream>
+#include <set>
 
 namespace dev
 {
@@ -35,7 +36,7 @@ llvm::Function* Array::createArrayPushFunc()
 	func->setDoesNotCapture(1);
 
 	llvm::Type* reallocArgTypes[] = {Type::BytePtr, Type::Size};
-	auto reallocFunc = llvm::Function::Create(llvm::FunctionType::get(Type::BytePtr, reallocArgTypes, false), llvm::Function::ExternalLinkage, "realloc", getModule());
+	auto reallocFunc = llvm::Function::Create(llvm::FunctionType::get(Type::BytePtr, reallocArgTypes, false), llvm::Function::ExternalLinkage, "ext_realloc", getModule());
 	reallocFunc->setDoesNotThrow();
 	reallocFunc->setDoesNotAlias(0);
 	reallocFunc->setDoesNotCapture(1);
@@ -130,11 +131,35 @@ llvm::Function* Array::createArrayGetFunc()
 	return func;
 }
 
+llvm::Function* Array::createFreeFunc()
+{
+	auto func = llvm::Function::Create(llvm::FunctionType::get(Type::Void, m_array->getType(), false), llvm::Function::PrivateLinkage, "array.free", getModule());
+	func->setDoesNotThrow();
+	func->setDoesNotCapture(1);
+
+	auto freeFunc = llvm::Function::Create(llvm::FunctionType::get(Type::Void, Type::BytePtr, false), llvm::Function::ExternalLinkage, "ext_free", getModule());
+	freeFunc->setDoesNotThrow();
+	freeFunc->setDoesNotCapture(1);
+
+	auto arrayPtr = &func->getArgumentList().front();
+	arrayPtr->setName("arrayPtr");
+
+	InsertPointGuard guard{m_builder};
+	m_builder.SetInsertPoint(llvm::BasicBlock::Create(m_builder.getContext(), {}, func));
+	auto dataPtr = m_builder.CreateStructGEP(arrayPtr, 0, "dataPtr");
+	auto data = m_builder.CreateLoad(dataPtr, "data");
+	auto mem = m_builder.CreateBitCast(data, Type::BytePtr, "mem");
+	m_builder.CreateCall(freeFunc, mem);
+	m_builder.CreateRetVoid();
+	return func;
+}
+
 Array::Array(llvm::IRBuilder<>& _builder, char const* _name) :
 	CompilerHelper(_builder),
 	m_pushFunc([this](){ return createArrayPushFunc(); }),
 	m_setFunc([this](){ return createArraySetFunc(); }),
-	m_getFunc([this](){ return createArrayGetFunc(); })
+	m_getFunc([this](){ return createArrayGetFunc(); }),
+	m_freeFunc([this](){ return createFreeFunc(); })
 {
 	llvm::Type* elementTys[] = {Type::WordPtr, Type::Size, Type::Size};
 	static auto arrayTy = llvm::StructType::create(elementTys, "Array");
@@ -319,6 +344,22 @@ void Stack::push(llvm::Value* _value)
 }
 }
 
+namespace
+{
+	struct AllocatedMemoryWatchdog
+	{
+		std::set<void*> allocatedMemory;
+
+		~AllocatedMemoryWatchdog()
+		{
+			if (!allocatedMemory.empty())
+				std::cerr << allocatedMemory.size() << " MEM LEAKS!" << std::endl;
+		}
+	};
+
+	AllocatedMemoryWatchdog watchdog;
+}
+
 extern "C"
 {
 	using namespace dev::eth::jit;
@@ -349,9 +390,24 @@ extern "C"
 	{
 		//std::cerr << "REALLOC: " << _data << " [" << _size << "]" << std::endl;
 		auto newData = std::realloc(_data, _size);
-		//std::cerr << "REALLOC: " << _data << " -> " << newData << " [" << _size << "]" << std::endl;
+		if (_data != newData)
+		{
+			std::cerr << "REALLOC: " << _data << " -> " << newData << " [" << _size << "]" << std::endl;
+			watchdog.allocatedMemory.erase(_data);
+			watchdog.allocatedMemory.insert(newData);
+		}
 		return newData;
 	}
 
+	EXPORT void ext_free(void* _data)
+	{
+		std::free(_data);
+		if (_data)
+		{
+			std::cerr << "FREE   : " << _data << std::endl;
+			watchdog.allocatedMemory.erase(_data);
+		}
+	}
+
 } // extern "C"
 
diff --git a/libevmjit/Stack.h b/libevmjit/Stack.h
index 242063962..ac64d427e 100644
--- a/libevmjit/Stack.h
+++ b/libevmjit/Stack.h
@@ -38,6 +38,7 @@ public:
 	llvm::Value* get(llvm::Value* _index) { return m_getFunc.call(m_builder, {m_array, _index}); }
 	void pop(llvm::Value* _count);
 	llvm::Value* size();
+	void free() { m_freeFunc.call(m_builder, {m_array}); }
 
 	llvm::Value* getPointerTo() const { return m_array; }
 
@@ -47,10 +48,12 @@ private:
 	LazyFunction m_pushFunc;
 	LazyFunction m_setFunc;
 	LazyFunction m_getFunc;
+	LazyFunction m_freeFunc;
 
 	llvm::Function* createArrayPushFunc();
 	llvm::Function* createArraySetFunc();
 	llvm::Function* createArrayGetFunc();
+	llvm::Function* createFreeFunc();
 };
 
 class Stack : public CompilerHelper
@@ -62,6 +65,7 @@ public:
 	void set(size_t _index, llvm::Value* _value);
 	void pop(size_t _count);
 	void push(llvm::Value* _value);
+	void free() { m_stack.free(); }
 
 private:
 	llvm::Function* getPopFunc();