preserves/implementations/rust/preserves-schema-macros/src/lib.rs

use preserves::value::Map;
use preserves_schema::compiler::*;
use preserves_schema::compiler::types::Purpose;
use preserves_schema::gen::schema::Schema;
use proc_macro2::Span;
use quote::ToTokens;
use quote::quote;
use std::fmt::Display;
use syn::LitStr;
use syn::Token;
use syn::parenthesized;
use syn::parse::Parser;
use syn::punctuated::Punctuated;

mod kw {
    use syn::custom_keyword;
    custom_keyword!(load);
    custom_keyword!(cross_reference);
    custom_keyword!(external_module);
}

#[derive(Debug)]
enum Instruction {
    Namespace(String),
    Load(LitStr),
    CrossReference {
        namespace: String,
        bundle_path: LitStr,
    },
    ExternalModule {
        module_path: ModulePath,
        namespace: String,
    }
}

fn syn_path_string(p: syn::Path) -> String {
    p.to_token_stream().to_string().replace(&[' ', '\t', '\n', '\r'], "")
}

fn syn_litstr_resolve(s: &syn::LitStr) -> String {
    let s: String = s.value();
    match s.chars().nth(0) {
        Some('/') => s.into(),
        Some('<') => match &s[1..].split_once('>') {
            Some((envvar, remainder)) => match std::env::var(envvar) {
                Ok(p) => p + "/" + remainder,
                Err(_) => panic!("No such environment variable: {:?}", s),
            }
            None => panic!("Invalid relative path syntax: {:?}", s),
        },
        _ => panic!("Invalid path syntax: {:?}", s)
    }
}

impl syn::parse::Parse for Instruction {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let lookahead = input.lookahead1();
        if lookahead.peek(kw::load) {
            let _: kw::load = input.parse()?;
            let content;
            let _ = parenthesized!(content in input);
            let bundle_path: syn::LitStr = content.parse()?;
            Ok(Instruction::Load(bundle_path))
        } else if lookahead.peek(kw::cross_reference) {
            let _: kw::cross_reference = input.parse()?;
            let content;
            let _ = parenthesized!(content in input);
            let namespace: syn::Path = content.parse()?;
            let _: Token![=] = content.parse()?;
            let bundle_path: syn::LitStr = content.parse()?;
            Ok(Instruction::CrossReference {
                namespace: syn_path_string(namespace),
                bundle_path,
            })
        } else if lookahead.peek(kw::external_module) {
            let _: kw::external_module = input.parse()?;
            let content;
            let _ = parenthesized!(content in input);
            let module_path = Punctuated::<syn::Ident, syn::Token![.]>::parse_separated_nonempty(&content)?;
            let _: Token![=] = content.parse()?;
            let namespace: syn::Path = content.parse()?;
            Ok(Instruction::ExternalModule {
                module_path: module_path.into_iter().map(|p| p.to_string()).collect(),
                namespace: syn_path_string(namespace),
            })
        } else {
            let ns: syn::Path = input.parse()?;
            Ok(Instruction::Namespace(syn_path_string(ns)))
        }
    }
}

struct ModuleTree {
    own_body: String,
    children: Map<String, ModuleTree>,
}

impl Default for ModuleTree {
    fn default() -> Self {
        ModuleTree {
            own_body: String::new(),
            children: Map::default(),
        }
    }
}

impl ModuleTree {
    fn build(outputs: Map<Option<ModulePath>, String>) -> Self {
        let mut mt = ModuleTree::default();
        for (p, c) in outputs.into_iter() {
            match p {
                None => mt.own_body = c,
                Some(k) => {
                    let mut r = &mut mt;
                    for e in k { r = mt.children.entry(names::render_modname(&e)).or_default(); }
                    r.own_body = c;
                }
            }
        }
        mt
    }
}

impl Display for ModuleTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.own_body)?;
        for (label, mt) in self.children.iter() {
            write!(f, "\npub mod {} {{ ", label)?;
            mt.fmt(f)?;
            write!(f, "}}")?;
        }
        Ok(())
    }
}

#[proc_macro]
pub fn compile_preserves_schemas(src: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let instructions = Punctuated::<Instruction, syn::Token![,]>::parse_terminated
        .parse(src)
        .expect("valid sequence of compile_preserves_schemas instructions");

    let mut namespace = None::<String>;
    let mut bundles_to_load = Vec::<LitStr>::new();
    let mut bundles_to_xref = Vec::<LitStr>::new();
    let mut external_modules = Vec::<ExternalModule>::new();

    for i in instructions.into_iter() {
        match i {
            Instruction::Namespace(n) => {
                if namespace.is_some() {
                    panic!("Only one namespace is permitted")
                }
                namespace = Some(n)
            }
            Instruction::Load(p) => bundles_to_load.push(p),
            Instruction::ExternalModule { module_path, namespace } => {
                external_modules.push(ExternalModule::new(module_path, &namespace));
            }
            Instruction::CrossReference { namespace, bundle_path } => {
                let mut bundle = Map::<ModulePath, Schema>::new();
                let is_schema = load_schema_or_bundle(&mut bundle, &syn_litstr_resolve(&bundle_path).into())
                    .expect("Invalid schema/bundle binary");
                bundles_to_xref.push(bundle_path);
                for (k, _v) in bundle.into_iter() {
                    external_modules.push(if is_schema {
                        ExternalModule::new(k, &namespace)
                    } else {
                        let ns = namespace.clone();
                        let mut pieces = vec![ns.clone()];
                        pieces.extend(
                            k.iter().map(|p| names::render_modname(&p)).collect::<Vec<_>>());
                        ExternalModule::new(k, &pieces.join("::"))
                            .set_fallback_language_types(
                                move |v| vec![format!("{}::Language<{}>", ns, v)].into_iter().collect())
                    });
                }
            }
        }
    }

    let namespace = namespace.expect("Missing namespace");

    let mut dependency_paths = Vec::<syn::LitStr>::new();

    let mut c = CompilerConfig::new(namespace.clone());
    for b in bundles_to_load.into_iter() {
        dependency_paths.push(syn::LitStr::new(&syn_litstr_resolve(&b), Span::call_site()));
        load_schema_or_bundle_with_purpose(&mut c.bundle, &syn_litstr_resolve(&b).into(), Purpose::Codegen)
            .expect(&b.value());
    }
    for b in bundles_to_xref.into_iter() {
        dependency_paths.push(syn::LitStr::new(&syn_litstr_resolve(&b), Span::call_site()));
        load_schema_or_bundle_with_purpose(&mut c.bundle, &syn_litstr_resolve(&b).into(), Purpose::Xref)
            .expect(&b.value());
    }
    for m in external_modules.into_iter() {
        c.add_external_module(m);
    }

    let mut outputs = Map::<Option<ModulePath>, String>::new();
    let mut collector = CodeCollector {
        emit_mod_declarations: false,
        collect_module: CodeModuleCollector::Custom {
            collect_output: &mut |p, c| {
                outputs.insert(p.cloned(), c.to_owned());
                Ok(())
            },
        },
    };
    compile(&c, &mut collector).expect("Compilation failed");

    let top_module_source = format!(
        "pub mod {} {{ {} }}",
        names::render_modname(namespace.split("::").last().unwrap()),
        ModuleTree::build(outputs));
    let top_module: syn::Item = syn::parse_str(&top_module_source)
        .expect("Invalid generated code");

    quote!{
        // TODO: this is ugly, but makes the code depend on the actual schema bundle:
        // See https://doc.rust-lang.org/nightly/proc_macro/tracked_path/fn.path.html
        // and https://github.com/rust-lang/rust/issues/99515
        #( const _: &'static [u8] = include_bytes!(#dependency_paths); )*

        #top_module
    }.into()
}