strict_encoding/
writer.rs

1// Strict encoding library for deterministic binary serialization.
2//
3// SPDX-License-Identifier: Apache-2.0
4//
5// Designed in 2019-2025 by Dr Maxim Orlovsky <orlovsky@ubideco.org>
6// Written in 2024-2025 by Dr Maxim Orlovsky <orlovsky@ubideco.org>
7//
8// Copyright (C) 2022-2025 Laboratories for Ubiquitous Deterministic Computing (UBIDECO),
9//                         Institute for Distributed and Cognitive Systems (InDCS), Switzerland.
10// Copyright (C) 2022-2025 Dr Maxim Orlovsky.
11// All rights under the above copyrights are reserved.
12//
13// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
14// in compliance with the License. You may obtain a copy of the License at
15//
16//        http://www.apache.org/licenses/LICENSE-2.0
17//
18// Unless required by applicable law or agreed to in writing, software distributed under the License
19// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
20// or implied. See the License for the specific language governing permissions and limitations under
21// the License.
22
23use std::collections::{BTreeMap, BTreeSet};
24use std::io;
25use std::io::Sink;
26use std::marker::PhantomData;
27
28use amplify::confinement::U64 as U64MAX;
29use amplify::WriteCounter;
30
31use crate::{
32    DefineEnum, DefineStruct, DefineTuple, DefineUnion, FieldName, LibName, StrictEncode,
33    StrictEnum, StrictStruct, StrictSum, StrictTuple, StrictUnion, TypeName, TypedParent,
34    TypedWrite, Variant, VariantName, WriteEnum, WriteRaw, WriteStruct, WriteTuple, WriteUnion,
35    LIB_EMBEDDED,
36};
37
38// TODO: Move to amplify crate
39#[derive(Clone, Debug)]
40pub struct ConfinedWriter<W: io::Write> {
41    count: usize,
42    limit: usize,
43    writer: W,
44}
45
46impl<W: io::Write> From<W> for ConfinedWriter<W> {
47    fn from(writer: W) -> Self {
48        Self {
49            count: 0,
50            limit: usize::MAX,
51            writer,
52        }
53    }
54}
55
56impl<W: io::Write> ConfinedWriter<W> {
57    pub fn with(limit: usize, writer: W) -> Self {
58        Self {
59            count: 0,
60            limit,
61            writer,
62        }
63    }
64
65    pub fn unconfine(self) -> W { self.writer }
66}
67
68impl<W: io::Write> io::Write for ConfinedWriter<W> {
69    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
70        if self.count + buf.len() > self.limit {
71            return Err(io::ErrorKind::InvalidInput.into());
72        }
73        let count = self.writer.write(buf)?;
74        self.count += count;
75        Ok(count)
76    }
77
78    fn flush(&mut self) -> io::Result<()> { self.writer.flush() }
79}
80
81#[derive(Clone, Debug)]
82pub struct StreamWriter<W: io::Write>(ConfinedWriter<W>);
83
84impl<W: io::Write> StreamWriter<W> {
85    pub fn new<const MAX: usize>(inner: W) -> Self { Self(ConfinedWriter::with(MAX, inner)) }
86    pub fn unconfine(self) -> W { self.0.unconfine() }
87}
88
89impl<W: io::Write> WriteRaw for StreamWriter<W> {
90    fn write_raw<const MAX_LEN: usize>(&mut self, bytes: impl AsRef<[u8]>) -> io::Result<()> {
91        use io::Write;
92        self.0.write_all(bytes.as_ref())?;
93        Ok(())
94    }
95}
96
97impl StreamWriter<Vec<u8>> {
98    pub fn in_memory<const MAX: usize>() -> Self { Self::new::<MAX>(vec![]) }
99}
100
101impl StreamWriter<WriteCounter> {
102    pub fn counter<const MAX: usize>() -> Self { Self::new::<MAX>(WriteCounter::default()) }
103}
104
105impl StreamWriter<Sink> {
106    pub fn sink<const MAX: usize>() -> Self { Self::new::<MAX>(Sink::default()) }
107}
108
109#[derive(Debug, From)]
110pub struct StrictWriter<W: WriteRaw>(W);
111
112impl StrictWriter<StreamWriter<Vec<u8>>> {
113    pub fn in_memory<const MAX: usize>() -> Self { Self(StreamWriter::in_memory::<MAX>()) }
114}
115
116impl StrictWriter<StreamWriter<WriteCounter>> {
117    pub fn counter<const MAX: usize>() -> Self { Self(StreamWriter::counter::<MAX>()) }
118}
119
120impl StrictWriter<StreamWriter<Sink>> {
121    pub fn sink<const MAX: usize>() -> Self { Self(StreamWriter::sink::<MAX>()) }
122}
123
124impl<W: WriteRaw> StrictWriter<W> {
125    pub fn with(writer: W) -> Self { Self(writer) }
126    pub fn unbox(self) -> W { self.0 }
127}
128
129impl<W: WriteRaw> TypedWrite for StrictWriter<W> {
130    type TupleWriter = StructWriter<W, Self>;
131    type StructWriter = StructWriter<W, Self>;
132    type UnionDefiner = UnionWriter<W>;
133    type RawWriter = W;
134
135    unsafe fn raw_writer(&mut self) -> &mut Self::RawWriter { &mut self.0 }
136
137    fn write_union<T: StrictUnion>(
138        self,
139        inner: impl FnOnce(Self::UnionDefiner) -> io::Result<Self>,
140    ) -> io::Result<Self> {
141        let writer = UnionWriter::with::<T>(self);
142        inner(writer)
143    }
144
145    fn write_enum<T: StrictEnum>(self, value: T) -> io::Result<Self>
146    where u8: From<T> {
147        let mut writer = UnionWriter::with::<T>(self);
148        for (_, name) in T::ALL_VARIANTS {
149            writer = writer.define_variant(vname!(*name));
150        }
151        writer = DefineEnum::complete(writer);
152        writer = writer.write_variant(vname!(value.variant_name()))?;
153        Ok(WriteEnum::complete(writer))
154    }
155
156    fn write_tuple<T: StrictTuple>(
157        self,
158        inner: impl FnOnce(Self::TupleWriter) -> io::Result<Self>,
159    ) -> io::Result<Self> {
160        let writer = StructWriter::tuple::<T>(self);
161        inner(writer)
162    }
163
164    fn write_struct<T: StrictStruct>(
165        self,
166        inner: impl FnOnce(Self::StructWriter) -> io::Result<Self>,
167    ) -> io::Result<Self> {
168        let writer = StructWriter::structure::<T>(self);
169        inner(writer)
170    }
171}
172
173#[derive(Debug)]
174pub struct StructWriter<W: WriteRaw, P: StrictParent<W>> {
175    lib: LibName,
176    name: Option<TypeName>,
177    named_fields: Vec<FieldName>,
178    tuple_fields: Option<u8>,
179    parent: P,
180    cursor: usize,
181    _phantom: PhantomData<W>,
182}
183
184impl<W: WriteRaw, P: StrictParent<W>> StructWriter<W, P> {
185    pub fn structure<T: StrictStruct>(parent: P) -> Self {
186        StructWriter {
187            lib: libname!(T::STRICT_LIB_NAME),
188            name: T::strict_name(),
189            named_fields: T::ALL_FIELDS.iter().map(|name| fname!(*name)).collect(),
190            tuple_fields: None,
191            parent,
192            cursor: 0,
193            _phantom: default!(),
194        }
195    }
196
197    pub fn tuple<T: StrictTuple>(parent: P) -> Self {
198        StructWriter {
199            lib: libname!(T::STRICT_LIB_NAME),
200            name: T::strict_name(),
201            named_fields: empty!(),
202            tuple_fields: Some(T::FIELD_COUNT),
203            parent,
204            cursor: 0,
205            _phantom: default!(),
206        }
207    }
208
209    pub fn unnamed(parent: P, tuple: bool) -> Self {
210        StructWriter {
211            lib: libname!(LIB_EMBEDDED),
212            name: None,
213            named_fields: empty!(),
214            tuple_fields: if tuple { Some(0) } else { None },
215            parent,
216            cursor: 0,
217            _phantom: default!(),
218        }
219    }
220
221    pub fn is_tuple(&self) -> bool { self.tuple_fields.is_some() }
222
223    pub fn is_struct(&self) -> bool { !self.is_tuple() }
224
225    pub fn named_fields(&self) -> &[FieldName] {
226        debug_assert!(self.tuple_fields.is_none(), "tuples do not contain named fields");
227        self.named_fields.as_slice()
228    }
229
230    pub fn fields_count(&self) -> u8 { self.tuple_fields.unwrap_or(self.named_fields.len() as u8) }
231
232    pub fn name(&self) -> &str { self.name.as_ref().map(|n| n.as_str()).unwrap_or("<unnamed>") }
233
234    pub fn into_parent(self) -> P { self.parent }
235
236    fn write_value(mut self, value: &impl StrictEncode) -> io::Result<Self> {
237        let (mut writer, remnant) = self.parent.into_write_split();
238        writer = value.strict_encode(writer)?;
239        self.parent = P::from_write_split(writer, remnant);
240        Ok(self)
241    }
242}
243
244impl<W: WriteRaw, P: StrictParent<W>> DefineStruct for StructWriter<W, P> {
245    type Parent = P;
246    fn define_field<T: StrictEncode>(mut self, field: FieldName) -> Self {
247        assert!(
248            !self.named_fields.contains(&field),
249            "field '{:#}' is already defined as a part of '{}'",
250            field,
251            self.name()
252        );
253        self.named_fields.push(field);
254        self
255    }
256    fn complete(self) -> P {
257        assert!(
258            !self.named_fields.is_empty(),
259            "struct '{}' does not have fields defined",
260            self.name()
261        );
262        self.parent
263    }
264}
265
266impl<W: WriteRaw, P: StrictParent<W>> WriteStruct for StructWriter<W, P> {
267    type Parent = P;
268    fn write_field(mut self, _field: FieldName, value: &impl StrictEncode) -> io::Result<Self> {
269        debug_assert!(self.tuple_fields.is_none(), "using struct method on tuple");
270        /* TODO: Propagate information about the fields at the parent
271        debug_assert!(
272            !self.named_fields.is_empty(),
273            "struct without fields {} asks to write value for field '{field}'",
274            self.name()
275        );
276        assert_eq!(
277            &self.named_fields[self.cursor],
278            &field,
279            "field '{:#}' was not defined for '{}' or is written outside of the order",
280            field,
281            self.name()
282        );
283         */
284        self.cursor += 1;
285        self.write_value(value)
286    }
287    fn complete(self) -> P {
288        /* TODO: Propagate information about the fields at the parent
289        assert_eq!(
290            self.cursor,
291            self.named_fields.len(),
292            "not all fields were written for '{}'",
293            self.name()
294        );
295         */
296        self.parent
297    }
298}
299
300impl<W: WriteRaw, P: StrictParent<W>> DefineTuple for StructWriter<W, P> {
301    type Parent = P;
302    fn define_field<T: StrictEncode>(mut self) -> Self {
303        self.tuple_fields
304            .as_mut()
305            .map(|count| *count += 1)
306            .expect("calling tuple method on struct");
307        self
308    }
309    fn complete(self) -> P {
310        assert_ne!(
311            self.tuple_fields.expect("tuple defined as struct"),
312            0,
313            "tuple '{}' does not have fields defined",
314            self.name()
315        );
316        debug_assert!(self.named_fields.is_empty(), "tuple '{}' defined as struct", self.name());
317        self.parent
318    }
319}
320
321impl<W: WriteRaw, P: StrictParent<W>> WriteTuple for StructWriter<W, P> {
322    type Parent = P;
323    fn write_field(mut self, value: &impl StrictEncode) -> io::Result<Self> {
324        /* TODO: Propagate information about number of fields at the parent
325        assert!(
326            self.tuple_fields.expect("writing tuple field to structure") as usize > self.cursor,
327            "writing more unnamed fields to the tuple {} than was defined",
328            self.name()
329        );
330         */
331        self.cursor += 1;
332        self.write_value(value)
333    }
334    fn complete(self) -> P {
335        assert_ne!(self.cursor, 0, "tuple '{}' does not have any fields written", self.name());
336        /* TODO: Propagate information about number of fields at the parent
337        assert_eq!(
338            Some(self.cursor as u8),
339            self.unnamed_fields,
340            "not all fields were written for '{}'",
341            self.name()
342        );
343         */
344        debug_assert!(self.named_fields.is_empty(), "tuple '{}' written as struct", self.name());
345        self.parent
346    }
347}
348
349#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
350pub enum VariantType {
351    Unit,
352    Tuple,
353    Struct,
354}
355
356// TODO: Collect data about defined variant types and check them on write
357#[derive(Debug)]
358pub struct UnionWriter<W: WriteRaw> {
359    lib: LibName,
360    name: Option<TypeName>,
361    declared_variants: BTreeMap<u8, VariantName>,
362    declared_index: BTreeMap<VariantName, u8>,
363    defined_variant: BTreeMap<Variant, VariantType>,
364    parent: StrictWriter<W>,
365    written: bool,
366    parent_ident: Option<TypeName>,
367}
368
369impl UnionWriter<StreamWriter<Sink>> {
370    pub fn sink() -> Self {
371        UnionWriter {
372            lib: libname!(LIB_EMBEDDED),
373            name: None,
374            declared_variants: empty!(),
375            declared_index: empty!(),
376            defined_variant: empty!(),
377            parent: StrictWriter::sink::<U64MAX>(),
378            written: false,
379            parent_ident: None,
380        }
381    }
382}
383
384impl<W: WriteRaw> UnionWriter<W> {
385    pub fn with<T: StrictSum>(parent: StrictWriter<W>) -> Self {
386        UnionWriter {
387            lib: libname!(T::STRICT_LIB_NAME),
388            name: T::strict_name(),
389            declared_variants: T::ALL_VARIANTS
390                .iter()
391                .map(|(tag, name)| (*tag, vname!(*name)))
392                .collect(),
393            declared_index: T::ALL_VARIANTS
394                .iter()
395                .map(|(tag, name)| (vname!(*name), *tag))
396                .collect(),
397            defined_variant: empty!(),
398            parent,
399            written: false,
400            parent_ident: None,
401        }
402    }
403
404    pub fn is_written(&self) -> bool { self.written }
405
406    pub fn variants(&self) -> &BTreeMap<Variant, VariantType> { &self.defined_variant }
407
408    pub fn name(&self) -> &str { self.name.as_ref().map(|n| n.as_str()).unwrap_or("<unnamed>") }
409
410    pub fn tag_by_name(&self, name: &VariantName) -> u8 {
411        *self
412            .declared_index
413            .get(name)
414            .unwrap_or_else(|| panic!("unknown variant `{name}` for the enum `{}`", self.name()))
415    }
416
417    fn _define_variant(mut self, name: VariantName, variant_type: VariantType) -> Self {
418        let tag = self.tag_by_name(&name);
419        let variant = Variant::named(tag, name);
420        assert!(
421            self.defined_variant.insert(variant.clone(), variant_type).is_none(),
422            "variant '{:#}' is already defined as a part of '{}'",
423            &variant,
424            self.name()
425        );
426        self
427    }
428
429    fn _write_variant(mut self, name: VariantName, variant_type: VariantType) -> io::Result<Self> {
430        let (variant, t) =
431            self.defined_variant.iter().find(|(f, _)| f.name == name).unwrap_or_else(|| {
432                panic!("variant '{:#}' was not defined in '{}'", &name, self.name())
433            });
434        assert_eq!(
435            *t,
436            variant_type,
437            "variant '{:#}' in '{}' must be a {:?} while it is written as {:?}",
438            &variant,
439            self.name(),
440            t,
441            variant_type
442        );
443        assert!(!self.written, "multiple attempts to write variants of '{}'", self.name());
444        self.written = true;
445        self.parent = variant.tag.strict_encode(self.parent)?;
446        Ok(self)
447    }
448
449    fn _complete_definition(self) -> Self {
450        let declared = self.declared_variants.values().map(|v| v.as_str()).collect::<BTreeSet<_>>();
451        let defined = self.defined_variant.keys().map(|v| v.name.as_str()).collect::<BTreeSet<_>>();
452        assert_eq!(
453            declared,
454            defined,
455            "unit or enum '{}' hasn't defined all of its declared variants. Elements skipped: {:?}",
456            self.name(),
457            declared.difference(&defined)
458        );
459        assert!(
460            !self.defined_variant.is_empty(),
461            "unit or enum '{}' does not have any fields defined",
462            self.name()
463        );
464        self
465    }
466
467    fn _complete_write(self) -> StrictWriter<W> {
468        assert!(self.written, "not a single variant is written for '{}'", self.name());
469        self.parent
470    }
471}
472
473impl<W: WriteRaw> DefineUnion for UnionWriter<W> {
474    type Parent = StrictWriter<W>;
475    type TupleDefiner = StructWriter<W, Self>;
476    type StructDefiner = StructWriter<W, Self>;
477    type UnionWriter = UnionWriter<W>;
478
479    fn define_unit(self, name: VariantName) -> Self {
480        self._define_variant(name, VariantType::Unit)
481    }
482    fn define_tuple(
483        mut self,
484        name: VariantName,
485        inner: impl FnOnce(Self::TupleDefiner) -> Self,
486    ) -> Self {
487        self = self._define_variant(name, VariantType::Tuple);
488        let definer = StructWriter::unnamed(self, true);
489        inner(definer)
490    }
491    fn define_struct(
492        mut self,
493        name: VariantName,
494        inner: impl FnOnce(Self::StructDefiner) -> Self,
495    ) -> Self {
496        self = self._define_variant(name, VariantType::Struct);
497        let definer = StructWriter::unnamed(self, false);
498        inner(definer)
499    }
500    fn complete(self) -> Self::UnionWriter { self._complete_definition() }
501}
502
503impl<W: WriteRaw> WriteUnion for UnionWriter<W> {
504    type Parent = StrictWriter<W>;
505    type TupleWriter = StructWriter<W, Self>;
506    type StructWriter = StructWriter<W, Self>;
507
508    fn write_unit(self, name: VariantName) -> io::Result<Self> {
509        self._write_variant(name, VariantType::Unit)
510    }
511    fn write_tuple(
512        mut self,
513        name: VariantName,
514        inner: impl FnOnce(Self::TupleWriter) -> io::Result<Self>,
515    ) -> io::Result<Self> {
516        self = self._write_variant(name, VariantType::Tuple)?;
517        let writer = StructWriter::unnamed(self, true);
518        inner(writer)
519    }
520    fn write_struct(
521        mut self,
522        name: VariantName,
523        inner: impl FnOnce(Self::StructWriter) -> io::Result<Self>,
524    ) -> io::Result<Self> {
525        self = self._write_variant(name, VariantType::Struct)?;
526        let writer = StructWriter::unnamed(self, false);
527        inner(writer)
528    }
529    fn complete(self) -> Self::Parent { self._complete_write() }
530}
531
532impl<W: WriteRaw> DefineEnum for UnionWriter<W> {
533    type Parent = StrictWriter<W>;
534    type EnumWriter = UnionWriter<W>;
535    fn define_variant(self, name: VariantName) -> Self {
536        self._define_variant(name, VariantType::Unit)
537    }
538    fn complete(self) -> Self::EnumWriter { self._complete_definition() }
539}
540
541impl<W: WriteRaw> WriteEnum for UnionWriter<W> {
542    type Parent = StrictWriter<W>;
543    fn write_variant(self, name: VariantName) -> io::Result<Self> {
544        self._write_variant(name, VariantType::Unit)
545    }
546    fn complete(self) -> Self::Parent { self._complete_write() }
547}
548
549pub trait StrictParent<W: WriteRaw>: TypedParent {
550    type Remnant;
551    fn from_write_split(writer: StrictWriter<W>, remnant: Self::Remnant) -> Self;
552    fn into_write_split(self) -> (StrictWriter<W>, Self::Remnant);
553}
554impl<W: WriteRaw> TypedParent for StrictWriter<W> {}
555impl<W: WriteRaw> TypedParent for UnionWriter<W> {}
556impl<W: WriteRaw> StrictParent<W> for StrictWriter<W> {
557    type Remnant = ();
558    fn from_write_split(writer: StrictWriter<W>, _: Self::Remnant) -> Self { writer }
559    fn into_write_split(self) -> (StrictWriter<W>, Self::Remnant) { (self, ()) }
560}
561impl<W: WriteRaw> StrictParent<W> for UnionWriter<W> {
562    type Remnant = UnionWriter<StreamWriter<Sink>>;
563    fn from_write_split(writer: StrictWriter<W>, remnant: Self::Remnant) -> Self {
564        Self {
565            lib: remnant.lib,
566            name: remnant.name,
567            declared_variants: remnant.declared_variants,
568            declared_index: remnant.declared_index,
569            defined_variant: remnant.defined_variant,
570            parent: writer,
571            written: remnant.written,
572            parent_ident: remnant.parent_ident,
573        }
574    }
575    fn into_write_split(self) -> (StrictWriter<W>, Self::Remnant) {
576        let remnant = UnionWriter {
577            lib: self.lib,
578            name: self.name,
579            declared_variants: self.declared_variants,
580            declared_index: self.declared_index,
581            defined_variant: self.defined_variant,
582            parent: StrictWriter::sink::<U64MAX>(),
583            written: self.written,
584            parent_ident: self.parent_ident,
585        };
586        (self.parent, remnant)
587    }
588}
589
590pub trait SplitParent {
591    type Parent: TypedParent;
592    type Remnant;
593    fn from_parent_split(parent: Self::Parent, remnant: Self::Remnant) -> Self;
594    fn into_parent_split(self) -> (Self::Parent, Self::Remnant);
595}
596impl<W: WriteRaw, P: StrictParent<W>> SplitParent for StructWriter<W, P> {
597    type Parent = P;
598    type Remnant = StructWriter<StreamWriter<Sink>, ParentDumb>;
599    fn from_parent_split(parent: P, remnant: Self::Remnant) -> Self {
600        Self {
601            lib: remnant.lib,
602            name: remnant.name,
603            named_fields: remnant.named_fields,
604            tuple_fields: remnant.tuple_fields,
605            parent,
606            cursor: remnant.cursor,
607            _phantom: none!(),
608        }
609    }
610    fn into_parent_split(self) -> (P, Self::Remnant) {
611        let remnant = StructWriter::<StreamWriter<Sink>, ParentDumb> {
612            lib: self.lib,
613            name: self.name,
614            named_fields: self.named_fields,
615            tuple_fields: self.tuple_fields,
616            parent: none!(),
617            cursor: self.cursor,
618            _phantom: none!(),
619        };
620        (self.parent, remnant)
621    }
622}
623
624#[derive(Default)]
625pub struct ParentDumb;
626impl TypedParent for ParentDumb {}
627impl<W: WriteRaw> StrictParent<W> for ParentDumb {
628    type Remnant = ();
629    fn from_write_split(_: StrictWriter<W>, _: Self::Remnant) -> Self { unreachable!() }
630    fn into_write_split(self) -> (StrictWriter<W>, Self::Remnant) { unreachable!() }
631}
strict_encoding/writer.rs

strict_encoding/
writer.rs
