package FiveStage
import chisel3._
import chisel3.util._
import chisel3.experimental.MultiIOModule


class MemoryFetch() extends MultiIOModule {


  // Don't touch the test harness
  val testHarness = IO(
    new Bundle {
      val DMEMsetup      = Input(new DMEMsetupSignals)
      val DMEMpeek       = Output(UInt(32.W))

      val testUpdates    = Output(new MemUpdates)
    })

  val io = IO(
    new Bundle {
      val ALUResult = Input(UInt(32.W))
      val writeData = Input(UInt(32.W))
      val readMem = Input(Bool())
      val writeMem = Input(Bool())
      val dataOut = Output(UInt(32.W))
      val jump = Input(Bool())
      val returnAddr = Input(UInt(32.W))
    })

  val DMEM = Module(new DMEM)

  /**
    * Setup. You should not change this code
    */
  DMEM.testHarness.setup  := testHarness.DMEMsetup
  testHarness.DMEMpeek    := DMEM.io.dataOut
  testHarness.testUpdates := DMEM.testHarness.testUpdates

  /**
    * Your code here.
    */

  DMEM.io.dataIn      := io.writeData
  DMEM.io.writeEnable := io.writeMem
  DMEM.io.dataAddress := io.ALUResult

  val readMem = RegInit(Bool(), false.B)
  readMem := io.readMem

  io.dataOut := Mux(readMem, DMEM.io.dataOut, Mux(io.jump, io.returnAddr, io.ALUResult))
}