04

Author: caozheng

sorted/src/lib.rs

@@ -1,36 +1,53 @@
-use proc_macro2::TokenStream;
+use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 use syn::{Item, parse_macro_input};
 use syn::spanned::Spanned;
 
-fn check(args: TokenStream, input: Item) -> TokenStream{
+fn check_enum_order(args: &TokenStream, input: &Item) -> syn::Result<()>{
+    let mut idents = vec![];
     match input{
         Item::Enum(ref item) =>{
-            let mut names: Vec<String> = vec![];
             for variant in item.variants.iter() {
-                let cur_name = variant.ident.to_string();
-                for name in names.iter() {
-                    if &cur_name < name{
-                        let msg = format!("{} should sort before {}", cur_name, name);
-                        return syn::Error::new(variant.span(), msg).into_compile_error();
-                    }
-                }
-                names.push(cur_name);
+                idents.push(&variant.ident);
             }
-            return quote! { #input }
+            check_order(idents)
         },
         _=>{
-            return syn::Error::new(args.span(), "expected enum or match expression").into_compile_error();
+            Err(syn::Error::new(args.span(), "expected enum or match expression"))
         }
     }
 }
 
+fn check_order(idents: Vec<&Ident>)-> syn::Result<()> {
+    if idents.is_empty(){
+        return Ok(());
+    }
+    let mut sorted = idents.clone();
+    sorted.sort_by(|ident1, ident2| ident1.cmp(ident2));
+    for (a, b) in idents.iter().zip(sorted.iter()) {
+        if a.to_string() != b.to_string(){
+            let msg = format!("{} should sort before {}", b.to_string(), a.to_string());
+            return Err(syn::Error::new(b.span(), msg));
+        }
+    }
+
+    Ok(())
+}
+
 
 #[proc_macro_attribute]
 pub fn sorted(args: proc_macro::TokenStream, input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let input = parse_macro_input!(input as Item);
     let args: TokenStream = args.into();
-    check(args.into(), input).into()
+    let mut result = TokenStream::new();
+    match check_enum_order(&args, &input) {
+        Ok(_) =>{},
+        Err(err) => {
+            result.extend(err.into_compile_error());
+        }
+    }
+    result.extend(quote! { #input });
+    result.into()
 }
 
 

sorted/tests/progress.rs

@@ -4,7 +4,7 @@ fn tests() {
     t.pass("tests/01-parse-enum.rs");
     t.compile_fail("tests/02-not-enum.rs");
     t.compile_fail("tests/03-out-of-order.rs");
-    //t.compile_fail("tests/04-variants-with-data.rs");
+    t.compile_fail("tests/04-variants-with-data.rs");
     //t.compile_fail("tests/05-match-expr.rs");
     //t.compile_fail("tests/06-pattern-path.rs");
     //t.compile_fail("tests/07-unrecognized-pattern.rs");