maxminddb/
encoder.rs

1use std::{collections::BTreeMap, fmt};
2
3use serde::{
4    de::{self, Deserialize, Deserializer, MapAccess, SeqAccess, Visitor},
5    ser::{self, Serialize, SerializeMap, SerializeSeq, SerializeStruct, SerializeTupleStruct},
6};
7
8pub type EncodeResult<T> = Result<T, String>;
9
10#[derive(Debug)]
11pub struct SerError(pub String);
12
13impl fmt::Display for SerError {
14    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
15        f.write_str(&self.0)
16    }
17}
18
19impl std::error::Error for SerError {}
20
21impl ser::Error for SerError {
22    fn custom<T: fmt::Display>(msg: T) -> Self {
23        SerError(msg.to_string())
24    }
25}
26
27const FIRST_SIZE: usize = 29;
28const SECOND_SIZE: usize = FIRST_SIZE + 256;
29const THIRD_SIZE: usize = SECOND_SIZE + (1 << 16);
30
31#[derive(Debug, Clone, PartialEq)]
32pub enum Value {
33    Pointer(u32),
34    String(String),
35    Float64(f64),
36    Bytes(Vec<u8>),
37    Uint16(u16),
38    Uint32(u32),
39    Int32(i32),
40    Map(BTreeMap<String, Value>),
41    Uint64(u64),
42    Uint128(u128),
43    Slice(Vec<Value>),
44    Bool(bool),
45    Float32(f32),
46}
47
48impl From<&str> for Value {
49    fn from(s: &str) -> Self {
50        Value::String(s.to_string())
51    }
52}
53
54impl From<String> for Value {
55    fn from(s: String) -> Self {
56        Value::String(s)
57    }
58}
59
60impl From<Vec<u8>> for Value {
61    fn from(v: Vec<u8>) -> Self {
62        Value::Bytes(v)
63    }
64}
65
66impl From<u16> for Value {
67    fn from(v: u16) -> Self {
68        Value::Uint16(v)
69    }
70}
71
72impl From<u32> for Value {
73    fn from(v: u32) -> Self {
74        Value::Uint32(v)
75    }
76}
77
78impl From<i32> for Value {
79    fn from(v: i32) -> Self {
80        Value::Int32(v)
81    }
82}
83
84impl From<u64> for Value {
85    fn from(v: u64) -> Self {
86        Value::Uint64(v)
87    }
88}
89
90impl From<u128> for Value {
91    fn from(v: u128) -> Self {
92        Value::Uint128(v)
93    }
94}
95
96impl From<f64> for Value {
97    fn from(v: f64) -> Self {
98        Value::Float64(v)
99    }
100}
101
102impl From<f32> for Value {
103    fn from(v: f32) -> Self {
104        Value::Float32(v)
105    }
106}
107
108impl From<bool> for Value {
109    fn from(v: bool) -> Self {
110        Value::Bool(v)
111    }
112}
113
114#[macro_export]
115macro_rules! map {
116    ($($key:expr => $val:expr),* $(,)?) => {{
117        let mut __map = ::std::collections::BTreeMap::new();
118        $(
119            __map.insert(
120                ($key).to_string(),
121                ::std::convert::Into::<$crate::encoder::Value>::into($val),
122            );
123        )*
124        $crate::encoder::Value::Map(__map)
125    }};
126}
127
128impl<'de> Deserialize<'de> for Value {
129    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
130        struct ValueVisitor;
131
132        impl<'de> Visitor<'de> for ValueVisitor {
133            type Value = Value;
134
135            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136                f.write_str("a MaxMind DB value")
137            }
138
139            fn visit_bool<E: de::Error>(self, v: bool) -> Result<Value, E> {
140                Ok(Value::Bool(v))
141            }
142
143            fn visit_i64<E: de::Error>(self, v: i64) -> Result<Value, E> {
144                if let Ok(v) = i32::try_from(v) {
145                    Ok(Value::Int32(v))
146                } else {
147                    Ok(Value::Uint64(v as u64))
148                }
149            }
150
151            fn visit_u64<E: de::Error>(self, v: u64) -> Result<Value, E> {
152                Ok(Value::Uint64(v))
153            }
154
155            fn visit_f64<E: de::Error>(self, v: f64) -> Result<Value, E> {
156                Ok(Value::Float64(v))
157            }
158
159            fn visit_str<E: de::Error>(self, v: &str) -> Result<Value, E> {
160                Ok(Value::String(v.to_string()))
161            }
162
163            fn visit_string<E: de::Error>(self, v: String) -> Result<Value, E> {
164                Ok(Value::String(v))
165            }
166
167            fn visit_bytes<E: de::Error>(self, v: &[u8]) -> Result<Value, E> {
168                Ok(Value::Bytes(v.to_vec()))
169            }
170
171            fn visit_byte_buf<E: de::Error>(self, v: Vec<u8>) -> Result<Value, E> {
172                Ok(Value::Bytes(v))
173            }
174
175            fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Value, A::Error> {
176                let mut items = Vec::with_capacity(seq.size_hint().unwrap_or(1));
177                while let Some(elem) = seq.next_element::<Value>()? {
178                    items.push(elem);
179                }
180                Ok(Value::Slice(items))
181            }
182
183            fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Value, A::Error> {
184                let mut entries = BTreeMap::new();
185                while let Some((key, val)) = map.next_entry::<String, Value>()? {
186                    entries.insert(key, val);
187                }
188                Ok(Value::Map(entries))
189            }
190        }
191
192        deserializer.deserialize_any(ValueVisitor)
193    }
194}
195
196// ---------------------------------------------------------------------------
197// ValueSerializer — converts Serialize into owned Value (used by insert_value)
198// ---------------------------------------------------------------------------
199
200pub struct ValueSerializer;
201
202impl<'a> ser::Serializer for &'a ValueSerializer {
203    type Ok = Value;
204    type Error = SerError;
205
206    type SerializeSeq = OwnedMapOrSlice;
207    type SerializeTuple = OwnedMapOrSlice;
208    type SerializeTupleStruct = OwnedMapOrSlice;
209    type SerializeTupleVariant = ser::Impossible<Value, SerError>;
210    type SerializeMap = OwnedMapOrSlice;
211    type SerializeStruct = OwnedMapOrSlice;
212    type SerializeStructVariant = ser::Impossible<Value, SerError>;
213
214    fn serialize_bool(self, v: bool) -> Result<Value, SerError> {
215        Ok(Value::Bool(v))
216    }
217
218    fn serialize_i8(self, v: i8) -> Result<Value, SerError> {
219        Ok(Value::Int32(v as i32))
220    }
221
222    fn serialize_i16(self, v: i16) -> Result<Value, SerError> {
223        Ok(Value::Int32(v as i32))
224    }
225
226    fn serialize_i32(self, v: i32) -> Result<Value, SerError> {
227        Ok(Value::Int32(v))
228    }
229
230    fn serialize_i64(self, v: i64) -> Result<Value, SerError> {
231        i32::try_from(v)
232            .map(Value::Int32)
233            .or_else(|_| Ok(Value::Uint64(v as u64)))
234            .map_err(SerError)
235    }
236
237    fn serialize_u8(self, v: u8) -> Result<Value, SerError> {
238        Ok(Value::Uint16(v as u16))
239    }
240
241    fn serialize_u16(self, v: u16) -> Result<Value, SerError> {
242        Ok(Value::Uint16(v))
243    }
244
245    fn serialize_u32(self, v: u32) -> Result<Value, SerError> {
246        Ok(Value::Uint32(v))
247    }
248
249    fn serialize_u64(self, v: u64) -> Result<Value, SerError> {
250        Ok(Value::Uint64(v))
251    }
252
253    fn serialize_u128(self, v: u128) -> Result<Value, SerError> {
254        Ok(Value::Uint128(v))
255    }
256
257    fn serialize_f32(self, v: f32) -> Result<Value, SerError> {
258        Ok(Value::Float32(v))
259    }
260
261    fn serialize_f64(self, v: f64) -> Result<Value, SerError> {
262        Ok(Value::Float64(v))
263    }
264
265    fn serialize_char(self, v: char) -> Result<Value, SerError> {
266        Ok(Value::String(v.to_string()))
267    }
268
269    fn serialize_str(self, v: &str) -> Result<Value, SerError> {
270        Ok(Value::String(v.to_string()))
271    }
272
273    fn serialize_bytes(self, v: &[u8]) -> Result<Value, SerError> {
274        Ok(Value::Bytes(v.to_vec()))
275    }
276
277    fn serialize_none(self) -> Result<Value, SerError> {
278        Err(SerError("MaxMind DB has no null type".into()))
279    }
280
281    fn serialize_some<T: ?Sized + Serialize>(self, value: &T) -> Result<Value, SerError> {
282        value.serialize(self)
283    }
284
285    fn serialize_unit(self) -> Result<Value, SerError> {
286        Err(SerError("MaxMind DB has no null type".into()))
287    }
288
289    fn serialize_unit_struct(self, _name: &'static str) -> Result<Value, SerError> {
290        Err(SerError("MaxMind DB has no null type".into()))
291    }
292
293    fn serialize_unit_variant(
294        self,
295        _name: &'static str,
296        _variant_index: u32,
297        variant: &'static str,
298    ) -> Result<Value, SerError> {
299        Ok(Value::String(variant.to_string()))
300    }
301
302    fn serialize_newtype_struct<T: ?Sized + Serialize>(
303        self,
304        _name: &'static str,
305        value: &T,
306    ) -> Result<Value, SerError> {
307        value.serialize(self)
308    }
309
310    fn serialize_newtype_variant<T: ?Sized + Serialize>(
311        self,
312        _name: &'static str,
313        _variant_index: u32,
314        _variant: &'static str,
315        value: &T,
316    ) -> Result<Value, SerError> {
317        value.serialize(self)
318    }
319
320    fn serialize_seq(self, _len: Option<usize>) -> Result<OwnedMapOrSlice, SerError> {
321        Ok(OwnedMapOrSlice::slice())
322    }
323
324    fn serialize_tuple(self, len: usize) -> Result<OwnedMapOrSlice, SerError> {
325        self.serialize_seq(Some(len))
326    }
327
328    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<OwnedMapOrSlice, SerError> {
329        self.serialize_seq(Some(len))
330    }
331
332    fn serialize_map(self, _len: Option<usize>) -> Result<OwnedMapOrSlice, SerError> {
333        Ok(OwnedMapOrSlice::map())
334    }
335
336    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<OwnedMapOrSlice, SerError> {
337        self.serialize_map(Some(len))
338    }
339
340    fn serialize_tuple_variant(
341        self,
342        _name: &'static str,
343        _variant_index: u32,
344        _variant: &'static str,
345        _len: usize,
346    ) -> Result<Self::SerializeTupleVariant, SerError> {
347        Err(SerError("tuple variants not supported".into()))
348    }
349
350    fn serialize_struct_variant(
351        self,
352        _name: &'static str,
353        _variant_index: u32,
354        _variant: &'static str,
355        _len: usize,
356    ) -> Result<Self::SerializeStructVariant, SerError> {
357        Err(SerError("struct variants not supported".into()))
358    }
359}
360
361pub struct OwnedMapOrSlice {
362    pending_key: Option<String>,
363    inner: OwnedMapOrSliceInner,
364}
365
366enum OwnedMapOrSliceInner {
367    Map(BTreeMap<String, Value>),
368    Slice(Vec<Value>),
369}
370
371impl OwnedMapOrSlice {
372    fn map() -> Self {
373        OwnedMapOrSlice {
374            pending_key: None,
375            inner: OwnedMapOrSliceInner::Map(BTreeMap::new()),
376        }
377    }
378    fn slice() -> Self {
379        OwnedMapOrSlice {
380            pending_key: None,
381            inner: OwnedMapOrSliceInner::Slice(Vec::new()),
382        }
383    }
384}
385
386impl SerializeSeq for OwnedMapOrSlice {
387    type Ok = Value;
388    type Error = SerError;
389
390    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
391        match self.inner {
392            OwnedMapOrSliceInner::Slice(ref mut v) => {
393                v.push(value.serialize(&ValueSerializer)?);
394                Ok(())
395            }
396            OwnedMapOrSliceInner::Map(_) => Err(SerError("expected map entry, got element".into())),
397        }
398    }
399
400    fn end(self) -> Result<Value, SerError> {
401        match self.inner {
402            OwnedMapOrSliceInner::Slice(v) => Ok(Value::Slice(v)),
403            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
404        }
405    }
406}
407
408impl ser::SerializeTuple for OwnedMapOrSlice {
409    type Ok = Value;
410    type Error = SerError;
411
412    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
413        SerializeSeq::serialize_element(self, value)
414    }
415
416    fn end(self) -> Result<Value, SerError> {
417        SerializeSeq::end(self)
418    }
419}
420
421impl SerializeTupleStruct for OwnedMapOrSlice {
422    type Ok = Value;
423    type Error = SerError;
424
425    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
426        SerializeSeq::serialize_element(self, value)
427    }
428
429    fn end(self) -> Result<Value, SerError> {
430        SerializeSeq::end(self)
431    }
432}
433
434impl SerializeMap for OwnedMapOrSlice {
435    type Ok = Value;
436    type Error = SerError;
437
438    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), SerError> {
439        match self.inner {
440            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected element, got map key".into())),
441            OwnedMapOrSliceInner::Map(_) => {
442                let k = key.serialize(&ValueSerializer)?;
443                match k {
444                    Value::String(s) => {
445                        self.pending_key = Some(s);
446                        Ok(())
447                    }
448                    _ => Err(SerError("map key must be a string".into())),
449                }
450            }
451        }
452    }
453
454    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
455        match self.inner {
456            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected element, got map value".into())),
457            OwnedMapOrSliceInner::Map(ref mut m) => {
458                let key = self
459                    .pending_key
460                    .take()
461                    .ok_or_else(|| SerError("missing map key".into()))?;
462                let v = value.serialize(&ValueSerializer)?;
463                m.insert(key, v);
464                Ok(())
465            }
466        }
467    }
468
469    fn end(self) -> Result<Value, SerError> {
470        match self.inner {
471            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
472            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got slice".into())),
473        }
474    }
475}
476
477impl SerializeStruct for OwnedMapOrSlice {
478    type Ok = Value;
479    type Error = SerError;
480
481    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) -> Result<(), SerError> {
482        match self.inner {
483            OwnedMapOrSliceInner::Map(ref mut m) => {
484                let v = value.serialize(&ValueSerializer)?;
485                m.insert(key.to_string(), v);
486                Ok(())
487            }
488            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got struct".into())),
489        }
490    }
491
492    fn end(self) -> Result<Value, SerError> {
493        match self.inner {
494            OwnedMapOrSliceInner::Map(m) => Ok(Value::Map(m)),
495            OwnedMapOrSliceInner::Slice(_) => Err(SerError("expected map, got slice".into())),
496        }
497    }
498}
499
500// ---------------------------------------------------------------------------
501// ByteSerializer — writes MaxMind DB binary directly into a Vec<u8>,
502//                  avoiding the intermediate Value allocation.
503// ---------------------------------------------------------------------------
504
505pub struct ByteSerializer {
506    buf: Vec<u8>,
507}
508
509impl ByteSerializer {
510    pub fn new() -> Self {
511        ByteSerializer {
512            buf: Vec::new(),
513        }
514    }
515
516    pub fn into_bytes(self) -> Vec<u8> {
517        self.buf
518    }
519}
520
521fn ctrl_and_size(w: &mut Vec<u8>, type_num: u8, payload_len: usize) {
522    let first_byte: u8;
523    let second_byte: u8;
524
525    if type_num < 8 {
526        first_byte = type_num << 5;
527        second_byte = 0;
528    } else {
529        first_byte = 0;
530        second_byte = type_num - 7;
531    }
532
533    let size_val = match payload_len {
534        s if s < FIRST_SIZE => s as u8,
535        s if s < SECOND_SIZE => 29u8,
536        s if s < THIRD_SIZE => 30u8,
537        _ => 31u8,
538    };
539    w.push(first_byte | size_val);
540
541    if second_byte != 0 {
542        w.push(second_byte);
543    }
544
545    match payload_len {
546        s if s < FIRST_SIZE => {}
547        s if s < SECOND_SIZE => {
548            w.push((s - FIRST_SIZE) as u8);
549        }
550        s if s < THIRD_SIZE => {
551            let v = s - SECOND_SIZE;
552            w.push((v >> 8) as u8);
553            w.push((v & 0xFF) as u8);
554        }
555        s => {
556            let v = s - THIRD_SIZE;
557            w.push((v >> 16) as u8);
558            w.push((v >> 8) as u8);
559            w.push((v & 0xFF) as u8);
560        }
561    }
562}
563
564struct ByteMapOrSlice<'a> {
565    ser: &'a mut ByteSerializer,
566    state: ByteMOS,
567    pending_key: Option<String>,
568}
569
570enum ByteMOS {
571    Map(BTreeMap<String, Vec<u8>>),
572    Slice(Vec<Vec<u8>>),
573}
574
575impl ByteMapOrSlice<'_> {
576    fn flush_map(self) -> Result<(), SerError> {
577        match self.state {
578            ByteMOS::Map(m) => {
579                ctrl_and_size(&mut self.ser.buf, 7, m.len());
580                for (k, v) in m {
581                    write_str_raw(&mut self.ser.buf, &k);
582                    self.ser.buf.extend_from_slice(&v);
583                }
584                Ok(())
585            }
586            ByteMOS::Slice(s) => {
587                ctrl_and_size(&mut self.ser.buf, 11, s.len());
588                for v in s {
589                    self.ser.buf.extend_from_slice(&v);
590                }
591                Ok(())
592            }
593        }
594    }
595}
596
597fn write_str_raw(w: &mut Vec<u8>, s: &str) {
598    ctrl_and_size(w, 2, s.len());
599    w.extend_from_slice(s.as_bytes());
600}
601
602fn write_bytes_raw(w: &mut Vec<u8>, b: &[u8]) {
603    ctrl_and_size(w, 4, b.len());
604    w.extend_from_slice(b);
605}
606
607impl SerializeSeq for ByteMapOrSlice<'_> {
608    type Ok = ();
609    type Error = SerError;
610
611    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
612        let mut sub = ByteSerializer::new();
613        value.serialize(ByteSerializerRef(&mut sub))?;
614        match self.state {
615            ByteMOS::Slice(ref mut v) => v.push(sub.buf),
616            ByteMOS::Map(_) => return Err(SerError("expected map entry, got element".into())),
617        }
618        Ok(())
619    }
620
621    fn end(self) -> Result<(), SerError> {
622        self.flush_map()
623    }
624}
625
626impl ser::SerializeTuple for ByteMapOrSlice<'_> {
627    type Ok = ();
628    type Error = SerError;
629
630    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
631        SerializeSeq::serialize_element(self, value)
632    }
633
634    fn end(self) -> Result<(), SerError> {
635        SerializeSeq::end(self)
636    }
637}
638
639impl SerializeTupleStruct for ByteMapOrSlice<'_> {
640    type Ok = ();
641    type Error = SerError;
642
643    fn serialize_field<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
644        SerializeSeq::serialize_element(self, value)
645    }
646
647    fn end(self) -> Result<(), SerError> {
648        SerializeSeq::end(self)
649    }
650}
651
652impl SerializeMap for ByteMapOrSlice<'_> {
653    type Ok = ();
654    type Error = SerError;
655
656    fn serialize_key<T: ?Sized + Serialize>(&mut self, key: &T) -> Result<(), SerError> {
657        match self.state {
658            ByteMOS::Slice(_) => Err(SerError("expected element, got map key".into())),
659            ByteMOS::Map(_) => {
660                let mut sub = ByteSerializer::new();
661                key.serialize(ByteSerializerRef(&mut sub))?;
662                // let raw = sub.buf;
663                // let s = String::new();
664                // decode the key from the encoded bytes
665                // simplest: serialize through ValueSerializer then extract
666                let k = key.serialize(&ValueSerializer)?;
667                match k {
668                    Value::String(s_val) => {
669                        self.pending_key = Some(s_val);
670                        Ok(())
671                    }
672                    _ => Err(SerError("map key must be a string".into())),
673                }
674            }
675        }
676    }
677
678    fn serialize_value<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<(), SerError> {
679        match self.state {
680            ByteMOS::Slice(_) => Err(SerError("expected element, got map value".into())),
681            ByteMOS::Map(ref mut m) => {
682                let key = self
683                    .pending_key
684                    .take()
685                    .ok_or_else(|| SerError("missing map key".into()))?;
686                let mut sub = ByteSerializer::new();
687                value.serialize(ByteSerializerRef(&mut sub))?;
688                m.insert(key, sub.buf);
689                Ok(())
690            }
691        }
692    }
693
694    fn end(self) -> Result<(), SerError> {
695        self.flush_map()
696    }
697}
698
699impl SerializeStruct for ByteMapOrSlice<'_> {
700    type Ok = ();
701    type Error = SerError;
702
703    fn serialize_field<T: ?Sized + Serialize>(&mut self, key: &'static str, value: &T) -> Result<(), SerError> {
704        match self.state {
705            ByteMOS::Map(ref mut m) => {
706                let mut sub = ByteSerializer::new();
707                value.serialize(ByteSerializerRef(&mut sub))?;
708                m.insert(key.to_string(), sub.buf);
709                Ok(())
710            }
711            ByteMOS::Slice(_) => Err(SerError("expected map, got struct".into())),
712        }
713    }
714
715    fn end(self) -> Result<(), SerError> {
716        self.flush_map()
717    }
718}
719
720struct ByteSerializerRef<'a>(&'a mut ByteSerializer);
721
722impl<'a> ser::Serializer for ByteSerializerRef<'a> {
723    type Ok = ();
724    type Error = SerError;
725
726    type SerializeSeq = ByteMapOrSlice<'a>;
727    type SerializeTuple = ByteMapOrSlice<'a>;
728    type SerializeTupleStruct = ByteMapOrSlice<'a>;
729    type SerializeTupleVariant = ser::Impossible<(), SerError>;
730    type SerializeMap = ByteMapOrSlice<'a>;
731    type SerializeStruct = ByteMapOrSlice<'a>;
732    type SerializeStructVariant = ser::Impossible<(), SerError>;
733
734    fn serialize_bool(self, v: bool) -> Result<(), SerError> {
735        let payload_len = if v { 1 } else { 0 };
736        ctrl_and_size(&mut self.0.buf, 14, payload_len);
737        Ok(())
738    }
739
740    fn serialize_i8(self, v: i8) -> Result<(), SerError> {
741        self.serialize_i32(v as i32)
742    }
743
744    fn serialize_i16(self, v: i16) -> Result<(), SerError> {
745        self.serialize_i32(v as i32)
746    }
747
748    fn serialize_i32(self, v: i32) -> Result<(), SerError> {
749        let u = v as u32;
750        let size = 4 - (u.leading_zeros() as usize / 8);
751        ctrl_and_size(&mut self.0.buf, 8, size);
752        for i in (0..size).rev() {
753            self.0.buf.push((u >> (8 * i)) as u8);
754        }
755        Ok(())
756    }
757
758    fn serialize_i64(self, v: i64) -> Result<(), SerError> {
759        if let Ok(v) = i32::try_from(v) {
760            self.serialize_i32(v)
761        } else {
762            self.serialize_u64(v as u64)
763        }
764    }
765
766    fn serialize_u8(self, v: u8) -> Result<(), SerError> {
767        self.serialize_u16(v as u16)
768    }
769
770    fn serialize_u16(self, v: u16) -> Result<(), SerError> {
771        let size = 2 - (v.leading_zeros() as usize / 8);
772        ctrl_and_size(&mut self.0.buf, 5, size);
773        for i in (0..size).rev() {
774            self.0.buf.push((v >> (8 * i)) as u8);
775        }
776        Ok(())
777    }
778
779    fn serialize_u32(self, v: u32) -> Result<(), SerError> {
780        let size = 4 - (v.leading_zeros() as usize / 8);
781        ctrl_and_size(&mut self.0.buf, 6, size);
782        for i in (0..size).rev() {
783            self.0.buf.push((v >> (8 * i)) as u8);
784        }
785        Ok(())
786    }
787
788    fn serialize_u64(self, v: u64) -> Result<(), SerError> {
789        let size = 8 - (v.leading_zeros() as usize / 8);
790        ctrl_and_size(&mut self.0.buf, 9, size);
791        for i in (0..size).rev() {
792            self.0.buf.push((v >> (8 * i)) as u8);
793        }
794        Ok(())
795    }
796
797    fn serialize_u128(self, v: u128) -> Result<(), SerError> {
798        let size = 16 - (v.leading_zeros() as usize / 8);
799        ctrl_and_size(&mut self.0.buf, 10, size);
800        for i in (0..size).rev() {
801            self.0.buf.push((v >> (8 * i)) as u8);
802        }
803        Ok(())
804    }
805
806    fn serialize_f32(self, v: f32) -> Result<(), SerError> {
807        ctrl_and_size(&mut self.0.buf, 15, 4);
808        self.0.buf.extend_from_slice(&v.to_bits().to_be_bytes());
809        Ok(())
810    }
811
812    fn serialize_f64(self, v: f64) -> Result<(), SerError> {
813        ctrl_and_size(&mut self.0.buf, 3, 8);
814        self.0.buf.extend_from_slice(&v.to_bits().to_be_bytes());
815        Ok(())
816    }
817
818    fn serialize_char(self, v: char) -> Result<(), SerError> {
819        self.serialize_str(&v.to_string())
820    }
821
822    fn serialize_str(self, v: &str) -> Result<(), SerError> {
823        write_str_raw(&mut self.0.buf, v);
824        Ok(())
825    }
826
827    fn serialize_bytes(self, v: &[u8]) -> Result<(), SerError> {
828        write_bytes_raw(&mut self.0.buf, v);
829        Ok(())
830    }
831
832    fn serialize_none(self) -> Result<(), SerError> {
833        Err(SerError("MaxMind DB has no null type".into()))
834    }
835
836    fn serialize_some<T: ?Sized + Serialize>(self, value: &T) -> Result<(), SerError> {
837        value.serialize(self)
838    }
839
840    fn serialize_unit(self) -> Result<(), SerError> {
841        Err(SerError("MaxMind DB has no null type".into()))
842    }
843
844    fn serialize_unit_struct(self, _name: &'static str) -> Result<(), SerError> {
845        Err(SerError("MaxMind DB has no null type".into()))
846    }
847
848    fn serialize_unit_variant(
849        self,
850        _name: &'static str,
851        _variant_index: u32,
852        variant: &'static str,
853    ) -> Result<(), SerError> {
854        self.serialize_str(variant)
855    }
856
857    fn serialize_newtype_struct<T: ?Sized + Serialize>(self, _name: &'static str, value: &T) -> Result<(), SerError> {
858        value.serialize(self)
859    }
860
861    fn serialize_newtype_variant<T: ?Sized + Serialize>(
862        self,
863        _name: &'static str,
864        _variant_index: u32,
865        _variant: &'static str,
866        value: &T,
867    ) -> Result<(), SerError> {
868        value.serialize(self)
869    }
870
871    fn serialize_seq(self, _len: Option<usize>) -> Result<ByteMapOrSlice<'a>, SerError> {
872        Ok(ByteMapOrSlice {
873            ser: self.0,
874            state: ByteMOS::Slice(Vec::new()),
875            pending_key: None,
876        })
877    }
878
879    fn serialize_tuple(self, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
880        self.serialize_seq(Some(len))
881    }
882
883    fn serialize_tuple_struct(self, _name: &'static str, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
884        self.serialize_seq(Some(len))
885    }
886
887    fn serialize_map(self, _len: Option<usize>) -> Result<ByteMapOrSlice<'a>, SerError> {
888        Ok(ByteMapOrSlice {
889            ser: self.0,
890            state: ByteMOS::Map(BTreeMap::new()),
891            pending_key: None,
892        })
893    }
894
895    fn serialize_struct(self, _name: &'static str, len: usize) -> Result<ByteMapOrSlice<'a>, SerError> {
896        self.serialize_map(Some(len))
897    }
898
899    fn serialize_tuple_variant(
900        self,
901        _name: &'static str,
902        _variant_index: u32,
903        _variant: &'static str,
904        _len: usize,
905    ) -> Result<Self::SerializeTupleVariant, SerError> {
906        Err(SerError("tuple variants not supported".into()))
907    }
908
909    fn serialize_struct_variant(
910        self,
911        _name: &'static str,
912        _variant_index: u32,
913        _variant: &'static str,
914        _len: usize,
915    ) -> Result<Self::SerializeStructVariant, SerError> {
916        Err(SerError("struct variants not supported".into()))
917    }
918}
919
920/// Serialize a `Serialize` value directly into MaxMind DB binary format.
921/// Returns the encoded bytes. This avoids allocating an intermediate `Value`.
922pub fn encode_serialize(v: &impl Serialize) -> Result<Vec<u8>, SerError> {
923    let mut ser = ByteSerializer::new();
924    v.serialize(ByteSerializerRef(&mut ser))?;
925    Ok(ser.buf)
926}
927
928fn encode_pointer_raw(w: &mut Vec<u8>, pointer: u32) {
929    match pointer {
930        p if p < 2048 => {
931            w.push(0b00100000 | ((p >> 8) as u8 & 0x07));
932            w.push((p & 0xFF) as u8);
933        }
934        p if p < 526336 => {
935            let v = p - 2048;
936            w.push(0b00101000 | ((v >> 16) as u8 & 0x07));
937            w.push((v >> 8) as u8);
938            w.push((v & 0xFF) as u8);
939        }
940        p if p < 134744064 => {
941            let v = p - 526336;
942            w.push(0b00110000 | ((v >> 24) as u8 & 0x07));
943            w.push((v >> 16) as u8);
944            w.push((v >> 8) as u8);
945            w.push((v & 0xFF) as u8);
946        }
947        p => {
948            w.push(0b00111000);
949            w.push((p >> 24) as u8);
950            w.push((p >> 16) as u8);
951            w.push((p >> 8) as u8);
952            w.push((p & 0xFF) as u8);
953        }
954    }
955}
956
957impl ByteSerializerRef<'_> {
958    fn serialize_value(self, v: &Value) -> Result<(), SerError> {
959        match v {
960            Value::Pointer(p) => {
961                encode_pointer_raw(&mut self.0.buf, *p);
962                Ok(())
963            }
964            Value::String(s) => {
965                write_str_raw(&mut self.0.buf, s);
966                Ok(())
967            }
968            Value::Float64(f) => {
969                ctrl_and_size(&mut self.0.buf, 3, 8);
970                self.0.buf.extend_from_slice(&f.to_bits().to_be_bytes());
971                Ok(())
972            }
973            Value::Bytes(b) => {
974                write_bytes_raw(&mut self.0.buf, b);
975                Ok(())
976            }
977            Value::Uint16(n) => {
978                let size = 2 - (n.leading_zeros() as usize / 8);
979                ctrl_and_size(&mut self.0.buf, 5, size);
980                for i in (0..size).rev() {
981                    self.0.buf.push((n >> (8 * i)) as u8);
982                }
983                Ok(())
984            }
985            Value::Uint32(n) => {
986                let size = 4 - (n.leading_zeros() as usize / 8);
987                ctrl_and_size(&mut self.0.buf, 6, size);
988                for i in (0..size).rev() {
989                    self.0.buf.push((n >> (8 * i)) as u8);
990                }
991                Ok(())
992            }
993            Value::Int32(n) => {
994                let u = *n as u32;
995                let size = 4 - (u.leading_zeros() as usize / 8);
996                ctrl_and_size(&mut self.0.buf, 8, size);
997                for i in (0..size).rev() {
998                    self.0.buf.push((u >> (8 * i)) as u8);
999                }
1000                Ok(())
1001            }
1002            Value::Uint64(n) => {
1003                let size = 8 - (n.leading_zeros() as usize / 8);
1004                ctrl_and_size(&mut self.0.buf, 9, size);
1005                for i in (0..size).rev() {
1006                    self.0.buf.push((n >> (8 * i)) as u8);
1007                }
1008                Ok(())
1009            }
1010            Value::Uint128(n) => {
1011                let size = 16 - (n.leading_zeros() as usize / 8);
1012                ctrl_and_size(&mut self.0.buf, 10, size);
1013                for i in (0..size).rev() {
1014                    self.0.buf.push((n >> (8 * i)) as u8);
1015                }
1016                Ok(())
1017            }
1018            Value::Map(m) => {
1019                ctrl_and_size(&mut self.0.buf, 7, m.len());
1020                for (k, val) in m.iter() {
1021                    write_str_raw(&mut self.0.buf, k);
1022                    ByteSerializerRef(&mut *self.0).serialize_value(val)?;
1023                }
1024                Ok(())
1025            }
1026            Value::Slice(s) => {
1027                ctrl_and_size(&mut self.0.buf, 11, s.len());
1028                for val in s {
1029                    ByteSerializerRef(&mut *self.0).serialize_value(val)?;
1030                }
1031                Ok(())
1032            }
1033            Value::Bool(b) => {
1034                let payload_len = if *b { 1 } else { 0 };
1035                ctrl_and_size(&mut self.0.buf, 14, payload_len);
1036                Ok(())
1037            }
1038            Value::Float32(f) => {
1039                ctrl_and_size(&mut self.0.buf, 15, 4);
1040                self.0.buf.extend_from_slice(&f.to_bits().to_be_bytes());
1041                Ok(())
1042            }
1043        }
1044    }
1045}
1046
1047pub fn encode_value(v: &Value) -> EncodeResult<Vec<u8>> {
1048    let mut ser = ByteSerializer::new();
1049    ByteSerializerRef(&mut ser).serialize_value(v).map_err(|e| e.0)?;
1050    Ok(ser.into_bytes())
1051}
1052
1053pub fn encoded_size(v: &Value) -> usize {
1054    encode_value(v).map(|b| b.len()).unwrap_or(0)
1055}
maxminddb/encoder.rs

maxminddb/
encoder.rs
