1use reifydb_codec::key::{
5 deserializer::KeyDeserializer,
6 encoded::{EncodedKey, EncodedKeyRange},
7 serializer::KeySerializer,
8};
9use reifydb_value::value::sumtype::SumTypeId;
10
11use crate::{
12 interface::catalog::id::{HandlerId, NamespaceId},
13 key::{EncodableKey, KeyKind},
14};
15
16#[derive(Debug, Clone, PartialEq)]
17pub struct VariantHandlerKey {
18 pub namespace: NamespaceId,
19 pub sumtype: SumTypeId,
20 pub variant_tag: u8,
21 pub handler: HandlerId,
22}
23
24impl VariantHandlerKey {
25 pub fn new(namespace: NamespaceId, sumtype: SumTypeId, variant_tag: u8, handler: HandlerId) -> Self {
26 Self {
27 namespace,
28 sumtype,
29 variant_tag,
30 handler,
31 }
32 }
33
34 pub fn encoded(
35 namespace: impl Into<NamespaceId>,
36 sumtype: impl Into<SumTypeId>,
37 variant_tag: u8,
38 handler: impl Into<HandlerId>,
39 ) -> EncodedKey {
40 Self::new(namespace.into(), sumtype.into(), variant_tag, handler.into()).encode()
41 }
42
43 pub fn variant_scan(namespace: NamespaceId, sumtype: SumTypeId, variant_tag: u8) -> EncodedKeyRange {
44 EncodedKeyRange::start_end(
45 Some(Self::variant_start(namespace, sumtype, variant_tag)),
46 Some(Self::variant_end(namespace, sumtype, variant_tag)),
47 )
48 }
49
50 fn variant_start(namespace: NamespaceId, sumtype: SumTypeId, variant_tag: u8) -> EncodedKey {
51 let mut serializer = KeySerializer::with_capacity(18);
52 serializer.extend_u8(Self::KIND as u8).extend_u64(namespace).extend_u64(sumtype).extend_u8(variant_tag);
53 serializer.to_encoded_key()
54 }
55
56 fn variant_end(namespace: NamespaceId, sumtype: SumTypeId, variant_tag: u8) -> EncodedKey {
57 let mut serializer = KeySerializer::with_capacity(18);
58 serializer
59 .extend_u8(Self::KIND as u8)
60 .extend_u64(namespace)
61 .extend_u64(sumtype)
62 .extend_u8(variant_tag.wrapping_sub(1));
63 serializer.to_encoded_key()
64 }
65}
66
67impl EncodableKey for VariantHandlerKey {
68 const KIND: KeyKind = KeyKind::VariantHandler;
69
70 fn encode(&self) -> EncodedKey {
71 let mut serializer = KeySerializer::with_capacity(26);
72 serializer
73 .extend_u8(Self::KIND as u8)
74 .extend_u64(self.namespace)
75 .extend_u64(self.sumtype)
76 .extend_u8(self.variant_tag)
77 .extend_u64(self.handler);
78 serializer.to_encoded_key()
79 }
80
81 fn decode(key: &EncodedKey) -> Option<Self> {
82 let mut de = KeyDeserializer::from_bytes(key.as_slice());
83
84 let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
85 if kind != Self::KIND {
86 return None;
87 }
88
89 let namespace = de.read_u64().ok()?;
90 let sumtype = de.read_u64().ok()?;
91 let variant_tag = de.read_u8().ok()?;
92 let handler = de.read_u64().ok()?;
93
94 Some(Self {
95 namespace: NamespaceId(namespace),
96 sumtype: SumTypeId(sumtype),
97 variant_tag,
98 handler: HandlerId(handler),
99 })
100 }
101}
102
103#[cfg(test)]
104pub mod tests {
105 use std::ops::Bound;
106
107 use reifydb_value::value::sumtype::SumTypeId;
108
109 use super::{EncodableKey, VariantHandlerKey};
110 use crate::interface::catalog::id::{HandlerId, NamespaceId};
111
112 #[test]
113 fn test_encode_decode() {
114 let key = VariantHandlerKey {
115 namespace: NamespaceId(0xABCD),
116 sumtype: SumTypeId(0x1234),
117 variant_tag: 5,
118 handler: HandlerId(0x6789),
119 };
120 let encoded = key.encode();
121 let expected: Vec<u8> = vec![
122 0xD2, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x54, 0x32, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xED,
123 0xCB, 0xFA, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x98, 0x76,
124 ];
125 assert_eq!(encoded.as_slice(), expected);
126
127 let decoded = VariantHandlerKey::decode(&encoded).unwrap();
128 assert_eq!(decoded.namespace, NamespaceId(0xABCD));
129 assert_eq!(decoded.sumtype, SumTypeId(0x1234));
130 assert_eq!(decoded.variant_tag, 5);
131 assert_eq!(decoded.handler, HandlerId(0x6789));
132 }
133
134 #[test]
135 fn test_order_preserving() {
136 let key1 = VariantHandlerKey {
137 namespace: NamespaceId::SYSTEM,
138 sumtype: SumTypeId(5),
139 variant_tag: 3,
140 handler: HandlerId(100),
141 };
142 let key2 = VariantHandlerKey {
143 namespace: NamespaceId::SYSTEM,
144 sumtype: SumTypeId(5),
145 variant_tag: 3,
146 handler: HandlerId(200),
147 };
148 let key3 = VariantHandlerKey {
149 namespace: NamespaceId::SYSTEM,
150 sumtype: SumTypeId(5),
151 variant_tag: 4,
152 handler: HandlerId(1),
153 };
154 let key4 = VariantHandlerKey {
155 namespace: NamespaceId::DEFAULT,
156 sumtype: SumTypeId(1),
157 variant_tag: 0,
158 handler: HandlerId(1),
159 };
160
161 let encoded1 = key1.encode();
162 let encoded2 = key2.encode();
163 let encoded3 = key3.encode();
164 let encoded4 = key4.encode();
165
166 assert!(encoded4 < encoded3, "ordering not preserved");
167 assert!(encoded3 < encoded2, "ordering not preserved");
168 assert!(encoded2 < encoded1, "ordering not preserved");
169 }
170
171 #[test]
172 fn test_variant_scan() {
173 let ns = NamespaceId::SYSTEM;
174 let st = SumTypeId(10);
175 let tag = 5u8;
176
177 let range = VariantHandlerKey::variant_scan(ns, st, tag);
178 let start = match &range.start {
179 Bound::Included(k) | Bound::Excluded(k) => k,
180 Bound::Unbounded => panic!("expected bounded start"),
181 };
182 let end = match &range.end {
183 Bound::Included(k) | Bound::Excluded(k) => k,
184 Bound::Unbounded => panic!("expected bounded end"),
185 };
186
187 let key = VariantHandlerKey {
188 namespace: ns,
189 sumtype: st,
190 variant_tag: tag,
191 handler: HandlerId(42),
192 };
193 let encoded = key.encode();
194 assert!(encoded.as_slice() >= start.as_slice());
195 assert!(encoded.as_slice() <= end.as_slice());
196
197 let other = VariantHandlerKey {
198 namespace: ns,
199 sumtype: st,
200 variant_tag: tag + 1,
201 handler: HandlerId(42),
202 };
203 let other_encoded = other.encode();
204 assert!(other_encoded.as_slice() < start.as_slice());
205 }
206}