1use alloc::boxed::Box;
2use alloc::format;
3use alloc::vec::Vec;
4use core::mem::ManuallyDrop;
5
6use deser_core::__format::extend;
7use deser_core::State;
8use deser_core::ext::{BigInt, ExtValue, Timestamp};
9use deser_core::ser::{self, PausableSink, SerializeDriver, Written};
10use deser_core::{Atom, ContainerShape, Error, ErrorKind, Event, Serialize};
11
12use crate::ext::{Ext, TIMESTAMP, encode_timestamp};
13
14#[derive(Debug, Clone, Default, PartialEq, Eq)]
35pub struct SerializerConfig {
36 canonical: bool,
37}
38
39#[derive(Clone, Copy)]
44struct Frame {
45 remaining: u64,
50 info: usize,
53}
54
55const IS_MAP: usize = 1;
56const UNKNOWN_LEN: usize = 2;
57
58impl Frame {
59 const TOP: Frame = Frame {
61 remaining: u64::MAX,
62 info: UNKNOWN_LEN,
63 };
64
65 #[inline(always)]
66 fn is_map(self) -> bool {
67 self.info & IS_MAP != 0
68 }
69
70 #[inline(always)]
71 fn header(self) -> usize {
72 self.info >> 2
73 }
74}
75
76struct CanonicalMap {
78 body: usize,
80 offsets_start: usize,
83}
84
85pub(crate) struct Writer {
87 pub(crate) out: Vec<u8>,
88 canonical: bool,
89 frame: Frame,
92 stack: Vec<Frame>,
93 maps: Vec<CanonicalMap>,
96 offsets: Vec<usize>,
97 insertions: Vec<Insertion>,
99 open_unknown: usize,
101 limit: usize,
103}
104
105struct Insertion {
108 offset: usize,
109 len: u8,
110 bytes: [u8; 4],
111}
112
113impl ser::EventSink for Writer {
114 #[inline(always)]
115 fn event(&mut self, event: Event, _state: &mut State) -> Result<(), Error> {
116 Writer::event(self, event)
117 }
118}
119
120impl PausableSink for Writer {
121 #[inline(always)]
122 fn event(
123 &mut self,
124 event: Event<'_>,
125 _value: &dyn Serialize,
126 _state: &mut State,
127 ) -> Result<(), Error> {
128 Writer::event(self, event)
129 }
130
131 #[inline]
132 fn pause(&mut self) -> bool {
133 if self.out.len() < self.limit || self.open_unknown > 0 || !self.maps.is_empty() {
136 return false;
137 }
138 self.finish();
139 true
140 }
141}
142
143impl Writer {
144 pub(crate) fn new(canonical: bool, out: Vec<u8>) -> Writer {
146 Writer {
147 out,
148 canonical,
149 frame: Frame::TOP,
150 stack: Vec::new(),
151 maps: Vec::new(),
152 offsets: Vec::new(),
153 insertions: Vec::new(),
154 open_unknown: 0,
155 limit: usize::MAX,
156 }
157 }
158
159 pub(crate) fn finish(&mut self) {
162 if !self.insertions.is_empty() {
163 self.apply_insertions();
164 self.insertions.clear();
165 }
166 }
167
168 pub(crate) fn drive(
174 &mut self,
175 driver: &mut SerializeDriver<'_>,
176 limit: usize,
177 ) -> Result<bool, Error> {
178 if limit == usize::MAX {
179 return self.drive_whole(driver).map(|()| true);
180 }
181 self.limit = limit;
182 let done = driver.drive_until(self)?;
183 if done {
184 self.finish();
185 }
186 Ok(done)
187 }
188
189 pub(crate) fn drive_whole(&mut self, driver: &mut SerializeDriver<'_>) -> Result<(), Error> {
194 driver.drive_sink(self)?;
195 self.finish();
196 Ok(())
197 }
198
199 #[inline(always)]
200 fn event(&mut self, event: Event) -> Result<(), Error> {
201 match event {
202 Event::Atom(atom) => {
203 self.begin_item();
204 self.write_atom(atom)
205 }
206 Event::MapStart(shape) => self.start(true, shape),
207 Event::SeqStart(shape) => self.start(false, shape),
208 Event::MapEnd | Event::SeqEnd => self.end(),
209 }
210 }
211
212 #[inline(always)]
214 fn begin_item(&mut self) {
215 self.frame.remaining = self.frame.remaining.wrapping_sub(1);
216 if self.canonical && self.frame.is_map() {
217 self.offsets.push(self.out.len());
218 }
219 }
220
221 #[inline(always)]
222 fn start(&mut self, is_map: bool, shape: ContainerShape) -> Result<(), Error> {
223 self.begin_item();
224 let header = self.out.len();
225 let mut info = (header << 2) | if is_map { IS_MAP } else { 0 };
226 let remaining = match shape.len() {
229 Some(len) => {
230 let len = check_len(len)?;
231 let mut buf = [0u8; 5];
232 extend(&mut self.out, encode_container_head(&mut buf, is_map, len));
233 if is_map {
234 u64::from(len) * 2
235 } else {
236 u64::from(len)
237 }
238 }
239 None => {
240 self.out.push(if is_map { 0x80 } else { 0x90 });
241 info |= UNKNOWN_LEN;
242 self.open_unknown += 1;
243 u64::MAX
244 }
245 };
246 if self.canonical && is_map {
247 self.maps.push(CanonicalMap {
248 body: self.out.len(),
249 offsets_start: self.offsets.len(),
250 });
251 }
252 self.stack.push(core::mem::replace(
253 &mut self.frame,
254 Frame { remaining, info },
255 ));
256 Ok(())
257 }
258
259 #[inline(always)]
260 fn end(&mut self) -> Result<(), Error> {
261 let Some(parent) = self.stack.pop() else {
262 return Err(Error::new(ErrorKind::Unexpected, "unexpected end"));
263 };
264 let frame = core::mem::replace(&mut self.frame, parent);
265 if frame.remaining == 0 && !self.canonical {
266 Ok(())
267 } else {
268 self.end_slow(frame)
269 }
270 }
271
272 #[inline(never)]
276 fn end_slow(&mut self, frame: Frame) -> Result<(), Error> {
277 let unknown = frame.info & UNKNOWN_LEN != 0;
278 if !unknown && frame.remaining != 0 {
279 return Err(Error::new(
280 ErrorKind::Unexpected,
281 "number of items does not match the length of the container",
282 ));
283 }
284 let items = if unknown {
285 u64::MAX - frame.remaining
286 } else {
287 0
288 };
289 if frame.is_map() {
290 if !items.is_multiple_of(2) {
291 return Err(Error::new(ErrorKind::Unexpected, "map without value"));
292 }
293 if self.canonical {
294 let map = self.maps.pop().unwrap();
295 self.sort_entries(&map)?;
296 }
297 }
298 if unknown {
299 self.open_unknown -= 1;
300 let count = if frame.is_map() { items / 2 } else { items };
301 let count = u32::try_from(count).map_err(|_| too_long())?;
302 self.patch_length(frame.header(), frame.is_map(), count);
303 }
304 Ok(())
305 }
306
307 fn patch_length(&mut self, header: usize, is_map: bool, count: u32) {
315 let mut buf = [0u8; 5];
316 let head = encode_container_head(&mut buf, is_map, count);
317 self.out[header] = head[0];
318 let extra = head.len() - 1;
319 if extra == 0 {
320 return;
321 }
322 if self.canonical {
323 let len = self.out.len();
324 self.out.resize(len + extra, 0);
325 self.out.copy_within(header + 1..len, header + 1 + extra);
326 self.out[header + 1..header + head.len()].copy_from_slice(&head[1..]);
327 } else {
328 let mut bytes = [0; 4];
329 bytes[..extra].copy_from_slice(&head[1..]);
330 self.insertions.push(Insertion {
331 offset: header + 1,
332 len: extra as u8,
333 bytes,
334 });
335 }
336 }
337
338 #[cold]
340 fn apply_insertions(&mut self) {
341 self.insertions
342 .sort_unstable_by_key(|insertion| insertion.offset);
343 let extra: usize = self.insertions.iter().map(|x| x.len as usize).sum();
344 let mut src_end = self.out.len();
345 self.out.resize(src_end + extra, 0);
346 let mut dst_end = self.out.len();
347 for insertion in self.insertions.iter().rev() {
350 let segment = src_end - insertion.offset;
351 self.out
352 .copy_within(insertion.offset..src_end, dst_end - segment);
353 dst_end -= segment + insertion.len as usize;
354 self.out[dst_end..dst_end + insertion.len as usize]
355 .copy_from_slice(&insertion.bytes[..insertion.len as usize]);
356 src_end = insertion.offset;
357 }
358 debug_assert_eq!(src_end, dst_end);
359 }
360
361 #[cold]
363 fn sort_entries(&mut self, map: &CanonicalMap) -> Result<(), Error> {
364 let offsets = self.offsets.split_off(map.offsets_start);
365 let body_start = map.body;
366 let body_end = self.out.len();
367 let mut entries: Vec<(usize, usize, usize)> = (0..offsets.len())
369 .step_by(2)
370 .map(|idx| {
371 let end = offsets.get(idx + 2).copied().unwrap_or(body_end);
372 (offsets[idx], offsets[idx + 1], end)
373 })
374 .collect();
375 let out = &self.out;
376 entries.sort_by(|a, b| out[a.0..a.1].cmp(&out[b.0..b.1]));
377 if entries
378 .windows(2)
379 .any(|pair| out[pair[0].0..pair[0].1] == out[pair[1].0..pair[1].1])
380 {
381 return Err(Error::new(ErrorKind::Unexpected, "duplicate map key"));
382 }
383 let mut body = Vec::with_capacity(body_end - body_start);
384 for (start, _, end) in entries {
385 body.extend_from_slice(&out[start..end]);
386 }
387 self.out[body_start..body_end].copy_from_slice(&body);
388 Ok(())
389 }
390
391 #[inline(always)]
392 fn write_atom(&mut self, atom: Atom) -> Result<(), Error> {
393 let atom = ManuallyDrop::new(atom);
396 match *atom {
397 Atom::Null => self.out.push(0xc0),
398 Atom::Bool(false) => self.out.push(0xc2),
399 Atom::Bool(true) => self.out.push(0xc3),
400 Atom::Str(ref val) if val.is_borrowed() => self.write_str(val)?,
401 Atom::Bytes(ref val) if val.is_borrowed() => self.write_bin(val)?,
402 Atom::Char(c) => self.write_str(c.encode_utf8(&mut [0u8; 4]))?,
403 Atom::U64(val) => self.write_u64(val),
404 Atom::I64(val) => self.write_i64(val),
405 Atom::F64(val) => {
406 let mut buf = [0xcb; 9];
407 buf[1..].copy_from_slice(&val.to_be_bytes());
408 extend(&mut self.out, &buf);
409 }
410 Atom::F32(val) => {
411 let mut buf = [0xca; 5];
412 buf[1..].copy_from_slice(&val.to_be_bytes());
413 extend(&mut self.out, &buf);
414 }
415 _ => return self.write_other_atom(ManuallyDrop::into_inner(atom)),
416 }
417 Ok(())
418 }
419
420 #[inline(never)]
421 fn write_other_atom(&mut self, atom: Atom) -> Result<(), Error> {
422 match atom {
423 Atom::Str(ref val) | Atom::Lexical(ref val) => self.write_str(val),
424 Atom::Bytes(ref val) => self.write_bin(val),
425 Atom::Ext(ref ext) => self.write_ext(ext),
426 Atom::Implicit(ref val) => self.write_atom(val.value().to_atom()),
428 _ => Err(Error::new(ErrorKind::UnsupportedType, "unknown atom")),
429 }
430 }
431
432 #[inline(always)]
436 fn write_len(&mut self, heads: [u8; 3], len: usize) -> Result<(), Error> {
437 let len = check_len(len)?;
438 if len <= u32::from(u8::MAX) {
439 extend(&mut self.out, &[heads[0], len as u8]);
440 } else if len <= u32::from(u16::MAX) {
441 let [a, b] = (len as u16).to_be_bytes();
442 extend(&mut self.out, &[heads[1], a, b]);
443 } else {
444 let [a, b, c, d] = len.to_be_bytes();
445 extend(&mut self.out, &[heads[2], a, b, c, d]);
446 }
447 Ok(())
448 }
449
450 #[inline(always)]
451 fn write_bin(&mut self, val: &[u8]) -> Result<(), Error> {
452 self.write_len([0xc4, 0xc5, 0xc6], val.len())?;
453 extend(&mut self.out, val);
454 Ok(())
455 }
456
457 #[inline(always)]
458 fn write_str(&mut self, val: &str) -> Result<(), Error> {
459 if val.len() < 32 {
460 self.out.push(0xa0 | val.len() as u8);
461 } else {
462 self.write_len([0xd9, 0xda, 0xdb], val.len())?;
463 }
464 extend(&mut self.out, val.as_bytes());
465 Ok(())
466 }
467
468 #[inline(always)]
469 fn write_u64(&mut self, val: u64) {
470 if val < 128 {
471 self.out.push(val as u8);
472 } else if val <= u64::from(u8::MAX) {
473 extend(&mut self.out, &[0xcc, val as u8]);
474 } else if val <= u64::from(u16::MAX) {
475 let [a, b] = (val as u16).to_be_bytes();
476 extend(&mut self.out, &[0xcd, a, b]);
477 } else if val <= u64::from(u32::MAX) {
478 let [a, b, c, d] = (val as u32).to_be_bytes();
479 extend(&mut self.out, &[0xce, a, b, c, d]);
480 } else {
481 let mut buf = [0xcf; 9];
482 buf[1..].copy_from_slice(&val.to_be_bytes());
483 extend(&mut self.out, &buf);
484 }
485 }
486
487 #[inline(always)]
488 fn write_i64(&mut self, val: i64) {
489 if val >= 0 {
490 self.write_u64(val as u64);
491 } else if val >= -32 {
492 self.out.push(val as u8);
493 } else if val >= i64::from(i8::MIN) {
494 extend(&mut self.out, &[0xd0, val as u8]);
495 } else if val >= i64::from(i16::MIN) {
496 let [a, b] = (val as i16).to_be_bytes();
497 extend(&mut self.out, &[0xd1, a, b]);
498 } else if val >= i64::from(i32::MIN) {
499 let [a, b, c, d] = (val as i32).to_be_bytes();
500 extend(&mut self.out, &[0xd2, a, b, c, d]);
501 } else {
502 let mut buf = [0xd3; 9];
503 buf[1..].copy_from_slice(&val.to_be_bytes());
504 extend(&mut self.out, &buf);
505 }
506 }
507
508 fn write_ext_data(&mut self, kind: i8, data: &[u8]) -> Result<(), Error> {
510 match data.len() {
511 1 => self.out.push(0xd4),
512 2 => self.out.push(0xd5),
513 4 => self.out.push(0xd6),
514 8 => self.out.push(0xd7),
515 16 => self.out.push(0xd8),
516 len => self.write_len([0xc7, 0xc8, 0xc9], len)?,
517 }
518 self.out.push(kind as u8);
519 extend(&mut self.out, data);
520 Ok(())
521 }
522
523 #[cold]
524 fn write_ext(&mut self, ext: &ExtValue) -> Result<(), Error> {
525 if let Some(val) = ext.downcast_ref::<Ext>() {
526 self.write_ext_data(val.kind, &val.data)
527 } else if let Some(val) = ext.downcast_ref::<Timestamp>() {
528 let mut buf = [0; 12];
529 self.write_ext_data(TIMESTAMP, encode_timestamp(val, &mut buf))
530 } else if let Some(&val) = ext.downcast_ref::<u128>() {
531 let val = u64::try_from(val).map_err(|_| int_out_of_range())?;
532 self.write_u64(val);
533 Ok(())
534 } else if let Some(&val) = ext.downcast_ref::<i128>() {
535 if let Ok(val) = i64::try_from(val) {
536 self.write_i64(val);
537 } else {
538 let val = u64::try_from(val).map_err(|_| int_out_of_range())?;
539 self.write_u64(val);
540 }
541 Ok(())
542 } else if let Some(val) = ext
543 .downcast_ref::<BigInt>()
544 .and_then(|x| x.to_i128())
545 .filter(|&x| i64::try_from(x).is_ok() || u64::try_from(x).is_ok())
546 {
547 match i64::try_from(val) {
549 Ok(val) => self.write_i64(val),
550 Err(_) => self.write_u64(val as u64),
551 }
552 Ok(())
553 } else {
554 match ext.fallback() {
555 Atom::Ext(_) => Err(Error::new(
556 ErrorKind::UnsupportedType,
557 format!("MessagePack does not support {}", ext.name()),
558 )),
559 fallback => self.write_atom(fallback),
560 }
561 }
562 }
563}
564
565#[inline(always)]
567fn check_len(len: usize) -> Result<u32, Error> {
568 u32::try_from(len).map_err(|_| too_long())
569}
570
571#[cold]
572fn too_long() -> Error {
573 Error::new(ErrorKind::OutOfRange, "length out of range for MessagePack")
574}
575
576#[cold]
577fn int_out_of_range() -> Error {
578 Error::new(
579 ErrorKind::OutOfRange,
580 "integer out of range for MessagePack",
581 )
582}
583
584#[inline]
586fn encode_container_head(buf: &mut [u8; 5], is_map: bool, len: u32) -> &[u8] {
587 if len < 16 {
588 buf[0] = if is_map { 0x80 } else { 0x90 } | len as u8;
589 &buf[..1]
590 } else if len <= u32::from(u16::MAX) {
591 buf[0] = if is_map { 0xde } else { 0xdc };
592 buf[1..3].copy_from_slice(&(len as u16).to_be_bytes());
593 &buf[..3]
594 } else {
595 buf[0] = if is_map { 0xdf } else { 0xdd };
596 buf[1..5].copy_from_slice(&len.to_be_bytes());
597 &buf[..5]
598 }
599}
600
601impl SerializerConfig {
602 pub const fn new() -> SerializerConfig {
604 SerializerConfig { canonical: false }
605 }
606
607 pub const fn canonical(mut self, yes: bool) -> SerializerConfig {
614 self.canonical = yes;
615 self
616 }
617
618 pub fn to_vec(&self, value: &dyn Serialize) -> Result<Vec<u8>, Error> {
620 self.to_vec_with(value, |_| {})
621 }
622
623 pub fn to_vec_with<F>(&self, value: &dyn Serialize, setup: F) -> Result<Vec<u8>, Error>
628 where
629 F: FnOnce(&mut SerializeDriver<'_>),
630 {
631 let mut driver = SerializeDriver::new(value);
632 setup(&mut driver);
633 self.serialize_driver(&mut driver)
634 }
635
636 pub(crate) fn serialize_part(
644 &self,
645 item: &mut Option<Box<Writer>>,
646 driver: &mut SerializeDriver<'_>,
647 out: &mut Vec<u8>,
648 limit: usize,
649 ) -> Result<bool, Error> {
650 let len = out.len();
651 if item.is_none() && limit == usize::MAX {
654 let mut writer = Writer::new(self.canonical, core::mem::take(out));
655 let rv = writer.drive(driver, usize::MAX);
656 *out = writer.out;
657 if rv.is_err() {
658 out.truncate(len);
659 }
660 return rv;
661 }
662 let adopt = out.is_empty();
665 let mut writer = item
666 .take()
667 .unwrap_or_else(|| Box::new(Writer::new(self.canonical, Vec::new())));
668 if adopt {
669 writer.out = core::mem::take(out);
670 }
671 let rv = writer.drive(driver, limit);
673 let output = core::mem::take(&mut writer.out);
674 if adopt {
675 *out = output;
676 } else if rv.is_ok() {
677 out.extend_from_slice(&output);
678 }
679 let done = match rv {
680 Ok(done) => done,
681 Err(err) => {
682 out.truncate(len);
683 return Err(err);
684 }
685 };
686 if !done {
687 *item = Some(writer);
688 }
689 Ok(done)
690 }
691
692 pub(crate) fn serialize_driver(
694 &self,
695 driver: &mut SerializeDriver<'_>,
696 ) -> Result<Vec<u8>, Error> {
697 let mut writer = Writer::new(self.canonical, Vec::with_capacity(128));
698 writer.drive_whole(driver)?;
699 Ok(writer.out)
700 }
701}
702
703pub struct Serializer {
725 config: SerializerConfig,
726 out: Vec<u8>,
727 written: usize,
728 item: Option<Box<Writer>>,
730 in_progress: bool,
732}
733
734impl Default for Serializer {
735 fn default() -> Serializer {
736 Serializer::new()
737 }
738}
739
740impl Clone for Serializer {
741 fn clone(&self) -> Serializer {
747 Serializer {
748 config: self.config.clone(),
749 out: self.out.clone(),
750 written: self.written,
751 item: None,
752 in_progress: self.in_progress,
753 }
754 }
755}
756
757impl core::fmt::Debug for Serializer {
758 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
759 f.debug_struct("Serializer")
760 .field("config", &self.config)
761 .field("output", &self.out)
762 .field("written", &self.written)
763 .field("in_progress", &self.in_progress)
764 .finish()
765 }
766}
767
768impl Serializer {
769 pub fn new() -> Serializer {
771 Serializer::with_config(&SerializerConfig::new())
772 }
773
774 pub fn with_config(config: &SerializerConfig) -> Serializer {
776 Serializer {
777 config: config.clone(),
778 out: Vec::new(),
779 written: 0,
780 item: None,
781 in_progress: false,
782 }
783 }
784
785 pub fn config(&self) -> &SerializerConfig {
787 &self.config
788 }
789
790 pub fn written(&self) -> usize {
792 self.written
793 }
794
795 pub fn serialize(&mut self, value: &dyn Serialize) -> Result<(), Error> {
799 ser::Serializer::serialize(self, value)
800 }
801
802 pub fn serialize_with<F>(&mut self, value: &dyn Serialize, setup: F) -> Result<(), Error>
807 where
808 F: FnOnce(&mut SerializeDriver<'_>),
809 {
810 ser::Serializer::serialize_with(self, value, setup)
811 }
812
813 pub fn output(&self) -> &[u8] {
815 &self.out
816 }
817
818 pub fn finish(self) -> Vec<u8> {
820 self.out
821 }
822}
823
824impl ser::Serializer for Serializer {
825 fn drive(&mut self, driver: &mut SerializeDriver<'_>) -> Result<(), Error> {
826 if self.in_progress {
828 return Err(Error::in_progress());
829 }
830 ser::StreamSerializer::drive_partial(self, driver, usize::MAX).map(|_| ())
831 }
832}
833
834impl ser::StreamSerializer for Serializer {
835 fn output(&self) -> &[u8] {
836 &self.out
837 }
838
839 fn clear_output(&mut self) {
840 self.out.clear();
841 }
842
843 fn supports_partial(&self) -> bool {
844 true
845 }
846
847 fn drive_partial(
848 &mut self,
849 driver: &mut SerializeDriver<'_>,
850 limit: usize,
851 ) -> Result<Written, Error> {
852 if self.item.is_none() && self.in_progress {
853 return Err(Error::in_progress());
854 }
855 if !self
858 .config
859 .serialize_part(&mut self.item, driver, &mut self.out, limit)?
860 {
861 self.in_progress = true;
862 return Ok(Written::Partial);
863 }
864 self.in_progress = false;
865 self.written += 1;
866 Ok(Written::Done)
867 }
868
869 fn in_progress(&self) -> bool {
870 self.in_progress
871 }
872}
873
874#[cfg(feature = "io")]
875impl SerializerConfig {
876 pub fn writer<W: std::io::Write>(&self, writer: W) -> deser_core::io::Writer<W, Serializer> {
890 deser_core::io::Writer::new(writer, Serializer::with_config(self))
891 }
892
893 pub fn to_writer<W: std::io::Write>(
897 &self,
898 writer: W,
899 value: &dyn Serialize,
900 ) -> Result<(), Error> {
901 self.writer(writer).write(value)
902 }
903}
904
905#[cfg(feature = "io")]
917pub fn to_writer<W: std::io::Write>(writer: W, value: &dyn Serialize) -> Result<(), Error> {
918 SerializerConfig::new().to_writer(writer, value)
919}
920
921pub fn to_vec(value: &dyn Serialize) -> Result<Vec<u8>, Error> {
925 SerializerConfig::new().to_vec(value)
926}