1use yo_common::{Code, Error, Result};
29
30use crate::head::{self, ARRAY, COUNT_MAX, DEPTH_MAX, INTERNED, OFFSETS, SORTED, Tag};
31use crate::layout;
32use crate::read::{Value, key_order};
33
34#[derive(Debug, Default)]
40pub struct Builder {
41 out: Vec<u8>,
44 open: Vec<Open>,
46 members: Vec<Member>,
49 keys: Vec<u8>,
51 scratch: Vec<u8>,
54 pending: Option<Member>,
56 seq: u32,
59}
60
61#[derive(Debug)]
63struct Open {
64 at: usize,
66 flags: u32,
68 first: usize,
70 keys_at: usize,
73 key: Member,
75}
76
77#[derive(Debug, Default, Clone, Copy)]
79struct Member {
80 head: u32,
82 at: u32,
84 len: u32,
86 key_at: u32,
88 key_len: u32,
89 id: u16,
91 seq: u32,
93}
94
95impl Builder {
96 #[must_use]
98 pub fn new() -> Builder {
99 Builder::default()
100 }
101
102 #[must_use]
104 pub fn with_capacity(bytes: usize) -> Builder {
105 Builder {
106 out: Vec::with_capacity(bytes),
107 ..Builder::default()
108 }
109 }
110
111 pub fn clear(&mut self) {
113 self.out.clear();
114 self.open.clear();
115 self.members.clear();
116 self.keys.clear();
117 self.pending = None;
118 self.seq = 0;
119 }
120
121 pub fn finish(&mut self) -> Result<&[u8]> {
127 if let Some(open) = self.open.last() {
128 let what = if open.flags & ARRAY != 0 {
129 "array"
130 } else {
131 "object"
132 };
133 return Err(Error::fmt(
134 Code::Invalid,
135 format_args!("the document ends inside an unclosed {what}"),
136 ));
137 }
138 if self.pending.is_some() {
139 return Err(Error::new(Code::Invalid, "a key with no value after it"));
140 }
141 if self.out.is_empty() {
142 return Err(Error::new(Code::Invalid, "the document holds no value"));
143 }
144 Ok(&self.out)
145 }
146
147 pub fn null(&mut self) -> Result<()> {
149 self.scalar(Tag::Null, &[])
150 }
151
152 pub fn bool(&mut self, v: bool) -> Result<()> {
154 self.scalar(if v { Tag::True } else { Tag::False }, &[])
155 }
156
157 pub fn int(&mut self, v: i64) -> Result<()> {
159 let raw = v.to_le_bytes();
160 self.scalar(Tag::Int, &raw[..int_width(v)])
161 }
162
163 pub fn float(&mut self, v: f64) -> Result<()> {
165 self.scalar(Tag::Float, &v.to_le_bytes())
166 }
167
168 pub fn text(&mut self, v: &str) -> Result<()> {
170 self.scalar(Tag::Text, v.as_bytes())
171 }
172
173 pub fn text_bytes(&mut self, v: &[u8]) -> Result<()> {
181 self.scalar(Tag::Text, v)
182 }
183
184 pub fn embed(&mut self, v: &Value<'_>) -> Result<()> {
190 let bytes = v
191 .as_bytes()
192 .ok_or_else(|| Error::new(Code::Corrupt, "the value being copied is not readable"))?;
193 self.start()?;
194 let at = self.out.len();
195 self.out.extend_from_slice(bytes);
196 self.record(at)
197 }
198
199 pub fn begin_object(&mut self) -> Result<()> {
201 self.begin(0)
202 }
203
204 pub fn begin_object_interned(&mut self) -> Result<()> {
211 self.begin(INTERNED)
212 }
213
214 pub fn begin_array(&mut self) -> Result<()> {
216 self.begin(ARRAY)
217 }
218
219 pub fn end_object(&mut self) -> Result<()> {
221 self.end(false)
222 }
223
224 pub fn end_array(&mut self) -> Result<()> {
226 self.end(true)
227 }
228
229 pub fn key(&mut self, key: &[u8]) -> Result<()> {
235 let open = self.expect_object()?;
236 if open.flags & INTERNED != 0 {
237 return Err(Error::new(
238 Code::Invalid,
239 "this object takes key ids, not key bytes",
240 ));
241 }
242 if key.len() > COUNT_MAX {
243 return Err(Error::new(Code::Full, "the key is longer than 16 MiB"));
244 }
245 self.stash(Member {
246 key_at: u32::try_from(self.keys.len()).map_err(|_| too_big())?,
247 key_len: key.len() as u32,
248 ..Member::default()
249 })?;
250 self.keys.extend_from_slice(key);
251 Ok(())
252 }
253
254 pub fn key_id(&mut self, id: u16) -> Result<()> {
256 let open = self.expect_object()?;
257 if open.flags & INTERNED == 0 {
258 return Err(Error::new(
259 Code::Invalid,
260 "this object takes key bytes, not key ids",
261 ));
262 }
263 self.stash(Member {
264 id,
265 ..Member::default()
266 })
267 }
268
269 fn expect_object(&self) -> Result<&Open> {
271 match self.open.last() {
272 Some(open) if open.flags & ARRAY == 0 => Ok(open),
273 Some(_) => Err(Error::new(Code::Invalid, "an array element has no key")),
274 None => Err(Error::new(Code::Invalid, "there is no object open")),
275 }
276 }
277
278 fn stash(&mut self, key: Member) -> Result<()> {
280 if self.pending.is_some() {
281 return Err(Error::new(Code::Invalid, "two keys in a row"));
282 }
283 self.pending = Some(key);
284 Ok(())
285 }
286
287 fn scalar(&mut self, tag: Tag, payload: &[u8]) -> Result<()> {
289 if payload.len() > COUNT_MAX {
290 return Err(Error::new(Code::Full, "the value is longer than 16 MiB"));
291 }
292 self.start()?;
293 let at = self.out.len();
294 let h = head::head(tag, 0, payload.len());
295 self.out.extend_from_slice(&h.to_le_bytes());
296 self.out.extend_from_slice(payload);
297 self.record(at)
298 }
299
300 fn start(&mut self) -> Result<()> {
303 match self.open.last() {
304 Some(open) if open.flags & ARRAY == 0 && self.pending.is_none() => Err(Error::new(
305 Code::Invalid,
306 "an object member needs a key before its value",
307 )),
308 Some(_) => Ok(()),
309 None if self.out.is_empty() => Ok(()),
310 None => Err(Error::new(
311 Code::Invalid,
312 "a document holds one value, and it is already written",
313 )),
314 }
315 }
316
317 fn record(&mut self, at: usize) -> Result<()> {
320 if self.open.is_empty() {
321 return Ok(());
322 }
323 let mut m = self.pending.take().unwrap_or_default();
324 m.head = head::read(&self.out, at).expect("the header was just written");
325 m.at = u32::try_from(at).map_err(|_| too_big())?;
326 m.len = u32::try_from(self.out.len() - at).map_err(|_| too_big())?;
327 m.seq = self.seq;
328 self.seq += 1;
329 self.members.push(m);
330 Ok(())
331 }
332
333 fn begin(&mut self, flags: u32) -> Result<()> {
335 if self.open.len() >= DEPTH_MAX {
336 return Err(Error::fmt(
337 Code::Full,
338 format_args!("a document nests at most {DEPTH_MAX} deep"),
339 ));
340 }
341 self.start()?;
342 let at = self.out.len();
343 self.out.extend_from_slice(&[0; 4]);
344 self.open.push(Open {
345 at,
346 flags,
347 first: self.members.len(),
348 keys_at: self.keys.len(),
349 key: self.pending.take().unwrap_or_default(),
350 });
351 Ok(())
352 }
353
354 fn end(&mut self, array: bool) -> Result<()> {
364 let Some(open) = self.open.pop() else {
365 return Err(Error::new(Code::Invalid, "nothing is open"));
366 };
367 if array != (open.flags & ARRAY != 0) {
368 return Err(Error::new(
369 Code::Invalid,
370 "an object is not ended by ending an array, or the other way round",
371 ));
372 }
373 if self.pending.is_some() {
374 return Err(Error::new(Code::Invalid, "a key with no value after it"));
375 }
376 if !array {
377 self.sort_members(&open);
378 }
379 let n = self.members.len() - open.first;
380 if n > COUNT_MAX {
381 return Err(Error::fmt(
382 Code::Full,
383 format_args!("a container holds at most {COUNT_MAX} elements"),
384 ));
385 }
386
387 let sorted = if array { 0 } else { SORTED };
388 let h = head::head(Tag::Container, open.flags | OFFSETS | sorted, n);
389 let entries_end = 4 + layout::keys_area(h, n) + n * 8;
390 let key_bytes: usize = self.members[open.first..]
391 .iter()
392 .map(|m| m.key_len as usize)
393 .sum();
394
395 let children_at = open.at + 4;
396 self.scratch.clear();
397 self.scratch.extend_from_slice(&self.out[children_at..]);
398 self.out.truncate(children_at);
399 self.out[open.at..children_at].copy_from_slice(&h.to_le_bytes());
400
401 if !array {
402 if open.flags & INTERNED != 0 {
403 for i in open.first..self.members.len() {
404 self.out
405 .extend_from_slice(&self.members[i].id.to_le_bytes());
406 }
407 if n % 2 == 1 {
411 self.out.extend_from_slice(&[0; 2]);
412 }
413 } else {
414 let mut key_at = entries_end;
415 for i in open.first..self.members.len() {
416 let off = u32::try_from(key_at).map_err(|_| too_big())?;
417 self.out.extend_from_slice(&off.to_le_bytes());
418 key_at += self.members[i].key_len as usize;
419 }
420 }
421 }
422
423 let mut val_at = entries_end + key_bytes;
424 for i in open.first..self.members.len() {
425 let m = self.members[i];
426 self.out.extend_from_slice(&m.head.to_le_bytes());
427 let off = u32::try_from(val_at).map_err(|_| too_big())?;
428 self.out.extend_from_slice(&off.to_le_bytes());
429 val_at += m.len as usize;
430 }
431
432 if !array && open.flags & INTERNED == 0 {
433 for i in open.first..self.members.len() {
434 let m = self.members[i];
435 let at = m.key_at as usize;
436 self.out
437 .extend_from_slice(&self.keys[at..at + m.key_len as usize]);
438 }
439 }
440
441 for i in open.first..self.members.len() {
446 let m = self.members[i];
447 let from = m.at as usize - children_at;
448 self.out
449 .extend_from_slice(&self.scratch[from..from + m.len as usize]);
450 }
451
452 self.members.truncate(open.first);
453 self.keys.truncate(open.keys_at);
454 if !self.open.is_empty() {
455 self.pending = Some(open.key);
456 }
457 self.record(open.at)
458 }
459
460 fn sort_members(&mut self, open: &Open) {
463 let interned = open.flags & INTERNED != 0;
464 let keys = &self.keys;
465 let key_of = |m: &Member| {
466 let at = m.key_at as usize;
467 &keys[at..at + m.key_len as usize]
468 };
469 self.members[open.first..].sort_by(|a, b| {
470 if interned {
471 a.id.cmp(&b.id).then(a.seq.cmp(&b.seq))
472 } else {
473 key_order(key_of(a), key_of(b)).then(a.seq.cmp(&b.seq))
474 }
475 });
476
477 let same = |a: &Member, b: &Member| {
478 if interned {
479 a.id == b.id
480 } else {
481 key_of(a) == key_of(b)
482 }
483 };
484 let mut write = open.first;
485 let mut read = open.first;
486 while read < self.members.len() {
487 let mut run = read + 1;
488 while run < self.members.len() && same(&self.members[read], &self.members[run]) {
489 run += 1;
490 }
491 self.members[write] = self.members[run - 1];
496 write += 1;
497 read = run;
498 }
499 self.members.truncate(write);
500 }
501}
502
503fn int_width(v: i64) -> usize {
505 if i64::from(v as i8) == v {
506 1
507 } else if i64::from(v as i16) == v {
508 2
509 } else if i64::from(v as i32) == v {
510 4
511 } else {
512 8
513 }
514}
515
516fn too_big() -> Error {
517 Error::new(Code::Full, "a document is at most four gigabytes")
518}
519
520#[cfg(test)]
521mod tests {
522 use super::*;
523 use crate::head::Kind;
524
525 fn built(f: impl FnOnce(&mut Builder) -> Result<()>) -> Vec<u8> {
528 let mut b = Builder::new();
529 f(&mut b).expect("the builder accepted every call");
530 let bytes = b.finish().expect("the value is finished").to_vec();
531 let v = Value::new(&bytes).expect("the reader accepts it");
532 assert!(v.validate(), "the value is self consistent");
533 assert_eq!(
534 v.encoded_len(),
535 Some(bytes.len()),
536 "the value is exactly as long as the buffer"
537 );
538 bytes
539 }
540
541 #[test]
542 fn every_scalar_comes_back_as_itself() {
543 let cases: Vec<(Vec<u8>, Kind)> = vec![
544 (built(|b| b.null()), Kind::Null),
545 (built(|b| b.bool(true)), Kind::Bool),
546 (built(|b| b.bool(false)), Kind::Bool),
547 (built(|b| b.int(-9)), Kind::Int),
548 (built(|b| b.float(1.5)), Kind::Float),
549 (built(|b| b.text("hello")), Kind::Text),
550 ];
551 for (bytes, kind) in &cases {
552 assert_eq!(Value::new(bytes).expect("readable").kind(), *kind);
553 }
554 assert!(Value::new(&cases[0].0).expect("readable").is_null());
555 assert_eq!(
556 Value::new(&cases[1].0).expect("readable").as_bool(),
557 Some(true)
558 );
559 assert_eq!(
560 Value::new(&cases[2].0).expect("readable").as_bool(),
561 Some(false)
562 );
563 assert_eq!(
564 Value::new(&cases[3].0).expect("readable").as_int(),
565 Some(-9)
566 );
567 assert_eq!(
568 Value::new(&cases[4].0).expect("readable").as_float(),
569 Some(1.5)
570 );
571 assert_eq!(
572 Value::new(&cases[5].0).expect("readable").as_text(),
573 Some("hello")
574 );
575 }
576
577 #[test]
578 fn an_integer_takes_as_few_bytes_as_it_fits_in() {
579 let cases = [
582 (0i64, 1usize),
583 (127, 1),
584 (-128, 1),
585 (128, 2),
586 (-129, 2),
587 (32_767, 2),
588 (-32_768, 2),
589 (32_768, 4),
590 (2_147_483_647, 4),
591 (-2_147_483_648, 4),
592 (2_147_483_648, 8),
593 (i64::MIN, 8),
594 (i64::MAX, 8),
595 ];
596 for (v, width) in cases {
597 let bytes = built(|b| b.int(v));
598 assert_eq!(bytes.len(), 4 + width, "{v} takes {width} bytes");
599 assert_eq!(Value::new(&bytes).expect("readable").as_int(), Some(v));
600 }
601 }
602
603 #[test]
604 fn an_object_comes_back_in_key_order_whatever_order_it_went_in() {
605 let bytes = built(|b| {
606 b.begin_object()?;
607 for k in ["zebra", "b", "aa", "a", "yak"] {
608 b.key(k.as_bytes())?;
609 b.text(k)?;
610 }
611 b.end_object()
612 });
613 let v = Value::new(&bytes).expect("readable");
614 let keys: Vec<&[u8]> = v.members().map(|(k, _)| k).collect();
615 assert_eq!(keys, [&b"a"[..], b"b", b"aa", b"yak", b"zebra"]);
617 for k in ["zebra", "b", "aa", "a", "yak"] {
618 assert_eq!(v.get(k.as_bytes()).expect("found").as_text(), Some(k));
619 }
620 assert!(v.get(b"nope").is_none());
621 assert!(v.get(b"").is_none());
622 }
623
624 #[test]
625 fn writing_a_key_twice_keeps_the_last_one_and_leaves_no_dead_bytes() {
626 let bytes = built(|b| {
627 b.begin_object()?;
628 b.key(b"a")?;
629 b.int(1)?;
630 b.key(b"b")?;
631 b.int(2)?;
632 b.key(b"a")?;
633 b.text("the winner")?;
634 b.key(b"a")?;
635 b.int(3)?;
636 b.end_object()
637 });
638 let v = Value::new(&bytes).expect("readable");
639 assert_eq!(v.len(), 2, "two keys, however many times they were written");
640 assert_eq!(v.get(b"a").expect("found").as_int(), Some(3));
641 assert_eq!(v.get(b"b").expect("found").as_int(), Some(2));
642 assert_eq!(bytes.len(), 4 + 2 * 4 + 2 * 8 + 2 + 5 + 5);
645 }
646
647 #[test]
648 fn a_nested_document_reads_at_every_level() {
649 let bytes = built(|b| {
650 b.begin_object()?;
651 b.key(b"id")?;
652 b.int(7)?;
653 b.key(b"lines")?;
654 b.begin_array()?;
655 for i in 0..3i64 {
656 b.begin_object()?;
657 b.key(b"sku")?;
658 b.int(i)?;
659 b.key(b"note")?;
660 b.text("a line of some length so the offsets move")?;
661 b.end_object()?;
662 }
663 b.end_array()?;
664 b.key(b"open")?;
665 b.bool(true)?;
666 b.end_object()
667 });
668 let v = Value::new(&bytes).expect("readable");
669 assert_eq!(v.get(b"id").expect("found").as_int(), Some(7));
670 assert_eq!(v.get(b"open").expect("found").as_bool(), Some(true));
671 let lines = v.get(b"lines").expect("found");
672 assert_eq!(lines.kind(), Kind::Array);
673 assert_eq!(lines.len(), 3);
674 for i in 0..3i64 {
675 let line = lines.at(i as usize).expect("an element");
676 assert_eq!(line.get(b"sku").expect("found").as_int(), Some(i));
677 assert!(line.get(b"note").expect("found").as_text().is_some());
678 let alone = line.as_bytes().expect("a length");
681 let again = Value::new(alone).expect("readable on its own");
682 assert!(again.validate());
683 assert_eq!(again.get(b"sku").expect("found").as_int(), Some(i));
684 }
685 }
686
687 #[test]
688 fn an_empty_container_is_four_bytes() {
689 let obj = built(|b| {
690 b.begin_object()?;
691 b.end_object()
692 });
693 assert_eq!(obj.len(), 4);
694 let v = Value::new(&obj).expect("readable");
695 assert_eq!(v.kind(), Kind::Object);
696 assert!(v.is_empty());
697 assert!(v.get(b"a").is_none());
698
699 let arr = built(|b| {
700 b.begin_array()?;
701 b.end_array()
702 });
703 assert_eq!(arr.len(), 4);
704 let v = Value::new(&arr).expect("readable");
705 assert_eq!(v.kind(), Kind::Array);
706 assert!(v.is_empty());
707 assert!(v.at(0).is_none());
708 }
709
710 #[test]
711 fn an_interned_object_looks_up_by_id() {
712 for n in [1u16, 2, 3, 8, 9] {
715 let bytes = built(|b| {
716 b.begin_object_interned()?;
717 for id in (0..n).rev() {
718 b.key_id(id * 3)?;
719 b.int(i64::from(id))?;
720 }
721 b.end_object()
722 });
723 let v = Value::new(&bytes).expect("readable");
724 assert!(v.is_interned());
725 assert_eq!(v.len(), usize::from(n));
726 for id in 0..n {
727 assert_eq!(
728 v.get_id(id * 3).expect("found").as_int(),
729 Some(i64::from(id))
730 );
731 }
732 assert!(v.get_id(1).is_none(), "1 is not a multiple of 3");
733 assert!(v.key_at(0).is_none(), "the names are not in the document");
734 assert_eq!(v.key_id_at(0), Some(0));
735 }
736 }
737
738 #[test]
739 fn a_thousand_keys_are_all_findable() {
740 let names: Vec<String> = (0..1_000).map(|i| format!("field{i}")).collect();
741 let bytes = built(|b| {
742 b.begin_object()?;
743 for (i, name) in names.iter().enumerate() {
744 b.key(name.as_bytes())?;
745 b.int(i as i64)?;
746 }
747 b.end_object()
748 });
749 let v = Value::new(&bytes).expect("readable");
750 assert_eq!(v.len(), 1_000);
751 for (i, name) in names.iter().enumerate() {
752 assert_eq!(
753 v.get(name.as_bytes()).expect("found").as_int(),
754 Some(i as i64)
755 );
756 }
757 assert!(v.get(b"field1000").is_none());
758 }
759
760 #[test]
761 fn a_value_that_is_already_encoded_can_be_copied_in() {
762 let inner = built(|b| {
763 b.begin_object()?;
764 b.key(b"x")?;
765 b.int(3)?;
766 b.end_object()
767 });
768 let bytes = built(|b| {
769 b.begin_array()?;
770 b.int(1)?;
771 b.embed(&Value::new(&inner).expect("readable"))?;
772 b.int(2)?;
773 b.end_array()
774 });
775 let v = Value::new(&bytes).expect("readable");
776 assert_eq!(v.len(), 3);
777 assert_eq!(
778 v.at(1)
779 .expect("an element")
780 .get(b"x")
781 .expect("found")
782 .as_int(),
783 Some(3)
784 );
785 }
786
787 #[test]
788 fn a_builder_can_be_used_again() {
789 let mut b = Builder::new();
790 b.int(1).expect("a value");
791 assert_eq!(b.finish().expect("finished").len(), 5);
792 b.clear();
793 b.text("hello").expect("a value");
794 let bytes = b.finish().expect("finished");
795 assert_eq!(
796 Value::new(bytes).expect("readable").as_text(),
797 Some("hello")
798 );
799 }
800
801 #[test]
802 fn the_builder_says_no_to_every_way_of_getting_it_wrong() {
803 let bad = |f: fn(&mut Builder) -> Result<()>| {
804 let mut b = Builder::new();
805 f(&mut b).unwrap_err()
806 };
807
808 assert!(
810 bad(|b| {
811 b.begin_object()?;
812 b.key(b"a")?;
813 b.end_object()
814 })
815 .message()
816 .contains("no value")
817 );
818 assert!(
820 bad(|b| {
821 b.begin_object()?;
822 b.key(b"a")?;
823 b.key(b"b")
824 })
825 .message()
826 .contains("two keys")
827 );
828 assert!(
830 bad(|b| {
831 b.begin_object()?;
832 b.int(1)
833 })
834 .message()
835 .contains("needs a key")
836 );
837 assert!(
839 bad(|b| {
840 b.begin_array()?;
841 b.key(b"a")
842 })
843 .message()
844 .contains("no key")
845 );
846 assert!(bad(|b| b.key(b"a")).message().contains("no object open"));
848 assert!(
850 bad(|b| {
851 b.begin_object()?;
852 b.end_array()
853 })
854 .message()
855 .contains("not ended by")
856 );
857 assert!(
859 bad(|b| b.end_object())
860 .message()
861 .contains("nothing is open")
862 );
863 assert!(
865 bad(|b| {
866 b.int(1)?;
867 b.int(2)
868 })
869 .message()
870 .contains("already written")
871 );
872 assert!(
874 bad(|b| {
875 b.begin_object_interned()?;
876 b.key(b"a")
877 })
878 .message()
879 .contains("key ids")
880 );
881 assert!(
882 bad(|b| {
883 b.begin_object()?;
884 b.key_id(1)
885 })
886 .message()
887 .contains("key bytes")
888 );
889 }
890
891 #[test]
892 fn finishing_early_is_an_error_and_not_a_short_document() {
893 let mut b = Builder::new();
894 assert!(b.finish().unwrap_err().message().contains("no value"));
895 b.begin_array().expect("open");
896 assert!(b.finish().unwrap_err().message().contains("unclosed array"));
897 b.end_array().expect("close");
898 b.finish().expect("finished now");
899
900 let mut b = Builder::new();
901 b.begin_object().expect("open");
902 assert!(
903 b.finish()
904 .unwrap_err()
905 .message()
906 .contains("unclosed object")
907 );
908 }
909
910 #[test]
911 fn a_document_nests_as_deep_as_the_reader_will_walk_and_no_deeper() {
912 let mut b = Builder::new();
913 for _ in 0..DEPTH_MAX {
914 b.begin_array().expect("within the limit");
915 }
916 assert!(
917 b.begin_array()
918 .unwrap_err()
919 .message()
920 .contains("nests at most"),
921 "one past the limit is refused"
922 );
923 for _ in 0..DEPTH_MAX {
924 b.end_array().expect("close");
925 }
926 let bytes = b.finish().expect("finished").to_vec();
927 let v = Value::new(&bytes).expect("readable");
928 assert!(v.validate(), "the reader walks all of it");
929 assert_eq!(v.encoded_len(), Some(bytes.len()));
930 }
931}