// Copyright 2017 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package freechips.rocketchip.system

import chisel3._
import chiffre.{ChiffreParameters, LeChiffre, BuildChiffre}
import freechips.rocketchip.config.{Config, Parameters}
import freechips.rocketchip.tile.{BuildRoCC, OpcodeSet}
import freechips.rocketchip.subsystem.RocketTilesKey
import freechips.rocketchip.diplomacy.LazyModule

class WithLeChiffre extends Config (
  (site, here, up) => {
    case BuildChiffre => ChiffreParameters()
    case BuildRoCC => List(
      (p: Parameters) => {
        val chiffre = LazyModule(new LeChiffre(OpcodeSet.custom2, "main")(p))
        chiffre })
  })

class LeChiffreConfig extends Config(new WithLeChiffre ++ new DefaultConfig)


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre

case class InjectorInfoBindingException(msg: String) extends Exception(msg)

trait InjectorInfo extends HasName with HasWidth with Equals {
  /* All configurable fields for this specific injector */
  val fields: Seq[ScanField]

  val name: String = this.getClass.getSimpleName

  /* The width of this injector's scan chain configuration */
  lazy val width: Int = fields.foldLeft(0)( (l, r) => l + r.width )

  /* Prety print */
  def serialize(tab: String = ""): String = {
    s"""|${tab}class: $name
        |${tab}width: $width

      .stripMargin
  }

  def toBits(): String = fields.map(_.toBits()).mkString

  def bind(in: Seq[ScanField]): InjectorInfo = {
    val widthIn = in.foldLeft(0)( (l, r) => l + r.width )
    if (width != widthIn) {
      throw new InjectorInfoBindingException(s"Unable to bind, widths differ ($width vs. $widthIn) between $name and $in") }
    if (!(fields.zip(in).map{ case (l, r) => l canEqual r }).foldLeft(true)( (l, r) => l & r )) {
      throw new InjectorInfoBindingException(s"Unable to bind, types differ between $name and $in") }
    fields.zip(in).map{ case (l, r) => l.bind(r.value) }
    this
  }

  def isBound(): Boolean = fields.map(_.isBound).reduceOption(_ && _).getOrElse(true)

  def unbind(): InjectorInfo = {
    fields.map(_.unbind)
    this
  }

  override def equals(that: Any): Boolean = that match {
    case that: InjectorInfo => (this canEqual that) &&
        width == that.width &&
        fields.zip(that.fields).foldLeft(true){ case (x, (l, r)) => x & (l == r) }
    case _ => false
  }

  override def hashCode = fields.hashCode
}

|${fields.map(a => s"${a.serialize(tab + "  ")}").mkString("\n")}"""


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre

case class ScanFieldException(msg: String) extends Exception(msg)
case class ScanFieldUnboundException(msg: String) extends Exception(msg)

trait HasWidth {
  val width: Int
}

trait HasName {
  val name: String
}

/** A configurable field of the scan chain */
trait ScanField extends HasName with HasWidth with Equals {
  var value: Option[BigInt] = None

  val name: String = this.getClass.getSimpleName

  if (width < 1) {
    throw new ScanFieldException(s"ScanField '$name' width must be greater than 0") }

  def bind(in: BigInt): ScanField = {
    if (in < 0 || in > maxValue) {
      throw new ScanFieldBindingException(
        s"Cannot bind ScanField '$name' to value $in must be on domain [0, $maxValue], but would be ${in}")
    }
    value = Some(in)
    this
  }

  def bind(in: Option[BigInt]): ScanField =
    if (in.nonEmpty) {
      bind(in.get)
    } else {
      value = None
      this
    }

  lazy val maxValue = BigInt(2).pow(width) - 1

  def toBits(): String = s"%${width}s"
    .format(value.getOrElse(BigInt(0)).toString(2))
    .replace(' ', '0')

  def serialize(indent: String = ""): String = {
    s"""|${indent}name: $name
        |${indent}  width: $width
        |${indent}  value: ${toBits}"""
      .stripMargin
  }

  def unbind(): ScanField = {
    value = None
    this
  }

  def isBound(): Boolean = value.nonEmpty

  override def equals(that: Any): Boolean = that match {
    case that: ScanField => (this canEqual that) && (width == that.width) && (value == that.value)
  }

  override def hashCode = value.hashCode
}

trait ProbabilityBind { this: ScanField =>
  def bind(probability: Double): ScanField = bind(BigDecimal((math.pow(2, width) - 1) * probability).toBigInt)
  def bind(probability: Option[Double]): ScanField =
    if (probability.nonEmpty) {
      bind(probability.get)
    } else {
      value = None
      this
    }
}

case class ScanFieldBindingException(msg: String) extends Exception(msg)


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre.inject

import chisel3._
import chiffre.ChiffreInjector

import chiffre.{InjectorInfo, ScanField, HasWidth}

case class CycleInject(width: Int) extends ScanField

case class CycleInjectorInfo(bitWidth: Int, cycleWidth: Int) extends InjectorInfo {
  val fields = Seq(Cycle(cycleWidth), CycleInject(bitWidth))
}

class CycleInjector(bitWidth: Int, val cycleWidth: Int) extends Injector(bitWidth) {
  val cycleTarget = Reg(UInt(cycleWidth.W))
  val cycleCounter = Reg(UInt(cycleWidth.W))
  val flipMask = Reg(UInt(bitWidth.W))

  lazy val info = CycleInjectorInfo(bitWidth, cycleWidth)

  val fire = enabled & (cycleCounter === cycleTarget)
  io.out := Mux(fire, io.in ^ flipMask, io.in)

  when (enabled) {
    cycleCounter := cycleCounter + 1.U
  }

  when (io.scan.clk) {
    enabled := false.B
    cycleCounter := 0.U
    cycleTarget := (io.scan.in ## cycleTarget) >> 1
    flipMask := (cycleTarget(0) ## flipMask) >> 1
  }
  io.scan.out := flipMask(0)

  when (enabled && RegNext(!enabled)) {
    printf(s"""|[info] $name enabled
               |[info]   - target: 0x%x
               |[info]   - mask: 0x%x
               |""".stripMargin, cycleTarget, flipMask)
  }

  when (!enabled && RegNext(enabled)) {
    printf(s"[info] $name disabled\n")
  }

  when (fire) {
    printf(s"[info] $name injecting 0x%x into 0x%x to output 0x%x!\n", flipMask, io.in, io.out)
    enabled := false.B
  }
}

// scalastyle:off magic.number
class CycleInjector32(bitWidth: Int, val scanId: String) extends CycleInjector(bitWidth, 32) with ChiffreInjector
// scalastyle:on magic.number

case class Cycle(width: Int) extends ScanField


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre.inject

import chisel3._
import chiffre.ChiffreInjector

import chiffre.{ScanField, InjectorInfo, ProbabilityBind}

case class Seed(width: Int) extends ScanField

case class LfsrInjectorInfo(bitWidth: Int, lfsrWidth: Int) extends InjectorInfo {
  val fields = Seq.fill(bitWidth)(Seq(Seed(lfsrWidth), Difficulty(lfsrWidth))).flatten
}

sealed class LfsrInjector(val lfsrWidth: Int) extends Injector(1) {
  val difficulty = RegInit(0.U(lfsrWidth.W))
  val seed = RegInit(1.U(lfsrWidth.W))
  lazy val info = LfsrInjectorInfo(1, lfsrWidth)

  val lfsr = Module(new perfect.random.Lfsr(lfsrWidth))
  lfsr.io.seed.valid := io.scan.en
  lfsr.io.seed.bits := seed

  val fire = enabled && (lfsr.io.y <= difficulty)
  io.out := Mux(fire, ~io.in, io.in)

  when (io.scan.clk) {
    enabled := false.B
    seed := (io.scan.in ## seed) >> 1
    difficulty := (seed(0) ## difficulty) >> 1
  }

  io.scan.out := difficulty(0)

  when (enabled && RegNext(!enabled)) {
    printf(s"""|[info] $name enabled
               |[info]   - seed: 0x%x
               |[info]   - difficulty: 0x%x
               |""".stripMargin, seed, difficulty)
  }

  when (!enabled && RegNext(enabled)) {
    printf(s"[info] $name disabled\n")
  }
}

class LfsrInjectorN(bitWidth: Int, val lfsrWidth: Int, val scanId: String) extends
    InjectorBitwise(bitWidth, new LfsrInjector(lfsrWidth)) with ChiffreInjector {
  lazy val info = LfsrInjectorInfo(bitWidth, lfsrWidth)
}

class LfsrInjector32(bitWidth: Int, scanId: String) extends LfsrInjectorN(bitWidth, 32, scanId) // scalastyle:ignore

case class Difficulty(width: Int) extends ScanField with ProbabilityBind


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre.inject

import chisel3._
import chisel3.util.Fill
import chiffre.ChiffreInjector

import chiffre.{ScanField, InjectorInfo}

case class Mask(width: Int) extends ScanField
case class StuckAt(width: Int) extends ScanField

case class StuckAtInjectorInfo(bitWidth: Int) extends InjectorInfo {
  val fields = Seq(Mask(bitWidth), StuckAt(bitWidth))
}

class StuckAtInjector(val bitWidth: Int) extends Injector(bitWidth) {
  val mask = Reg(UInt(bitWidth.W))
  val value = Reg(UInt(bitWidth.W))

  lazy val info = StuckAtInjectorInfo(bitWidth)

  val select = mask & Fill(mask.getWidth, enabled)
  io.out := (io.in & ~select) | (value & select)

  when (io.scan.clk) {
    enabled := false.B
    mask := (io.scan.in ## mask) >> 1
    value := (mask(0) ## value) >> 1
  }

  io.scan.out := value(0)

  when (io.scan.en && !enabled) {
    printf(s"""|[info] $name enabled
               |[info]   - mask: 0x%x
               |[info]   - value: 0x%x
               |""".stripMargin, mask, value)
  }

  when (io.scan.en && enabled) {
    printf(s"[info] $name disabled\n")
  }
}

class StuckAtInjectorWithId(bitWidth: Int, val scanId: String) extends StuckAtInjector(bitWidth) with ChiffreInjector


// Copyright 2017 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre.passes

import chiffre.inject.Injector
import chiffre.util.removeZeroWidth
import chisel3.experimental.{annotate, ChiselAnnotation}
import chisel3.stage.{ChiselGeneratorAnnotation, CircuitSerializationAnnotation}
import firrtl._
import firrtl.ir._
import firrtl.passes.{PassException, ToWorkingIR}
import firrtl.passes.wiring.SinkAnnotation
import firrtl.stage.FirrtlCircuitAnnotation
import firrtl.annotations.{Annotation, ComponentName, ModuleName, CircuitName}
import firrtl.annotations.AnnotationUtils._
import scala.collection.mutable

case class FaultInstrumentationException(msg: String) extends PassException(msg)

case class FaultInstrumentationInfo(orig: ComponentName, conn: ComponentName, repl: ComponentName)

case class Modifications(
  defines: Seq[Statement] = Seq.empty,
  connects: Seq[Statement] = Seq.empty,
  modules: Seq[DefModule] = Seq.empty,
  annotations: Seq[Annotation] = Seq.empty,
  renames: Map[String, String] = Map.empty) {

  override def toString: String = serialize("")

  def serialize(indent: String): String =
    s"""|${indent}defines:
        |${defines.map(a => s"$indent  - ${a.serialize}").mkString("\n")}
        |${indent}connects:
        |${connects.map(a => s"$indent  - ${a.serialize}").mkString("\n")}
        |${indent}modules:
        |${modules.map(a => s"$indent  - ${a.name}").mkString("\n")}
        |${indent}annotations:
        |${annotations.map(a => s"$indent  - ${a.serialize}").mkString("\n")}
        |${indent}renames:
        |${renames.map{case (a, b) => s"$indent  - $a: $b"}.mkString("\n")}"""
      .stripMargin

  def serializeInMemory(indent: String): String =
    s"""|${indent}defines:
        |${defines.map(a => s"$indent  - $a").mkString("\n")}
        |${indent}connects:
        |${connects.map(a => s"$indent  - $a").mkString("\n")}
        |${indent}modules:
        |${modules.map(a => s"$indent  - ${a.name}").mkString("\n")}
        |${indent}annotations:
        |${annotations.map(a => s"$indent  - $a").mkString("\n")}

        |${renames.map{case (a, b) => s"$indent  - $a: $b"}.mkString("\n")}"""
      .stripMargin
}

class FaultInstrumentation(compMap: Map[String, Seq[(ComponentName, String, Class[_ <: Injector])]]) extends Transform {
  def inputForm: CircuitForm = MidForm
  def outputForm: CircuitForm = MidForm
  def execute(state: CircuitState): CircuitState = {
    val modifications = analyze(state.circuit)

    val (mxx, ax) = modifications
      .foldLeft((Seq[DefModule](), Seq[Annotation]())){
        case ((m, a), (_, Modifications(_,_,mm,aa,_))) => (m ++ mm, a ++ aa) }

    val mx = state.circuit.modules map onModule(modifications)
    val cx = ToWorkingIR.run(state.circuit.copy(modules = mxx ++ mx))

    val inAnno: Seq[Annotation] = state.annotations.toSeq
    state.copy(circuit = cx,
               annotations = AnnotationSeq(inAnno ++ ax))
  }

  private def inlineCompile(gen: () => chisel3.Module, ns: Option[Namespace] = None): CircuitState = {
    def genName(name: String, n: Option[Namespace] = None): String = n.map(_.newName(name)).getOrElse(name)

    val (chirrtl: firrtl.ir.Circuit, inlineAnnos: AnnotationSeq) = {
      val outputAnnotations = (new chisel3.stage.ChiselStage).execute(
        args = Array("--no-run-firrtl"),
        annotations = Seq(ChiselGeneratorAnnotation(gen))
      )
      val (circuitAnnos, otherAnnos) = outputAnnotations.partition {
        case _: FirrtlCircuitAnnotation => true
        case _ => false
      }
      require(circuitAnnos.size == 1)
      (circuitAnnos.head.asInstanceOf[FirrtlCircuitAnnotation].circuit, AnnotationSeq(otherAnnos))
    }
    val midFirrtl = (new MiddleFirrtlCompiler)
      .compileAndEmit(CircuitState(chirrtl, ChirrtlForm))
      .circuit
      .mapModule(
        _ match {
          case m: Module    => m.copy(name = genName(m.name, ns))
          case m: ExtModule => m.copy(name = genName(m.name, ns))
        })
    CircuitState(
      circuit = midFirrtl,
      form = MidForm,
      annotations = AnnotationSeq(inlineAnnos))
  }

  private def analyze(c: Circuit): Map[String, Modifications] = {
    val mods = new mutable.HashMap[String, Modifications].withDefaultValue(Modifications())
    val cmods = new mutable.HashMap[ComponentName, CircuitState]()
    val circuitNamespace = Namespace(c)

    c.modules
      .filter((m: DefModule) => compMap.contains(m.name))
      .foreach {
        case m: Module =>
          val moduleNamespace = Namespace(m)
          var scanIn: Option[String] = None
          var scanOut: String = ""
          compMap(m.name) map { case (comp, id, gen)  =>
            val t = passes.wiring.WiringUtils.getType(c, m.name, comp.name)
            val width = bitWidth(t)
            if (width == 0)
              throw new FaultInstrumentationException("Cannot instrument zero-width signals")

            val args = Array[AnyRef](new java.lang.Integer(width.toInt), id)
            val dut = () => {
              val injector = gen.getConstructors()(0)
                .newInstance(args: _*)
                .asInstanceOf[Injector]
              annotate(new ChiselAnnotation {
                def toFirrtl: ScanChainDescriptionAnnotation = ScanChainDescriptionAnnotation(comp, id, injector.info)
              })
              injector
            }
            val (subcir, defms, annosx) = if (cmods.contains(comp)) {
              (cmods(comp).circuit, Seq.empty, Seq.empty)
            } else {
              try {
                cmods(comp) = inlineCompile(dut, Some(circuitNamespace))
              } catch {
                case e: java.lang.IllegalArgumentException => throw new FaultInstrumentationException(
                  s"Did not find '(Int, String)' constructor for injector '${gen.getName}' (Did you forget to specify it?)")
              }
              (cmods(comp).circuit, cmods(comp).circuit.modules,
               if (cmods(comp).annotations.isEmpty) { Seq.empty }
               else { cmods(comp).annotations.toSeq   } )
            }
            val injector = defms.last // the top-level module (the injector) should be the last module of the subcircuit
            val inst = DefInstance(NoInfo, moduleNamespace.newName(s"${comp.name}_injector"), injector.name)
            val rename = moduleNamespace.newName(s"${comp.name}_fault")

            val Seq(scanEn, scanClk, scanIn, scanOut) =
              Seq("en", "clk", "in", "out").map( s =>
                ComponentName(s"${inst.name}.io.scan.$s",
                              ModuleName(m.name, CircuitName(c.main))) )

            val faulty = DefWire(NoInfo, rename, t)
            val data = fromBits(WRef(faulty).mapType(removeZeroWidth.apply), toExp(s"${inst.name}.io.out")) match {
              case Block(stmts: Seq[Statement]) =>
                stmts :+ Connect(NoInfo, toExp(s"${inst.name}.io.in"), toBits(WRef(comp.name, t, RegKind, UnknownFlow).mapType(removeZeroWidth.apply)))
            }
            val x = mods(m.name)
            mods(m.name) = x.copy(
              defines = inst +: x.defines :+ faulty,
              connects = x.connects ++ Seq(
                Connect(NoInfo, toExp(s"${inst.name}.clock"), toExp(s"clock")),
                Connect(NoInfo, toExp(s"${inst.name}.reset"), toExp(s"reset")),
                IsInvalid(NoInfo, toExp(s"${inst.name}.io.scan.en")),
                IsInvalid(NoInfo, toExp(s"${inst.name}.io.scan.clk")),
                IsInvalid(NoInfo, toExp(s"${inst.name}.io.scan.in"))
              ) ++ data,
              modules = x.modules ++ defms,
              annotations = x.annotations ++ annosx ++ Seq(
                SinkAnnotation(scanEn, "scan_en"),
                SinkAnnotation(scanClk, "scan_clk"),
                ScanChainInjectorAnnotation(comp, id, injector.name),
                ScanChainAnnotation(scanIn, "slave", "in", id, Some(comp)),
                ScanChainAnnotation(scanOut, "slave", "out", id, Some(comp))
              ),
              renames = x.renames ++ Map(comp.name -> rename)
            )
          }
        case m: ExtModule =>
          throw new FaultInstrumentationException(
            "Tried to instrument an ExtModule")
      }

    mods.map{ case (k, v) =>
      logger.info(s"[info] $k")
      logger.info(v.serialize("[info]   "))
    }

    mods.toMap
  }

  private def onModule(mods: Map[String, Modifications])
                      (m: DefModule): DefModule = {
    mods.get(m.name) match {
      case None => m
      case Some(l) => m match {
        case m: Module =>
          val x = mods(m.name)
          val mx = m.copy(
            body = Block(
              (x.defines :+ (m.body mapStmt onStmt(x.renames))) ++ x.connects))
          mx
        case _ => m
      }
    }
  }

  private def onStmt(renames: Map[String, String])(s: Statement): Statement = {
    s mapStmt onStmt(renames) match {
      case Connect(i, l, e) =>
        Connect(i, l, replace(renames)(e))
      case PartialConnect(i, l, e) =>
        PartialConnect(i, l, replace(renames)(e))
      case s => s mapExpr replace(renames)
    }
  }

  private def replace(renames: Map[String, String])
                     (e: Expression): Expression = {
    e match {
      case ex: WRef => ex.name match {
        case name if renames.contains(name) => ex.copy(name=renames(name))
        case _ => ex
      }
      case ex: Reference => ex.name match {
        case name if renames.contains(name) => ex.copy(name=renames(name))
        case _ => ex
      }
      case _ => e mapExpr replace(renames)
    }
  }
}


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffre.passes

import firrtl._
import firrtl.ir._
import firrtl.passes.PassException
import firrtl.passes.wiring.{SinkAnnotation, SourceAnnotation}
import firrtl.annotations.{ComponentName, ModuleName, CircuitName, SingleTargetAnnotation, Annotation}
import chiffre.{InjectorInfo, FaultyComponent}
import chiffre.scan.{ScanChain, JsonProtocol}

import scala.collection.mutable
import scala.collection.immutable.ListMap
import java.io.{File, FileWriter}

case class ScanChainException(msg: String) extends PassException(msg)

case class ScanChainInfo(
  masterScan: ComponentName,
  /* Everything here is keyed by the injector component */
  slaveIn: Map[ComponentName, ComponentName] = Map.empty,
  slaveOut: Map[ComponentName, ComponentName] = Map.empty,
  injectors: ListMap[ComponentName, ModuleName] = ListMap.empty,
  /* This is keyed by the injector module name */
  description: Map[ComponentName, InjectorInfo] = Map.empty) {
  // scalastyle:off line.size.limit
  def serialize(tab: String): String =
    s"""|${tab}master:
        |${tab}  scan: ${masterScan}
        |${tab}slaves:

        |${tab}description:
        |${description.map{case(k,v)=>s"${tab}  ${k.name}: $v"}.mkString("\n")}"""
      .stripMargin
  // scalastyle:on line.size.limit

  def toFaultyComponent: Seq[FaultyComponent] = injectors
    .map{ case(c, m) => FaultyComponent(c.serialize, description(c)) }
    .toSeq
}

sealed trait ScanAnnos

case class ScanChainAnnotation(
  target: ComponentName,
  ctrl: String,
  dir: String,
  id: String,
  // [todo] remove key, is this used?
  key: Option[ComponentName]) extends SingleTargetAnnotation[ComponentName]
    with ScanAnnos {
  def duplicate(x: ComponentName): ScanChainAnnotation = this.copy(target = x)
}

case class ScanChainInjectorAnnotation(
  target: ComponentName,
  id: String,
  moduleName: String) extends SingleTargetAnnotation[ComponentName]
    with ScanAnnos {
  def duplicate(x: ComponentName): ScanChainInjectorAnnotation =
    this.copy(target = x)
}

case class ScanChainDescriptionAnnotation(
  target: ComponentName,
  id: String,
  d: InjectorInfo) extends SingleTargetAnnotation[ComponentName]
    with ScanAnnos {
  def duplicate(x: ComponentName): ScanChainDescriptionAnnotation =
    this.copy(target = x)
}

class ScanChainTransform extends Transform {
  def inputForm: CircuitForm = MidForm
  def outputForm: CircuitForm = HighForm

  // scalastyle:off cyclomatic.complexity
  def analyze(circuit: Circuit, annos: Seq[Annotation]): Map[String, ScanChainInfo] = {
    val s = mutable.HashMap[String, ScanChainInfo]()
    annos.foreach {
      case ScanChainAnnotation(comp, ctrl, dir, id, key) => (ctrl, dir) match {
        case ("master", "scan") => s(id) = ScanChainInfo(masterScan = comp)
        case _ =>
      }
      case _ =>
    }
    def exceptionIfUnknownId(id: String): Unit = if (!s.contains(id)) {
      throw new ScanChainException(
        s"No known scan chain master '$id' (Did you misspell it? Did you not include an injector?)") }
    annos.foreach {
      case ScanChainAnnotation(comp, ctrl, dir, id, key) =>
        exceptionIfUnknownId(id)
        s(id) =
          (ctrl, dir) match {
            case ("slave", "in")  => s(id).copy(slaveIn  = s(id).slaveIn  ++ Map(key.get -> comp))
            case ("slave", "out") => s(id).copy(slaveOut = s(id).slaveOut ++ Map(key.get -> comp))
            case _                => s(id)
          }
      case ScanChainInjectorAnnotation(comp, id, mod) =>
        exceptionIfUnknownId(id)
        s(id) = s(id).copy(injectors = s(id).injectors ++ Map(comp -> ModuleName(mod, comp.module.circuit)))
      case _ =>
    }
    annos.foreach {
      case ScanChainDescriptionAnnotation(comp, id, d) =>
        exceptionIfUnknownId(id)
        s(id) = s(id).copy(description = s(id).description ++ Map(comp -> d))
      case _ =>
    }
    s.toMap
  } // scalastyle:on cyclomatic.complexity

  /** Run the transform
    *
    * @param state a circuit
    * @todo order the scan chain based on distance
    */
  def execute(state: CircuitState): CircuitState = {
    val targetDir = new File(state.annotations.collectFirst{ case a: TargetDirAnnotation => a.directory }.getOrElse("."))
    val myAnnos = state.annotations.collect{ case a: ScanAnnos => a }
    myAnnos match {
      case Nil => state
      case p =>
        // s is a map of scan chains indexed by scan chain id
        val s: Map[String, ScanChainInfo] = analyze(state.circuit, p)

        s.foreach{ case (k, v) => logger.info(
                    s"""|[info] scan chain:
                        |[info]   name: ${k}
                        |${v.serialize("[info]   ")}""".stripMargin) }

        val sc = s.flatMap{ case(k, v) => Map(k -> v.toFaultyComponent) }
        if (!targetDir.exists()) { targetDir.mkdirs() }
        val jsonFile = new FileWriter(targetDir + "/scan-chain.json")
        jsonFile.write(JsonProtocol.serialize(sc))
        jsonFile.close()

        val ax = s.foldLeft(Seq[Annotation]()){ case (a, (name, v)) =>
          val masterIn = v.masterScan.copy(name=v.masterScan.name + ".in")
          val masterOut = v.masterScan.copy(name=v.masterScan.name + ".out")
          val masterClk = v.masterScan.copy(name=v.masterScan.name + ".clk")
          val masterEn = v.masterScan.copy(name=v.masterScan.name + ".en")

          val masterAnnotations = Seq(
            SourceAnnotation(masterClk, "scan_clk"),
            SourceAnnotation(masterEn, "scan_en"))

          // [todo] This is not deterministic
          val chain: Seq[ComponentName] = masterOut +: v.injectors
            .flatMap{ case (k, _) => Seq(v.slaveIn(k), v.slaveOut(k)) }
            .toSeq :+ masterIn

          val annotations = chain
            .grouped(2).zipWithIndex
            .flatMap{ case (Seq(l, r), i) =>
              Seq(SourceAnnotation(l, s"scan_${name}_$i"),
                  SinkAnnotation(r, s"scan_${name}_$i")) }

          a ++ masterAnnotations ++ annotations
        }

        state.copy(annotations = ((state.annotations.toSeq ++ ax).toSet -- myAnnos.toSet).toSeq)
    }
  }
}

|${injectors.map{ case (k, v) => s"${tab}  ${v.name}, ${k.module.name}, ${slaveIn(k).module.name}.${slaveIn(k).name}, ${slaveOut(k).module.name}.${slaveOut(k).name}"}.mkString("\n")}


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffreTests

import chisel3._
import chisel3.iotesters.ChiselFlatSpec
import chisel3.stage.ChiselStage
import chiffre.{InjectorInfo, ChiffreController, ChiffreInjector, ChiffreInjectee}
import chiffre.passes.{ScanChainAnnotation, ScanChainDescriptionAnnotation, FaultInjectionAnnotation}
import chiffre.inject.{Injector, LfsrInjector32}

class DummyController extends Module with ChiffreController {
  val io = IO(new Bundle{})
  val scanId = "dummy"
}

class DummyInjector extends Injector(1) with ChiffreInjector {
  val scanId = "dummy"
  val info = EmptyInjectorInfo
}

class DummyInjectee extends Module with ChiffreInjectee {
  val io = IO(new Bundle{})
  val x = Reg(UInt(1.W))
  isFaulty(x, "dummy", classOf[LfsrInjector32])
}

class InstrumentationSpec extends ChiselFlatSpec {
  behavior of "ChiffreController annotation"

  it should "emit a ScanChainAnnotation" in {
    val circuit = ChiselStage.elaborate(new DummyController)
    circuit.annotations.map(_.toFirrtl).collect{ case a: ScanChainAnnotation => a }.size should be (1)
  }

  behavior of "Chiffree Injectee annotation"

  it should "emit an annotation" in {
    val circuit = ChiselStage.elaborate(new DummyInjectee)


// Copyright 2018 IBM
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package chiffreTests.inject

import chiffre.inject.{StuckAtInjector, StuckAtInjectorInfo, Mask, StuckAt}
import chisel3.iotesters.{ChiselFlatSpec, Driver}

class StuckAtTester(dut: StuckAtInjector) extends InjectorTester(dut) {
  require(dut.bitWidth == 8)
  val mask = BigInt("11110000", 2)
  val stuckAt = BigInt("00110000", 2)
  val in = BigInt("01010101", 2)
  val faulty = BigInt("00110101", 2)

  dut.info.fields.foreach {
    case m: Mask => m.bind(mask)
    case s: StuckAt => s.bind(stuckAt)
  }

  poke(dut.io.scan.en, 0)
  poke(dut.io.in, in)
  load(dut.info)
  poke(dut.io.scan.en, 1)
  step(1)
  poke(dut.io.scan.en, 0)
  step(1)

  val fault = peek(dut.io.out)
  assert(fault == faulty, s"Expected to see $faulty, but got $fault")

  poke(dut.io.scan.en, 1)
  step(1)
  poke(dut.io.scan.en, 0)
  val noFault = peek(dut.io.out)
  assert(noFault == in, s"Expected to see $in, but got $noFault")
}

class StuckAtInjectorSpec extends ChiselFlatSpec {
  behavior of "StuckAtInjectorInfo"

  it should "be the expected width" in {
    val x = StuckAtInjectorInfo(1337)
    x.width should be (1337 * 2)
  }

  behavior of "StuckAtInjector"

  it should "be able to cycle a configuration" in {
    Driver(() => new StuckAtInjector(13)) { dut => new InjectorCycleTester(dut) }
  }

  it should "make a signal stuck when enabled" in {

Driver(() => new StuckAtInjector(8)) { dut => new StuckAtTester(dut) }


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer

import java.io.{File, PrintWriter}

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration.Duration
import scala.concurrent.{blocking, Await, Future, TimeoutException}
import scala.sys.process._

class RenderSvg(dotProgram: String = "dot", dotTimeOut: Int = 7) {
  def render(fileName: String): Unit = {
    if (dotProgram != "none") {
      if (fileName.isEmpty) {
        println(s"Tried to call render program $dotProgram without a filename")
      } else if (!new File(fileName).exists()) {
        println(s"Tried to call render program $dotProgram on non existent file $fileName")
      } else {
        val dotProcessString = s"$dotProgram -Tsvg -O $fileName"
        val process = Process(dotProcessString).run()
        val processFuture = Future(blocking(process.exitValue()))
        try {
          Await.result(processFuture, Duration(dotTimeOut, "sec"))
        } catch {
          case _: TimeoutException =>
            println(s"Rendering timed out after $dotTimeOut seconds on $fileName with command $dotProcessString")
            println(s"You can try increasing with the --dot-timeout seconds flag")
            process.destroy()
            val printWriter = new PrintWriter(new File(fileName + ".svg"))
            printWriter.print(s"""

                                 | -->
                                 |<!-- Title: TopLevel Pages: 1 -->
                                 |<svg width="351pt" height="44pt"
                                 | viewBox="0.00 0.00 351.26 44.00" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
                                 |<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 40)">
                                 |<title>TopLevel</title>
                                 |<polygon fill="#ffffff" stroke="transparent" points="-4,4 -4,-40 347.2637,-40 347.2637,4 -4,4"/>
                                 |<!-- Sorry, Rendering timed out on this file, Use Back to return -->
                                 |<g id="node1" class="node">
                                 |<title>
                                 |Sorry, Rendering timed out on this file, Use Back to return.
                                 |You can try increasing the timeout with the --dot-timeout seconds flag
                                 |</title>
                                 |<polygon fill="none" stroke="#000000" points="343.3956,-36 -.1319,-36 -.1319,0 343.3956,0 343.3956,-36"/>
                                 |<text text-anchor="middle" x="171.6318" y="-13.8" font-family="Times,serif" font-size="14.00" fill="#000000">Sorry, Rendering timed out on this file, Use Back to return</text>
                                 |</g>
                                 |</g>
                                 |</svg>
              """.stripMargin)
            printWriter.close()

        }
      }
    }
  }

}

                                 |<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
                                 | "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer

import firrtl.CompilerUtils.getLoweringTransforms
import firrtl.Mappers._
import firrtl.PrimOps.Dshl
import firrtl._
import firrtl.ir._
import firrtl.options.{Dependency, Phase}
import firrtl.stage.{FirrtlCircuitAnnotation, Forms}
import firrtl.transforms.BlackBoxSourceHelper

/**
  * Use these lowering transforms to prepare circuit for compiling
  */
class ToLoFirrtl extends Phase {
  private val targets = Forms.LowFormOptimized ++ Seq(
    Dependency(passes.RemoveValidIf),
    Dependency(passes.memlib.VerilogMemDelays),
    Dependency(passes.SplitExpressions),
    Dependency[firrtl.transforms.LegalizeAndReductionsTransform],
    Dependency[firrtl.transforms.ConstantPropagation],
    Dependency[firrtl.transforms.CombineCats],
    Dependency(passes.CommonSubexpressionElimination),
    Dependency[firrtl.transforms.DeadCodeElimination]
  )

  private def compiler = new firrtl.stage.transforms.Compiler(targets, currentState = Nil)
  private val transforms = compiler.flattenedTransformOrder

  override def transform(annotationSeq: AnnotationSeq): AnnotationSeq = {

    annotationSeq.flatMap {
      case FirrtlCircuitAnnotation(circuit) =>
        val state = CircuitState(circuit, annotationSeq)
        val newState = transforms.foldLeft(state) {
          case (prevState, transform) => transform.runTransform(prevState)
        }
        Some(FirrtlCircuitAnnotation(newState.circuit))
      case other =>
        Some(other)
    }
  }
}

/**
  *  Workaround for https://github.com/freechipsproject/firrtl/issues/498 from @jackkoenig
  */
class FixupOps extends Transform with DependencyAPIMigration {
  override def prerequisites = Seq.empty

  override def optionalPrerequisites = Seq.empty

  override def optionalPrerequisiteOf = Seq.empty

  override def invalidates(a: Transform) = false

  private def onExpr(expr: Expression): Expression =
    expr.map(onExpr) match {
      case prim @ DoPrim(Dshlw, _, _, _) => prim.copy(op = Dshl)
      case other                         => other
    }
  private def onStmt(stmt: Statement): Statement = stmt.map(onStmt).map(onExpr)

  def execute(state: CircuitState): CircuitState = {
    state.copy(circuit = state.circuit.map(onMod))
  }
}

private def onMod(mod:   DefModule): DefModule = mod.map(onStmt)


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.dotnodes

case class LiteralNode(name: String, value: BigInt, parentOpt: Option[DotNode]) extends DotNode {

  def render: String = {
    s"""$absoluteName [shape="circle" style="filled" BGCOLOR="#C0C0C0" label="$value"]
     """.stripMargin
  }


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.dotnodes

object PrimOpNode {
  /* the pseudoHash is necessary because case classes with identical args have identical hashes */
  var pseudoHash: Long = 0
  def hash: Long = {
    pseudoHash += 1
    pseudoHash
  }

  val BlackDot = "●"
}

case class BinaryOpNode(
  name:         String,
  parentOpt:    Option[DotNode],
  arg0ValueOpt: Option[String],
  arg1ValueOpt: Option[String])
    extends DotNode {

  def in1:                   String = s"$absoluteName:in1"
  def in2:                   String = s"$absoluteName:in2"

override val asRhs:        String = s"$absoluteName:out"

  def render: String = {
    s"""
       |$absoluteName [shape = "plaintext" label=<
       |

       |

       |    ${arg0ValueOpt.getOrElse(PrimOpNode.BlackDot)}
       |     $name
       |    ${PrimOpNode.BlackDot}
       |

       |

       |    ${arg1ValueOpt.getOrElse(PrimOpNode.BlackDot)}
       |

       |
>];
       """.stripMargin
  }
}

case class UnaryOpNode(name: String, parentOpt: Option[DotNode]) extends DotNode {
  def in1:                   String = s"$absoluteName:in1"
  override val absoluteName: String = s"op_${name}_${PrimOpNode.hash}"
  override val asRhs:        String = s"$absoluteName:out"

  def render: String = {
    s"""
       |$absoluteName [shape = "plaintext" label=<
       |

       |

       |    ●
       |     $name
       |    ●
       |

       |
>];
       """.stripMargin
  }
}

case class OneArgOneParamOpNode(name: String, parentOpt: Option[DotNode], value: BigInt) extends DotNode {
  def in1:                   String = s"$absoluteName:in1"
  override val absoluteName: String = s"op_${name}_${PrimOpNode.hash}"
  override val asRhs:        String = s"$absoluteName:out"

  def render: String = {
    s"""
       |$absoluteName [shape = "plaintext" label=<
       |

       |

       |    ●
       |     $name
       |    ●
       |

       |

       |    $value
       |

       |
>];
       """.stripMargin
  }
}
case class OneArgTwoParamOpNode(
  name:      String,
  parentOpt: Option[DotNode],
  value1:    BigInt,
  value2:    BigInt)
    extends DotNode {
  def in1:                   String = s"$absoluteName:in1"
  override val absoluteName: String = s"op_${name}_${PrimOpNode.hash}"
  override val asRhs:        String = s"$absoluteName:out"

  def render: String = {
    s"""
       |$absoluteName [shape = "plaintext" label=<
       |

       |

       |    ●
       |     $name
       |    ●
       |

       |

       |    ($value1, $value2)
       |

       |
>];
       """.stripMargin
  }
}


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.stage

import dotvisualizer.stage.phase.{
  CheckPhase,
  GenerateDotFilePhase,
  GetFirrtlCircuitPhase,
  OptionallyBuildTargetDirPhase
}
import firrtl.options.phases.DeletedWrapper
import firrtl.options.{Dependency, Phase, PhaseManager}

class DiagrammerPhase extends PhaseManager(DiagrammerPhase.targets) {

  override val wrappers = Seq((a: Phase) => DeletedWrapper(a))

}

object DiagrammerPhase {

  val targets: Seq[PhaseManager.PhaseDependency] = Seq(
    Dependency[CheckPhase],
    Dependency[GetFirrtlCircuitPhase],


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.stage.phase

import java.io.{File, PrintWriter}

import dotvisualizer.RenderSvg
import dotvisualizer.stage.{
  DotTimeoutSecondsAnnotation,
  JustTopLevelAnnotation,
  OpenCommandAnnotation,
  SetRenderProgramAnnotation
}
import dotvisualizer.transforms.{MakeDiagramGroup, ModuleLevelDiagrammer}
import firrtl.options.{Dependency, Phase, TargetDirAnnotation}
import firrtl.stage.FirrtlCircuitAnnotation
import firrtl.{AnnotationSeq, CircuitState}

import scala.sys.process._

class GenerateDotFilePhase extends Phase {

  override def optionalPrerequisites = Seq.empty

  override def optionalPrerequisiteOf = Seq.empty

  override def invalidates(a: Phase) = false

  /**
    * Make a simple css file that controls highlighting
    *
    * @param targetDir where to put the css
    */
  def addCss(targetDir: String): Unit = {
    val file = new File(s"$targetDir/styles.css")
    val printWriter = new PrintWriter(file)
    printWriter.println(
      """
        |.edge:hover * {
        |  stroke: #ff0000;
        |}
        |.edge:hover polygon {
        |  fill: #ff0000;
        |}
      """.stripMargin
    )
    printWriter.close()
  }

  /**
    * Open an svg file using the open program
    *
    * @param fileName    file to be opened
    * @param openProgram program to use
    */
  def show(fileName: String, openProgram: String): Unit = {
    if (openProgram.nonEmpty && openProgram != "none") {
      val openProcessString = s"$openProgram $fileName.svg"
      openProcessString.!!
    } else {
      println(s"""There is no program identified which will render the svg files.""")
      println(s"""The file to start with is $fileName.svg, open it in the appropriate viewer""")
      println(s"""Specific module views should be in the same directory as $fileName.svg""")
    }
  }

  override def transform(annotationSeq: AnnotationSeq): AnnotationSeq = {
    val targetDir = annotationSeq.collectFirst { case TargetDirAnnotation(dir) => dir }.get
    val firrtlCircuit = annotationSeq.collectFirst { case FirrtlCircuitAnnotation(circuit) => circuit }.get
    val circuitState = CircuitState(firrtlCircuit, annotationSeq)

    addCss(targetDir)

    val dotProgram = annotationSeq.collectFirst {
      case SetRenderProgramAnnotation(program) => program
    }.getOrElse("dot")

    val dotTimeOut = annotationSeq.collectFirst {
      case DotTimeoutSecondsAnnotation(secs) => secs
    }.getOrElse(7)

    val renderer = new RenderSvg(dotProgram, dotTimeOut)

    if (annotationSeq.exists { case JustTopLevelAnnotation => true; case _ => false }) {
      val justTopLevelTransform = new ModuleLevelDiagrammer(renderer)
      justTopLevelTransform.execute(circuitState)
    } else {
      val transform = new MakeDiagramGroup(renderer)
      transform.execute(circuitState)
    }

    val fileName = s"$targetDir/${circuitState.circuit.main}_hierarchy.dot"
    val openProgram = annotationSeq.collectFirst {
      case OpenCommandAnnotation(program) => program
    }.getOrElse("open")

    show(fileName, openProgram)

    annotationSeq
  }
}

override def prerequisites = Seq(Dependency[OptionallyBuildTargetDirPhase])


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.stage.phase

import java.io.File

import dotvisualizer.stage.DiagrammerException
import firrtl.options.{Dependency, Phase, TargetDirAnnotation}
import firrtl.stage.FirrtlCircuitAnnotation
import firrtl.{AnnotationSeq, FileUtils}

class OptionallyBuildTargetDirPhase extends Phase {
  override def prerequisites = Seq(Dependency[GetFirrtlCircuitPhase])

  override def optionalPrerequisites = Seq.empty

  override def optionalPrerequisiteOf = Seq.empty

  override def invalidates(a: Phase) = false

  override def transform(annotationSeq: AnnotationSeq): AnnotationSeq = {
    val newAnnotations: AnnotationSeq = if (annotationSeq.exists(_.isInstanceOf[TargetDirAnnotation])) {
      annotationSeq
    } else {
      val circuit = annotationSeq.collectFirst { case FirrtlCircuitAnnotation(circuit) => circuit }.get
      val targetDir = s"test_run_dir/${circuit.main}"
      annotationSeq :+ TargetDirAnnotation(targetDir)
    }

    newAnnotations.foreach {
      case TargetDirAnnotation(targetDir) =>
        val targetDirFile = new File(targetDir)
        if (targetDirFile.exists()) {
          if (!targetDirFile.isDirectory) {

          }
        } else {
          FileUtils.makeDirectory(targetDir)
          if (!targetDirFile.exists()) {
<TARGET>
          }
        }
      case _ =>
    }
    newAnnotations
  }
}

throw new DiagrammerException(s"Error: Target dir ${targetDir} exists and is not a directory")


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.transforms

import dotvisualizer._
import dotvisualizer.stage.StartModuleNameAnnotation
import firrtl.options.TargetDirAnnotation
import firrtl.{CircuitForm, CircuitState, DependencyAPIMigration, LowForm, Transform}

import scala.collection.mutable

class MakeDiagramGroup(renderSvg: RenderSvg) extends Transform with DependencyAPIMigration {
  override def prerequisites = Seq.empty

  override def invalidates(a: Transform) = false

  /**
    * Creates a series of diagrams starting with the startModule and continuing
    * through all descendant sub-modules.
    * @param state the state to be diagrammed
    * @return
    */

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

    val targetDir = state.annotations.collectFirst { case TargetDirAnnotation(dir) => dir }.getOrElse {
      s"test_run_dir/${state.circuit.main}/"
    }

    val startModule = state.annotations.collectFirst {
      case StartModuleNameAnnotation(moduleName) => moduleName
    }.getOrElse(state.circuit.main)

    val queue = new mutable.Queue[String]()
    val modulesSeen = new mutable.HashSet[String]()

    val pass_remove_gen = new RemoveTempWires()
    var cleanedState = pass_remove_gen.execute(state)

    val pass_top_level = new ModuleLevelDiagrammer(renderSvg)
    pass_top_level.execute(cleanedState)

    queue += startModule // set top level of diagram tree

    while (queue.nonEmpty) {
      val moduleName = queue.dequeue()
      if (!modulesSeen.contains(moduleName)) {

        val updatedAnnotations = {
          state.annotations.filterNot { x =>
            x.isInstanceOf[StartModuleNameAnnotation]
          } :+ StartModuleNameAnnotation(moduleName)
        }
        val stateToDiagram = CircuitState(cleanedState.circuit, state.form, updatedAnnotations)

        val pass = new MakeOneDiagram(renderSvg)
        pass.execute(stateToDiagram)

        queue ++= pass.subModulesFound.map(module => module.name)
        renderSvg.render(s"$targetDir/$moduleName.dot")
      }
      modulesSeen += moduleName
    }

    // we return the original state, all transform work is just in the interest of diagramming
    state
  }
}

override def optionalPrerequisites = Seq.empty


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer.transforms

import dotvisualizer.ToLoFirrtl
import dotvisualizer.stage.DiagrammerException
import firrtl.ir._
import firrtl.stage.FirrtlCircuitAnnotation
import firrtl.{
  CircuitForm,
  CircuitState,
  DependencyAPIMigration,
  FileUtils,
  LowForm,
  Parser,
  Transform,
  WRef,
  WSubField,
  WSubIndex
}

import scala.collection.mutable

class RemoveTempWires extends Transform with DependencyAPIMigration {
  override def prerequisites = Seq.empty

  override def optionalPrerequisites = Seq.empty

  override def optionalPrerequisiteOf = Seq.empty

  override def invalidates(a: Transform) = false

  /**
    * Foreach Module in a firrtl circuit
    *   Find all the DefNodes with temp names and render their expression
    *   Remove all the found dev nodes
    *   recursively replace their references with their associated expression
    * @param state to be altered
    * @return
    */

  def execute(state: CircuitState): CircuitState = {

    val c = state.circuit

    /**
      * removes all references to temp wires in module
      * @param module the module to be altered
      * @return
      */
    def removeTempWiresFromModule(module: Module): Module = {

      val toRemove = new mutable.HashMap[String, Expression]()

      /**
        * Saves reference to the expression associated
        * with a temp wire associated with a Node statement
        * @param s statement to be checked
        */
      def collectTempExpressions(s: Statement): Unit = s match {
        case block: Block =>
          block.stmts.foreach { substatement =>
            collectTempExpressions(substatement)
          }

        case node: DefNode =>
          if (node.name.startsWith(RemoveTempWires.GenPrefix) || node.name.startsWith(RemoveTempWires.TempPrefix)) {

            if (toRemove.contains(node.name)) {
              println(s"Houston we have a problem, ${node.name} already marked for removal")
            }
            toRemove(node.name) = node.value
            None
          }
        case _ => //do nothing
      }

      /**
        * recursively find any references to temp wires in the expression and replace the
        * references with the associated expression
        * @param e expression to be altered
        * @return
        */
      def removeGen(e: Expression): Expression = {
        e match {
          case wire: WRef =>
            if (
              (wire.name.startsWith(RemoveTempWires.GenPrefix) ||
              wire.name.startsWith(RemoveTempWires.TempPrefix)) && toRemove.contains(wire.name)
            ) {
              val new_node = toRemove(wire.name)
              removeGen(new_node)
            } else {
              wire
            }
          case wire: WSubField =>
            if (
              (wire.name.startsWith(RemoveTempWires.GenPrefix) ||
              wire.name.startsWith(RemoveTempWires.TempPrefix)) && toRemove.contains(wire.name)
            ) {
              val new_node = toRemove(wire.name)
              removeGen(new_node)
            } else {
              wire
            }
          case wire: WSubIndex =>
            wire.mapExpr(removeGen)
          case ee => ee.mapExpr(removeGen)
        }
      }

      /**
        * Removes node definition statements for temp wires
        * @param s statement to be altered
        * @return
        */
      def removeGenStatement(s: Statement): Option[Statement] = {
        s match {
          case block: Block =>

            if (node.name.startsWith(RemoveTempWires.GenPrefix) || node.name.startsWith(RemoveTempWires.TempPrefix)) {
              None
            } else {
              Some(node.mapExpr(removeGen))
            }
          case other: Statement =>
            Some(other.mapExpr(removeGen))
          case _ => Some(s) //do nothing
        }
      }

      collectTempExpressions(module.body)
      val moduleWithTempExpressionsRemmoved = removeGenStatement(module.body)
      module.copy(body = moduleWithTempExpressionsRemmoved.get)
    }

    val newModules = c.modules.map {
      case m: Module => removeTempWiresFromModule(m)
      case otherMod => otherMod
    }

    state.copy(circuit = Circuit(c.info, newModules, c.main))
  }
}

object RemoveTempWires {
  val GenPrefix = "_GEN_"
  val TempPrefix = "_T_"

  def main(args: Array[String]): Unit = {
    args.headOption match {
      case Some(fileName) =>
        val firrtlSource = FileUtils.getLines(fileName).mkString("\n")
        val annos = (new ToLoFirrtl).transform(Seq(FirrtlCircuitAnnotation(Parser.parse(firrtlSource))))
        val firrtl = annos.collectFirst { case FirrtlCircuitAnnotation(c) => c }.getOrElse {
          throw new DiagrammerException("Error: could not process supplied firrtl")
        }

        val grapher = new RemoveTempWires

        val newState = grapher.execute(CircuitState(firrtl, Seq.empty))
        println(s"${newState.circuit.serialize}")

      case _ =>
    }
  }
}

            val result = Some(Block(block.stmts.flatMap { substatement =>
              removeGenStatement(substatement)
            }))


// SPDX-License-Identifier: Apache-2.0

package dotvisualizer

import chisel3._
import chisel3.stage.ChiselGeneratorAnnotation
import dotvisualizer.stage._
import firrtl.FileUtils
import firrtl.annotations.Annotation
import firrtl.options.TargetDirAnnotation
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.should.Matchers

import java.io.{ByteArrayOutputStream, PrintStream}

class HasPrintf extends Module {
  val in = IO(Input(Bool()))
  val out = IO(Output(Bool()))
  out := in
  printf("in %d, out %d\n", in, out)
}
class PrintfSpec extends AnyFreeSpec with Matchers {

  "printfs can now be rendered" - {

    def makeDotFile(showPrintfs: Boolean): String = {
      val targetDir = s"test_run_dir/has_printf_$showPrintfs"
      val annos = Seq(
        TargetDirAnnotation(targetDir),
        ChiselGeneratorAnnotation(() => new HasPrintf),
        RankElementsAnnotation,
        RankDirAnnotation("TB"),
        OpenCommandAnnotation("")

            else { Seq.empty[Annotation] })
      val outputBuf = new ByteArrayOutputStream()
      Console.withOut(new PrintStream(outputBuf)) {
        (new DiagrammerStage).transform(annos)
      }
      val output = outputBuf.toString
      output should include(s"creating dot file test_run_dir/has_printf_$showPrintfs/HasPrintf.dot")

      FileUtils.getText(s"$targetDir/HasPrintf.dot")
    }

    "showPrintfs=true will render printfs in dot file" in {
      val dotText = makeDotFile(showPrintfs = true)

      dotText should include("struct_cluster_HasPrintf_printf_")
      dotText should include("""printf("in %d, out %d\n")""")
    }

    "default behavior will not render printfs in dot file" in {
      val dotText = makeDotFile(showPrintfs = false)

      dotText.contains("struct_cluster_HasPrintf_printf_") should be(false)
      dotText.contains("""printf("in %d, out %d\n")""") should be(false)
    }
  }
}

) ++ (if (showPrintfs) { Seq(ShowPrintfsAnnotation) }


// See LICENSE.BU for license details.
// See LICENSE.IBM for license details.

package dana

import chisel3._
import chisel3.util._
import cde._

import dana.abi._

class ActivationFunctionReq(implicit p: Parameters) extends DanaBundle()(p) {
  val decimal            = UInt(decimalPointWidth.W)
  val steepness          = UInt(steepnessWidth.W)
  // Differentiate activation function requests vs. error function
  // requests
  val activationFunction = UInt(log2Up(18).W) // [TODO] fragile
  val in                 = SInt(elementWidth.W)
}

class ActivationFunctionReqLearn(implicit p: Parameters)
    extends ActivationFunctionReq()(p) {
  val afType             = UInt(log2Up(2).W) // [TODO] fragile
  val errorFunction      = UInt(p(GlobalInfo).error_function.W)
}

class ActivationFunctionResp(implicit p: Parameters) extends DanaBundle()(p) {
  val out                = SInt(elementWidth.W)
}

class ActivationFunctionInterface(implicit p: Parameters) extends DanaBundle()(p) {
  val req = Flipped(Valid(new ActivationFunctionReq))
  val resp = Valid(new ActivationFunctionResp)
}

class ActivationFunctionInterfaceLearn(implicit p: Parameters)
    extends ActivationFunctionInterface()(p) {
  override val req = Flipped(Valid(new ActivationFunctionReqLearn))
}

class DSP(implicit p: Parameters) extends DanaModule()(p) {
  val io = IO(new Bundle {
    val a = Input(SInt(elementWidth.W))
    val b = Input(SInt(elementWidth.W))
    val c = Input(UInt(elementWidth.W))
    val d = Output(SInt(elementWidth.W))
  })
  io.d := (((io.a * io.b) >> io.c)(elementWidth - 1, 0)).asSInt
}

class ActivationFunction(id: Int = 0)(implicit p: Parameters) extends DanaModule()(p) {
  lazy val io = (new ActivationFunctionInterface)
  override val printfSigil = "dana.PE[" + id + "]: "

  // Temporary values
  val inD0 = Wire(SInt(elementWidth.W))
  val decimal = Wire(UInt())
  val out = Reg(init = 0.S(elementWidth.W))

  val _xmin      = -1420910720.S // -54B16080
  val _x1        =  -790391808.S // -2F1C6C00
  val _x2        =  -294906496.S // -1193EA80
  val _x3        =   294906496.S //  1193EA80
  val _x4        =   790391808.S //  2F1C6C00
  val _xmax      =  1420910720.S //  54b16080
  val _sigy0     =    10737418.S //  00A3D70A
  val _sigy1     =   107374184.S //  06666668
  val _sigy2     =   536870912.S //  20000000
  val _sigy3     =  1610612736.S //  60000000
  val _sigy4     =  2040109440.S //  79999980
  val _sigy5     =  2136746240.S //  7F5C2900
  val _slope1    =   164567525.S //  09CF19E5
  val _slope2    =   930741212.S //  3779FBDC
  val _slope3    =  1954723819.S //  7482B7EB
  val _slope4    =   930741280.S //  3779FC20
  val _slope5    =   164567500.S //  09CF19CC
  val _symy0     = -2126008831.S // -7EB851FF
  val _symy1     = -1932735232.S // -73333300
  val _symy2     = -1073741824.S // -40000000
  val _symy3     =  1073741824.S //  40000000
  val _symy4     =  1932735232.S //  73333300
  val _symy5     =  2126008831.S //  7EB851FF
  // Chisel (or Scala) has a problem with creating UInts that use the
  // full bit width when specifying the value as an integer, e.g., the
  // string assignment that will be interpreted as hex works, but the
  // below assignment to the full width sslope3 does not work
  val _sslope1 = "h139E343F".U   // 139E343F
  val _sslope2 = "h6EF3F751".U  // 6EF3F751
  val _sslope3 = "hE9056FD7".U  // E9056FD7
  val _sslope4 = "h6EF3F751".U  // 6EF3F751
  val _sslope5 = "h139E343F".U   // 139E343F

  val xmin   = _xmin >>  (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val x1     = _x1 >>    (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val x2     = _x2 >>    (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val x3     = _x3 >>    (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val x4     = _x4 >>    (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val xmax   = _xmax >>  (29.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy0  = _sigy0 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy1  = _sigy1 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy2  = _sigy2 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy3  = _sigy3 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy4  = _sigy4 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sigy5  = _sigy5 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val slope1 = _slope1 >>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val slope2 = _slope2 >>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val slope3 = _slope3 >>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val slope4 = _slope4 >>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val slope5 = _slope5 >>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy0  = _symy0 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy1  = _symy1 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy2  = _symy2 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy3  = _symy3 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy4  = _symy4 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val symy5  = _symy5 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val sslope1= _sslope1>>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val sslope2= _sslope2>>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val sslope3= _sslope3>>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val sslope4= _sslope4>>(32.U-io.req.bits.decimal-decimalPointOffset.U)
  val sslope5= _sslope5>>(32.U-io.req.bits.decimal-decimalPointOffset.U)

  decimal := decimalPointOffset.U((decimalPointWidth + log2Up(decimalPointOffset)).W) + io.req.bits.decimal

  // DSP Unit
  val dsp = Module(new DSP).io
  def DSP(a: SInt, b: SInt, c: UInt) {
    dsp.a := a
    dsp.b := b
    dsp.c := c
  }

  // All activation functions currently take two cycles, so the output
  // valid signal is delayed by two cycles.
  val ioVal_d0 = Reg(next = io.req.valid)
  val ioVal_d1 = Reg(next = ioVal_d0)
  io.resp.bits.out := out
  io.resp.valid := ioVal_d1
  val dataIn = RegNext(io.req.bits.in)

  when (ioVal_d1) {
    printfInfo("af(0x%x) = 0x%x\n", dataIn, out)
  }

  def applySteepness(x: SInt, steepness: UInt): SInt = {
    val tmp = Wire(SInt())
    when (steepness < steepnessOffset.U) {
      tmp := x >> (steepnessOffset.U - steepness)
    } .elsewhen (steepness === steepnessOffset.U) {
      tmp := x
    } .otherwise {
      tmp := x << (steepness - steepnessOffset.U)
    }
    tmp
  }

  val one          = 1.S(elementWidth.W) << decimal
  val negOne       = -1.S(elementWidth.W) << decimal
  val seventeen    = 17.S(elementWidth.W) << decimal
  val negSeventeen = -17.S(elementWidth.W) << decimal
  val offsetX      = Wire(SInt(elementWidth.W))
  val offsetSigY   = Wire(SInt(elementWidth.W))
  val offsetSymY   = Wire(SInt(elementWidth.W))
  val slopeSig     = Wire(SInt(elementWidth.W))
  val slopeSym     = Wire(SInt(elementWidth.W))
  when(inD0 < xmin) {
    offsetX    := 0.S
    offsetSigY := 0.S
    offsetSymY := negOne
    slopeSig   := 0.S
    slopeSym   := 0.S
  } .elsewhen((xmin <= inD0) && (inD0 < x1)) {
    offsetX    := xmin
    offsetSigY := sigy0
    offsetSymY := symy0
    slopeSig   := slope1
    slopeSym   := sslope1.asSInt
  } .elsewhen((x1 <= inD0) && (inD0 < x2)) {
    offsetX    := x1
    offsetSigY := sigy1
    offsetSymY := symy1
    slopeSig   := slope2
    slopeSym   := sslope2.asSInt
  } .elsewhen((x2 <= inD0) && (inD0 < x3)) {
    offsetX    := x2
    offsetSigY := sigy2
    offsetSymY := symy2
    slopeSig   := slope3
    slopeSym   := sslope3.asSInt
  } .elsewhen((x3 <= inD0) && (inD0 < x4)) {
    offsetX    := x3
    offsetSigY := sigy3
    offsetSymY := symy3
    slopeSig   := slope4
    slopeSym   := sslope4.asSInt
  } .elsewhen((xmin <= inD0) && (inD0 < xmax)) {
    offsetX    := x4
    offsetSigY := sigy4
    offsetSymY := symy4
    slopeSig   := slope5
    slopeSym   := sslope5.asSInt
  } .otherwise {
    offsetX    := 0.S
    offsetSigY := one
    offsetSymY := one
    slopeSig   := 0.S
    slopeSym   := 0.S
  }

  // [TODO] You can probably remove this---by default `out` gets a
  // garbage value (a very big integer) which should be visibile in
  // the output
  DSP(slopeSym, inD0-offsetX, decimal)
  out := dsp.d + offsetSymY
  inD0 := applySteepness(dataIn, io.req.bits.steepness)
  // FANN_LINEAR
  switch (io.req.bits.activationFunction) {
    is (e_FANN_LINEAR) {
      out := inD0
    } // FANN_THRESHOLD
    is (e_FANN_THRESHOLD) {
      when (inD0 <= 0.S) { out := 0.S

    } // FANN_THRESHOLD_SYMMETRIC
    is (e_FANN_THRESHOLD_SYMMETRIC) {
      when (inD0 < 0.S) {
        out := negOne
      } .elsewhen(inD0 === 0.S) {
        out := 0.S
      } .otherwise {
        out := one }
    } // FANN_SIGMOID and STEPWISE
    is (e_FANN_SIGMOID) {
      DSP(slopeSig, inD0-offsetX, decimal)
      out := dsp.d + offsetSigY
    }
    is (e_FANN_SIGMOID_STEPWISE) {
      DSP(slopeSig, inD0-offsetX, decimal)
      out := dsp.d + offsetSigY
    }
  }
}

class ActivationFunctionLearn(id: Int = 0)(implicit p: Parameters)
    extends ActivationFunction(id)(p) {
  override lazy val io = (new ActivationFunctionInterfaceLearn)

  // atanh specific
  // Binary Point: 31
  val _atanh_x0 = -2147483433.S // -0.9999999
  val _atanh_x1 = -2138482164.S // -0.9958083573500538
  val _atanh_x2 = -1876440389.S // -0.8737856469351458
  val _atanh_x3 =  1876440379.S // 0.873785642278533
  val _atanh_x4 =  2138482162.S // 0.9958083564187312
  val _atanh_x5 =  2147483433.S // 0.9999999
  // Binary Point: 26
  val _atanh_y0 = -1128183406.S // -16.81124278204462
  val _atanh_y1 =  -413773911.S // -6.165711741243977
  val _atanh_y2 =  -181041891.S // -2.6977343972924133
  val _atanh_y3 =   181041888.S // 2.697734357912798
  val _atanh_y4 =   413773896.S // 6.165711518591778
  val _atanh_y5 =  1128183406.S // 16.81124278204462
  // Binary Point: 19
  val _atanh_s1 =  1331568027.U // 2539.7644566343943
  val _atanh_s2 =    14900660.U // 28.42075325289503
  val _atanh_s3 =     1618692.U // 3.087409817560523
  val _atanh_s4 =    14900659.U // 28.420750883269488
  val _atanh_s5 =  1331567759.U // 2539.763945441386
  // Binary Point: 31
  val atanh_x0 = _atanh_x0 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_x1 = _atanh_x1 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_x2 = _atanh_x2 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_x3 = _atanh_x3 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_x4 = _atanh_x4 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_x5 = _atanh_x5 >> (31.U-io.req.bits.decimal-decimalPointOffset.U)
  // Binary Point: 26
  val atanh_y0 = _atanh_y0 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_y1 = _atanh_y1 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_y2 = _atanh_y2 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_y3 = _atanh_y3 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_y4 = _atanh_y4 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_y5 = _atanh_y5 >> (26.U-io.req.bits.decimal-decimalPointOffset.U)
  // Binary Point: 19
  val atanh_s1 = _atanh_s1 >> (19.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_s2 = _atanh_s2 >> (19.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_s3 = _atanh_s3 >> (19.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_s4 = _atanh_s4 >> (19.U-io.req.bits.decimal-decimalPointOffset.U)
  val atanh_s5 = _atanh_s5 >> (19.U-io.req.bits.decimal-decimalPointOffset.U)

  // Atanh error function
  val atanhOffsetX = Wire(SInt(elementWidth.W))
  val atanhOffsetY = Wire(SInt(elementWidth.W))
  val atanhSlope   = Wire(SInt(elementWidth.W))
  when (dataIn < atanh_x0) {
    atanhOffsetX := 0.S
    atanhOffsetY := negSeventeen
    atanhSlope   := 0.S
  } .elsewhen ((atanh_x0 <= dataIn) && (dataIn < atanh_x1)) {
    atanhOffsetX := atanh_x0
    atanhOffsetY := atanh_y0
    atanhSlope   := atanh_s1.asSInt
  } .elsewhen ((atanh_x1 <= dataIn) && (dataIn < atanh_x2)) {
    atanhOffsetX := atanh_x1
    atanhOffsetY := atanh_y1
    atanhSlope   := atanh_s2.asSInt
  } .elsewhen ((atanh_x2 <= dataIn) && (dataIn < atanh_x3)) {
    atanhOffsetX := atanh_x2
    atanhOffsetY := atanh_y2
    atanhSlope   := atanh_s3.asSInt
  } .elsewhen ((atanh_x3 <= dataIn) && (dataIn < atanh_x4)) {
    atanhOffsetX := atanh_x3
    atanhOffsetY := atanh_y3
    atanhSlope   := atanh_s4.asSInt
  } .elsewhen ((atanh_x4 <= dataIn) && (dataIn < atanh_x5)) {
    atanhOffsetX := atanh_x4
    atanhOffsetY := atanh_y4
    atanhSlope   := atanh_s5.asSInt
  } .otherwise {
    atanhOffsetX := 0.S
    atanhOffsetY := seventeen
    atanhSlope   := 0.S
  }

  when (io.req.bits.afType === e_AF_DO_ERROR_FUNCTION) {
    switch (io.req.bits.errorFunction) {
      is (e_FANN_ERRORFUNC_LINEAR) {
        out := dataIn
      }
      is (e_FANN_ERRORFUNC_TANH) {
        DSP(atanhSlope, dataIn-atanhOffsetX, decimal)
        out := dsp.d + atanhOffsetY
      }
    }
  }
}


// See LICENSE.BU for license details.
// See LICENSE.IBM for license details.

package dana

import Chisel._
import cde._

import dana.abi._

class ProcessingElementReq(implicit p: Parameters) extends DanaBundle()(p) {
  val numWeights         = UInt(p(NeuronInfo).num_weights.W)
  val index              = UInt(log2Up(peTableNumEntries).W)
  val decimalPoint       = UInt(decimalPointWidth.W)
  val steepness          = UInt(steepnessWidth.W)
  val activationFunction = UInt(activationFunctionWidth.W)
  val iBlock             = Vec(elementsPerBlock, SInt(elementWidth.W))
  val wBlock             = Vec(elementsPerBlock, SInt(elementWidth.W))
  val bias               = SInt(elementWidth.W)
}

class ProcessingElementReqLearn(implicit p: Parameters)
    extends ProcessingElementReq()(p) {
  val errorFunction      = UInt(p(GlobalInfo).error_function.W)
  val learningRate       = UInt(elementWidth.W)
  val weightDecay        = SInt(elementWidth.W)
  val learnReg           = SInt(elementWidth.W)
  val stateLearn         = UInt(log2Up(8).W) // [TODO] fragile
  val inLast             = Bool()
  val inFirst            = Bool()
  val dw_in              = SInt(elementWidth.W)
  val tType              = UInt(log2Up(3).W)
}

class ProcessingElementResp(implicit p: Parameters) extends DanaBundle()(p) {
  val data               = SInt(elementWidth.W)
  val state              = UInt(log2Up(PE_states.size).W)
  val index              = UInt(log2Up(peTableNumEntries).W)
  val incWriteCount      = Bool()
}

class ProcessingElementRespLearn(implicit p: Parameters)
    extends ProcessingElementResp()(p) {
  val dataBlock          = Vec(elementsPerBlock, SInt(elementWidth.W))
  val error              = SInt(elementWidth.W)
  val resetWeightPtr     = Bool()
}

class ProcessingElementInterface(implicit p: Parameters) extends DanaBundle()(p) {
  // The Processing Element Interface consists of three main
  // components: requests from the PE Table (really kicks to do
  // something), responses to the PE Table, and semi-static data which
  // th PE Table manages and is used by the PEs for computation.
  val req                = Flipped(Decoupled(new ProcessingElementReq))
  val resp               = Decoupled(new ProcessingElementResp)
}

class ProcessingElementInterfaceLearn(implicit p: Parameters)
    extends ProcessingElementInterface()(p) {
  override val req       = Flipped(Decoupled(new ProcessingElementReqLearn))
  override val resp      = Decoupled(new ProcessingElementRespLearn)
}

class ProcessingElement(id: Int = 0)(implicit p: Parameters) extends DanaModule()(p) {
  override val printfSigil = "dana.PE[" + id + "]: "

  // Interface to the PE Table
  lazy val io = IO(new ProcessingElementInterface)

  // Activation Function module
  lazy val af = Module(new ActivationFunction(id))

  val index    = Reg(UInt(p(NeuronInfo).num_weights.W))
  val acc      = Reg(SInt(elementWidth.W))
  val dataOut  = Reg(SInt(elementWidth.W))
  val reqSent  = Reg(Bool())
  val eleIndex = index(log2Up(elementsPerBlock) - 1, 0)

  val state = Reg(UInt(log2Up(PE_states.size).W),
    init = PE_states('e_PE_UNALLOCATED))
  val nextState = Wire(UInt(log2Up(PE_states.size).W))
  nextState := state

  // Local state storage. Any and all of these are possible kludges
  // which could be implemented more cleanly.
  val hasBias = Reg(Bool())
  val steepness = steepnessOffset.U - io.req.bits.steepness
  val decimal = decimalPointOffset.U((decimalPointWidth + 1).W) +
    io.req.bits.decimalPoint
  val one = 1.S << decimal

  def applySteepness(x: SInt, steepness: UInt): SInt = {
    val tmp = Wire(SInt())
    when (steepness < (steepnessOffset).U) {
      tmp := x >> ((steepnessOffset).U - steepness)
    } .elsewhen (steepness === (steepnessOffset).U) {
      tmp := x
    } .otherwise {
      tmp := x << (steepness - (steepnessOffset).U)
    }
    tmp
  }

  // DSP Unit
  val dsp = Module(new DSP).io
  def DSP(a: SInt, b: SInt, c: UInt) {
    dsp.a := a
    dsp.b := b
    dsp.c := c
  }
  DSP(0.S, 0.S, 0.U)

  // Default values
  acc := acc
  reqSent := reqSent
  io.req.ready := Bool(false)
  io.resp.valid := Bool(false)
  io.resp.bits.state := state
  io.resp.bits.index := io.req.bits.index
  io.resp.bits.data := dataOut
  io.resp.bits.incWriteCount := Bool(false)
  index := index
  // Activation function unit default values
  af.io.req.valid := Bool(false)
  af.io.req.bits.in := (0).S
  af.io.req.bits.decimal := io.req.bits.decimalPoint
  af.io.req.bits.steepness := io.req.bits.steepness
  af.io.req.bits.activationFunction := io.req.bits.activationFunction

  state := state
  // Jump to nextState when we are able to get a request out
  def reqNoResp() = {
    when (io.req.valid) { state := nextState }
  }
  // Jump to nextState when we get a request out and we receive a
  // response
  def reqWaitForResp() = {
      when (!reqSent) {
        io.resp.valid := true.B
        reqSent := io.req.valid
      } .elsewhen (io.req.valid) {
        reqSent := false.B
        state := nextState }}
  def reqAf() = {
    when (!reqSent) {
      af.io.req.valid := true.B
      reqSent := true.B
    } .elsewhen (af.io.resp.valid) {
      reqSent := false.B
      state := nextState }}

  // State-driven logic
  switch (state) {
    is (PE_states('e_PE_UNALLOCATED)) {
      nextState := PE_states('e_PE_GET_INFO)
      reqNoResp()
      io.req.ready := true.B
      hasBias := false.B
      index := 0.U
      reqSent := false.B
    }
    is (PE_states('e_PE_GET_INFO)) {
      dataOut := 0.S
      index := 0.U
      nextState := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
      reqWaitForResp()
    }
    is (PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)) {
      // If hasBias is false, then this is the first time we're in this
      // state and we need to load the bias into the accumulator
      hasBias := true.B
      when (hasBias === false.B) {
        acc := io.req.bits.bias
      }
      nextState := PE_states('e_PE_RUN)
      reqWaitForResp()
    }
    is (PE_states('e_PE_RUN)) {
      when (index === (io.req.bits.numWeights - 1.U)) {
        state := PE_states('e_PE_ACTIVATION_FUNCTION)
      } .elsewhen (eleIndex === (elementsPerBlock - 1).U) {
        state := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
      }
      DSP(io.req.bits.iBlock(eleIndex), io.req.bits.wBlock(eleIndex), decimal)
      acc := acc + dsp.d
      index := index + 1.U
      printfInfo("run 0x%x + (0x%x * 0x%x) >> 0x%x = 0x%x\n",
        acc, io.req.bits.iBlock(eleIndex), io.req.bits.wBlock(eleIndex),
        decimal, acc + dsp.d)

      dataOut := Mux(af.io.resp.valid, af.io.resp.bits.out, dataOut)
    }
    is (PE_states('e_PE_DONE)) {
      nextState := PE_states('e_PE_UNALLOCATED)
      reqNoResp()
      io.resp.bits.incWriteCount := true.B
      io.resp.valid := true.B
    }
  }

  assert (!(state === PE_states('e_PE_ERROR)), printfSigil ++
    "is in error state\n")
}

class ProcessingElementLearn(id: Int = 0)(implicit p: Parameters)
    extends ProcessingElement(id)(p) {
  override lazy val io = IO(new ProcessingElementInterfaceLearn)
  override lazy val af = Module(new ActivationFunctionLearn(id))

  val weightWB        = Reg(Vec(elementsPerBlock, SInt(elementWidth.W)))
  val derivative      = Reg(SInt(elementWidth.W)) //delta
  val errorOut        = Reg(SInt(elementWidth.W)) //ek
  val mse             = Reg(SInt(elementWidth.W))
  val dwWritebackDone = Reg(Bool())
  val stateLearn = io.req.bits.stateLearn

  // Defaults
  derivative := derivative
  io.resp.bits.dataBlock := weightWB
  io.resp.bits.error := errorOut
  io.resp.bits.resetWeightPtr := false.B
  af.io.req.bits.afType := e_AF_DO_ACTIVATION_FUNCTION
  af.io.req.bits.errorFunction := io.req.bits.errorFunction

  switch (state) {
    is (PE_states('e_PE_UNALLOCATED)) {
      dwWritebackDone := false.B
    }
    is (PE_states('e_PE_GET_INFO)) {
      nextState := PE_states('e_PE_REQUEST_OUTPUTS_ERROR_BACKPROP)
      when (stateLearn === e_TTABLE_STATE_FEEDFORWARD ||
        stateLearn === e_TTABLE_STATE_LEARN_FEEDFORWARD ||
        ((stateLearn === e_TTABLE_STATE_LEARN_WEIGHT_UPDATE) &&
          (io.req.bits.tType === e_TTYPE_BATCH))) {
        nextState := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
      }
      reqWaitForResp()
    }
    is (PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)) {
      when (io.req.bits.tType === e_TTYPE_BATCH) {
        when (stateLearn === e_TTABLE_STATE_LEARN_ERROR_BACKPROP ||
          dwWritebackDone) {
          nextState := PE_states('e_PE_RUN_UPDATE_SLOPE)
        } .elsewhen (stateLearn === e_TTABLE_STATE_LEARN_WEIGHT_UPDATE) {
          nextState := PE_states('e_PE_RUN_WEIGHT_UPDATE) }
      } .elsewhen (io.req.bits.tType === e_TTYPE_INCREMENTAL) {
        when (stateLearn === e_TTABLE_STATE_LEARN_ERROR_BACKPROP ||
          dwWritebackDone) {
          nextState := PE_states('e_PE_RUN_UPDATE_SLOPE) }}
      reqWaitForResp()
    }
    is (PE_states('e_PE_RUN)) {
      when (index === (io.req.bits.numWeights - 1.U)) {
        state := PE_states('e_PE_ACTIVATION_FUNCTION)
      } .elsewhen (eleIndex === (elementsPerBlock - 1).U) {
        state := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
      }
      DSP(io.req.bits.iBlock(eleIndex), io.req.bits.wBlock(eleIndex), decimal)
      acc := acc + dsp.d
      index := index + 1.U
    }
    is (PE_states('e_PE_ACTIVATION_FUNCTION)) {
      af.io.req.bits.afType := e_AF_DO_ACTIVATION_FUNCTION
    }
    is (PE_states('e_PE_DONE)) {
      when (io.req.bits.inLast &&
        stateLearn === e_TTABLE_STATE_LEARN_FEEDFORWARD) {
        nextState := PE_states('e_PE_COMPUTE_DERIVATIVE)
      } .otherwise {
        nextState := PE_states('e_PE_UNALLOCATED)
      }
      reqNoResp()
      io.resp.bits.resetWeightPtr := true.B
    }
    is (PE_states('e_PE_COMPUTE_DERIVATIVE)){
      state := Mux((stateLearn === e_TTABLE_STATE_LEARN_ERROR_BACKPROP),
        PE_states('e_PE_COMPUTE_DELTA), PE_states('e_PE_REQUEST_EXPECTED_OUTPUT))

      val af = io.req.bits.activationFunction
      when (af === e_FANN_LINEAR) {
        printfWarn("Linear activation function untested\n")
        when (steepness < steepnessOffset.U) {
          derivative := one >> (steepnessOffset.U - steepness)
          printfInfo("derivative linear: 0x%x\n",
            one >> (steepnessOffset.U - steepness))
        } .elsewhen (steepness === steepnessOffset.U) {
          derivative := one
          printfInfo("derivative linear: 0x%x\n", one)
        } .otherwise {
          derivative := one << (steepness - steepnessOffset.U)
          printfInfo("derivative linear: 0x%x\n",
            one << (steepness - steepnessOffset.U))
        }
      } .elsewhen (af === e_FANN_SIGMOID || af === e_FANN_SIGMOID_STEPWISE) {
        DSP(dataOut, one - dataOut, decimal + steepness - 1.U)
        derivative := dsp.d
        printfInfo("derivative sigmoid: 0x%x\n", dsp.d)
      } .elsewhen (af === e_FANN_SIGMOID_SYMMETRIC ||
        af === e_FANN_SIGMOID_SYMMETRIC_STEPWISE) {
        val steepness = io.req.bits.steepness
        DSP(dataOut, dataOut, decimal)
        when (steepness < steepnessOffset.U) {
          derivative := (one - dsp.d) >> (steepnessOffset.U - steepness)
          printfInfo("derivative sigmoid symmetric: 0x%x\n",
            (one - dsp.d) >> (steepnessOffset.U - steepness))
        } .elsewhen (steepness === steepnessOffset.U) {
          derivative := one - dsp.d
          printfInfo("derivative sigmoid symmetric: 0x%x\n",
            one - dsp.d)
        } .otherwise {
          derivative := (one - dsp.d) << (steepness - steepnessOffset.U)
          printfInfo("derivative sigmoid symmetric: 0x%x\n",
            (one - dsp.d) << (steepness - steepnessOffset.U))
        }
      }
    }
    is (PE_states('e_PE_REQUEST_EXPECTED_OUTPUT)) {
      nextState := PE_states('e_PE_COMPUTE_ERROR)
      reqWaitForResp()
    }
    is (PE_states('e_PE_COMPUTE_ERROR)) {
      val error = io.req.bits.learnReg - dataOut
      val errorSym = error >> 1.U
      // Divide by 2 should be conditional on being in a symmetric
      when (io.req.bits.activationFunction === e_FANN_SIGMOID_SYMMETRIC ||
        io.req.bits.activationFunction === e_FANN_SIGMOID_SYMMETRIC_STEPWISE) {
        errorOut := errorSym
        DSP(errorSym, errorSym, decimal)
        mse := dsp.d
        // mse := (errorSym * errorSym) >> decimal
        printfInfo("errorOut and Error square set to 0x%x and  0x%x\n",
          errorSym, dsp.d)
      } .otherwise {
        errorOut := error
        DSP(error, error, decimal)
        mse := dsp.d
        // mse := (error * error) >> decimal
        printfInfo("errorOut and Error square set to 0x%x and  0x%x\n",
          error, dsp.d)
      }
      state := PE_states('e_PE_ERROR_FUNCTION)
    }
    is (PE_states('e_PE_ERROR_FUNCTION)) {
      reqAf()
      nextState := PE_states('e_PE_COMPUTE_DELTA)
      af.io.req.bits.in := errorOut
      af.io.req.bits.afType := e_AF_DO_ERROR_FUNCTION
      errorOut := Mux(af.io.resp.valid, af.io.resp.bits.out, errorOut)
      when (af.io.resp.valid) {
        printfInfo("errFn(0x%x) = 0x%x\n", errorOut, af.io.resp.bits.out)
      }
    }
    is (PE_states('e_PE_COMPUTE_DELTA)){
      // Reset the index in preparation for doing the deltas in the
      // previous layer
      val der = Mux(derivative === 0.S, 1.S, derivative)
      index := 0.U

      switch(stateLearn){
        is(e_TTABLE_STATE_LEARN_FEEDFORWARD){
          DSP(der, errorOut, decimal)
          errorOut := dsp.d
          // errorOut := (der * errorOut) >> decimal
          printfInfo("delta (output) 0x%x * 0x%x = 0x%x\n",
            der, errorOut, dsp.d)
          state := PE_states('e_PE_ERROR_BACKPROP_REQUEST_WEIGHTS)
        }
        is(e_TTABLE_STATE_LEARN_ERROR_BACKPROP){
          DSP(der, io.req.bits.dw_in, decimal)
          errorOut := dsp.d
          // errorOut := (der * io.req.bits.dw_in) >> decimal
          printfInfo("sees errFn*derivative 0x%x\n", dsp.d)
          when(io.req.bits.inFirst) {
            state := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
          }.otherwise {
            state := PE_states('e_PE_ERROR_BACKPROP_REQUEST_WEIGHTS)
          }
        }
      }
    }
    is (PE_states('e_PE_ERROR_BACKPROP_REQUEST_WEIGHTS)) {
      nextState := PE_states('e_PE_ERROR_BACKPROP_DELTA_WEIGHT_MUL)
      reqWaitForResp()
    }
    is (PE_states('e_PE_ERROR_BACKPROP_DELTA_WEIGHT_MUL)) {
      // Loop over all the weights in the weight buffer, multiplying
      // these by their delta
      when (index === (io.req.bits.numWeights - 1.U) ||
        eleIndex === (elementsPerBlock - 1).U) {
        state := PE_states('e_PE_ERROR_BACKPROP_WEIGHT_WB)
      }
      DSP(errorOut, io.req.bits.wBlock(eleIndex), decimal)
      weightWB(eleIndex) := dsp.d
      // weightWB(eleIndex) := (errorOut * io.req.bits.wBlock(eleIndex)) >>
      //   decimal
      printfInfo("d*weight (0x%x * 0x%x) >> 0x%x = 0x%x\n",
        errorOut, io.req.bits.wBlock(eleIndex), decimal,
        dsp.d)
      index := index + 1.U
    }
    is (PE_states('e_PE_ERROR_BACKPROP_WEIGHT_WB)) {
      when (io.req.valid) {
        when (index === io.req.bits.numWeights) {
          index := 0.U
          state := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
          io.resp.bits.resetWeightPtr := true.B
        } .otherwise {
          state := PE_states('e_PE_ERROR_BACKPROP_REQUEST_WEIGHTS)
        }
      }
      dwWritebackDone := true.B
      io.resp.valid := true.B
    }
    is (PE_states('e_PE_REQUEST_OUTPUTS_ERROR_BACKPROP)) {
      nextState := PE_states('e_PE_REQUEST_DELTA_WEIGHT_PRODUCT_ERROR_BACKPROP)
      reqWaitForResp()
    }
    is (PE_states('e_PE_REQUEST_DELTA_WEIGHT_PRODUCT_ERROR_BACKPROP)) {
      dataOut := io.req.bits.learnReg
      nextState := PE_states('e_PE_COMPUTE_DERIVATIVE)
      reqWaitForResp()
    }
    is(PE_states('e_PE_RUN_UPDATE_SLOPE)){
      // val delta = Mux(io.req.bits.inFirst, errorOut, io.req.bits.learnReg)
      val delta =  errorOut
      // [TODO] This is actually an element index!
      val inLastElement = (index === (io.req.bits.numWeights - 1.U) ||
        eleIndex === (elementsPerBlock - 1).U)
      index := index + 1.U
      when (inLastElement) {
        state := PE_states('e_PE_SLOPE_WB)
        when (io.req.bits.tType === e_TTYPE_INCREMENTAL) {
          state := PE_states('e_PE_RUN_WEIGHT_UPDATE)
          // Reset any modifications that have been made to the index
          // as this will be reused during weight update
          // [TODO] #54: Does this actually work?
          index := index(index.getWidth - 1, log2Up(elementsPerBlock)) ##
            0.U(log2Up(elementsPerBlock).W)
        }
      }
      DSP(delta, io.req.bits.iBlock(eleIndex), decimal)
      weightWB(eleIndex) := dsp.d
      printfInfo("update slope 0x%x * 0x%x = 0x%x\n",
        delta, io.req.bits.iBlock(eleIndex),
        dsp.d)
    }
    is(PE_states('e_PE_SLOPE_WB)){
      // val delta = Mux(io.req.bits.inFirst, errorOut, io.req.bits.learnReg)
      val delta = errorOut
      nextState := Mux(index === io.req.bits.numWeights,
        PE_states('e_PE_SLOPE_BIAS_WB),
        PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS))
      reqNoResp()
      io.resp.valid := true.B
      // Setup the bias to be written back (bias is in delta!)
      dataOut := delta
    }
    is (PE_states('e_PE_SLOPE_BIAS_WB)) {
      nextState := PE_states('e_PE_UNALLOCATED)
      reqNoResp()
      io.resp.valid := true.B
      printfInfo("bias wb: 0x%x\n", dataOut)
    }
    is (PE_states('e_PE_RUN_WEIGHT_UPDATE)){
      when (index === (io.req.bits.numWeights - 1.U) ||
        eleIndex === (elementsPerBlock - 1).U) {
        state := PE_states('e_PE_WEIGHT_UPDATE_WRITE_BACK)
      }
      val weightDecay = (-io.req.bits.wBlock(eleIndex) * io.req.bits.weightDecay) >>
        decimal
      printfInfo("weight decay %d: -1 * 0x%x * 0x%x = 0x%x\n",
        eleIndex,
        io.req.bits.wBlock(eleIndex), io.req.bits.weightDecay, weightDecay)

      when (io.req.bits.tType === e_TTYPE_BATCH) {
        DSP(io.req.bits.iBlock(eleIndex), io.req.bits.learningRate.asSInt,
          decimal)
        weightWB(eleIndex) := dsp.d + weightDecay
        printfInfo("weight update %d: 0x%x * 0x%x + 0x%x = 0x%x\n",
          eleIndex,
          io.req.bits.iBlock(eleIndex), io.req.bits.learningRate.asSInt,
          weightDecay, dsp.d + weightDecay)
      } .otherwise {
        // [TODO] #54: we're doing incremental learning here, so the
        // source is coming from weightWB
        DSP(weightWB(eleIndex), io.req.bits.learningRate.asSInt, decimal)
        weightWB(eleIndex) := dsp.d + weightDecay
        printfInfo("weight update %d: 0x%x * 0x%x + 0x%x = 0x%x\n",
          eleIndex,
          weightWB(eleIndex), io.req.bits.learningRate, weightDecay,
          dsp.d + weightDecay)
      }
      index := index + 1.U
    }
    is (PE_states('e_PE_WEIGHT_UPDATE_WRITE_BACK)){
      when (io.req.valid) {
        printfInfo("weight update writeback index/numWeights 0x%x/0x%x\n",
          index, io.req.bits.numWeights)
        // [TODO] #54: cleanup all this mux nonsense to make it
        // readable. Also, the transition here for the bias update is
        // likely incorrect when doing incremental learning. I think we
        // need to jump into the _d0 state and make sure that the bias
        // is loaded into dw_in prior to that.
        state := PE_states('e_PE_REQUEST_INPUTS_AND_WEIGHTS)
        when (index === io.req.bits.numWeights) {
          when (io.req.bits.tType === e_TTYPE_BATCH) {
            state := PE_states('e_PE_WEIGHT_UPDATE_REQUEST_BIAS)
          } .otherwise {
            state := PE_states('e_PE_WEIGHT_UPDATE_COMPUTE_BIAS)
          }
        }
      }
      io.resp.valid := true.B
    }
    is (PE_states('e_PE_WEIGHT_UPDATE_REQUEST_BIAS)) {
      when (!reqSent) {
        io.resp.valid := true.B
        reqSent := io.req.valid
      } .elsewhen (io.req.valid) {
        reqSent := false.B
        state := PE_states('e_PE_WEIGHT_UPDATE_COMPUTE_BIAS)
      }
    }
    is (PE_states('e_PE_WEIGHT_UPDATE_COMPUTE_BIAS)) {
      state := PE_states('e_PE_WEIGHT_UPDATE_WRITE_BIAS)
      val biasSlope = Mux(io.req.bits.tType === e_TTYPE_BATCH,
        io.req.bits.dw_in, errorOut)
      DSP(biasSlope, io.req.bits.learningRate.asSInt, decimal)
      dataOut := dsp.d
      printfInfo("biasSlope scale 0x%x * 0x%x = 0x%x\n",
        biasSlope, io.req.bits.learningRate.asSInt, dsp.d)
    }
    is (PE_states('e_PE_WEIGHT_UPDATE_WRITE_BIAS)) {
      when (io.req.valid) { state := PE_states('e_PE_UNALLOCATED) }
      io.resp.valid := true.B
    }
  }

  // Assertions

  // [TODO] #54: disallow specific states for certain transaction types
  assert(!(io.req.bits.tType === e_TTYPE_INCREMENTAL &&
    (state === PE_states('e_PE_SLOPE_WB) ||
      state === PE_states('e_PE_SLOPE_BIAS_WB))),
    printfSigil ++ "PE entered a disallowed state for incremental learning")
}

    }
    is (PE_states('e_PE_ACTIVATION_FUNCTION)) {
      reqAf()


// See LICENSE.BU for license details.
// See LICENSE.IBM for license details.

package dana

import chisel3._
import chisel3.util._
import scala.math.pow
import xfiles.{TransactionTableNumEntries}
import cde._

import dana.abi._

class RegisterFileInterface(implicit p: Parameters) extends DanaStatusIO()(p) {
  lazy val pe = Flipped(new PERegisterFileInterface)
  val control = Flipped(new ControlRegisterFileInterface)
  val tTable  = Flipped(new TTableRegisterFileInterface)
}

class RegisterFileInterfaceLearn(implicit p: Parameters)
    extends RegisterFileInterface()(p) {
  override lazy val pe = Flipped(new PERegisterFileInterfaceLearn)
}

class RegisterFileState(implicit p: Parameters) extends DanaBundle()(p) {
  val valid = Bool()
  val totalWrites = UInt(p(GlobalInfo).total_neurons.W)
  val countWrites = UInt(p(GlobalInfo).total_neurons.W)
}

class RegisterFileBase[SramIf <: SRAMElementInterface](
  genSram: => Vec[SramIf])(implicit p: Parameters)
    extends DanaModule()(p) {
  lazy val io = IO(new RegisterFileInterface)
  override val printfSigil = "dana.RegFile: "
  val mem = genSram

  val state = Reg(Vec(transactionTableNumEntries * 2, new RegisterFileState))
  val stateToggle = Reg(Vec(transactionTableNumEntries, UInt(1.W)))
  val tTableRespValid = Reg(Bool())
  val tTableRespTIdx = Reg(UInt(log2Up(transactionTableNumEntries).W))
  val tTableRespAddr = Reg(UInt(log2Up(p(ScratchpadElements)).W))

  // Default values for SRAMs
  for (transaction <- 0 until transactionTableNumEntries) {
    for (port <- 0 until mem(transaction).numPorts) {
      mem(transaction).we(port) := false.B
      mem(transaction).re(port) := false.B
      mem(transaction).din(port) := 0.U
      mem(transaction).dinElement(port) := 0.U
      mem(transaction).addr(port) := 0.U
    }
  }
  // Default Control interface values
  io.control.req.ready := true.B
  io.control.resp.valid := false.B
  io.control.resp.bits.tIdx := 0.U
  // Default Transaction Table interface values
  io.tTable.resp.valid := false.B
  io.tTable.resp.bits.data := 0.U

  // Requests from the Processing Element Table
  when (io.pe.req.valid) {
    // Take action based on whether this is a write or a read
    val tIdx = io.pe.req.bits.tIdx
    val sIdx = io.pe.req.bits.tIdx ## io.pe.req.bits.location
    when (io.pe.req.bits.isWrite) { // This is a Write
      mem(tIdx).we(0) := true.B
      mem(tIdx).addr(0) := io.pe.req.bits.addr
      mem(tIdx).dinElement(0) := io.pe.req.bits.data.asUInt
      printfInfo("PE write element tIdx/Addr/Data 0x%x/0x%x/0x%x\n",
        tIdx, io.pe.req.bits.addr, io.pe.req.bits.data)
      // Increment the write count and generate a response to the
      // control module if this puts us at the write count
      when (io.pe.req.bits.incWriteCount) {
        state(sIdx).countWrites := state(sIdx).countWrites + 1.U
        printfInfo("write count loc/seen/expected 0x%x/0x%x/0x%x\n",
          sIdx, state(sIdx).countWrites + 1.U, state(sIdx).totalWrites)
        when (state(sIdx).countWrites === state(sIdx).totalWrites - 1.U) {
          io.control.resp.valid := true.B
          io.control.resp.bits.tIdx := tIdx
        }
      }
    } .otherwise {                  // This is a read
      mem(tIdx).re(0) := true.B
      mem(tIdx).addr(0) := io.pe.req.bits.addr
      printfInfo("PE read tIdx/Addr 0x%x/0x%x\n", tIdx,
        io.pe.req.bits.addr)
    }
  }

  // Requests from the Control module
  when (io.control.req.valid) {
    val tIdx = io.control.req.bits.tIdx
    val location = io.control.req.bits.location
    state(tIdx << 1.U | location).valid := true.B
    state(tIdx << 1.U | location).totalWrites :=
      io.control.req.bits.totalWrites
    state(tIdx << 1.U | location).countWrites := 0.U
    printfInfo("Control req tIdx/location/totalWrites 0x%x/0x%x/0x%x\n",
      tIdx, location, io.control.req.bits.totalWrites)
  }

  // Requests form the Transaction Table
  when (io.tTable.req.valid) {
    val tIdx = io.tTable.req.bits.tidIdx
    switch (io.tTable.req.bits.reqType) {
      is (e_TTABLE_REGFILE_WRITE) {
        printfInfo("Saw TTable write idx/Addr/Data 0x%x/0x%x/0x%x\n",
          tIdx, io.tTable.req.bits.addr, io.tTable.req.bits.data)
        mem(tIdx).we(0) := true.B
        mem(tIdx).dinElement(0) := io.tTable.req.bits.data
        mem(tIdx).addr(0) := io.tTable.req.bits.addr
      }
      is (e_TTABLE_REGFILE_READ) {
        printfInfo("Saw TTable read Addr 0x%x\n",
          io.tTable.req.bits.addr)
        mem(tIdx).re(0) := true.B
        mem(tIdx).dinElement(0) := io.tTable.req.bits.data
        mem(tIdx).addr(0) := io.tTable.req.bits.addr
      }
    }
  }

  val readReqValid_d0 = Reg(next = io.pe.req.valid && !io.pe.req.bits.isWrite)
  val peIndex_d0 = Reg(next = io.pe.req.bits.peIndex)
  val tIndex_d0 = Reg(next = io.pe.req.bits.tIdx)

  io.pe.resp.valid := readReqValid_d0
  io.pe.resp.bits.peIndex := peIndex_d0
  io.pe.resp.bits.data := mem(tIndex_d0).dout(0)
  when (io.pe.resp.valid) {
    printfInfo("PE resp PE/Data 0x%x/0x%x\n",
      io.pe.resp.bits.peIndex, io.pe.resp.bits.data);
  }
  // Transaction Table Response
  tTableRespValid := io.tTable.req.valid &&
    io.tTable.req.bits.reqType === e_TTABLE_REGFILE_READ
  tTableRespTIdx := io.tTable.req.bits.tidIdx
  tTableRespAddr := io.tTable.req.bits.addr
  when (tTableRespValid) {
  val memDataVec = Vec((0 until elementsPerBlock).map(i =>
    (mem(tTableRespTIdx).dout(0))(elementWidth * (i + 1) - 1, elementWidth * i)))
    io.tTable.resp.valid := true.B
    io.tTable.resp.bits.data :=
      memDataVec(tTableRespAddr(log2Up(elementsPerBlock)-1,0))
    printfInfo("Returning data to TTable 0x%x\n",
      io.tTable.resp.bits.data)
  }

  // Reset
  when (reset) {for (i <- 0 until transactionTableNumEntries * 2) {
    state(i).valid := false.B}}

  // Assertions

  // The number of writes that we've seen should never be greater than

    "The total writes to a Regsiter File entry exceeded the number expected")

  // A request to change the total number of writes should only happen
  // to a state entry marked as valid if it's countWrites is equal to
  // the totalWrites.
  assert(!(io.control.req.valid &&
    state(io.control.req.bits.tIdx << 1.U |
      io.control.req.bits.location).valid &&
    (state(io.control.req.bits.tIdx << 1.U |
      io.control.req.bits.location).countWrites =/=
      state(io.control.req.bits.tIdx << 1.U |
        io.control.req.bits.location).totalWrites)), printfSigil ++
    "RegFile totalWrites being changed when valid && (countWrites != totalWrites)")

  // We shouldn't be trying to write data outside of the bounds of the
  // memory
  (0 until transactionTableNumEntries).map(i =>
    assert(!(mem(i).addr(0) >= p(ScratchpadElements).U), printfSigil ++
      "RegFile address (read or write) is out of bounds"))

}

class RegisterFile(implicit p: Parameters) extends RegisterFileBase (
  Vec.fill(p(TransactionTableNumEntries))(Module(new SRAMElement(
    dataWidth = p(BitsPerBlock),
    sramDepth = pow(2, log2Up(p(RegFileNumBlocks))).toInt,
    numPorts = 1,
    elementWidth = p(DanaDataBits))).io))(p)

class RegisterFileLearn(implicit p: Parameters) extends RegisterFileBase (
  Vec.fill(p(TransactionTableNumEntries))(Module(new SRAMElementIncrement(
    dataWidth = p(BitsPerBlock),
    sramDepth = pow(2, log2Up(p(RegFileNumBlocks))).toInt,
    numPorts = 1,
    elementWidth = p(DanaDataBits))).io))(p) {
  override lazy val io = IO(new RegisterFileInterfaceLearn)

  val readReqType_d0 = Reg(next = io.pe.req.bits.reqType)
  io.pe.resp.bits.reqType := readReqType_d0

  for (transaction <- 0 until transactionTableNumEntries) {
    for (port <- 0 until mem(transaction).numPorts) {
      mem(transaction).wType(port) := 0.U
    }
  }

  when (io.pe.req.valid) {
    // Take action based on whether this is a write or a read
    val tIdx = io.pe.req.bits.tIdx
    val sIdx = io.pe.req.bits.tIdx ## io.pe.req.bits.location
    when (io.pe.req.bits.isWrite) { // This is a Write
      mem(tIdx).addr(0) := io.pe.req.bits.addr
      switch (io.pe.req.bits.reqType) {
        is (e_PE_WRITE_ELEMENT) {
          mem(tIdx).wType(0) := 0.U
          mem(tIdx).dinElement(0) := io.pe.req.bits.data.asUInt
          printfInfo("PE write element tIdx/Addr/Data 0x%x/0x%x/0x%x\n",
            tIdx, io.pe.req.bits.addr, io.pe.req.bits.data)
        }
        is (e_PE_WRITE_BLOCK_NEW) {
          mem(tIdx).wType(0) := 1.U
          mem(tIdx).din(0) := io.pe.req.bits.dataBlock
          printfInfo("PE write block new tIdx/Addr/Data 0x%x/0x%x/0x%x\n",
            tIdx, io.pe.req.bits.addr, io.pe.req.bits.dataBlock)
        }
        is (e_PE_WRITE_BLOCK_ACC) {
          mem(tIdx).wType(0) := 2.U
          mem(tIdx).din(0) := io.pe.req.bits.dataBlock
          printfInfo("PE write block inc tIdx/Addr/Data 0x%x/0x%x/0x%x\n",
            tIdx, io.pe.req.bits.addr, io.pe.req.bits.dataBlock)
        }
        // Kludge to kill the write _if_ we're just incrementing the
        // write count
        is (e_PE_INCREMENT_WRITE_COUNT) {
          printfInfo("PE increment write count\n")
          mem(tIdx).we(0) := false.B
        }
      }
    }
  }
}

  // the number of expected writes.
  assert(!Vec((0 until transactionTableNumEntries * 2).map(
    i => (state(i).valid &&


// See LICENSE.IBM for license details.

package dana.util

import chisel3._

object divUp {
  def apply(dividend: Int, divisor: Int): Int = {
    (dividend + divisor - 1) / divisor}
}

object packInfo {

  def apply(epb: Int, pes: Int, cache: Int): Int = {
    var x = epb << (6 + 4);
    x = x | pes << 4;
    x = x | cache;
    x}