How to use the pyrtl.mux function in pyrtl

def test_verilog_testbench_does_not_throw_error(self):
        zero = pyrtl.Input(1, 'zero')
        counter_output = pyrtl.Output(3, 'counter_output')
        counter = pyrtl.Register(3, 'counter')
        counter.next <<= pyrtl.mux(zero, counter + 1, 0)
        counter_output <<= counter
        sim_trace = pyrtl.SimulationTrace([counter_output, zero])
        sim = pyrtl.Simulation(tracer=sim_trace)
        for cycle in range(15):
            sim.step({zero: random.choice([0, 0, 0, 1])})
        with io.StringIO() as tbfile:
            pyrtl.output_verilog_testbench(tbfile, sim_trace)

def test_as_graph_memory(self):
        m = pyrtl.MemBlock(addrwidth=2, bitwidth=2, name='m', max_read_ports=None)
        i = pyrtl.Register(bitwidth=2, name='i')
        o = pyrtl.WireVector(bitwidth=2, name='o')
        i.next <<= i + 1
        m[i] <<= pyrtl.mux((m[i] != 0), 0, m[i])
        o <<= m[i]

        b = pyrtl.working_block()
        src_g, dst_g = b.net_connections(False)
        self.check_graph_correctness(src_g, dst_g)

        src_g, dst_g = b.net_connections(True)
        self.check_graph_correctness(src_g, dst_g, True)

def setUp(self):
        pyrtl.reset_working_block()
        bitwidth = 3
        self.a = pyrtl.Input(bitwidth=bitwidth)
        self.b = pyrtl.Input(bitwidth=bitwidth)
        self.sel = pyrtl.Input(bitwidth=1)
        self.muxout = pyrtl.Output(bitwidth=bitwidth, name='muxout')
        self.muxout <<= pyrtl.mux(self.sel, self.a, self.b)

        # build the actual simulation environment
        self.sim_trace = pyrtl.SimulationTrace()
        self.sim = self.sim(tracer=self.sim_trace)

def setUp(self):
        pyrtl.reset_working_block()
        bitwidth = 3
        self.a = pyrtl.Input(bitwidth=bitwidth)
        self.b = pyrtl.Input(bitwidth=bitwidth)
        self.sel = pyrtl.Input(bitwidth=1)
        self.muxout = pyrtl.Output(bitwidth=bitwidth, name='muxout')
        self.muxout <<= pyrtl.mux(self.sel, self.a, self.b)

        # build the actual simulation environment
        self.sim_trace = pyrtl.SimulationTrace()
        self.sim = pyrtl.FastSimulation(tracer=self.sim_trace)

#  L--->|	   |
#	|__________|
#
#
north_in=pyrtl.Input(9,'north_in')	#Input request from VC 
south_in=pyrtl.Input(9,'south_in')
east_in=pyrtl.Input(9,'east_in')
west_in=pyrtl.Input(9,'west_in')
self_in=pyrtl.Input(9,'self_in')
surf_sch = pyrtl.Input(1,'surf_sch')
sel_read_buffer = pyrtl.Output(2,'sel_read_buffer')	
vc = pyrtl.WireVector(2,'vc')
out_port = pyrtl.Output(4,'out_port')
port_north = pyrtl.mux((north_in[4:8]==15),truecase=north_in[4:8],falsecase=0)	#checking if local is requested by the input signals
port_south = pyrtl.mux((south_in[4:8]==15),truecase=south_in[4:8],falsecase=0)
port_east = pyrtl.mux((east_in[4:8]==15),truecase=east_in[4:8],falsecase=0)
port_west = pyrtl.mux((west_in[4:8]==15),truecase=west_in[4:8],falsecase=0)
port_self = pyrtl.mux((self_in[4:8]==15),truecase=self_in[4:8],falsecase=0)
with pyrtl.ConditionalUpdate() as condition:	#taking out the VC and port from the request
    with condition((self_in[8]==1)&(port_self!=0)):
	vc |= self_in[4]
	out_port |= self_in[4:8]
    with condition((north_in[8]==1)&(port_north!=0)):
	vc |= north_in[4]
	out_port |= north_in[4:8]
    with condition((south_in[8]==1)&(port_south!=0)):
	vc |= south_in[4]
	out_port |= south_in[4:8]
    with condition((east_in[8]==1)&(port_east!=0)):
	vc |= east_in[4]
	out_port |= east_in[4:8]
    with condition((west_in[8]==1)&(port_west!=0)):

head = pyrtl.Register(addrwidth) # write pointer into the circular buffer
    tail = pyrtl.Register(addrwidth) # read pointer into the circular buffer
    count = pyrtl.Register(addrwidth+1)  # number of elements currently stored in buffer
    
    full = pyrtl.mux(count >= 2**addrwidth, truecase=1, falsecase=0)
    do_write = pyrtl.mux(full, truecase=0, falsecase=write_enable)
    empty = (~do_write) & (count==0)
    do_read = pyrtl.mux(empty, truecase=0, falsecase=read_enable)

    buffer_memory[head] <<= pyrtl.MemBlock.EnabledWrite(data, do_write)

    head.next <<= pyrtl.mux(do_write, truecase=head+1, falsecase=head)
    tail.next <<= pyrtl.mux(do_read, truecase=tail+1, falsecase=tail)
    count.next <<= count + do_write - do_read

    read_output = pyrtl.mux(do_read & do_write & (head==tail), truecase=data, falsecase=buffer_memory[tail])
    return (read_output, do_read, full)

# 	-------------------------
#	|G|N|S|E|W|port 3bits|VC|		9bits
#	-------------------------
#	port N = 0 | S = 1 | E = 2 | W= 3 | Self = 4 |
#
port = pyrtl.Input(3,'port')	#port which is giving Input signals

go_x=pyrtl.Output(9,'go_x')	#the X co-ordinate Request	##select X or Y request || X first then Y
go_y=pyrtl.Output(9,'go_y')	#the Y co-ordinate Request
x=pyrtl.WireVector(2,'x')	#X co-ordinate of the route
y=pyrtl.WireVector(2,'y')	#Y co-ordinate of the route
vc = pyrtl.Input(1,'vc')	# Virtual channel

x<<=pyrtl.mux((router_id[0:2]>dest_id[0:2]),truecase=router_id[0:2]-dest_id[0:2],falsecase=dest_id[0:2]-router_id[0:2])
y<<=pyrtl.mux((router_id[2:4]>dest_id[2:4]),truecase=router_id[2:4]-dest_id[2:4],falsecase=dest_id[2:4]-router_id[2:4])
x1 = pyrtl.mux(((x==3)|(dest_id[0:2]>router_id[0:2])),truecase=18,falsecase=pyrtl.mux(router_id[0:2]>dest_id[0:2],truecase=17,falsecase=pyrtl.mux(((x==0)&(y==0)),truecase=31, falsecase=00000)))
y1 = pyrtl.mux(((y==3)|(dest_id[2:4]>router_id[2:4])),truecase=24,falsecase=pyrtl.mux(router_id[2:4]>dest_id[2:4],truecase=20,falsecase=pyrtl.mux(((y==0)&(x==0)),truecase=31,falsecase=00000)))


go_x <<= pyrtl.concat(x1,port,vc)
go_y <<= pyrtl.concat(y1,port,vc)

simvals ={
        router_id:[3,9,15,15,2,0,0],
	dest_id:[3,3,0,4,0,1,4],
	port:[0,1,2,3,4,2,1],
	vc: [0,1,0,1,0,1,1] 
}
sim_trace = pyrtl.SimulationTrace()
sim = pyrtl.Simulation(tracer=sim_trace)
for cycle in xrange(len(simvals[dest_id])):
    sim.step({k: v[cycle] for k,v in simvals.items()    })

def surfnoc_buffer(bitwidth, addrwidth, data, write_enable, read_enable):
    """ """

    buffer_memory = pyrtl.MemBlock(bitwidth=bitwidth, addrwidth=addrwidth)

    head = pyrtl.Register(addrwidth) # write pointer into the circular buffer
    tail = pyrtl.Register(addrwidth) # read pointer into the circular buffer
    count = pyrtl.Register(addrwidth+1)  # number of elements currently stored in buffer
    
    full = pyrtl.mux(count >= 2**addrwidth, truecase=1, falsecase=0)
    do_write = pyrtl.mux(full, truecase=0, falsecase=write_enable)
    empty = (~do_write) & (count==0)
    do_read = pyrtl.mux(empty, truecase=0, falsecase=read_enable)

    buffer_memory[head] <<= pyrtl.MemBlock.EnabledWrite(data, do_write)

    head.next <<= pyrtl.mux(do_write, truecase=head+1, falsecase=head)
    tail.next <<= pyrtl.mux(do_read, truecase=tail+1, falsecase=tail)
    count.next <<= count + do_write - do_read

    read_output = pyrtl.mux(do_read & do_write & (head==tail), truecase=data, falsecase=buffer_memory[tail])
    return (read_output, do_read, full)