From b163c9f7ea5ef9213b5378a2162ae186c1ad7b70 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Mon, 3 Nov 2014 18:44:30 +0800
Subject: [PATCH] test: add optimization transform stack

---
 artiq/coredevice/core.py | 48 +++++++++++++++++++++++-----------------
 test/transforms.py       | 45 +++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 20 deletions(-)
 create mode 100644 test/transforms.py

diff --git a/artiq/coredevice/core.py b/artiq/coredevice/core.py
index 022b3d968..5d4ae8725 100644
--- a/artiq/coredevice/core.py
+++ b/artiq/coredevice/core.py
@@ -35,6 +35,10 @@ def _make_debug_unparse(final):
         return _filtered_unparse
 
 
+def _no_debug_unparse(label, node):
+    pass
+
+
 class Core:
     def __init__(self, core_com, runtime_env=None):
         if runtime_env is None:
@@ -42,51 +46,55 @@ class Core:
         self.runtime_env = runtime_env
         self.core_com = core_com
 
-    def run(self, k_function, k_args, k_kwargs):
-        # transform/simplify AST
-        _debug_unparse = _make_debug_unparse("remove_dead_code_2")
-
-        func_def, rpc_map, exception_map = inline(
-            self, k_function, k_args, k_kwargs)
-        _debug_unparse("inline", func_def)
-
+    def transform_stack(self, func_def, rpc_map, exception_map,
+                        debug_unparse=_no_debug_unparse):
         lower_units(func_def, rpc_map)
-        _debug_unparse("lower_units", func_def)
+        debug_unparse("lower_units", func_def)
 
         remove_inter_assigns(func_def)
-        _debug_unparse("remove_inter_assigns_1", func_def)
+        debug_unparse("remove_inter_assigns_1", func_def)
 
         fold_constants(func_def)
-        _debug_unparse("fold_constants_1", func_def)
+        debug_unparse("fold_constants_1", func_def)
 
         unroll_loops(func_def, 500)
-        _debug_unparse("unroll_loops", func_def)
+        debug_unparse("unroll_loops", func_def)
 
         interleave(func_def)
-        _debug_unparse("interleave", func_def)
+        debug_unparse("interleave", func_def)
 
         lower_time(func_def,
                    getattr(self.runtime_env, "initial_time", 0),
                    self.runtime_env.ref_period)
-        _debug_unparse("lower_time", func_def)
+        debug_unparse("lower_time", func_def)
 
         remove_inter_assigns(func_def)
-        _debug_unparse("remove_inter_assigns_2", func_def)
+        debug_unparse("remove_inter_assigns_2", func_def)
 
         fold_constants(func_def)
-        _debug_unparse("fold_constants_2", func_def)
+        debug_unparse("fold_constants_2", func_def)
 
         remove_dead_code(func_def)
-        _debug_unparse("remove_dead_code_1", func_def)
+        debug_unparse("remove_dead_code_1", func_def)
 
         remove_inter_assigns(func_def)
-        _debug_unparse("remove_inter_assigns_3", func_def)
+        debug_unparse("remove_inter_assigns_3", func_def)
 
         fold_constants(func_def)
-        _debug_unparse("fold_constants_3", func_def)
+        debug_unparse("fold_constants_3", func_def)
 
         remove_dead_code(func_def)
-        _debug_unparse("remove_dead_code_2", func_def)
+        debug_unparse("remove_dead_code_2", func_def)
+
+    def run(self, k_function, k_args, k_kwargs):
+        # transform/simplify AST
+        debug_unparse = _make_debug_unparse("remove_dead_code_2")
+
+        func_def, rpc_map, exception_map = inline(
+            self, k_function, k_args, k_kwargs)
+        debug_unparse("inline", func_def)
+
+        self.transform_stack(func_def, rpc_map, exception_map, debug_unparse)
 
         # compile to machine code and run
         binary = get_runtime_binary(self.runtime_env, func_def)
diff --git a/test/transforms.py b/test/transforms.py
new file mode 100644
index 000000000..ea89a7322
--- /dev/null
+++ b/test/transforms.py
@@ -0,0 +1,45 @@
+import unittest
+import ast
+
+from artiq.coredevice import comm_dummy, core
+from artiq.transforms.unparse import unparse
+
+
+# Original code before inline:
+#
+# n = time_to_cycles(1.2345*ns)
+# ftw = self.dds.frequency_to_ftw(345*MHz)
+# f = self.dds.ftw_to_frequency(ftw)
+# phi = 1000*cycles_to_time(n)*f
+# do_someting(int(phi))
+#
+optimize_in = """
+
+def run():
+    dds_sysclk = Quantity(Fraction(1000000000, 1), 'Hz')
+    n = time_to_cycles((1.2345 * Quantity(Fraction(1, 1000000000), 's')))
+    with sequential:
+        frequency = (345 * Quantity(Fraction(1000000, 1), 'Hz'))
+        frequency_to_ftw_return = int((((2 ** 32) * frequency) / dds_sysclk))
+    ftw = frequency_to_ftw_return
+    with sequential:
+        ftw2 = ftw
+        ftw_to_frequency_return = ((ftw2 * dds_sysclk) / (2 ** 32))
+    f = ftw_to_frequency_return
+    phi = ((1000 * cycles_to_time(n)) * f)
+    do_something(int(phi))
+"""
+
+optimize_out = """
+
+def run():
+    do_something(344)
+"""
+
+
+class OptimizeCase(unittest.TestCase):
+    def test_optimize(self):
+        coredev = core.Core(comm_dummy.Comm())
+        func_def = ast.parse(optimize_in).body[0]
+        coredev.transform_stack(func_def, dict(), dict())
+        self.assertEqual(unparse(func_def), optimize_out)