chore: get_all_paths returns Iterator instead of HashSet

ynqa · Dec 22, 2024 · 584743a · 584743a
1 parent d8253b0
commit 584743a
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 26 deletions.
diff --git a/promkit/src/jsonz.rs b/promkit/src/jsonz.rs
@@ -1,5 +1,3 @@
-use std::collections::HashSet;
-
 use rayon::prelude::*;
 
 pub mod format;
@@ -392,36 +390,49 @@ pub fn create_rows<'a, T: IntoIterator<Item = &'a serde_json::Value>>(iter: T) -
     rows
 }
 
-fn collect_paths(value: &serde_json::Value, current_path: &str, paths: &mut HashSet<String>) {
-    paths.insert(current_path.to_string());
+#[derive(Debug)]
+pub struct PathIterator<'a> {
+    stack: Vec<(String, &'a serde_json::Value)>,
+}
 
-    match value {
-        serde_json::Value::Object(obj) => {
-            for (key, val) in obj {
-                let new_path = if current_path == "." {
-                    format!(".{}", key)
-                } else {
-                    format!("{}.{}", current_path, key)
-                };
-                collect_paths(val, &new_path, paths);
-            }
-        }
-        serde_json::Value::Array(arr) => {
-            for (i, val) in arr.iter().enumerate() {
-                let new_path = format!("{}[{}]", current_path, i);
-                collect_paths(val, &new_path, paths);
+impl<'a> Iterator for PathIterator<'a> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some((current_path, value)) = self.stack.pop() {
+            match value {
+                serde_json::Value::Object(obj) => {
+                    for (key, val) in obj.iter() {
+                        let new_path = if current_path == "." {
+                            format!(".{}", key)
+                        } else {
+                            format!("{}.{}", current_path, key)
+                        };
+                        self.stack.push((new_path, val));
+                    }
+                }
+                serde_json::Value::Array(arr) => {
+                    for (i, val) in arr.iter().enumerate() {
+                        let new_path = format!("{}[{}]", current_path, i);
+                        self.stack.push((new_path, val));
+                    }
+                }
+                _ => {}
             }
+
+            Some(current_path)
+        } else {
+            None
         }
-        _ => {}
     }
 }
 
 pub fn get_all_paths<'a, T: IntoIterator<Item = &'a serde_json::Value>>(
     iter: T,
-) -> HashSet<String> {
-    let mut paths = HashSet::new();
+) -> impl Iterator<Item = String> + 'a {
+    let mut stack = Vec::new();
     for value in iter {
-        collect_paths(value, ".", &mut paths);
+        stack.push((".".to_string(), value));
     }
-    paths
+    PathIterator { stack }
 }
diff --git a/promkit/tests/jsonz_get_all_paths_test.rs b/promkit/tests/jsonz_get_all_paths_test.rs
@@ -53,7 +53,7 @@ mod get_all_paths {
         )
         .unwrap();
 
-        let actual = jsonz::get_all_paths([&v]);
+        let actual = jsonz::get_all_paths([&v]).collect::<HashSet<_>>();
         let expected = HashSet::from_iter(
             [
                 ".",
@@ -114,7 +114,7 @@ mod get_all_paths {
         .filter_map(serde_json::Result::ok)
         .collect::<Vec<_>>();
 
-        let actual = jsonz::get_all_paths(binding.iter());
+        let actual = jsonz::get_all_paths(binding.iter()).collect::<HashSet<_>>();
         let expected = HashSet::from_iter(
             [
                 ".",