MemStruct: dyn types (returned by mem*()) are now cached

2 files changed, 53 insertions, 26 deletions

diff --git a/miasm2/analysis/mem.py b/miasm2/analysis/mem.py
index da12f5ba..2e52ec1a 100644
--- a/miasm2/analysis/mem.py
+++ b/miasm2/analysis/mem.py
@@ -85,8 +85,11 @@ console_handler.setFormatter(logging.Formatter("%(levelname)-5s: %(message)s"))
 log.addHandler(console_handler)
 log.setLevel(logging.WARN)
 
-# allocator is a function(vm, size) -> allocated_address
-allocator = None
+# ALLOCATOR is a function(vm, size) -> allocated_address
+ALLOCATOR = None
+
+# Cache for dynamically generated MemStructs
+DYN_MEM_STRUCT_CACHE = {}
 
 def set_allocator(alloc_func):
     """Set an allocator for this module; allows to instanciate statically sized
@@ -96,8 +99,8 @@ def set_allocator(alloc_func):
     Args:
         alloc_func: func(VmMngr) -> integer_address
     """
-    global allocator
-    allocator = alloc_func
+    global ALLOCATOR
+    ALLOCATOR = alloc_func
 
 
 # Helpers
@@ -167,14 +170,17 @@ def set_str_utf16(vm, addr, s):
 
 # MemField to MemStruct helper
 
-# TODO: cache generated types
 def mem(field):
     """Generate a MemStruct subclass from a field. The field's value can
     be accessed through self.value or self.deref_value if field is a Ptr.
     """
+    if field in DYN_MEM_STRUCT_CACHE:
+        return DYN_MEM_STRUCT_CACHE[field]
+
     fields = [("value", field)]
     # Build a type to contain the field type
     mem_type = type("Mem%r" % field, (MemStruct,), {'fields': fields})
+    DYN_MEM_STRUCT_CACHE[field] = mem_type
     return mem_type
 
 
@@ -356,7 +362,7 @@ class Ptr(Num):
 
     def __hash__(self):
         return hash(super(Ptr, self).__hash__() + hash(self._dst_type) +
-                hash(self._type_args) + hash(self._type_kwargs))
+                hash(self._type_args))
 
 
 class Inline(MemField):
@@ -408,7 +414,7 @@ class Inline(MemField):
 
     def __hash__(self):
         return hash(hash(self.__class__) + hash(self._il_type) +
-                hash(self._type_args) + hash(self._type_kwargs))
+                hash(self._type_args))
 
 
 class Array(MemField):
@@ -521,7 +527,7 @@ class Union(MemField):
                 self.field_list == other.field_list
 
     def __hash__(self):
-        return hash(hash(self.__class__) + hash(self.field_list))
+        return hash(hash(self.__class__) + hash(tuple(self.field_list)))
 
 
 class Bits(MemField):
@@ -716,14 +722,15 @@ class MemStruct(object):
     # Classic usage methods
 
     def __init__(self, vm, addr=None, *args, **kwargs):
-        global allocator
+        global ALLOCATOR
         super(MemStruct, self).__init__(*args, **kwargs)
         self._vm = vm
         if addr is None:
-            if allocator is None:
+            if ALLOCATOR is None:
                 raise ValueError("Cannot provide None address to MemStruct() if"
-                                 "%s.allocator is not set." % __name__)
-            self._addr = allocator(vm, self.get_size())
+                                 "%s.set_allocator has not been called."
+                                 % __name__)
+            self._addr = ALLOCATOR(vm, self.get_size())
         else:
             self._addr = addr
 
@@ -1151,17 +1158,6 @@ class MemArray(MemStruct):
         return "[%r, ...] [%r]" % (self[0], self._field_type)
 
 
-def mem_array_type(field_type):
-    """Generate a MemArray subclass that has a fixed @field_type. It allows to
-    instanciate this class with only vm and addr argument, as are standard
-    MemStructs.
-    """
-    array_type = type('MemArray_%r' % (field_type,),
-                      (MemArray,),
-                      {'_field_type': field_type})
-    return array_type
-
-
 class MemSizedArray(MemArray):
     """A fixed size MemArray. Its additional arg represents the @array_len (in
     number of elements) of this array.
@@ -1221,11 +1217,31 @@ class MemSizedArray(MemArray):
         return "[%s] [%r; %s]" % (items, self._field_type, self._array_len)
 
 
+def mem_array_type(field_type):
+    """Generate a MemArray subclass that has a fixed @field_type. It allows to
+    instanciate this class with only vm and addr argument, as are standard
+    MemStructs.
+    """
+    cache_key = (field_type, None)
+    if cache_key in DYN_MEM_STRUCT_CACHE:
+        return DYN_MEM_STRUCT_CACHE[cache_key]
+
+    array_type = type('MemArray_%r' % (field_type,),
+                      (MemArray,),
+                      {'_field_type': field_type})
+    DYN_MEM_STRUCT_CACHE[cache_key] = array_type
+    return array_type
+
+
 def mem_sized_array_type(field_type, array_len):
     """Generate a MemSizedArray subclass that has a fixed @field_type and a
     fixed @array_len. This allows to instanciate the returned type with only
     the vm and addr arguments, as are standard MemStructs.
     """
+    cache_key = (field_type, array_len)
+    if cache_key in DYN_MEM_STRUCT_CACHE:
+        return DYN_MEM_STRUCT_CACHE[cache_key]
+
     @classmethod
     def sizeof(cls):
         return cls._field_type.size() * cls._array_len
@@ -1235,5 +1251,6 @@ def mem_sized_array_type(field_type, array_len):
                       {'_array_len': array_len,
                        '_field_type': field_type,
                        'sizeof': sizeof})
+    DYN_MEM_STRUCT_CACHE[cache_key] = array_type
     return array_type
 
diff --git a/test/analysis/mem.py b/test/analysis/mem.py
index a3642a4f..d0590ebc 100644
--- a/test/analysis/mem.py
+++ b/test/analysis/mem.py
@@ -292,9 +292,6 @@ assert cont.instruct.bar == 0x03
 assert cont.last == 0x04
 assert jitter.vm.get_mem(cont.get_addr(), len(cont)) == '\x01\x02\x03\x04'
 
-# Quick mem(MemField) test:
-assert mem(Num("f"))(jitter.vm, addr) == mem(Num("f"))(jitter.vm, addr)
-
 
 # Union test
 class UniStruct(MemStruct):
@@ -464,6 +461,19 @@ assert BitField(Num("B"), [("f1", 1), ("f2", 4), ("f3", 1)]) != \
         BitField(Num("B"), [("f1", 2), ("f2", 4), ("f3", 1)])
 
 
+# Quick mem(MemField)/MemField hash test:
+assert mem(Num("f"))(jitter.vm, addr) == mem(Num("f"))(jitter.vm, addr)
+# Types are cached
+assert mem(Num("f")) == mem(Num("f"))
+assert mem(Num("d")) != mem(Num("f"))
+assert mem(Union([("f1", Num("I")), ("f2", Num("H"))])) == \
+        mem(Union([("f1", Num("I")), ("f2", Num("H"))]))
+assert mem_array_type(Num("B")) == mem_array_type(Num("B"))
+assert mem_array_type(Num("I")) != mem_array_type(Num("B"))
+assert mem_sized_array_type(Num("B"), 20) == mem_sized_array_type(Num("B"), 20)
+assert mem_sized_array_type(Num("B"), 19) != mem_sized_array_type(Num("B"), 20)
+
+
 # Repr tests
 
 print "Some struct reprs:\n"