// AXI4-Stream building blocks (32-bit data, TKEEP, TLAST)

// Counts the bytes of each packet on its input stream and sends the count on its output stream
// when TLAST arrives. Only bytes with TKEEP high are counted.
module axis_pkt_count (
  input             aclk,
  input             aresetn,
  input      [31:0] s_tdata,
  input      [3:0]  s_tkeep,
  input             s_tlast,
  input             s_tvalid,
  output            s_tready,
  output reg [15:0] m_tdata,
  output reg        m_tvalid,
  input             m_tready
);
  reg [15:0] bytes;

  wire [2:0] kept = s_tkeep[0] + s_tkeep[1] + s_tkeep[2] + s_tkeep[3];

  // Stall the input while a count is waiting to be taken on the output
  assign s_tready = !m_tvalid || m_tready;

  always @(posedge aclk) begin
    if (!aresetn) begin
      bytes    <= 16'd0;
      m_tvalid <= 1'b0;
    end
    else begin
      if (m_tvalid && m_tready)
        m_tvalid <= 1'b0;
      if (s_tvalid && s_tready) begin
        if (s_tlast) begin
          m_tdata  <= bytes + kept;
          m_tvalid <= 1'b1;
          bytes    <= 16'd0;
        end
        else begin
          bytes <= bytes + kept;
        end
      end
    end
  end
endmodule

// A simple register stage: registers the data, but only accepts a new beat when it is empty.
// s_tready depends only on the stage's own state, but throughput drops to one beat every two cycles.
module axis_reg_simple (
  input             aclk,
  input             aresetn,
  input      [31:0] s_tdata,
  input             s_tvalid,
  output            s_tready,
  output reg [31:0] m_tdata,
  output reg        m_tvalid,
  input             m_tready
);
  assign s_tready = !m_tvalid;

  always @(posedge aclk) begin
    if (!aresetn) begin
      m_tvalid <= 1'b0;
    end
    else begin
      if (m_tvalid && m_tready)
        m_tvalid <= 1'b0;
      if (s_tvalid && s_tready) begin
        m_tdata  <= s_tdata;
        m_tvalid <= 1'b1;
      end
    end
  end
endmodule

// A skid buffer: s_tready is a register, so there is no combinational path from m_tready to s_tready,
// yet a beat can pass every cycle. When the output stalls, the one beat already on its way is caught
// in the skid register.
module axis_skid (
  input             aclk,
  input             aresetn,
  input      [31:0] s_tdata,
  input             s_tvalid,
  output reg        s_tready,
  output reg [31:0] m_tdata,
  output reg        m_tvalid,
  input             m_tready
);
  reg [31:0] skid_data;
  reg        skid_full;

  always @(posedge aclk) begin
    if (!aresetn) begin
      s_tready  <= 1'b1;
      m_tvalid  <= 1'b0;
      skid_full <= 1'b0;
    end
    else begin
      if (!m_tvalid || m_tready) begin
        // The output register is free (or being emptied): refill it
        if (skid_full) begin
          m_tdata   <= skid_data;
          m_tvalid  <= 1'b1;
          skid_full <= 1'b0;
          s_tready  <= 1'b1;
        end
        else begin
          m_tdata  <= s_tdata;
          m_tvalid <= s_tvalid && s_tready;
        end
      end
      else if (s_tvalid && s_tready) begin
        // The output is stalled and a beat arrives: park it in the skid register
        skid_data <= s_tdata;
        skid_full <= 1'b1;
        s_tready  <= 1'b0;
      end
    end
  end
endmodule

// A synchronous stream FIFO with DEPTH entries (DEPTH a power of two)
module axis_fifo #(
  parameter DEPTH = 4,
  parameter AW    = 2
) (
  input             aclk,
  input             aresetn,
  input      [31:0] s_tdata,
  input             s_tvalid,
  output            s_tready,
  output     [31:0] m_tdata,
  output            m_tvalid,
  input             m_tready
);
  reg [31:0] mem [0:DEPTH-1];
  reg [AW:0] wr_ptr;
  reg [AW:0] rd_ptr;

  wire [AW:0] count = wr_ptr - rd_ptr;

  assign s_tready = (count != DEPTH);
  assign m_tvalid = (count != 0);
  assign m_tdata  = mem[rd_ptr[AW-1:0]];

  always @(posedge aclk) begin
    if (!aresetn) begin
      wr_ptr <= 0;
      rd_ptr <= 0;
    end
    else begin
      if (s_tvalid && s_tready) begin
        mem[wr_ptr[AW-1:0]] <= s_tdata;
        wr_ptr <= wr_ptr + 1'b1;
      end
      if (m_tvalid && m_tready)
        rd_ptr <= rd_ptr + 1'b1;
    end
  end
endmodule

// Testbench: send three packets into the byte counter, with gaps and backpressure
module tb;
  reg         aclk    = 0;
  reg         aresetn = 0;
  reg  [31:0] s_tdata  = 0;
  reg  [3:0]  s_tkeep  = 0;
  reg         s_tlast  = 0;
  reg         s_tvalid = 0;
  wire        s_tready;
  wire [15:0] m_tdata;
  wire        m_tvalid;
  reg         m_tready = 1;

  axis_pkt_count dut (
    .aclk     (aclk),
    .aresetn  (aresetn),
    .s_tdata  (s_tdata),
    .s_tkeep  (s_tkeep),
    .s_tlast  (s_tlast),
    .s_tvalid (s_tvalid),
    .s_tready (s_tready),
    .m_tdata  (m_tdata),
    .m_tvalid (m_tvalid),
    .m_tready (m_tready)
  );

  always #5 aclk = ~aclk;

  // The receiver of the counts is busy (TREADY low) from 80 ns to 120 ns
  initial begin
    #80 m_tready = 0;
    #40 m_tready = 1;
  end

  // Send one beat: hold it until the handshake, as the protocol requires
  task beat(input [31:0] d, input [3:0] k, input l);
    begin
      s_tdata  <= d;
      s_tkeep  <= k;
      s_tlast  <= l;
      s_tvalid <= 1;
      @(posedge aclk);
      while (!s_tready)
        @(posedge aclk);
      s_tvalid <= 0;
    end
  endtask

  always @(posedge aclk) begin
    if (s_tvalid && s_tready)
      $display("%0t  in   data=0x%h keep=%b%0s", $time, s_tdata, s_tkeep, s_tlast ? " last" : "");
    if (m_tvalid && m_tready)
      $display("%0t  out  packet length = %0d bytes", $time, m_tdata);
  end

  initial begin
    $dumpfile("dump.vcd");
    $dumpvars(0, tb);
    #22 aresetn = 1;
    @(posedge aclk);
    // Packet 1: three full beats and a last beat with two bytes
    beat(32'h03020100, 4'b1111, 0);
    beat(32'h07060504, 4'b1111, 0);
    beat(32'h0B0A0908, 4'b1111, 0);
    beat(32'h00000D0C, 4'b0011, 1);
    // Packet 2: a single byte
    beat(32'h000000EE, 4'b0001, 1);
    // Packet 3: two beats with a gap between them
    beat(32'h13121110, 4'b1111, 0);
    repeat (2) @(posedge aclk);
    beat(32'h17161514, 4'b1111, 1);
    repeat (6) @(posedge aclk);
    $finish;
  end
endmodule