Demilade Sonuga's blog

All posts

Refactoring IV

2022-12-21

In the previous post, we finally got block.bmp on screen, but in the process of doing this, we convoluted our efi_main, again. What we're going to do now is to rectify that.

We start by taking a critical look at the code:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let block_bytes = include_bytes!("./block.bmp");

    let block_bytes_ptr: *const u8 = block_bytes.as_ptr();
    let file_header_ptr = block_bytes_ptr as *const FileHeader;
    let file_header = unsafe { &(*file_header_ptr) };
    const FILE_HEADER_SIZE: usize = core::mem::size_of::<FileHeader>();
    const DIB_HEADER_OFFSET: isize = FILE_HEADER_SIZE as isize;
    let dib_header_ptr = unsafe { block_bytes_ptr.offset(DIB_HEADER_OFFSET) as *const DIBHeader };
    let dib_header = unsafe { &(*dib_header_ptr) };
    const DIB_HEADER_SIZE: usize = core::mem::size_of::<DIBHeader>();
    const COLOR_TABLE_OFFSET: isize = (FILE_HEADER_SIZE + DIB_HEADER_SIZE) as isize;
    let color_table_ptr = unsafe { block_bytes_ptr.offset(COLOR_TABLE_OFFSET) as *const ColorTable };
    let color_table = unsafe { &(*color_table_ptr) };

    let pixel_array = &block_bytes[file_header.pixel_array_offset as usize..];

    let image_height = dib_header.image_height as usize;
    let image_width = dib_header.image_width as usize;
    for row in 0..image_height {
        for col in 0..image_width {
            let inverted_row = image_height - row - 1;
            let color_table_index = pixel_array[inverted_row * image_width + col];
            let color = color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }

    0
}

The next step is to ask ourselves: what exactly do these lines of code mean? You should try to state the instructions this code is carrying out in at most 3 lines before reading on.

We can break down this code into 3 sections:

The first is this line: let block_bytes = include_bytes!("./block.bmp"); which simply means "get the bitmap".

The second are these lines:

let block_bytes_ptr: *const u8 = block_bytes.as_ptr();
let file_header_ptr = block_bytes_ptr as *const FileHeader;
let file_header = unsafe { &(*file_header_ptr) };
const FILE_HEADER_SIZE: usize = core::mem::size_of::<FileHeader>();
const DIB_HEADER_OFFSET: isize = FILE_HEADER_SIZE as isize;
let dib_header_ptr = unsafe { block_bytes_ptr.offset(DIB_HEADER_OFFSET) as *const DIBHeader };
let dib_header = unsafe { &(*dib_header_ptr) };
const DIB_HEADER_SIZE: usize = core::mem::size_of::<DIBHeader>();
const COLOR_TABLE_OFFSET: isize = (FILE_HEADER_SIZE + DIB_HEADER_SIZE) as isize;
let color_table_ptr = unsafe { block_bytes_ptr.offset(COLOR_TABLE_OFFSET) as *const ColorTable };
let color_table = unsafe { &(*color_table_ptr) };

let pixel_array = &block_bytes[file_header.pixel_array_offset as usize..];

let image_height = dib_header.image_height as usize;
let image_width = dib_header.image_width as usize;

Taking a very close look at this, you'll see that these lines of code all perform a similar task. They are initializing structures in memory which will make it easier to deal with the bitmap. The meaning of these lines is simply "initialize structures for manipulating the bitmap".

The last section:

for row in 0..image_height {
    for col in 0..image_width {
        let inverted_row = image_height - row - 1;
        let color_table_index = pixel_array[inverted_row * image_width + col];
        let color = color_table.0[color_table_index as usize];
        screen.pixels[row][col] = Pixel {
            red: color.red,
            green: color.green,
            blue: color.blue,
            reserved: 0
        };
    }
}

This code takes the bytes in the pixel array, translates them to pixels and puts them on screen. The essence of this code is "draw bitmap on screen".

From this breakdown, all those lines of code are simply saying these 3 things:

  1. Get the bitmap's bytes
  2. Initialize structures for dealing with the bitmap.
  3. Draw the bitmap on the screen.

Now, have to think of how to transform those many lines of code from telling how things are done to telling what things are done instead. The meaning of the code should be clear.

The first section, let block_bytes = include_bytes!("./block.bmp"); means get bitmap's bytes. I think this line effectively communicates that, so we don't need to change it. But, you do need to remember that there are no absolutes here. There is no absolute "right" or "wrong" way of doing this. As long as what you're doing makes sense. If that line seems not to adequately communicate meaning to you, then change it.

The second section initializes a bunch of structures for dealing with bitmaps. Instead of dealing with all these structures directly, we can combine them into a single structure Bitmap, which we'll then use for initializing and handling bitmaps.

In bitmap.rs:

// A structure for handling bitmaps
pub struct Bitmap {
    // The bitmap's file header
    pub file_header: &'static FileHeader,
    // The bitmap's DIB header
    pub dib_header: &'static DIBHeader,
    // The bitmap's color table
    pub color_table: &'static ColorTable,
    // The bitmap's pixel array
    pub pixel_array: &'static [u8]
}

In this structure, there is a field for all the structures we initialized. Notice that there is no #[repr(C)] on top of this struct. This is because the validity of our code doesn't depend on the field order. We aren't using this to interpret anything in memory. We're just grouping structures to create a single structure that will be easier to handle.

In Rust, all references in struct definitions must have a lifetime so the compiler can use those lifetimes to ensure we don't use a reference when it's no longer valid. All the fields here have a 'static lifetime because they are all valid throughout the lifetime of the program. The reason why they're always valid is because of include_bytes!. Remember that this macro takes the bytes in the file specified and places it in the output binary, returning a slice to the bytes in the binary. Since from the beginning to the end of a run of the program, those bytes will be in the same location in the binary, anything that refers to them will also remain valid throughout the program run.

To fully understand this 'static lifetime, we have to properly understand lifetimes and why they're needed.

Lifetimes

Consider the following scenario:

fn main() {
    let another_array = x();
    y();
    println!("{:?}", another_array);
}

fn x() -> &[u8; 5] {
    let another_array: [u8; 5] = [1, 2, 3, 4, 5];
    return &another_array
}

fn y() {
    let yet_another_array: [u8; 5] = [0; 5];
}

The above doesn't compile and for very good reasons. To understand why imagine that the Rust compiler accepts such code and let's see what that will lead to.

main is executing and x is called.

Let's say that before x is called, the stack looks like this:

Address Stack
1000    0       <- Stack top
999     0
998     0
...

And the top of the stack is at address 1000. The next values which will be going on the stack will go into addresses 1000, 999, 988, and so on.

During main's execution, when x is called, another_array will be initialized on the stack and look something like this:

Address Stack
1000    1
999     2
998     3
997     4
996     5       <- Stack top

The &another_array gives the location in memory where another_array dwells, so this will evaluate to address 1000 in our example because that is where the base of the array is. So, address 1000 is returned from x as the reference to another_array. But at the end of x's execution, the stack pointer is reset to what it was before x began executing. After x's execution, the stack will look like this:

Address Stack
1000    1       <- Stack top
999     2
998     3
997     4
996     5

The another_array's data is still there, but the stack top is now back at address 1000, so the next values going on the stack will overwrite whatever is in 1000, 999, 998, and so on.

The address 1000 is saved in main's another_array local variable and y is called.

y's yet_another_array gets initialized on the stack:

Address Stack
1000    0
999     0
998     0
997     0
996     0       <- Stack top

This completely overwrites x's another_array value. The array starting at address 1000 is now y's yet_another_array, but main is still holding the pointer 1000, thinking that it is a pointer to the array defined in x. In fact, that pointer is no longer valid, because it is no longer pointing to data that we're expecting it to point to. This is because the array was initialized on the stack and stack-allocated variables become invalid once the function has returned. This is because when a function is returned, the stack pointer is restored to its previous state and other functions will have their own local variables overwrite whatever was on the stack before.

On y's return, the stack will look like this:

Address Stack
1000    0       <- Stack top
999     0
998     0
997     0
996     0

Any function that is called now will overwrite these stack values with their local variables. Upon calling println!("{:?}", another_array);, [0, 0, 0, 0, 0] will be printed instead of [1, 2, 3, 4, 5]. You can imagine how sickening it will be if such an error was made in a large program of hundreds of thousands of lines. Finding that bug will be a pain.

This is where Rust steps in with lifetimes. Rust will never let this code compile because it will smell from a mile away that if it allows a reference to another_array to be returned, then when x finishes executing, that reference will become invalid and unsafe to dereference -- an anathema.

If you try compiling this code, Rust won't give you all these explanations. Instead, it will tell you to give x's return type an explicit lifetime. If you do that and try compiling again, it will tell you something like: "Can't return a reference to data owned by the current function".

That's Rust doing its work of protecting you from yourself. The borrow checker, which is in charge of making sure all these references are valid, was able to identify the problem with the code through the borrow checker rules. With this example, I think I can believe that it's clear why such a rule even exists.

Now, going back to the 'static lifetime, 'static is just a lifetime that tells Rust that a reference is valid throughout the whole program's execution. From the previous example, it's clear that references to local variables in a function are only valid while that function is executing. This is because of the stack allocation. But, if data is hardcoded into the binary, and not on the stack, then that data is sure to be in the one location it was hard coded into throughout the program's execution time.

Any reference to such data will have a 'static lifetime because the data the reference is pointing to will remain valid throughout the program's execution.

Remember that include_bytes! hardcodes a file's bytes into the compiler's output binary. So when that binary is loaded into memory for execution, the bytes of the file will remain in memory throughout execution time. Since Bitmap will only hold references to such a file, then all its references are 'static.

Getting On With Our Bitmap Structure

The purpose of this Bitmap structure is to make the code in our efi_main clearer and less convoluted, effectively communicating the meaning of what we intend to do. Simply defining the structure is not enough. We need to move the initialization code elsewhere.

To do this, we create an associated function for Bitmap which will be called whenever we want to handle any bitmap bytes. This initialization function will be in charge of doing all the initial interpretation work that we did in the second section of the code.

In bitmap.rs

impl Bitmap {
    // Initializes the Bitmap structure from the bitmap bytes
    pub fn new(bitmap_bytes: &'static [u8]) -> Self {
        // Retrieving a pointer to the block.bmp's bytes
        let bitmap_bytes_ptr: *const u8 = bitmap_bytes.as_ptr();
        // Reinterpreting a pointer to bytes as a pointer to a FileHeader instance
        let file_header_ptr = bitmap_bytes_ptr as *const FileHeader;
        // Interpreting the first section of the bitmap as the file header
        let file_header = unsafe { &(*file_header_ptr) };
        // The number of bytes that make up the FileHeader
        const FILE_HEADER_SIZE: usize = core::mem::size_of::<FileHeader>();
        // The DIB header comes immediately after the file header
        const DIB_HEADER_OFFSET: isize = FILE_HEADER_SIZE as isize;
        // Reinterpreting a pointer to the bytes at offset DIB_HEADER_OFFSET
        // as a pointer to the DIB header
        let dib_header_ptr = unsafe { bitmap_bytes_ptr.offset(DIB_HEADER_OFFSET) as *const DIBHeader };
        // Interpreting the second section of the bitmap as the DIB header
        let dib_header = unsafe { &(*dib_header_ptr) };
        // The number of bytes that make up the DIB header
        const DIB_HEADER_SIZE: usize = core::mem::size_of::<DIBHeader>();
        // The color table comes immediately after the file header and the DIB header
        const COLOR_TABLE_OFFSET: isize = (FILE_HEADER_SIZE + DIB_HEADER_SIZE) as isize;
        // Reinterpreting a pointer to the bytes at offset COLOR_TABLE_OFFSET as a pointer
        // to the color table
        let color_table_ptr = unsafe { bitmap_bytes_ptr.offset(COLOR_TABLE_OFFSET) as *const ColorTable };
        // Interpreting the bytes at `COLOR_TABLE_OFFSET` as the color table
        let color_table = unsafe { &(*color_table_ptr) };

        // Get a slice to the bitmap's pixel array 
        let pixel_array = &bitmap_bytes[file_header.pixel_array_offset as usize..];

        Self {
            file_header,
            dib_header,
            color_table,
            pixel_array
        }
    }
}

The Self keyword is just an alias for the struct Bitmap.

Now, we have a function that takes a slice to the bitmap bytes and returns all the structures we need to work with the bitmap, hiding the process of initialization. We can now refactor efi_main like so:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let block_bytes = include_bytes!("./block.bmp");

    /* DELETED:
    let block_bytes_ptr: *const u8 = block_bytes.as_ptr();
    let file_header_ptr = block_bytes_ptr as *const FileHeader;
    let file_header = unsafe { &(*file_header_ptr) };
    const FILE_HEADER_SIZE: usize = core::mem::size_of::<FileHeader>();
    const DIB_HEADER_OFFSET: isize = FILE_HEADER_SIZE as isize;
    let dib_header_ptr = unsafe { block_bytes_ptr.offset(DIB_HEADER_OFFSET) as *const DIBHeader };
    let dib_header = unsafe { &(*dib_header_ptr) };
    const DIB_HEADER_SIZE: usize = core::mem::size_of::<DIBHeader>();
    const COLOR_TABLE_OFFSET: isize = (FILE_HEADER_SIZE + DIB_HEADER_SIZE) as isize;
    let color_table_ptr = unsafe { block_bytes_ptr.offset(COLOR_TABLE_OFFSET) as *const ColorTable };
    let color_table = unsafe { &(*color_table_ptr) };

    let pixel_array = &block_bytes[file_header.pixel_array_offset as usize..];
    */
    
    let block = Bitmap::new(block_bytes); // NEW

    // DELETED: let image_height = dib_header.image_height as usize;
    let image_height = block.dib_header.image_height as usize; // NEW
    // DELETED: let image_width = dib_header.image_width as usize;
    let image_width = block.dib_header.image_width as usize; // NEW
    for row in 0..image_height {
        for col in 0..image_width {
            let inverted_row = image_height - row - 1;
            // DELETED: let color_table_index = pixel_array[inverted_row * image_width + col];
            let color_table_index = block.pixel_array[inverted_row * image_width + col]; // NEW
            // DELETED: let color = color_table.0[color_table_index as usize];
            let color = block.color_table.0[color_table_index as usize]; // NEW
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }

    0
}

We also have to import Bitmap into main.rs:

#![no_std]
#![no_main]
#![feature(abi_efiapi)]

mod font;
use font::FONT;

mod uefi;
use uefi::{SystemTable, Screen, PixelFormat, Pixel, pre_graphics_print_str, print_str, printint};

mod bitmap;
// DELETED: use bitmap::{FileHeader, DIBHeader, ColorTable, Color};
use bitmap::{FileHeader, DIBHeader, ColorTable, Color, Bitmap}; // NEW

We can also add convenience functions to Bitmap to get us the width and height:

impl Bitmap {
    fn new(bitmap_bytes: &'static [u8]) -> Self {
        /* OTHERS */
    }

    // Returns the height of the image
    pub fn height(&self) -> usize {
        // Retrieving the image height from the DIB header
        // Casting it as a usize so that it can be used to index into arrays and slices
        self.dib_header.image_height as usize
    }

    // Returns the width of the image
    pub fn width(&self) -> usize {
        // Retrieving the image width from the DIB header
        // Casting it as a usize so that it can be used to index into arrays and slices
        self.dib_header.image_width as usize
    }
}

Refactoring efi_main again:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let block_bytes = include_bytes!("./block.bmp");
    
    let block = Bitmap::new(block_bytes);

    /* DELETED:
    let image_height = block.dib_header.image_height as usize;
    let image_width = block.dib_header.image_width as usize;
    */

    // DELETED: for row in 0..image_height {
    for row in 0..block.height() { // NEW
        // DELETED: for col in 0..image_width {
        for col in 0..block.width() { // NEW
            // DELETED: let inverted_row = image_height - row - 1;
            let inverted_row = block.height() - row - 1; // NEW
            // DELETED: let color_table_index = block.pixel_array[inverted_row * image_width + col];
            let color_table_index = block.pixel_array[inverted_row * block.width() + col]; // NEW
            let color = block.color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }

    0
}

Now, as for the last section we're refactoring, we want to turn that double nested for loop into a more semantically accurate statement that effectively communicates "draw bitmap on screen".

This is a prime candidate for a function:

In main.rs

// Draws the bitmap on the screen
fn draw_bitmap(screen: &mut Screen, bitmap: &Bitmap) {
    for row in 0..bitmap.height() {
        for col in 0..bitmap.width() {
            // The image is upside down in the pixel array,
            // so the pixel array's rows have to retrieved from the bottom up
            let inverted_row = bitmap.height() - row - 1;
            let color_table_index = bitmap.pixel_array[inverted_row * bitmap.width() + col];
            let color = bitmap.color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }
}

Refactoring with this new function:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let block_bytes = include_bytes!("./block.bmp");
    
    let block = Bitmap::new(block_bytes);
    
    /* DELETED:
    for row in 0..block.height() {
        for col in 0..block.width() {
            let inverted_row = block.height() - row - 1;
            let color_table_index = block.pixel_array[inverted_row * block.width() + col];
            let color = block.color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }
    */

    // Draw the block on screen
    draw_bitmap(screen, &block);

    0
}

There are still quite some improvements we can make to this code, but we'll leave it here because we need to get to other things. We'll look at just one last thing: the draw_bitmap function.

The purpose of this main.rs is to hold driver code, that is, code that drives the execution of the main program. The purpose of uefi.rs is to hold code that is used for interfacing with the UEFI firmware. The purpose of the font.rs is to hold our font definitions and the purpose of bitmap.rs is to hold code that is used for working with bitmaps.

The draw_bitmap function is definitely not supposed to be defined in the main.rs file, because it isn't driver code. It seems like the most logically reasonable place to put it in is bitmap.rs because it has something to do with manipulating bitmaps, so we put it there.

main.rs:

/* DELETED:
fn draw_bitmap(screen: &mut Screen, bitmap: &Bitmap) {
    for row in 0..bitmap.height() {
        for col in 0..bitmap.width() {
            let inverted_row = bitmap.height() - row - 1;
            let color_table_index = bitmap.pixel_array[inverted_row * bitmap.width() + col];
            let color = bitmap.color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }
}
*/

bitmap.rs:

use crate::{Screen, Pixel};

pub fn draw_bitmap(screen: &mut Screen, bitmap: &Bitmap) {
    for row in 0..bitmap.height() {
        for col in 0..bitmap.width() {
            let inverted_row = bitmap.height() - row - 1;
            let color_table_index = bitmap.pixel_array[inverted_row * bitmap.width() + col];
            let color = bitmap.color_table.0[color_table_index as usize];
            screen.pixels[row][col] = Pixel {
                red: color.red,
                green: color.green,
                blue: color.blue,
                reserved: 0
            };
        }
    }
}

We now import it into main.rs

#![no_std]
#![no_main]
#![feature(abi_efiapi)]

mod font;
use font::FONT;

mod uefi;
use uefi::{SystemTable, Screen, PixelFormat, Pixel, pre_graphics_print_str, print_str, printint};

mod bitmap;
// DELETED: use bitmap::{FileHeader, DIBHeader, ColorTable, Color, Bitmap};
use bitmap::{FileHeader, DIBHeader, ColorTable, Color, Bitmap, draw_bitmap}; // NEW

Compiling and running the code now will still draw block.bmp on the screen.

Take Away

  • Rust lifetimes is the mechanism used by Rust to stop us from screwing ourselves up with invalid references.
  • A reference with the 'static lifetime is valid throughout the execution of the whole program.
  • The Self keyword is an alias for the name of the struct whose definition or implementation block it was used in.

For the full code, go to the repo

In The Next Post

We'll be automating our build process.

References

  • https://doc.rust-lang.org/book/ch10-03-lifetime-syntax.html
  • https://doc.rust-lang.org/nomicon/lifetimes.html
  • https://doc.rust-lang.org/reference/paths.html#self-1

← Previous Post

Drawing Bitmaps III

Next Post →

Automating The Build Process