Skip to main content

hashtree_core/hashtree/
stream.rs

1use std::pin::Pin;
2
3use futures::stream::{self, Stream};
4
5use crate::codec::{decode_tree_node, is_tree_node};
6use crate::store::Store;
7use crate::types::{Cid, Hash, Link};
8
9use super::{HashTree, HashTreeError};
10
11impl<S: Store> HashTree<S> {
12    /// Read content as a stream of chunks by Cid (handles decryption automatically)
13    ///
14    /// Returns an async stream that yields chunks as they are read.
15    /// Useful for large files or when you want to process data incrementally.
16    pub fn get_stream(
17        &self,
18        cid: &Cid,
19    ) -> Pin<Box<dyn Stream<Item = Result<Vec<u8>, HashTreeError>> + Send + '_>> {
20        Box::pin(stream::unfold(
21            ReadStreamState::Init {
22                cid: cid.clone(),
23                tree: self,
24            },
25            |state| async move {
26                match state {
27                    ReadStreamState::Init { cid, tree } => {
28                        let decrypted = match tree.get_cid_root_bytes(&cid).await {
29                            Ok(Some(data)) => data,
30                            Ok(None) => return None,
31                            Err(e) => return Some((Err(e), ReadStreamState::Done)),
32                        };
33
34                        if !is_tree_node(&decrypted) {
35                            // Single blob - yield decrypted data
36                            return Some((Ok(decrypted), ReadStreamState::Done));
37                        }
38
39                        // Tree node - parse and traverse
40                        let node = match decode_tree_node(&decrypted) {
41                            Ok(n) => n,
42                            Err(e) => {
43                                return Some((Err(HashTreeError::Codec(e)), ReadStreamState::Done))
44                            }
45                        };
46
47                        let mut stack = node.links.into_iter().rev().collect();
48
49                        tree.process_stream_stack(&mut stack).await
50                    }
51                    ReadStreamState::Processing { mut stack, tree } => {
52                        tree.process_stream_stack(&mut stack).await
53                    }
54                    ReadStreamState::Done => None,
55                }
56            },
57        ))
58    }
59
60    async fn process_stream_stack<'a>(
61        &'a self,
62        stack: &mut Vec<Link>,
63    ) -> Option<(Result<Vec<u8>, HashTreeError>, ReadStreamState<'a, S>)> {
64        while let Some(item) = stack.pop() {
65            let decrypted = match self.read_link_bytes(&item).await {
66                Ok(bytes) => bytes,
67                Err(e) => return Some((Err(e), ReadStreamState::Done)),
68            };
69            match Self::decode_linked_file_node(&item, &decrypted) {
70                Ok(Some(node)) => stack.extend(node.links.into_iter().rev()),
71                Ok(None) => {
72                    return Some((
73                        Ok(decrypted),
74                        ReadStreamState::Processing {
75                            stack: std::mem::take(stack),
76                            tree: self,
77                        },
78                    ))
79                }
80                Err(e) => return Some((Err(e), ReadStreamState::Done)),
81            }
82        }
83        None
84    }
85
86    /// Read a file as stream of chunks
87    /// Returns an async stream that yields chunks as they are read
88    pub fn read_file_stream(
89        &self,
90        hash: Hash,
91    ) -> Pin<Box<dyn Stream<Item = Result<Vec<u8>, HashTreeError>> + Send + '_>> {
92        self.get_stream(&Cid::public(hash))
93    }
94}
95
96// Internal state types for streaming
97
98enum ReadStreamState<'a, S: Store> {
99    Init {
100        cid: Cid,
101        tree: &'a HashTree<S>,
102    },
103    Processing {
104        stack: Vec<Link>,
105        tree: &'a HashTree<S>,
106    },
107    Done,
108}