From 0f7616a5a9d28aeb6b23d299a95ac248544ad37b Mon Sep 17 00:00:00 2001
From: Chris Mounce <christophermounce@gmail.com>
Date: Sat, 14 Jun 2025 13:09:52 -0700
Subject: [PATCH] Add ability to walk the capture group tree

---
 src/peg/peg.rs | 49 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/peg/peg.rs b/src/peg/peg.rs
index c7a9317..49a4897 100644
--- a/src/peg/peg.rs
+++ b/src/peg/peg.rs
@@ -10,6 +10,7 @@ pub struct Captures<'a, T: Clone> {
     input: &'a str,
     raw: &'a [RawCapture<T>],
     index: usize,
+    walk: bool,
 }
 
 pub struct Capture<'a, T: Clone> {
@@ -38,8 +39,15 @@ impl<T: Clone> ParseState<T> {
             input: &self.input,
             raw: &self.captures,
             index: 0,
+            walk: false,
         }
     }
+
+    pub fn walk_captures<'a>(&'a self) -> Captures<'a, T> {
+        let mut result = self.captures();
+        result.walk = true;
+        result
+    }
 }
 
 pub mod backend {
@@ -159,7 +167,7 @@ impl<'a, T: Clone> Iterator for Captures<'a, T> {
         let head = self.raw.get(self.index)?;
         let subtree_len = head.subtree_len.unwrap().get();
         let subtree_slice = &self.raw[self.index..self.index + subtree_len];
-        self.index += subtree_len;
+        self.index += if self.walk { 1 } else { subtree_len };
         Some(Capture {
             input: &self.input,
             raw: subtree_slice,
@@ -173,9 +181,16 @@ impl<'a, T: Clone> Capture<'a, T> {
             input: &self.input,
             raw: &self.raw[1..],
             index: 0,
+            walk: false,
         }
     }
 
+    pub fn walk_children(&self) -> Captures<'a, T> {
+        let mut result = self.children();
+        result.walk = true;
+        result
+    }
+
     pub fn kind(&self) -> T {
         self.raw[0].kind.clone()
     }
@@ -223,6 +238,7 @@ mod tests {
         hex_config = "let " var_name " = 0x" ('a'..'f' / '0'..'9')+;
         var_name = "foo" / "bar"; // case must match
 
+        email_text = (!email ANY / email)*;
         email = #Email:(#User:user "@" #Domain:domain);
         user = ('a'..'z'i)+;
         domain = user+ ("." user)+;
@@ -358,4 +374,35 @@ mod tests {
         ]
         "#);
     }
+
+    #[test]
+    fn test_walk_captures() {
+        let mut p = ParseState::new("Contact alice@foo.com or bob@bar.net.");
+        assert!(email_text(&mut p));
+        let extract_by_tag = |tag| {
+            p.walk_captures()
+                .filter_map(|c| {
+                    if c.kind() == tag {
+                        Some(c.text())
+                    } else {
+                        None
+                    }
+                })
+                .collect::<Vec<_>>()
+        };
+        let users = extract_by_tag(Tag::User);
+        let domains = extract_by_tag(Tag::Domain);
+        assert_debug_snapshot!(users, @r#"
+        [
+            "alice",
+            "bob",
+        ]
+        "#);
+        assert_debug_snapshot!(domains, @r#"
+        [
+            "foo.com",
+            "bar.net",
+        ]
+        "#);
+    }
 }