From fb63698ef470992dc53570b98c2341db098d3dc0 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sebastien@milkymist.org>
Date: Sat, 10 Nov 2012 23:30:14 +0100
Subject: [PATCH] pytholite/compiler: clean up visit_statement

---
 migen/pytholite/compiler.py | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/migen/pytholite/compiler.py b/migen/pytholite/compiler.py
index 5023f93b..5dc001cd 100644
--- a/migen/pytholite/compiler.py
+++ b/migen/pytholite/compiler.py
@@ -92,17 +92,17 @@ class _Compiler:
 	
 	# entry state is first state returned
 	def visit_statement(self, statement):
-		states = []
-		exit_states = []
 		if isinstance(statement, ast.Assign):
 			op = self.visit_assign(statement)
 			if op:
-				states.append(op)
-				exit_states.append(op)
+				return [op], [op]
+			else:
+				return [], []
 		elif isinstance(statement, ast.If):
 			test = self.visit_expr(statement.test)
 			states_t, exit_states_t = self.visit_block(statement.body)
 			states_f, exit_states_f = self.visit_block(statement.orelse)
+			exit_states = exit_states_t + exit_states_f
 			
 			test_state_stmt = If(test, _AbstractNextState(states_t[0]))
 			test_state = [test_state_stmt]
@@ -111,9 +111,7 @@ class _Compiler:
 			else:
 				exit_states.append(test_state)
 			
-			states.append(test_state)
-			states += states_t + states_f
-			exit_states += exit_states_t + exit_states_f
+			return [test_state] + states_t + states_f, exit_states
 		elif isinstance(statement, ast.While):
 			test = self.visit_expr(statement.test)
 			states_b, exit_states_b = self.visit_block(statement.body)
@@ -122,9 +120,7 @@ class _Compiler:
 			for exit_state in exit_states_b:
 				exit_state.insert(0, _AbstractNextState(test_state))
 			
-			exit_states.append(test_state)
-			states += states_b
-			states.append(test_state)
+			return [test_state] + states_b, [test_state]
 		elif isinstance(statement, ast.For):
 			if not isinstance(statement.target, ast.Name):
 				raise NotImplementedError
@@ -132,6 +128,7 @@ class _Compiler:
 			if target in self.symdict:
 				raise NotImplementedError("For loop target must use an available name")
 			it = self.visit_iterator(statement.iter)
+			states = []
 			last_exit_states = []
 			for iteration in it:
 				self.symdict[target] = iteration
@@ -140,17 +137,15 @@ class _Compiler:
 					exit_state.insert(0, _AbstractNextState(states_b[0]))
 				last_exit_states = exit_states_b
 				states += states_b
-			exit_states += last_exit_states
 			del self.symdict[target]
+			return states, last_exit_states
 		elif isinstance(statement, ast.Expr):
 			if isinstance(statement.value, ast.Yield):
 				yvalue = statement.value.value
 				if not isinstance(yvalue, ast.Call) or not isinstance(yvalue.func, ast.Name):
 					raise NotImplementedError("Unrecognized I/O sequence")
 				callee = self.symdict[yvalue.func.id]
-				states_i, exit_states_i = gen_io(self, callee, yvalue.args, [])
-				states += states_i
-				exit_states += exit_states_i
+				return gen_io(self, callee, yvalue.args, [])
 			else:
 				raise NotImplementedError
 		else:
-- 
2.30.2