reoxide-proc 0.7.0

use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};

macro_rules! ident {
    ($ident:expr, $span:expr) => {
        TokenTree::Ident(Ident::new($ident, $span))
    };
    ($ident:expr) => {
        TokenTree::Ident(Ident::new($ident, Span::call_site()))
    };
}

macro_rules! punct {
    ($ch:literal, $span:expr) => {{
        let mut punct = Punct::new($ch, Spacing::Alone);
        punct.set_span($span);
        TokenTree::Punct(punct)
    }};
    ($ch:literal) => {
        TokenTree::Punct(Punct::new($ch, Spacing::Alone))
    };
}

macro_rules! jpunct {
    ($ch:literal, $span:expr) => {{
        let mut punct = Punct::new($ch, Spacing::Joint);
        punct.set_span($span);
        TokenTree::Punct(punct)
    }};
    ($ch:literal) => {
        TokenTree::Punct(Punct::new($ch, Spacing::Joint))
    };
}

macro_rules! literal {
    ($lit:expr, $span:expr) => {{
        let mut literal = Literal::string($lit);
        literal.set_span($span);
        TokenTree::Literal(literal)
    }};
    ($lit:expr) => {
        TokenTree::Literal(Literal::string($lit))
    };
}

macro_rules! group {
    ($delimiter:expr, $array: expr, $span:expr) => {{
        let mut group = Group::new($delimiter, TokenStream::from_iter($array));
        group.set_span($span);
        TokenTree::Group(group)
    }};
    ($delimiter:expr, $array: expr) => {
        TokenTree::Group(Group::new($delimiter, TokenStream::from_iter($array)))
    };
}

macro_rules! braces {
    ($delimiter:expr, $span:expr) => {{
        let mut group = Group::new($delimiter, TokenStream::new());
        group.set_span($span);
        TokenTree::Group(group)
    }};
    ($delimiter:expr) => {
        TokenTree::Group(Group::new($delimiter, TokenStream::new()))
    };
}

struct Error {
    message: String,
    span: Span,
}

impl Error {
    fn into_token_stream(self) -> TokenStream {
        let span = self.span;
        let message = self.message;

        TokenStream::from_iter([
            jpunct!(':', span),
            punct!(':', span),
            ident!("core", span),
            jpunct!(':', span),
            punct!(':', span),
            ident!("compile_error", span),
            punct!('!', span),
            group!(Delimiter::Brace, [literal!(&message, span)], span),
        ])
    }

    fn into_combined_token_stream(self, other: TokenStream) -> TokenStream {
        let mut stream = self.into_token_stream();
        stream.extend(other);
        stream
    }
}

type Result<T> = std::result::Result<T, Error>;

macro_rules! error {
    ($msg:expr, $span:expr) => {
        Err(Error {
            message: $msg,
            span: $span,
        })
    };
}

fn parse_ident<T>(iter: &mut T, span: Span, ident: &str) -> Result<Span>
where
    T: Iterator<Item = TokenTree>,
{
    match iter.next() {
        Some(TokenTree::Ident(l)) if l.to_string() == ident => Ok(l.span().end()),
        Some(t) => error!(format!(r#"Expected "{}", got "{}""#, ident, t), t.span()),
        None => error!(format!(r#"Expected "{}""#, ident), span),
    }
}

fn parse_ident_name<T>(iter: &mut T, span: Span) -> Result<Ident>
where
    T: Iterator<Item = TokenTree>,
{
    match iter.next() {
        Some(TokenTree::Ident(ident)) => Ok(ident),
        Some(t) => error!(format!(r#"Expected identifier got "{}""#, t), t.span()),
        None => error!("Expected identifier".into(), span),
    }
}

fn parse_punct<T>(iter: &mut T, span: Span, ch: char) -> Result<Span>
where
    T: Iterator<Item = TokenTree>,
{
    match iter.next() {
        Some(TokenTree::Punct(p)) if p.as_char() == ch => Ok(p.span().end()),
        Some(t) => error!(format!(r#"Expected "{}", got "{}""#, ch, t), t.span()),
        None => error!(format!(r#"Expected "{}""#, ch), span),
    }
}

fn parse_literal<T>(iter: &mut T, span: Span) -> Result<Literal>
where
    T: Iterator<Item = TokenTree>,
{
    match iter.next() {
        Some(TokenTree::Literal(l)) => Ok(l),
        Some(t) => error!("Expected string literal in argument".to_owned(), t.span()),
        None => error!("Expected string literal in argument".to_owned(), span),
    }
}

fn parse_name(literal: &Literal) -> Result<String> {
    let name_string = literal.to_string();
    if name_string.len() <= 2 || !name_string.starts_with('"') || !name_string.ends_with('"') {
        return error!(
            "Provided name is not a string literal or is empty".into(),
            literal.span()
        );
    }

    let substr = &name_string[1..name_string.len() - 1];
    if let Some(_) = substr.find('\0') {
        error!(
            "Provided name is not a valid C string, because it contains internal null bytes".into(),
            literal.span()
        )
    } else {
        Ok(substr.into())
    }
}

fn parse_list_group<T>(iter: &mut T, span: Span) -> Result<Group>
where
    T: Iterator<Item = TokenTree>,
{
    match iter.next() {
        Some(TokenTree::Group(g)) if matches!(g.delimiter(), Delimiter::Bracket) => Ok(g),
        Some(t) => error!("Expected list of identifiers".into(), t.span()),
        _ => error!(
            "Expected list of identifiers, got nothing".into(),
            span.end()
        ),
    }
}

fn parse_ident_list<T>(iter: &mut T, span: Span) -> Result<(Span, Vec<String>)>
where
    T: Iterator<Item = TokenTree>,
{
    let list_group = parse_list_group(iter, span)?;
    let mut span = list_group.span().start();

    let mut list_iter = list_group.stream().into_iter().peekable();
    let mut ident_list = Vec::new();

    while let Some(_) = list_iter.peek() {
        let ident_name = parse_ident_name(&mut list_iter, span)?;
        span = ident_name.span().end();
        ident_list.push(ident_name.to_string());

        if let Some(_) = list_iter.peek() {
            span = parse_punct(&mut list_iter, span, ',')?;
        }
    }

    Ok((span, ident_list))
}

/// Skip outer attributes, these should start with '#' according to
/// https://doc.rust-lang.org/reference/items.html#items
fn skip_outer_attributes<T>(item: &mut std::iter::Peekable<T>, mut span: Span) -> Result<Span>
where
    T: Iterator<Item = TokenTree>,
{
    while let Some(TokenTree::Punct(p)) = item.peek() {
        span = p.span();
        if p.as_char() != '#' {
            return error!("Unexpected token".into(), span);
        }

        item.next();

        // The '#' should be followed by a bracket group
        if let Some(TokenTree::Group(_)) = item.peek() {
            item.next();
        } else {
            return error!("Unexpected token".into(), span.end());
        }
    }

    Ok(span.end())
}

fn skip_visibility_modifiers<T>(item: &mut std::iter::Peekable<T>, mut span: Span) -> Span
where
    T: Iterator<Item = TokenTree>,
{
    // Skip visibility modifiers
    if let Some(TokenTree::Ident(ident)) = item.peek() {
        if ident.to_string() == "pub" {
            span = ident.span();
            item.next();

            // If we have modifier on the visibility, also consume it
            if let Some(TokenTree::Group(g)) = item.peek() {
                span = g.span();
                item.next();
            }
        }
    }

    span.end()
}

fn check_struct_specifier<T>(item: &mut std::iter::Peekable<T>, span: Span) -> Result<Span>
where
    T: Iterator<Item = TokenTree>,
{
    match item.next() {
        Some(TokenTree::Ident(ident)) if ident.to_string() == "struct" => Ok(ident.span()),
        Some(TokenTree::Ident(ident)) => {
            error!("Only \"struct\" item supported".into(), ident.span())
        }
        Some(t) => error!("Unexpected token as item specifier".into(), t.span()),
        _ => error!("Missing item specifier".into(), span),
    }
}

fn parse_name_from_args<T>(args: &mut T, mut span: Span) -> Result<(String, Span)>
where
    T: Iterator<Item = TokenTree>,
{
    span = parse_ident(args, span, "name")?;
    span = parse_punct(args, span, '=')?;
    let name_literal = parse_literal(args, span)?;
    Ok((parse_name(&name_literal)?, name_literal.span().end()))
}

fn parse_item_name<T>(item: &mut std::iter::Peekable<T>, mut span: Span) -> Result<(String, Span)>
where
    T: Iterator<Item = TokenTree>,
{
    span = skip_outer_attributes(item, span)?;
    span = skip_visibility_modifiers(item, span);
    span = check_struct_specifier(item, span)?;
    let item_name = parse_ident_name(item, span)?;
    Ok((item_name.to_string(), item_name.span()))
}

fn rule_impl(args: TokenStream, item: TokenStream) -> Result<TokenStream> {
    let span = Span::call_site();
    let mut args_iter = args.into_iter();
    let (name, span) = parse_name_from_args(&mut args_iter, span)?;
    let mut item_iter = item.clone().into_iter().peekable();
    let (item_name, _) = parse_item_name(&mut item_iter, span)?;
    let name_chars: Vec<i8> = name
        .bytes()
        .map(|b| b as i8)
        .chain(std::iter::once(0i8))
        .collect();

    let mut stream = format!(
        r#"
        #[unsafe(no_mangle)]
        pub static reoxide_rule_name_{item_name}: [i8; {rule_name_chars_len}] = {rule_name_chars:?};

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_rule_new_{item_name}(
            context: *const ::reoxide::plugin::Context,
            plugin: *mut ::reoxide::plugin::OpaquePlugin,
            args: *const ::reoxide::plugin::ExtraArgs,
        ) -> *mut ::reoxide::plugin::OpaqueRule {{
            unsafe {{
                let args_ref = &*args;
                let group_name_cstr = ::core::ffi::CStr::from_ptr(args_ref.group_name);
                let group_name = match group_name_cstr.to_str() {{
                    Ok(s) => s,
                    Err(_) => return ::core::ptr::null_mut(),
                }};
                let context_ref = &*context;
                let arch = &mut *context_ref.arch;
                let stack_space = &mut *context_ref.stackspace;
                let reoxide = &mut *context_ref.reoxide;
                let plugin_ref = &mut *plugin.cast();
                Box::into_raw(
                        Box::new(
                            {item_name}::new(::reoxide::plugin::CreationContext {{
                                arch,
                                stack_space,
                                reoxide,
                                plugin: plugin_ref,
                                group_name,
                                extra_args: args_ref.extra_arg
                            }}),
                        ),
                    )
                    .cast()
            }}
        }}

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_rule_delete_{item_name}(rule: *mut ::reoxide::plugin::OpaqueRule) {{
            if !rule.is_null() {{
                let _ = unsafe {{
                    let concrete_rule: *mut {item_name} = rule.cast();
                    Box::from_raw(concrete_rule)
                }};
            }}
        }}

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_rule_oplist_{item_name}(
            this: *const ::reoxide::plugin::OpaqueRule,
            oplist: *mut ::reoxide::cpp::StdVector<::reoxide::plugin::OpCode>,
        ) {{
            let rule = unsafe {{ &*this.cast::<{item_name}>() }};
            let oplist: &mut ::reoxide::cpp::StdVector<::reoxide::plugin::OpCode> = unsafe {{ &mut *oplist }};
            oplist.extend(&rule.op_list());
        }}


        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_rule_apply_{item_name}(
            this: *mut ::reoxide::plugin::OpaqueRule,
            op: *mut ::reoxide::plugin::PcodeOp,
            data: *mut ::reoxide::plugin::Funcdata,
        ) -> i32 {{
            unsafe {{
                let rule = &mut *this.cast::<{item_name}>();
                let op_ref = &mut *op;
                let data_ref = &mut *data;
                rule.apply(op_ref, data_ref) as i32
            }}
        }}
        "#,
        item_name = item_name,
        rule_name_chars = name_chars,
        rule_name_chars_len = name_chars.len()
    )
    .parse::<TokenStream>()
    .expect("Proc macro failed to produce output stream");
    stream.extend(item);
    Ok(stream)
}

fn action_impl(args: TokenStream, item: TokenStream) -> Result<TokenStream> {
    let span = Span::call_site();
    let mut args_iter = args.into_iter();
    let (name, span) = parse_name_from_args(&mut args_iter, span)?;
    let mut item_iter = item.clone().into_iter().peekable();
    let (item_name, _) = parse_item_name(&mut item_iter, span)?;
    let name_chars: Vec<i8> = name
        .bytes()
        .map(|b| b as i8)
        .chain(std::iter::once(0i8))
        .collect();

    let mut stream = format!(
        r#"
        #[unsafe(no_mangle)]
        pub static reoxide_action_name_{item_name}: [i8; {action_name_chars_len}] = {action_name_chars:?};

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_action_new_{item_name}(
            context: *const ::reoxide::plugin::Context,
            plugin: *mut ::reoxide::plugin::OpaquePlugin,
            args: *const ::reoxide::plugin::ExtraArgs,
        ) -> *mut ::reoxide::plugin::OpaqueAction {{
            unsafe {{
                let args_ref = &*args;
                let group_name_cstr = ::core::ffi::CStr::from_ptr(args_ref.group_name);
                let group_name = match group_name_cstr.to_str() {{
                    Ok(s) => s,
                    Err(_) => return ::core::ptr::null_mut(),
                }};
                let context_ref = &*context;
                let arch = &mut *context_ref.arch;
                let stack_space = &mut *context_ref.stackspace;
                let reoxide = &mut *context_ref.reoxide;
                let plugin_ref = &mut *plugin.cast();
                Box::into_raw(
                        Box::new(
                            {item_name}::new(::reoxide::plugin::CreationContext {{
                                arch,
                                stack_space,
                                reoxide,
                                plugin: plugin_ref,
                                group_name,
                                extra_args: args_ref.extra_arg
                            }}),
                        ),
                    )
                    .cast()
            }}
        }}

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_action_delete_{item_name}(action: *mut ::reoxide::plugin::OpaqueAction) {{
            if !action.is_null() {{
                let _ = unsafe {{
                    let concrete_action: *mut {item_name} = action.cast();
                    Box::from_raw(concrete_action)
                }};
            }}
        }}

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_action_apply_{item_name}(
            this: *mut ::reoxide::plugin::OpaqueAction,
            data: *mut ::reoxide::plugin::Funcdata,
        ) -> i32 {{
            unsafe {{
                let action = &mut *this.cast::<{item_name}>();
                let data_ref = &mut *data;
                action.apply(data_ref) as i32
            }}
        }}"#,
        item_name = item_name,
        action_name_chars = name_chars,
        action_name_chars_len = name_chars.len()
    )
    .parse::<TokenStream>()
    .expect("Proc macro failed to produce output stream");
    stream.extend(item);
    Ok(stream)
}

#[proc_macro_attribute]
pub fn rule(args: TokenStream, item: TokenStream) -> TokenStream {
    match rule_impl(args, item.clone()) {
        Ok(s) => s,
        Err(e) => e.into_combined_token_stream(item),
    }
}

#[proc_macro_attribute]
pub fn action(args: TokenStream, item: TokenStream) -> TokenStream {
    match action_impl(args, item.clone()) {
        Ok(s) => s,
        Err(e) => e.into_combined_token_stream(item),
    }
}

fn rule_definition_to_stream(ident: &str, span: Span) -> TokenStream {
    let definition = [
        ident!("name", span),
        punct!(':', span),
        ident!(&format!("reoxide_rule_name_{}", ident), span),
        punct!('.', span),
        ident!("as_ptr", span),
        braces!(Delimiter::Parenthesis, span),
        punct!(',', span),
        ident!("constructor", span),
        punct!(':', span),
        ident!(&format!("reoxide_rule_new_{}", ident), span),
        punct!(',', span),
        ident!("destructor", span),
        punct!(':', span),
        ident!(&format!("reoxide_rule_delete_{}", ident), span),
        punct!(',', span),
        ident!("oplist", span),
        punct!(':', span),
        ident!(&format!("reoxide_rule_oplist_{}", ident), span),
        punct!(',', span),
        ident!("apply", span),
        punct!(':', span),
        ident!(&format!("reoxide_rule_apply_{}", ident), span),
    ];

    TokenStream::from_iter([
        jpunct!(':', span),
        punct!(':', span),
        ident!("reoxide", span),
        jpunct!(':', span),
        punct!(':', span),
        ident!("plugin", span),
        jpunct!(':', span),
        punct!(':', span),
        ident!("RuleDefinition", span),
        group!(Delimiter::Brace, definition, span),
        punct!(',', span),
    ])
}

fn action_definition_to_stream(ident: &str, span: Span) -> TokenStream {
    let definition = [
        ident!("name", span),
        punct!(':', span),
        ident!(&format!("reoxide_action_name_{}", ident), span),
        punct!('.', span),
        ident!("as_ptr", span),
        braces!(Delimiter::Parenthesis, span),
        punct!(',', span),
        ident!("constructor", span),
        punct!(':', span),
        ident!(&format!("reoxide_action_new_{}", ident), span),
        punct!(',', span),
        ident!("destructor", span),
        punct!(':', span),
        ident!(&format!("reoxide_action_delete_{}", ident), span),
        punct!(',', span),
        ident!("apply", span),
        punct!(':', span),
        ident!(&format!("reoxide_action_apply_{}", ident), span),
    ];

    TokenStream::from_iter([
        jpunct!(':', span),
        punct!(':', span),
        ident!("reoxide", span),
        jpunct!(':', span),
        punct!(':', span),
        ident!("plugin", span),
        jpunct!(':', span),
        punct!(':', span),
        ident!("ActionDefinition", span),
        group!(Delimiter::Brace, definition, span),
        punct!(',', span),
    ])
}

fn plugin_impl(input: TokenStream) -> Result<TokenStream> {
    let mut span = Span::call_site();
    let mut args_iter = input.into_iter();

    span = parse_ident(&mut args_iter, span, "context")?;
    span = parse_punct(&mut args_iter, span, '=')?;
    let plugin_ident = parse_ident_name(&mut args_iter, span)?;
    span = parse_punct(&mut args_iter, span, ',')?;
    span = parse_ident(&mut args_iter, span, "rules")?;
    span = parse_punct(&mut args_iter, span, '=')?;
    let (mut span, rules) = parse_ident_list(&mut args_iter, span)?;
    span = parse_punct(&mut args_iter, span, ',')?;
    span = parse_ident(&mut args_iter, span, "actions")?;
    span = parse_punct(&mut args_iter, span, '=')?;
    let (_, actions) = parse_ident_list(&mut args_iter, span)?;

    let mut out_stream = TokenStream::new();

    out_stream.extend(
        format!(
            r#"
        #[unsafe(no_mangle)]
        pub static reoxide_c_abi: usize = 1;

        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_plugin_new() -> *mut ::reoxide::plugin::OpaquePlugin {{
            ::std::panic::set_hook(Box::new(reoxide::plugin::panic_hook));
            Box::into_raw(Box::new({plugin_name} {{}})).cast()
        }}
        
        #[unsafe(no_mangle)]
        pub extern "C-unwind" fn reoxide_plugin_delete(plugin: *mut {plugin_name}) {{
            if !plugin.is_null() {{
                let _ = unsafe {{ Box::from_raw(plugin) }};
            }}
        }}
        "#,
            plugin_name = plugin_ident.to_string()
        )
        .parse::<TokenStream>()
        .expect("Proc macro failed to produce output stream"),
    );

    out_stream.extend(
        format!(
            r#"
            #[unsafe(no_mangle)]
            pub static reoxide_rule_count: usize = {rule_count};

            #[unsafe(no_mangle)]
            pub static reoxide_rule_defs: ::reoxide::plugin::DefinitionTable<
                [::reoxide::plugin::RuleDefinition; {rule_count}]
            > = ::reoxide::plugin::DefinitionTable::<
                [::reoxide::plugin::RuleDefinition; {rule_count}]
            >"#,
            rule_count = rules.len()
        )
        .parse::<TokenStream>()
        .expect("Proc macro failed to produce output stream"),
    );

    let mut rule_stream = TokenStream::new();
    rules
        .into_iter()
        .for_each(|ident| rule_stream.extend(rule_definition_to_stream(&ident, Span::call_site())));

    out_stream.extend([
        group!(
            Delimiter::Brace,
            [
                ident!("table"),
                punct!(':'),
                group!(Delimiter::Bracket, rule_stream),
            ]
        ),
        punct!(';'),
    ]);

    out_stream.extend(
        format!(
            r#"
            #[unsafe(no_mangle)]
            pub static reoxide_action_count: usize = {action_count};

            #[unsafe(no_mangle)]
            pub static reoxide_action_defs: ::reoxide::plugin::DefinitionTable<
                [::reoxide::plugin::ActionDefinition; {action_count}]
            > = ::reoxide::plugin::DefinitionTable::<
                [::reoxide::plugin::ActionDefinition; {action_count}]
            >"#,
            action_count = actions.len()
        )
        .parse::<TokenStream>()
        .expect("Proc macro failed to produce output stream"),
    );

    let mut action_stream = TokenStream::new();
    actions.into_iter().for_each(|ident| {
        action_stream.extend(action_definition_to_stream(&ident, Span::call_site()))
    });

    out_stream.extend([
        group!(
            Delimiter::Brace,
            [
                ident!("table"),
                punct!(':'),
                group!(Delimiter::Bracket, action_stream),
            ]
        ),
        punct!(';'),
    ]);

    Ok(out_stream)
}

#[proc_macro]
pub fn plugin(input: TokenStream) -> TokenStream {
    match plugin_impl(input) {
        Ok(s) => s,
        Err(e) => e.into_token_stream(),
    }
}