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