dodb_core/
key.rs

1use std::fmt;
2
3#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
4pub struct PrimaryKey(Vec<u8>);
5
6#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
7pub struct SortKey(Vec<u8>);
8
9#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
10pub struct DocumentKey {
11    pub pk: PrimaryKey,
12    pub sk: SortKey,
13}
14
15impl PrimaryKey {
16    pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
17        Self(bytes.into())
18    }
19
20    pub fn as_bytes(&self) -> &[u8] {
21        &self.0
22    }
23
24    pub fn into_bytes(self) -> Vec<u8> {
25        self.0
26    }
27}
28
29impl SortKey {
30    pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
31        Self(bytes.into())
32    }
33
34    pub fn as_bytes(&self) -> &[u8] {
35        &self.0
36    }
37
38    pub fn into_bytes(self) -> Vec<u8> {
39        self.0
40    }
41}
42
43impl DocumentKey {
44    pub fn new(pk: impl Into<Vec<u8>>, sk: impl Into<Vec<u8>>) -> Self {
45        Self {
46            pk: PrimaryKey::new(pk),
47            sk: SortKey::new(sk),
48        }
49    }
50
51    pub fn from_parts(pk: PrimaryKey, sk: SortKey) -> Self {
52        Self { pk, sk }
53    }
54
55    /// Encodes `(pk, sk)` as two escaped components.
56    ///
57    /// A zero byte in a component is encoded as `00 FF`; the component
58    /// terminator is `00 00`. This makes both arbitrary binary data and empty
59    /// components order-preserving under bytewise lexicographic comparison.
60    pub fn encode(&self) -> Vec<u8> {
61        let mut encoded = Vec::with_capacity(self.pk.0.len() + self.sk.0.len() + 4);
62        encode_component(&self.pk.0, &mut encoded);
63        encode_component(&self.sk.0, &mut encoded);
64        encoded
65    }
66
67    /// Returns the length of the canonical encoding without allocating it.
68    pub fn encoded_len(&self) -> usize {
69        encoded_component_len(&self.pk.0) + encoded_component_len(&self.sk.0)
70    }
71
72    /// Decodes the exact canonical encoding produced by [`Self::encode`].
73    pub fn decode(encoded: &[u8]) -> Result<Self, KeyCodecError> {
74        let (pk, next) = decode_component(encoded, 0)?;
75        let (sk, end) = decode_component(encoded, next)?;
76        if end != encoded.len() {
77            return Err(KeyCodecError::TrailingBytes { offset: end });
78        }
79        Ok(Self::from_parts(PrimaryKey::new(pk), SortKey::new(sk)))
80    }
81}
82
83impl fmt::Debug for PrimaryKey {
84    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
85        formatter.debug_tuple("PrimaryKey").field(&self.0).finish()
86    }
87}
88
89impl fmt::Debug for SortKey {
90    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
91        formatter.debug_tuple("SortKey").field(&self.0).finish()
92    }
93}
94
95impl fmt::Debug for DocumentKey {
96    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
97        formatter
98            .debug_struct("DocumentKey")
99            .field("pk", &self.pk)
100            .field("sk", &self.sk)
101            .finish()
102    }
103}
104
105#[derive(Clone, Debug, Eq, PartialEq)]
106pub enum KeyCodecError {
107    TruncatedEscape { offset: usize },
108    InvalidEscape { offset: usize, byte: u8 },
109    MissingComponent { offset: usize },
110    TrailingBytes { offset: usize },
111}
112
113impl fmt::Display for KeyCodecError {
114    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
115        match self {
116            Self::TruncatedEscape { offset } => write!(formatter, "truncated escape at {offset}"),
117            Self::InvalidEscape { offset, byte } => {
118                write!(formatter, "invalid escape byte {byte:#04x} at {offset}")
119            }
120            Self::MissingComponent { offset } => write!(formatter, "missing component at {offset}"),
121            Self::TrailingBytes { offset } => write!(formatter, "trailing bytes at {offset}"),
122        }
123    }
124}
125
126impl std::error::Error for KeyCodecError {}
127
128fn encode_component(bytes: &[u8], output: &mut Vec<u8>) {
129    for &byte in bytes {
130        if byte == 0 {
131            output.extend_from_slice(&[0, 0xff]);
132        } else {
133            output.push(byte);
134        }
135    }
136    output.extend_from_slice(&[0, 0]);
137}
138
139fn encoded_component_len(bytes: &[u8]) -> usize {
140    bytes.len() + bytes.iter().filter(|byte| **byte == 0).count() + 2
141}
142
143fn decode_component(encoded: &[u8], start: usize) -> Result<(Vec<u8>, usize), KeyCodecError> {
144    if start >= encoded.len() {
145        return Err(KeyCodecError::MissingComponent { offset: start });
146    }
147
148    let mut bytes = Vec::new();
149    let mut offset = start;
150    while offset < encoded.len() {
151        let byte = encoded[offset];
152        if byte != 0 {
153            bytes.push(byte);
154            offset += 1;
155            continue;
156        }
157
158        let Some(&escape) = encoded.get(offset + 1) else {
159            return Err(KeyCodecError::TruncatedEscape { offset });
160        };
161        match escape {
162            0 => return Ok((bytes, offset + 2)),
163            0xff => {
164                bytes.push(0);
165                offset += 2;
166            }
167            other => {
168                return Err(KeyCodecError::InvalidEscape {
169                    offset,
170                    byte: other,
171                });
172            }
173        }
174    }
175
176    Err(KeyCodecError::TruncatedEscape { offset })
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use proptest::prelude::*;
183
184    proptest! {
185        #[test]
186        fn round_trip_arbitrary_binary(pk in proptest::collection::vec(any::<u8>(), 0..80), sk in proptest::collection::vec(any::<u8>(), 0..80)) {
187            let key = DocumentKey::new(pk, sk);
188            let encoded = key.encode();
189            prop_assert_eq!(DocumentKey::decode(&encoded), Ok(key));
190        }
191
192        #[test]
193        fn encoding_preserves_document_order(
194            apk in proptest::collection::vec(any::<u8>(), 0..20),
195            ask in proptest::collection::vec(any::<u8>(), 0..20),
196            bpk in proptest::collection::vec(any::<u8>(), 0..20),
197            bsk in proptest::collection::vec(any::<u8>(), 0..20),
198        ) {
199            let left = DocumentKey::new(apk, ask);
200            let right = DocumentKey::new(bpk, bsk);
201            prop_assert_eq!(left.cmp(&right), left.encode().cmp(&right.encode()));
202        }
203    }
204
205    #[test]
206    fn malformed_encodings_are_rejected() {
207        for input in [
208            vec![],
209            vec![0],
210            vec![0, 1],
211            vec![0, 0],
212            vec![0, 0, 0],
213            vec![1, 0, 0, 0, 0, 9],
214        ] {
215            assert!(DocumentKey::decode(&input).is_err(), "accepted {input:?}");
216        }
217        assert!(DocumentKey::decode(&[0, 0, 0, 0]).is_ok());
218    }
219
220    #[test]
221    fn encoded_len_matches_canonical_encoding_and_counts_zero_escapes() {
222        let key = DocumentKey::new(vec![0; 1_994], Vec::new());
223        assert_eq!(key.encoded_len(), 3_992);
224        assert_eq!(key.encoded_len(), key.encode().len());
225        let oversized = DocumentKey::new(vec![0; 1_995], Vec::new());
226        assert_eq!(oversized.encoded_len(), 3_994);
227    }
228}
dodb_core/key.rs

dodb_core/
key.rs
