iroh_bytes/format/
collection.rs

1//! The collection type used by iroh
2use std::collections::BTreeMap;
3
4use anyhow::Context;
5use bao_tree::blake3;
6use bytes::Bytes;
7use iroh_io::AsyncSliceReaderExt;
8use serde::{Deserialize, Serialize};
9
10use crate::{
11    get::{fsm, Stats},
12    hashseq::HashSeq,
13    store::MapEntry,
14    util::TempTag,
15    BlobFormat, Hash,
16};
17
18/// A collection of blobs
19///
20/// Note that the format is subject to change.
21#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Default)]
22pub struct Collection {
23    /// Links to the blobs in this collection
24    blobs: Vec<(String, Hash)>,
25}
26
27impl std::ops::Index<usize> for Collection {
28    type Output = (String, Hash);
29
30    fn index(&self, index: usize) -> &Self::Output {
31        &self.blobs[index]
32    }
33}
34
35impl<K, V> Extend<(K, V)> for Collection
36where
37    K: Into<String>,
38    V: Into<Hash>,
39{
40    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
41        self.blobs
42            .extend(iter.into_iter().map(|(k, v)| (k.into(), v.into())));
43    }
44}
45
46impl<K, V> FromIterator<(K, V)> for Collection
47where
48    K: Into<String>,
49    V: Into<Hash>,
50{
51    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
52        let mut res = Self::default();
53        res.extend(iter);
54        res
55    }
56}
57
58impl IntoIterator for Collection {
59    type Item = (String, Hash);
60    type IntoIter = std::vec::IntoIter<Self::Item>;
61
62    fn into_iter(self) -> Self::IntoIter {
63        self.blobs.into_iter()
64    }
65}
66
67/// Metadata for a collection
68///
69/// This is the wire format for the metadata blob.
70#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
71struct CollectionMeta {
72    header: [u8; 13], // Must contain "CollectionV0."
73    names: Vec<String>,
74}
75
76impl Collection {
77    /// The header for the collection format.
78    ///
79    /// This is the start of the metadata blob.
80    pub const HEADER: &'static [u8; 13] = b"CollectionV0.";
81
82    /// Convert the collection to an iterator of blobs, with the last being the
83    /// root blob.
84    ///
85    /// To persist the collection, write all the blobs to storage, and use the
86    /// hash of the last blob as the collection hash.
87    pub fn to_blobs(&self) -> impl Iterator<Item = Bytes> {
88        let meta = CollectionMeta {
89            header: *Self::HEADER,
90            names: self.names(),
91        };
92        let meta_bytes = postcard::to_stdvec(&meta).unwrap();
93        let meta_bytes_hash = blake3::hash(&meta_bytes).into();
94        let links = std::iter::once(meta_bytes_hash)
95            .chain(self.links())
96            .collect::<HashSeq>();
97        let links_bytes = links.into_inner();
98        [meta_bytes.into(), links_bytes].into_iter()
99    }
100
101    /// Read the collection from a get fsm.
102    ///
103    /// Returns the fsm at the start of the first child blob (if any),
104    /// the links array, and the collection.
105    pub async fn read_fsm(
106        fsm_at_start_root: fsm::AtStartRoot,
107    ) -> anyhow::Result<(fsm::EndBlobNext, HashSeq, Collection)> {
108        let (next, links) = {
109            let curr = fsm_at_start_root.next();
110            let (curr, data) = curr.concatenate_into_vec().await?;
111            let links = HashSeq::new(data.into()).context("links could not be parsed")?;
112            (curr.next(), links)
113        };
114        let fsm::EndBlobNext::MoreChildren(at_meta) = next else {
115            anyhow::bail!("expected meta");
116        };
117        let (next, collection) = {
118            let mut children = links.clone();
119            let meta_link = children.pop_front().context("meta link not found")?;
120            let curr = at_meta.next(meta_link);
121            let (curr, names) = curr.concatenate_into_vec().await?;
122            let names = postcard::from_bytes::<CollectionMeta>(&names)?;
123            anyhow::ensure!(
124                names.header == *Self::HEADER,
125                "expected header {:?}, got {:?}",
126                Self::HEADER,
127                names.header
128            );
129            let collection = Collection::from_parts(children, names);
130            (curr.next(), collection)
131        };
132        Ok((next, links, collection))
133    }
134
135    /// Read the collection and all it's children from a get fsm.
136    ///
137    /// Returns the collection, a map from blob offsets to bytes, and the stats.
138    pub async fn read_fsm_all(
139        fsm_at_start_root: crate::get::fsm::AtStartRoot,
140    ) -> anyhow::Result<(Collection, BTreeMap<u64, Bytes>, Stats)> {
141        let (next, links, collection) = Self::read_fsm(fsm_at_start_root).await?;
142        let mut res = BTreeMap::new();
143        let mut curr = next;
144        let end = loop {
145            match curr {
146                fsm::EndBlobNext::MoreChildren(more) => {
147                    let child_offset = more.child_offset();
148                    let Some(hash) = links.get(usize::try_from(child_offset)?) else {
149                        break more.finish();
150                    };
151                    let header = more.next(hash);
152                    let (next, blob) = header.concatenate_into_vec().await?;
153                    res.insert(child_offset - 1, blob.into());
154                    curr = next.next();
155                }
156                fsm::EndBlobNext::Closing(closing) => break closing,
157            }
158        };
159        let stats = end.next().await?;
160        Ok((collection, res, stats))
161    }
162
163    /// Load a collection from a store given a root hash
164    ///
165    /// This assumes that both the links and the metadata of the collection is stored in the store.
166    /// It does not require that all child blobs are stored in the store.
167    pub async fn load<D>(db: &D, root: &Hash) -> anyhow::Result<Self>
168    where
169        D: crate::store::Map,
170    {
171        let links_entry = db.get(root).await?.context("links not found")?;
172        anyhow::ensure!(links_entry.is_complete(), "links not complete");
173        let links_bytes = links_entry.data_reader().await?.read_to_end().await?;
174        let mut links = HashSeq::try_from(links_bytes)?;
175        let meta_hash = links.pop_front().context("meta hash not found")?;
176        let meta_entry = db.get(&meta_hash).await?.context("meta not found")?;
177        anyhow::ensure!(links_entry.is_complete(), "links not complete");
178        let meta_bytes = meta_entry.data_reader().await?.read_to_end().await?;
179        let meta: CollectionMeta = postcard::from_bytes(&meta_bytes)?;
180        anyhow::ensure!(
181            meta.names.len() == links.len(),
182            "names and links length mismatch"
183        );
184        Ok(Self::from_parts(links, meta))
185    }
186
187    /// Store a collection in a store. returns the root hash of the collection
188    /// as a TempTag.
189    pub async fn store<D>(self, db: &D) -> anyhow::Result<TempTag>
190    where
191        D: crate::store::Store,
192    {
193        let (links, meta) = self.into_parts();
194        let meta_bytes = postcard::to_stdvec(&meta)?;
195        let meta_tag = db.import_bytes(meta_bytes.into(), BlobFormat::Raw).await?;
196        let links_bytes = std::iter::once(*meta_tag.hash())
197            .chain(links)
198            .collect::<HashSeq>();
199        let links_tag = db
200            .import_bytes(links_bytes.into(), BlobFormat::HashSeq)
201            .await?;
202        Ok(links_tag)
203    }
204
205    /// Split a collection into a sequence of links and metadata
206    fn into_parts(self) -> (Vec<Hash>, CollectionMeta) {
207        let mut names = Vec::with_capacity(self.blobs.len());
208        let mut links = Vec::with_capacity(self.blobs.len());
209        for (name, hash) in self.blobs {
210            names.push(name);
211            links.push(hash);
212        }
213        let meta = CollectionMeta {
214            header: *Self::HEADER,
215            names,
216        };
217        (links, meta)
218    }
219
220    /// Create a new collection from a list of hashes and metadata
221    fn from_parts(links: impl IntoIterator<Item = Hash>, meta: CollectionMeta) -> Self {
222        meta.names.into_iter().zip(links).collect()
223    }
224
225    /// Get the links to the blobs in this collection
226    fn links(&self) -> impl Iterator<Item = Hash> + '_ {
227        self.blobs.iter().map(|(_name, hash)| *hash)
228    }
229
230    /// Get the names of the blobs in this collection
231    fn names(&self) -> Vec<String> {
232        self.blobs.iter().map(|(name, _)| name.clone()).collect()
233    }
234
235    /// Iterate over the blobs in this collection
236    pub fn iter(&self) -> impl Iterator<Item = &(String, Hash)> {
237        self.blobs.iter()
238    }
239
240    /// Get the number of blobs in this collection
241    pub fn len(&self) -> usize {
242        self.blobs.len()
243    }
244
245    /// Check if this collection is empty
246    pub fn is_empty(&self) -> bool {
247        self.blobs.is_empty()
248    }
249
250    /// Add the given blob to the collection.
251    pub fn push(&mut self, name: String, hash: Hash) {
252        self.blobs.push((name, hash));
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn roundtrip_blob() {
262        let b = (
263            "test".to_string(),
264            blake3::Hash::from_hex(
265                "3aa61c409fd7717c9d9c639202af2fae470c0ef669be7ba2caea5779cb534e9d",
266            )
267            .unwrap()
268            .into(),
269        );
270
271        let mut buf = bytes::BytesMut::zeroed(1024);
272        postcard::to_slice(&b, &mut buf).unwrap();
273        let deserialize_b: (String, Hash) = postcard::from_bytes(&buf).unwrap();
274        assert_eq!(b, deserialize_b);
275    }
276
277    #[test]
278    fn roundtrip_collection_meta() {
279        let expected = CollectionMeta {
280            header: *Collection::HEADER,
281            names: vec!["test".to_string(), "a".to_string(), "b".to_string()],
282        };
283        let mut buf = bytes::BytesMut::zeroed(1024);
284        postcard::to_slice(&expected, &mut buf).unwrap();
285        let actual: CollectionMeta = postcard::from_bytes(&buf).unwrap();
286        assert_eq!(expected, actual);
287    }
288}
iroh_bytes/format/collection.rs

iroh_bytes/format/
collection.rs
