/*
              _____
     x[0] ==>|  A  |
 x_stb[0] -->|  D  |
 x_ack[0] <--|  D  |==> y
             |  E  |--> y_stb
     x[1] ==>|  R  |<-- y_ack
 x_stb[1] -->|     |
 x_ack[1] <--|_____|

*/
module cadder#(parameter N=32)(clk, rst, x, x_ack, x_stb, y, y_ack, y_stb);
    input logic clk;
    input logic rst;
    input wire [N-1:0] x [1:0];
    output logic [N-1:0] y;
    output x_ack[1:0];
    input x_stb[1:0];
    input y_ack;
    output y_stb;

    wire left_ack, left_stb;
    assign x_ack[0] = left_ack;
    assign x_ack[1] = left_ack;
    assign left_stb = x_stb[0] & x_stb[1];

    adder add0 (
        .clk(clk),
        .rst(rst),
        .input_a(x[0]),
        .input_b(x[1]),
        .input_stb(left_stb),
        .input_ack(left_ack),
        .output_z(y),
        .output_z_ack(y_ack),
        .output_z_stb(y_stb)
    );
endmodule : cadder

/*
  K layers of cascade adder
  
  Example of K=3 adder:
  
IN | K3 |  K2  |  K1  | OUT
     _
  ->| |    _
  ->|_|-->| |
     _    | |--\
  ->| |-->|_|  |   _
  ->|_|        \->| |
     _            | |->
  ->| |    _   /->|_|
  ->|_|-->| |  |
     _    | |--/
  ->| |-->|_|
  ->|_|
  
[inputs]
x size: 2**K
left io size: 2**K
*/

module adder_casc#(parameter K,N=32)(clk, rst, x, y, left, right);
    input logic clk;
    input logic rst;
    input wire [N-1:0] x [2**K-1:0];
    output logic [N-1:0] y;
    
    abus_io right;
    abus_io left[2**K-1:0];

    wire [N-1:0] layer_w [2**K-3:0];
    wire ack_w [2**K-3:0];
    wire stb_w [2**K-3:0];

    genvar i,j;
    generate
        for(i=0; i<K; i++) begin : generate_layers
            // First layers
            if(i == 0) begin
                for(j=0; j<2**(K-1); j++) begin : generate_casc0
                    cadder a(
                      .clk(clk),
                      .rst(rst),
                      .x(x[j*2+:2]),
                      .y(layer_w[j]),
                      .x_ack({left[j*2].ack, left[j*2+1].ack}),
                      .x_stb({left[j*2].stb, left[j*2+1].stb}),
                      .y_ack(ack_w[j]),
                      .y_stb(stb_w[j])
                    );
                end
            end
            // Last layer
            else if((K-i) <= 1) begin
                localparam s0 = 2**K-4;
                localparam s1 = 2**K-3;
                cadder c(
                    .clk(clk),
                    .rst(rst),
                    .x(layer_w[s0+:2]),
                    .y(y),
                    .x_ack({ack_w[s0], ack_w[s1]}),
                    .x_stb({stb_w[s0], stb_w[s1]}),
                    .y_ack(right.ack),
                    .y_stb(right.stb)
                );
            end
            // Middle layers
            else begin
                for(j=0; j<2**(K-i-1); j++) begin : generate_casc1
                    localparam s = $floor((2.0**(K-1.0) * (2.0**(i-1)-1.0)/2.0**(i-1))+j);
                    localparam ix = s*2;
                    localparam iy = s+2**(K-1);

                    cadder b(
                        .clk(clk),
                        .rst(rst),
                        .x(layer_w[ix+:2]),
                        .y(layer_w[iy]),
                        .x_ack(ack_w[ix+:2]),
                        .x_stb(stb_w[ix+:2]),
                        .y_ack(ack_w[iy]),
                        .y_stb(stb_w[iy])
                    );
                end
            end
        end
    endgenerate

endmodule : adder_casc

module adder_casc_tb();
    logic clk, rst;
    
    localparam K=4;
    logic [31:0] x [2**K-1:0];
    logic [31:0] y;
    logic ack [2**K-1:0];
    logic stb [2**K-1:0];

    abus_io input_ios[2**K-1:0]();
    abus_io output_io();
    
    genvar k;
    generate
        for(k=0; k<2**K; k++) begin : io_mapper
            assign input_ios[k].stb = stb[k];
            assign ack[k] = input_ios[k].ack;
        end
    endgenerate
    
    
    adder_casc#(.K(K)) adder_casc0(.clk(clk), .rst(rst), .x(x), .y(y), .left(input_ios), .right(output_io.left));
    initial forever #5 clk = ~clk;
    initial begin
        
        $display("Testing adder_casc");
        clk = 0;
        rst = 1;
        
        foreach(stb[i]) stb[i] = 0;
        output_io.ack = 0;
        #20
        rst = 0;
        // Initialise with floating point 2**i
        foreach(x[i]) x[i] = ('h400 + (i*8)) << 20;
        foreach(stb[i]) stb[i] = 1;

        fork
            foreach(ack[i]) begin
                fork
                    wait(ack[i] == 1);
                    #20
                    stb[i] = 0;
                join
            end
        join
        #20
        
        wait(output_io.stb == 1);
        output_io.ack = 1;
        assert(y[0] == 'h47ffff00);
        wait(output_io.stb == 0);
        output_io.ack = 0;
    end
    
endmodule : adder_casc_tb