3 files changed, 72 insertions, 1 deletions

diff --git a/miasm2/ir/ir.py b/miasm2/ir/ir.py
index 0995b86a..caf2c0ed 100644
--- a/miasm2/ir/ir.py
+++ b/miasm2/ir/ir.py
@@ -30,10 +30,21 @@ from miasm2.core.graph import DiGraph
 
 
 class AssignBlock(object):
+    """Represent parallel IR assignment, such as:
+    EAX = EBX
+    EBX = EAX
+
+    Also provides common manipulation on this assignments
+    """
     __slots__ = ["_assigns", "_instr"]
 
     def __init__(self, irs=None, instr=None):
-        """@irs seq"""
+        """Create a new AssignBlock
+        @irs: (optional) sequence of ExprAff, or dictionnary dst (Expr) -> src
+              (Expr)
+        @instr: (optional) associate an instruction with this AssignBlock
+
+        """
         if irs is None:
             irs = []
         self._instr = instr
@@ -49,6 +60,7 @@ class AssignBlock(object):
 
     @property
     def instr(self):
+        """Return the associated instruction, if any"""
         return self._instr
 
     def _set(self, dst, src):
@@ -215,6 +227,18 @@ class AssignBlock(object):
         @dst: Expr instance"""
         return m2_expr.ExprAff(dst, self[dst])
 
+    def simplify(self, simplifier):
+        """Return a new AssignBlock with expression simplified
+        @simplifier: ExpressionSimplifier instance"""
+        new_assignblk = {}
+        for dst, src in self.iteritems():
+            if dst == src:
+                continue
+            src = simplifier(src)
+            dst = simplifier(dst)
+            new_assignblk[dst] = src
+        return AssignBlock(irs=new_assignblk, instr=self.instr)
+
 
 class IRBlock(object):
     """Intermediate representation block object.
diff --git a/test/ir/ir.py b/test/ir/ir.py
new file mode 100644
index 00000000..5c428a94
--- /dev/null
+++ b/test/ir/ir.py
@@ -0,0 +1,46 @@
+from miasm2.expression.expression import *
+from miasm2.ir.ir import AssignBlock
+from miasm2.expression.simplifications import expr_simp
+
+id_a = ExprId("a")
+id_b = ExprId("b")
+int0 = ExprInt(0, id_a.size)
+
+# Test AssignBlock
+## Constructors
+assignblk1 = AssignBlock([ExprAff(id_a, id_b)])
+assignblk2 = AssignBlock({id_a: id_b})
+
+## Equality
+assignblk1_bis = AssignBlock([ExprAff(id_a, id_b)])
+assert assignblk1 == assignblk1_bis
+assert assignblk1 == assignblk2
+
+## Immutability
+try:
+    assignblk1[id_a] = id_a
+except RuntimeError:
+    pass
+else:
+    raise RuntimeError("An error was expected")
+try:
+    del assignblk1[id_a]
+except RuntimeError:
+    pass
+else:
+    raise RuntimeError("An error was expected")
+
+## Basic APIs
+assert assignblk1.get_r() == set([id_b])
+assert assignblk1.get_w() == set([id_a])
+assert assignblk1.get_rw() == {id_a: set([id_b])}
+assert assignblk1.keys() == [id_a]
+assert dict(assignblk1) == {id_a: id_b}
+assert assignblk1[id_a] == id_b
+
+## Simplify
+assignblk3 = AssignBlock({id_a: id_b - id_b})
+assert assignblk3[id_a] != int0
+assignblk4 = assignblk3.simplify(expr_simp)
+assert assignblk3[id_a] != int0
+assert assignblk4[id_a] == int0
diff --git a/test/test_all.py b/test/test_all.py
index 9b3f2dc1..94779898 100755
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -245,6 +245,7 @@ for script in ["modint.py",
 
 ## IR
 for script in ["symbexec.py",
+               "ir.py",
                ]:
     testset += RegressionTest([script], base_dir="ir")