yozuk_sdk/
model.rs

1use anyhow::Result;
2use bytes::Bytes;
3use std::{
4    io::{Cursor, Write},
5    ops::Range,
6};
7
8pub struct ModelSet {
9    data: Bytes,
10    keys: Vec<(String, Range<usize>)>,
11    header_len: usize,
12}
13
14impl ModelSet {
15    pub fn new<D, I>(data: D, keys: I) -> Self
16    where
17        D: Into<Bytes>,
18        I: IntoIterator<Item = (String, Range<usize>)>,
19    {
20        Self {
21            data: data.into(),
22            keys: keys.into_iter().collect(),
23            header_len: 0,
24        }
25    }
26
27    pub fn from_data<T: Into<Bytes>>(data: T) -> Result<Self> {
28        let data = data.into();
29
30        let mut cursor = Cursor::new(&data);
31        let keys = bincode::deserialize_from(&mut cursor)?;
32        let header_len = cursor.position() as _;
33        Ok(Self {
34            data,
35            keys,
36            header_len,
37        })
38    }
39
40    pub fn get(&self, key: &str) -> Option<Bytes> {
41        self.get_index(key)
42            .map(|index| self.keys[index].1.clone())
43            .filter(|range| !range.is_empty())
44            .map(|range| {
45                self.data
46                    .slice(range.start + self.header_len..range.end + self.header_len)
47            })
48    }
49
50    pub fn get_index(&self, key: &str) -> Option<usize> {
51        self.keys
52            .binary_search_by(|entry| entry.0.as_str().cmp(key))
53            .ok()
54    }
55
56    pub fn write<W: Write>(&self, mut dst: W) -> bincode::Result<()> {
57        bincode::serialize_into(&mut dst, &self.keys)?;
58        dst.write_all(&self.data)?;
59        Ok(())
60    }
61}