hashtree_core/hashtree/
walk.rs

1use super::*;
2
3impl<S: Store> HashTree<S> {
4    fn decode_node_or_blob(data: &[u8]) -> Result<Option<TreeNode>, HashTreeError> {
5        match decode_tree_node(data) {
6            Ok(node) => Ok(Some(node)),
7            Err(err) if is_tree_node(data) => Err(HashTreeError::Codec(err)),
8            Err(_) => Ok(None),
9        }
10    }
11
12    /// Walk entire tree depth-first (returns Vec)
13    pub async fn walk(&self, cid: &Cid, path: &str) -> Result<Vec<WalkEntry>, HashTreeError> {
14        let mut entries = Vec::new();
15        self.walk_recursive(cid, path, &mut entries).await?;
16        Ok(entries)
17    }
18
19    async fn walk_recursive(
20        &self,
21        cid: &Cid,
22        path: &str,
23        entries: &mut Vec<WalkEntry>,
24    ) -> Result<(), HashTreeError> {
25        let data = match self
26            .store
27            .get(&cid.hash)
28            .await
29            .map_err(|e| HashTreeError::Store(e.to_string()))?
30        {
31            Some(d) => d,
32            None => return Ok(()),
33        };
34
35        // Decrypt if key is present
36        let data = if let Some(key) = &cid.key {
37            decrypt_chk(&data, key).map_err(|e| HashTreeError::Decryption(e.to_string()))?
38        } else {
39            data
40        };
41
42        let node = match Self::decode_node_or_blob(&data)? {
43            Some(node) => node,
44            None => {
45                entries.push(WalkEntry {
46                    path: path.to_string(),
47                    hash: cid.hash,
48                    link_type: LinkType::Blob,
49                    size: data.len() as u64,
50                    key: cid.key,
51                });
52                return Ok(());
53            }
54        };
55
56        let node_size: u64 = node.links.iter().map(|l| l.size).sum();
57        entries.push(WalkEntry {
58            path: path.to_string(),
59            hash: cid.hash,
60            link_type: node.node_type,
61            size: node_size,
62            key: cid.key,
63        });
64
65        for link in &node.links {
66            let child_path = match &link.name {
67                Some(name) => {
68                    if Self::is_internal_directory_link(&node, link) {
69                        let sub_cid = Cid {
70                            hash: link.hash,
71                            key: link.key,
72                        };
73                        Box::pin(self.walk_recursive(&sub_cid, path, entries)).await?;
74                        continue;
75                    }
76                    if path.is_empty() {
77                        name.clone()
78                    } else {
79                        format!("{}/{}", path, name)
80                    }
81                }
82                None => path.to_string(),
83            };
84
85            // Child nodes use their own key from link
86            let child_cid = Cid {
87                hash: link.hash,
88                key: link.key,
89            };
90            Box::pin(self.walk_recursive(&child_cid, &child_path, entries)).await?;
91        }
92
93        Ok(())
94    }
95
96    /// Walk entire tree with parallel fetching
97    /// Uses a work-stealing approach: always keeps `concurrency` requests in flight
98    pub async fn walk_parallel(
99        &self,
100        cid: &Cid,
101        path: &str,
102        concurrency: usize,
103    ) -> Result<Vec<WalkEntry>, HashTreeError> {
104        self.walk_parallel_with_progress(cid, path, concurrency, None)
105            .await
106    }
107
108    /// Walk entire tree with parallel fetching and optional progress counter
109    /// The counter is incremented for each node fetched (not just entries found)
110    ///
111    /// OPTIMIZATION: Blobs are NOT fetched - their metadata (hash, size, link_type)
112    /// comes from the parent node's link, so we just add them directly to entries.
113    /// This avoids downloading file contents during tree traversal.
114    pub async fn walk_parallel_with_progress(
115        &self,
116        cid: &Cid,
117        path: &str,
118        concurrency: usize,
119        progress: Option<&std::sync::atomic::AtomicUsize>,
120    ) -> Result<Vec<WalkEntry>, HashTreeError> {
121        use futures::stream::{FuturesUnordered, StreamExt};
122        use std::collections::VecDeque;
123        use std::sync::atomic::Ordering;
124
125        let mut entries = Vec::new();
126        let mut pending: VecDeque<(Cid, String)> = VecDeque::new();
127        let mut active = FuturesUnordered::new();
128
129        // Seed with root
130        pending.push_back((cid.clone(), path.to_string()));
131
132        loop {
133            // Fill up to concurrency limit from pending queue
134            while active.len() < concurrency {
135                if let Some((node_cid, node_path)) = pending.pop_front() {
136                    let store = &self.store;
137                    let fut = async move {
138                        let data = store
139                            .get(&node_cid.hash)
140                            .await
141                            .map_err(|e| HashTreeError::Store(e.to_string()))?;
142                        Ok::<_, HashTreeError>((node_cid, node_path, data))
143                    };
144                    active.push(fut);
145                } else {
146                    break;
147                }
148            }
149
150            // If nothing active, we're done
151            if active.is_empty() {
152                break;
153            }
154
155            // Wait for any future to complete
156            if let Some(result) = active.next().await {
157                let (node_cid, node_path, data) = result?;
158
159                // Update progress counter
160                if let Some(counter) = progress {
161                    counter.fetch_add(1, Ordering::Relaxed);
162                }
163
164                let data = match data {
165                    Some(d) => d,
166                    None => continue,
167                };
168
169                // Decrypt if key is present
170                let data = if let Some(key) = &node_cid.key {
171                    decrypt_chk(&data, key).map_err(|e| {
172                        HashTreeError::Decryption(format!(
173                            "{} at path '{}' hash {} key {}",
174                            e,
175                            node_path,
176                            hex::encode(node_cid.hash),
177                            hex::encode(key)
178                        ))
179                    })?
180                } else {
181                    data
182                };
183
184                let node = match Self::decode_node_or_blob(&data)? {
185                    Some(node) => node,
186                    None => {
187                        // It's a blob/file - this case only happens for root
188                        entries.push(WalkEntry {
189                            path: node_path,
190                            hash: node_cid.hash,
191                            link_type: LinkType::Blob,
192                            size: data.len() as u64,
193                            key: node_cid.key,
194                        });
195                        continue;
196                    }
197                };
198
199                // It's a directory/file node
200                let node_size: u64 = node.links.iter().map(|l| l.size).sum();
201                entries.push(WalkEntry {
202                    path: node_path.clone(),
203                    hash: node_cid.hash,
204                    link_type: node.node_type,
205                    size: node_size,
206                    key: node_cid.key,
207                });
208
209                // Queue children - but DON'T fetch blobs, just add them directly
210                for link in &node.links {
211                    let child_path = match &link.name {
212                        Some(name) => {
213                            if Self::is_internal_directory_link(&node, link) {
214                                let sub_cid = Cid {
215                                    hash: link.hash,
216                                    key: link.key,
217                                };
218                                pending.push_back((sub_cid, node_path.clone()));
219                                continue;
220                            }
221                            if node_path.is_empty() {
222                                name.clone()
223                            } else {
224                                format!("{}/{}", node_path, name)
225                            }
226                        }
227                        None => node_path.clone(),
228                    };
229
230                    // OPTIMIZATION: If it's a blob, add entry directly without fetching
231                    // The link already contains all the metadata we need
232                    if link.link_type == LinkType::Blob {
233                        entries.push(WalkEntry {
234                            path: child_path,
235                            hash: link.hash,
236                            link_type: LinkType::Blob,
237                            size: link.size,
238                            key: link.key,
239                        });
240                        if let Some(counter) = progress {
241                            counter.fetch_add(1, Ordering::Relaxed);
242                        }
243                        continue;
244                    }
245
246                    // For tree nodes (File/Dir), we need to fetch to see their children
247                    let child_cid = Cid {
248                        hash: link.hash,
249                        key: link.key,
250                    };
251                    pending.push_back((child_cid, child_path));
252                }
253            }
254        }
255
256        Ok(entries)
257    }
258
259    /// Walk tree as stream
260    pub fn walk_stream(
261        &self,
262        cid: Cid,
263        initial_path: String,
264    ) -> Pin<Box<dyn Stream<Item = Result<WalkEntry, HashTreeError>> + Send + '_>> {
265        Box::pin(stream::unfold(
266            WalkStreamState::Init {
267                cid,
268                path: initial_path,
269                tree: self,
270            },
271            |state| async move {
272                match state {
273                    WalkStreamState::Init { cid, path, tree } => {
274                        let data = match tree.store.get(&cid.hash).await {
275                            Ok(Some(d)) => d,
276                            Ok(None) => return None,
277                            Err(e) => {
278                                return Some((
279                                    Err(HashTreeError::Store(e.to_string())),
280                                    WalkStreamState::Done,
281                                ))
282                            }
283                        };
284
285                        // Decrypt if key is present
286                        let data = if let Some(key) = &cid.key {
287                            match decrypt_chk(&data, key) {
288                                Ok(d) => d,
289                                Err(e) => {
290                                    return Some((
291                                        Err(HashTreeError::Decryption(format!(
292                                            "{} at path '{}' hash {} key {}",
293                                            e,
294                                            path,
295                                            hex::encode(cid.hash),
296                                            hex::encode(key)
297                                        ))),
298                                        WalkStreamState::Done,
299                                    ))
300                                }
301                            }
302                        } else {
303                            data
304                        };
305
306                        let node = match Self::decode_node_or_blob(&data) {
307                            Ok(Some(node)) => node,
308                            Ok(None) => {
309                                // Blob data
310                                let entry = WalkEntry {
311                                    path,
312                                    hash: cid.hash,
313                                    link_type: LinkType::Blob,
314                                    size: data.len() as u64,
315                                    key: cid.key,
316                                };
317                                return Some((Ok(entry), WalkStreamState::Done));
318                            }
319                            Err(err) => return Some((Err(err), WalkStreamState::Done)),
320                        };
321
322                        let node_size: u64 = node.links.iter().map(|l| l.size).sum();
323                        let entry = WalkEntry {
324                            path: path.clone(),
325                            hash: cid.hash,
326                            link_type: node.node_type,
327                            size: node_size,
328                            key: cid.key,
329                        };
330
331                        // Create stack with children to process
332                        let mut stack: Vec<WalkStackItem> = Vec::new();
333                        for link in node.links.iter().rev() {
334                            let is_internal = Self::is_internal_directory_link(&node, link);
335                            let child_path = match &link.name {
336                                Some(name) if !is_internal => {
337                                    if path.is_empty() {
338                                        name.clone()
339                                    } else {
340                                        format!("{}/{}", path, name)
341                                    }
342                                }
343                                _ => path.clone(),
344                            };
345                            // Child nodes use their own key from link
346                            stack.push(WalkStackItem {
347                                hash: link.hash,
348                                path: child_path,
349                                key: link.key,
350                            });
351                        }
352
353                        Some((Ok(entry), WalkStreamState::Processing { stack, tree }))
354                    }
355                    WalkStreamState::Processing { mut stack, tree } => {
356                        tree.process_walk_stack(&mut stack).await
357                    }
358                    WalkStreamState::Done => None,
359                }
360            },
361        ))
362    }
363
364    async fn process_walk_stack<'a>(
365        &'a self,
366        stack: &mut Vec<WalkStackItem>,
367    ) -> Option<(Result<WalkEntry, HashTreeError>, WalkStreamState<'a, S>)> {
368        while let Some(item) = stack.pop() {
369            let data = match self.store.get(&item.hash).await {
370                Ok(Some(d)) => d,
371                Ok(None) => continue,
372                Err(e) => {
373                    return Some((
374                        Err(HashTreeError::Store(e.to_string())),
375                        WalkStreamState::Done,
376                    ))
377                }
378            };
379
380            let data = if let Some(key) = &item.key {
381                match decrypt_chk(&data, key) {
382                    Ok(d) => d,
383                    Err(e) => {
384                        return Some((
385                            Err(HashTreeError::Decryption(format!(
386                                "{} at path '{}' hash {} key {}",
387                                e,
388                                item.path,
389                                hex::encode(item.hash),
390                                hex::encode(key)
391                            ))),
392                            WalkStreamState::Done,
393                        ))
394                    }
395                }
396            } else {
397                data
398            };
399
400            let node = match Self::decode_node_or_blob(&data) {
401                Ok(Some(node)) => node,
402                Ok(None) => {
403                    // Blob data
404                    let entry = WalkEntry {
405                        path: item.path,
406                        hash: item.hash,
407                        link_type: LinkType::Blob,
408                        size: data.len() as u64,
409                        key: item.key,
410                    };
411                    return Some((
412                        Ok(entry),
413                        WalkStreamState::Processing {
414                            stack: std::mem::take(stack),
415                            tree: self,
416                        },
417                    ));
418                }
419                Err(err) => return Some((Err(err), WalkStreamState::Done)),
420            };
421
422            let node_size: u64 = node.links.iter().map(|l| l.size).sum();
423            let entry = WalkEntry {
424                path: item.path.clone(),
425                hash: item.hash,
426                link_type: node.node_type,
427                size: node_size,
428                key: item.key,
429            };
430
431            // Push children to stack
432            for link in node.links.iter().rev() {
433                let is_internal = Self::is_internal_directory_link(&node, link);
434                let child_path = match &link.name {
435                    Some(name) if !is_internal => {
436                        if item.path.is_empty() {
437                            name.clone()
438                        } else {
439                            format!("{}/{}", item.path, name)
440                        }
441                    }
442                    _ => item.path.clone(),
443                };
444                stack.push(WalkStackItem {
445                    hash: link.hash,
446                    path: child_path,
447                    key: link.key,
448                });
449            }
450
451            return Some((
452                Ok(entry),
453                WalkStreamState::Processing {
454                    stack: std::mem::take(stack),
455                    tree: self,
456                },
457            ));
458        }
459        None
460    }
461}
462
463struct WalkStackItem {
464    hash: Hash,
465    path: String,
466    key: Option<[u8; 32]>,
467}
468
469enum WalkStreamState<'a, S: Store> {
470    Init {
471        cid: Cid,
472        path: String,
473        tree: &'a HashTree<S>,
474    },
475    Processing {
476        stack: Vec<WalkStackItem>,
477        tree: &'a HashTree<S>,
478    },
479    Done,
480}
hashtree_core/hashtree/walk.rs

hashtree_core/hashtree/
walk.rs
