2 files changed, 39 insertions, 24 deletions

diff --git a/graph_macro/src/lib.rs b/graph_macro/src/lib.rs
index 0f56f57..3b7742a 100644
--- a/graph_macro/src/lib.rs
+++ b/graph_macro/src/lib.rs
@@ -84,20 +84,20 @@ impl CompileError {
         // I cannot use the parse method of the type TokenStream, as
         // that will not set the spans properly.
 
-        let compile_error_ident = TokenTree::Ident(Ident::new("compile_error", self.span.clone()));
+        let compile_error_ident = TokenTree::Ident(Ident::new("compile_error", self.span));
         let mut exclamation_punct = TokenTree::Punct(Punct::new('!', Spacing::Alone));
 
-        exclamation_punct.set_span(self.span.clone());
+        exclamation_punct.set_span(self.span);
 
         let mut arg_mes_literal = TokenTree::Literal(Literal::string(&self.mes));
 
-        arg_mes_literal.set_span(self.span.clone());
+        arg_mes_literal.set_span(self.span);
 
         let arg_mes_stream = [arg_mes_literal].into_iter().collect();
 
         let mut arg_group = TokenTree::Group(Group::new(Delimiter::Parenthesis, arg_mes_stream));
 
-        arg_group.set_span(self.span.clone());
+        arg_group.set_span(self.span);
 
         let mut semi_colon_punct = TokenTree::Punct(Punct::new(';', Spacing::Alone));
 
@@ -158,7 +158,7 @@ pub fn graph_derive(input: TokenStream) -> TokenStream {
         let mut result = String::new();
 
         if !generics.is_empty() {
-            result.push_str("<");
+            result.push('<');
         }
 
         for generic in generics.iter() {
@@ -166,7 +166,7 @@ pub fn graph_derive(input: TokenStream) -> TokenStream {
         }
 
         if !generics.is_empty() {
-            result.push_str(">");
+            result.push('>');
         }
 
         result
@@ -176,7 +176,7 @@ pub fn graph_derive(input: TokenStream) -> TokenStream {
         let mut result = String::new();
 
         if !generics.is_empty() {
-            result.push_str("<");
+            result.push('<');
         }
 
         for generic in generics.iter() {
@@ -184,7 +184,7 @@ pub fn graph_derive(input: TokenStream) -> TokenStream {
         }
 
         if !generics.is_empty() {
-            result.push_str(">");
+            result.push('>');
         }
 
         result
@@ -236,6 +236,11 @@ self.{field_name}.has_edge(nodea, nodeb)
 fn replace_by_builder(&mut self, _builder: impl graph::builder::Builder<Result = Self>) {{
 unimplemented!()
 }}
+
+#[inline]
+fn print_viz(&self, filename: &str) -> Result<(), std::io::Error> {{
+self.{field_name}.print_viz(filename)
+}}
 }}"
     );
 
@@ -386,6 +391,10 @@ fn get_graph_field_name_type(
 
             let mut cloned_body = body.clone();
 
+            while move_attributes(&mut cloned_body).is_ok() {}
+
+            let _ = get_visibility_mod(&mut cloned_body)?;
+
             let mut result_name = get_first_ident(&mut cloned_body, g.span())?.to_string();
 
             move_punct(&mut cloned_body, ':')?;
@@ -418,12 +427,19 @@ fn get_graph_field_name_type(
                                     break;
                                 }
 
-                                let _ = move_punct(&mut attr_stream, ',');
+                                get_until(
+                                    &mut attr_stream,
+                                    |tree| matches!(tree, TokenTree::Punct(p) if p.as_char() == ','),
+                                    true,
+                                    true,
+                                )?;
                             }
 
                             body.next();
 
                             if found_attribute {
+                                while move_attributes(&mut body).is_ok() {}
+
                                 get_visibility_mod(&mut body)?;
 
                                 result_name = get_ident(&mut body)?.to_string();
@@ -557,12 +573,9 @@ fn get_first_ident(
     input: &mut Peekable<impl Iterator<Item = TokenTree>>,
     span: Span,
 ) -> Result<Ident, CompileError> {
-    while let Some(tree) = input.next() {
-        match tree {
-            TokenTree::Ident(id) => {
-                return Ok(id);
-            }
-            _ => {}
+    for tree in input {
+        if let TokenTree::Ident(id) = tree {
+            return Ok(id);
         }
     }
 
@@ -580,7 +593,7 @@ fn get_until(
     let mut found_predicate = false;
 
     if consume_boundary {
-        while let Some(tree) = input.next() {
+        for tree in input.by_ref() {
             if predicate(&tree) {
                 found_predicate = true;
                 break;
@@ -647,16 +660,14 @@ fn move_attributes(
             match input.peek() {
                 Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
                     input.next();
+
+                    Ok(())
                 }
                 Some(_) => {
-                    return myerror!("expected a group in square brackets", input);
-                }
-                _ => {
-                    return end_of_input;
+                    myerror!("expected a group in square brackets", input)
                 }
+                _ => end_of_input,
             }
-
-            Ok(())
         }
         Some(_) => myerror!(error_mes, input),
         _ => end_of_input,
diff --git a/graph_macro/tests/works.rs b/graph_macro/tests/works.rs
index a57e866..dc0f75a 100644
--- a/graph_macro/tests/works.rs
+++ b/graph_macro/tests/works.rs
@@ -11,13 +11,17 @@ pub(crate) struct Haha {
 
 #[derive(Debug)]
 /// Testing docs
-#[allow(unused)]
+#[repr(C)]
 #[derive(Default, Graph)]
+#[allow(unused)]
 pub struct HahaU(pub(crate) graph::ALGraph);
 
 #[derive(Debug, Graph)]
 pub struct HahaG<T: graph::GraphLabel> {
-    graph: graph::DLGraph<T>,
+    #[graph(haha)]
+    #[allow(unused)]
+    #[inline]
+    pub(crate) graph: graph::DLGraph<T>,
 }
 
 impl<T: graph::GraphLabel> Default for HahaG<T> {