`include "VX_platform.vh"

// Fast encoder using parallel prefix computation
// Adapter from BaseJump STL: http://bjump.org/data_out.html

module VX_onehot_encoder #(
    parameter N       = 1,    
    parameter REVERSE = 0,
    parameter FAST    = 1,
    parameter LN      = `LOG2UP(N)
) (
    input wire [N-1:0]   data_in,    
    output wire [LN-1:0] data_out,
    output wire          valid
); 
    if (N == 1) begin

        assign data_out = data_in;
        assign valid    = data_in;

    end else if (N == 2) begin

        assign data_out = data_in[!REVERSE];
        assign valid    = (| data_in);

    end else if (N == 4)  begin            
            
        reg [LN-1:0] index_r;

        if (REVERSE) begin
            always @(*) begin
                casez (data_in)
                4'b1000: index_r = LN'(0);
                4'b?100: index_r = LN'(1);
                4'b??10: index_r = LN'(2);
                4'b???1: index_r = LN'(3);
                default: index_r = 'x;
                endcase
            end
        end else begin        
            always @(*) begin
                casez (data_in)
                4'b0001: index_r = LN'(0);
                4'b001?: index_r = LN'(1);
                4'b01??: index_r = LN'(2);
                4'b1???: index_r = LN'(3);
                default: index_r = 'x;
                endcase
            end
        end

        assign data_out = index_r;
        assign valid    = (| data_in);

    end else if (N == 8)  begin

        reg [LN-1:0] index_r;

        if (REVERSE) begin
            always @(*) begin
                casez (data_in)
                8'b10000000: index_r = LN'(0);
                8'b?1000000: index_r = LN'(1);
                8'b??100000: index_r = LN'(2);
                8'b???10000: index_r = LN'(3);
                8'b????1000: index_r = LN'(4);
                8'b?????100: index_r = LN'(5);
                8'b??????10: index_r = LN'(6);
                8'b???????1: index_r = LN'(7);
                default:     index_r = 'x;
                endcase
            end
        end else begin        
            always @(*) begin
                casez (data_in)
                8'b00000001: index_r = LN'(0);
                8'b0000001?: index_r = LN'(1);
                8'b000001??: index_r = LN'(2);
                8'b00001???: index_r = LN'(3);
                8'b0001????: index_r = LN'(4);
                8'b001?????: index_r = LN'(5);
                8'b01??????: index_r = LN'(6);
                8'b1???????: index_r = LN'(7);
                default:     index_r = 'x;
                endcase
            end
        end

        assign data_out = index_r;
        assign valid    = (| data_in);

    end else if (N == 16)  begin

        reg [LN-1:0] index_r;

        if (REVERSE) begin
            always @(*) begin
                casez (data_in)
                16'b1000000000000000: index_r = LN'(0);
                16'b?100000000000000: index_r = LN'(1);
                16'b??10000000000000: index_r = LN'(2);
                16'b???1000000000000: index_r = LN'(3);
                16'b????100000000000: index_r = LN'(4);
                16'b?????10000000000: index_r = LN'(5);
                16'b??????1000000000: index_r = LN'(6);
                16'b???????100000000: index_r = LN'(7);
                16'b????????10000000: index_r = LN'(8);
                16'b?????????1000000: index_r = LN'(9);
                16'b??????????100000: index_r = LN'(10);
                16'b???????????10000: index_r = LN'(11);
                16'b????????????1000: index_r = LN'(12);
                16'b?????????????100: index_r = LN'(13);
                16'b??????????????10: index_r = LN'(14);
                16'b???????????????1: index_r = LN'(15);
                default:              index_r = 'x;
                endcase
            end
        end else begin        
            always @(*) begin
                casez (data_in)
                16'b0000000000000001: index_r = LN'(0);
                16'b000000000000001?: index_r = LN'(1);
                16'b00000000000001??: index_r = LN'(2);
                16'b0000000000001???: index_r = LN'(3);
                16'b000000000001????: index_r = LN'(4);
                16'b00000000001?????: index_r = LN'(5);
                16'b0000000001??????: index_r = LN'(6);
                16'b000000001???????: index_r = LN'(7);
                16'b00000001????????: index_r = LN'(8);
                16'b0000001?????????: index_r = LN'(9);
                16'b000001??????????: index_r = LN'(10);
                16'b00001???????????: index_r = LN'(11);
                16'b0001????????????: index_r = LN'(12);
                16'b001?????????????: index_r = LN'(13);
                16'b01??????????????: index_r = LN'(14);
                16'b1???????????????: index_r = LN'(15);
                default:              index_r = 'x;
                endcase
            end
        end

        assign data_out = index_r;
        assign valid    = (| data_in);

    end if (FAST) begin
    `IGNORE_WARNINGS_BEGIN
        localparam levels_lp = $clog2(N);
        localparam aligned_width_lp = 1 << $clog2(N);
    
        wire [levels_lp:0][aligned_width_lp-1:0] addr;
        wire [levels_lp:0][aligned_width_lp-1:0] v; 
    
        // base case, also handle padding for non-power of two inputs
        assign v[0] = REVERSE ? (data_in << (aligned_width_lp - N)) : ((aligned_width_lp)'(data_in));
        assign addr[0] = 'x;
    
        for (genvar level = 1; level < levels_lp+1; level=level+1) begin
            localparam segments_lp      = 2**(levels_lp-level);
            localparam segment_slot_lp  = aligned_width_lp/segments_lp;
            localparam segment_width_lp = level; // how many bits are needed at each level
        
            for (genvar segment = 0; segment < segments_lp; segment=segment+1) begin
                wire [1:0] vs = {
                    v[level-1][segment*segment_slot_lp+(segment_slot_lp >> 1)],
                    v[level-1][segment*segment_slot_lp]
                };
            
                assign v[level][segment*segment_slot_lp] = (| vs);

                if (level == 1) begin
                    assign addr[level][(segment*segment_slot_lp)+:segment_width_lp] = vs[!REVERSE]; 
                end else begin
                    assign addr[level][(segment*segment_slot_lp)+:segment_width_lp] = { 
                        vs[!REVERSE],
                        addr[level-1][segment*segment_slot_lp+:segment_width_lp-1] | addr[level-1][segment*segment_slot_lp+(segment_slot_lp >> 1)+:segment_width_lp-1]
                    };
                end        
            end  
        end	
    
        assign data_out = addr[levels_lp][`LOG2UP(N)-1:0];
        assign valid = v[levels_lp][0];
    `IGNORE_WARNINGS_END
    end else begin 

        reg [LN-1:0] index_r;

        if (REVERSE) begin

            always @(*) begin        
                index_r = 'x; 
                for (integer i = N-1; i >= 0; --i) begin
                    if (data_in[i]) begin                
                        index_r = `LOG2UP(N)'(i);
                    end
                end
            end

        end else begin
            always @(*) begin        
                index_r = 'x; 
                for (integer i = 0; i < N; i++) begin
                    if (data_in[i]) begin                
                        index_r = `LOG2UP(N)'(i);
                    end
                end
            end
        end

        assign data_out = index_r;
        assign valid    = (| data_in);
    end

endmodule