MMIO Abstraction Example

Example.rs

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

mod register;
use register::*;

#[lang = "eh_personality"]
extern fn eh_personality() {}

#[panic_handler]
fn panic(_: &core::panic::PanicInfo) -> ! {
    loop {}
}

const CPU_FREQ: u32 = 16_000_000;
/// The baud rate we'll use for serial.
const BAUD_RATE: u32 = 9600;
/// The calculated value to put into the UBBR register to set the baud rate.
const UBBR_VAL: u16 = ((CPU_FREQ / 8 / BAUD_RATE) - 1) as u16;

reg! {
    UDR0: u8 {
        addr: 0xC6,
        write mask: 0xFF,
    }
}

reg! {
    UCSR0A: u8 {
        addr: 0xC0,
        write mask: 0b1110_0011,
        bits: {
            MPCM0 = 0, RW;
            U2X0 = 1, RW;
            UPE0 = 2, R;
            DOR0 = 3, R;
            FE0 = 4, R;
            UDRE0 = 5, R;
            TXC0 = 6, RW;
            RXC0 = 7, R;
        }
    }
}

reg! {
    UCSR0B: u8 {
        addr: 0xC1,
        write mask: 0b1111_1111,
        bits: {
            TXB80 = 0, RW;
            RXB80 = 1, R;
            UCSZ02 = 2, RW;
            TXEN0 = 3, RW;
            RXEN0 = 4, RW;
            UDRIE0 = 5, RW;
            TXCIE0 = 6, RW;
            RXCIE0 = 7, RW;
        }
    }
}

reg! {
    UCSR0C: u8 {
        addr: 0xC2,
        write mask: 0b1111_1111,
        bits: {
            UCPOL0 = 0, RW;
            UCSZ00 = 1, RW;
            UCSZ01 = 2, RW;
            USBS0 = 3, RW;
            UPM00 = 4, RW;
            UPM01 = 5, RW;
            UMSEL00 = 6, RW;
            UMSEL01 = 7, RW;
        }
    }
}

reg! {
    UBRR0: u16 {
        addr: 0xC4,
        write mask: 0x0FFF,
    }
}

#[no_mangle]
extern "C" fn main() {
    unsafe {
        // Set our baud rate.
        UBRR0::set_raw_value(UBBR_VAL);

        // Configure for:
        // // * 2x speed
        // // * 8-bit characters
        // // * 1 stop bit
        // // * No parity
        // // * Async mode,
        // // * Enable RX/TX
        UCSR0A::set_value(U2X0 | MPCM0);
        UCSR0B::set_value(RXEN0 | TXEN0);
        UCSR0C::set_value(UCSZ01 | UCSZ00);

        let message = "Hello World!";

        message.bytes().for_each(|b| {
            // Wait for the data register to become available.
            while !UCSR0A::get_bit(UDRE0) {}

            // Stick the byte into the buffer to send it.
            UDR0::set_raw_value(b);
        });
    }
}

registers.rs

use core::{
    ops::{BitAnd, BitOr, BitOrAssign, Shl, Not},
    marker::PhantomData,
};

// Used to restrict what data types a register can be.
pub trait RegisterType: Copy + BitAnd<Output=Self> + BitOr<Output=Self> + Shl<Output=Self> + Not<Output=Self> + Eq + PartialEq {
    const ZERO: Self;
    const ONE: Self;
}

impl RegisterType for u8 {
    const ZERO: Self = 0;
    const ONE: Self = 1;
}

impl RegisterType for u16 {
    const ZERO: Self = 0;
    const ONE: Self = 1;
}

// This trait and types represent whether a bit is readable and/or writable in the type system.
pub trait Access {}
pub struct Readable;
impl Access for Readable {}
pub struct NotReadable;
impl Access for NotReadable {}

pub struct Writable;
impl Access for Writable {}
pub struct NotWritable;
impl Access for NotWritable {}


// This allows us to say that access is only allowed if both bits have that access.
// It is, essentially, a boolean AND operation, hence the name.
pub trait AccessAnd<Rhs> {
    type Output: Access;
}

impl AccessAnd<Readable> for Readable {
    type Output = Readable;
}
impl AccessAnd<NotReadable> for Readable {
    type Output = NotReadable;
}
impl AccessAnd<NotReadable> for NotReadable {
    type Output = NotReadable;
}
impl AccessAnd<Readable> for NotReadable {
    type Output = NotReadable;
}

impl AccessAnd<Writable> for Writable {
    type Output = Writable;
}
impl AccessAnd<NotWritable> for Writable {
    type Output = NotWritable;
}
impl AccessAnd<NotWritable> for NotWritable {
    type Output = NotWritable;
}
impl AccessAnd<Writable> for NotWritable {
    type Output = NotWritable;
}

// A bit needs to be associated with its parent register, and what read/write access the bit has, so incorrect usage
// is prevented at compile-time.
pub trait Bit {
    type ReadAccess: Access;
    type WriteAccess: Access;
    type Register: Register;
    fn bit_id(&self) -> <Self::Register as Register>::DataType;
}

// The BitBuilder allows the user to write expressions like `Bit1 | Bit2` to build up a bit pattern, while
// still retaining the connection to the register and the restrictions on read/write access.
pub struct BitBuilder<DataType: RegisterType, Reg: Register<DataType = DataType>, R: Access, W: Access>{
    data: DataType,
    _p: PhantomData<(Reg, R, W)>
}

impl<DataType: RegisterType, Reg: Register<DataType = DataType>> BitBuilder<DataType, Reg, Readable, Writable> {
    pub fn new() -> Self {
        Self {
            data: DataType::ZERO,
            _p: PhantomData,
        }
    }
}

impl <DataType: RegisterType, Reg: Register<DataType = DataType>, R: Access, W: Access> BitBuilder<DataType, Reg, R, W> {
    pub fn raw_value(&self) -> DataType {
        self.data
    }
}

// This one's pretty gnarly. Basically, this lets the user do `bits | bits`, while continuing the connection
// with the register. The resulting BitBuilder will have the most restrictive read/write access.
// E.g. A BitBuilder with read/write access being ORed with one with only read access will result in a BitBuilder
// with only read access.
impl<DataType, Reg, R1, W1, R2, W2> BitOr<BitBuilder<DataType, Reg, R2, W2>> for BitBuilder<DataType, Reg, R1, W1>
    where R1: Access,
          R2: Access,
          W1: Access,
          W2: Access,
          DataType: RegisterType,
          Reg: Register<DataType = DataType>,
          R2: AccessAnd<R1>,
          W2: AccessAnd<W1>,
{
    type Output = BitBuilder<DataType, Reg, <R2 as AccessAnd<R1>>::Output, <W2 as AccessAnd<W1>>::Output>;
    fn bitor(self, rhs: BitBuilder<DataType, Reg, R2, W2>) -> Self::Output {
        BitBuilder {
            data: self.data | rhs.data,
            _p: PhantomData,
        }
    }
}

// This just lets us OR a BitBuilder with a Bit from the same register. As with above, the most restrictive access
// applies.
impl<DataType, Reg, R1, W1, B> BitOr<B> for BitBuilder<DataType, Reg, R1, W1>
    where R1: Access,
          W1: Access,
          DataType: RegisterType,
          Reg: Register<DataType = DataType>,
          B: Bit<Register=Reg>,
          B::ReadAccess: AccessAnd<R1>,
          B::WriteAccess: AccessAnd<W1>,
{
    type Output = BitBuilder<DataType, Reg, <B::ReadAccess as AccessAnd<R1>>::Output, <B::WriteAccess as AccessAnd<W1>>::Output>;
    fn bitor(self, rhs: B) -> Self::Output {
        BitBuilder {
            data: self.data | (DataType::ONE << rhs.bit_id()),
            _p: PhantomData,
        }
    }
}

// This one's a little restrictive, but lets us do things like:
//
// let mut bits = TWEN | TWIE | TWINT;
// if ack {
//     bits |= TWEA;
// }
// TWCR::set_value(bits);
//
// One restriction is that both read and write access of the new bit must exactly match the access of the BitBuilder.
impl<DataType, Reg, R1, W1, B> BitOrAssign<B> for BitBuilder<DataType, Reg, R1, W1>
    where R1: Access,
          W1: Access,
          DataType: RegisterType,
          Reg: Register<DataType = DataType>,
          B: Bit<Register=Reg, ReadAccess=R1, WriteAccess=W1>,
{
    fn bitor_assign(&mut self, rhs: B) {
        self.data = self.data | (DataType::ONE << rhs.bit_id());
    }
}

// This trait is to fix an irritating papercut. If we don't have it and have set_value take a BitBuilder, it would
// mean that we could do this:
//
// TWCR::set_value(TWEN | TWIE);
//
// But not this:
//
// TWCR::set_value(TWEN);
//
// Which feels incredibly inconsistent.
pub trait SetValueType {
    type DataType: RegisterType;
    type Register: Register<DataType = Self::DataType>;
    type WriteAccess: Access;

    fn value(&self) -> Self::DataType;
}

impl<DataType, Reg, R, W> SetValueType for BitBuilder<DataType, Reg, R, W>
    where DataType: RegisterType,
          Reg: Register<DataType = DataType>,
          R: Access,
          W: Access,
{
    type DataType = DataType;
    type Register = Reg;
    type WriteAccess = W;

    fn value(&self) -> DataType {
        self.data
    }
}

// Abstracts away the messy volatile pointer accessses and bit-twiddling needed for an MMIO.
// The type needs to be tracked, as does a write mask, as some bits should never be written to.
// One thing that does need to be provided is the ability to set a raw value, as some registers
// are data registers, not control/status registers. Bit-twiddling these makes no sense.
pub trait Register: Sized {
    type DataType: RegisterType;
    const ADDR: *mut Self::DataType;

    // Some bits should always be written 0 when writing, this allows that.
    const WRITE_MASK: Self::DataType;

    unsafe fn set_raw_value(val: Self::DataType) {
        Self::ADDR.write_volatile(val & Self::WRITE_MASK);
    }

    unsafe fn set_value<V>(val: V)
        where V: SetValueType<DataType=Self::DataType, Register=Self, WriteAccess=Writable>
    {
        let value = val.value();
        Self::set_raw_value(value);
    }

    unsafe fn get_value() -> Self::DataType {
        Self::ADDR.read_volatile()
    }

    unsafe fn get_bit<B>(bit: B) -> bool
        where B: Bit<Register=Self, ReadAccess=Readable>
    {
        let bit = Self::DataType::ONE << bit.bit_id();

        (Self::get_value() & bit) != Self::DataType::ZERO
    }

    unsafe fn set_bits<V>(bits: V)
        where V: SetValueType<DataType=Self::DataType, Register=Self, WriteAccess=Writable>,
    {
        let val = Self::get_value();
        Self::set_raw_value(val | bits.value());
    }

    unsafe fn clear_bits<V>(bits: V)
        where V: SetValueType<DataType=Self::DataType, Register=Self, WriteAccess=Writable>,
    {
        let val = Self::get_value();
        Self::set_raw_value(val & !bits.value());
    }

    unsafe fn replace_bits(
        mask: BitBuilder<Self::DataType, Self, Readable, Writable>,
        new_val: BitBuilder<Self::DataType, Self, Readable, Writable>,
    ) {
        let reg_val = Self::get_value() & !mask.data;
        let masked_val = new_val.data & mask.data;
        Self::set_raw_value(reg_val | masked_val);
    }
}

// Unfortunately, all the traits and above make actually defining a register and its bits something that no sane person
// should do. Fortunately, because it's the same for all registers/bits, macros can be used.


// These two are just for expanding the read/write access type defs. It allows the other macros to just match
// the access flags as a token tree, and have these two expand it.
#[macro_export]
macro_rules! expand_read_access {
    (R) => {
        type ReadAccess = crate::register::Readable;
    };
    (W) => {
        type ReadAccess = crate::register::NotReadable;
    };
    (RW) => {
        type ReadAccess = crate::register::Readable;
    };
}

#[macro_export]
macro_rules! expand_write_access {
    (R) => {
        type WriteAccess = crate::register::NotWritable;
    };
    (W) => {
        type WriteAccess = crate::register::Writable;
    };
    (RW) => {
        type WriteAccess = crate::register::Writable;
    };
}

// This one is what creates all the types representing the bits.
#[macro_export]
macro_rules! reg_named_bits {
    ( 
        $name:ident : $type:ty {
            $( $(#[$bit_doc:meta])* $bit:ident = $id:expr, $acc:tt;)+
        }
    ) => {
        $(
            $(#[$bit_doc])*
            #[derive(Copy, Clone)]
            pub struct $bit;
            impl crate::register::Bit for $bit {
                type Register = $name;
                expand_read_access!{$acc}
                expand_write_access!{$acc}
                fn bit_id(&self) -> $type {
                    $id
                }
            }

            impl crate::register::SetValueType for $bit {
                type DataType = $type;
                type Register = $name;
                expand_write_access!{$acc}

                fn value(&self) -> $type {
                    use crate::register::Bit;
                    1 << self.bit_id()
                }
            }

            impl<B2> core::ops::BitOr<B2> for $bit
                where Self: crate::register::Bit,
                    B2: crate::register::Bit<Register = <Self as crate::register::Bit>::Register>,
                    B2::ReadAccess: crate::register::AccessAnd<<Self as crate::register::Bit>::ReadAccess>,
                    B2::WriteAccess: crate::register::AccessAnd<<Self as crate::register::Bit>::WriteAccess>,
            {
                type Output = crate::register::BitBuilder<
                    <<Self as crate::register::Bit>::Register as crate::register::Register>::DataType,
                    <Self as crate::register::Bit>::Register,
                    <B2::ReadAccess as crate::register::AccessAnd<<Self as crate::register::Bit>::ReadAccess>>::Output,
                    <B2::WriteAccess as crate::register::AccessAnd<<Self as crate::register::Bit>::WriteAccess>>::Output,
                >;

                fn bitor(self, rhs: B2) -> Self::Output {
                    crate::register::BitBuilder::new() | self | rhs
                }
            }
        )*
    };
}

// There are a lot of bits, and it would be useful to namespace them. This macro allows the user to access the bits
// like this: TWCR::TWEN;
#[macro_export]
macro_rules! reg_bit_consts {
    (
        $struct_name:ident {
            $( $(#[$bit_doc:meta])* $bit:ident ),+ $(,)*
        }
    ) => {
        impl $struct_name {
            $( 
                $(#[$bit_doc])*
                #[allow(dead_code)]
                pub const $bit: $bit = $bit;
            )*
        }
    }
}

#[macro_export]
macro_rules! reg {
    (
        $(#[$reg_doc:meta])*
        $name:ident : $type:ty {
            addr: $addr:expr,
            write mask: $mask:expr $(,)*
        }
    ) => {
        $(#[$reg_doc])*
        #[allow(dead_code)]
        pub struct $name;
        impl crate::register::Register for $name {
            type DataType = $type;
            const WRITE_MASK: $type = $mask;
            const ADDR: *mut $type = $addr as *mut $type;
        }
    };

    (
        $(#[$reg_doc:meta])*
        $name:ident: $type:ty {
            addr: $addr:expr,
            write mask: $mask:expr,
            bits: {
                $( $(#[$bit_doc:meta])* $bit:ident = $id:expr, $acc:tt;)+
            }
        }
    ) => {
        reg_named_bits! {
            $name: $type {
                $( $(#[$bit_doc])* $bit = $id, $acc; )+
            }
        }

        reg! {
            $(#[$reg_doc])*
            $name: $type {
                addr: $addr,
                write mask: $mask,
            }
        }

        reg_bit_consts! {
            $name {
                $( $bit ),+
            }
        }
    };
}