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}