core: fix add_label api for asm_bloc (louis granboulan)

6 files changed, 51 insertions, 46 deletions

diff --git a/example/asm_arm.py b/example/asm_arm.py
index b146a362..2072f519 100755
--- a/example/asm_arm.py
+++ b/example/asm_arm.py
@@ -44,7 +44,7 @@ open("graph.txt" , "w").write(g)
 
 for b in all_bloc[0]:
     print b
-symbol_pool.add(asmbloc.asm_label('base_address', 0x0))
+symbol_pool.add_label('base_address', 0x0)
 symbol_pool.getby_name("toto").offset = 0x0
 
 resolved_b, patches = asmbloc.asm_resolve_final(my_mn, all_bloc[0], symbol_pool)
diff --git a/example/asm_box.py b/example/asm_box.py
index 86ff7e71..f1e4bde5 100644
--- a/example/asm_box.py
+++ b/example/asm_box.py
@@ -26,7 +26,7 @@ title:
 msg:
 .string "World!"
 ''')
-symbol_pool.add(asmbloc.asm_label('base_address', 0))
+symbol_pool.add_label('base_address', 0)
 symbol_pool.getby_name("MessageBoxA").offset = e.DirImport.get_funcvirt('MessageBoxA')
 symbol_pool.getby_name("main").offset = e.rva2virt(s_text.addr)
 resolved_b, patches = asmbloc.asm_resolve_final(x86_mn, all_bloc[0], symbol_pool)
diff --git a/example/expression/manip_expression4.py b/example/expression/manip_expression4.py
index 2cc4abf9..65f1f2f4 100644
--- a/example/expression/manip_expression4.py
+++ b/example/expression/manip_expression4.py
@@ -91,7 +91,7 @@ in_str = bin_stream(data)
 
 job_done = set()
 symbol_pool = asmbloc.asm_symbol_pool()
-l = asmbloc.asm_label('toto')
+l = symbol_pool.add_label('toto')
 b = asmbloc.asm_bloc(l)
 
 ad = 0x2E
diff --git a/miasm/core/asmbloc.py b/miasm/core/asmbloc.py
index 63d49ad7..f2aa9cba 100644
--- a/miasm/core/asmbloc.py
+++ b/miasm/core/asmbloc.py
@@ -130,25 +130,32 @@ class asm_bloc:
         return [x.offset for x in self.lines]
 
 class asm_symbol_pool:
-    def __init__(self):
+    def __init__(self, no_collision = True):
+        self.labels = []
         self.s = {}
         self.s_offset = {}
+        self.no_collision = no_collision
 
-    def add(self, obj):
+    def add_label(self, name = "", offset = None):
         """
-        obj can be an asm_label or an offset
+        This should be the only method to create new asm_label objects
         """
-        if not isinstance(obj, asm_label):
-            # whatever obj is, it will be an offset
-            offset = int(obj)
-            if offset in self.s_offset:
-                raise ValueError('symbol offset already exist %r'%offset)
-            obj = asm_label(offset, offset)
-        if obj.name in self.s:
-            raise ValueError('symbol already exist %r'%obj)
-        self.s[obj.name] = obj
-        if obj.offset != None:
-            self.s_offset[obj.offset] = obj
+        l = asm_label(name, offset)
+        collision = None
+        if l.offset in self.s_offset and l != self.s_offset[l.offset]:
+            collision = 'offset'
+        if l.name in self.s and l != self.s[l.name]:
+            collision = 'name'
+        if self.no_collision and collision == 'offset':
+            raise ValueError('symbol %s has same offset as %s'%(l, self.s_offset[l.offset]))
+        if self.no_collision and collision == 'name':
+            raise ValueError('symbol %s has same name as %s'%(l, self.s[l.name]))
+        self.labels.append(l)
+        if l.offset != None and collision == None:
+            self.s_offset[l.offset] = l
+        if l.name != "" and collision == None:
+            self.s[l.name] = l
+        return l
 
     def remove(self, obj):
         """
@@ -178,29 +185,27 @@ class asm_symbol_pool:
                 self.s[l].offset = None
 
     def getby_offset(self, offset):
-        if offset in self.s_offset:
-            return self.s_offset[offset]
+        l = asm_label(offset, offset)
+        if l.offset in self.s_offset:
+            return self.s_offset[l.offset]
         return None
 
     def getby_name(self, name):
-        if name in self.s:
-            return self.s[name]
+        l = asm_label(name)
+        if l.name in self.s:
+            return self.s[l.name]
         return None
 
     def getby_name_create(self, name):
-        if name in self.s:
-            return self.s[name]
-        else:
-            l = asm_label(name)
-            self.add(l)
+        l = self.getby_name(name)
+        if l == None:
+            l = self.add_label(name)
         return l
 
     def getby_offset_create(self, offset):
-        if offset in self.s_offset:
-            return self.s_offset[offset]
-        else:
-            l = asm_label(offset, offset)
-            self.add(l)
+        l = self.getby_offset(offset)
+        if l == None:
+            l = self.add_label(offset, offset)
         return l
 
     def rename(self, s, newname):
@@ -211,6 +216,9 @@ class asm_symbol_pool:
         s.name = newname
         self.s[s.name] = s
 
+    def items(self):
+        return self.labels[:]
+
     def __str__(self):
         return reduce(lambda x,y: x+str(y)+'\n', [self.s[l] for l in self.s], "")
 
@@ -535,8 +543,9 @@ def gen_free_space_intervals(f):
     return interval
 
 def add_dont_erase(f, dont_erase = []):
+    tmp_symbol_pool = asm_symbol_pool()
     for a,b in dont_erase:
-        l = asm_label(a, a)
+        l = tmp_symbol_pool.add_label(a, a)
         l.offset_free = a
         f[l] = b-a
     return
@@ -784,8 +793,7 @@ def calc_symbol_offset(symbol_pool):
 
     s_dependent = {}
 
-    for l in symbol_pool.s:
-        label = symbol_pool.s[l]
+    for label in symbol_pool.items():
         if label.offset == None:
             raise ValueError("symbol missing?", l)
         if not is_int(label.offset):
@@ -911,13 +919,12 @@ def asmbloc_final(mnemo, all_blocs, symbol_pool, symb_reloc_off = {}):
                 if my_s!=None:
                     my_symb_reloc_off[b.label].append(offset_i+my_s)
                 offset_i+=len(instr.data)
-    for l in symbol_pool.s:
-        if symbol_pool.s[l].offset_g ==None:
-            fdfd
     #we have fixed all relative values
     #recompute good offsets
-    for l in symbol_pool.s:
-        symbol_pool.s[l].offset = symbol_pool.s[l].offset_g
+    for label in symbol_pool.items():
+        if label.offset_g == None:
+            fdfd
+        label.offset = label.offset_g
 
     for a, b in my_symb_reloc_off.items():
         symb_reloc_off[a] = b
diff --git a/miasm/core/parse_asm.py b/miasm/core/parse_asm.py
index 430df537..7a1dc242 100644
--- a/miasm/core/parse_asm.py
+++ b/miasm/core/parse_asm.py
@@ -23,10 +23,10 @@ declarator = {'byte':'B', 'long':'I'}
 def guess_next_new_label(symbol_pool, gen_label_index = 0):
     i = 0
     while True:
-        l = asm_label(i)
+        l = symbol_pool.getby_name(i)
+        if l == None:
+            return symbol_pool.add_label(i)
         i+=1
-        if not l.name in symbol_pool.s:
-            return l
         
     
 def parse_txt(mnemo, txt, symbol_pool = None, gen_label_index = 0):
@@ -154,7 +154,6 @@ def parse_txt(mnemo, txt, symbol_pool = None, gen_label_index = 0):
             if state == 0:
                 if not isinstance(lines[i], asm_label):
                     l = guess_next_new_label(symbol_pool)
-                    symbol_pool.add(l)
                     lines[i:i] = [l]
                 else:
                     l = lines[i]
@@ -208,7 +207,6 @@ def parse_txt(mnemo, txt, symbol_pool = None, gen_label_index = 0):
                                 l = lines[i+1]
                             else:
                                 l = guess_next_new_label(symbol_pool)
-                                symbol_pool.add(l)
                                 lines[i+1:i+1] = [l]
                         else:
                             state = 0
diff --git a/miasm/tools/to_c_helper.py b/miasm/tools/to_c_helper.py
index 901e3179..38f21272 100644
--- a/miasm/tools/to_c_helper.py
+++ b/miasm/tools/to_c_helper.py
@@ -688,7 +688,7 @@ def asm2C(f_name, known_mems, dyn_func, in_str, x86_mn, symbol_pool, func_to_dis
                 if not l_name:
                     continue
 
-                label = asmbloc.asm_label(l_name, i)
+                label = symbol_pool.add_label(l_name, i)
                 a[x86_afs.symb] = {label:1}
                 del a[x86_afs.imm]
 
@@ -707,7 +707,7 @@ def asm2C(f_name, known_mems, dyn_func, in_str, x86_mn, symbol_pool, func_to_dis
         l_name = "loc_%.16X"%(f&mask_int)
         funcs_code[-1:-1] = [l_name+":"]
         funcs_code[-1:-1] = f_code.split('\n')
-        l = asmbloc.asm_label(l_name, f)
+        l = symbol_pool.add_label(l_name, f)
         b = asmbloc.asm_bloc(l)
         #all_bloc_funcs[f] = [b]
         all_bloc_funcs += [b]