// See LICENSE for license details.

package firrtlTests

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import FirrtlCheckers._

import collection.mutable

class RemoveWiresSpec extends FirrtlFlatSpec {
  def compile(input: String): CircuitState =
    (new LowFirrtlCompiler).compileAndEmit(CircuitState(parse(input), ChirrtlForm), List.empty)
  def compileBody(body: String) = {
    val str = """
      |circuit Test :
      |  module Test :
      |""".stripMargin + body.split("\n").mkString("    ", "\n    ", "")
    compile(str)
  }

  def getNodesAndWires(circuit: Circuit): (Seq[DefNode], Seq[DefWire]) = {
    require(circuit.modules.size == 1)

    val nodes = mutable.ArrayBuffer.empty[DefNode]
    val wires = mutable.ArrayBuffer.empty[DefWire]
    def onStmt(stmt: Statement): Statement = {
      stmt map onStmt match {
        case node: DefNode => nodes += node
        case wire: DefWire => wires += wire
        case _ =>
      }
      stmt
    }

    circuit.modules.head match {
      case Module(_,_,_, body) => onStmt(body)
    }
    (nodes, wires)
  }

  "Remove Wires" should "turn wires and their single connect into nodes" in {
    val result = compileBody(s"""
      |input a : UInt<8>
      |output b : UInt<8>
      |wire w : UInt<8>
      |w <= a
      |b <= w""".stripMargin
    )
    val (nodes, wires) = getNodesAndWires(result.circuit)
    wires.size should be (0)

    nodes.map(_.serialize) should be (Seq("node w = a"))
  }

  it should "order nodes in a legal, flow-forward way" in {
    val result = compileBody(s"""
      |input a : UInt<8>
      |output b : UInt<8>
      |wire w : UInt<8>
      |wire x : UInt<8>
      |node y = x
      |x <= w
      |w <= a
      |b <= y""".stripMargin
    )
    val (nodes, wires) = getNodesAndWires(result.circuit)
    wires.size should be (0)
    nodes.map(_.serialize) should be (
      Seq("node w = a",
          "node x = w",
          "node y = x")
    )
  }

  it should "properly pad rhs of introduced nodes if necessary" in {
    val result = compileBody(s"""
      |output b : UInt<8>
      |wire w : UInt<8>
      |w <= UInt(2)
      |b <= w""".stripMargin
    )
    val (nodes, wires) = getNodesAndWires(result.circuit)
    wires.size should be (0)
    nodes.map(_.serialize) should be (
      Seq("""node w = pad(UInt<2>("h2"), 8)""")
    )
  }

  it should "support arbitrary expression for wire connection rhs" in {
    val result = compileBody(s"""
      |input a : UInt<8>
      |input b : UInt<8>
      |output c : UInt<8>
      |wire w : UInt<8>
      |w <= tail(add(a, b), 1)
      |c <= w""".stripMargin
    )
    val (nodes, wires) = getNodesAndWires(result.circuit)
    wires.size should be (0)
    nodes.map(_.serialize) should be (
      Seq("""node w = tail(add(a, b), 1)""")
    )
  }

  it should "do a reasonable job preserving input order for unrelatd logic" in {
    val result = compileBody(s"""
      |input a : UInt<8>
      |input b : UInt<8>
      |output z : UInt<8>
      |node x = not(a)
      |node y = not(b)
      |z <= and(x, y)""".stripMargin
    )
    val (nodes, wires) = getNodesAndWires(result.circuit)
    wires.size should be (0)
    nodes.map(_.serialize) should be (
      Seq("node x = not(a)",
          "node y = not(b)")
    )
  }
}