CNRV

import chisel3._
import chisel3.util._
import chisel3.stage.ChiselStage

// z = s ? i1 : i0
class Mux2[T <: Data](gen: T) extends BlackBox {
  val io = IO(new Bundle {
    val i0 = Input(gen)
    val i1 = Input(gen)
    val s = Input(Bool())
    val z = Output(gen)
  })
}

object Mux2 {
  def apply[T <: Data](s: Bool, i1: T, i0: T): T = {
    val c = Module(new Mux2(chiselTypeOf(i0)))
    c.io.s := s
    c.io.i1 := i1
    c.io.i0 := i0
    c.io.z
  }
}

// r = d[s], N: d width, 1<s<N
object MuxN {
  private def iter[T <: Data](d: Vec[T], s: UInt, dw: Int): T = dw match {
    case 1 => d(0)
    case 2 => Mux2(s(0), d(1), d(0))
    case _ =>
      val sw = s.getWidth
      val half = 1 << (sw - 1)
      val sd = d.grouped(half).toSeq
      Mux2(
        s(sw - 1),
        iter(VecInit(sd(1)), s(sw - 2, 0), sd(1).size),
        iter(VecInit(sd(0)), s(sw - 2, 0), sd(0).size)
      )
  }

  def apply[T <: Data](d: Vec[T], s: UInt): T = iter(d, s, d.size)
}