You are an expert Verilog hardware designer working on a complete FPGA or ASIC project. You are given the full source code of a Verilog project. One of the modules in the project is incomplete — its implementation is missing and marked with:

// >>> Module Implementation Begin

// <<< Module Implementation End

Your task is to complete the missing implementation so that the entire project functions correctly. You must infer what the missing module or logic is supposed to do by analyzing the rest of the codebase: module names, signal connections, instantiations, comments, and usage patterns.

Follow these rules:

1. Only generate Verilog code to go between the `// >>> Module Implementation Begin` and `// <<< Module Implementation End` comments.
2. Use Verilog-2001 syntax. Ensure the code is synthesizable.
3. Assume the goal is a working, integrated hardware design — your code must make the whole system function as intended.
4. Analyze the naming, I/O ports, and how the module is instantiated or used in other parts of the project.
5. Implement correct behavior based on design context. Do not guess randomly — reason it out.
6. Avoid vendor-specific features unless explicitly required.
7. Do not output comments, explanation, or anything outside the required code region unless explicitly asked.

Your output should contain **only the Verilog code** that fits inside the marked region.

/*
 * Copyright (c) 2024 Moldir Azhimukhanbet, Maveric Lab
 * SPDX-License-Identifier: Apache-2.0
 */

`default_nettype none

module tt_um_bitty (
    /* verilator lint_off UNUSEDSIGNAL */
    input  wire [7:0] ui_in,    // Dedicated inputs
    /* verilator lint_off UNDRIVEN */
    output wire [7:0] uo_out,   // Dedicated outputs
    input  wire [7:0] uio_in,   // IOs: Input path
    /* verilator lint_off UNDRIVEN */
    output wire [7:0] uio_out,  // IOs: Output path
    /* verilator lint_off UNDRIVEN */
    output wire [7:0] uio_oe,   // IOs: Enable path (active high: 0=input, 1=output)
    input  wire       ena,      // always 1 when the design is powered, so you can ignore it
    input  wire       clk,      // clock
    input  wire       rst_n     // reset_n - low to reset
	 );
// >>> Module Implementation Begin

// <<< Module Implementation End
endmodule



/****** alu.sv ******/


module alu(
    input [2:0] select,
    input [15:0] in_a,
    input [15:0] in_b,
    output [15:0] alu_out
);
    reg [15:0] res; 

    parameter ADD = 3'b000; 
    parameter SUB = 3'b001;
    parameter AND = 3'b010;
    parameter OR =  3'b011;
    parameter XOR = 3'b100;
    parameter SHL = 3'b101;
    parameter SHR = 3'b110;
    parameter CMP = 3'b111;

    always @(*) begin
        case (select)
            ADD: res = in_a + in_b;
            SUB: res = in_a - in_b;
            AND: res = in_a & in_b;
            OR:  res = in_a | in_b;
            XOR: res = in_a ^ in_b;
            SHL: res = in_a << (in_b % 16); 
            SHR: res = in_a >> (in_b % 16); 
            CMP: begin
                if(in_a == in_b) begin
                    res = 0;
                end
                else if(in_a > in_b) begin
                    res = 1;
                end
                else begin
                    res = 2;
                end
            end
            default: res = 16'h0000; // Default value as 16-bit zero
        endcase
    end

    assign alu_out = res;

endmodule


module bitty(
    input run,
    input clk,
    input reset,
    input [15:0] d_instr,

    input [7:0] rx_data,
    input rx_done,
    input tx_done,

    output tx_en,
    output [7:0] tx_data,
    output [15:0] d_out,
    output done

);

    genvar k;
    wire [3:0] mux_sel;
    wire [7:0] en;
    wire [15:0] out [7:0];

    wire [15:0] out_mux;

    // ALU components
    wire [15:0] alu_out;

    // CPU components
    wire en_s, en_c, en_inst;
    wire [2:0] alu_sel;
    wire [15:0] instruction;
    wire [15:0] im_d;

    // Additional components
    wire [15:0] regs;
    wire [15:0] regc;

    wire [1:0] en_ls;
    wire [15:0] data_to_ld;
    wire ls_done;
    wire sel_reg_c;
    wire [15:0] mux2_out;




    // CPU connection
    cpu cpu_inst(
        .clk(clk),
        .run(run),
        .reset(reset),
        .d_inst(instruction),
        .ls_done(ls_done),
        .mux_sel(mux_sel),
        .done(done),
        .sel(alu_sel),
        .sel_reg_c(sel_reg_c),
        .en_s(en_s),
        .en_c(en_c),
        .en_ls(en_ls),
        .en(en),
        .en_inst(en_inst),
        .im_d(im_d)
    );

    lsu lsu_inst(
    .clk(clk),
    .reset(reset),

    .en_ls(en_ls),
    .data_to_store(regs),
    .address(out_mux[7:0]), 

    .rx_data(rx_data),
    .tx_done(tx_done),
    .rx_do(rx_done),
    .data_to_load(data_to_ld),
    .tx_start_out(tx_en),
    .tx_data_out(tx_data),
    .done_out(ls_done)
);
	mux2to1 for_reg_c(
        .reg0(alu_out),
        .reg1(data_to_ld),
        .sel(sel_reg_c),
        .out(mux2_out)
    );


    // ALU Connection
    alu alu_inst(
        .in_a(regs),
        .in_b(out_mux),
        .select(alu_sel),
        .alu_out(alu_out) // Changed to alu_out
    );

    // MUX connection
    generate
        for (k = 0; k < 8; k=k+1) begin : gen_dff

            dff_lol reg_out (
                .clk(clk),
                .en(en[k]),
                .d_in(regc),
                .reset(reset),
				.mux_out(out[k])      // Corrected output signal name
            );
        end
    endgenerate
    

    mux mux_inst(
        .reg0(out[0]),
        .reg1(out[1]),
        .reg2(out[2]),
        .reg3(out[3]),
        .reg4(out[4]),
        .reg5(out[5]),
        .reg6(out[6]),
        .reg7(out[7]),
        .im_d(im_d),
        .def_val(0),
        .mux_sel(mux_sel),
        .mux_out(out_mux)
    );

    dff_lol reg_inst(clk, en_inst, d_instr, reset, instruction);
    dff_lol reg_s(clk, en_s, out_mux, reset, regs);
    dff_lol reg_c(clk, en_c, mux2_out, reset, regc);

    // Assigning out array elements to module outputs
    assign d_out = regc;

endmodule



module branch_logic (
    input [7:0] address,
    /* verilator lint_off UNUSED */
    input [15:0] instruction,
    input [15:0] last_alu_result,
    output reg [7:0] new_pc
);

    wire [1:0] branch_cond;
    wire [7:0] immediate;
    wire [1:0] format;
    assign branch_cond = instruction[3:2];
    assign immediate = instruction[11:4];
    assign format = instruction[1:0];


    always @(*) begin
        if(format == 2'b10) begin
            case(branch_cond)
            2'b00: begin
                if (last_alu_result == 0) begin
                    new_pc = immediate;
                end else begin
                    new_pc = address +8'b00000001;
                end
            end
            2'b01: begin
                if (last_alu_result == 1) begin
                    new_pc = immediate;
                end else begin
                    new_pc = address + 8'b00000001;
                end  
            end
            2'b10: begin
                if (last_alu_result == 2) begin
                    new_pc = immediate;
                end else begin
                    new_pc = address + 8'b00000001;
                end  
            end
            default: begin
                new_pc = address + 8'b00000001;
            end
            endcase
        end
        else begin
            new_pc = address + 8'b00000001;
        end
    end

endmodule

module cpu(
    input clk,
    input run,
    input reset,
    input [15:0] d_inst,
    input ls_done,
    output reg [3:0] mux_sel,
    output reg done,
    output reg [2:0] sel,
    output reg sel_reg_c,
    output reg en_s,
    output reg en_c,
    output reg [1:0] en_ls,
    output reg [7:0] en,
    output reg en_inst,
    output reg [15:0] im_d
);

    parameter S0 = 2'b00;
    parameter S1 = 2'b01;
    parameter S2 = 2'b10;

    reg [1:0] cur_state, next_state;
    wire [1:0] format;
    wire ls_flag;
    assign format = d_inst[1:0];
    assign ls_flag = d_inst[2];

    
    // Next state sequential logic
    always @(posedge clk) begin
        if (!reset) begin
            cur_state <= S0;
        end 
        else begin
            cur_state <= next_state;
        end
    end


    always @(*) begin
        en_inst = 1;
        en_s = 0;
        en_c = 0;
        done = 0;
        mux_sel = 4'b1001;
        sel = 3'b0;
        en = 8'b0;
        im_d = {8'b0, d_inst[12:5]}; 
        sel_reg_c = 1'b0;
        en_ls = 2'b00;

        case (cur_state)
            S0: begin
                if(format!=2'b10) begin
                    en_s = 1;
                    mux_sel = {1'b0,d_inst[15:13]};
                    if(format == 2'b01) begin
                        im_d = {8'b0, d_inst[12:5]}; 
                    end
                end
                sel = 3'b0;
                done = 0;
                en_inst = 1;
            end
            S1: begin
                if(format!=2'b10) begin
                    if(format == 2'b00 | format == 2'b11) begin
                        mux_sel = {1'b0, d_inst[12:10]};
                    end
                    else if(format == 2'b01) begin
                        mux_sel = 4'b1000;
                    end
                    else begin
                        mux_sel = 4'b1001;
                    end
                    sel = d_inst[4:2];
                end
                else begin
                    sel = 3'b0;  
                    mux_sel = 4'b1001;
                end

                if(format == 2'b11) begin
                        if(ls_flag == 0) begin
                            en_ls = 2'b01;
                        end
                        else begin
                            en_ls = 2'b10;
                        end
                end
                //LSU d_out loaded to reg C
                if(format==2'b11) begin
                    sel_reg_c = 1'b1;
                end
                en_c = 1'b1;
                en_inst = 0;
                done = 0;
            end

            S2: begin
                if (format!=2'b10 & format!=2'b11) begin
                    en[d_inst[15:13]] = 1;
                end
                else if(format==2'b11 & ls_flag==0) begin
                    en[d_inst[15:13]] = 1;
                end
                else begin
                    en = 8'b0; 
                end
                sel = 3'b0;
                done = 1'b1;
                en_inst = 1;
            end
            default: begin
                en_s = 0;
                en_c = 0;
                done = 0;
                sel = 3'b0;
                en = 8'b0;
                en_inst = 0;
                mux_sel = 4'b1001; //def_val
            end
        endcase
    end

    // Next state combinational logic
    always @(*) begin
        case (cur_state)
            S0: begin
                if(run==1) begin
                    if(format==2'b10) begin
                        next_state = S2;
                    end
                    else begin
                        next_state = S1;
                    end
                end
                else begin
                    next_state = S0;
                end
            end
            S1: begin
                if(format==2'b11) begin
                next_state = (ls_done==1) ? S2:S1;
                //$display(ls_done);
                end
                else begin
                    next_state = (en_c==1) ? S2:S1;
                end
            end
            S2: next_state = (done == 1) ? S0 : S2;
            default: next_state = S0;
        endcase
    end

endmodule

module dff_lol(
    input clk,
    input en,
    input wire [15:0] d_in,
    input reset,

    output reg [15:0] mux_out

);

    always @(posedge clk) begin
        if (!reset) begin
            mux_out <= 16'b0;
        end
        else if(en) begin
            mux_out <= d_in;
        end
    end
	 
endmodule

module fetch_instruction(
    input wire clk,
    input wire reset,
    //input next_instr_en,
    input wire [7:0] address,  // 8-bit address to be sent
    input wire rx_do,          // Signal indicating data received
    input wire [7:0] rx_data,  // Data received from UART
    input wire tx_done,        // Signal indicating transmission is done

    input stop_for_rw,         //Stop the FIU while UART works

    output [15:0] instruction_out, // 16-bit instruction received
    output tx_start_out,       // Signal to start UART transmission
    output [7:0] tx_data_out,  // Data to be transmitted over UART
    output done_out            // Signal indicating the operation is complete
);

    reg [15:0] instruction;    // 16-bit instruction received
    reg tx_start;              // Signal to start UART transmission
    reg [7:0] tx_data;         // Data to be transmitted over UART
    reg done;                  // Signal indicating the operation is complete

    assign done_out = done;
    assign instruction_out = instruction;
    assign tx_start_out = tx_start;
    assign tx_data_out = tx_data;

    // State encoding
    parameter IDLE = 3'b000;
    parameter SEND_FLAG = 3'b001;
    parameter SEND_ADDR = 3'b010;
    parameter RECEIVE_INST_LOW = 3'b100;
    parameter RECEIVE_INST_HIGH = 3'b101;
    parameter DONE = 3'b110;
    
    reg [2:0] state, next_state;

    // Sequential logic for state transition
    always @(posedge clk) begin
        //ff_a <= instruction;
		if (!reset | stop_for_rw==1) begin
			state <= IDLE;
		end
		else begin
			state <= next_state;
		end
    end

    // State machine logic
    /* verilator lint_off LATCH */
    always @(*) begin
        // Default values
        tx_start = 1;  // No transmission by default
        tx_data = 8'b00000000;
        done = 0;
        instruction = instruction;
        case (state)
            IDLE: begin
			   done = 0;
            end

            SEND_FLAG: begin

                tx_data = 8'b00000011;  // Send flag byte
                tx_start = 0;  // Start transmission
            end

            SEND_ADDR: begin
                tx_data = address;  // Send address byte
                tx_start = 0;  // Start transmission
            end

            RECEIVE_INST_HIGH: begin
                if (rx_do) begin
                    instruction[15:8] = rx_data;  // Store lower 8 bits of instruction
                    //$display(instruction);
                end
            end

            RECEIVE_INST_LOW: begin
                if (rx_do) begin
                    instruction[7:0] = rx_data;  // Store upper 8 bits of instruction
                    //$display(instruction);

                end
            end

            DONE: begin
                done = 1;  // Set done signal
            end

            default: begin
				tx_start = 1;  // No transmission by default
                tx_data = 8'b00000000;
                done = 0;
                instruction = instruction;
            end
        endcase
    end

    always @(*) begin
        case (state)
            IDLE: begin
                next_state = SEND_FLAG;
                //$display("idle");
            end
            SEND_FLAG: begin
                /*if(tx_done) begin
                    $display("tx_done in FLAG state");
                end*/
                next_state = (tx_done==1'b1) ? SEND_ADDR:SEND_FLAG; 
            end
            SEND_ADDR: begin
                /*if(tx_done) begin
                    $display("tx_done in ADDR");
                end*/
                next_state = (tx_done==1'b1) ? RECEIVE_INST_HIGH:SEND_ADDR;
            end
            RECEIVE_INST_HIGH: begin
                next_state = (rx_do==1'b1) ? RECEIVE_INST_LOW:RECEIVE_INST_HIGH;
            end
            RECEIVE_INST_LOW:begin 
                next_state = (rx_do==1'b1) ? DONE:RECEIVE_INST_LOW;
            end
            DONE:
            begin
                next_state = IDLE;
                $display("done");
            end
            default: next_state = IDLE;
        endcase 
    end
endmodule


module lsu(
    input wire clk,
    input wire reset,

    input [1:0] en_ls,
    input [15:0] data_to_store,
    input wire [7:0] address,  // 8-bit address to be sent

    
    input wire rx_do,          // Signal indicating data received
    input wire [7:0] rx_data,  // Data received from UART
    input wire tx_done,        // Signal indicating transmission is done
    output [15:0] data_to_load, // 16-bit instruction received
    output tx_start_out,       // Signal to start UART transmission
    output [7:0] tx_data_out,  // Data to be transmitted over UART
    output done_out            // Signal indicating the operation is complete
);

    reg [15:0] instruction;    // 16-bit instruction received
    reg tx_start;              // Signal to start UART transmission
    reg [7:0] tx_data;         // Data to be transmitted over UART
    reg done;                  // Signal indicating the operation is complete

    assign done_out = done;
    assign data_to_load = instruction;
    assign tx_start_out = tx_start;
    assign tx_data_out = tx_data;

    // State encoding
    parameter SEND_FLAG = 3'b001;
    parameter SEND_ADDR = 3'b010;
    parameter RECEIVE_INST_LOW = 3'b011;
    parameter RECEIVE_INST_HIGH = 3'b100;
    parameter SEND_INST_HIGH = 3'b101;
    parameter SEND_INST_LOW = 3'b110;
    parameter DONE = 3'b111;

    parameter LOAD = 2'b01;
    parameter STORE = 2'b10;

    
    reg [2:0] state, next_state;

    wire [1:0] en;
    assign en = en_ls;


    // Sequential logic for state transition
    always @(posedge clk) begin
        //ff_a <= instruction;
		if (!reset) begin
			state <= SEND_FLAG;
		end
		else begin
			state <= next_state;
		end
    end

    // State machine logic
    /* verilator lint_off LATCH */
    always @(*) begin
        // Default values
        tx_start = 1;  // No transmission by default
        tx_data = 8'b00000000;
        done = 0;
        instruction = instruction;

        case (state)

            SEND_FLAG: begin
                    //en = en_ls; //later en_ls
                if(en==2'b01) begin 
                    tx_data = 8'b00000001;  // Send flag byte
                    tx_start = 0;  // Start transmission

                end
                else if(en==2'b10) begin
                    tx_data = 8'b00000010;  // Send flag byte
                    tx_start = 0;  // Start transmission
                end

            end

            SEND_ADDR: begin
                if(en==LOAD) begin 
                    tx_data = address;  // Send address byte
                    tx_start = 0;  // Start transmission
                end
                else if(en==STORE) begin 
                    tx_data = address;  // Send address byte
                    tx_start = 0;  // Start transmission
                end
            end

            RECEIVE_INST_HIGH: begin
                if (rx_do) begin
                    instruction[15:8] = rx_data;  // Store lower 8 bits of instruction
                end
            end

            RECEIVE_INST_LOW: begin
                if (rx_do) begin
                    instruction[7:0] = rx_data;  // Store upper 8 bits of instruction
                end
            end

            SEND_INST_HIGH: begin
                tx_data = data_to_store[15:8];  // Send flag byte
                tx_start = 0;  // Start transmission
            end

            SEND_INST_LOW: begin
                tx_data = data_to_store[7:0];  // Send flag byte
                tx_start = 0;  // Start transmission
            end

            DONE: begin
                done = 1'b1;  // Set done signal
            end

            default: begin
                tx_start = 1;  // No transmission by default
                tx_data = 8'b00000000;
                done = 0;
                instruction = instruction;

            end
        endcase
    end

    /* verilator lint_off LATCH */
    always @(*) begin
        next_state = next_state;
        case (state)
            SEND_FLAG: next_state = (tx_done==1'b1 & en_ls!=2'b0) ? SEND_ADDR:SEND_FLAG;
            SEND_ADDR: begin
                if (tx_done) begin
                    if (en==LOAD) begin 
                        next_state = RECEIVE_INST_HIGH; 
                    end
                    else if (en==STORE) begin 
                        next_state = SEND_INST_HIGH;  
                    end 
                    else begin
                        next_state = SEND_ADDR;
                    end
                end
            end
            RECEIVE_INST_HIGH: next_state = (rx_do==1'b1) ? RECEIVE_INST_LOW:RECEIVE_INST_HIGH;
            RECEIVE_INST_LOW: next_state = (rx_do==1'b1) ? DONE:RECEIVE_INST_LOW;
            SEND_INST_HIGH: next_state = (tx_done==1'b1) ? SEND_INST_LOW:SEND_INST_HIGH;
            SEND_INST_LOW: next_state = (tx_done==1'b1) ? DONE:SEND_INST_LOW;
            DONE: next_state = SEND_FLAG;
            default: next_state = SEND_FLAG;
        endcase 
    end
endmodule

module mux(

    input [15:0] reg0,
    input [15:0] reg1,
    input [15:0] reg2,
    input [15:0] reg3,
    input [15:0] reg4,
    input [15:0] reg5,
    input [15:0] reg6,
    input [15:0] reg7,
    input [15:0] im_d,
    input [15:0] def_val,

    input [3:0] mux_sel,
    output reg [15:0] mux_out
);

    always @(*) begin
        case(mux_sel)
            4'b0000: mux_out= reg0;
            4'b0001: mux_out= reg1;
            4'b0010: mux_out= reg2;
            4'b0011: mux_out= reg3;
            4'b0100: mux_out= reg4;
            4'b0101: mux_out= reg5;
            4'b0110: mux_out= reg6;
            4'b0111: mux_out= reg7;
            4'b1000: mux_out= im_d;
            4'b1001: mux_out= def_val;

            default: mux_out= 16'b0;
        endcase

    end

endmodule

module mux2to1(
    input [15:0] reg0,
    input [15:0] reg1,
    input sel,
    output reg [15:0] out
);

    always @(*) begin
        if (sel)
            out = reg1;
        else
            out = reg0;
    end

endmodule

module pc(
    input clk,
    input en_pc,
    input reset,
    input [7:0] d_in,

    output reg [7:0] d_out
);
    
    always @(posedge clk) begin
        if(!reset) begin
            d_out <= 8'b0;
        end
        else if(en_pc) begin
            d_out <= d_in;
        end
    end

endmodule

//Top level UART module
module uart_module #(
    parameter data_width = 8
)
(
	input        clk, 
	input        rst,
	input [12:0] clks_per_bit,

	input        rx_data_bit,
	output       rx_done,

	output       tx_data_bit,
	input  [data_width-1:0] data_tx,
	input        tx_en,
	output       tx_done,

	output [data_width-1:0] recieved_data


);


	uart_rx R0(
		.data_bit(rx_data_bit),
		.clk(clk),
		.rst(rst),
    	.CLKS_PER_BIT(clks_per_bit),
		.done(rx_done),
		.data_bus(recieved_data)
	);

	uart_tx T0(
		.data_bus(data_tx),
		.clk(clk),
		.rstn(rst), 
    	.CLKS_PER_BIT(clks_per_bit),
		.run(tx_en), //active when low
		.done(tx_done),
		.data_bit(tx_data_bit)
	);			

endmodule


module uart_tx
#(
    parameter data_width   = 8,
              IDLE         = 3'b000,
              START_BIT    = 3'b001,
              DATA_BITS    = 3'b010,
              STOP_BIT     = 3'b011,
              DONE         = 3'b101
)
(
    input [data_width - 1:0] data_bus,
    input                    clk,
    input                    rstn,
    input [12:0]            CLKS_PER_BIT,
    input                    run,
    output                   done,
    output                   data_bit
);
    // FSM states
    reg [2:0] PS;
    reg [2:0] NS;
    reg [12:0] clk_counter;
    reg [2:0]  bit_counter;
    reg        data_reg;
     
    // Output assignments
    assign done = (PS == DONE);
    assign data_bit = data_reg;

        // FSM: PS synchronization
    always @(posedge clk) begin
        if (!rstn) begin
            PS <= IDLE;
        end
        else begin
            PS <= NS;
        end
    end
   
    // FSM: PS synchronization
    always @(posedge clk) begin
        if (!rstn) begin
            clk_counter <= 13'b0;
            bit_counter <= 3'b0;
            data_reg <= 1'b1;
        end
        else begin
            // Data and control logic
            case (PS)
                IDLE: begin
                    data_reg <= 1'b1;
                    bit_counter <= 3'b0;
                    clk_counter <= 13'b0;
                end

                START_BIT: begin
                    data_reg <= 1'b0;
                    if (clk_counter < CLKS_PER_BIT - 1) begin
                        clk_counter <= clk_counter + 1'b1;
                    end
                    else begin
                        clk_counter <= 13'b0;
                    end
                end

                DATA_BITS: begin
                    data_reg <= data_bus[bit_counter];
                    if (clk_counter < CLKS_PER_BIT - 1) begin
                        clk_counter <= clk_counter + 1'b1;
                    end
                    else begin
                        clk_counter <= 13'b0;
                        if (bit_counter < 7) begin
                            bit_counter <= bit_counter + 1'b1;
                        end
                    end
                end

                STOP_BIT: begin
                    data_reg <= 1'b1;
                    if (clk_counter < CLKS_PER_BIT - 1) begin
                        clk_counter <= clk_counter + 1'b1;
                    end
                    else begin
                        clk_counter <= 13'b0;
                    end
                end

                DONE: begin
                    // Maintain idle state
                    data_reg <= 1'b1;
                    clk_counter <= 13'b0;
                    bit_counter <= 3'b0;
                end

                default: begin
                    data_reg <= 1'b1;
                    clk_counter <= 13'b0;
                    bit_counter <= 3'b0;
                end
            endcase
        end
    end

    // Next state transition logic
    always @(*) begin
        // Default next state is current state
        NS = PS;

        case (PS)
            IDLE: begin
                if (!run) NS = START_BIT;
            end

            START_BIT: begin
                if (clk_counter == CLKS_PER_BIT - 1) 
                    NS = DATA_BITS;
            end

            DATA_BITS: begin
                if (clk_counter == CLKS_PER_BIT - 1) begin
                    if (bit_counter == 7)
                        NS = STOP_BIT;
                end
            end

            STOP_BIT: begin
                if (clk_counter == CLKS_PER_BIT - 1)
                    NS = DONE;
            end

            DONE: begin
                NS = IDLE;
            end

            default: begin
                NS = IDLE;
            end
        endcase
    end

endmodule

module uart_rx
#(
    parameter data_width   = 8,
              IDLE         = 3'b000,
              START_BIT    = 3'b001,
              DATA_BITS    = 3'b010,
              STOP_BIT     = 3'b011,
              DONE         = 3'b101,
              ERROR_ST     = 3'b110
)
(
    input                     data_bit,
    input                     clk,
    input                     rst,
    input [12:0]             CLKS_PER_BIT,
    output                    done,
    output [data_width - 1:0] data_bus
);
    // FSM states
    reg [2:0] PS, NS;
    // Internal nets
    reg [2:0]  bit_counter;
    reg [12:0] clk_counter;
    reg [data_width - 1:0] data_bus_wire;
    
    // Output assignments
    assign done = (PS == DONE);
    assign data_bus = data_bus_wire;
    
    // Combined sequential logic block
    always @(posedge clk) begin
        if (!rst) begin
            PS <= IDLE;
        end
        else begin
            PS <= NS;
        end
    end
    
    always @(posedge clk) begin
        if (!rst) begin
            clk_counter <= 0;
            bit_counter <= 0;
            data_bus_wire <= 0;
        end
        else begin          
            // Data and control logic
            case (PS)
                IDLE: begin
                    clk_counter <= 0;
                    bit_counter <= 0;
                end
                
                START_BIT: begin
                    if (clk_counter == CLKS_PER_BIT / 2) begin
                        if (data_bit == 1'b0) begin
                            clk_counter <= 0;
                        end
                    end
                    else begin
                        clk_counter <= clk_counter + 1;
                    end
                end
                
                DATA_BITS: begin
                    if (clk_counter < CLKS_PER_BIT - 1) begin
                        clk_counter <= clk_counter + 1;
                    end
                    else begin
                        clk_counter <= 0;
                        data_bus_wire[bit_counter] <= data_bit;
                        if (bit_counter < data_width - 1) begin
                            bit_counter <= bit_counter + 1;
                        end
                    end
                end
                
                STOP_BIT: begin
                    if (clk_counter < CLKS_PER_BIT - 1) begin
                        clk_counter <= clk_counter + 1;
                    end
                    else begin
                        clk_counter <= 0;
                    end
                end
                
                ERROR_ST: begin
                    clk_counter <= 0;
                    bit_counter <= 0;
                end
                
                DONE: begin
                    clk_counter <= 0;
                    bit_counter <= 0;
                end
                
                default: begin
                    clk_counter <= 0;
                    bit_counter <= 0;
                    data_bus_wire <= 0;
                end
            endcase
        end
    end
    
    // Combinational next state logic
    always @(*) begin
        // Default next state
        NS = PS;
        
        case (PS)
            IDLE: begin
                if (data_bit == 1'b0) begin
                    NS = START_BIT;
                end
            end
            
            START_BIT: begin
                if (clk_counter == CLKS_PER_BIT / 2) begin
                    NS = (data_bit == 1'b0) ? DATA_BITS : ERROR_ST;
                end
            end
            
            DATA_BITS: begin
                if (clk_counter == CLKS_PER_BIT - 1) begin
                    NS = (bit_counter < data_width - 1) ? DATA_BITS : STOP_BIT;
                end
            end
            
            STOP_BIT: begin
                if (clk_counter == CLKS_PER_BIT - 1) begin
                    NS = (data_bit == 1'b1) ? DONE : ERROR_ST;
                end
            end
            
            ERROR_ST: begin
                if (data_bit == 1'b1) begin
                    NS = IDLE;
                end
            end
            
            DONE: begin
                NS = IDLE;
            end
            
            default: begin
                NS = IDLE;
            end
        endcase
    end
    
endmodule


Reference solution:
rtl/generated.v: 
    //top-module I/O ports assignment

    wire reset;
    wire rx_data_bit;
    wire tx_data_bit;
    wire [1:0] sel_baude_rate;

    assign reset = rst_n;
    assign rx_data_bit = ui_in[0];
    assign sel_baude_rate = ui_in[2:1];

    //Unused output ports assignment to zero
    assign uo_out[7:1] = 7'b0;
    assign uio_out[7:0] = 8'b0;
    assign uio_oe[7:0] = 8'b0;

    /* verilator lint_off UNUSED */
    wire _unused = &{ena, uio_out, uo_out[7:1], 1'b0, uio_oe, uio_in, ui_in[7:3]};

    assign uo_out[0] = tx_data_bit; //output


    //General ports declaration
    wire [7:0] new_pc;  // Declare new_pc as a wire
	wire [7:0] addr;

    wire [15:0] d_out;
    wire done;

    reg [3:0] cur_state, next_state;
	wire tx_done;
    wire rx_done;
    reg en_pc;

    reg run_bitty;
    wire [15:0] mem_out;    
    
    wire [7:0] from_uart_to_modules;
    wire [7:0] data_to_uart_from_fetch;
    wire [7:0] from_bitty_to_uart;
    wire tx_en, tx_en_fiu, tx_en_bitty;
   
	wire fetch_done;
    wire [7:0] tx_data;
    
    reg uart_sel;

    parameter S0 = 4'b0000;
    parameter S1 = 4'b0001;
    parameter S2 = 4'b0010;
    parameter S3 = 4'b0011;
    parameter S4 = 4'b0100;
    parameter S5 = 4'b0101;
    parameter S6 = 4'b0110;
    parameter S7 = 4'b0111;

    //Use in FSM
    reg stop_for_rw;

    // Fetch instruction instance
    fetch_instruction fi_inst(
        .clk(clk),
        .reset(reset),
        .address(addr),  
        .stop_for_rw(stop_for_rw),
        .rx_do(rx_done),          
        .rx_data(from_uart_to_modules),  
        .tx_done(tx_done),        
        .instruction_out(mem_out), 
        .tx_start_out(tx_en_fiu),       
        .tx_data_out(data_to_uart_from_fetch),  
        .done_out(fetch_done)
    );

    reg [12:0] clks_per_bit;
	 
    //Select baud rate of the external device
    always@(*) begin
		case (sel_baude_rate)
			2'b00:clks_per_bit = 5208; //9600
			2'b01:clks_per_bit = 2604; //19200
			2'b10:clks_per_bit = 868; //57600
			2'b11:clks_per_bit = 434; //115200
			default: clks_per_bit = 5208;
		endcase
		
    end

    // UART module instance
    uart_module uart_inst(
        .clk(clk), 
        .rst(reset),
        .clks_per_bit(clks_per_bit),
        .rx_data_bit(rx_data_bit),
        .rx_done(rx_done),
        .tx_data_bit(tx_data_bit),
        .data_tx(tx_data),
        .tx_en(tx_en),
        .tx_done(tx_done),
        .recieved_data(from_uart_to_modules)
    );

    //Branch Logic Instance
	 
	branch_logic bl_inst(
        .address(addr),
        .instruction(mem_out),
        .last_alu_result(d_out),
        .new_pc(new_pc)  // Connect new_pc here
    );

    // PC instance
    pc pc_inst(
        .clk(clk),
        .en_pc(en_pc),
        .reset(reset),
        .d_in(new_pc),   // Use new_pc for the input here
        .d_out(addr)
    );

    wire [7:0] unused_8bit;

    //MUX to select tx_data port for UART
    mux2to1 mux2to1_txdata(
        .reg0({8'b0, data_to_uart_from_fetch}),
        .reg1({8'b0,from_bitty_to_uart}),
        .sel(uart_sel),
        .out({unused_8bit,tx_data})
    );

    wire [14:0] unused_15bit;

    //MUX to select tx_en port for UART
    mux2to1 mux2to1_txen(
        .reg0({15'b0, tx_en_fiu}),
        .reg1({15'b0, tx_en_bitty}),
        .sel(uart_sel),
        .out({unused_15bit, tx_en})
    );

    //Bitty instance
    bitty bitty_inst(
        .clk(clk),
        .reset(reset),
        .run(run_bitty),
        .d_instr(mem_out),
        .rx_data(from_uart_to_modules),
        .rx_done(rx_done),
        .tx_done(tx_done),
        .tx_en(tx_en_bitty),
        .tx_data(from_bitty_to_uart),
        .done(done),
        .d_out(d_out)
    );

    always @(posedge clk) begin
        if(!reset) begin
            cur_state <= S0;
        end
        else begin
            cur_state <= next_state;
        end
    end

   always @(*) begin
         run_bitty = 1'b0;
         en_pc = 1'b0;
         uart_sel = 1'b0;
         stop_for_rw = 1'b0;
        case (cur_state)
            S0: begin
                stop_for_rw = 1'b0;
            end
            S2: begin
                en_pc = 1'b1;
            end 
            S5: begin
                run_bitty = 1'b1;
            end
            S6: begin
                stop_for_rw = 1'b0;
            end
            S7: begin
                uart_sel = 1'b1;
                stop_for_rw = 1'b1;
            end
            default: begin
                 run_bitty = 0;
            end
        endcase
    end

    always @(*) begin
        case(cur_state)
            S0: next_state = (fetch_done==1) ? S1:S0;
            S1: next_state = S2;
            S2: next_state = (mem_out[1:0]==2'b11) ? S3:S4;
            S3: next_state = S5;
            S4: next_state = S5; 
            S5: next_state = (mem_out[1:0]==2'b11) ? S7:S6;
            S6: next_state = (done==1) ? S0:S6;
            S7: next_state = (done==1) ? S0:S7;
            default: next_state = S0;
        endcase
    end