// SPDX-License-Identifier: Apache-2.0

package firrtl.passes
package memlib

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import MemPortUtils._
import MemTransformUtils._

/** Changes memory port names to standard port names (i.e. RW0 instead T_408)
  */
object RenameAnnotatedMemoryPorts extends Pass {

  /** Renames memory ports to a standard naming scheme:
    *    - R0, R1, ... for each read port
    *    - W0, W1, ... for each write port
    *    - RW0, RW1, ... for each readwrite port
    */
  def createMemProto(m: DefAnnotatedMemory): DefAnnotatedMemory = {
    val rports = m.readers.indices.map(i => s"R$i")
    val wports = m.writers.indices.map(i => s"W$i")
    val rwports = m.readwriters.indices.map(i => s"RW$i")
    m.copy(readers = rports, writers = wports, readwriters = rwports)
  }

  /** Maps the serialized form of all memory port field names to the
    *    corresponding new memory port field Expression.
    *  E.g.:
    *    - ("m.read.addr") becomes (m.R0.addr)
    */
  def getMemPortMap(m: DefAnnotatedMemory, memPortMap: MemPortMap): Unit = {
    val defaultFields = Seq("addr", "en", "clk")
    val rFields = defaultFields :+ "data"
    val wFields = rFields :+ "mask"
    val rwFields = defaultFields ++ Seq("wmode", "wdata", "rdata", "wmask")

    def updateMemPortMap(ports: Seq[String], fields: Seq[String], newPortKind: String): Unit =
      for {
        (p, i) <- ports.zipWithIndex
        f <- fields
      } {
        val newPort = WSubField(WRef(m.name), newPortKind + i)
        val field = WSubField(newPort, f)
        memPortMap(s"${m.name}.$p.$f") = field
      }
    updateMemPortMap(m.readers, rFields, "R")
    updateMemPortMap(m.writers, wFields, "W")
    updateMemPortMap(m.readwriters, rwFields, "RW")
  }

  /** Replaces candidate memories with memories with standard port names
    * Does not update the references (this is done via updateStmtRefs)
    */
  def updateMemStmts(memPortMap: MemPortMap)(s: Statement): Statement = s match {
    case m: DefAnnotatedMemory =>
      val updatedMem = createMemProto(m)
      getMemPortMap(m, memPortMap)
      updatedMem
    case s => s.map(updateMemStmts(memPortMap))
  }

  /** Replaces candidate memories and their references with standard port names
    */
  def updateMemMods(m: DefModule) = {
    val memPortMap = new MemPortMap
    (m.map(updateMemStmts(memPortMap))
      .map(updateStmtRefs(memPortMap)))
  }

  def run(c: Circuit) = c.copy(modules = c.modules.map(updateMemMods))
}