// poly_arith_sync.v - Polynomial modular add/sub for ML-KEM
//
// Performs element-wise modular addition or subtraction on two streaming
// 256-coefficient polynomials over Z_q (q=3329).
//
// mode=0: coeff_out = (coeff_a_in + coeff_b_in) mod Q
// mode=1: coeff_out = (coeff_a_in - coeff_b_in) mod Q
//
// Each cycle processes one coefficient pair. Pure streaming — no internal
// storage; the caller sequences all 256 coefficients in order.

`include "sync_rtl/common/defines.vh"

module poly_arith_sync (
    input         clk,
    input         rst_n,
    input  [11:0] coeff_a_in,
    input  [11:0] coeff_b_in,
    input         mode,       // 0=add, 1=sub
    input         valid_i,
    output        ready_o,
    output [11:0] coeff_out,
    output        valid_o,
    input         ready_i
);

    //--------------------------------------------------------------
    // Combinational modular arithmetic
    //--------------------------------------------------------------
    wire [12:0] add_raw;        // a + b (13-bit to catch overflow)
    wire [11:0] add_sub_q;      // (a + b) - Q
    wire [11:0] add_result;     // (a + b) mod Q

    assign add_raw   = {1'b0, coeff_a_in} + {1'b0, coeff_b_in};
    assign add_sub_q = add_raw[11:0] - `Q;
    assign add_result = (add_raw < `Q) ? add_raw[11:0] : add_sub_q;

    wire [11:0] sub_result;     // (a - b) mod Q
    // When a >= b: diff = a - b (already in [0, Q-2])
    // When a <  b: diff = a + Q - b (borrow correction)
    assign sub_result = (coeff_a_in < coeff_b_in)
        ? (coeff_a_in + `Q - coeff_b_in)
        : (coeff_a_in - coeff_b_in);

    wire [11:0] mod_result = mode ? sub_result : add_result;

    //--------------------------------------------------------------
    // Pipeline the result through pipeline_reg for valid/ready
    //--------------------------------------------------------------
    pipeline_reg #(.DW(12)) u_pipe (
        .clk    (clk),
        .rst_n  (rst_n),
        .data_i (mod_result),
        .valid_i(valid_i),
        .ready_o(ready_o),
        .data_o (coeff_out),
        .valid_o(valid_o),
        .ready_i(ready_i)
    );

endmodule