From b28fdf5fb087116d111584019444e104e293d37f Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 21 Jun 2014 15:06:15 +0200
Subject: [PATCH] compiler: add unroll_loops transform

---
 artiq/compiler/unroll_loops.py | 46 ++++++++++++++++++++++++++++++++++
 artiq/devices/core.py          |  3 ++-
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 artiq/compiler/unroll_loops.py

diff --git a/artiq/compiler/unroll_loops.py b/artiq/compiler/unroll_loops.py
new file mode 100644
index 000000000..553277eab
--- /dev/null
+++ b/artiq/compiler/unroll_loops.py
@@ -0,0 +1,46 @@
+import ast
+
+from artiq.compiler.tools import eval_ast, make_stmt_transformer
+
+def _count_stmts(node):
+	if isinstance(node, (ast.For, ast.While, ast.If)):
+		return 1 + _count_stmts(node.body) + _count_stmts(node.orelse)
+	elif isinstance(node, ast.With):
+		return 1 + _count_stmts(node.body)
+	elif isinstance(node, list):
+		return sum(map(_count_stmts, node))
+	else:
+		return 1
+
+class _LoopUnroller(ast.NodeTransformer):
+	def __init__(self, limit):
+		self.limit = limit
+
+	def visit_For(self, node):
+		self.generic_visit(node)
+		try:
+			it = eval_ast(node.iter)
+		except:
+			return node
+		l_it = len(it)
+		if l_it:
+			n = l_it*_count_stmts(node.body)
+			if n < self.limit:
+				replacement = []
+				for i in it:
+					if not isinstance(i, int):
+						replacement = None
+						break
+					replacement.append(ast.copy_location(
+						ast.Assign(targets=[node.target], value=ast.Num(n=i)), node))
+					replacement += node.body
+				if replacement is not None:
+					return replacement
+				else:
+					return node
+			else:
+				return node
+		else:
+			return node.orelse
+
+unroll_loops = make_stmt_transformer(_LoopUnroller)
diff --git a/artiq/devices/core.py b/artiq/devices/core.py
index ca225520a..3e924adc0 100644
--- a/artiq/devices/core.py
+++ b/artiq/devices/core.py
@@ -2,13 +2,14 @@ from operator import itemgetter
 
 from artiq.compiler.inline import inline
 from artiq.compiler.fold_constants import fold_constants
+from artiq.compiler.unroll_loops import unroll_loops
 from artiq.compiler.unparse import Unparser
 
 class Core:
 	def run(self, k_function, k_args, k_kwargs):
 		stmts, rpc_map = inline(self, k_function, k_args, k_kwargs)
 		fold_constants(stmts)
-
+		unroll_loops(stmts, 50)
 
 		print("=========================")
 		print(" Inlined")