Analysis: fix unssa algorithm

2 files changed, 549 insertions, 113 deletions

diff --git a/example/disasm/full.py b/example/disasm/full.py
index f4c50d9a..8ee4f05c 100644
--- a/example/disasm/full.py
+++ b/example/disasm/full.py
@@ -10,9 +10,11 @@ from miasm2.analysis.data_flow import dead_simp, DiGraphDefUse, \
     ReachingDefinitions, merge_blocks, remove_empty_assignblks, \
     PropagateExpr, replace_stack_vars, load_from_int
 from miasm2.expression.simplifications import expr_simp
-from miasm2.analysis.ssa import SSADiGraph, remove_phi
+from miasm2.analysis.ssa import SSADiGraph, UnSSADiGraph, DiGraphLivenessSSA
 from miasm2.ir.ir import AssignBlock, IRBlock
 
+
+
 log = logging.getLogger("dis")
 console_handler = logging.StreamHandler()
 console_handler.setFormatter(logging.Formatter("%(levelname)-5s: %(message)s"))
@@ -63,6 +65,10 @@ parser.add_argument('-y', "--stack2var", action="store_true",
                     help="*Try* to do transform stack accesses into variables. "
                     "Use only with --propagexpr option. "
                     "WARNING: not reliable, may fail.")
+parser.add_argument('-e', "--loadint", action="store_true",
+                    help="Load integers from binary in fixed memory lookup.")
+parser.add_argument('-j', "--calldontmodstack", action="store_true",
+                    help="Consider stack high is not modified in subcalls")
 
 
 args = parser.parse_args()
@@ -202,12 +208,29 @@ log.info('total lines %s' % total_l)
 if args.propagexpr:
     args.gen_ir = True
 
+
+class IRADelModCallStack(ira):
+
+        def call_effects(self, addr, instr):
+            assignblks, extra = super(IRADelModCallStack, self).call_effects(addr, instr)
+            if not args.calldontmodstack:
+                return assignblks, extra
+            out = []
+            for assignblk in assignblks:
+                dct = dict(assignblk)
+                dct = {
+                    dst:src for (dst, src) in dct.iteritems() if dst != self.sp
+                }
+                out.append(AssignBlock(dct, assignblk.instr))
+            return out, extra
+
+
 # Bonus, generate IR graph
 if args.gen_ir:
     log.info("generating IR and IR analysis")
 
     ir_arch = ir(mdis.loc_db)
-    ir_arch_a = ira(mdis.loc_db)
+    ir_arch_a = IRADelModCallStack(mdis.loc_db)
 
     ircfg = ir_arch.new_ircfg()
     ircfg_a = ir_arch.new_ircfg()
@@ -231,7 +254,9 @@ if args.gen_ir:
         print block
 
     if args.simplify > 0:
+        log.info("dead simp...")
         dead_simp(ir_arch_a, ircfg_a)
+        log.info("ok...")
 
     if args.defuse:
         reachings = ReachingDefinitions(ircfg_a)
@@ -281,7 +306,6 @@ if args.propagexpr:
                         out[reg] = dst
             return set(out.values())
 
-
     # Add dummy dependency to uncover out regs affectation
     for loc in ircfg_a.leaves():
         irblock = ircfg_a.blocks.get(loc)
@@ -327,23 +351,18 @@ if args.propagexpr:
 
     head = list(entry_points)[0]
     heads = set([head])
-    all_ssa_vars = set()
+    all_ssa_vars = {}
 
     propagate_expr = PropagateExpr()
-
-
+    ssa_variable_to_expr = {}
 
     while modified:
         ssa = SSADiGraph(ircfg_a)
         ssa.immutable_ids.update(ssa_forbidden_regs)
-
+        ssa.ssa_variable_to_expr.update(all_ssa_vars)
         ssa.transform(head)
-
         all_ssa_vars.update(ssa.ssa_variable_to_expr)
 
-        ssa_regs = [reg for reg in ssa.expressions if reg.is_id()]
-        ssa_forbidden_regs.update(ssa_regs)
-
         if args.verbose > 3:
             open("ssa_%d.dot" % index, "wb").write(ircfg_a.dot())
 
@@ -369,13 +388,17 @@ if args.propagexpr:
                 if args.verbose > 3:
                     open('tmp_before_%d.dot' % index, 'w').write(ircfg_a.dot())
                 simp_modified = False
+                log.info("dead simp...")
                 simp_modified |= dead_simp(ir_arch_a, ircfg_a)
+                log.info("ok...")
+
                 index += 1
                 if args.verbose > 3:
                     open('tmp_after_%d.dot' % index, 'w').write(ircfg_a.dot())
                 simp_modified |= remove_empty_assignblks(ircfg_a)
                 simp_modified |= merge_blocks(ircfg_a, heads)
-                simp_modified |= load_from_int(ircfg_a, bs, is_addr_ro_variable)
+                if args.loadint:
+                    simp_modified |= load_from_int(ircfg_a, bs, is_addr_ro_variable)
                 modified |= simp_modified
                 index += 1
         if args.verbose > 3:
@@ -391,14 +414,42 @@ if args.propagexpr:
         open('final_merge.dot', 'w').write(ircfg_a.dot())
     ssa = SSADiGraph(ircfg_a)
     ssa.immutable_ids.update(ssa_forbidden_regs)
+    ssa.ssa_variable_to_expr.update(all_ssa_vars)
     ssa.transform(head)
-    all_ssa_vars.update(ssa.ssa_variable_to_expr)
     print '*'*80, "Remove phi"
-    ssa.ssa_variable_to_expr = all_ssa_vars
     if args.verbose > 3:
         open('final_ssa.dot', 'w').write(ircfg_a.dot())
-    remove_phi(ssa, head)
+
+    cfg_liveness = DiGraphLivenessSSA(ircfg_a)
+    cfg_liveness.init_var_info(ir_arch_a)
+    cfg_liveness.compute_liveness()
+
+    unssa = UnSSADiGraph(ssa, head, cfg_liveness)
+
     if args.verbose > 3:
         open('final_no_phi.dot', 'w').write(ircfg_a.dot())
-    dead_simp(ir_arch_a, ircfg_a)
+
+    modified = True
+    while modified:
+        log.debug('Loop %d', index)
+        index += 1
+        modified = False
+        modified |= ircfg_a.simplify(expr_simp)
+        if args.verbose > 3:
+            open('tmp_simp_%d.dot' % index, 'w').write(ircfg_a.dot())
+        simp_modified = True
+        while simp_modified:
+            index += 1
+            if args.verbose > 3:
+                open('tmp_before_%d.dot' % index, 'w').write(ircfg_a.dot())
+            simp_modified = False
+            simp_modified |= dead_simp(ir_arch_a, ircfg_a)
+            index += 1
+            if args.verbose > 3:
+                open('tmp_after_%d.dot' % index, 'w').write(ircfg_a.dot())
+            simp_modified |= remove_empty_assignblks(ircfg_a)
+            simp_modified |= merge_blocks(ircfg_a, heads)
+            modified |= simp_modified
+            index += 1
+
     open('final.dot', 'w').write(ircfg_a.dot())
diff --git a/miasm2/analysis/ssa.py b/miasm2/analysis/ssa.py
index aa978953..e4a6684f 100644
--- a/miasm2/analysis/ssa.py
+++ b/miasm2/analysis/ssa.py
@@ -1,10 +1,10 @@
 from collections import deque
 
+from miasm2.analysis.data_flow import DiGraphLivenessIRA, remove_empty_assignblks
 from miasm2.expression.expression import ExprId, ExprAssign, ExprOp, get_expr_ids
 from miasm2.ir.ir import AssignBlock, IRBlock
 
 
-
 class SSA(object):
     """
     Generic class for static single assignment (SSA) transformation
@@ -250,7 +250,7 @@ class SSA(object):
 
             # transform LHS
             for expr in instructions:
-                if expr.dst in self.immutable_ids:
+                if expr.dst in self.immutable_ids or expr.dst in self.ssa_variable_to_expr:
                     dst_ssa = expr.dst
                 else:
                     dst_ssa = self._transform_expression_lhs(expr.dst)
@@ -372,6 +372,7 @@ class SSADiGraph(SSA):
         self._rename(head)
         self._insert_phi()
         self._convert_phi()
+        self._fix_no_def_var(head)
 
     def reset(self):
         """Resets SSA transformation"""
@@ -400,7 +401,7 @@ class SSADiGraph(SSA):
                     # enforce ExprId
                     if dst.is_id():
                         # exclude immutable ids
-                        if dst in self.immutable_ids:
+                        if dst in self.immutable_ids or dst in self.ssa_variable_to_expr:
                             continue
                         # map variable definition to blocks
                         self.defs.setdefault(dst, set()).add(irblock.loc_key)
@@ -438,13 +439,11 @@ class SSADiGraph(SSA):
                 for node in frontier.get(loc_key, []):
                     if node in done:
                         continue
-
                     # place empty phi functions for a variable
                     empty_phi = self._gen_empty_phi(variable)
 
                     # add empty phi node for variable in node
                     self._phinodes.setdefault(node, {})[variable] = empty_phi.src
-
                     done.add(node)
 
                     if node not in intodo:
@@ -575,120 +574,506 @@ class SSADiGraph(SSA):
             new_irs = IRBlock(loc_key, [assignblk] + list(irblock.assignblks))
             self.ircfg.blocks[loc_key] = new_irs
 
+    def _fix_no_def_var(self, head):
+        """
+        Replace phi source variables which are not ssa vars by ssa vars.
+        @head: loc_key of the graph head
+        """
+        var_to_insert = set()
+        for loc_key in self._phinodes:
+            for dst, sources in self._phinodes[loc_key].iteritems():
+                for src in sources.args:
+                    if src in self.ssa_variable_to_expr:
+                        continue
+                    var_to_insert.add(src)
+        var_to_newname = {}
+        newname_to_var = {}
+        for var in var_to_insert:
+            new_var = self._transform_var_lhs(var)
+            var_to_newname[var] = new_var
+            newname_to_var[new_var] = var
+
+        # Replace non modified node used in phi with new variable
+        self.ircfg.simplify(lambda expr:expr.replace_expr(var_to_newname))
+
+        irblock = self.ircfg.blocks[head]
+        assignblks = list(irblock)
+        assignblks[0:0] = [AssignBlock(newname_to_var, assignblks[0].instr)]
+        self.ircfg.blocks[head] = IRBlock(head, assignblks)
+
+        # Updt structure
+        for loc_key in self._phinodes:
+            for dst, sources in self._phinodes[loc_key].items():
+                self._phinodes[loc_key][dst] = sources.replace_expr(var_to_newname)
+
+        for var, (loc_key, index) in self.ssa_to_location.items():
+            if loc_key == head:
+                self.ssa_to_location[var] = loc_key, index + 1
 
+        for newname, var in newname_to_var.iteritems():
+            self.ssa_to_location[newname] = head, 0
+            self.ssa_variable_to_expr[newname] = var
+            self.immutable_ids.add(newname)
+            self.expressions[newname] = var
 
-def get_assignblk(graph, loc, index):
+
+def irblock_has_phi(irblock):
     """
-    Return the dictionary of the AssignBlock from @graph at location @loc at
-    @index
-    @graph: IRCFG instance
-    @loc: Location instance
-    @index: assignblock index
+    Return True if @irblock has Phi assignments
+    @irblock: IRBlock instance
     """
+    if not irblock.assignblks:
+        return False
+    for src in irblock[0].itervalues():
+        return src.is_op('Phi')
+    return False
 
-    irblock = graph.blocks[loc]
-    assignblks = irblock.assignblks
-    assignblk = assignblks[index]
-    assignblk_dct = dict(assignblk)
-    return assignblk_dct
 
+class VarInfos(object):
+    """Store liveness information for a variable"""
+    __slots__ = ["live_index", "loc_key", "index"]
+
+    def __init__(self, live_index, loc_key, index):
+        self.live_index = live_index
+        self.loc_key = loc_key
+        self.index = index
 
-def set_assignblk(graph, loc, index, assignblk_dct):
-    """
-    Set the Assignblock in @graph at location @loc at @index using dictionary
-    @assignblk_dct
 
-    @graph: IRCFG instance
-    @loc: Location instance
-    @index: assignblock index
-    @assignblk_dct: dictionary representing the AssignBlock
+def get_var_assignment_src(ircfg, node, variables):
     """
+    Return the variable of @variables which is written by the irblock at @node
+    @node: Location
+    @variables: a set of variable to test
+    """
+    irblock = ircfg.blocks[node]
+    for assignblk in irblock:
+        result = set(assignblk).intersection(variables)
+        if not result:
+            continue
+        assert len(result) == 1
+        return list(result)[0]
+    return None
 
-    irblock = graph.blocks[loc]
-    assignblks = list(irblock.assignblks)
-    assignblk = assignblks[index]
 
-    assignblks[index] = AssignBlock(
-        assignblk_dct,
-        assignblk.instr
-    )
-    new_irblock = IRBlock(loc, assignblks)
-    return new_irblock
+def get_phi_sources_parent_block(ircfg, loc_key, sources):
+    """
+    Return a dictionary linking a variable to it's direct parent label
+    which belong to a path which affects the node.
+    @loc_key: the starting node
+    @sources: set of variables to resolve
+    """
+    source_to_parent = {}
+    for parent in ircfg.predecessors(loc_key):
+        done = set()
+        todo = set([parent])
+        found = False
+        while todo:
+            node = todo.pop()
+            if node in done:
+                continue
+            done.add(node)
+            ret = get_var_assignment_src(ircfg, node, sources)
+            if ret:
+                source_to_parent.setdefault(ret, set()).add(parent)
+                found = True
+                break
+            for pred in ircfg.predecessors(node):
+                todo.add(pred)
+        assert found
+    return source_to_parent
 
 
-def remove_phi(ssa, head):
+class UnSSADiGraph(object):
     """
-    Remove Phi using naive algorithm
-    Note: The _ssa_variable_to_expr must be up to date
-
-    @ssa: a SSADiGraph instance
-    @head: the loc_key of the graph head
+    Implements unssa algorithm
+    Revisiting Out-of-SSA Translation for Correctness, Code Quality, and
+    Efficiency
     """
 
-    phivar2var = {}
+    def __init__(self, ssa, head, cfg_liveness):
+        self.cfg_liveness = cfg_liveness
+        self.ssa = ssa
+        self.head = head
+
+        self.copy_vars = set()
+        self.parallel_copy_block_at_end = {}
+        self.parallel_copy_block_at_phi = {}
+        self.phi_new_var = {}
+        self.end_new_var = {}
+
+        self.isolate_phi_nodes_block()
+        self.coalesce_phi_nodes_block()
+        self.order_ssa_var_dom()
+        self.aggressive_coalesce_block()
+        self.insert_parallel_copy()
+        self.replace_merge_sets()
+        self.remove_assign_eq()
+        remove_empty_assignblks(self.ssa.graph)
+
+    def insert_parallel_copy(self):
+        """Insert generated parallel copies"""
+        ircfg = self.ssa.graph
+
+        # Run coalesce on the post phi parallel copy
+        for irblock in ircfg.blocks.values():
+            if not irblock_has_phi(irblock):
+                continue
+            parallel_copies = {}
+            for dst in self.parallel_copy_block_at_phi[irblock.loc_key]:
+                new_var = self.phi_new_var[(irblock.loc_key, dst)]
+                parallel_copies[dst] = new_var
+
+            assignblks = list(irblock)
+            assignblks[0] = AssignBlock(parallel_copies, irblock[0].instr)
+            new_irblock = IRBlock(irblock.loc_key, assignblks)
+            ircfg.blocks[irblock.loc_key] = new_irblock
+
+
+            parent_to_parallel_copies = {}
+            parallel_copies = {}
+            for dst in irblock[0]:
+                new_var = self.phi_new_var[(irblock.loc_key, dst)]
+                for parent, src in self.parallel_copy_block_at_end[(irblock.loc_key, dst)]:
+                    parent_to_parallel_copies.setdefault(parent, {})[new_var] = src
+
+            for parent, parallel_copies in parent_to_parallel_copies.iteritems():
+                parent = ircfg.blocks[parent]
+                assignblks = list(parent)
+                assignblks.append(AssignBlock(parallel_copies, parent[-1].instr))
+                new_irblock = IRBlock(parent.loc_key, assignblks)
+                ircfg.blocks[parent.loc_key] = new_irblock
+
+    def create_copy_var(self, var):
+        """
+        Generate a new var in @phivar2var replacing @var
+        @phivar2var: dictionary associating variable with their replacing one.
+        @var: variable to replace
+        """
+        new_phi_var = ExprId('var%d' % len(self.copy_vars), var.size)
+        self.copy_vars.add(new_phi_var)
+        return new_phi_var
+
+    def isolate_phi_nodes_block(self):
+        """Insert parallel copy before/after each phi node"""
+
+        ircfg = self.ssa.graph
+        for irblock in ircfg.blocks.values():
+            if not irblock_has_phi(irblock):
+                continue
+            all_dst = set()
+            for dst, sources in irblock[0].items():
+                all_dst.add(dst)
+                assert sources.is_op('Phi')
+                new_var = self.create_copy_var(dst)
+                self.phi_new_var[(irblock.loc_key, dst)] = new_var
+
+                var_to_parents = get_phi_sources_parent_block(
+                    self.ssa.graph,
+                    irblock.loc_key,
+                    sources.args
+                )
+                for src in sources.args:
+                    parents = var_to_parents[src]
+                    for parent in parents:
+                        self.parallel_copy_block_at_end.setdefault((irblock.loc_key, dst), set()).add((parent, src))
+                        self.end_new_var.setdefault(new_var, set()).add(parent)
+
+            self.parallel_copy_block_at_phi.setdefault(irblock.loc_key, set()).update(all_dst)
+
+    def coalesce_phi_nodes_block(self):
+        """
+        Generate trivial coalescing of phi variable and its destination
+        """
+        merge_state = {}
+        for phi_new_var in self.phi_new_var.itervalues():
+            merge_state.setdefault(phi_new_var, set([phi_new_var]))
+        self.merge_state = merge_state
+
+    def order_ssa_var_dom(self):
+        """Compute dominance order of each ssa variable"""
+        ircfg = self.ssa.graph
+
+        # compute dominator tree
+        dominator_tree = ircfg.compute_dominator_tree(self.head)
+
+        self.var_to_infos = {}
+        live_index = 0
+        # walk in DFS over the dominator tree
+        for loc_key in dominator_tree.walk_depth_first_forward(self.head):
+            irblock = ircfg.blocks[loc_key]
+
+            # Create live index for phi new vars
+            if irblock_has_phi(irblock):
+                for dst in irblock[0]:
+                    if not dst.is_id():
+                        continue
+                    new_var = self.phi_new_var[(irblock.loc_key, dst)]
+                    self.var_to_infos[new_var] = VarInfos(live_index, loc_key, None)
+
+                live_index += 1
+
+            # Create live index for remaining assignments
+            for index, assignblk in enumerate(irblock):
+                used = False
+                for dst in assignblk:
+                    if not dst.is_id():
+                        continue
+                    if dst in self.ssa.immutable_ids:
+                        continue
+
+                    assert dst not in self.var_to_infos
+                    self.var_to_infos[dst] = VarInfos(live_index, loc_key, index)
+                    used = True
+                if used:
+                    live_index += 1
+
+
+    def ssa_def_dominates(self, node_a, node_b):
+        """
+        Return living index order of @node_a and @node_b
+        @node_a: VarInfos instance
+        @node_b: VarInfos instance
+        """
+        ret = self.var_to_infos[node_a].live_index <= self.var_to_infos[node_b].live_index
+        return ret
+
+    def merge_set_sort(self, merge_set):
+        """
+        Sort the list of nodes @merge_set in dominance order
+        @node_a: VarInfos instance
+        @node_b: VarInfos instance
+        """
+        out = [(self.var_to_infos[var].live_index, var) for var in merge_set]
+        out.sort()
+        return out
+
+    def ssa_def_is_live_at(self, node_a, node_b, parent):
+        """
+        Return true if @node_a is live during @node_b definition
+        if @parent is None, this is a liveness test for a post phi variable;
+        Else, it is a liveness test for a variable source of the phi node
+
+        @node_a: VarInfos instance
+        @node_b: VarInfos instance
+        @parent: Optional parent location of the phi source
+        """
+        loc_key, index = self.var_to_infos[node_b].loc_key, self.var_to_infos[node_b].index
+        if parent and index is None:
+            index = 0
+        if node_a not in self.end_new_var:
+            info = self.cfg_liveness.blocks[loc_key].infos[index]
+            ret = node_a in info.var_out
+            return ret
+
+        for def_loc_key in self.end_new_var[node_a]:
+            if def_loc_key == parent:
+                continue
+            info = self.cfg_liveness.blocks[def_loc_key].infos[-1]
+            if node_b in info.var_out:
+                return True
+        return False
+
+    def merge_nodes_interfere(self, node_a, node_b, parent):
+        """
+        Return True if @node_a and @node_b interfere
+        @node_a: VarInfos instance
+        @node_b: VarInfos instance
+        @parent: Optional parent location of the phi source for liveness tests
+        """
+        if self.var_to_infos[node_a].live_index == self.var_to_infos[node_b].live_index:
+            return True
 
-    all_ssa_vars = ssa._ssa_variable_to_expr
+        if self.var_to_infos[node_a].live_index < self.var_to_infos[node_b].live_index:
+            return self.ssa_def_is_live_at(node_a, node_b, parent)
+        return self.ssa_def_is_live_at(node_b, node_a, parent)
 
-    # Retrieve Phi nodes
-    phi_nodes = []
-    for irblock in ssa.graph.blocks.itervalues():
-        for index, assignblk in enumerate(irblock):
-            for dst, src in assignblk.iteritems():
+    def merge_sets_interfere(self, merge_a, merge_b, parent):
+        """
+        Return True if no variable in @merge_a and @merge_b interferes.
+        @merge_a: a dom ordered list of equivalent variables
+        @merge_b: a dom ordered list of equivalent variables
+        @parent: Optional parent location of the phi source for liveness tests
+        """
+        merge_a_list = self.merge_set_sort(merge_a)
+        merge_b_list = self.merge_set_sort(merge_b)
+        dom = []
+        while merge_a_list or merge_b_list:
+            if not merge_a_list:
+                _, current = merge_b_list.pop(0)
+            elif not merge_b_list:
+                _, current = merge_a_list.pop(0)
+            else:
+                if merge_a_list[-1] < merge_b_list[-1]:
+                    _, current = merge_a_list.pop(0)
+                else:
+                    _, current = merge_b_list.pop(0)
+            while dom and not self.ssa_def_dominates(dom[-1], current):
+                dom.pop()
+
+            # Don't test node in same merge_set
+            if (dom and
+                not (dom[-1] in merge_a and current in merge_a) and
+                not (dom[-1] in merge_b and current in merge_b)):
+                if dom and self.merge_nodes_interfere(current, dom[-1], parent):
+                    return True
+            dom.append(current)
+        return False
+
+    def aggressive_coalesce_parallel_copy(self, parallel_copies, parent):
+        """
+        Try to coalesce variables each dst/src couple together from
+        @parallel_copies
+
+        @parallel_copies: a dictionary representing dst/src parallel
+        assignments.
+        @parent: Optional parent location of the phi source for liveness tests
+        """
+        for dst, src in parallel_copies.iteritems():
+            dst_merge = self.merge_state.setdefault(dst, set([dst]))
+            src_merge = self.merge_state.setdefault(src, set([src]))
+            if dst_merge == src_merge:
+                continue
+            if not self.merge_sets_interfere(dst_merge, src_merge, parent):
+                dst_merge.update(src_merge)
+                for node in dst_merge:
+                    self.merge_state[node] = dst_merge
+
+    def aggressive_coalesce_block(self):
+        """Try to coalesce phi var with their pre/post variables"""
+
+        ircfg = self.ssa.graph
+
+        # Run coalesce on the post phi parallel copy
+        for irblock in ircfg.blocks.values():
+            if not irblock_has_phi(irblock):
+                continue
+            parallel_copies = {}
+            for dst in self.parallel_copy_block_at_phi[irblock.loc_key]:
+                parallel_copies[dst] = self.phi_new_var[(irblock.loc_key, dst)]
+            self.aggressive_coalesce_parallel_copy(parallel_copies, None)
+
+            # Run coalesce on the pre phi parallel copy
+            parent_to_parallel_copies = {}
+            parallel_copies = {}
+            for dst in irblock[0]:
+                new_var = self.phi_new_var[(irblock.loc_key, dst)]
+                for parent, src in self.parallel_copy_block_at_end[(irblock.loc_key, dst)]:
+                    parent_to_parallel_copies.setdefault(parent, {})[new_var] = src
+                    #self.end_new_var.setdefault(new_var, set()).add(parent)
+
+            for parent, parallel_copies in parent_to_parallel_copies.iteritems():
+                self.aggressive_coalesce_parallel_copy(parallel_copies, parent)
+
+    def get_best_merge_set_name(self, merge_set):
+        """
+        For a given @merge_set, prefer an original SSA variable instead of a
+        created copy. In other case, take a random name.
+        @merge_set: set of equivalent expressions
+        """
+        if not merge_set:
+            raise RuntimeError("Merge set should not be empty")
+        for var in merge_set:
+            if var not in self.copy_vars:
+                return var
+        # Get random name
+        return var
+
+
+    def replace_merge_sets(self):
+        """
+        Replace all variables in a merge state by a representative variable
+        """
+        replace = {}
+        merge_sets = set()
+        merge_set_to_name = {}
+        for merge_set in self.merge_state.itervalues():
+            frozen_merge_set = frozenset(merge_set)
+            merge_sets.add(frozen_merge_set)
+            var_name = self.get_best_merge_set_name(merge_set)
+            merge_set_to_name[frozen_merge_set] = var_name
+
+        for merge_set in merge_sets:
+            var_name = merge_set_to_name[merge_set]
+            merge_set = list(merge_set)
+            for var in merge_set:
+                replace[var] = var_name
+
+        self.ssa.graph.simplify(lambda x: x.replace_expr(replace))
+
+    def remove_phi(self):
+        """
+        Remove phi operators in @ifcfg
+        @ircfg: IRDiGraph instance
+        """
+
+        for irblock in self.ssa.graph.blocks.values():
+            assignblks = list(irblock)
+            out = {}
+            for dst, src in assignblks[0].iteritems():
                 if src.is_op('Phi'):
-                    phi_nodes.append((irblock.loc_key, index))
+                    assert set([dst]) == set(src.args)
+                    continue
+                out[dst] = src
+            assignblks[0] = AssignBlock(out, assignblks[0].instr)
+            self.ssa.graph.blocks[irblock.loc_key] = IRBlock(irblock.loc_key, assignblks)
 
+    def remove_assign_eq(self):
+        """
+        Remove trivial expressions (a=a) in @ifcfg
+        @ircfg: IRDiGraph instance
+        """
 
-    for phi_loc, phi_index in phi_nodes:
-        assignblk_dct = get_assignblk(ssa.graph, phi_loc, phi_index)
-        for dst, src in assignblk_dct.iteritems():
-            if src.is_op('Phi'):
-                break
-        else:
-            raise RuntimeError('Cannot find phi?')
-        node_src = src
-        var = dst
-
-        # Create new variable
-        new_var = ExprId('var%d' % len(phivar2var), var.size)
-        phivar2var[var] = new_var
-        phi_sources = set(node_src.args)
-
-        # Place var init for non ssa variables
-        to_remove = set()
-        for phi_source in list(phi_sources):
-            if phi_source not in all_ssa_vars.union(phivar2var.values()):
-                assignblk_dct = get_assignblk(ssa.graph, head, 0)
-                assignblk_dct[new_var] = phi_source
-                new_irblock = set_assignblk(ssa.graph, head, 0, assignblk_dct)
-                ssa.graph.blocks[head] = new_irblock
-                to_remove.add(phi_source)
-        phi_sources.difference_update(to_remove)
-
-        var_to_replace = set([var])
-        var_to_replace.update(phi_sources)
-
-
-
-
-        # Replace variables
-        to_replace_dct = {x:new_var for x in var_to_replace}
-        for loc in ssa.graph.blocks:
-            irblock = ssa.graph.blocks[loc]
-            assignblks = []
-            for assignblk in irblock:
-                assignblk_dct = {}
+        for irblock in self.ssa.graph.blocks.values():
+            assignblks = list(irblock)
+            for i, assignblk in enumerate(assignblks):
+                out = {}
                 for dst, src in assignblk.iteritems():
-                    dst = dst.replace_expr(to_replace_dct)
-                    src = src.replace_expr(to_replace_dct)
-                    assignblk_dct[dst] = src
-                assignblks.append(AssignBlock(assignblk_dct, assignblk.instr))
-            new_irblock = IRBlock(loc, assignblks)
-            ssa.graph.blocks[loc] = new_irblock
+                    if dst == src:
+                        continue
+                    out[dst] = src
+                assignblks[i] = AssignBlock(out, assignblk.instr)
+            self.ssa.graph.blocks[irblock.loc_key] = IRBlock(irblock.loc_key, assignblks)
+
+
+class DiGraphLivenessSSA(DiGraphLivenessIRA):
+    """
+    DiGraph representing variable liveness is a SSA graph
+    """
+    def __init__(self, ircfg):
+        super(DiGraphLivenessSSA, self).__init__(ircfg)
+
+        self.loc_key_to_phi_parents = {}
+        for irblock in self.blocks.values():
+            if not irblock_has_phi(irblock):
+                continue
+            out = {}
+            for sources in irblock[0].itervalues():
+                var_to_parents = get_phi_sources_parent_block(self, irblock.loc_key, sources.args)
+                for var, var_parents in var_to_parents.iteritems():
+                    out.setdefault(var, set()).update(var_parents)
+            self.loc_key_to_phi_parents[irblock.loc_key] = out
+
+    def back_propagate_to_parent(self, todo, node, parent):
+        parent_block = self.blocks[parent]
+        cur_block = self.blocks[node]
+        irblock = self.ircfg.blocks[node]
+        if cur_block.infos[0].var_in == parent_block.infos[-1].var_out:
+            return
+        var_info = cur_block.infos[0].var_in.union(parent_block.infos[-1].var_out)
+
+        if irblock_has_phi(irblock):
+            # Remove phi special case
+            out = set()
+            phi_sources = self.loc_key_to_phi_parents[irblock.loc_key]
+            for var in var_info:
+                if var not in phi_sources:
+                    out.add(var)
+                    continue
+                if parent in phi_sources[var]:
+                    out.add(var)
+            var_info = out
 
-        # Remove phi
-        assignblk_dct = get_assignblk(ssa.graph, phi_loc, phi_index)
-        del assignblk_dct[new_var]
+        parent_block.infos[-1].var_out = var_info
+        todo.add(parent)
 
 
-        new_irblock = set_assignblk(ssa.graph, phi_loc, phi_index, assignblk_dct)
-        ssa.graph.blocks[phi_loc] = new_irblock