Core/cpu: improve asm fromstring

Determine label/integer sizes using context
Default size in other cases

7 files changed, 128 insertions, 114 deletions

diff --git a/miasm2/arch/aarch64/arch.py b/miasm2/arch/aarch64/arch.py
index c875d787..9f265e75 100644
--- a/miasm2/arch/aarch64/arch.py
+++ b/miasm2/arch/aarch64/arch.py
@@ -74,34 +74,29 @@ _, _, base_expr64 = gen_base_expr()
 
 
 def ast_id2expr32(t):
-    if not t in mn_aarch64.regs.all_regs_ids_byname:
-        r = m2_expr.ExprId(asm_label(t))
-    else:
-        r = mn_aarch64.regs.all_regs_ids_byname[t]
-    if not r.size == 32:
-        raise StopIteration
-    return r
-
+    if t in mn_aarch64.regs.all_regs_ids_byname:
+        t = mn_aarch64.regs.all_regs_ids_byname[t]
+        if not t.size == 32:
+            raise StopIteration
+    return t
 
 def ast_int2expr32(a):
     return m2_expr.ExprInt32(a)
 
 
 def ast_id2expr64(t):
-    if not t in mn_aarch64.regs.all_regs_ids_byname:
-        r = m2_expr.ExprId(asm_label(t))
-    else:
-        r = mn_aarch64.regs.all_regs_ids_byname[t]
-    if not r.size == 64:
-        raise StopIteration
-    return r
+    if t in mn_aarch64.regs.all_regs_ids_byname:
+        t = mn_aarch64.regs.all_regs_ids_byname[t]
+        if not t.size == 64:
+            raise StopIteration
+    return t
 
 
 def ast_int2expr64(a):
     return m2_expr.ExprInt64(a)
 
-my_var_parser32 = parse_ast(ast_id2expr32, ast_int2expr32)
-my_var_parser64 = parse_ast(ast_id2expr64, ast_int2expr64)
+my_var_parser32 = parse_ast(ast_id2expr32, ast_int2expr32, default_size=32)
+my_var_parser64 = parse_ast(ast_id2expr64, ast_int2expr64, default_size=64)
 
 base_expr32.setParseAction(my_var_parser32)
 base_expr64.setParseAction(my_var_parser64)
diff --git a/miasm2/arch/arm/arch.py b/miasm2/arch/arm/arch.py
index 0e58008d..9a19d03b 100644
--- a/miasm2/arch/arm/arch.py
+++ b/miasm2/arch/arm/arch.py
@@ -181,11 +181,7 @@ int_or_expr = base_expr
 
 
 def ast_id2expr(t):
-    if not t in mn_arm.regs.all_regs_ids_byname:
-        r = ExprId(asm_label(t))
-    else:
-        r = mn_arm.regs.all_regs_ids_byname[t]
-    return r
+    return mn_arm.regs.all_regs_ids_byname.get(t, t)
 
 
 def ast_int2expr(a):
diff --git a/miasm2/arch/mips32/arch.py b/miasm2/arch/mips32/arch.py
index 2ac16770..68665153 100644
--- a/miasm2/arch/mips32/arch.py
+++ b/miasm2/arch/mips32/arch.py
@@ -53,11 +53,7 @@ int_or_expr = base_expr
 
 
 def ast_id2expr(t):
-    if not t in mn_mips32.regs.all_regs_ids_byname:
-        r = ExprId(asm_label(t))
-    else:
-        r = mn_mips32.regs.all_regs_ids_byname[t]
-    return r
+    return mn_mips32.regs.all_regs_ids_byname.get(t, t)
 
 
 def ast_int2expr(a):
diff --git a/miasm2/arch/msp430/arch.py b/miasm2/arch/msp430/arch.py
index a9f695ec..8d4ae956 100644
--- a/miasm2/arch/msp430/arch.py
+++ b/miasm2/arch/msp430/arch.py
@@ -72,11 +72,7 @@ PINC = Suppress("+")
 
 
 def ast_id2expr(t):
-    if not t in mn_msp430.regs.all_regs_ids_byname:
-        r = ExprId(asm_label(t), 16)
-    else:
-        r = mn_msp430.regs.all_regs_ids_byname[t]
-    return r
+    return mn_msp430.regs.all_regs_ids_byname.get(t, t)
 
 
 def ast_int2expr(a):
diff --git a/miasm2/arch/sh4/arch.py b/miasm2/arch/sh4/arch.py
index d72e6945..c4e462b9 100644
--- a/miasm2/arch/sh4/arch.py
+++ b/miasm2/arch/sh4/arch.py
@@ -36,11 +36,7 @@ def parse_pcandimmimm(t):
     return (t[0] & t[1]) + t[2]
 
 def ast_id2expr(t):
-    if not t in mn_sh4.regs.all_regs_ids_byname:
-        r = ExprId(asm_label(t))
-    else:
-        r = mn_sh4.regs.all_regs_ids_byname[t]
-    return r
+    return mn_sh4.regs.all_regs_ids_byname.get(t, t)
 
 def ast_int2expr(a):
     return ExprInt32(a)
diff --git a/miasm2/arch/x86/arch.py b/miasm2/arch/x86/arch.py
index 20fdc1cf..5b31ba55 100644
--- a/miasm2/arch/x86/arch.py
+++ b/miasm2/arch/x86/arch.py
@@ -223,11 +223,7 @@ variable, operand, base_expr = gen_base_expr()
 
 
 def ast_id2expr(t):
-    if not t in mn_x86.regs.all_regs_ids_byname:
-        r = ExprId(asm_label(t))
-    else:
-        r = mn_x86.regs.all_regs_ids_byname[t]
-    return r
+    return mn_x86.regs.all_regs_ids_byname.get(t, t)
 
 
 def ast_int2expr(a):
diff --git a/miasm2/core/cpu.py b/miasm2/core/cpu.py
index 22f4c8ab..4dd7fba3 100644
--- a/miasm2/core/cpu.py
+++ b/miasm2/core/cpu.py
@@ -199,84 +199,123 @@ def ast_int2expr(a):
     return m2_expr.ExprInt32(a)
 
 
-def ast_raw2expr(a, my_id2expr, my_int2expr):
-    assert(isinstance(a, tuple))
-    if a[0] is m2_expr.ExprId:
-        e = my_id2expr(a[1])
-    elif a[0] is m2_expr.ExprInt:
-        e = my_int2expr(a[1])
-    elif a[0] is m2_expr.ExprOp:
-        out = []
-        for x in a[1]:
-            if isinstance(x, tuple):
-                x = ast_raw2expr(x, my_id2expr, my_int2expr)
-            out.append(x)
-        e = ast_parse_op(out)
-    else:
-        raise TypeError('unknown type')
-    return e
-
-
-def ast_get_ids(a):
-    assert(isinstance(a, tuple))
-    if a[0] is m2_expr.ExprId:
-        return set([a[1]])
-    elif a[0] is m2_expr.ExprInt:
-        return set()
-    elif a[0] is m2_expr.ExprOp:
-        out = set()
-        for x in a[1]:
-            if isinstance(x, tuple):
-                out.update(ast_get_ids(x))
-        return out
-    raise TypeError('unknown type')
-
-
-def _extract_ast_core(a):
-    assert(isinstance(a, tuple))
-    if a[0] in [m2_expr.ExprInt, m2_expr.ExprId]:
-        return a
-    elif a[0] is m2_expr.ExprOp:
-        out = []
-        for x in a[1]:
-            if isinstance(x, tuple):
-                x = _extract_ast_core(x)
-            out.append(x)
-        return tuple([a[0]] + [out])
-    else:
-        raise TypeError('unknown type')
 
+class parse_ast(object):
 
-def extract_ast_core(v, my_id2expr, my_int2expr):
-    ast_tokens = _extract_ast_core(v)
-    ids = ast_get_ids(ast_tokens)
-    ids_expr = [my_id2expr(x) for x in ids]
-    sizes = set([i.size for i in ids_expr])
-
-    if len(sizes) == 0:
-        pass
-    elif len(sizes) == 1:
-        size = sizes.pop()
-        my_int2expr = lambda x: m2_expr.ExprInt(x, size)
-    else:
-        # Multiple sizes in ids
-        raise StopIteration
-    e = ast_raw2expr(ast_tokens, my_id2expr, my_int2expr)
-    return e
+    def __init__(self, id2expr, int2expr, default_size=32):
+        self.id2expr = id2expr
+        self.int2expr = int2expr
+        self.default_size = default_size
+
+    def int_from_size(self, size, value):
+        """Transform a string into ExprInt.
+        * if @size is None, use provided int2expr
+        * else, use @size to generate integer
+        @size: size of int; None if not forced.
+        @value: string representing an integer
+        """
+        if size is None:
+            return self.int2expr(value)
+        else:
+            return m2_expr.ExprInt(value, size)
+
+    def id_from_size(self, size, value):
+        """Transform a string into ExprId.
+        * if @size is None, use provided id2expr
+        * else, use @size to generate id
+        @size: size of id; None if not forced.
+        @value: string representing the id
+        """
+        value = self.id2expr(value)
+        if isinstance(value, m2_expr.Expr):
+            return value
+        if size is None:
+            size = self.default_size
+        assert value is not None
+        return m2_expr.ExprId(asmbloc.asm_label(value), size)
+
+    def ast_to_expr(self, size, ast):
+        """Transform a typed ast into a Miasm expression
+        @size: default size
+        @ast: typed ast
+        """
+        assert(isinstance(ast, tuple))
+        if ast[0] is m2_expr.ExprId:
+            expr = self.id_from_size(size, ast[1])
+            if isinstance(expr, str):
+                expr = self.id_from_size(size, expr)
+        elif ast[0] is m2_expr.ExprInt:
+            expr = self.int_from_size(size, ast[1])
+        elif ast[0] is m2_expr.ExprOp:
+            out = []
+            for arg in ast[1]:
+                if isinstance(arg, tuple):
+                    arg = self.ast_to_expr(size, arg)
+                out.append(arg)
+            expr = ast_parse_op(out)
+        else:
+            raise TypeError('unknown type')
+        return expr
 
+    def ast_get_ids(self, ast):
+        """Retrieve every node of type ExprId in @ast
+        @ast: typed ast
+        """
+        assert(isinstance(ast, tuple))
+        if ast[0] is m2_expr.ExprId:
+            return set([ast[1]])
+        elif ast[0] is m2_expr.ExprInt:
+            return set()
+        elif ast[0] is m2_expr.ExprOp:
+            out = set()
+            for x in ast[1]:
+                if isinstance(x, tuple):
+                    out.update(self.ast_get_ids(x))
+            return out
+        raise TypeError('unknown type')
 
-class parse_ast:
+    def _extract_ast_core(self, ast):
+        assert(isinstance(ast, tuple))
+        if ast[0] in [m2_expr.ExprInt, m2_expr.ExprId]:
+            return ast
+        elif ast[0] is m2_expr.ExprOp:
+            out = []
+            for arg in ast[1]:
+                if isinstance(arg, tuple):
+                    arg = self._extract_ast_core(arg)
+                out.append(arg)
+            return tuple([ast[0]] + [out])
+        else:
+            raise TypeError('unknown type')
 
-    def __init__(self, id2expr, int2expr, extract_ast=extract_ast_core):
-        self.id2expr = id2expr
-        self.int2expr = int2expr
-        self.extract_ast_core = extract_ast
+    def extract_ast_core(self, ast):
+        """
+        Trasform an @ast into a Miasm expression.
+        Use registers size to deduce label and integers sizes.
+        """
+        ast = self._extract_ast_core(ast)
+        ids = self.ast_get_ids(ast)
+        ids_expr = [self.id2expr(x) for x in ids]
+        sizes = set([expr.size for expr in ids_expr
+                     if isinstance(expr, m2_expr.Expr)])
+        if not sizes:
+            size = None
+        elif len(sizes) == 1:
+            size = sizes.pop()
+        else:
+            # Multiple sizes in ids
+            raise StopIteration
+        return self.ast_to_expr(size, ast)
 
-    def __call__(self, v):
-        v = v[0]
-        if isinstance(v, m2_expr.Expr):
-            return v
-        return self.extract_ast_core(v, self.id2expr, self.int2expr)
+    def __call__(self, ast):
+        """
+        Trasform an @ast into a Miasm expression.
+        Use registers size to deduce label and integers sizes.
+        """
+        ast = ast[0]
+        if isinstance(ast, m2_expr.Expr):
+            return ast
+        return self.extract_ast_core(ast)
 
 
 def neg_int(t):