iota_storage/
blob.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    io::{Read, Write},
7    marker::PhantomData,
8};
9
10use anyhow::{Result, anyhow};
11use byteorder::ReadBytesExt;
12use integer_encoding::{VarInt, VarIntReader};
13use num_enum::{IntoPrimitive, TryFromPrimitive};
14use serde::{Serialize, de::DeserializeOwned};
15
16pub const MAX_VARINT_LENGTH: usize = 10;
17pub const BLOB_ENCODING_BYTES: usize = 1;
18
19#[derive(Copy, Clone, Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
20#[repr(u8)]
21pub enum BlobEncoding {
22    Bcs = 1,
23}
24
25pub struct Blob {
26    pub data: Vec<u8>,
27    pub encoding: BlobEncoding,
28}
29
30impl Blob {
31    pub fn encode<T: Serialize>(value: &T, encoding: BlobEncoding) -> Result<Self> {
32        let value_buf = bcs::to_bytes(value)?;
33        let (data, encoding) = match encoding {
34            BlobEncoding::Bcs => (value_buf, encoding),
35        };
36        Ok(Blob { data, encoding })
37    }
38    pub fn decode<T: DeserializeOwned>(self) -> Result<T> {
39        let data = match &self.encoding {
40            BlobEncoding::Bcs => self.data,
41        };
42        let res = bcs::from_bytes(&data)?;
43        Ok(res)
44    }
45    pub fn read<R: Read>(rbuf: &mut R) -> Result<Blob> {
46        let len = rbuf.read_varint::<u64>()? as usize;
47        if len == 0 {
48            return Err(anyhow!("Invalid object length of 0 in file"));
49        }
50        let encoding = rbuf.read_u8()?;
51        let mut data = vec![0u8; len];
52        rbuf.read_exact(&mut data)?;
53        let blob = Blob {
54            data,
55            encoding: BlobEncoding::try_from(encoding)?,
56        };
57        Ok(blob)
58    }
59    pub fn write<W: Write>(&self, wbuf: &mut W) -> Result<usize> {
60        let mut buf = [0u8; MAX_VARINT_LENGTH];
61        let mut counter = 0;
62        let n = (self.data.len() as u64).encode_var(&mut buf);
63        wbuf.write_all(&buf[0..n])?;
64        counter += n;
65        buf[0] = self.encoding.into();
66        wbuf.write_all(&buf[0..BLOB_ENCODING_BYTES])?;
67        counter += 1;
68        wbuf.write_all(&self.data)?;
69        counter += self.data.len();
70        Ok(counter)
71    }
72    pub fn size(&self) -> usize {
73        let mut blob_size = self.data.len().required_space();
74        blob_size += BLOB_ENCODING_BYTES;
75        blob_size += self.data.len();
76        blob_size
77    }
78    pub fn to_bytes(&self) -> Vec<u8> {
79        [vec![self.encoding.into()], self.data.clone()].concat()
80    }
81    pub fn from_bytes<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
82        let (encoding, data) = bytes.split_first().ok_or(anyhow!("empty bytes"))?;
83        Blob {
84            data: data.to_vec(),
85            encoding: BlobEncoding::try_from(*encoding)?,
86        }
87        .decode()
88    }
89}
90
91/// An iterator over blobs in a blob file.
92pub struct BlobIter<T> {
93    reader: Box<dyn Read>,
94    _phantom: PhantomData<T>,
95}
96
97impl<T: DeserializeOwned> BlobIter<T> {
98    pub fn new(reader: Box<dyn Read>) -> Self {
99        Self {
100            reader,
101            _phantom: PhantomData,
102        }
103    }
104    fn next_blob(&mut self) -> Result<T> {
105        let blob = Blob::read(&mut self.reader)?;
106        blob.decode()
107    }
108}
109
110impl<T: DeserializeOwned> Iterator for BlobIter<T> {
111    type Item = T;
112    fn next(&mut self) -> Option<Self::Item> {
113        self.next_blob().ok()
114    }
115}