regs-macros 0.1.0

use proc_macro::TokenStream;
use proc_macro2::Span;
use proc_macro_crate::{crate_name, FoundCrate};
use quote::{format_ident, quote};
use syn::parse::{Parse, ParseStream, Result};
use syn::punctuated::Punctuated;
use syn::{braced, bracketed, Error, Ident, LitInt, Token, Type};

fn get_regs_crate() -> Result<proc_macro2::TokenStream> {
    match crate_name("regs") {
        Ok(FoundCrate::Itself) => Ok(quote!(crate)),
        Ok(FoundCrate::Name(name)) => {
            let ident = Ident::new(&name, Span::call_site());
            Ok(quote!(#ident))
        }
        Err(_) => Err(Error::new(Span::call_site(), "cannot find crate `regs`")),
    }
}

struct Access {
    r: bool,
    w: bool,
}

impl Default for Access {
    fn default() -> Self {
        Self { r: true, w: true }
    }
}

enum RegBitfieldOptions {
    Access(Access),
}

impl Parse for RegBitfieldOptions {
    fn parse(input: ParseStream) -> Result<Self> {
        let ident: Ident = input.parse()?;
        match ident.to_string().as_str() {
            "ro" => Ok(RegBitfieldOptions::Access(Access { r: true, w: false })),
            "rw" => Ok(RegBitfieldOptions::Access(Access { r: true, w: true })),
            "wo" => Ok(RegBitfieldOptions::Access(Access { r: false, w: true })),
            other => Err(Error::new(
                ident.span(),
                &format!("Unexpected bitfield option: {other}"),
            )),
        }
    }
}

pub struct RegBitfield {
    ident: Ident,
    shift: LitInt,
    width: LitInt,
    access: Access,
}

impl Parse for RegBitfield {
    fn parse(input: ParseStream) -> Result<Self> {
        let ident: Ident = input.parse()?;
        let content;
        bracketed!(content in input);

        let shift: LitInt = content.parse()?;
        let start: u8 = shift.base10_parse()?;

        let width = if content.peek(Token![..]) {
            let _: Token![..] = content.parse()?;
            let end: LitInt = content.parse()?;
            end.base10_parse::<u8>()? - start + 1
        } else if content.peek(Token![;]) {
            let _: Token![;] = content.parse()?;
            let width: LitInt = content.parse()?;
            width.base10_parse::<u8>()?
        } else {
            1u8
        };

        let options = if content.peek(Token![,]) {
            let _: Token![,] = content.parse()?;
            Some(Punctuated::<RegBitfieldOptions, Token![,]>::parse_separated_nonempty(&content)?)
        } else {
            None
        };

        let mut access: Option<Access> = None;
        if let Some(options) = options {
            for opt in options {
                match opt {
                    RegBitfieldOptions::Access(a) => {
                        if access.is_some() {
                            return Err(Error::new(
                                input.span(),
                                "Duplicate bitfield option: access",
                            ));
                        }

                        access.replace(a);
                    }
                }
            }
        }

        let width = LitInt::new(&width.to_string(), Span::call_site());

        Ok(Self {
            ident,
            shift,
            width,
            access: access.unwrap_or_default(),
        })
    }
}

pub enum RegSpecOption {
    Bits(Punctuated<RegBitfield, Token![,]>),
}

impl RegSpecOption {
    fn parse_bits(input: ParseStream) -> Result<Self> {
        let bits;
        braced!(bits in input);

        let fields = Punctuated::<RegBitfield, Token![,]>::parse_terminated(&bits)?;
        Ok(Self::Bits(fields))
    }
}

impl Parse for RegSpecOption {
    fn parse(input: ParseStream) -> Result<Self> {
        let ident: Ident = input.parse()?;
        match ident.to_string().as_str() {
            "bits" => Self::parse_bits(input),
            other => Err(Error::new(
                ident.span(),
                format!("Unexpected option: {other}"),
            )),
        }
    }
}

pub struct RegSpec {
    ident: Ident,
    width: Type,
    fields: Vec<RegBitfield>,
    regs_crate: proc_macro2::TokenStream,
}

impl Parse for RegSpec {
    fn parse(input: ParseStream) -> Result<Self> {
        let ident: Ident = input.parse()?;
        let _as: Token![as] = input.parse()?;
        let width: Type = input.parse()?;

        let mut fields: Option<Vec<RegBitfield>> = None;

        // If there is more to parse, at this point we need to encounter a comma.
        if !input.is_empty() {
            let _: Token![,] = input.parse()?;

            // If there is nothing more after the comma, we reject the input.
            if input.is_empty() {
                return Err(input.error("Unexpected trailing comma."));
            }

            let items = Punctuated::<RegSpecOption, Token![,]>::parse_separated_nonempty(input)?;
            for item in items {
                match item {
                    RegSpecOption::Bits(punctuated) => {
                        if fields.is_some() {
                            return Err(Error::new(Span::call_site(), "Duplicate option: bits"));
                        }

                        let mut res = Vec::new();
                        for f in punctuated {
                            res.push(f);
                        }
                        fields.replace(res);
                    }
                }
            }
        }

        Ok(Self {
            ident,
            width,
            fields: fields.unwrap_or_default(),
            regs_crate: get_regs_crate()?,
        })
    }
}

impl RegSpec {
    pub fn into_token_stream(self) -> TokenStream {
        let regs_crate = self.regs_crate;

        let reg_type = format_ident!("{}_Reg", self.ident);
        let val_type = format_ident!("{}_Val", self.ident);

        let width = self.width;

        let mut methods = Vec::new();
        for f in self.fields {
            let shift = f.shift;
            let len = f.width;

            if f.access.r {
                let ident = Ident::new(&f.ident.to_string().to_lowercase(), f.ident.span());
                let read = quote! {
                    pub fn #ident(&self) -> #width {
                        (self.0 >> #shift) & ((1 << #len) - 1)
                    }
                };

                methods.push(read);
            }

            if f.access.w {
                let ident = &f.ident;
                let ident = Ident::new(
                    &format_ident!("with_{ident}").to_string().to_lowercase(),
                    f.ident.span(),
                );

                let write = quote! {
                    pub fn #ident(self, value: #width) -> Self {
                        Self(self.0 | ((value & ((1 << #len) - 1)) << #shift))
                    }
                };

                methods.push(write);
            }
        }

        let regspec_def = quote! {
            #[allow(non_camel_case_types)]
            pub struct #reg_type(usize);

            impl #reg_type {
                pub const fn address(&self) -> usize {
                    self.0
                }

                pub unsafe fn write(&self, value: #val_type) {
                    ::#regs_crate::write::<#width>(self.0, value.bits());
                }

                pub unsafe fn read(&self) -> #val_type {
                    #val_type(::#regs_crate::read::<#width>(self.0))
                }

                pub unsafe fn modify<F: FnOnce(#val_type) -> #val_type>(&self, f: F) {
                    self.write(f(self.read()));
                }
            }

            #[allow(non_camel_case_types)]
            pub struct #val_type(#width);

            impl #val_type {
                pub fn bits(&self) -> #width {
                    self.0
                }

                #(#methods)*
            }
        };

        regspec_def.into()
    }
}