`default_nettype none
`timescale 1ns/1ps

module datapath
(
	input wire clk, rst
);

// Pipeline Signals
wire stall_f;
wire stall_d;
wire flush_d;
wire flush_e;
wire [1:0] r1forward_e, r2forward_e;

// Fetch Signals
wire [31:0] pc_f;
wire [31:0] pc_4_f = pc_f + 4;
wire [31:0] pc_new = branch_pc ? pc_b_j : pc_4_f;

// Decode Signals
// From fetch
reg [31:0] instr_d;
reg [31:0] pc_d, pc_4_d;
// Controller Decode Signals
wire wb_pc_d, wb_mem_d;
wire wr_mem_d, wr_reg_d;
wire [4:0] alu_op_d;
wire alu_1_d, alu_2_d;
wire branch_d, jump_d, jump_r_d;
wire [3:0] i_b_j_u;
// Register Signals
wire [31:0] r1data_d, r2data_d;
wire [4:0] waddr_d = instr_d[11:7];
wire [4:0] r1addr_d = instr_d[19:15];
wire [4:0] r2addr_d = instr_d[24:20];
// Immediate Generation
wire [31:0] imm_d;

// Execute Signals
// PC and PC+4
reg [31:0] pc_e, pc_4_e;
// Control Signals
reg wb_pc_e, wb_mem_e;
reg wr_mem_e, wr_reg_e;
reg [4:0] alu_op_e;
reg alu_1_e, alu_2_e;
reg branch_e, jump_e, jump_r_e;
// Registers
reg [31:0] r1data_e, r2data_e;
reg [4:0] waddr_e, r1addr_e, r2addr_e;
//Immediate
reg [31:0] imm_e;
// ALU signals
wire [31:0] r1forward_reg_e = r1forward_e[1] ? reg_write_w : (r1forward_e[0] ? alu_out_m : r1data_e);
wire [31:0] r2forward_reg_e = r2forward_e[1] ? reg_write_w : (r2forward_e[0] ? alu_out_m : r2data_e);
wire [31:0] alu_in_1 = alu_1_e ? r1forward_reg_e : pc_e;
wire [31:0] alu_in_2 = alu_2_e ? r2forward_reg_e : imm_e;
wire [31:0] alu_out_e;
// Control flow signals
wire alu_branch;
wire branch_pc = jump_e | (branch_e & alu_branch);
wire [31:0] pc_b_j = imm_e + (jump_r_e ? r1data_e : pc_e);

// Memory Signals
// PC+4
reg [31:0] pc_4_m;
// Control Signals
reg wb_pc_m, wb_mem_m;
reg wr_mem_m, wr_reg_m;
//Registers
reg [31:0] r2data_m;
reg [4:0] waddr_m;
// ALU Out
reg [31:0] alu_out_m;

// WB signals
// Reg write back
reg wb_pc_w, wb_mem_w;
reg wr_reg_w;
reg [4:0] waddr_w;
reg [31:0] pc_4_w;
reg [31:0] alu_out_w, mem_out_w;
wire [31:0] reg_write_w = wb_pc_w ? pc_4_w : (wb_mem_w ? mem_out_w : alu_out_w);


always @ (posedge clk) begin

	// IF/ID Pipeline Stage
	if (flush_d) begin
		// Insert NOP
		instr_d <= 32'b0000000000000000000000010011;
		pc_d    <= 32'hXXXXXXXX;
		pc_4_d  <= 32'hXXXXXXXX;
	end
	else if (stall_d) begin
		instr_d <= instr_d;
		pc_d    <= pc_d;
		pc_4_d  <= pc_4_d;
	end
	else begin
		pc_d    <= pc_f;
		pc_4_d  <= pc_4_f;
	end
	
	// ID/IE stage
	// These Control signals don't affect arch state if left alone
	// So don't need to be flushed
	pc_e <= pc_d;
	pc_4_e <= pc_4_d;
	wb_pc_e <= wb_pc_d;
	wb_mem_e <= wb_mem_d;
	alu_op_e <= alu_op_d;
	alu_1_e <= alu_1_d;
	alu_2_e <= alu_2_d;
	r1data_e <= r1data_d;
	r2data_e <= r2data_d;
	r1addr_e <= r1addr_d;
	r2addr_e <= r2addr_d;
	waddr_e <= waddr_d;
	imm_e <= imm_d;
	if (flush_e) begin

		wr_mem_e <= 1'b0;
		wr_reg_e <= 1'b0;
		branch_e <= 1'b0;
		jump_e <= 1'b0;
		jump_r_e <= 1'b0;
	end
	else begin
		wr_mem_e <= wr_mem_d;
		wr_reg_e <= wr_reg_d;
		branch_e <= branch_d;
		jump_e <= jump_d;
		jump_r_e <= jump_r_d;
	end

	// IE/M Stage
	// Never flushed or stalled
	pc_4_m <= pc_4_e;
	wb_pc_m <= wb_pc_e;
	wb_mem_m <= wb_mem_e;
	wr_reg_m <= wr_reg_e;
	wr_mem_m <= wr_mem_e;
	waddr_m <= waddr_e;
	r2data_m <= r2forward_reg_e; 
	alu_out_m <= alu_out_e;

	// M/W Stage
	// Never flushed or stalled
	pc_4_w <= pc_4_m;
	wb_pc_w <= wb_pc_m;
	wb_mem_w <= wb_mem_m;
	wr_reg_w <= wr_reg_m;
	waddr_w <= waddr_m;
	alu_out_w <= alu_out_m;
end


// Pipeline Control
hazard hazard0
(
	.r1addr_e(r1addr_e), .r2addr_e(r2addr_e),
	.waddr_e(waddr_e), .waddr_m(waddr_m), .waddr_w(waddr_w),
	.wr_reg_m(wr_reg_m), .wr_reg_w(wr_reg_w),
	.r1forward_e(r1forward_e), .r2forward_e(r2forward_e),
	.r1addr_d(r1addr_d), .r2addr_d(r2addr_d),
	.wb_mem_e(wb_mem_e),
	.stall_f(stall_f), .stall_d(stall_d), .flush_e(flush_e), .flush_d(flush_d),
	.branch_pc(branch_pc)
);

// FETCH STAGE

pc pc0(
	.clk(clk), .we(~stall_f), .rst(rst),
	.pc_new(pc_new), .pc(pc_f)
);

memory #(.ADDR_WIDTH(8), .FILE_NAME("instr_mem.mem")) instr_mem
(
	.clk(clk), .we(1'b0),
	.be(4'b1111),
	.addr(pc_f[7:0]),
	.din(32'b0),
	.dout(instr_d)
);

// DECODE STAGE

decode dec0
(
	.opcode(instr_d[6:2]), .func3(instr_d[14:12]), .func7(instr_d[31]),
	.wb_pc_d(wb_pc_d), .wb_mem_d(wb_mem_d),
	.wr_mem_d(wr_mem_d), .wr_reg_d(wr_reg_d),
	.alu_op_d(alu_op_d), .alu_1_d(alu_1_d), .alu_2_d(alu_2_d),
	.branch_d(branch_d), .jump_d(jump_d), .jump_r_d(jump_r_d),
	.i_b_j_u(i_b_j_u)
);

regfile reg0
(
	.clk(clk), .we(wr_reg_w),
	.waddr(waddr_w), .wdata(reg_write_w),
	.r1addr(r1addr_d), .r2addr(r2addr_d), .r1data(r1data_d), .r2data(r2data_d)
);

immediate imm0
(
	.instr(instr_d[31:7]),
	.i(i_b_j_u[3]), .b(i_b_j_u[2]), .j(i_b_j_u[1]), .u(i_b_j_u[0]),
	.imm(imm_d)
);

// EXECUTE STAGE

alu alu0
(
	.A(alu_in_1),
	.B(alu_in_2),
	.C(alu_out_e),
	.OP(alu_op_e),
	.Branch(alu_branch)
);

// MEMORY STAGE
memory #(.ADDR_WIDTH(8), .FILE_NAME("data_mem.mem")) data_mem
(
	.clk(clk), .we(wr_mem_m),
	.be(4'b1111), // TODO Add support for smaller byte writes
	.addr(alu_out_m[7:0]),
	.din(r2data_m),
	.dout(mem_out_w)
);

// WB Stage has no modules only a mux


// Debug Waveforms
`ifdef verilator

initial begin
	$dumpfile("./log_datapath/datapath.vcd");
	$dumpvars(0, datapath);
end

`endif

endmodule