// 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: a simple register stage versus a skid buffer, both fed by a source that always has data
module tb;
  reg         aclk    = 0;
  reg         aresetn = 0;
  reg  [31:0] a_data  = 0;               // source for the simple stage
  reg  [31:0] b_data  = 0;               // source for the skid buffer
  wire        a_ready;
  wire        b_ready;
  wire [31:0] a_out;
  wire [31:0] b_out;
  wire        a_valid;
  wire        b_valid;
  reg         m_tready = 1;
  integer     a_beats  = 0;
  integer     b_beats  = 0;
  integer     cyc      = 0;

  axis_reg_simple stage (
    .aclk     (aclk),
    .aresetn  (aresetn),
    .s_tdata  (a_data),
    .s_tvalid (aresetn),
    .s_tready (a_ready),
    .m_tdata  (a_out),
    .m_tvalid (a_valid),
    .m_tready (m_tready)
  );

  axis_skid skid (
    .aclk     (aclk),
    .aresetn  (aresetn),
    .s_tdata  (b_data),
    .s_tvalid (aresetn),
    .s_tready (b_ready),
    .m_tdata  (b_out),
    .m_tvalid (b_valid),
    .m_tready (m_tready)
  );

  always #5 aclk = ~aclk;

  always @(posedge aclk) begin
    if (aresetn) begin
      cyc <= cyc + 1;
      // each source moves to its next value when its beat is accepted
      if (a_ready)
        a_data <= a_data + 1;
      if (b_ready)
        b_data <= b_data + 1;
      if (a_valid && m_tready)
        a_beats <= a_beats + 1;
      if (b_valid && m_tready)
        b_beats <= b_beats + 1;
      // the sink stalls in cycles 10 and 11
      m_tready <= !(cyc == 9 || cyc == 10);
    end
  end

  initial begin
    $dumpfile("dump.vcd");
    $dumpvars(0, tb);
    #22 aresetn = 1;
    repeat (20) @(posedge aclk);
    #1;
    $display("after 20 cycles: simple stage %0d beats, skid buffer %0d beats", a_beats, b_beats);
    $display("last beats out:  simple stage 0x%0h, skid buffer 0x%0h", a_out, b_out);
    $finish;
  end
endmodule