reifydb_core/value/index/
layout.rs1use std::{ops::Deref, sync::Arc};
5
6use reifydb_type::{error, util::cowvec::CowVec, value::r#type::Type};
7
8use crate::{
9 error::diagnostic::catalog::{index_types_directions_mismatch, index_variable_length_not_supported},
10 sort::SortDirection,
11 value::index::encoded::EncodedIndexKey,
12};
13
14#[derive(Debug, Clone)]
15pub struct EncodedIndexLayout(Arc<EncodedIndexLayoutInner>);
16
17impl Deref for EncodedIndexLayout {
18 type Target = EncodedIndexLayoutInner;
19
20 fn deref(&self) -> &Self::Target {
21 &self.0
22 }
23}
24
25impl EncodedIndexLayout {
26 pub fn new(types: &[Type], directions: &[SortDirection]) -> reifydb_type::Result<Self> {
27 if types.len() != directions.len() {
28 return Err(error!(index_types_directions_mismatch(types.len(), directions.len())));
29 }
30
31 for typ in types {
32 if matches!(typ, Type::Utf8 | Type::Blob) {
33 return Err(error!(index_variable_length_not_supported()));
34 }
35 }
36
37 Ok(Self(Arc::new(EncodedIndexLayoutInner::new(types, directions))))
38 }
39}
40
41#[derive(Debug)]
42pub struct EncodedIndexLayoutInner {
43 pub fields: Vec<IndexField>,
44 pub total_size: usize,
45 pub bitvec_size: usize,
46 pub alignment: usize,
47}
48
49#[derive(Debug)]
50pub struct IndexField {
51 pub offset: usize,
52 pub size: usize,
53 pub align: usize,
54 pub value: Type,
55 pub direction: SortDirection,
56}
57
58impl EncodedIndexLayoutInner {
59 fn new(types: &[Type], directions: &[SortDirection]) -> Self {
60 assert!(!types.is_empty());
61 assert_eq!(types.len(), directions.len());
62
63 let num_fields = types.len();
64 let bitvec_bytes = (num_fields + 7) / 8;
65
66 let mut offset = bitvec_bytes;
67 let mut fields = Vec::with_capacity(num_fields);
68 let mut max_align = 1;
69
70 for (i, value) in types.iter().enumerate() {
71 let size = value.size();
72 let align = value.alignment();
73
74 offset = align_up(offset, align);
75 fields.push(IndexField {
76 offset,
77 size,
78 align,
79 value: value.clone(),
80 direction: directions[i].clone(),
81 });
82
83 offset += size;
84 max_align = max_align.max(align);
85 }
86
87 let total_size = align_up(offset, max_align);
88
89 EncodedIndexLayoutInner {
90 fields,
91 total_size,
92 alignment: max_align,
93 bitvec_size: bitvec_bytes,
94 }
95 }
96
97 pub fn allocate_key(&self) -> EncodedIndexKey {
98 let layout = std::alloc::Layout::from_size_align(self.total_size, self.alignment).unwrap();
99 unsafe {
100 let ptr = std::alloc::alloc_zeroed(layout);
101 if ptr.is_null() {
102 std::alloc::handle_alloc_error(layout);
103 }
104 let vec = Vec::from_raw_parts(ptr, self.total_size, self.total_size);
105 EncodedIndexKey(CowVec::new(vec))
106 }
107 }
108
109 pub const fn data_offset(&self) -> usize {
110 self.bitvec_size
111 }
112
113 pub fn all_defined(&self, key: &EncodedIndexKey) -> bool {
114 let bits = self.fields.len();
115 if bits == 0 {
116 return false;
117 }
118
119 let bitvec_slice = &key[..self.bitvec_size];
120 for (i, &byte) in bitvec_slice.iter().enumerate() {
121 let bits_in_byte = if i == self.bitvec_size - 1 && bits % 8 != 0 {
122 bits % 8
123 } else {
124 8
125 };
126
127 let mask = if bits_in_byte == 8 {
128 0xFF
129 } else {
130 (1u8 << bits_in_byte) - 1
131 };
132 if (byte & mask) != mask {
133 return false;
134 }
135 }
136
137 true
138 }
139
140 pub fn value(&self, index: usize) -> Type {
141 self.fields[index].value.clone()
142 }
143
144 pub fn direction(&self, index: usize) -> &SortDirection {
145 &self.fields[index].direction
146 }
147}
148
149fn align_up(offset: usize, align: usize) -> usize {
150 (offset + align - 1) & !(align - 1)
151}
152
153#[cfg(test)]
154pub mod tests {
155 use super::*;
156 use crate::sort::SortDirection;
157
158 #[test]
159 fn test_single_field_int() {
160 let layout = EncodedIndexLayout::new(&[Type::Int4], &[SortDirection::Asc]).unwrap();
161
162 assert_eq!(layout.bitvec_size, 1);
163 assert_eq!(layout.fields.len(), 1);
164 assert_eq!(layout.fields[0].offset, 4); assert_eq!(layout.fields[0].value, Type::Int4);
166 assert_eq!(layout.fields[0].direction, SortDirection::Asc);
167 assert_eq!(layout.alignment, 4);
168 assert_eq!(layout.total_size, 8); }
170
171 #[test]
172 fn test_multiple_fields_mixed_directions() {
173 let layout = EncodedIndexLayout::new(
174 &[Type::Int4, Type::Int8, Type::Uint8],
175 &[SortDirection::Desc, SortDirection::Asc, SortDirection::Asc],
176 )
177 .unwrap();
178
179 assert_eq!(layout.bitvec_size, 1);
180 assert_eq!(layout.fields.len(), 3);
181
182 assert_eq!(layout.fields[0].value, Type::Int4);
183 assert_eq!(layout.fields[0].direction, SortDirection::Desc);
184
185 assert_eq!(layout.fields[1].value, Type::Int8);
186 assert_eq!(layout.fields[1].direction, SortDirection::Asc);
187
188 assert_eq!(layout.fields[2].value, Type::Uint8);
189 assert_eq!(layout.fields[2].direction, SortDirection::Asc);
190
191 assert_eq!(layout.alignment, 8);
192 }
193
194 #[test]
195 fn test_reject_variable_length_types() {
196 let result =
197 EncodedIndexLayout::new(&[Type::Int4, Type::Utf8], &[SortDirection::Asc, SortDirection::Asc]);
198 assert!(result.is_err());
199
200 let result = EncodedIndexLayout::new(&[Type::Blob], &[SortDirection::Desc]);
201 assert!(result.is_err());
202 }
203
204 #[test]
205 fn test_allocate_key() {
206 let layout = EncodedIndexLayout::new(
207 &[Type::Boolean, Type::Int4],
208 &[SortDirection::Asc, SortDirection::Desc],
209 )
210 .unwrap();
211
212 let key = layout.allocate_key();
213 assert_eq!(key.len(), layout.total_size);
214
215 for byte in key.as_slice() {
216 assert_eq!(*byte, 0);
217 }
218 }
219}