From 603d49dffa99dfa1e86fd5bc03894fb885b8373d Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 18 Jul 2015 20:48:52 +0300
Subject: [PATCH] Add a dominator analysis.

---
 artiq/compiler/analyses/__init__.py   |  1 +
 artiq/compiler/analyses/domination.py | 48 +++++++++++++++++++++++++++
 artiq/compiler/ir.py                  | 10 ++++++
 3 files changed, 59 insertions(+)
 create mode 100644 artiq/compiler/analyses/__init__.py
 create mode 100644 artiq/compiler/analyses/domination.py

diff --git a/artiq/compiler/analyses/__init__.py b/artiq/compiler/analyses/__init__.py
new file mode 100644
index 000000000..f7f69f651
--- /dev/null
+++ b/artiq/compiler/analyses/__init__.py
@@ -0,0 +1 @@
+from .domination import DominatorTree
diff --git a/artiq/compiler/analyses/domination.py b/artiq/compiler/analyses/domination.py
new file mode 100644
index 000000000..8963cf189
--- /dev/null
+++ b/artiq/compiler/analyses/domination.py
@@ -0,0 +1,48 @@
+"""
+:class:`DominatorTree` computes the dominance relation over
+control flow graphs.
+
+See http://www.cs.colostate.edu/~mstrout/CS553/slides/lecture04.pdf.
+"""
+
+from functools import reduce, cmp_to_key
+
+# Key Idea
+#   If a node dominates all
+#   predecessors of node n, then it
+#   also dominates node n
+class DominatorTree:
+    def __init__(self, func):
+        entry = func.get_entry()
+
+        self.dominated_by = { entry: {entry} }
+        for block in func.basic_blocks:
+            if block != entry:
+                self.dominated_by[block] = set(func.basic_blocks)
+
+        predecessors = {block: block.predecessors() for block in func.basic_blocks}
+        while True:
+            changed = False
+
+            for block in func.basic_blocks:
+                if block == entry:
+                    continue
+
+                if not any(predecessors[block]):
+                    # Unreachable blocks are dominated by everything
+                    continue
+
+                new_dominated_by = {block}.union(
+                    reduce(lambda a, b: a.intersection(b),
+                           (self.dominated_by[pred] for pred in predecessors[block])))
+                if new_dominated_by != self.dominated_by[block]:
+                    self.dominated_by[block] = new_dominated_by
+                    changed = True
+
+            if not changed:
+                break
+
+    def in_domination_order(self):
+        blocks = list(self.dominated_by.keys())
+        blocks.sort(key=cmp_to_key(lambda a, b: a in self.dominated_by[b]))
+        return blocks
diff --git a/artiq/compiler/ir.py b/artiq/compiler/ir.py
index 3b77db9b7..ad36fbbeb 100644
--- a/artiq/compiler/ir.py
+++ b/artiq/compiler/ir.py
@@ -335,6 +335,9 @@ class BasicBlock(NamedValue):
 
         return "\n".join(lines)
 
+    def __repr__(self):
+        return "<BasicBlock '{}'>".format(self.name)
+
 class Argument(NamedValue):
     """
     A function argument.
@@ -383,6 +386,13 @@ class Function:
         basic_block._detach()
         self.basic_blocks.remove(basic_block)
 
+    def get_entry(self):
+        assert any(self.basic_blocks)
+        return self.basic_blocks[0]
+
+    def instructions(self):
+        for basic_block in self.basic_blocks:
+            yield from iter(basic_block.instructions)
 
     def __str__(self):
         printer = types.TypePrinter()