Huffman_Tree_constr.sv

`timescale 1ns / 1ps

// You MUST NOT change the module name or the ports declarations

module Huffman_Tree_constr #(
    parameter CHAR_SPACE = 5,
    parameter CHAR_LEN = 3,
    parameter CODE_LEN = 4,
    parameter FREQ_LEN = 4,
    parameter CODE_LEN_LEN = 3
)(
    input clk,
    input rst_n,
    input [CHAR_SPACE-1: 0][FREQ_LEN-1: 0] character_freq_unsort,
    input start,

    output reg start_nxt,
    output reg [CHAR_SPACE-1: 0][CODE_LEN-1: 0] huffman_lib_out, 
    output reg [CHAR_SPACE-1: 0][CODE_LEN_LEN-1: 0] huffman_len_out  
);
    // These declarations are provided as hints for a possible implementation.
    // You are free to modify, add, or remove them as long as the module's
    // external behavior is correct

    reg [CHAR_SPACE-1:0][CODE_LEN-1:0] char_code;
    reg [CHAR_SPACE-1:0][CODE_LEN_LEN-1:0] char_len;
    reg [CHAR_SPACE-1:0][CHAR_LEN-1:0] char_pool_id;
    
    localparam  INIT        = 0,
                LOAD        = 1,
                SORT        = 2,
                WAIT_SORT   = 3,
                SUM         = 4,
                POOL_UPDATE = 5,
                CODE_UPDATE = 6,
                DONE        = 7;

    localparam MAX_NUM = 2 ** FREQ_LEN;

    reg[2:0] state;
    reg [FREQ_LEN:0] sum_min_nxtmin;
    reg [CHAR_LEN-1:0] loop_counter;


    // DO NOT change the module instantiation

    reg start_sort;
    wire done_sort;

    reg [CHAR_SPACE-1:0][FREQ_LEN:0] arr_pool_freq;
    wire [CHAR_LEN-1:0] min_pool_id;
    wire [CHAR_LEN-1:0] nxtmin_pool_id;

    Huffman_Sort #(
        .CHAR_SPACE(CHAR_SPACE),
        .FREQ_LEN(FREQ_LEN),
        .CHAR_LEN(CHAR_LEN)
    ) Huffman_Sort_inst(
        .start(start_sort),
        .clk(clk),
        .rst_n(rst_n),
        .character_pool(arr_pool_freq),

        .start_nxt(done_sort),
        .min_pool_id(min_pool_id),
        .nxtmin_pool_id(nxtmin_pool_id)
    );
    
    // ----------- Insert your codes below --------------//
    
    
    /*
      * ChatGPT assistance in Huffman_Tree_constr :
      * - FSM logic for pool and code update states 
      * - Handshake asserts with sub-modules
      * - some comments
     */
     
     
    // Latch selected pools for stable updates across states
    reg [CHAR_LEN-1:0] min_id_reg, nxt_id_reg;

    integer i;

    always_ff @(posedge clk) begin
        if (!rst_n) begin
            start_nxt       <= 1'b0;
            start_sort      <= 1'b0;
            state           <= INIT;

            arr_pool_freq   <= '0;
            char_code       <= '0;
            char_len        <= '0;
            char_pool_id    <= '0;

            huffman_lib_out <= '0;
            huffman_len_out <= '0;

            sum_min_nxtmin  <= '0;
            loop_counter    <= '0;

            min_id_reg      <= '0;
            nxt_id_reg      <= '0;
            
        end else begin
            case (state)

                INIT: begin
                    start_nxt  <= 1'b0;
                    start_sort <= 1'b0;
                    if (start) state <= LOAD;
                end

                LOAD: begin
                    // init pool freqs and symbol structures
                    for (i = 0; i < CHAR_SPACE; i++) begin
                        arr_pool_freq[i] <= {1'b0, character_freq_unsort[i]};
                        char_pool_id[i] <= i;    // integer i treated as 32 bit vector
                        char_code[i] <= '0;
                        char_len[i] <= '0;
                    end
                    loop_counter <= '0;
                    state <= SORT;
                end

                SORT: begin
                    // start_sort high until done_sort is observed
                    start_sort <= 1'b1;
                    state <= WAIT_SORT;
                end

                WAIT_SORT: begin
                    start_sort <= 1'b1; 
                    if (done_sort) begin
                        // latch the selected pool ids
                        min_id_reg <= min_pool_id;
                        nxt_id_reg <= nxtmin_pool_id;

                        // now we can drop start_sort next state
                        start_sort <= 1'b0;
                        state <= SUM;
                    end
                end

                SUM: begin
                    sum_min_nxtmin <= arr_pool_freq[min_id_reg] + arr_pool_freq[nxt_id_reg];
                    state <= POOL_UPDATE;
                end

                POOL_UPDATE: begin
                    // Per alorithm :
                    // - smallest becomes inactive
                    // - next-smallest becomes merged node
                    arr_pool_freq[min_id_reg] <= MAX_NUM[FREQ_LEN:0];
                    arr_pool_freq[nxt_id_reg] <= sum_min_nxtmin;
                    state <= CODE_UPDATE;
                end

                CODE_UPDATE: begin
                    // Append 0/1 at LSB (leaf->root storage)
                    // Use current char_pool_id BEFORE merging pool ids.
                    for (i = 0; i < CHAR_SPACE; i++) begin
                        if (char_pool_id[i] == min_id_reg) begin
                            // smallest -> append 0
                            char_code[i] <= (char_code[i] << 1);          // | 0
                            char_len[i]  <= char_len[i] + 1'b1;
                            // merge: move this symbol into nxtmin pool
                            char_pool_id[i] <= nxt_id_reg;
                        end else if (char_pool_id[i] == nxt_id_reg) begin
                            // second smallest -> append 1
                            char_code[i] <= (char_code[i] << 1) | 1'b1;
                            char_len[i]  <= char_len[i] + 1'b1;
                            // stays in nxtmin pool
                        end
                    end

                    // one merge completed
                    loop_counter <= loop_counter + 1'b1;

                    if (loop_counter == (CHAR_SPACE-2)) begin
                        state <= DONE;   // after CHAR_SPACE-1 merges
                    end else begin
                        state <= SORT;
                    end
                end

                DONE: begin
                    huffman_lib_out <= char_code;
                    huffman_len_out <= char_len;
                
                    start_nxt <= 1'b1;   // assert done for 1 full cycle
                    state     <= INIT;   // go back next cycle

                end

                default: state <= INIT;
            endcase
        end
    end
        

endmodule