From 7ea28642e4ffcb499b93195c27ba98ac6dc92ca7 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Mon, 18 Aug 2014 23:06:34 +0800
Subject: [PATCH] compiler: add lower_units transform

---
 artiq/compiler/interleave.py  | 16 ++++---------
 artiq/compiler/lower_time.py  | 10 +--------
 artiq/compiler/lower_units.py | 42 +++++++++++++++++++++++++++++++++++
 artiq/devices/core.py         |  2 ++
 4 files changed, 49 insertions(+), 21 deletions(-)
 create mode 100644 artiq/compiler/lower_units.py

diff --git a/artiq/compiler/interleave.py b/artiq/compiler/interleave.py
index d739777dd..07947a2b3 100644
--- a/artiq/compiler/interleave.py
+++ b/artiq/compiler/interleave.py
@@ -1,12 +1,11 @@
 import ast, types
 
-from artiq.language import units
 from artiq.compiler.tools import *
 
 # -1 statement duration could not be pre-determined
 #  0 statement has no effect on timeline
 # >0 statement is a static delay that advances the timeline
-#     by the given amount (in seconds)
+#     by the given amount (in microcycles)
 def _get_duration(stmt):
 	if isinstance(stmt, (ast.Expr, ast.Assign)):
 		return _get_duration(stmt.value)
@@ -21,10 +20,8 @@ def _get_duration(stmt):
 			try:
 				da = eval_constant(stmt.args[0])
 			except NotConstant:
-				return -1
-			if da.unit != units.s_unit:
-				raise units.DimensionError("Delay not expressed in seconds")
-			return da.amount
+				da = -1
+			return da
 		else:
 			return 0
 	else:
@@ -54,14 +51,9 @@ def _interleave_timelines(timelines):
 				stmt.delay -= dt
 				if stmt.delay == 0:
 					ref_stmt = stmt.stmt
-			da_expr = ast.copy_location(
-					ast.Call(func=ast.Name("Quantity", ast.Load()),
-					args=[value_to_ast(dt), ast.Name("s_unit", ast.Load())],
-					keywords=[], starargs=[], kwargs=[]),
-				ref_stmt)
 			delay_stmt = ast.copy_location(
 				ast.Expr(ast.Call(func=ast.Name("delay", ast.Load()),
-					args=[da_expr],
+					args=[value_to_ast(dt)],
 					keywords=[], starargs=[], kwargs=[])),
 				ref_stmt)
 			r.append(delay_stmt)
diff --git a/artiq/compiler/lower_time.py b/artiq/compiler/lower_time.py
index 1dbdae853..78104677e 100644
--- a/artiq/compiler/lower_time.py
+++ b/artiq/compiler/lower_time.py
@@ -10,16 +10,8 @@ def _insert_int64(node):
 		keywords=[], starargs=[], kwargs=[]), node)
 
 class _TimeLowerer(ast.NodeTransformer):
-	def __init__(self, ref_period):
-		self.ref_period = ref_period
-
 	def visit_Call(self, node):
-		if isinstance(node.func, ast.Name) \
-		  and node.func.id == "Quantity" and node.args[1].id == "s_unit":
-			return ast.copy_location(
-				ast.BinOp(left=node.args[0], op=ast.Div(), right=value_to_ast(self.ref_period.amount)),
-				node)
-		elif isinstance(node.func, ast.Name) and node.func.id == "now":
+		if isinstance(node.func, ast.Name) and node.func.id == "now":
 			return ast.copy_location(ast.Name("now", ast.Load()), node)
 		else:
 			self.generic_visit(node)
diff --git a/artiq/compiler/lower_units.py b/artiq/compiler/lower_units.py
new file mode 100644
index 000000000..f151f5534
--- /dev/null
+++ b/artiq/compiler/lower_units.py
@@ -0,0 +1,42 @@
+import ast
+
+from artiq.compiler.tools import value_to_ast
+from artiq.language import units
+
+# TODO:
+#  * track variable and expression dimensions
+#  * raise exception on dimension errors in expressions
+#  * modify RPC map to reintroduce units
+#  * handle core time conversion outside of delay/at,
+#    e.g. foo = now() + 1*us [...] at(foo)
+
+class _UnitsLowerer(ast.NodeTransformer):
+	def __init__(self, ref_period):
+		self.ref_period = ref_period
+		self.in_core_time = False
+
+	def visit_Call(self, node):
+		fn = node.func.id
+		if fn in ("delay", "at"):
+			old_in_core_time = self.in_core_time
+			self.in_core_time = True
+			self.generic_visit(node)
+			self.in_core_time = old_in_core_time
+		elif fn == "Quantity":
+			if self.in_core_time:
+				if node.args[1].id == "microcycle_units":
+					node = node.args[0]
+				else:
+					node = ast.copy_location(
+						ast.BinOp(left=node.args[0], op=ast.Div(), right=value_to_ast(self.ref_period)),
+						node)
+			else:
+				node = node.args[0]
+		else:
+			self.generic_visit(node)
+		return node
+
+def lower_units(funcdef, ref_period):
+	if not isinstance(ref_period, units.Quantity) or ref_period.unit is not units.s_unit:
+		raise units.DimensionError("Reference period not expressed in seconds")
+	_UnitsLowerer(ref_period.amount).visit(funcdef)
diff --git a/artiq/devices/core.py b/artiq/devices/core.py
index 0d7ff3bf6..e829160a4 100644
--- a/artiq/devices/core.py
+++ b/artiq/devices/core.py
@@ -1,4 +1,5 @@
 from artiq.compiler.inline import inline
+from artiq.compiler.lower_units import lower_units
 from artiq.compiler.fold_constants import fold_constants
 from artiq.compiler.unroll_loops import unroll_loops
 from artiq.compiler.interleave import interleave
@@ -14,6 +15,7 @@ class Core:
 
 	def run(self, k_function, k_args, k_kwargs):
 		funcdef, rpc_map = inline(self, k_function, k_args, k_kwargs)
+		lower_units(funcdef, self.runtime_env.ref_period)
 		fold_constants(funcdef)
 		unroll_loops(funcdef, 50)
 		interleave(funcdef)