// SPDX-License-Identifier: Apache-2.0

package chiselTests

import org.scalacheck._

import chisel3._
import chisel3.stage.ChiselStage
import chisel3.testers.BasicTester
import chisel3.util._
import org.scalacheck.Shrink

class LitTesterMod(vecSize: Int) extends Module {
  val io = IO(new Bundle {
    val out = Output(Vec(vecSize, UInt()))
  })
  io.out := VecInit(Seq.fill(vecSize){0.U})
}

class RegTesterMod(vecSize: Int) extends Module {
  val io = IO(new Bundle {
    val in = Input(Vec(vecSize, UInt()))
    val out = Output(Vec(vecSize, UInt()))
  })
  val vecReg = RegNext(io.in, VecInit(Seq.fill(vecSize){0.U}))
  io.out := vecReg
}

class IOTesterMod(vecSize: Int) extends Module {
  val io = IO(new Bundle {
    val in = Input(Vec(vecSize, UInt()))
    val out = Output(Vec(vecSize, UInt()))
  })
  io.out := io.in
}

class OneBitUnitRegVec extends Module {
  val io = IO(new Bundle {
    val out = Output(UInt(1.W))
  })
  val oneBitUnitRegVec = Reg(Vec(1, UInt(1.W)))
  oneBitUnitRegVec(0) := 1.U(1.W)
  io.out := oneBitUnitRegVec(0)
}

class LitTester(w: Int, values: List[Int]) extends BasicTester {
  val dut = Module(new LitTesterMod(values.length))
  for (a <- dut.io.out)
    assert(a === 0.U)
  stop()
}

class RegTester(w: Int, values: List[Int]) extends BasicTester {
  val v = VecInit(values.map(_.U(w.W)))
  val dut = Module(new RegTesterMod(values.length))
  val doneReg = RegInit(false.B)
  dut.io.in := v
  when (doneReg) {
    for ((a,b) <- dut.io.out.zip(values))
      assert(a === b.U)
    stop()
  } .otherwise {
    doneReg := true.B
    for (a <- dut.io.out)
      assert(a === 0.U)
  }
}

class IOTester(w: Int, values: List[Int]) extends BasicTester {
  val v = VecInit(values.map(_.U(w.W))) // Does this need a Wire? No. It's a Vec of Lits and hence synthesizeable.
  val dut = Module(new IOTesterMod(values.length))
  dut.io.in := v
  for ((a,b) <- dut.io.out.zip(values)) {
    assert(a === b.U)
  }
  stop()
}

class IOTesterModFill(vecSize: Int) extends Module {
  // This should generate a BindingException when we attempt to wire up the Vec.fill elements
  //  since they're pure types and hence unsynthesizeable.
  val io = IO(new Bundle {
    val in = Input(VecInit(Seq.fill(vecSize) {UInt()}))
    val out = Output(VecInit(Seq.fill(vecSize) {UInt()}))
  })
  io.out := io.in
}

class ValueTester(w: Int, values: List[Int]) extends BasicTester {
  val v = VecInit(values.map(_.asUInt(w.W)))
  for ((a,b) <- v.zip(values)) {
    assert(a === b.asUInt)
  }
  stop()
}

class TabulateTester(n: Int) extends BasicTester {
  val v = VecInit(Range(0, n).map(i => (i*2).asUInt))
  val x = VecInit(Array.tabulate(n){ i => (i*2).asUInt })
  val u = VecInit.tabulate(n)(i => (i*2).asUInt)

  assert(v.asUInt() === x.asUInt())
  assert(v.asUInt() === u.asUInt())
  assert(x.asUInt() === u.asUInt())

  stop()
}

class FillTester(n: Int, value: Int) extends BasicTester {
  val x = VecInit(Array.fill(n)(value.U))
  val u = VecInit.fill(n)(value.U)

  assert(x.asUInt() === u.asUInt(), s"Expected Vec to be filled like $x, instead VecInit.fill created $u")
  stop()
}

class IterateTester(start: Int, len: Int)(f: UInt => UInt) extends BasicTester {
  val controlVec = VecInit(Seq.iterate(start.U, len)(f))
  val testVec = VecInit.iterate(start.U, len)(f)
  assert(controlVec.asUInt() === testVec.asUInt(), s"Expected Vec to be filled like $controlVec, instead creaeted $testVec\n")
  stop()
}

class ShiftRegisterTester(n: Int) extends BasicTester {
  val (cnt, wrap) = Counter(true.B, n*2)
  val shifter = Reg(Vec(n, UInt((log2Ceil(n) max 1).W)))
  (shifter, shifter drop 1).zipped.foreach(_ := _)
  shifter(n-1) := cnt
  when (cnt >= n.asUInt) {
    val expected = cnt - n.asUInt
    assert(shifter(0) === expected)
  }
  when (wrap) {
    stop()
  }
}

class HugeVecTester(n: Int) extends BasicTester {
  require(n > 0)
  val myVec = Wire(Vec(n, UInt()))
  myVec.foreach { x =>
    x := 123.U
    assert(x === 123.U)
  }
  stop()
}

class OneBitUnitRegVecTester extends BasicTester {
  val dut = Module(new OneBitUnitRegVec)
  assert(dut.io.out === 1.U)
  stop()
}

class ZeroEntryVecTester extends BasicTester {
  require(Vec(0, Bool()).getWidth == 0)

  val bundleWithZeroEntryVec = new Bundle {
    val foo = Bool()
    val bar = Vec(0, Bool())
  }
  require(0.U.asTypeOf(bundleWithZeroEntryVec).getWidth == 1)
  require(bundleWithZeroEntryVec.asUInt.getWidth == 1)

  val m = Module(new Module {
    val io = IO(Output(bundleWithZeroEntryVec))
    io.foo := false.B
  })
  WireDefault(m.io.bar)

  stop()
}

class PassthroughModuleTester extends Module {
  val io = IO(Flipped(new PassthroughModuleIO))
  // This drives the input of a PassthroughModule
  io.in := 123.U
  assert(io.out === 123.U)
}

class ModuleIODynamicIndexTester(n: Int) extends BasicTester {
  val duts = VecInit(Seq.fill(n)(Module(new PassthroughModule).io))
  val tester = Module(new PassthroughModuleTester)

  val (cycle, done) = Counter(true.B, n)
  for ((m, i) <- duts.zipWithIndex) {
    when (cycle =/= i.U) {
      m.in := 0.U  // default
      assert(m.out === 0.U)
    } .otherwise {
      m.in := DontCare
    }
  }
  // only connect one dut per cycle
  duts(cycle) <> tester.io
  assert(duts(cycle).out === 123.U)

  when (done) { stop() }
}

class VecSpec extends ChiselPropSpec with Utils {
  // Disable shrinking on error.
  implicit val noShrinkListVal = Shrink[List[Int]](_ => Stream.empty)
  implicit val noShrinkInt = Shrink[Int](_ => Stream.empty)

  property("Vecs should be assignable") {
    forAll(safeUIntN(8)) { case(w: Int, v: List[Int]) =>
      assertTesterPasses{ new ValueTester(w, v) }
    }
  }

  property("Vecs should be passed through vec IO") {
    forAll(safeUIntN(8)) { case(w: Int, v: List[Int]) =>
      assertTesterPasses{ new IOTester(w, v) }
    }
  }

  property("Vec.fill with a pure type should generate an exception") {
    // We don't really need a sequence of random widths here, since any should throw an exception.
    forAll(safeUIntWidth) { case(w: Int) =>
      an[BindingException] should be thrownBy extractCause[BindingException] {
        ChiselStage.elaborate(new IOTesterModFill(w))
      }
    }
  }

  property("A Reg of a Vec should operate correctly") {
    forAll(safeUIntN(8)) { case(w: Int, v: List[Int]) =>
      assertTesterPasses{ new RegTester(w, v) }
    }
  }

  property("A Vec of lit should operate correctly") {
    forAll(safeUIntN(8)) { case(w: Int, v: List[Int]) =>
      assertTesterPasses{ new LitTester(w, v) }
    }
  }

  property("VecInit should tabulate correctly") {
    forAll(smallPosInts) { (n: Int) => assertTesterPasses{ new TabulateTester(n) } }
  }

  property("VecInit should fill correctly") {
    forAll(smallPosInts, Gen.choose(0, 50)) { (n: Int, value: Int) => assertTesterPasses{ new FillTester(n, value) } }
  }

  property("VecInit should iterate correctly") {
    forAll(Gen.choose(1, 10), smallPosInts) { (start: Int, len: Int) => assertTesterPasses{ new IterateTester(start, len)(x => x + 50.U)}}
  }

  property("Regs of vecs should be usable as shift registers") {
    forAll(smallPosInts) { (n: Int) => assertTesterPasses{ new ShiftRegisterTester(n) } }
  }

  property("Infering widths on huge Vecs should not cause a stack overflow") {
    assertTesterPasses { new HugeVecTester(10000) }
  }

  property("A Reg of a Vec of a single 1 bit element should compile and work") {
    assertTesterPasses{ new OneBitUnitRegVecTester }
  }

  property("A Vec with zero entries should compile and have zero width") {
    assertTesterPasses{ new ZeroEntryVecTester }
  }

  property("Dynamic indexing of a Vec of Module IOs should work") {
    assertTesterPasses{ new ModuleIODynamicIndexTester(4) }
  }

  property("It should be possible to bulk connect a Vec and a Seq") {
    ChiselStage.elaborate(new Module {
      val io = IO(new Bundle {
        val out = Output(Vec(4, UInt(8.W)))
      })
      val seq = Seq.fill(4)(0.U)
      io.out <> seq
    })
  }

  property("Bulk connecting a Vec and Seq of different sizes should report a ChiselException") {
    a [ChiselException] should be thrownBy extractCause[ChiselException] {
      ChiselStage.elaborate(new Module {
        val io = IO(new Bundle {
          val out = Output(Vec(4, UInt(8.W)))
        })
        val seq = Seq.fill(5)(0.U)
        io.out <> seq
      })
    }
  }

  property("It should be possible to initialize a Vec with DontCare") {
    ChiselStage.elaborate(new Module {
      val io = IO(new Bundle {
        val out = Output(Vec(4, UInt(8.W)))
      })
      io.out := VecInit(Seq(4.U, 5.U, DontCare, 2.U))
    })
  }

  property("Indexing a Chisel type Vec by a hardware type should give a sane error message") {
    a [ExpectedHardwareException] should be thrownBy extractCause[ChiselException] {
      ChiselStage.elaborate{
        new Module {
          val io = IO(new Bundle{})
          val foo = Vec(2, Bool())
          foo(0.U) := false.B
        }}
    }
  }
}