1use crate::error::{DecodeError, EncodeError};
2use crate::tag;
3
4#[derive(Debug, Copy, Clone, PartialEq, Eq)]
5pub enum VectorType {
6 Integer8,
7 Integer16,
8 Integer32,
9 Integer64,
10 Float32,
11 Float64,
12}
13
14impl VectorType {
15 pub fn code(self) -> u8 {
16 match self {
17 Self::Integer8 => 0x01,
18 Self::Integer16 => 0x02,
19 Self::Integer32 => 0x03,
20 Self::Integer64 => 0x04,
21 Self::Float32 => 0x11,
22 Self::Float64 => 0x12,
23 }
24 }
25
26 pub fn from_code(code: u8) -> Result<Self, DecodeError> {
27 match code {
28 0x01 => Ok(Self::Integer8),
29 0x02 => Ok(Self::Integer16),
30 0x03 => Ok(Self::Integer32),
31 0x04 => Ok(Self::Integer64),
32 0x11 => Ok(Self::Float32),
33 0x12 => Ok(Self::Float64),
34 _ => Err(DecodeError::InvalidVector(format!(
35 "unknown vector coordinate type 0x{code:02x}"
36 ))),
37 }
38 }
39}
40
41#[derive(Debug, Clone, PartialEq)]
42pub enum VectorStorage {
43 I8(Vec<i8>),
44 I16(Vec<i16>),
45 I32(Vec<i32>),
46 I64(Vec<i64>),
47 F32(Vec<f32>),
48 F64(Vec<f64>),
49}
50
51#[derive(Debug, Copy, Clone, PartialEq)]
52pub enum VectorElem {
53 I8(i8),
54 I16(i16),
55 I32(i32),
56 I64(i64),
57 F32(f32),
58 F64(f64),
59}
60
61#[derive(Debug, Clone, PartialEq)]
62pub struct VectorValue {
63 pub coord_type: VectorType,
64 pub values: VectorStorage,
65}
66
67impl VectorValue {
68 pub fn len(&self) -> usize {
69 match &self.values {
70 VectorStorage::I8(v) => v.len(),
71 VectorStorage::I16(v) => v.len(),
72 VectorStorage::I32(v) => v.len(),
73 VectorStorage::I64(v) => v.len(),
74 VectorStorage::F32(v) => v.len(),
75 VectorStorage::F64(v) => v.len(),
76 }
77 }
78
79 pub fn is_empty(&self) -> bool {
80 self.len() == 0
81 }
82
83 pub fn iter(&self) -> VectorIter<'_> {
84 match &self.values {
85 VectorStorage::I8(v) => VectorIter::I8(v.iter()),
86 VectorStorage::I16(v) => VectorIter::I16(v.iter()),
87 VectorStorage::I32(v) => VectorIter::I32(v.iter()),
88 VectorStorage::I64(v) => VectorIter::I64(v.iter()),
89 VectorStorage::F32(v) => VectorIter::F32(v.iter()),
90 VectorStorage::F64(v) => VectorIter::F64(v.iter()),
91 }
92 }
93
94 pub fn validate(&self) -> Result<(), EncodeError> {
95 let ok = matches!(
96 (&self.coord_type, &self.values),
97 (VectorType::Integer8, VectorStorage::I8(_))
98 | (VectorType::Integer16, VectorStorage::I16(_))
99 | (VectorType::Integer32, VectorStorage::I32(_))
100 | (VectorType::Integer64, VectorStorage::I64(_))
101 | (VectorType::Float32, VectorStorage::F32(_))
102 | (VectorType::Float64, VectorStorage::F64(_))
103 );
104 if ok {
105 Ok(())
106 } else {
107 Err(EncodeError::InvalidVector(
108 "coord_type does not match vector storage",
109 ))
110 }
111 }
112
113 pub fn encode_blob(&self) -> Result<Vec<u8>, EncodeError> {
114 self.validate()?;
115 if self.len() > u16::MAX as usize {
116 return Err(EncodeError::InvalidVector("vector length exceeds u16::MAX"));
117 }
118 let mut out = Vec::with_capacity(1 + 1 + 2 + self.len() * 8);
119 out.push(tag::VECTOR);
120 out.push(self.coord_type.code());
121 out.extend_from_slice(&(self.len() as u16).to_le_bytes());
122 match &self.values {
123 VectorStorage::I8(values) => {
124 for value in values {
125 out.push(*value as u8);
126 }
127 }
128 VectorStorage::I16(values) => {
129 for value in values {
130 out.extend_from_slice(&value.to_le_bytes());
131 }
132 }
133 VectorStorage::I32(values) => {
134 for value in values {
135 out.extend_from_slice(&value.to_le_bytes());
136 }
137 }
138 VectorStorage::I64(values) => {
139 for value in values {
140 out.extend_from_slice(&value.to_le_bytes());
141 }
142 }
143 VectorStorage::F32(values) => {
144 for value in values {
145 if !value.is_finite() {
146 return Err(EncodeError::NonCanonicalValue(
147 "non-finite float in Float32 vector",
148 ));
149 }
150 out.extend_from_slice(&value.to_le_bytes());
151 }
152 }
153 VectorStorage::F64(values) => {
154 for value in values {
155 if !value.is_finite() {
156 return Err(EncodeError::NonCanonicalValue(
157 "non-finite float in Float64 vector",
158 ));
159 }
160 out.extend_from_slice(&value.to_le_bytes());
161 }
162 }
163 }
164 Ok(out)
165 }
166
167 pub fn decode_blob(blob: &[u8]) -> Result<Self, DecodeError> {
168 if blob.is_empty() {
169 return Err(DecodeError::EmptyInput);
170 }
171 if blob[0] != tag::VECTOR {
172 return Err(DecodeError::UnexpectedTag {
173 expected: "vector tag",
174 actual: blob[0],
175 });
176 }
177 if blob.len() < 4 {
178 return Err(DecodeError::Truncated);
179 }
180 let coord_type = VectorType::from_code(blob[1])?;
181 let len = u16::from_le_bytes([blob[2], blob[3]]) as usize;
182 let payload = &blob[4..];
183 let values = match coord_type {
184 VectorType::Integer8 => {
185 if payload.len() != len {
186 return Err(DecodeError::InvalidVector(
187 "invalid i8 payload length".into(),
188 ));
189 }
190 VectorStorage::I8(payload.iter().map(|b| *b as i8).collect())
191 }
192 VectorType::Integer16 => {
193 if payload.len() != len * 2 {
194 return Err(DecodeError::InvalidVector(
195 "invalid i16 payload length".into(),
196 ));
197 }
198 let mut out = Vec::with_capacity(len);
199 for chunk in payload.chunks_exact(2) {
200 out.push(i16::from_le_bytes([chunk[0], chunk[1]]));
201 }
202 VectorStorage::I16(out)
203 }
204 VectorType::Integer32 => {
205 if payload.len() != len * 4 {
206 return Err(DecodeError::InvalidVector(
207 "invalid i32 payload length".into(),
208 ));
209 }
210 let mut out = Vec::with_capacity(len);
211 for chunk in payload.chunks_exact(4) {
212 out.push(i32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
213 }
214 VectorStorage::I32(out)
215 }
216 VectorType::Integer64 => {
217 if payload.len() != len * 8 {
218 return Err(DecodeError::InvalidVector(
219 "invalid i64 payload length".into(),
220 ));
221 }
222 let mut out = Vec::with_capacity(len);
223 for chunk in payload.chunks_exact(8) {
224 out.push(i64::from_le_bytes([
225 chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6],
226 chunk[7],
227 ]));
228 }
229 VectorStorage::I64(out)
230 }
231 VectorType::Float32 => {
232 if payload.len() != len * 4 {
233 return Err(DecodeError::InvalidVector(
234 "invalid f32 payload length".into(),
235 ));
236 }
237 let mut out = Vec::with_capacity(len);
238 for chunk in payload.chunks_exact(4) {
239 let value = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
240 if !value.is_finite() {
241 return Err(DecodeError::NonCanonical(
242 "non-finite float in Float32 vector",
243 ));
244 }
245 out.push(value);
246 }
247 VectorStorage::F32(out)
248 }
249 VectorType::Float64 => {
250 if payload.len() != len * 8 {
251 return Err(DecodeError::InvalidVector(
252 "invalid f64 payload length".into(),
253 ));
254 }
255 let mut out = Vec::with_capacity(len);
256 for chunk in payload.chunks_exact(8) {
257 let value = f64::from_le_bytes([
258 chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6],
259 chunk[7],
260 ]);
261 if !value.is_finite() {
262 return Err(DecodeError::NonCanonical(
263 "non-finite float in Float64 vector",
264 ));
265 }
266 out.push(value);
267 }
268 VectorStorage::F64(out)
269 }
270 };
271 Ok(Self { coord_type, values })
272 }
273}
274
275pub enum VectorIter<'a> {
276 I8(std::slice::Iter<'a, i8>),
277 I16(std::slice::Iter<'a, i16>),
278 I32(std::slice::Iter<'a, i32>),
279 I64(std::slice::Iter<'a, i64>),
280 F32(std::slice::Iter<'a, f32>),
281 F64(std::slice::Iter<'a, f64>),
282}
283
284#[cfg(test)]
285mod tests {
286 use super::*;
287
288 fn assert_vector_roundtrip(value: VectorValue) {
289 let decoded = VectorValue::decode_blob(&value.encode_blob().unwrap()).unwrap();
290 assert_eq!(decoded, value);
291 }
292
293 #[test]
294 fn roundtrip_empty_vector() {
295 assert_vector_roundtrip(VectorValue {
296 coord_type: VectorType::Integer16,
297 values: VectorStorage::I16(vec![]),
298 });
299 }
300
301 #[test]
302 fn roundtrip_single_value_vector() {
303 assert_vector_roundtrip(VectorValue {
304 coord_type: VectorType::Integer8,
305 values: VectorStorage::I8(vec![-7]),
306 });
307 }
308
309 #[test]
310 fn roundtrip_longer_vector() {
311 let values = (0..257).map(|v| v as i32 - 128).collect::<Vec<_>>();
312 assert_vector_roundtrip(VectorValue {
313 coord_type: VectorType::Integer32,
314 values: VectorStorage::I32(values),
315 });
316 }
317
318 #[test]
319 fn rejects_oversized_vector() {
320 let value = VectorValue {
321 coord_type: VectorType::Integer8,
322 values: VectorStorage::I8(vec![0; u16::MAX as usize + 1]),
323 };
324 assert!(matches!(
325 value.encode_blob(),
326 Err(EncodeError::InvalidVector("vector length exceeds u16::MAX"))
327 ));
328 }
329
330 #[test]
331 fn rejects_mismatched_coord_type_and_storage() {
332 let value = VectorValue {
333 coord_type: VectorType::Integer16,
334 values: VectorStorage::I8(vec![1, 2, 3]),
335 };
336 assert!(matches!(
337 value.encode_blob(),
338 Err(EncodeError::InvalidVector(
339 "coord_type does not match vector storage"
340 ))
341 ));
342 }
343
344 #[test]
345 fn rejects_non_finite_float_vector_on_encode() {
346 let value = VectorValue {
347 coord_type: VectorType::Float64,
348 values: VectorStorage::F64(vec![1.0, f64::NAN]),
349 };
350 assert!(matches!(
351 value.encode_blob(),
352 Err(EncodeError::NonCanonicalValue(
353 "non-finite float in Float64 vector"
354 ))
355 ));
356 }
357
358 #[test]
359 fn rejects_non_finite_float_vector_on_decode() {
360 let mut blob = vec![tag::VECTOR, VectorType::Float64.code()];
361 blob.extend_from_slice(&(2u16).to_le_bytes());
362 blob.extend_from_slice(&1.0f64.to_le_bytes());
363 blob.extend_from_slice(&f64::INFINITY.to_le_bytes());
364 assert!(matches!(
365 VectorValue::decode_blob(&blob),
366 Err(DecodeError::NonCanonical(
367 "non-finite float in Float64 vector"
368 ))
369 ));
370 }
371}
372
373impl<'a> Iterator for VectorIter<'a> {
374 type Item = VectorElem;
375
376 fn next(&mut self) -> Option<Self::Item> {
377 match self {
378 Self::I8(iter) => iter.next().copied().map(VectorElem::I8),
379 Self::I16(iter) => iter.next().copied().map(VectorElem::I16),
380 Self::I32(iter) => iter.next().copied().map(VectorElem::I32),
381 Self::I64(iter) => iter.next().copied().map(VectorElem::I64),
382 Self::F32(iter) => iter.next().copied().map(VectorElem::F32),
383 Self::F64(iter) => iter.next().copied().map(VectorElem::F64),
384 }
385 }
386}