pytholite: visit AST and list registers

migen/pytholite/compiler.py

@@ -0,0 +1,88 @@
+import inspect
+import ast
+
+from migen.fhdl.structure import *
+from migen.pytholite import transel
+
+class FinalizeError(Exception):
+	pass
+
+class _AbstractLoad:
+	def __init__(self, target, source):
+		self.target = target
+		self.source = source
+
+class _Register:
+	def __init__(self, name, nbits):
+		self.storage = Signal(BV(nbits), name=name)
+		self.source_encoding = {}
+		self.finalized = False
+	
+	def load(self, source):
+		if source not in self.source_encoding:
+			self.source_encoding[source] = len(self.source_encoding) + 1
+		return _AbstractLoad(self, source)
+	
+	def finalize(self):
+		if self.finalized:
+			raise FinalizeError
+		self.sel = Signal(BV(bits_for(len(self.source_encoding) + 1)))
+		self.finalized = True
+	
+	def get_fragment(self):
+		if not self.finalized:
+			raise FinalizeError
+		# do nothing when sel == 0
+		cases = [(v, self.storage.eq(k)) for k, v in self.source_encoding.items()]
+		sync = [Case(self.sel, *cases)]
+		return Fragment(sync=sync)
+
+class _AnonymousRegister:
+	def __init__(self, nbits):
+		self.nbits = nbits
+
+class _CompileVisitor(ast.NodeVisitor):
+	def __init__(self, symdict, registers):
+		self.symdict = symdict
+		self.registers = registers
+	
+	def visit_Assign(self, node):
+		value = self.visit(node.value)
+		if isinstance(value, _AnonymousRegister):
+			if isinstance(node.targets[0], ast.Name):
+				name = node.targets[0].id
+			else:
+				raise NotImplementedError
+			value = _Register(name, value.nbits)
+			self.registers.append(value)
+			for target in node.targets:
+				if isinstance(target, ast.Name):
+					self.symdict[target.id] = value
+				else:
+					raise NotImplementedError
+	
+	def visit_Call(self, node):
+		if isinstance(node.func, ast.Name):
+			callee = self.symdict[node.func.id]
+		else:
+			raise NotImplementedError
+		if callee == transel.Register:
+			if len(node.args) != 1:
+				raise TypeError("Register() takes exactly 1 argument")
+			nbits = ast.literal_eval(node.args[0])
+			return _AnonymousRegister(nbits)
+		else:
+			raise NotImplementedError
+	
+def make_pytholite(func):
+	tree = ast.parse(inspect.getsource(func))
+	symdict = func.__globals__.copy()
+	registers = []
+	
+	cv = _CompileVisitor(symdict, registers)
+	cv.visit(tree)
+	
+	print(registers)
+	print(symdict)
+
+	#print(ast.dump(tree))