refactor: make ErasedSegment::recursive_crawl_all return Iterator…

… instead of `Vec`
quarylabs · Feb 2, 2025 · 955f950 · 955f950
1 parent dbe8152
commit 955f950
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 18 deletions.
diff --git a/crates/lib-core/src/parser/segments/base.rs b/crates/lib-core/src/parser/segments/base.rs
@@ -172,7 +172,6 @@ impl ErasedSegment {
 
     pub fn get_raw_segments(&self) -> Vec<ErasedSegment> {
         self.recursive_crawl_all(false)
-            .into_iter()
             .filter(|it| it.segments().is_empty())
             .collect()
     }
@@ -618,22 +617,26 @@ impl ErasedSegment {
         )
     }
 
-    pub fn recursive_crawl_all(&self, reverse: bool) -> Vec<ErasedSegment> {
-        let mut result = Vec::with_capacity(self.segments().len() + 1);
+    pub fn recursive_crawl_all(&self, reverse: bool) -> impl Iterator<Item = ErasedSegment> + '_ {
+        let self_clone = self.clone();
 
         if reverse {
-            for seg in self.segments().iter().rev() {
-                result.append(&mut seg.recursive_crawl_all(reverse));
-            }
-            result.push(self.clone());
+            Box::new(
+                self.segments()
+                    .iter()
+                    .rev()
+                    .flat_map(move |seg| seg.recursive_crawl_all(reverse))
+                    .chain(std::iter::once(self_clone)),
+            ) as Box<dyn Iterator<Item = _>>
         } else {
-            result.push(self.clone());
-            for seg in self.segments() {
-                result.append(&mut seg.recursive_crawl_all(reverse));
-            }
+            Box::new(
+                std::iter::once(self_clone).chain(
+                    self.segments()
+                        .iter()
+                        .flat_map(move |seg| seg.recursive_crawl_all(reverse)),
+                ),
+            ) as Box<dyn Iterator<Item = _>>
         }
-
-        result
     }
 
     pub fn raw_segments_with_ancestors(&self) -> &[(ErasedSegment, Vec<PathStep>)] {

diff --git a/crates/lib/src/rules/structure/st05.rs b/crates/lib/src/rules/structure/st05.rs
@@ -574,11 +574,7 @@ impl CTEBuilder {
         clone_map: &SegmentCloneMap,
     ) {
         for (idx, cte) in enumerate(&self.ctes) {
-            if cte
-                .recursive_crawl_all(false)
-                .into_iter()
-                .any(|seg| segment.is(&seg))
-            {
+            if cte.recursive_crawl_all(false).any(|seg| segment.is(&seg)) {
                 self.ctes[idx] = clone_map[&self.ctes[idx]].clone();
                 return;
             }