iris_chat_protocol/storage/
checkpoint.rs

1use super::{StorageError, StorageResult};
2use nostr::hashes::{sha256, Hash};
3use std::collections::HashMap;
4use std::sync::Arc;
5
6/// Immutable local checkpoint bytes. Its content identity is computed once,
7/// never published, and relies on SHA-256 collision resistance. Keeping the
8/// identity and bytes private prevents a caller from reusing an ID for new data.
9pub struct StorageCheckpointBody {
10    pub(super) id: [u8; 32],
11    json: Arc<serde_json::value::RawValue>,
12}
13
14impl StorageCheckpointBody {
15    pub fn new(json: Arc<serde_json::value::RawValue>) -> Self {
16        Self {
17            id: sha256::Hash::hash(json.get().as_bytes()).to_byte_array(),
18            json,
19        }
20    }
21    pub fn json(&self) -> &str {
22        self.json.get()
23    }
24}
25
26impl serde::Serialize for StorageCheckpointBody {
27    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
28        serde::Serialize::serialize(self.json.as_ref(), serializer)
29    }
30}
31
32pub(super) enum PartBytes {
33    Inline(String),
34    Immutable(Arc<StorageCheckpointBody>),
35}
36
37impl PartBytes {
38    pub(super) fn bytes(&self) -> &[u8] {
39        match self {
40            Self::Inline(value) => value.as_bytes(),
41            Self::Immutable(body) => body.json().as_bytes(),
42        }
43    }
44}
45
46pub(super) struct CheckpointPart {
47    pub key: Vec<u8>,
48    pub payload_key: Vec<u8>,
49    pub bytes: PartBytes,
50    pub prefix: String,
51    pub suffix: String,
52    pub padding: usize,
53}
54
55impl CheckpointPart {
56    pub(super) fn len(&self) -> usize {
57        self.prefix
58            .len()
59            .saturating_add(self.bytes.bytes().len())
60            .saturating_add(self.suffix.len())
61            .saturating_add(self.padding)
62    }
63}
64
65/// A lossless description of a logical String. SQLite can store its immutable
66/// records independently; other adapters materialize the same bytes normally.
67pub struct StorageCheckpoint {
68    pub(super) parts: Vec<CheckpointPart>,
69    occurrences: HashMap<(u8, [u8; 32]), u64>,
70    bytes: usize,
71}
72
73impl Default for StorageCheckpoint {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl StorageCheckpoint {
80    pub fn new() -> Self {
81        Self {
82            parts: Vec::new(),
83            occurrences: HashMap::new(),
84            bytes: 0,
85        }
86    }
87    pub fn len(&self) -> usize {
88        self.bytes
89    }
90    pub fn is_empty(&self) -> bool {
91        self.bytes == 0
92    }
93
94    pub fn push_inline(&mut self, name: &str, text: String, padding: usize) {
95        let mut key = vec![0];
96        key.extend_from_slice(name.as_bytes());
97        let part = CheckpointPart {
98            payload_key: key.clone(),
99            key,
100            bytes: PartBytes::Inline(text),
101            prefix: String::new(),
102            suffix: String::new(),
103            padding,
104        };
105        self.bytes = self.bytes.saturating_add(part.len());
106        self.parts.push(part);
107    }
108
109    pub fn push_record(
110        &mut self,
111        collection: u8,
112        body: Arc<StorageCheckpointBody>,
113        prefix: String,
114        suffix: String,
115        padding: usize,
116    ) {
117        let occurrence = self.occurrences.entry((collection, body.id)).or_default();
118        let mut key = vec![1, collection];
119        key.extend_from_slice(&body.id);
120        key.extend_from_slice(&occurrence.to_be_bytes());
121        *occurrence += 1;
122        let mut payload_key = vec![1];
123        payload_key.extend_from_slice(&body.id);
124        let part = CheckpointPart {
125            key,
126            payload_key,
127            bytes: PartBytes::Immutable(body),
128            prefix,
129            suffix,
130            padding,
131        };
132        self.bytes = self.bytes.saturating_add(part.len());
133        self.parts.push(part);
134    }
135
136    pub fn pad_records(&mut self, page: usize, minimum_record: usize) {
137        let page = page.max(1);
138        for part in &mut self.parts {
139            if matches!(&part.bytes, PartBytes::Immutable(_)) && part.len() >= minimum_record {
140                let padding = part
141                    .len()
142                    .div_ceil(page)
143                    .saturating_mul(page)
144                    .saturating_sub(part.len());
145                part.padding = part.padding.saturating_add(padding);
146                self.bytes = self.bytes.saturating_add(padding);
147            }
148        }
149    }
150
151    pub fn materialize(&self) -> StorageResult<String> {
152        self.validate()?;
153        let mut result = String::new();
154        result
155            .try_reserve_exact(self.bytes)
156            .map_err(|error| StorageError::new(error.to_string()))?;
157        for part in &self.parts {
158            result.push_str(&part.prefix);
159            // Both variants originate in String or validated RawValue.
160            result.push_str(
161                std::str::from_utf8(part.bytes.bytes())
162                    .map_err(|_| StorageError::new("invalid checkpoint UTF-8"))?,
163            );
164            result.push_str(&part.suffix);
165            append_padding(&mut result, part.padding);
166        }
167        Ok(result)
168    }
169
170    pub(super) fn validate(&self) -> StorageResult<()> {
171        let total = self.parts.iter().try_fold(0usize, |total, part| {
172            part.prefix
173                .len()
174                .checked_add(part.bytes.bytes().len())
175                .and_then(|n| n.checked_add(part.suffix.len()))
176                .and_then(|n| n.checked_add(part.padding))
177                .and_then(|n| total.checked_add(n))
178        });
179        if total != Some(self.bytes) || i64::try_from(self.bytes).is_err() {
180            return Err(StorageError::new("invalid checkpoint size"));
181        }
182        Ok(())
183    }
184}
185
186pub(super) fn append_padding(output: &mut String, mut bytes: usize) {
187    const SPACES: &str = "                                                                ";
188    while bytes >= SPACES.len() {
189        output.push_str(SPACES);
190        bytes -= SPACES.len();
191    }
192    output.push_str(&SPACES[..bytes]);
193}
iris_chat_protocol/storage/checkpoint.rs

iris_chat_protocol/storage/
checkpoint.rs
