Demilade Sonuga's blog

All posts

Refactoring V

2023-01-08

In the process of pulling the screen reference out of efi_main, we introduced a lot of unsafety in the code, which exposes us to more opportunities to screw ourselves up. We start our next refactor by hiding this unsafety behind a safe interface.

The site of our unsafe madness:

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

    let framebuffer_base = mode.framebuffer_base;

    let screen = framebuffer_base as *mut Screen;
    
    let screen = unsafe { &mut *screen };

    unsafe { SCREEN = Some(screen); }

    let screen = unsafe { SCREEN.as_mut().unwrap() };
    
    let block_bytes = include_bytes!("./block.bmp");

    let block = Bitmap::new(block_bytes);
}

Upon taking a good look at those three lines, their purpose is clear:

  1. Initialize the SCREEN static.
  2. Get the screen reference now stored in the static.

From this information, it seems like a sensible way to refactor this is to create two safe functions: init_screen and get_screen for initializing the screen static and retrieving the reference from the screen static, respectively.

But there's a problem with doing it that way. Suppose we create a safe function init_screen:

fn init_screen(screen: *mut Screen) {
    let screen = unsafe { &mut *screen };
    unsafe { SCREEN = Some(screen); }
}

If we create such a function, then we'll be able to call it from safe code without putting any unsafe blocks anywhere. The problem here is we could call this function from safe code, passing a pointer that is definitely not the correct pointer to Screen, and the SCREEN static will end up holding a wrong reference.

To understand what I'm talking about, take a look at this code sample:

let screen_ptr = 0 as *mut Screen;
init_screen(screen_ptr);
// This may be a safe reference, but it is definitely not safe to use
let screen = get_screen();
// Using screen will now lead to undefined behavior

In the sample, screen_ptr is set to an address that is definitely not the screen. This pointer is passed to init_screen which then initializes SCREEN to the wrong address. get_screen will now return a Rust reference which, although is meant to be safe at all times, is not safe here.

To be more precise, I mean the code in this sample will render a Rust reference unsafe. To understand why this is so, we need to understand how the safety of Rust's references works.

When Rust encounters a reference, it makes a few assumptions about them and analyzes the code for safety violations based on those assumptions. From the docs:

a reference is just a pointer that is assumed to be aligned, not null, and pointing to memory containing a valid value of T

That is, the assumptions are: aligned, not null and pointing to a valid value of T (for a &T or &mut T). The reason the Rust reference returned by get_screen in the sample is unsafe is that the last assumption of references has been violated by our code. The init_screen function creates a reference from a pointer that does not point to a valid value of Screen. This violation renders Rust's static checks useless because the reference itself does not hold up its conditions to be a reference.

The problem here is that the only way to create a safe reference to a Screen is from a raw screen pointer that must already be pointing to a valid value of the screen. But there is absolutely no way for the compiler to verify at compile time that the raw pointer we're passing to init_screen is safe.

This implies that the only way to verify the safety of a call to init_screen is to manually check that the pointer passed is indeed pointing to a valid Screen value. Because of this compulsory manual verification and the inability of the compiler to do this for us, we have to mark the init_screen function as unsafe:

// Initializes the screen static
// This function is unsafe because the caller has to
// manually verify that the argument passed is a valid
// pointer to Screen
unsafe fn init_screen(screen: *mut Screen) {
    let screen = &mut *screen;
    SCREEN = Some(screen);
}

This tells us whenever we call the function that there are some assumptions that we have to verify ourselves, that the compiler can't verify for us. In this case, the assumption we need to verify is that the raw pointer passed as an argument to the function is pointing to a valid value of Screen.

In other words, the unsafety of casting a random pointer to a Screen reference is unescapable because one way or the other, the validity of that pointer still needs to be verified and the compiler has no way of doing it for us.

The get_screen function can still be safe because all it's doing is getting a reference from a static. There is no manual verification that we need to do to know that the SCREEN static is safe to access and that the reference it contains is valid.

In main.rs:

fn get_screen() -> Option<&'static mut &'static mut Screen> {
    unsafe { SCREEN.as_mut() }
}

We now refactor efi_main with our new functions:

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

    let framebuffer_base = mode.framebuffer_base;

    let screen = framebuffer_base as *mut Screen;
    
    /* DELETED:
    let screen = unsafe { &mut *screen };

    unsafe { SCREEN = Some(screen); }

    let screen = unsafe { SCREEN.as_mut().unwrap() };
    */
    // NEW:
    unsafe { init_screen(screen); }
    
    let screen = get_screen().unwrap();
    
    let block_bytes = include_bytes!("./block.bmp");

    let block = Bitmap::new(block_bytes);
}

Since we've reduced the amount of unsafety here, we're making progress. But there is still some unsafety here that could have been avoided.

To understand how we could have avoided this, let's take a look at some of our previous abstractions:

In uefi.rs

impl GraphicsOutput {
    // Retrieves a reference to the GraphicsMode
    pub fn mode(&self) -> &GraphicsMode {
        unsafe { &*self.mode }
    }

    // ... Others
}

This is an implementation block of GraphicsOutput. It dereferences a raw pointer and returns a reference to the underlying data. We also have the same thing in our SystemTable implementation:

impl SystemTable {
    pub fn boot_services(&self) -> &BootServices {
        unsafe { &*self.boot_services }
    }

    pub fn simple_text_output(&self) -> &SimpleTextOutput {
        unsafe { &*self.simple_text_output }
    }
}

These functions dereference raw pointers and return references to the underlying data.

The question that this demands now is why is this safe when the init_screen function is unsafe even though they all do the same thing?

To answer this question, we first need to remember that the reason the init_screen function is unsafe is that we have to manually verify that the raw pointer being passed meets all the assumptions expected from a reference.

The SystemTable's functions are safe because the pointers they're dereferencing are assured to meet all the assumptions. And how do we know they are assured? Well, we know this because the UEFI spec dictates that the System Table looks like this:

#[repr(C)]
pub struct SystemTable {
    unneeded: [u8; 60],
    simple_text_output: *mut SimpleTextOutput,
    unneeded2: [u8; 24],
    boot_services: *const BootServices,
}

So, once we get a pointer to a valid SystemTable value, we are assured that all these other pointers contained in it are valid. But how do we get a valid pointer to the SystemTable?

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

According to the UEFI spec, the second argument in this function is guaranteed to be a pointer to a valid System Table. This manual verification justifies our taking a reference to the underlying data. So, all the pointer dereferences in the SystemTable's safe function implementations are safe. The same thing goes for the GraphicsOutput struct. Its raw pointer dereferences are safe because the GraphicsOutput instances that we're using only come from places that are guaranteed to hold valid instances.

In other words, the reason why these functions are safe and init_screen is unsafe is that the pointers dereferenced in the safe functions are already guaranteed to be safe to dereference, while the pointer dereferenced in the init_screen function has to be manually verified on every call to it.

With this information, we can make our initialization of the SCREEN static safe. Do this yourself before you continue.

Take a good look at the code:

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

    let framebuffer_base = mode.framebuffer_base;

    let screen = framebuffer_base as *mut Screen;

    // ... Others
}

And we see that the only address to be interpreted as the Screen is the GraphicsMode's framebuffer_base. Rather than getting that address, interpreting it as a Screen pointer and passing it to init_screen, we can do:

In uefi.rs

#[repr(C)]
pub struct GraphicsMode {
    pub max_mode: u32,
    pub mode: u32,
    pub info: *const GraphicsModeInfo,
    pub size_of_info: usize,
    // DELETED: pub framebuffer_base: PhysAddr,
    pub framebuffer_base: *mut Screen, // NEW
    pub framebuffer_size: usize,
}

Now, the compiler does the pointer interpretation for us. As long as the SystemTable we're making use of is valid, this interpretation will be correct.

Now, to get the screen:

impl GraphicsMode {
    // Returns a reference to the screen
    pub fn screen(&self) -> &mut Screen {
        unsafe { &mut *self.framebuffer_base }
    }
}

The value this pointer is pointing to is guaranteed to be a valid Screen, so it is safe to take a reference to it.

With this new safe function, we can refactor efi_main:

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

    let screen = framebuffer_base as *mut Screen;
    */

    // NEW:
    // A reference to the screen
    let screen = mode.screen();

    // DELETED: unsafe { init_screen(screen); }
    init_screen(screen); // NEW
    
    let screen = get_screen().unwrap();
    
    let block_bytes = include_bytes!("./block.bmp");

    let block = Bitmap::new(block_bytes);
}

The screen function returns a reference to the screen, which is already guaranteed to be safe. Passing this reference instead of a raw pointer to init_screen removes the need for manual verification on every call to init_screen. So, the only unsafe thing init_screen does now handling a mutable static:

// DELETED: unsafe fn init_screen(screen: *mut Screen) {
fn init_screen(screen: &'static mut Screen) { // NEW
    // DELETED: let screen = &mut *screen;
    unsafe { SCREEN = Some(screen); }
}

This assignment to the SCREEN static is safe because we don't need to manually verify that what is being passed to it is valid.

We can also eliminate the unsafety from the panic handler:

#[panic_handler]
fn panic_handler(panic_info: &core::panic::PanicInfo) -> ! {
    // DELETED: let screen = unsafe { SCREEN.as_mut().unwrap() };
    let screen = get_screen().unwrap(); // NEW
    write!(screen, "{}", panic_info);
    loop {}
}

Next, let's take another look at efi_main:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    let sys_table = unsafe { &*sys_table };
    let boot_services = sys_table.boot_services();
    let gop = boot_services.locate_gop();
    if gop.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }
    let gop = gop.unwrap();
    let mode = gop.mode();
    let max_mode = mode.max_mode;
    let mut desired_mode = 0;
    for mode_number in 0..max_mode {
        let mode = gop.query_mode(mode_number);
        if mode.is_err() {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to query mode\n");
            loop {}
        }
        let mode = mode.unwrap();

        let horizontal_resolution = mode.horizontal_resolution;
        let vertical_resolution = mode.vertical_resolution;
        let pixel_format = mode.pixel_format;
        if horizontal_resolution == 640
            && vertical_resolution == 480
            && pixel_format == PixelFormat::BlueGreenRedReserved
        {
            desired_mode = mode_number;
            break;
        }
        if mode_number == max_mode - 1 {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to find desired mode\n");
            loop {}
        }
    }

    let set_mode_result = gop.set_mode(desired_mode);

    if set_mode_result.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }

    let screen = mode.screen();

    init_screen(screen);
    
    let screen = get_screen().unwrap();

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

    let block = Bitmap::new(block_bytes);
    
    let mut block_position = (0, 0); // (row, column)
    while block_position.1 < NO_OF_PIXELS_IN_A_ROW {
        draw_bitmap(screen, &block, block_position);
        erase_bitmap(screen, &block, block_position);
        block_position.1 += 1;
    }
    
    0
}

It's not very clear what the intent of this function is from just looking at it. But this function does essentially just three things:

  1. Get the System Table and Boot Services.
  2. Set up graphics mode to our desired mode.
  3. Animate a bitmap on screen.

The clutter in this efi_main function is step 2. To clean this up, we create a new function:

fn init_graphics() {

}

The role of this function will be to initialize our graphics mode.

We begin:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    let sys_table = unsafe { &*sys_table };
    let boot_services = sys_table.boot_services();
    /* DELETED:
    let gop = boot_services.locate_gop();
    if gop.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }
    let gop = gop.unwrap();
    let mode = gop.mode();
    let max_mode = mode.max_mode;
    let mut desired_mode = 0;
    for mode_number in 0..max_mode {
        let mode = gop.query_mode(mode_number);
        if mode.is_err() {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to query mode\n");
            loop {}
        }
        let mode = mode.unwrap();

        let horizontal_resolution = mode.horizontal_resolution;
        let vertical_resolution = mode.vertical_resolution;
        let pixel_format = mode.pixel_format;
        if horizontal_resolution == 640
            && vertical_resolution == 480
            && pixel_format == PixelFormat::BlueGreenRedReserved
        {
            desired_mode = mode_number;
            break;
        }
        if mode_number == max_mode - 1 {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to find desired mode\n");
            loop {}
        }
    }

    let set_mode_result = gop.set_mode(desired_mode);

    if set_mode_result.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }

    let screen = mode.screen();
    */


    init_screen(screen);
    
    let screen = get_screen().unwrap();

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

    let block = Bitmap::new(block_bytes);

    // ... Others
}

fn init_graphics() {
    // NEW:
    let gop = boot_services.locate_gop();
    if gop.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }
    let gop = gop.unwrap();
    let mode = gop.mode();
    let max_mode = mode.max_mode;
    let mut desired_mode = 0;
    for mode_number in 0..max_mode {
        let mode = gop.query_mode(mode_number);
        if mode.is_err() {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to query mode\n");
            loop {}
        }
        let mode = mode.unwrap();

        let horizontal_resolution = mode.horizontal_resolution;
        let vertical_resolution = mode.vertical_resolution;
        let pixel_format = mode.pixel_format;
        if horizontal_resolution == 640
            && vertical_resolution == 480
            && pixel_format == PixelFormat::BlueGreenRedReserved
        {
            desired_mode = mode_number;
            break;
        }
        if mode_number == max_mode - 1 {
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to find desired mode\n");
            loop {}
        }
    }

    let set_mode_result = gop.set_mode(desired_mode);

    if set_mode_result.is_err() {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }
    let screen = mode.screen();
}

After the initialization of the graphics mode, we make use of the GraphicsMode instance to get a reference to the screen. So, our init_graphics function needs to return this Screen reference. But there is also a possibility that the function could fail:

fn init_graphics() {
    let gop = boot_services.locate_gop();
    if gop.is_err() {
        // POTENTIAL FAILURE
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }
    // ... Others
    for mode_number in 0..max_mode {
        let mode = gop.query_mode(mode_number);
        if mode.is_err() {
            // POTENTIAL FAILURE
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to query mode\n");
            loop {}
        }
        let mode = mode.unwrap();

        // ... Others

        if mode_number == max_mode - 1 {
            // POTENTIAL FAILURE
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to find desired mode\n");
            loop {}
        }
    }

    let set_mode_result = gop.set_mode(desired_mode);

    if set_mode_result.is_err() {
        // POTENTIAL FAILURE
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }
}

To indicate that this function could fail, instead of returning a GraphicsMode instance, it will return a Result that holds a GraphicsMode on success or an error message on failure:

// DELETED: fn init_graphics() {
fn init_graphics() -> Result<&'static mut Screen, &'static str> {
    let gop = boot_services.locate_gop();
    // If location failed, return error
    if gop.is_err() {
        /* DELETED:
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
        */
        return Err("Failed to locate GOP\n"); // NEW
    }
    // ... Others
    for mode_number in 0..max_mode {
        let mode = gop.query_mode(mode_number);
        if mode.is_err() {
            /* DELETED:
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to query mode\n");
            loop {}
            */
            return Err("Failed to query mode\n"); // NEW
        }
        let mode = mode.unwrap();

        // ... Others

        // Return an error if the desired mode wasn't found
        if mode_number == max_mode - 1 {
            /* DELETED:
            let simple_text_output = sys_table.simple_text_output();
            pre_graphics_print_str(simple_text_output, "Failed to find desired mode\n");
            loop {}
            */
            return Err("Failed to find desired mode\n"); // NEW
        }
    }

    let set_mode_result = gop.set_mode(desired_mode);

    if set_mode_result.is_err() {
        /* DELETED:
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
        */
        return Err("Failed to set the desired mode\n"); // NEW
    }
    
    // DELETED: let screen = mode.screen();
    Ok(mode.screen()) // NEW
}

To complete the init_graphics function, there's just one more thing: that is the only external thing this function depends on to work is a reference to BootServices:

// DELETED: fn init_graphics() -> Result<&'static mut Screen, &'static str> {
fn init_graphics(boot_services: &BootServices) -> Result<&'static mut Screen, &'static str> { // NEW
    // ... Others
}

Refactoring efi_main with this function:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    let sys_table = unsafe { &*sys_table };
    let boot_services = sys_table.boot_services();
    
    // NEW:
    let init_graphics_result = init_graphics(boot_services);
    // Halt with error if graphics initialization failed
    if let Err(msg) = init_graphics_result {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, msg);
        loop {}
    }
    let screen = init_graphics_result.unwrap();

    init_screen(screen);
    
    let screen = get_screen().unwrap();

    // ... Others
}

For the code to compile now:

In main.rs

mod font;
use font::FONT;

mod uefi;
use uefi::{SystemTable, Screen, PixelFormat, Pixel, pre_graphics_print_str, print_str, printint, 
// DELETED: NO_OF_PIXELS_IN_A_ROW};
NO_OF_PIXELS_IN_A_ROW, BootServices}; // NEW

mod bitmap;
use bitmap::{FileHeader, DIBHeader, ColorTable, Color, Bitmap, draw_bitmap, erase_bitmap};

// This needs to be imported so that the functions it defines may
// be used
use core::fmt::Write;

If you try running now, you'll get a lifetime error:

error: lifetime may not live long enough
   --> blasterball/src/main.rs:117:5
    |
66  | fn init_graphics(boot_services: &BootServices) -> Result<&'static mut Screen, &'static str> {
    |                                 - let's call the lifetime of this reference `'1`
...
117 |     Ok(mode.screen())
    |     ^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'static`

warning: `blasterball` (bin "blasterball") generated 3 warnings
error: could not compile `blasterball` due to previous error; 3 warnings emitted
Failed to build the game

It tells us that the BootServices reference has to have a 'static lifetime, so we give it one:

// DELETED: fn init_graphics(boot_services: &BootServices) -> Result<&'static mut Screen, &'static str> {
fn init_graphics(boot_services: &'static BootServices) -> Result<&'static mut Screen, &'static str> { // NEW
    // ... Others
}

The code will compile now and run as expected. This works because the BootServices reference we passed here is valid for the full duration of the program's execution, that is, it actually has a 'static lifetime.

The efi_main now looks clearer and easily readable.

To finalize this refactor, we will make use of traits. We've already seen the power of traits when we used the write! macro to print panic messages. Now, let's see where else we can use them:

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

    if let Err(msg) = init_graphics_result {
        let simple_text_output = sys_table.simple_text_output();
        pre_graphics_print_str(simple_text_output, msg);
        loop {}
    }

    // ... Others
}

Let's start from here. Remember the core::fmt::Write trait defines behavior for anything that characters can be written into. Clearly, the SimpleTextOutput indicates that behavior. We can implement Write for it:

In uefi.rs

impl Write for SimpleTextOutput {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        pre_graphics_print_str(self, s);
        Ok(())
    }
}

In main.rs

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

    if let Err(msg) = init_graphics_result {
        let simple_text_output = sys_table.simple_text_output();
        // DELETED: pre_graphics_print_str(simple_text_output, msg);
        write!(simple_text_output, "{}", msg);
        loop {}
    }

    // ... Others
}

The only problem now is that write! uses mutable references but SystemTable's simple_text_output function returns an immutable reference to SimpleTextOutput. To resolve this:

In uefi.rs

impl SystemTable {
    // ... Others

    // DELETED: pub fn simple_text_output(&self) -> &SimpleTextOutput {
    pub fn simple_text_output(&self) -> &mut SimpleTextOutput {
        // DELETED: unsafe { &*self.simple_text_output }
        unsafe { &mut *self.simple_text_output }
    }
}

This simple change may not have improved the clarity of our code, but it makes the code more generic: The abstractions being used to print are the same everywhere.

Similarly to Write, there is another trait, core::ops::Index, that defines behavior for anything that can be indexed.

In bitmap.rs

#[repr(transparent)]
pub struct ColorTable(pub [Color; 256]);

This ColorTable has the notion of something that can be indexed. Rather than referring to its contents with color_table.0[idx], we can implement Index:

impl Index<usize> for ColorTable {
    type Output = Color;

    fn index(&self, idx: usize) -> &Self::Output {
        &self.0[idx]
    }
}

This implementation allows us to index directly into the color table without requiring us to access the ColorTable's fields.

We can now make ColorTable's field private because we don't need it to be public:

#[repr(transparent)]
// DELETED: pub struct ColorTable(pub [Color; 256]);
pub struct ColorTable([Color; 256]); // NEW

To reflect this change in draw_bitmap:

pub fn draw_bitmap(screen: &mut Screen, bitmap: &Bitmap, pos: (usize, usize)) {
    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];
            // DELETED: let color = bitmap.color_table.0[color_table_index as usize];
            let color = bitmap.color_table[color_table_index as usize]; // NEW
            if row + pos.0 < NO_OF_PIXELS_IN_A_COLUMN && col + pos.1 < NO_OF_PIXELS_IN_A_ROW {
                screen.pixels[row + pos.0][col + pos.1] = Pixel {
                    red: color.red,
                    green: color.green,
                    blue: color.blue,
                    reserved: 0
                };
            }
        }
    }
}

Another structure that has this behavior of something being indexed is Screen. We index into a screen to get the pixel at a position or to color in a pixel.

In uefi.rs

impl Index<usize> for Screen {
    type Output = [Pixel; NO_OF_PIXELS_IN_A_ROW];

    fn index(&self, idx: usize) -> &Self::Output {
        &self.pixels[idx]
    }
}

impl IndexMut<usize> for Screen {
    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
        &mut self.pixels[idx]
    }
}

We also implement IndexMut because we need to be able to mutate the screen.

Refactoring to reflect this new ability:

In uefi.rs

pub fn print_char(
    screen: &mut Screen,
    font_description: &[[bool; 16]; 16],
    curr_screen_pos: (usize, usize),
) {
    for i in 0..16 {
        for j in 0..16 {
            if font_description[i][j] {
                // Red and green is yellow (which we're using as our foreground color here)
                // DELETED: screen.pixels[curr_screen_pos.0 + i][curr_screen_pos.1 + j] = Pixel {
                screen[curr_screen_pos.0 + i][curr_screen_pos.1 + j] = Pixel { //  NEW
                    red: 255,
                    green: 255,
                    blue: 0,
                    reserved: 0,
                };
            } else {
                // DELETED: screen.pixels[curr_screen_pos.0 + i][curr_screen_pos.1 + j] = Pixel {
                screen[curr_screen_pos.0 + i][curr_screen_pos.1 + j] = Pixel { // NEW
                    red: 0,
                    green: 0,
                    blue: 0,
                    reserved: 0,
                };
            }
        }
    }
}

In bitmap.rs

pub fn draw_bitmap(screen: &mut Screen, bitmap: &Bitmap, pos: (usize, usize)) {
    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[color_table_index as usize];
            if row + pos.0 < NO_OF_PIXELS_IN_A_COLUMN && col + pos.1 < NO_OF_PIXELS_IN_A_ROW {
                // DELETED: screen.pixels[row + pos.0][col + pos.1] = Pixel {
                screen[row + pos.0][col + pos.1] = Pixel { // NEW
                    red: color.red,
                    green: color.green,
                    blue: color.blue,
                    reserved: 0
                };
            }
        }
    }
}

pub fn erase_bitmap(screen: &mut Screen, bitmap: &Bitmap, pos: (usize, usize)) {
    for row in 0..bitmap.width() {
        for col in 0..bitmap.height() {
            if row + pos.0 < NO_OF_PIXELS_IN_A_COLUMN && col + pos.1 < NO_OF_PIXELS_IN_A_ROW {
                // DELETED: screen.pixels[row + pos.0][col + pos.1] = Pixel {
                screen[row + pos.0][col + pos.1] = Pixel { // NEW
                    red: 0,
                    green: 0,
                    blue: 0,
                    reserved: 0
                };
            }
        }
    }
}

For the code to compile:

In bitmap.rs

use crate::{Screen, Pixel};
use crate::uefi::{NO_OF_PIXELS_IN_A_COLUMN, NO_OF_PIXELS_IN_A_ROW};
use core::ops::Index; // NEW

In uefi.rs

use crate::FONT;
use core::fmt::Write;
use core::sync::atomic::{AtomicUsize, Ordering};
use core::ops::{Index, IndexMut}; // NEW

And that's it for now. Running the code now will result in the same block animation.

Take Away

  • The core::ops::Index and core::ops::IndexMut are traits used to define indexing behavior.

For the full code up to this point, go to the repo

In The Next Post

We'll be drawing our game scene

References

  • https://doc.rust-lang.org/std/primitive.reference.html
  • https://doc.rust-lang.org/core/ops/trait.Index.html
  • https://doc.rust-lang.org/core/ops/trait.IndexMut.html