litex/migen/fhdl/namer.py

from collections import OrderedDict
from itertools import combinations

from migen.fhdl.structure import *

class _Node:
	def __init__(self):
		self.signal_count = 0
		self.numbers = set()
		self.use_name = False
		self.use_number = False
		self.children = OrderedDict()

def _display_tree(filename, tree):
	from migen.graph.treeviz import RenderNode
	
	def _to_render_node(name, node):
		children = [_to_render_node(k, v) for k, v in node.children.items()]
		if node.use_name:
			if node.use_number:
				color = (0.5, 0.9, 0.8)
			else:
				color = (0.8, 0.5, 0.9)
		else:
			if node.use_number:
				color = (0.9, 0.8, 0.5)
			else:
				color = (0.8, 0.8, 0.8)
		label = "{0}\n{1} signals\n{2}".format(name, node.signal_count, node.numbers)
		return RenderNode(label, children, color=color)

	top = _to_render_node("top", tree)
	top.to_svg(filename)

def _build_tree(signals, basic_tree=None):
	root = _Node()
	for signal in signals:
		current_b = basic_tree
		current = root
		current.signal_count += 1
		for name, number in signal.backtrace:
			if basic_tree is None:
				use_number = False
			else:
				current_b = current_b.children[name]
				use_number = current_b.use_number
			if use_number:
				key = (name, number)
			else:
				key = name
			try:
				current = current.children[key]
			except KeyError:
				new = _Node()
				current.children[key] = new
				current = new
			current.numbers.add(number)
			if use_number:
				current.all_numbers = sorted(current_b.numbers)
			current.signal_count += 1
	return root

def _set_use_name(node, node_name=""):
	if not node.children:
		node.use_name = True
		return {(node_name, )}
	else:
		cnames = [(k, _set_use_name(v, k)) for k, v in node.children.items()]
		for (c1_prefix, c1_names), (c2_prefix, c2_names) in combinations(cnames, 2):
			if not c1_names.isdisjoint(c2_names):
				node.children[c1_prefix].use_name = True
				node.children[c2_prefix].use_name = True
		r = set()
		for c_prefix, c_names in cnames:
			if node.children[c_prefix].use_name:
				for c_name in c_names:
					r.add((c_prefix, ) + c_name)
			else:
				r |= c_names
		return r

def _name_signal(tree, signal):
	elements = []
	treepos = tree
	for step_name, step_n in signal.backtrace:
		try:
			treepos = treepos.children[(step_name, step_n)]
			use_number = True
		except KeyError:
			treepos = treepos.children[step_name]
			use_number = False
		if treepos.use_name:
			elname = step_name
			if use_number:
				elname += str(treepos.all_numbers.index(step_n))
			elements.append(elname)
	return "_".join(elements)

def _build_pnd(tree, signals):
	return dict((signal, _name_signal(tree, signal)) for signal in signals)

def _invert_pnd(pnd):
	inv_pnd = dict()
	for k, v in pnd.items():
		inv_pnd[v] = inv_pnd.get(v, [])
		inv_pnd[v].append(k)
	return inv_pnd

def _list_conflicting_signals(pnd):
	inv_pnd = _invert_pnd(pnd)
	r = set()
	for k, v in inv_pnd.items():
		if len(v) > 1:
			r.update(v)
	return r

def _set_use_number(tree, signals):
	for signal in signals:
		current = tree
		for step_name, step_n in signal.backtrace:
			current = current.children[step_name]
			current.use_number = current.signal_count > len(current.numbers) and len(current.numbers) > 1

_debug = False

def build_namespace(signals):
	basic_tree = _build_tree(signals)
	_set_use_name(basic_tree)
	if _debug:
		_display_tree("tree_basic.svg", basic_tree)
	pnd = _build_pnd(basic_tree, signals)

	# If there are conflicts, try splitting the tree by numbers
	# on paths taken by conflicting signals.
	conflicting_signals = _list_conflicting_signals(pnd)
	if conflicting_signals:
		_set_use_number(basic_tree, conflicting_signals)
		if _debug:
			print("namer: using split-by-number strategy")
			_display_tree("tree_marked.svg", basic_tree)
		numbered_tree = _build_tree(signals, basic_tree)
		_set_use_name(numbered_tree)
		if _debug:
			_display_tree("tree_numbered.svg", numbered_tree)
		pnd = _build_pnd(numbered_tree, signals)
	else:
		if _debug:
			print("namer: using basic strategy")
	
	# ...then add number suffixes by HUID
	inv_pnd = _invert_pnd(pnd)
	huid_suffixed = False
	for name, signals in inv_pnd.items():
		if len(signals) > 1:
			huid_suffixed = True
			for n, signal in enumerate(sorted(signals, key=lambda x: x.huid)):
				pnd[signal] += str(n)
	if _debug and huid_suffixed:
		print("namer: using HUID suffixes")

	ns = Namespace(pnd)
	# register signals with name_override
	for signal in signals:
		if signal.name_override is not None:
			ns.get_name(signal)
	return ns

class Namespace:
	def __init__(self, pnd):
		self.counts = {}
		self.sigs = {}
		self.pnd = pnd
	
	def get_name(self, sig):
		if sig.name_override is not None:
			sig_name = sig.name_override
		else:
			sig_name = self.pnd[sig]
		try:
			n = self.sigs[sig]
		except KeyError:
			try:
				n = self.counts[sig_name]
			except KeyError:
				n = 0
			self.sigs[sig] = n
			self.counts[sig_name] = n + 1
		if n:
			return sig_name + "_" + str(n)
		else:
			return sig_name