Merge pull request #85 from commial/bounded-dict

Bounded dict

2 files changed, 18 insertions, 7 deletions

diff --git a/miasm2/jitter/jitcore.py b/miasm2/jitter/jitcore.py
index f1e34870..16f201c9 100644
--- a/miasm2/jitter/jitcore.py
+++ b/miasm2/jitter/jitcore.py
@@ -17,6 +17,7 @@
 #
 from miasm2.core import asmbloc
 from miasm2.core.interval import interval
+from miasm2.core.utils import BoundedDict
 from miasm2.jitter.csts import *
 
 
@@ -24,6 +25,9 @@ class JitCore(object):
 
     "JiT management. This is an abstract class"
 
+    jitted_block_delete_cb = None
+    jitted_block_max_size = 10000
+
     def __init__(self, ir_arch, bs=None):
         """Initialise a JitCore instance.
         @ir_arch: ir instance for current architecture
@@ -33,7 +37,8 @@ class JitCore(object):
         self.ir_arch = ir_arch
         self.bs = bs
         self.known_blocs = {}
-        self.lbl2jitbloc = {}
+        self.lbl2jitbloc = BoundedDict(self.jitted_block_max_size,
+                                       delete_cb=self.jitted_block_delete_cb)
         self.lbl2bloc = {}
         self.log_mn = False
         self.log_regs = False
@@ -67,7 +72,7 @@ class JitCore(object):
     def load(self, arch, attrib):
         "Initialise the Jitter according to arch and attrib"
 
-        raise Exception("DO NOT instanciate JitCore")
+        raise NotImplementedError("Abstract class")
 
     def get_bloc_min_max(self, cur_bloc):
         "Update cur_bloc to set min/max address"
@@ -91,7 +96,7 @@ class JitCore(object):
         @irblocs: a gorup of irblocs
         """
 
-        raise Exception("DO NOT instanciate JitCore")
+        raise NotImplementedError("Abstract class")
 
     def add_bloc(self, b):
         """Add a bloc to JiT and JiT it.
diff --git a/miasm2/jitter/jitcore_tcc.py b/miasm2/jitter/jitcore_tcc.py
index 46b46184..7ea77b15 100644
--- a/miasm2/jitter/jitcore_tcc.py
+++ b/miasm2/jitter/jitcore_tcc.py
@@ -92,12 +92,19 @@ class JitCore_Tcc(jitcore.JitCore):
     "JiT management, using LibTCC as backend"
 
     def __init__(self, ir_arch, bs=None):
+        self.jitted_block_delete_cb = self.deleteCB
         super(JitCore_Tcc, self).__init__(ir_arch, bs)
         self.resolver = resolver()
         self.exec_wrapper = Jittcc.tcc_exec_bloc
-        self.tcc_states =[]
+        self.tcc_states = {}
         self.ir_arch = ir_arch
 
+    def deleteCB(self, offset):
+        "Free the TCCState corresponding to @offset"
+        if offset in self.tcc_states:
+            Jittcc.tcc_end(self.tcc_states[offset])
+            del self.tcc_states[offset]
+
     def load(self):
         # os.path.join(os.path.dirname(os.path.realpath(__file__)), "jitter")
         lib_dir = os.path.dirname(os.path.realpath(__file__))
@@ -120,12 +127,11 @@ class JitCore_Tcc(jitcore.JitCore):
         include_files = [x[1:]
             for x in include_files if x.startswith(' /usr/include')]
         include_files += [include_dir, get_python_inc()]
-
         include_files = ";".join(include_files)
         Jittcc.tcc_set_emul_lib_path(include_files, libs)
 
     def __del__(self):
-        for tcc_state in self.tcc_states:
+        for tcc_state in self.tcc_states.values():
             Jittcc.tcc_end(tcc_state)
 
     def jitirblocs(self, label, irblocs):
@@ -143,8 +149,8 @@ class JitCore_Tcc(jitcore.JitCore):
         # open('tmp_%.4d.c'%self.jitcount, "w").write(func_code)
         self.jitcount += 1
         tcc_state, mcode = jit_tcc_compil(f_name, func_code)
-        self.tcc_states.append(tcc_state)
         jcode = jit_tcc_code(mcode)
         self.lbl2jitbloc[label.offset] = mcode
+        self.tcc_states[label.offset] = tcc_state
         self.addr2obj[label.offset] = jcode
         self.addr2objref[label.offset] = objref(jcode)