steps_rsn.py

from io import StringIO
from sys import stdout
from collections import defaultdict
from itertools import groupby
from typing import *

from py_verilog_parser.verilog import *
from steps_rsn_elements import *


class RsnModule(object):

	def __init__(self, module: str):
		assert isinstance(module, str)
		self.module: str = module
		self.elements: List[RsnElement] = []
		self.connections: Dict[Tuple[RsnElement, str], List[Tuple[RsnElement, str]]] = defaultdict(list)
		self.reverse_connections: Dict[Tuple[RsnElement, str], List[Tuple[RsnElement, str]]] = defaultdict(list)

	def _add_element(self, element: RsnElement) -> RsnElement:
		self.elements.append(element)
		return element

	def add_host_interface(self, instance_name: str, impelementation: str, target: Union[Module, ModuleInstance]) -> RsnHostInterface:
		element = self._add_element(RsnHostInterface(instance_name, impelementation, target))
		self.add_connection(element, 'host_scan_out', element, 'host_scan_in')
		return element

	def add_client_interface(self, instance_name: str, impelementation: str, target: Union[Module, ModuleInstance]) -> RsnClientInterface:
		return self._add_element(RsnClientInterface(instance_name, impelementation, target))

	def add_sib(self, instance_name: str, implementation: str) -> RsnSib:
		element = self._add_element(RsnSib(instance_name, implementation))
		self.add_connection(element, 'host_scan_out', element, 'host_scan_in')
		return element

	def add_data_register(self, instance_name: str, implementation: str) -> RsnDataRegister:
		return self._add_element(RsnDataRegister(instance_name, implementation))

	def add_internal_scan_cell(self, instance_name: str, implementation: str, target: ModuleInstance, variant: str = 'full') -> RsnInternalScanCell:
		return self._add_element(RsnInternalScanCell(instance_name, implementation, target, variant))

	def add_boundary_scan_cell(self, instance_name: str, implementation: str, target: IdentifierIndexed, direction: str = 'input') -> RsnBoundaryCell:
		return self._add_element(RsnBoundaryCell(instance_name, implementation, target, direction))

	def get_host_interfaces(self) -> List[RsnHostInterface]:
		return filter(lambda elem: isinstance(elem, RsnHostInterface), self.elements)

	def get_client_interfaces(self) -> List[RsnClientInterface]:
		return filter(lambda elem: isinstance(elem, RsnClientInterface), self.elements)

	def get_sibs(self) -> List[RsnSib]:
		return filter(lambda elem: isinstance(elem, RsnSib), self.elements)

	def get_data_registers(self) -> List[RsnDataRegister]:
		return filter(lambda elem: isinstance(elem, RsnDataRegister), self.elements)

	def get_internal_scan_cells(self) -> List[RsnInternalScanCell]:
		return filter(lambda elem: isinstance(elem, RsnInternalScanCell), self.elements)

	def get_boundary_scan_cells(self) -> List[RsnBoundaryCell]:
		return filter(lambda elem: isinstance(elem, RsnBoundaryCell), self.elements)

	def add_connection(self, elementA: RsnElement, portA: str, elementB: RsnElement, portB) -> None:
		assert elementA is None or portA in elementA.get_rsn_port_names(), f'{portA} was not found'
		assert elementB is None or portB in elementB.get_rsn_port_names(), f'{portB} was not found'

		#print(f'Add connection from {"module" if elementA is None else elementA.instance_name}:{portA} to {"module" if elementB is None else elementB.instance_name}:{portB}')
		self.connections[(elementA, portA)] += [(elementB, portB)]
		self.reverse_connections[(elementB, portB)] += [(elementA, portA)]

	def remove_connection(self, elementA: RsnElement, portA: str, elementB: RsnElement, portB) -> None:
		assert elementA is None or portA in elementA.get_rsn_port_names(), f'{portA} was not found'
		assert elementB is None or portB in elementB.get_rsn_port_names(), f'{portB} was not found'

		#print(f'Remove connection from {"module" if elementA is None else elementA.instance_name}:{portA} to {"module" if elementB is None else elementB.instance_name}:{portB}')
		self.connections[(elementA, portA)] = list(filter(
			lambda connection: (connection[0] != elementB) or (connection[1] != portB),
			self.connections[(elementA, portA)]
		))
		self.reverse_connections[(elementB, portB)] = list(filter(
			lambda connection: (connection[0] != elementA) or (connection[1] != portA),
			self.connections[(elementB, portB)]
		))

	def get_connections(self) -> Dict[Tuple[RsnElement, str], List[Tuple[RsnElement, str]]]:
		return self.connections

	def get_drivers(self) -> Dict[Tuple[RsnElement, str], List[Tuple[RsnElement, str]]]:
		return self.reverse_connections

	def get_port_connections(self, elementA: RsnElement, portA: str) -> List[Tuple[RsnElement, str]]:
		return self.connections[(elementA, portA)]

	def get_port_drivers(self, elementA: RsnElement, portA: str) -> List[Tuple[RsnElement, str]]:
		return self.reverse_connections[(elementA, portA)]

	def _chain_rsn_interfaces(self, elementA: RsnElement, interfaceA: str, elementB: RsnElement, interfaceB: str):
		assert interfaceB == 'client'

		if interfaceA == 'host':
			for inElem, inPort in self.get_port_connections(elementA, f'{interfaceA}_scan_out'):
				self.remove_connection(elementA, f'{interfaceA}_scan_out', inElem, inPort)

			self.add_connection(elementA, f'{interfaceA}_scan_in', elementB, f'{interfaceB}_scan_in')
			self.add_connection(elementB, f'{interfaceB}_scan_out', elementA, f'{interfaceA}_scan_out')

			self.add_connection(elementA, f'{interfaceA}_select', elementB, f'{interfaceB}_select')
			self.add_connection(elementA, f'{interfaceA}_reset', elementB, f'{interfaceB}_reset')
			self.add_connection(elementA, f'{interfaceA}_capture_en', elementB, f'{interfaceB}_capture_en')
			self.add_connection(elementA, f'{interfaceA}_shift_en', elementB, f'{interfaceB}_shift_en')
			self.add_connection(elementA, f'{interfaceA}_update_en', elementB, f'{interfaceB}_update_en')
			self.add_connection(elementA, f'{interfaceA}_mode', elementB, f'{interfaceB}_mode')
			self.add_connection(elementA, f'{interfaceA}_clock', elementB, f'{interfaceB}_clock')

		elif interfaceA == 'client':
			for outElem, outPort in self.get_port_connections(elementA, f'{interfaceA}_scan_out'):
				self.remove_connection(elementA, f'{interfaceA}_scan_out', outElem, outPort)
				self.add_connection(elementB, f'{interfaceB}_scan_out', outElem, outPort)
			self.add_connection(elementA, f'{interfaceA}_scan_out', elementB, f'{interfaceB}_scan_in')

			def _copy_drivers(elementA: RsnElement, portA: str, elementB: RsnElement, portB: str):
				for driverElement, driverPort in self.get_port_drivers(elementA, portA):
					self.add_connection(driverElement, driverPort, elementB, portB)

			_copy_drivers(elementA, f'{interfaceA}_select', elementB, f'{interfaceB}_select')
			_copy_drivers(elementA, f'{interfaceA}_reset', elementB, f'{interfaceB}_reset')
			_copy_drivers(elementA, f'{interfaceA}_capture_en', elementB, f'{interfaceB}_capture_en')
			_copy_drivers(elementA, f'{interfaceA}_shift_en', elementB, f'{interfaceB}_shift_en')
			_copy_drivers(elementA, f'{interfaceA}_update_en', elementB, f'{interfaceB}_update_en')
			_copy_drivers(elementA, f'{interfaceA}_mode', elementB, f'{interfaceB}_mode')
			_copy_drivers(elementA, f'{interfaceA}_clock', elementB, f'{interfaceB}_clock')

	def link_rsn_interfaces(self, interfaceA: RsnInterface, interfaceB: RsnInterface) -> None:
		self._chain_rsn_interfaces(interfaceA.element, interfaceA.interface, interfaceB.element, interfaceB.interface)

	def link_rsn_elements(self, elementA: RsnElement, elementB: RsnElement, interface: str=None) -> None:
		if interface is None:
			interface = 'host' if isinstance(elementA, RsnHostInterface) else 'client'
		self._chain_rsn_interfaces(elementA, interface, elementB, 'client')

	def get_dot_graph(self, output: StringIO, mode='simple') -> None:
		if mode == 'simple':
			client_ports = ['client_scan_in', 'client_scan_out']
			host_ports = ['host_scan_in', 'host_scan_out']
		elif mode == 'full':
			client_ports = ['client_scan_in', 'client_scan_out', 'client_select', 'client_reset', 'client_capture_en', 'client_shift_en', 'client_update_en', 'client_clock']
			host_ports = ['host_scan_in', 'host_scan_out', 'host_select', 'host_reset', 'host_capture_en', 'host_shift_en', 'host_update_en', 'host_clock']

		print(f'digraph {self.module} {{', file=output)
		for element in self.elements:
			client_interface = '|'.join([f'<{name}> {element.get_rsn_port_name(name)}' for name in client_ports if name in element.get_rsn_port_names()])
			host_interface = '|'.join([f'<{name}> {element.get_rsn_port_name(name)}' for name in host_ports if name in element.get_rsn_port_names()])
			label = element.instance_name
			if client_interface:
				label = f'{{{client_interface}}}|{label}'
			if host_interface:
				label = f'{label}|{{{host_interface}}}'

			print(f'    {element.instance_name} [shape=record label="{label}"];', file=output)

		for source_element in self.elements:
			for source_port in client_ports + host_ports:
				source_name = source_element.instance_name
				colour = 'crimson' if 'scan' in source_port else 'black'

				connections = self.get_port_connections(source_element, source_port)
				for target_element, target_port in connections:
					target_name = target_element.instance_name if target_element is not None else target_port
					print(f'    edge [color={colour}]; {source_name}:{source_port} -> {target_name}:{target_port};', file=output)

		print(f'}}', file=output)

	def apply_rsn(self, module: Module) -> None:
		self._rename_boundary_wires(module)
		self._replace_flipflops(module)
		self._insert_and_connect_elements(module)

	def _rename_boundary_wires(self, module: Module) -> None:
		# First iterate over all boundary elements and find the indices of the port wires to replace
		replace_targets: Dict[str, Dict[int, Expression]] = defaultdict(dict)
		for element in self.elements:
			if not isinstance(element, RsnBoundaryCell):
				continue

			assert element.target is not None
			target_name = element.target.name
			target_index = element.target.index.as_integer()
			replace_targets[target_name][target_index] = None

		# Now create the wires that will replace the original ports
		boundary_declarations: List[NetDeclaration] = []
		target_replacements: Dict[str, Union[Identifier, Concatenation]] = {}
		for target_name, replace_indices in replace_targets.items():
			target_port = module.get_port(target_name)
			target_indices = [0] if target_port.range is None else target_port.range.to_indices()

			if len(target_indices) == len(replace_indices.keys()):
				boundary_target = f'{target_name}_boundary'
				boundary_declarations += [NetDeclaration([Identifier(boundary_target)], target_port.range)]
				target_replacements[target_name] = Identifier(boundary_target)
				for index in target_indices:
					replace_indices[index] = IdentifierIndexed(boundary_target, Number(None, None, f'{index}'))
			else:
				concatenation = Concatenation()
				for index in target_indices:
					if index not in replace_indices:
						element = IdentifierIndexed(target_name, Number(None, None, f'{index}'))
					else:
						boundary_target = f'{target_name}_boundary_{index}'
						boundary_declarations += [NetDeclaration([boundary_target], None)]
						element = Identifier(boundary_target)
					concatenation.elements += [element]
					replace_indices[index] = element
				target_replacements[target_name] = concatenation

		module.module_items += boundary_declarations
		module.net_declarations += boundary_declarations
		module.update_net_index()

		# Now replace the wires with the boundary ones
		for assignment in module.assignments:
			assignment.left = module.replace_targets_in_expression(assignment.left, target_replacements)
			assignment.right = module.replace_targets_in_expression(assignment.right, target_replacements)
		module.update_assignment_index()

		for instance in module.module_instances:
			for port_name in instance.ports.keys():
				instance.ports[port_name] = module.replace_targets_in_expression(instance.ports[port_name], target_replacements)

		# Assing the wires to the cells
		for element in self.elements:
			if not isinstance(element, RsnBoundaryCell):
				continue

			assert element.target is not None
			target_name = element.target.name
			target_index = element.target.index.as_integer()
			if element.direction == 'input':
				self.add_connection(None, element.target, element, 'data_in')
				self.add_connection(element, 'data_out', None, replace_targets[target_name][target_index])
			elif element.direction == 'output':
				self.add_connection(None, replace_targets[target_name][target_index], element, 'data_in')
				self.add_connection(element, 'data_out', None, element.target)

	def _replace_flipflops(self, module: Module) -> None:
		flipflop_instances = []
		for element in self.elements:
			if not isinstance(element, RsnInternalScanCell):
				continue

			assert element.target is not None
			target = element.target
			flipflop_instances += [target.instance_name]

			connections = {}
			for port, port_name in element.get_target_port_names().items():
				if port_name is None:
					continue

				connections[port] = target.ports.get(Identifier(port_name), None)

			for port, port_name in element.get_rsn_port_names().items():
				connection = connections.get(port, None)
				if connection is None:
					continue

				self.add_connection(element, port, None, connection)

		module.module_instances = [instance for instance in module.module_instances if instance.instance_name not in flipflop_instances]
		module.update_instance_index()

	def _insert_and_connect_elements(self, module: Module) -> None:
		def _is_module(element: RsnElement):
			return isinstance(element, RsnInternalScanCell) \
				or isinstance(element, RsnBoundaryCell) \
				or isinstance(element, RsnDataRegister) \
				or isinstance(element, RsnSib)

		def _is_port(element: RsnElement):
			return isinstance(element, RsnHostInterface) \
				or isinstance(element, RsnClientInterface)

		# Create all instances and insert ports
		instances: Dict[RsnElement, ModuleInstance] = {}
		for element in self.elements:
			if _is_module(element):
				instance = ModuleInstance(
					Identifier(element.implementation),
					Identifier(element.instance_name),
					{}
				)
				module.module_items += [instance]
				module.module_instances += [instance]
			elif _is_port(element):
				instance = element.target
			instances[element] = instance

			port_identifiers = []
			port_ranges = []
			port_directions = []
			for port in element.get_rsn_port_names().keys():
				port_name = element.get_rsn_port_name(port)
				port_size = element.get_rsn_port_size(port)
				port_direction = element.get_rsn_port_direction(port)
				if port_name is None:
					continue

				port_identifier = Identifier(port_name)

				# Add ports for instances
				if isinstance(instance, ModuleInstance):
					#instance.ports[port_identifier] = Concatenation([])
					pass
				elif isinstance(instance, Module):
					port_ranges += [Range(Number(None, None, f'{port_size-1}'), Number(None, None, '0')) if port_size != 1 else None]
					port_directions += [port_direction]
					port_identifiers += [port_identifier]
					instance.port_list += [port_identifier]

			if isinstance(instance, Module):
				all_ports = list(zip(port_identifiers, port_directions, port_ranges))

				# Inputs are outputs and outputs are inputs when looking at the interface from the module inside.
				input_ports = list([(port_identifier, port_range) for port_identifier, port_direction, port_range in all_ports if port_direction == 'output'])
				output_ports = list([(port_identifier, port_range) for port_identifier, port_direction, port_range in all_ports if port_direction == 'input'])

				# input_ports.sort(key=lambda port: port[1])
				# output_ports.sort(key=lambda port: port[1])

				for port_range, grouper in groupby(input_ports, key=lambda port: port[1]):
					input_port = InputDeclaration([port_identifier for port_identifier, _ in grouper], port_range)
					instance.module_items += [input_port]
					instance.input_declarations += [input_port]

				for port_range, grouper in groupby(output_ports, key=lambda port: port[1]):
					output_port = OutputDeclaration([port_identifier for port_identifier, _ in grouper], port_range)
					instance.module_items += [output_port]
					instance.output_declarations += [output_port]

				instance.update_port_index()

		def _get_connections(element, port):
			sources = set()
			sinks = set()
			active_elements = set([(element, port)])

			while active_elements:
				element, port = active_elements.pop()
				new_sources = self.get_port_drivers(element, port)
				new_sinks = self.get_port_connections(element, port)

				active_elements.update([elem for elem in new_sources if elem not in sources and elem not in sinks])
				active_elements.update([elem for elem in new_sinks if elem not in sources and elem not in sinks])

				sinks.update(new_sinks)
				sources.update(new_sources)

			return sources, sinks

		# Connect the ports
		connected_ports: Set[Tuple[RsnElement,str]] = set()
		for element in self.elements:
			for port in element.get_rsn_port_names().keys():
				if (element, port) in connected_ports:
					continue

				if element.get_rsn_port_name(port) is None:
					continue

				sinks, sources = _get_connections(element, port)

				if isinstance(instances[element], Module):
					port_direction = element.get_rsn_port_direction(port)
					if port_direction == 'input':
						sources.add((None, Identifier(element.get_rsn_port_name(port))))
					elif port_direction == 'output':
						sinks.add((None, Identifier(element.get_rsn_port_name(port))))

				if len(sinks) == 0 and len(sources) == 0:
					continue

				sink_wires = [connection for element, connection in sinks if element is None]
				source_wires = [connection for element, connection in sources if element is None]

				assert len(sink_wires) + len(source_wires) <= 1, "Resolving multipe wires with assingments is not supported yet"

				# Create a wire if none exists
				all_wires = sink_wires + source_wires
				if len(all_wires) == 0:
					for driverElement, driverPort in [*sources, *sinks, (element, port)]:
						if driverElement is None:
							continue

						# Inputs are outputs and outputs are inputs when looking at the interface from the module inside.
						expected_direction = 'input' if isinstance(instances[driverElement], Module) else 'output'
						if driverElement.get_rsn_port_direction(driverPort) == expected_direction:
							net_name = f'{driverElement.instance_name}_{driverPort}'
							break
					else:
						assert False, "No output port driving the connection has been found"

					net = NetDeclaration([Identifier(net_name)], range=None)
					module.net_declarations += [net]
					all_wires += [Identifier(net_name)]

				# Connect the ports with the existing wire
				for otherElement, otherPort in [*sinks, *sources, (element, port)]:
					if otherElement is None:
						continue

					if (otherElement, otherPort) in connected_ports:
						continue

					connected_ports.add((otherElement, otherPort))
					if isinstance(instances[otherElement], Module):
						continue

					port_name = otherElement.get_rsn_port_name(otherPort)
					if port_name is None:
						continue

					instances[otherElement].ports[Identifier(port_name)] = all_wires[0]


class Rsn(object):

	def __init__(self):
		self.modules: Dict[str, RsnModule] = {}

	def add_module(self, module: RsnModule):
		self.modules[module.module] = module

	def get_module(self, name: str):
		return self.modules[name]

	def apply_rsn(self, modules: List[Module]):
		for module in modules:
			rsn_module = self.modules.get(module.module_name.name, None)
			if rsn_module is None:
				print(f'Skipping module {rsn_module.module}')
				continue

			print(f'Applying module {rsn_module.module}')
			rsn_module.apply_rsn(module)


if __name__ == '__main__':
	module = RsnModule('test')
	rsn_port = module.add_host_interface('port1')
	scan_cell1 = module.add_internal_scan_cell('cell1', 'rsn_scan_cell')
	scan_cell2 = module.add_internal_scan_cell('cell2', 'rsn_scan_cell')
	scan_cell3 = module.add_internal_scan_cell('cell3', 'rsn_scan_cell')
	scan_cell4 = module.add_internal_scan_cell('cell4', 'rsn_scan_cell')
	scan_cell5 = module.add_internal_scan_cell('cell5', 'rsn_scan_cell')
	scan_cell6 = module.add_internal_scan_cell('cell6', 'rsn_scan_cell')
	sib1 = module.add_sib('sib1', 'sib')
	module.link_rsn_elements(rsn_port, scan_cell1, interface='host')
	module.link_rsn_elements(scan_cell1, scan_cell2)
	module.link_rsn_elements(scan_cell2, scan_cell3)
	module.link_rsn_elements(scan_cell3, sib1)
	module.link_rsn_elements(sib1, scan_cell4, interface='host')
	module.link_rsn_elements(scan_cell4, scan_cell5)
	module.link_rsn_elements(scan_cell5, scan_cell6)

	module.get_dot_graph(stdout, mode='simple')