package FiveStage
import Chisel.MuxLookup
import chisel3._
import chisel3.util.{BitPat, MuxCase}
import chisel3.experimental.MultiIOModule


class Execute extends MultiIOModule {
  val io = IO(
    new Bundle {
      val op1 = Input(SInt(32.W))
      val op2 = Input(SInt(32.W))
      val rs1ValueIn = Input(SInt(32.W))
      val rs2ValueIn = Input(SInt(32.W))
      val rs2ValueOut = Output(SInt(32.W))
      val branchType = Input(UInt(3.W))
      val branch = Output(Bool())
      val ALUOp = Input(UInt(4.W))
      val ALUResult = Output(SInt(32.W))
    }
  )

  val ALUOpsMap = Array (
    ALUOps.ADD    -> (io.op1 + io.op2),
    ALUOps.SUB    -> (io.op1 - io.op2),
    ALUOps.AND    -> (io.op1 & io.op2),
    ALUOps.OR     -> (io.op1 | io.op2),
    ALUOps.XOR    -> (io.op1 ^ io.op2),
    ALUOps.SLT    -> Mux(io.op1 < io.op2, 1.S, 0.S),
    ALUOps.SLTU   -> Mux(io.op1.asUInt() < io.op2.asUInt(), 1.S, 0.S),
    ALUOps.SRA    -> (io.op1 >> io.op2(4, 0)),
    ALUOps.SRL    -> (io.op1.asUInt() >> io.op2(4, 0)).asSInt(),
    ALUOps.SLL    -> (io.op1.asUInt() << io.op2(4, 0)).asSInt(),
    ALUOps.COPY_A -> io.op1,
    ALUOps.COPY_B -> io.op2,
  )

  val BranchALUOpsMap = Array (
    branchType.beq -> (io.rs1ValueIn === io.rs2ValueIn),
    branchType.neq -> !(io.rs1ValueIn === io.rs2ValueIn),
    branchType.lt -> (io.rs1ValueIn < io.rs2ValueIn),
    branchType.gte -> (io.rs1ValueIn >= io.rs2ValueIn),
    branchType.ltu -> (io.rs1ValueIn.asUInt() < io.rs2ValueIn.asUInt()),
    branchType.gteu -> (io.rs1ValueIn.asUInt() >= io.rs2ValueIn.asUInt()),
    branchType.jump -> true.B,
  )

  io.rs2ValueOut := io.rs2ValueIn
  io.branch := MuxLookup(io.branchType, false.B, BranchALUOpsMap)
  io.ALUResult := MuxLookup(io.ALUOp, 0.S(32.W), ALUOpsMap)
}