GtkwaveTransform.scala

import chisel3.util.log2Ceil
import chisel3.experimental.EnumAnnotations.{EnumComponentAnnotation, EnumVecAnnotation, EnumDefAnnotation}
import firrtl.{CircuitState, DependencyAPIMigration, Transform}
import firrtl.options.TargetDirAnnotation
import firrtl.stage.{Forms, RunFirrtlTransformAnnotation}
import firrtl.analyses.InstanceKeyGraph
import firrtl.ir.{GroundType, SIntType, IntWidth, Port}
import chiseltest.internal.{VerilatorBackendAnnotation, WriteVcdAnnotation}
import logger.LazyLogging

object GtkwFlag extends Enumeration {
  val highlight, hex, dec, bin, oct, rjustify, invert, reverse, exclude, blank, signed, ascii, collapsed, ftranslated,
    ptranslated, analog_step, analog_interpolated, analog_blank_stretch, real, analog_fullscale, zerofill, onefill,
    closed, grp_begin, grp_end, bingray, graybin, real2bits, ttranslated, popcnt, fpdecshif = Value
}

case class GtkwaveFlags(flags: GtkwFlag.Value*) {
  override def toString =
    s"@${flags.map(BigInt(1) << _.id).reduce(_ | _).toString(16)}"

  def :+(flag: GtkwFlag.Value): GtkwaveFlags       = GtkwaveFlags(this.flags :+ flag: _*)
  def addIf(cond: Boolean, flags: GtkwFlag.Value*) = if (cond) GtkwaveFlags(this.flags ++ flags: _*) else this
}

class GtkwaveTransform extends Transform with DependencyAPIMigration with LazyLogging {
  override def prerequisites = Forms.BackendEmitters

  override def invalidates(a: Transform) = false

  val colors: Seq[String] = Seq.empty[String]

  def getColor(v: Int) =
    if (colors.nonEmpty) s"?${colors(v % colors.length)}?" else ""

  override def execute(state: CircuitState): CircuitState = {

    val targetDirOpt = state.annotations.collectFirst { case TargetDirAnnotation(path) =>
      path
    }

    val tdPath = os.pwd / os.RelPath(targetDirOpt.get)

    val superTop = state.annotations.collectFirst { case VerilatorBackendAnnotation =>
      logger.warn("[GtkwaveTransform] Verilator backend detected!")
      "Top"
    }.toSeq

    val vcdEnabled = state.annotations
      .collectFirst { case WriteVcdAnnotation | treadle.WriteVcdAnnotation =>
        true
      }
      .getOrElse(false)

    if (!vcdEnabled) {
      logger.warn(
        s"[GtkwaveTransform] No VCD generation Annotation detected. NOT generating GtkWave filters and savefile!"
      )
      return state
    }

    if (!os.exists(tdPath))
      os.makeDir.all(tdPath)

    val enumDefsMap = state.annotations.collect {
      case EnumVecAnnotation(target, typeName, fields) =>
        logger.error(
          s"[GtkwaveTransform] NOT HANDLED! Vec $target (${target.serialize}) -> $typeName  fields: $fields"
        ) // TODO FIXME need to figure it out
        target.serialize -> typeName
      case EnumComponentAnnotation(target, typeName) =>
        target.serialize -> typeName
    }.toMap

    val typToFilterMap = state.annotations.collect {
      case EnumDefAnnotation(typeName: String, definition: Map[String, BigInt]) =>
        val path = tdPath / s"${typeName}.gtkw-filter"
        logger.warn(s"[GtkwaveTransform] Generating GtkWave filter: ${path}")
        val numBits   = log2Ceil(definition.size)
        val hex       = numBits > 4
        val numDigits = if (hex) (numBits.toDouble / 4).ceil.toInt else numBits
        def formatValue(v: Int): String = {
          val s = if (hex) v.toHexString else v.toBinaryString
          ("0" * (numDigits - s.size)) + s
        }
        os.write.over(
          path,
          definition.toSeq
            .sortBy(_._2)
            .map { case (k, v) =>
              s"${formatValue(v.toInt)} ${getColor(v.toInt)}$k"
            }
            .mkString("\n")
        )
        typeName -> path.toString
    }.toMap

    val vcdPath = tdPath / (state.circuit.main + ".vcd")

    val ikg = InstanceKeyGraph(state.circuit)

    def portIntWidth(p: Port): Option[BigInt] = p.tpe match {
      case GroundType(IntWidth(w)) => Some(w)
      case _                       => None
    }

    def getPorts(instPath: Seq[InstanceKeyGraph.InstanceKey], module: String): Seq[String] = {
      var lastWidthFlag                = GtkwaveFlags()
      def portName(p: Port, w: BigInt) = (p.name, if (w > 1) s"[${w - 1}:0]" else "")
      def widthFlags(signed: Boolean, w: BigInt) =
        GtkwaveFlags(if (signed) GtkwFlag.signed else if (w >= 4) GtkwFlag.hex else GtkwFlag.bin, GtkwFlag.rjustify)

      def portNameAndFlags(p: Port) = p.tpe match {
        case SIntType(IntWidth(w)) =>
          (portName(p, w), widthFlags(true, w))
        case GroundType(IntWidth(w)) =>
          (portName(p, w), widthFlags(false, w)) // show everything else as hex/binary for now
        case _ =>
          (portName(p, 0), widthFlags(false, 0))
      }

      ikg.moduleMap
        .get(module)
        .toSeq
        .flatMap { mod =>
          mod.ports.map { p =>
            val ((name, vecSuffix), nwf) = portNameAndFlags(p)
            val flags = Option(nwf != lastWidthFlag).collect { case true =>
              lastWidthFlag = nwf
              Seq(nwf.toString)
            }
            val instanceName = (instPath.map(_.module) :+ name).mkString(".")
            val gtkwName     = (superTop ++ instPath.map(_.name) :+ (name + vecSuffix)).mkString(".")
            val filterFileFlags = enumDefsMap.get(instanceName).map(typToFilterMap.get(_)).flatten.map { filterFile =>
              lastWidthFlag = nwf :+ GtkwFlag.ftranslated
              Seq(
                lastWidthFlag.toString(), // TODO only add if this one is != lastWidthFlag, but should be ok
                "^1 " + filterFile
              )
            }
            (filterFileFlags.orElse(flags).getOrElse(Seq()) :+ gtkwName).mkString("\n")
          }
        }
    }

    def mkGroup(name: String, signals: Seq[String], closed: Boolean = false): Seq[String] = {
      Seq(
        GtkwaveFlags(GtkwFlag.grp_begin, GtkwFlag.blank).addIf(closed, GtkwFlag.closed).toString(),
        s"-${name}"
      ) ++ signals ++ Seq(
        GtkwaveFlags(GtkwFlag.grp_end, GtkwFlag.blank).addIf(closed, GtkwFlag.closed, GtkwFlag.collapsed).toString(),
        s"-${name}"
      )
    }

    val signals: Seq[String] = ikg.fullHierarchy
      .collect { case (m, paths) =>
        val sigs = paths.map(p => getPorts(p, m.module)).flatten
        mkGroup(m.name, sigs, m != ikg.top)
      }
      .flatten
      .toSeq

    val gtkwPath = tdPath / (state.circuit.main + ".gtkw")

    val gtkwLines = Seq(
      "[*]",
      "[*] Generated by GtkwaveTransform",
      "[*]",
      s"""[dumpfile] \"${vcdPath}\"""",
      s"""[savefile] \"${gtkwPath}\"""",
      "[timestart] 0",
      "[signals_width] 250",
      "[sst_width] 250",
      "[sst_expanded] 1",
      s"""[treeopen] \"${ikg.top.name}\""""
    ) ++ signals

    os.write.over(gtkwPath, gtkwLines.mkString("\n"))

    state
  }
}

object GtkwaveFiltersAnnotation extends RunFirrtlTransformAnnotation(new GtkwaveTransform)

object GtkwaveColoredFiltersAnnotation
    extends RunFirrtlTransformAnnotation(new GtkwaveTransform {
      override val colors = Seq(
        "tomato1",
        "cadet blue",
        "orange",
        "orchid",
        "dark sea green",
        "dark khaki",
        "yellow green",
        "blue violet",
        "indigo",
        "violet",
        "tan1",
        "turquoise",
        "brown",
        "dark orchid",
        "dark grey",
        "dark cyan",
        "hot pink",
        "steel blue",
        "tan",
        "coral4"
      )
    })