// SPDX-License-Identifier: Apache-2.0

package chiselTests

import chisel3._
import chisel3.experimental.doNotDedup
import chisel3.stage.{ChiselGeneratorAnnotation, ChiselStage}
import firrtl.stage.FirrtlCircuitAnnotation
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.should.Matchers


class MuchUsedModule extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(16.W))
    val out = Output(UInt(16.W))
  })
  io.out := io.in +% 1.U
}

class UsesMuchUsedModule(addAnnos: Boolean) extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(16.W))
    val out = Output(UInt(16.W))
  })

  val mod0 = Module(new MuchUsedModule)
  val mod1 = Module(new MuchUsedModule)
  val mod2 = Module(new MuchUsedModule)
  val mod3 = Module(new MuchUsedModule)

  mod0.io.in := io.in
  mod1.io.in := mod0.io.out
  mod2.io.in := mod1.io.out
  mod3.io.in := mod2.io.out
  io.out := mod3.io.out

  if(addAnnos) {
    doNotDedup(mod1)
    doNotDedup(mod3)
  }
}

class AnnotationNoDedup extends AnyFreeSpec with Matchers {
  val stage = new ChiselStage
  "Firrtl provides transform that reduces identical modules to a single instance" - {
    "Annotations can be added which will prevent this deduplication for specific modules instances" in {
      val lowFirrtl = stage
        .execute(Array("-X", "low", "--target-dir", "test_run_dir"),
                 Seq(ChiselGeneratorAnnotation(() => new UsesMuchUsedModule(addAnnos = true))))
        .collectFirst {
          case FirrtlCircuitAnnotation(circuit) => circuit.serialize
        }.getOrElse(fail)
      lowFirrtl should include ("module MuchUsedModule :")
      lowFirrtl should include ("module MuchUsedModule_1 :")
      lowFirrtl should include ("module MuchUsedModule_3 :")
      lowFirrtl should not include "module MuchUsedModule_2 :"
      lowFirrtl should not include "module MuchUsedModule_4 :"
    }
    "Turning off these annotations dedups all the occurrences" in {
      val lowFirrtl = stage
        .execute(Array("-X", "low", "--target-dir", "test_run_dir"),
                 Seq(ChiselGeneratorAnnotation(() => new UsesMuchUsedModule(addAnnos = false))))
        .collectFirst {
          case FirrtlCircuitAnnotation(circuit) => circuit.serialize
        }.getOrElse(fail)
      lowFirrtl should include ("module MuchUsedModule :")
      lowFirrtl should not include "module MuchUsedModule_1 :"
      lowFirrtl should not include "module MuchUsedModule_3 :"
      lowFirrtl should not include "module MuchUsedModule_2 :"
      lowFirrtl should not include "module MuchUsedModule_4 :"
    }
  }
}