From d55aa763d06410a9aea738f7b8d80e0c2d8317d2 Mon Sep 17 00:00:00 2001
From: Tomas Tauber <2410580+tomtau@users.noreply.github.com>
Date: Fri, 1 Mar 2024 14:15:32 +0800
Subject: [PATCH] fix: adjust tag association in optional and repeat
 expressions

fixes #984, fixes #982
---
 derive/tests/implicit.rs   | 10 ++++++
 derive/tests/opt.pest      |  8 +++++
 derive/tests/opt.rs        | 29 ++++++++++++++++
 generator/src/generator.rs | 70 ++++++++++++++++++++++++++++++++------
 4 files changed, 107 insertions(+), 10 deletions(-)
 create mode 100644 derive/tests/opt.pest
 create mode 100644 derive/tests/opt.rs

diff --git a/derive/tests/implicit.rs b/derive/tests/implicit.rs
index 48c88021..0ca7a43d 100644
--- a/derive/tests/implicit.rs
+++ b/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);
+}
diff --git a/derive/tests/opt.pest b/derive/tests/opt.pest
new file mode 100644
index 00000000..b5ff4661
--- /dev/null
+++ b/derive/tests/opt.pest
@@ -0,0 +1,8 @@
+expr = {
+  SOI ~
+  #prefix=(STAR)? ~ #suffix=DOT?
+  ~ EOI
+}
+
+STAR={"*"}
+DOT={"."}
diff --git a/derive/tests/opt.rs b/derive/tests/opt.rs
new file mode 100644
index 00000000..8ce316c8
--- /dev/null
+++ b/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());
+}
diff --git a/generator/src/generator.rs b/generator/src/generator.rs
index f5fa5a1f..042b2ff9 100644
--- a/generator/src/generator.rs
+++ b/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))
+                }
+            }
+        },
     }
 }