fix: adjust tag association in optional and repeat expressions

derive/tests/implicit.rs

@@ -35,3 +35,13 @@ fn test_implicit_whitespace() {
     assert_eq!(pairs.clone().find_tagged("one_comp").count(), 2);
     assert_eq!(pairs.find_tagged("one_array").count(), 2);
 }
+
+#[test]
+#[cfg(feature = "grammar-extras")]
+fn test_implicit_whitespace_multitag() {
+    let successful_parse = TestImplicitParser::parse(Rule::program, "a a a");
+    assert!(successful_parse.is_ok());
+    dbg!(&successful_parse);
+    let pairs = successful_parse.unwrap();
+    assert_eq!(pairs.clone().find_tagged("tail").count(), 2);
+}

derive/tests/opt.pest

@@ -0,0 +1,8 @@
+expr = {
+  SOI ~
+  #prefix=(STAR)? ~ #suffix=DOT?
+  ~ EOI
+}
+
+STAR={"*"}
+DOT={"."}

derive/tests/opt.rs

@@ -0,0 +1,29 @@
+// Licensed under the Apache License, Version 2.0
+// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
+// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. All files in the project carrying such notice may not be copied,
+// modified, or distributed except according to those terms.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+extern crate alloc;
+extern crate pest;
+extern crate pest_derive;
+
+#[cfg(feature = "grammar-extras")]
+use pest::Parser;
+use pest_derive::Parser;
+
+#[derive(Parser)]
+#[grammar = "../tests/opt.pest"]
+struct TestOptParser;
+
+#[test]
+#[cfg(feature = "grammar-extras")]
+fn test_opt_tag() {
+    let successful_parse = TestOptParser::parse(Rule::expr, "*");
+    assert!(successful_parse.is_ok());
+    dbg!(&successful_parse);
+    let pairs = successful_parse.unwrap();
+    assert!(pairs.find_first_tagged("prefix").is_some());
+    assert!(pairs.find_first_tagged("suffix").is_none());
+}

generator/src/generator.rs

@@ -571,12 +571,42 @@ fn generate_expr(expr: OptimizedExpr) -> TokenStream {
             }
         }
         #[cfg(feature = "grammar-extras")]
-        OptimizedExpr::NodeTag(expr, tag) => {
-            let expr = generate_expr(*expr);
-            quote! {
-                #expr.and_then(|state| state.tag_node(#tag))
+        OptimizedExpr::NodeTag(expr, tag) => match *expr {
+            OptimizedExpr::Opt(expr) => {
+                let expr = generate_expr(*expr);
+                quote! {
+                    state.optional(|state| {
+                        #expr.and_then(|state| state.tag_node(#tag))
+                    })
+                }
             }
-        }
+            OptimizedExpr::Rep(expr) => {
+                let expr = generate_expr(*expr);
+                quote! {
+                    state.sequence(|state| {
+                        state.optional(|state| {
+                            #expr.and_then(|state| {
+                                state.repeat(|state| {
+                                    state.sequence(|state| {
+                                        super::hidden::skip(
+                                            state
+                                        ).and_then(|state| {
+                                            #expr.and_then(|state| state.tag_node(#tag))
+                                        })
+                                    })
+                                })
+                            }).and_then(|state| state.tag_node(#tag))
+                        })
+                    })
+                }
+            }
+            expr => {
+                let expr = generate_expr(expr);
+                quote! {
+                    #expr.and_then(|state| state.tag_node(#tag))
+                }
+            }
+        },
     }
 }
 
@@ -726,12 +756,32 @@ fn generate_expr_atomic(expr: OptimizedExpr) -> TokenStream {
             }
         }
         #[cfg(feature = "grammar-extras")]
-        OptimizedExpr::NodeTag(expr, tag) => {
-            let expr = generate_expr_atomic(*expr);
-            quote! {
-                #expr.and_then(|state| state.tag_node(#tag))
+        OptimizedExpr::NodeTag(expr, tag) => match *expr {
+            OptimizedExpr::Opt(expr) => {
+                let expr = generate_expr_atomic(*expr);
+
+                quote! {
+                    state.optional(|state| {
+                        #expr.and_then(|state| state.tag_node(#tag))
+                    })
+                }
             }
-        }
+            OptimizedExpr::Rep(expr) => {
+                let expr = generate_expr_atomic(*expr);
+
+                quote! {
+                    state.repeat(|state| {
+                        #expr.and_then(|state| state.tag_node(#tag))
+                    })
+                }
+            }
+            expr => {
+                let expr = generate_expr_atomic(expr);
+                quote! {
+                    #expr.and_then(|state| state.tag_node(#tag))
+                }
+            }
+        },
     }
 }