from migen.fhdl.std import *

(SP_WITHDRAW, SP_CE) = range(2)

class RoundRobin(Module):
	def __init__(self, n, switch_policy=SP_WITHDRAW):
		self.request = Signal(n)
		self.grant = Signal(max=max(1, n))
		self.switch_policy = switch_policy
		if self.switch_policy == SP_CE:
			self.ce = Signal()
	
		###

		if n > 1:
			cases = {}
			for i in range(n):
				switch = []
				for j in reversed(range(i+1,i+n)):
					t = j % n
					switch = [
						If(self.request[t],
							self.grant.eq(t)
						).Else(
							*switch
						)
					]
				if self.switch_policy == SP_WITHDRAW:
					case = [If(~self.request[i], *switch)]
				else:
					case = switch
				cases[i] = case
			statement = Case(self.grant, cases)
			if self.switch_policy == SP_CE:
				statement = If(self.ce, statement)
			self.sync += statement
		else:
			self.comb += self.grant.eq(0)