nyanoblog/macros/src/define_json_ns.rs

use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
use std::iter::Peekable;
use std::sync::atomic::{AtomicIsize, Ordering};

use crate::util::*;

// "forgive me, rustc, for i have sinned"
//
// good luck to whoever is trying to read not to mention understand this lmao

/// This is a global counter for unique indices because every invocation
/// of this macro should yield an enum with unique numeric representations
static INDEX: AtomicIsize = AtomicIsize::new(1);

pub fn invoke(input: TokenStream) -> Result<TokenStream> {
    let mut stream = input.into_iter().peekable();

    // parse the thing

    let (mut meta, mut classes, mut props) = (None, None, None);
    while let Some(tt) = stream.next() {
        if let TokenTree::Ident(ident) = &tt {
            let name = ident.to_string();
            match name.as_str() {
                "meta" => {
                    let body = eat_group(&mut stream, Delimiter::Brace)?;
                    if meta.replace(eat_meta(body.stream())?).is_some() {
                        return err("Duplicate definition of metadata", ident.span());
                    }
                }
                "class" => {
                    let body = eat_group(&mut stream, Delimiter::Brace)?;
                    let body = eat_enum_members(body.stream())?;
                    if classes.replace((body, ident.span())).is_some() {
                        return err("Duplicate class definition", ident.span());
                    }
                }
                "prop" => {
                    let body = eat_group(&mut stream, Delimiter::Brace)?;
                    let body = eat_enum_members(body.stream())?;
                    if props.replace((body, ident.span())).is_some() {
                        return err("Duplicate props definition", ident.span());
                    }
                }
                name => return err(format!("Unknown section \"{name}\""), ident.span()),
            }
            eat_maybe_comma(&mut stream);
        }
    }

    // now spit out some new tokens

    let meta = meta.ok_or_else(|| error("Missing meta block", None))?;
    let mut ts = TokenStream::new();
    if let Some((classes, span)) = classes {
        ts.extend(compile_enum(
            &meta,
            make_ident_tt("__Class", None),
            &classes,
            span,
        ));
    }
    if let Some((props, span)) = props {
        ts.extend(compile_enum(
            &meta,
            make_ident_tt("__Prop", None),
            &props,
            span,
        ));
    }
    let mut module = rust_code!(pub mod );
    module.extend([meta.mod_name, make_group_tt(Delimiter::Brace, ts)]);
    //println!("{module}");
    Ok(module)
}

struct NsMeta {
    mod_name: TokenTree,
    iri_base: String,
    alias_base: Option<String>,
}

#[derive(Clone)]
struct EnumMember {
    name: Ident,
}

//
// parser
//

// meta { key1 = val1, key2 = val2, ... }
fn eat_meta(stream: TokenStream) -> Result<NsMeta> {
    let mut stream = stream.into_iter().peekable();
    let mut mod_name = None;
    let mut iri_base = None;
    let mut alias_base = None;
    let mut had_preceding_comma = true;
    while let Some(ident) = eat_ident_or_end(&mut stream)? {
        if !had_preceding_comma {
            return err("Expected a comma", ident.span());
        }

        let name = ident.to_string();
        match name.as_str() {
            "mod_name" => {
                if mod_name.replace(eat_assign_ident(&mut stream)?).is_some() {
                    return err("Duplicate definition of key \"mod_name\"", ident.span());
                }
            }
            "iri_base" => {
                if iri_base.replace(eat_assign_string(&mut stream)?).is_some() {
                    return err("Duplicate definition of key \"iri_base\"", ident.span());
                }
            }
            "alias_base" => {
                if alias_base
                    .replace(eat_assign_string(&mut stream)?)
                    .is_some()
                {
                    return err("Duplicate definition of key \"alias_base\"", ident.span());
                }
            }
            name => return err(format!("Unknown key \"{name}\""), ident.span()),
        }
        had_preceding_comma = eat_maybe_comma(&mut stream);
    }

    Ok(NsMeta {
        mod_name: mod_name.ok_or_else(|| error("Missing key \"mod_name\"", None))?,
        iri_base: iri_base.ok_or_else(|| error("Missing key \"iri_base\"", None))?,
        alias_base,
    })
}

// qualifier { Member1, Member2, ... }
fn eat_enum_members(stream: TokenStream) -> Result<Vec<EnumMember>> {
    let mut stream = stream.into_iter().peekable();
    let mut members = Vec::new();
    let mut had_preceding_comma = true;
    while let Some(member) = eat_enum_member(&mut stream)? {
        if !had_preceding_comma {
            return err_unexpected(",", &member.name, member.name.span());
        }

        members.push(member);
        had_preceding_comma = eat_maybe_comma(&mut stream);
    }
    Ok(members)
}

fn eat_enum_member(stream: &mut impl TtIter) -> Result<Option<EnumMember>> {
    match stream.next() {
        Some(tt) => {
            if let TokenTree::Ident(name) = tt {
                Ok(Some(EnumMember { name }))
            } else {
                err_unexpected("ident", &tt, tt.span())
            }
        }
        None => Ok(None),
    }
}

fn eat_assign_ident(stream: &mut Peekable<impl TtIter>) -> Result<TokenTree> {
    eat_eq(stream)?;
    Ok(TokenTree::Ident(eat_ident(stream)?))
}

fn eat_assign_string(stream: &mut impl TtIter) -> Result<String> {
    let _ = eat_eq(stream)?;
    match stream.next() {
        Some(TokenTree::Literal(literal)) => {
            let s = literal.to_string();
            let sb = s.as_bytes();
            if sb.len() >= 2 && sb[0] == b'"' && sb[sb.len() - 1] == b'"' {
                Ok(String::from_utf8(Vec::from(&sb[1..(sb.len() - 1)])).unwrap())
            } else {
                err_unexpected("string", &literal, literal.span())
            }
        }
        Some(tt) => err_unexpected("string", &tt, tt.span()),
        None => err_end(),
    }
}

fn eat_group(stream: &mut impl TtIter, delimiter: impl Into<Option<Delimiter>>) -> Result<Group> {
    match stream.next() {
        Some(TokenTree::Group(grp)) => {
            if let Some(delim) = delimiter.into() {
                if grp.delimiter() == delim {
                    Ok(grp)
                } else {
                    err("Expected a `{`", grp.span())
                }
            } else {
                Ok(grp)
            }
        }
        Some(tt) => err_unexpected("{", &tt, tt.span()),
        None => err_end(),
    }
}

fn eat_ident(stream: &mut impl TtIter) -> Result<Ident> {
    eat_ident_or_end(stream).and_then(|o| o.ok_or_else(error_end))
}

fn eat_ident_or_end(stream: &mut impl TtIter) -> Result<Option<Ident>> {
    match stream.next() {
        Some(TokenTree::Ident(ident)) => Ok(Some(ident)),
        Some(tt) => err_unexpected("ident", &tt, tt.span()),
        None => Ok(None),
    }
}

fn eat_eq(stream: &mut impl TtIter) -> Result<TokenTree> {
    match stream.next() {
        Some(TokenTree::Punct(punct)) if punct.as_char() == '=' => Ok(TokenTree::Punct(punct)),
        Some(tt) => err_unexpected("=", &tt, tt.span()),
        None => err_end(),
    }
}

fn eat_maybe_comma(stream: &mut Peekable<impl TtIter>) -> bool {
    stream
        .next_if(|tt| {
            if let TokenTree::Punct(punct) = tt {
                punct.as_char() == ','
            } else {
                false
            }
        })
        .is_some()
}

//
// compiler
//

fn compile_enum(
    meta: &NsMeta,
    name: TokenTree,
    members: &[EnumMember],
    kw_span: Span,
) -> TokenStream {
    // enum Name { Member1, Member2, ... }
    let mut ts = rust_code!(
        #[derive(Copy, Clone, Eq, PartialEq)]
        #[allow(non_camel_case_types)]
        pub
    );
    ts.extend([
        make_ident_tt("enum", kw_span),
        name.clone(),
        make_group_tt(
            Delimiter::Brace,
            members.iter().flat_map(|m| {
                [
                    TokenTree::Ident(m.name.clone()),
                    TokenTree::Punct(Punct::new('=', Spacing::Alone)),
                    TokenTree::Literal(Literal::isize_unsuffixed(
                        INDEX.fetch_add(1, Ordering::Relaxed),
                    )),
                    TokenTree::Punct(Punct::new(',', Spacing::Alone)),
                ]
            }),
        ),
    ]);

    ts.extend(rust_code!(pub use ));
    ts.extend([name.clone()]);
    ts.extend(rust_code!(::*;));

    // impl Enum { ... }
    ts.extend(compile_internal_impl(meta, &name, members));

    // impl HasContext for Enum { ... }
    ts.extend(compile_trait_impl(meta, &name, members));

    ts
}

fn compile_internal_impl(meta: &NsMeta, name: &TokenTree, members: &[EnumMember]) -> TokenStream {
    let mut ts = rust_code!(impl);
    ts.extend([name.clone()]);

    let impl_body = {
        let mut impl_body = rust_code!(const __MEMBERS: );
        impl_body.extend([
            // [(Self, Iri<'static>); $n]
            make_group_tt(
                Delimiter::Bracket,
                [
                    make_group_tt(
                        Delimiter::Parenthesis,
                        rust_code!(Self, ::iref::Iri<'static>),
                    ),
                    make_punct_alone(';'),
                    TokenTree::Literal(Literal::usize_unsuffixed(members.len())),
                ],
            ),
            make_punct_alone('='),
        ]);

        impl_body.extend([
            // [ (Enum::Name1, iri!("iri1")), (Enum::Name2, iri!("iri2")), ... ]
            make_group_tt(
                Delimiter::Bracket,
                members.iter().flat_map(|memb| {
                    [
                        make_group_tt(Delimiter::Parenthesis, {
                            let mut ts = TokenStream::from_iter([
                                name.clone(),
                                make_punct_joint(':'),
                                make_punct_alone(':'),
                                TokenTree::Ident(memb.name.clone()),
                                make_punct_alone(','),
                            ]);
                            ts.extend(compile_iri_macro_call(&meta.iri_base, Some(memb)));
                            ts
                        }),
                        make_punct_alone(','),
                    ]
                    .into_iter()
                }),
            ),
            TokenTree::Punct(Punct::new(';', Spacing::Alone)),
        ]);

        impl_body.extend(rust_code!(fn __from_str(s: &str) -> Option<Self>));
        let fn_body = {
            let mut fn_body = rust_code!(match s);

            let match_body = {
                let mut match_body = TokenStream::new();
                for memb in members {
                    let memb_name = memb.name.to_string();
                    match_body.extend([
                        TokenTree::Literal(Literal::string(&memb_name)),
                        make_punct_joint('='),
                        make_punct_alone('>'),
                        make_ident_tt("Some", None),
                        make_group_tt(
                            Delimiter::Parenthesis,
                            [
                                name.clone(),
                                make_punct_joint(':'),
                                make_punct_alone(':'),
                                TokenTree::Ident(memb.name.clone()),
                            ],
                        ),
                        make_punct_alone(','),
                    ]);
                }
                match_body.extend(rust_code!(_ => None,));
                make_group_ts(Delimiter::Brace, match_body)
            };
            fn_body.extend(match_body);

            make_group_ts(Delimiter::Brace, fn_body)
        };
        impl_body.extend(fn_body);

        make_group_ts(Delimiter::Brace, impl_body)
    };
    ts.extend(impl_body);

    ts
}

fn compile_trait_impl(meta: &NsMeta, name: &TokenTree, members: &[EnumMember]) -> TokenStream {
    // impl HasContext for Enum
    let mut ts = rust_code!(impl);
    ts.extend(absolute_path(None, rust_code!(ap::vocab::HasContext)));
    ts.extend(rust_code!(for));
    ts.extend([name.clone()]);

    let impl_body = {
        // const ... = ...; const ... = ...; ...
        let mut impl_body = compile_impl_const_members(meta, members);

        // fn from_index(...) -> ... { ... }
        impl_body.extend(compile_impl_from_index(name, members));

        impl_body.extend(rust_code!(
            fn as_index(&self) -> usize {
                *self as usize
            }
        ));

        // fn from_iri(...) -> ... { ... }
        impl_body.extend(compile_impl_from_iri());

        make_group_ts(Delimiter::Brace, impl_body)
    };
    ts.extend(impl_body);

    ts
}

fn compile_impl_const_members(meta: &NsMeta, members: &[EnumMember]) -> TokenStream {
    let mut ts = TokenStream::new();

    ts.extend(rust_code!(const COUNT: usize = ));
    ts.extend([
        TokenTree::Literal(Literal::usize_suffixed(members.len())),
        TokenTree::Punct(Punct::new(';', Spacing::Alone)),
    ]);

    ts.extend(rust_code!(const OFFSET: usize = ));
    ts.extend([
        TokenTree::Literal(Literal::usize_suffixed(
            INDEX.load(Ordering::Relaxed) as usize - members.len(),
        )),
        TokenTree::Punct(Punct::new(';', Spacing::Alone)),
    ]);

    ts.extend(rust_code!(const IRI_BASE: ::iref::Iri<'static> = ));
    ts.extend(compile_iri_macro_call(&meta.iri_base, None));
    ts.extend([make_punct_alone(';')]);

    ts.extend(rust_code!(const ALIAS_BASE: ::core::option::Option<::iref::Iri<'static>> =));
    match meta.alias_base.as_ref() {
        Some(s) => {
            ts.extend(rust_code!(Some));
            ts.extend(make_group_ts(
                Delimiter::Parenthesis,
                compile_iri_macro_call(s, None),
            ));
        }
        None => ts.extend(rust_code!(None)),
    }
    ts.extend(rust_code!(;));

    // this is so we get a ref with static lifetime; see compile_internal_impl()
    ts.extend(rust_code!(
        const MEMBERS: &'static [(Self, ::iref::Iri<'static>)] = &Self::__MEMBERS;
    ));

    ts
}

fn compile_impl_from_index(name: &TokenTree, members: &[EnumMember]) -> TokenStream {
    let mut ts = rust_code!(fn from_index(index: usize) -> Option<Self>);

    let fn_body = {
        let mut fn_body = rust_code!(match index);

        let match_body = {
            let mut match_body = TokenStream::new();

            for member in members {
                match_body.extend(rust_code!(x if x == ));
                match_body.extend([
                    name.clone(),
                    make_punct_joint(':'),
                    make_punct_alone(':'),
                    TokenTree::Ident(member.name.clone()),
                ]);
                match_body.extend(rust_code!(as usize => Some));
                match_body.extend(make_group_ts(
                    Delimiter::Parenthesis,
                    [
                        name.clone(),
                        make_punct_joint(':'),
                        make_punct_alone(':'),
                        TokenTree::Ident(member.name.clone()),
                    ],
                ));
                match_body.extend([make_punct_alone(',')]);
            }
            match_body.extend(rust_code!(_ => None,));

            make_group_ts(Delimiter::Brace, match_body)
        };
        fn_body.extend(match_body);

        make_group_ts(Delimiter::Brace, fn_body)
    };
    ts.extend(fn_body);

    ts
}

fn compile_impl_from_iri() -> TokenStream {
    rust_code!(
        fn from_iri(iri: &::iref::Iri) -> Option<Self> {
            let iri = iri.as_str();
            let iri_base = Self::IRI_BASE.as_str();

            if iri.starts_with(iri_base) {
                Self::__from_str(unsafe {
                    ::std::str::from_utf8_unchecked(&iri.as_bytes()[iri_base.len()..])
                })
            } else if let Some(alias_base) = Self::ALIAS_BASE.as_ref() {
                let alias_base = alias_base.as_str();
                if iri.starts_with(alias_base) {
                    Self::__from_str(unsafe {
                        ::std::str::from_utf8_unchecked(&iri.as_bytes()[alias_base.len()..])
                    })
                } else {
                    None
                }
            } else {
                None
            }
        }
    )
}

fn compile_iri_macro_call(iri_base: &str, member: Option<&EnumMember>) -> TokenStream {
    let mut ts = rust_code!(::static_iref::iri!);
    let iri = match member {
        Some(member) => format!("{iri_base}{}", member.name),
        None => String::from(iri_base),
    };
    ts.extend(make_group_ts(
        Delimiter::Parenthesis,
        [TokenTree::Literal(Literal::string(&iri))],
    ));
    ts
}