// coeff_bram_4way.v - 4-way conflict-free banked storage for NTT coefficients
//
// Replaces the 256x12 LUT register file (reg [11:0] cr [0:255]) with 4 BRAM
// banks (each 128x12 via sd_bram), removing the large multi-port LUT mux tree.
//
// Bank function: bank(a) = { a[0], ^a[7:1] }   (bit0 XOR parity of upper 7 bits)
//   local(a)   = a[7:1]   (7-bit index into the 128-deep bank)
//
// This function makes a dual-butterfly's four simultaneous coefficient accesses
// {j, j+1, j+len, j+1+len} land in four DISTINCT banks for every ML-KEM NTT
// layer_len in {2,4,...,128}: j is even and len is even, so
//   j        -> {0, P }
//   j+1      -> {1, P }
//   j+len    -> {0, P'}   (flipping bit m>=1 flips the upper-bits parity)
//   j+1+len  -> {1, P'}
// with P != P'. Four distinct banks => 4 conflict-free reads/cycle from four
// single-read-port banks, and likewise 4 conflict-free writes/cycle. Each bank
// therefore needs only 1 read + 1 write per cycle (an sd_bram), no ping-pong.
//
// Read latency is 1 cycle (sd_bram registered read): present rd_addr_k this
// cycle, rd_data_k valid next cycle. Writes commit on the presenting cycle.
//
// NOTE: the caller must guarantee the four read addresses map to four distinct
// banks (true for the NTT butterfly pattern) and the four write addresses too.
// Sequential access (load/output) uses a single port and is always safe.

module coeff_bram_4way (
    input         clk, rst_n,
    // Four read addresses (global 8-bit); data returns 1 cycle later.
    input  [7:0]  rd_addr_0, rd_addr_1, rd_addr_2, rd_addr_3,
    output [11:0] rd_data_0, rd_data_1, rd_data_2, rd_data_3,
    // Four write ports.
    input         wr_en_0, wr_en_1, wr_en_2, wr_en_3,
    input  [7:0]  wr_addr_0, wr_addr_1, wr_addr_2, wr_addr_3,
    input  [11:0] wr_data_0, wr_data_1, wr_data_2, wr_data_3
);

    // Bank id and local index for an address.
    function [1:0] bank_id;
        input [7:0] a;
        bank_id = {a[0], ^a[7:1]};
    endfunction

    wire [1:0] rb0 = bank_id(rd_addr_0);
    wire [1:0] rb1 = bank_id(rd_addr_1);
    wire [1:0] rb2 = bank_id(rd_addr_2);
    wire [1:0] rb3 = bank_id(rd_addr_3);

    wire [1:0] wb0 = bank_id(wr_addr_0);
    wire [1:0] wb1 = bank_id(wr_addr_1);
    wire [1:0] wb2 = bank_id(wr_addr_2);
    wire [1:0] wb3 = bank_id(wr_addr_3);

    // Per-bank read-address crossbar: bank b reads the local index of whichever
    // read port targets it (exactly one, by construction).
    wire [6:0] bank_rd_local [0:3];
    wire [6:0] bank_wr_local [0:3];
    wire       bank_wr_en    [0:3];
    wire [11:0] bank_wr_data  [0:3];
    wire [11:0] bank_rd_data  [0:3];

    genvar b;
    generate
        for (b = 0; b < 4; b = b + 1) begin : gen_xbar
            // read address route (priority: 3>2>1>0; ties impossible for NTT)
            assign bank_rd_local[b] =
                (rb3 == b[1:0]) ? rd_addr_3[7:1] :
                (rb2 == b[1:0]) ? rd_addr_2[7:1] :
                (rb1 == b[1:0]) ? rd_addr_1[7:1] :
                                  rd_addr_0[7:1];
            // write route
            wire m0 = wr_en_0 && (wb0 == b[1:0]);
            wire m1 = wr_en_1 && (wb1 == b[1:0]);
            wire m2 = wr_en_2 && (wb2 == b[1:0]);
            wire m3 = wr_en_3 && (wb3 == b[1:0]);
            assign bank_wr_en[b]   = m0 | m1 | m2 | m3;
            assign bank_wr_local[b] =
                m3 ? wr_addr_3[7:1] :
                m2 ? wr_addr_2[7:1] :
                m1 ? wr_addr_1[7:1] : wr_addr_0[7:1];
            assign bank_wr_data[b] =
                m3 ? wr_data_3 :
                m2 ? wr_data_2 :
                m1 ? wr_data_1 : wr_data_0;

            sd_bram #(.W(12), .D(128), .A(7)) u_bank (
                .clk(clk),
                .rd_addr(bank_rd_local[b]),
                .rd_data(bank_rd_data[b]),
                .wr_en(bank_wr_en[b]),
                .wr_addr(bank_wr_local[b]),
                .wr_data(bank_wr_data[b])
            );
        end
    endgenerate

    // Register the read bank-ids to align with sd_bram's 1-cycle read latency,
    // then route each bank's output back to the requesting read port.
    reg [1:0] rb0_r, rb1_r, rb2_r, rb3_r;
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            rb0_r <= 2'd0; rb1_r <= 2'd0; rb2_r <= 2'd0; rb3_r <= 2'd0;
        end else begin
            rb0_r <= rb0; rb1_r <= rb1; rb2_r <= rb2; rb3_r <= rb3;
        end
    end

    assign rd_data_0 = bank_rd_data[rb0_r];
    assign rd_data_1 = bank_rd_data[rb1_r];
    assign rd_data_2 = bank_rd_data[rb2_r];
    assign rd_data_3 = bank_rd_data[rb3_r];

endmodule