Merge pull request #112 from Nadrieril/improve-parser

Add features to parser macros

1 files changed, 274 insertions, 51 deletions

diff --git a/dhall_proc_macros/src/make_parser.rs b/dhall_proc_macros/src/make_parser.rs
index 268a639..a17ab61 100644
--- a/dhall_proc_macros/src/make_parser.rs
+++ b/dhall_proc_macros/src/make_parser.rs
@@ -1,70 +1,188 @@
+use std::collections::HashMap;
+use std::iter;
+
 use quote::quote;
-use syn::parse::{ParseStream, Result};
+use syn::parse::{Parse, ParseStream, Result};
 use syn::spanned::Spanned;
 use syn::{
     parse_quote, Error, Expr, FnArg, Ident, ImplItem, ImplItemMethod, ItemImpl,
-    Pat, Token,
+    LitBool, Pat, Token,
 };
 
-fn apply_special_attrs(function: &mut ImplItemMethod) -> Result<()> {
+mod kw {
+    syn::custom_keyword!(shortcut);
+}
+
+struct AliasArgs {
+    target: Ident,
+    is_shortcut: bool,
+}
+
+struct PrecClimbArgs {
+    child_rule: Ident,
+    climber: Expr,
+}
+
+struct AliasSrc {
+    ident: Ident,
+    is_shortcut: bool,
+}
+
+struct ParsedFn<'a> {
+    // Body of the function
+    function: &'a mut ImplItemMethod,
+    // Name of the function.
+    fn_name: Ident,
+    // Name of the first argument of the function, which should be of type `ParseInput`.
+    input_arg: Ident,
+    // List of aliases pointing to this function
+    alias_srcs: Vec<AliasSrc>,
+}
+
+impl Parse for AliasArgs {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let target = input.parse()?;
+        let is_shortcut = if input.peek(Token![,]) {
+            // #[alias(rule, shortcut = true)]
+            let _: Token![,] = input.parse()?;
+            let _: kw::shortcut = input.parse()?;
+            let _: Token![=] = input.parse()?;
+            let b: LitBool = input.parse()?;
+            b.value
+        } else {
+            // #[alias(rule)]
+            false
+        };
+        Ok(AliasArgs {
+            target,
+            is_shortcut,
+        })
+    }
+}
+
+impl Parse for PrecClimbArgs {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let child_rule = input.parse()?;
+        let _: Token![,] = input.parse()?;
+        let climber = input.parse()?;
+        Ok(PrecClimbArgs {
+            child_rule,
+            climber,
+        })
+    }
+}
+
+fn collect_aliases(
+    imp: &mut ItemImpl,
+) -> Result<HashMap<Ident, Vec<AliasSrc>>> {
+    let functions = imp.items.iter_mut().flat_map(|item| match item {
+        ImplItem::Method(m) => Some(m),
+        _ => None,
+    });
+
+    let mut alias_map = HashMap::new();
+    for function in functions {
+        let fn_name = function.sig.ident.clone();
+        let mut alias_attrs = function
+            .attrs
+            .drain_filter(|attr| attr.path.is_ident("alias"))
+            .collect::<Vec<_>>()
+            .into_iter();
+
+        if let Some(attr) = alias_attrs.next() {
+            let args: AliasArgs = attr.parse_args()?;
+            alias_map.entry(args.target).or_insert_with(Vec::new).push(
+                AliasSrc {
+                    ident: fn_name,
+                    is_shortcut: args.is_shortcut,
+                },
+            );
+        }
+        if let Some(attr) = alias_attrs.next() {
+            return Err(Error::new(
+                attr.span(),
+                "expected at most one alias attribute",
+            ));
+        }
+    }
+
+    Ok(alias_map)
+}
+
+fn parse_fn<'a>(
+    function: &'a mut ImplItemMethod,
+    alias_map: &mut HashMap<Ident, Vec<AliasSrc>>,
+) -> Result<ParsedFn<'a>> {
+    let fn_name = function.sig.ident.clone();
+    // Get the name of the first (`input`) function argument
+    let input_arg = function.sig.inputs.first().ok_or_else(|| {
+        Error::new(
+            function.sig.inputs.span(),
+            "a rule function needs an `input` argument",
+        )
+    })?;
+    let input_arg = match &input_arg {
+        FnArg::Receiver(_) => return Err(Error::new(
+            input_arg.span(),
+            "a rule function should not have a `self` argument",
+        )),
+        FnArg::Typed(input_arg) => match &*input_arg.pat{
+            Pat::Ident(ident) => ident.ident.clone(),
+            _ => return Err(Error::new(
+                input_arg.span(),
+                "this argument should be a plain identifier instead of a pattern",
+            )),
+        }
+    };
+
+    let alias_srcs = alias_map.remove(&fn_name).unwrap_or_else(Vec::new);
+
+    Ok(ParsedFn {
+        function,
+        fn_name,
+        input_arg,
+        alias_srcs,
+    })
+}
+
+fn apply_special_attrs(f: &mut ParsedFn, rule_enum: &Ident) -> Result<()> {
+    let function = &mut *f.function;
+    let fn_name = &f.fn_name;
+    let input_arg = &f.input_arg;
+
     *function = parse_quote!(
-        #[allow(non_snake_case, dead_code)]
+        #[allow(non_snake_case)]
         #function
     );
 
-    let recognized_attrs: Vec<_> = function
+    // `prec_climb` attr
+    let prec_climb_attrs: Vec<_> = function
         .attrs
         .drain_filter(|attr| attr.path.is_ident("prec_climb"))
         .collect();
 
-    let name = function.sig.ident.clone();
-
-    if recognized_attrs.is_empty() {
-        // do nothing
-    } else if recognized_attrs.len() > 1 {
+    if prec_climb_attrs.len() > 1 {
         return Err(Error::new(
-            recognized_attrs[1].span(),
-            "expected a single prec_climb attribute",
+            prec_climb_attrs[1].span(),
+            "expected at most one prec_climb attribute",
         ));
+    } else if prec_climb_attrs.is_empty() {
+        // do nothing
     } else {
-        let attr = recognized_attrs.into_iter().next().unwrap();
-        let (child_rule, climber) =
-            attr.parse_args_with(|input: ParseStream| {
-                let child_rule: Ident = input.parse()?;
-                let _: Token![,] = input.parse()?;
-                let climber: Expr = input.parse()?;
-                Ok((child_rule, climber))
-            })?;
-
-        // Get the name of the first (`input`) function argument
-        let first_arg = function.sig.inputs.first().ok_or_else(|| {
-            Error::new(
-                function.sig.inputs.span(),
-                "a prec_climb function needs 4 arguments",
-            )
-        })?;
-        let first_arg = match &first_arg {
-            FnArg::Receiver(_) => return Err(Error::new(
-                first_arg.span(),
-                "a prec_climb function should not have a `self` argument",
-            )),
-            FnArg::Typed(first_arg) =>   match &*first_arg.pat{
-                Pat::Ident(ident) => &ident.ident,
-                _ => return Err(Error::new(
-                    first_arg.span(),
-                    "this argument should be a plain identifier instead of a pattern",
-                )),
-            }
-        };
+        let attr = prec_climb_attrs.into_iter().next().unwrap();
+        let PrecClimbArgs {
+            child_rule,
+            climber,
+        } = attr.parse_args()?;
 
         function.block = parse_quote!({
             #function
 
             #climber.climb(
-                #first_arg.pair.clone().into_inner(),
-                |p| Self::#child_rule(#first_arg.with_pair(p)),
+                #input_arg.pair.clone().into_inner(),
+                |p| Self::#child_rule(#input_arg.with_pair(p)),
                 |l, op, r| {
-                    #name(#first_arg.clone(), l?, op, r?)
+                    #fn_name(#input_arg.clone(), l?, op, r?)
                 },
             )
         });
@@ -81,6 +199,37 @@ fn apply_special_attrs(function: &mut ImplItemMethod) -> Result<()> {
         })?;
     }
 
+    // `alias` attr
+    if !f.alias_srcs.is_empty() {
+        let aliases = f.alias_srcs.iter().map(|src| &src.ident);
+        let block = &function.block;
+        function.block = parse_quote!({
+            let mut #input_arg = #input_arg;
+            // While the current rule allows shortcutting, and there is a single child, and the
+            // child can still be parsed by the current function, then skip to that child.
+            while <Self as PestConsumer>::allows_shortcut(#input_arg.as_rule()) {
+                if let Some(child) = #input_arg.single_child() {
+                    if &<Self as PestConsumer>::rule_alias(child.as_rule())
+                            == stringify!(#fn_name) {
+                        #input_arg = child;
+                        continue;
+                    }
+                }
+                break
+            }
+
+            match #input_arg.as_rule() {
+                #(#rule_enum::#aliases => Self::#aliases(#input_arg),)*
+                #rule_enum::#fn_name => #block,
+                r => unreachable!(
+                    "make_parser: called {} on {:?}",
+                    stringify!(#fn_name),
+                    r
+                )
+            }
+        });
+    }
+
     Ok(())
 }
 
@@ -89,21 +238,95 @@ pub fn make_parser(
     input: proc_macro::TokenStream,
 ) -> Result<proc_macro2::TokenStream> {
     let rule_enum: Ident = syn::parse(attrs)?;
-
     let mut imp: ItemImpl = syn::parse(input)?;
-    imp.items
+
+    let mut alias_map = collect_aliases(&mut imp)?;
+    let rule_alias_branches: Vec<_> = alias_map
+        .iter()
+        .flat_map(|(tgt, srcs)| iter::repeat(tgt).zip(srcs))
+        .map(|(tgt, src)| {
+            let ident = &src.ident;
+            quote!(
+                #rule_enum::#ident => stringify!(#tgt).to_string(),
+            )
+        })
+        .collect();
+    let shortcut_branches: Vec<_> = alias_map
+        .iter()
+        .flat_map(|(_tgt, srcs)| srcs)
+        .map(|AliasSrc { ident, is_shortcut }| {
+            quote!(
+                #rule_enum::#ident => #is_shortcut,
+            )
+        })
+        .collect();
+
+    let fn_map: HashMap<Ident, ParsedFn> = imp
+        .items
         .iter_mut()
-        .map(|item| match item {
-            ImplItem::Method(m) => apply_special_attrs(m),
-            _ => Ok(()),
+        .flat_map(|item| match item {
+            ImplItem::Method(m) => Some(m),
+            _ => None,
+        })
+        .map(|method| {
+            let mut f = parse_fn(method, &mut alias_map)?;
+            apply_special_attrs(&mut f, &rule_enum)?;
+            Ok((f.fn_name.clone(), f))
+        })
+        .collect::<Result<_>>()?;
+
+    // Entries that remain in the alias map don't have a matching method, so we create one.
+    let extra_fns: Vec<_> = alias_map
+        .iter()
+        .map(|(tgt, srcs)| {
+            // Get the signature of one of the functions that has this alias. They should all have
+            // essentially the same signature anyways.
+            let f = fn_map.get(&srcs.first().unwrap().ident).unwrap();
+            let input_arg = f.input_arg.clone();
+            let mut sig = f.function.sig.clone();
+            sig.ident = tgt.clone();
+            let srcs = srcs.iter().map(|src| &src.ident);
+
+            Ok(parse_quote!(
+                #sig {
+                    match #input_arg.as_rule() {
+                        #(#rule_enum::#srcs => Self::#srcs(#input_arg),)*
+                        // We can't match on #rule_enum::#tgt since `tgt` might be an arbitrary
+                        // identifier.
+                        r if &format!("{:?}", r) == stringify!(#tgt) =>
+                            return Err(#input_arg.error(format!(
+                                "make_parser: missing method for rule {}",
+                                stringify!(#tgt),
+                            ))),
+                        r => unreachable!(
+                            "make_parser: called {} on {:?}",
+                            stringify!(#tgt),
+                            r
+                        )
+                    }
+                }
+            ))
         })
-        .collect::<Result<()>>()?;
+        .collect::<Result<_>>()?;
+    imp.items.extend(extra_fns);
 
     let ty = &imp.self_ty;
     let (impl_generics, _, where_clause) = imp.generics.split_for_impl();
     Ok(quote!(
         impl #impl_generics PestConsumer for #ty #where_clause {
-            type RuleEnum = #rule_enum;
+            type Rule = #rule_enum;
+            fn rule_alias(rule: Self::Rule) -> String {
+                match rule {
+                    #(#rule_alias_branches)*
+                    r => format!("{:?}", r),
+                }
+            }
+            fn allows_shortcut(rule: Self::Rule) -> bool {
+                match rule {
+                    #(#shortcut_branches)*
+                    _ => false,
+                }
+            }
         }
 
         #imp