1use std::collections::HashSet;
2
3use crate::{
4 block::{Block, BlockType, Properties, PropertyValue, Span},
5 document::ZoneMaps,
6 error::MqdbError,
7};
8
9fn invalid_data(message: impl Into<String>) -> MqdbError {
10 MqdbError::Storage(message.into())
11}
12
13fn as_u8(value: usize, field: &str) -> u8 {
14 u8::try_from(value).unwrap_or_else(|_| panic!("{field} exceeds u8 range"))
15}
16
17fn as_u16(value: usize, field: &str) -> u16 {
18 u16::try_from(value).unwrap_or_else(|_| panic!("{field} exceeds u16 range"))
19}
20
21fn as_u32(value: usize, field: &str) -> u32 {
22 u32::try_from(value).unwrap_or_else(|_| panic!("{field} exceeds u32 range"))
23}
24
25fn usize_from_u32(value: u32, field: &str) -> Result<usize, MqdbError> {
26 usize::try_from(value).map_err(|_| invalid_data(format!("{field} exceeds usize range")))
27}
28
29struct Decoder<'a> {
30 data: &'a [u8],
31 pos: usize,
32}
33
34impl<'a> Decoder<'a> {
35 fn new(data: &'a [u8]) -> Self {
36 Self { data, pos: 0 }
37 }
38
39 fn read_exact(&mut self, len: usize) -> Result<&'a [u8], MqdbError> {
40 let end = self
41 .pos
42 .checked_add(len)
43 .ok_or_else(|| invalid_data("byte offset overflow"))?;
44 if end > self.data.len() {
45 return Err(invalid_data("unexpected end of input"));
46 }
47 let bytes = &self.data[self.pos..end];
48 self.pos = end;
49 Ok(bytes)
50 }
51
52 fn read_u8(&mut self) -> Result<u8, MqdbError> {
53 Ok(self.read_exact(1)?[0])
54 }
55
56 fn read_u16(&mut self) -> Result<u16, MqdbError> {
57 let bytes: [u8; 2] = self
58 .read_exact(2)?
59 .try_into()
60 .map_err(|_| invalid_data("failed to read u16"))?;
61 Ok(u16::from_le_bytes(bytes))
62 }
63
64 fn read_u32(&mut self) -> Result<u32, MqdbError> {
65 let bytes: [u8; 4] = self
66 .read_exact(4)?
67 .try_into()
68 .map_err(|_| invalid_data("failed to read u32"))?;
69 Ok(u32::from_le_bytes(bytes))
70 }
71
72 fn read_i64(&mut self) -> Result<i64, MqdbError> {
73 let bytes: [u8; 8] = self
74 .read_exact(8)?
75 .try_into()
76 .map_err(|_| invalid_data("failed to read i64"))?;
77 Ok(i64::from_le_bytes(bytes))
78 }
79
80 fn read_f64(&mut self) -> Result<f64, MqdbError> {
81 let bytes: [u8; 8] = self
82 .read_exact(8)?
83 .try_into()
84 .map_err(|_| invalid_data("failed to read f64"))?;
85 Ok(f64::from_le_bytes(bytes))
86 }
87
88 fn read_string_u16(&mut self) -> Result<String, MqdbError> {
89 let len = usize::from(self.read_u16()?);
90 let bytes = self.read_exact(len)?;
91 String::from_utf8(bytes.to_vec())
92 .map_err(|e| invalid_data(format!("invalid UTF-8 string: {e}")))
93 }
94
95 fn read_string_u32(&mut self) -> Result<String, MqdbError> {
96 let len = usize_from_u32(self.read_u32()?, "string length")?;
97 let bytes = self.read_exact(len)?;
98 String::from_utf8(bytes.to_vec())
99 .map_err(|e| invalid_data(format!("invalid UTF-8 string: {e}")))
100 }
101}
102
103fn encode_block_type(block_type: &BlockType) -> u8 {
104 match block_type {
105 BlockType::Heading => 0,
106 BlockType::Paragraph => 1,
107 BlockType::Code => 2,
108 BlockType::List => 3,
109 BlockType::TableCell => 4,
110 BlockType::TableRow => 5,
111 BlockType::TableAlign => 6,
112 BlockType::Blockquote => 7,
113 BlockType::HorizontalRule => 8,
114 BlockType::Html => 9,
115 BlockType::Yaml => 10,
116 BlockType::Toml => 11,
117 BlockType::Math => 12,
118 BlockType::Definition => 13,
119 BlockType::Footnote => 14,
120 }
121}
122
123fn decode_block_type(value: u8) -> Result<BlockType, MqdbError> {
124 match value {
125 0 => Ok(BlockType::Heading),
126 1 => Ok(BlockType::Paragraph),
127 2 => Ok(BlockType::Code),
128 3 => Ok(BlockType::List),
129 4 => Ok(BlockType::TableCell),
130 5 => Ok(BlockType::TableRow),
131 6 => Ok(BlockType::TableAlign),
132 7 => Ok(BlockType::Blockquote),
133 8 => Ok(BlockType::HorizontalRule),
134 9 => Ok(BlockType::Html),
135 10 => Ok(BlockType::Yaml),
136 11 => Ok(BlockType::Toml),
137 12 => Ok(BlockType::Math),
138 13 => Ok(BlockType::Definition),
139 14 => Ok(BlockType::Footnote),
140 _ => Err(invalid_data(format!("unknown block type tag: {value}"))),
141 }
142}
143
144fn encode_property_value(value: &PropertyValue, out: &mut Vec<u8>) {
145 match value {
146 PropertyValue::Null => out.push(0x00),
147 PropertyValue::String(s) => {
148 out.push(0x01);
149 out.extend_from_slice(&as_u32(s.len(), "string length").to_le_bytes());
150 out.extend_from_slice(s.as_bytes());
151 }
152 PropertyValue::Int(i) => {
153 out.push(0x02);
154 out.extend_from_slice(&i.to_le_bytes());
155 }
156 PropertyValue::Float(f) => {
157 out.push(0x03);
158 out.extend_from_slice(&f.to_le_bytes());
159 }
160 PropertyValue::Bool(b) => {
161 out.push(0x04);
162 out.push(u8::from(*b));
163 }
164 PropertyValue::Array(values) => {
165 out.push(0x05);
166 out.extend_from_slice(&as_u16(values.len(), "array length").to_le_bytes());
167 for value in values {
168 encode_property_value(value, out);
169 }
170 }
171 }
172}
173
174fn decode_property_value(decoder: &mut Decoder<'_>) -> Result<PropertyValue, MqdbError> {
175 match decoder.read_u8()? {
176 0x00 => Ok(PropertyValue::Null),
177 0x01 => Ok(PropertyValue::String(decoder.read_string_u32()?)),
178 0x02 => Ok(PropertyValue::Int(decoder.read_i64()?)),
179 0x03 => Ok(PropertyValue::Float(decoder.read_f64()?)),
180 0x04 => match decoder.read_u8()? {
181 0 => Ok(PropertyValue::Bool(false)),
182 1 => Ok(PropertyValue::Bool(true)),
183 value => Err(invalid_data(format!("invalid bool tag: {value}"))),
184 },
185 0x05 => {
186 let count = usize::from(decoder.read_u16()?);
187 let mut values = Vec::with_capacity(count);
188 for _ in 0..count {
189 values.push(decode_property_value(decoder)?);
190 }
191 Ok(PropertyValue::Array(values))
192 }
193 value => Err(invalid_data(format!("unknown property value tag: {value}"))),
194 }
195}
196
197fn encode_len_prefixed_u16(value: &str, out: &mut Vec<u8>) {
198 out.extend_from_slice(&as_u16(value.len(), "string length").to_le_bytes());
199 out.extend_from_slice(value.as_bytes());
200}
201
202fn sorted_strings(set: &HashSet<String>) -> Vec<&str> {
203 let mut values: Vec<&str> = set.iter().map(String::as_str).collect();
204 values.sort_unstable();
205 values
206}
207
208pub fn encode_block(block: &Block) -> Vec<u8> {
209 let mut out = Vec::new();
210 out.extend_from_slice(&block.id.to_le_bytes());
211 out.extend_from_slice(&block.document_id.to_le_bytes());
212 out.push(encode_block_type(&block.block_type));
213 out.extend_from_slice(&block.pre.to_le_bytes());
214 out.extend_from_slice(&block.post.to_le_bytes());
215
216 match &block.span {
217 Some(span) => {
218 out.push(1);
219 out.extend_from_slice(&as_u32(span.start_line, "span.start_line").to_le_bytes());
220 out.extend_from_slice(&as_u32(span.start_col, "span.start_col").to_le_bytes());
221 out.extend_from_slice(&as_u32(span.end_line, "span.end_line").to_le_bytes());
222 out.extend_from_slice(&as_u32(span.end_col, "span.end_col").to_le_bytes());
223 }
224 None => out.push(0),
225 }
226
227 out.extend_from_slice(&as_u32(block.content.len(), "content length").to_le_bytes());
228 out.extend_from_slice(block.content.as_bytes());
229
230 let mut properties: Vec<(&String, &PropertyValue)> = block.properties.iter().collect();
231 properties.sort_unstable_by_key(|(left, _)| *left);
232
233 out.extend_from_slice(&as_u16(properties.len(), "property count").to_le_bytes());
234 for (key, value) in properties {
235 out.push(as_u8(key.len(), "property key length"));
236 out.extend_from_slice(key.as_bytes());
237 encode_property_value(value, &mut out);
238 }
239
240 out
241}
242
243pub fn decode_block(data: &[u8]) -> Result<(Block, usize), MqdbError> {
244 let mut decoder = Decoder::new(data);
245
246 let id = decoder.read_u32()?;
247 let document_id = decoder.read_u32()?;
248 let block_type = decode_block_type(decoder.read_u8()?)?;
249 let pre = decoder.read_u32()?;
250 let post = decoder.read_u32()?;
251 let span = match decoder.read_u8()? {
252 0 => None,
253 1 => Some(Span {
254 start_line: usize_from_u32(decoder.read_u32()?, "span.start_line")?,
255 start_col: usize_from_u32(decoder.read_u32()?, "span.start_col")?,
256 end_line: usize_from_u32(decoder.read_u32()?, "span.end_line")?,
257 end_col: usize_from_u32(decoder.read_u32()?, "span.end_col")?,
258 }),
259 value => return Err(invalid_data(format!("invalid span presence tag: {value}"))),
260 };
261
262 let content = decoder.read_string_u32()?;
263
264 let prop_count = usize::from(decoder.read_u16()?);
265 let mut properties = Properties::new();
266 for _ in 0..prop_count {
267 let key_len = usize::from(decoder.read_u8()?);
268 let key = String::from_utf8(decoder.read_exact(key_len)?.to_vec())
269 .map_err(|e| invalid_data(format!("invalid property key UTF-8: {e}")))?;
270 let value = decode_property_value(&mut decoder)?;
271 properties.set(key, value);
272 }
273
274 Ok((
275 Block {
276 id,
277 document_id,
278 block_type,
279 content,
280 span,
281 pre,
282 post,
283 properties,
284 },
285 decoder.pos,
286 ))
287}
288
289pub fn encode_table_rows(rows: &[Vec<String>]) -> Vec<u8> {
292 let mut out = Vec::new();
293 for row in rows {
294 for cell in row {
295 out.extend_from_slice(&as_u32(cell.len(), "cell length").to_le_bytes());
296 out.extend_from_slice(cell.as_bytes());
297 }
298 }
299 out
300}
301
302pub fn decode_table_rows(
303 data: &[u8],
304 num_rows: usize,
305 num_cols: usize,
306) -> Result<Vec<Vec<String>>, MqdbError> {
307 let mut decoder = Decoder::new(data);
308 let mut rows = Vec::with_capacity(num_rows);
309 for _ in 0..num_rows {
310 let mut row = Vec::with_capacity(num_cols);
311 for _ in 0..num_cols {
312 row.push(decoder.read_string_u32()?);
313 }
314 rows.push(row);
315 }
316 Ok(rows)
317}
318
319pub fn encode_zone_map(zm: &ZoneMaps) -> Vec<u8> {
320 let mut out = Vec::new();
321 out.push(zm.max_heading_depth);
322
323 let heading_slugs = sorted_strings(&zm.heading_slugs);
324 out.extend_from_slice(&as_u16(heading_slugs.len(), "heading slug count").to_le_bytes());
325 for value in heading_slugs {
326 encode_len_prefixed_u16(value, &mut out);
327 }
328
329 let heading_contents = sorted_strings(&zm.heading_contents);
330 out.extend_from_slice(&as_u16(heading_contents.len(), "heading content count").to_le_bytes());
331 for value in heading_contents {
332 encode_len_prefixed_u16(value, &mut out);
333 }
334
335 let code_languages = sorted_strings(&zm.code_languages);
336 out.extend_from_slice(&as_u16(code_languages.len(), "code language count").to_le_bytes());
337 for value in code_languages {
338 encode_len_prefixed_u16(value, &mut out);
339 }
340
341 let frontmatter_keys = sorted_strings(&zm.frontmatter_keys);
342 out.extend_from_slice(&as_u16(frontmatter_keys.len(), "frontmatter key count").to_le_bytes());
343 for value in frontmatter_keys {
344 encode_len_prefixed_u16(value, &mut out);
345 }
346
347 match &zm.title {
348 Some(title) => {
349 out.push(1);
350 encode_len_prefixed_u16(title, &mut out);
351 }
352 None => out.push(0),
353 }
354
355 out.extend_from_slice(&as_u16(zm.tags.len(), "tag count").to_le_bytes());
356 for tag in &zm.tags {
357 encode_len_prefixed_u16(tag, &mut out);
358 }
359
360 out
361}
362
363pub fn decode_zone_map(data: &[u8]) -> Result<ZoneMaps, MqdbError> {
364 let mut decoder = Decoder::new(data);
365
366 let max_heading_depth = decoder.read_u8()?;
367
368 let heading_slug_count = usize::from(decoder.read_u16()?);
369 let mut heading_slugs = HashSet::with_capacity(heading_slug_count);
370 for _ in 0..heading_slug_count {
371 heading_slugs.insert(decoder.read_string_u16()?);
372 }
373
374 let heading_content_count = usize::from(decoder.read_u16()?);
375 let mut heading_contents = HashSet::with_capacity(heading_content_count);
376 for _ in 0..heading_content_count {
377 heading_contents.insert(decoder.read_string_u16()?);
378 }
379
380 let code_language_count = usize::from(decoder.read_u16()?);
381 let mut code_languages = HashSet::with_capacity(code_language_count);
382 for _ in 0..code_language_count {
383 code_languages.insert(decoder.read_string_u16()?);
384 }
385
386 let frontmatter_key_count = usize::from(decoder.read_u16()?);
387 let mut frontmatter_keys = HashSet::with_capacity(frontmatter_key_count);
388 for _ in 0..frontmatter_key_count {
389 frontmatter_keys.insert(decoder.read_string_u16()?);
390 }
391
392 let title = match decoder.read_u8()? {
393 0 => None,
394 1 => Some(decoder.read_string_u16()?),
395 value => return Err(invalid_data(format!("invalid title presence tag: {value}"))),
396 };
397
398 let tag_count = usize::from(decoder.read_u16()?);
399 let mut tags = Vec::with_capacity(tag_count);
400 for _ in 0..tag_count {
401 tags.push(decoder.read_string_u16()?);
402 }
403
404 Ok(ZoneMaps {
405 max_heading_depth,
406 heading_slugs,
407 heading_contents,
408 code_languages,
409 frontmatter_keys,
410 title,
411 tags,
412 })
413}