Implement rule aliasing and simplify parser code

2 files changed, 156 insertions, 42 deletions

diff --git a/dhall_proc_macros/src/make_parser.rs b/dhall_proc_macros/src/make_parser.rs
index 268a639..3375c39 100644
--- a/dhall_proc_macros/src/make_parser.rs
+++ b/dhall_proc_macros/src/make_parser.rs
@@ -1,33 +1,101 @@
+use std::collections::HashMap;
+use std::iter;
+
 use quote::quote;
 use syn::parse::{ParseStream, Result};
 use syn::spanned::Spanned;
 use syn::{
     parse_quote, Error, Expr, FnArg, Ident, ImplItem, ImplItemMethod, ItemImpl,
-    Pat, Token,
+    Pat, Signature, Token,
 };
 
-fn apply_special_attrs(function: &mut ImplItemMethod) -> Result<()> {
+fn collect_aliases(
+    imp: &mut ItemImpl,
+) -> Result<HashMap<Ident, (Signature, Vec<Ident>)>> {
+    let mut alias_map = HashMap::new();
+
+    for item in &mut imp.items {
+        if let ImplItem::Method(function) = item {
+            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 tgt: Ident = attr.parse_args()?;
+                alias_map
+                    .entry(tgt)
+                    .or_insert_with(|| (function.sig.clone(), Vec::new()))
+                    .1
+                    .push(fn_name);
+            }
+            if let Some(attr) = alias_attrs.next() {
+                return Err(Error::new(
+                    attr.span(),
+                    "expected at most one alias attribute",
+                ));
+            }
+        }
+    }
+
+    Ok(alias_map)
+}
+
+fn parse_rulefn_sig(sig: &Signature) -> Result<(Ident, Ident)> {
+    let fn_name = sig.ident.clone();
+    // Get the name of the first (`input`) function argument
+    let input_arg = sig.inputs.first().ok_or_else(|| {
+        Error::new(
+            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",
+            )),
+        }
+    };
+
+    Ok((fn_name, input_arg))
+}
+
+fn apply_special_attrs(
+    function: &mut ImplItemMethod,
+    alias_map: &mut HashMap<Ident, (Signature, Vec<Ident>)>,
+    rule_enum: &Ident,
+) -> Result<()> {
     *function = parse_quote!(
         #[allow(non_snake_case, dead_code)]
         #function
     );
 
-    let recognized_attrs: Vec<_> = function
+    let (fn_name, input_arg) = parse_rulefn_sig(&function.sig)?;
+
+    // `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 attr = prec_climb_attrs.into_iter().next().unwrap();
         let (child_rule, climber) =
             attr.parse_args_with(|input: ParseStream| {
                 let child_rule: Ident = input.parse()?;
@@ -36,35 +104,14 @@ fn apply_special_attrs(function: &mut ImplItemMethod) -> Result<()> {
                 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",
-                )),
-            }
-        };
-
         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 +128,22 @@ fn apply_special_attrs(function: &mut ImplItemMethod) -> Result<()> {
         })?;
     }
 
+    // `alias` attr
+    if let Some((_, aliases)) = alias_map.remove(&fn_name) {
+        let block = &function.block;
+        function.block = parse_quote!({
+            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 +152,69 @@ 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)?;
+
+    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)| {
+            quote!(
+                #rule_enum::#src => stringify!(#tgt).to_string(),
+            )
+        })
+        .collect();
+
     imp.items
         .iter_mut()
         .map(|item| match item {
-            ImplItem::Method(m) => apply_special_attrs(m),
+            ImplItem::Method(m) => {
+                apply_special_attrs(m, &mut alias_map, &rule_enum)
+            }
             _ => Ok(()),
         })
         .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, (sig, srcs))| {
+            let mut sig = sig.clone();
+            sig.ident = tgt.clone();
+
+            let (_, input_arg) = parse_rulefn_sig(&sig)?;
+            Ok(ImplItem::Method(parse_quote!(
+                #sig {
+                    match #input_arg.as_rule() {
+                        #(#rule_enum::#srcs => Self::#srcs(#input_arg),)*
+                        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<_>>()?;
+    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),
+                }
+            }
         }
 
         #imp
diff --git a/dhall_proc_macros/src/parse_children.rs b/dhall_proc_macros/src/parse_children.rs
index b1d43fc..a35c03f 100644
--- a/dhall_proc_macros/src/parse_children.rs
+++ b/dhall_proc_macros/src/parse_children.rs
@@ -88,9 +88,7 @@ fn make_parser_branch(
     let i_variable_pattern =
         Ident::new("___variable_pattern", Span::call_site());
     let match_pat = branch.pattern.iter().map(|item| match item {
-        Single { rule_name, .. } => {
-            quote!(<<Self as PestConsumer>::RuleEnum>::#rule_name)
-        }
+        Single { rule_name, .. } => quote!(stringify!(#rule_name)),
         Multiple { .. } => quote!(#i_variable_pattern @ ..),
     });
     let match_filter = branch.pattern.iter().map(|item| match item {
@@ -101,7 +99,7 @@ fn make_parser_branch(
                 // https://github.com/rust-lang/rust/issues/59803.
                 let all_match = |slice: &[_]| {
                     slice.iter().all(|r|
-                        r == &<<Self as PestConsumer>::RuleEnum>::#rule_name
+                        *r == stringify!(#rule_name)
                     )
                 };
                 all_match(#i_variable_pattern)
@@ -192,6 +190,11 @@ pub fn parse_children(
             .clone()
             .into_inner()
             .map(|p| p.as_rule())
+            .map(<Self as PestConsumer>::rule_alias)
+            .collect();
+        let #i_children_rules: Vec<&str> = #i_children_rules
+            .iter()
+            .map(String::as_str)
             .collect();
 
         #[allow(unused_mut)]