litex/migen/genlib/record.py

from migen.fhdl.structure import *
from migen.fhdl.tracer import get_obj_var_name
from migen.genlib.misc import optree

(DIR_NONE, DIR_S_TO_M, DIR_M_TO_S) = range(3)

# Possible layout elements:
#   1. (name, size)
#   2. (name, size, direction)
#   3. (name, sublayout)
# size can be an int, or a (int, bool) tuple for signed numbers
# sublayout must be a list

def layout_len(layout):
	r = 0
	for f in layout:
		if isinstance(f[1], (int, tuple)): # cases 1/2
			if(len(f) == 3):
				fname, fsize, fdirection = f
			else:
				fname, fsize = f
		elif isinstance(f[1], list): # case 3
			fname, fsublayout = f
			fsize = layout_len(fsublayout)
		else:
			raise TypeError
		if isinstance(fsize, tuple):
			r += fsize[0]
		else:
			r += fsize
	return r

def layout_get(layout, name):
	for f in layout:
		if f[0] == name:
			return f
	raise KeyError(name)

def layout_partial(layout, *elements):
	r = []
	for path in elements:
		path_s = path.split("/")
		last = path_s.pop()
		copy_ref = layout
		insert_ref = r
		for hop in path_s:
			name, copy_ref = layout_get(copy_ref, hop)
			try:
				name, insert_ref = layout_get(insert_ref, hop)
			except KeyError:
				new_insert_ref = []
				insert_ref.append((hop, new_insert_ref))
				insert_ref = new_insert_ref
		insert_ref.append(layout_get(copy_ref, last))
	return r

class Record:
	def __init__(self, layout, name=None):
		self.name = get_obj_var_name(name, "")
		self.layout = layout

		if self.name:
			prefix = self.name + "_"
		else:
			prefix = ""
		for f in self.layout:
			if isinstance(f[1], (int, tuple)): # cases 1/2
				if(len(f) == 3):
					fname, fsize, fdirection = f
				else:
					fname, fsize = f
				finst = Signal(fsize, name=prefix + fname)
			elif isinstance(f[1], list): # case 3
				fname, fsublayout = f
				finst = Record(fsublayout, prefix + fname)
			else:
				raise TypeError
			setattr(self, fname, finst)

	def eq(self, other):
		return [getattr(self, f[0]).eq(getattr(other, f[0]))
		  for f in self.layout if hasattr(other, f[0])]
	
	def flatten(self):
		r = []
		for f in self.layout:			
			e = getattr(self, f[0])
			if isinstance(e, Signal):
				r.append(e)
			elif isinstance(e, Record):
				r += e.flatten()
			else:
				raise TypeError
		return r

	def raw_bits(self):
		return Cat(*self.flatten())
	
	def connect(self, *slaves, match_by_position=False):
		if match_by_position:
			iters = [iter(slave.layout) for slave in slaves]
		else:
			iters = [iter(self.layout) for slave in slaves]
		r = []
		for f in self.layout:
			field = f[0]
			self_e = getattr(self, field)
			if isinstance(self_e, Signal):
				direction = f[2]
				if direction == DIR_M_TO_S:
					r += [getattr(slave, next(it)[0]).eq(self_e) for slave, it in zip(slaves, iters)]
				elif direction == DIR_S_TO_M:
					r.append(self_e.eq(optree("|", [getattr(slave, next(it)[0])
					  for slave, it in zip(slaves, iters)])))
				else:
					raise TypeError
			else:
				for slave, it in zip(slaves, iters):
					r += self_e.connect(getattr(slave, next(it)[0]),
					  match_by_position=match_by_position)
		return r

	def __len__(self):
		return layout_len(self.layout)

	def __repr__(self):
		return "<Record " + ":".join(f[0] for f in self.layout) + " at " + hex(id(self)) + ">"