reifydb_engine/arena/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4//! Bumpalo-backed arena allocator used by the engine for per-request scratch storage. Plans, intermediate column
5//! buffers, and the bit-vectors that drive selection live in an arena tied to the request lifetime so the engine
6//! can release everything in one bump-reset rather than tracking individual allocations.
7
8use std::{
9	fmt::{self, Debug},
10	ops::Deref,
11};
12
13use bumpalo::Bump as BumpAlloc;
14use reifydb_type::storage::{DataBitVec, DataVec, Storage};
15
16type BumpaloVec<'bump, T> = bumpalo::collections::Vec<'bump, T>;
17
18pub mod convert;
19
20pub struct BumpVec<'bump, T: Clone + PartialEq> {
21	inner: BumpaloVec<'bump, T>,
22}
23
24impl<'bump, T: Clone + PartialEq + Debug> Debug for BumpVec<'bump, T> {
25	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26		f.debug_tuple("BumpVec").field(&self.inner.as_slice()).finish()
27	}
28}
29
30impl<'bump, T: Clone + PartialEq> BumpVec<'bump, T> {
31	pub fn with_capacity_in(capacity: usize, bump: &'bump BumpAlloc) -> Self {
32		Self {
33			inner: BumpaloVec::with_capacity_in(capacity, bump),
34		}
35	}
36
37	fn bump(&self) -> &'bump BumpAlloc {
38		self.inner.bump()
39	}
40}
41
42impl<'bump, T: Clone + PartialEq> Clone for BumpVec<'bump, T> {
43	fn clone(&self) -> Self {
44		let mut new = BumpaloVec::with_capacity_in(self.inner.len(), self.bump());
45		new.extend(self.inner.iter().cloned());
46		Self {
47			inner: new,
48		}
49	}
50}
51
52impl<'bump, T: Clone + PartialEq> PartialEq for BumpVec<'bump, T> {
53	fn eq(&self, other: &Self) -> bool {
54		self.inner.as_slice() == other.inner.as_slice()
55	}
56}
57
58impl<'bump, T: Clone + PartialEq> Deref for BumpVec<'bump, T> {
59	type Target = [T];
60
61	fn deref(&self) -> &Self::Target {
62		self.inner.as_slice()
63	}
64}
65
66impl<'bump, T: Clone + PartialEq> DataVec<T> for BumpVec<'bump, T> {
67	fn spawn(&self, capacity: usize) -> Self {
68		Self {
69			inner: BumpaloVec::with_capacity_in(capacity, self.bump()),
70		}
71	}
72
73	fn push(&mut self, value: T) {
74		self.inner.push(value);
75	}
76
77	fn clear(&mut self) {
78		self.inner.clear();
79	}
80
81	fn len(&self) -> usize {
82		self.inner.len()
83	}
84
85	fn as_slice(&self) -> &[T] {
86		self.inner.as_slice()
87	}
88
89	fn get(&self, idx: usize) -> Option<&T> {
90		self.inner.get(idx)
91	}
92
93	fn extend_from_slice(&mut self, other: &[T]) {
94		self.inner.extend(other.iter().cloned());
95	}
96
97	fn extend_iter(&mut self, iter: impl Iterator<Item = T>) {
98		self.inner.extend(iter);
99	}
100
101	fn capacity(&self) -> usize {
102		self.inner.capacity()
103	}
104}
105
106pub struct BumpBitVec<'bump> {
107	bits: BumpaloVec<'bump, u8>,
108	len: usize,
109}
110
111impl<'bump> Debug for BumpBitVec<'bump> {
112	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113		f.debug_struct("BumpBitVec").field("len", &self.len).finish()
114	}
115}
116
117impl<'bump> BumpBitVec<'bump> {
118	pub fn with_capacity_in(capacity: usize, bump: &'bump BumpAlloc) -> Self {
119		let byte_capacity = capacity.div_ceil(8);
120		Self {
121			bits: BumpaloVec::with_capacity_in(byte_capacity, bump),
122			len: 0,
123		}
124	}
125
126	fn bump(&self) -> &'bump BumpAlloc {
127		self.bits.bump()
128	}
129}
130
131impl<'bump> Clone for BumpBitVec<'bump> {
132	fn clone(&self) -> Self {
133		let mut new_bits = BumpaloVec::with_capacity_in(self.bits.len(), self.bump());
134		new_bits.extend(self.bits.iter().copied());
135		Self {
136			bits: new_bits,
137			len: self.len,
138		}
139	}
140}
141
142impl<'bump> PartialEq for BumpBitVec<'bump> {
143	fn eq(&self, other: &Self) -> bool {
144		if self.len != other.len {
145			return false;
146		}
147
148		let full_bytes = self.len / 8;
149		if self.bits[..full_bytes] != other.bits[..full_bytes] {
150			return false;
151		}
152
153		let remainder = self.len % 8;
154		if remainder > 0 {
155			let mask = (1u8 << remainder) - 1;
156			let a = self.bits[full_bytes] & mask;
157			let b = other.bits[full_bytes] & mask;
158			if a != b {
159				return false;
160			}
161		}
162		true
163	}
164}
165
166impl<'bump> DataBitVec for BumpBitVec<'bump> {
167	fn spawn(&self, capacity: usize) -> Self {
168		BumpBitVec::with_capacity_in(capacity, self.bump())
169	}
170
171	fn push(&mut self, bit: bool) {
172		let byte_index = self.len / 8;
173		let bit_index = self.len % 8;
174
175		if byte_index >= self.bits.len() {
176			self.bits.push(0);
177		}
178
179		if bit {
180			self.bits[byte_index] |= 1 << bit_index;
181		}
182
183		self.len += 1;
184	}
185
186	fn get(&self, idx: usize) -> bool {
187		assert!(idx < self.len);
188		let byte = self.bits[idx / 8];
189		let bit = idx % 8;
190		(byte >> bit) & 1 != 0
191	}
192
193	fn set(&mut self, idx: usize, value: bool) {
194		assert!(idx < self.len);
195		let byte = &mut self.bits[idx / 8];
196		let bit = idx % 8;
197		if value {
198			*byte |= 1 << bit;
199		} else {
200			*byte &= !(1 << bit);
201		}
202	}
203
204	fn len(&self) -> usize {
205		self.len
206	}
207
208	fn clear(&mut self) {
209		self.bits.clear();
210		self.len = 0;
211	}
212
213	fn extend_from(&mut self, other: &Self) {
214		for i in 0..other.len {
215			self.push(DataBitVec::get(other, i));
216		}
217	}
218
219	fn count_ones(&self) -> usize {
220		let mut count: usize = self.bits.iter().map(|&byte: &u8| byte.count_ones() as usize).sum();
221
222		let full_bytes = self.len / 8;
223		let remainder_bits = self.len % 8;
224
225		if remainder_bits > 0 && full_bytes < self.bits.len() {
226			let last_byte = self.bits[full_bytes];
227			let mask = (1u8 << remainder_bits) - 1;
228			count -= (last_byte & !mask).count_ones() as usize;
229		}
230
231		count
232	}
233
234	fn iter(&self) -> impl Iterator<Item = bool> + '_ {
235		(0..self.len).map(|i| {
236			let byte = self.bits[i / 8];
237			let bit = i % 8;
238			(byte >> bit) & 1 != 0
239		})
240	}
241
242	fn capacity(&self) -> usize {
243		self.bits.capacity() * 8
244	}
245}
246
247#[derive(Clone)]
248pub struct Bump<'bump>(&'bump BumpAlloc);
249
250impl<'bump> Bump<'bump> {
251	pub fn new(bump: &'bump BumpAlloc) -> Self {
252		Self(bump)
253	}
254
255	pub fn inner(&self) -> &'bump BumpAlloc {
256		self.0
257	}
258}
259
260impl<'bump> Storage for Bump<'bump> {
261	type Vec<T: Clone + PartialEq + 'static> = BumpVec<'bump, T>;
262	type BitVec = BumpBitVec<'bump>;
263}
264
265pub struct QueryArena {
266	bump: BumpAlloc,
267}
268
269impl QueryArena {
270	pub fn new() -> Self {
271		Self {
272			bump: BumpAlloc::with_capacity(64 * 1024),
273		}
274	}
275
276	pub fn reset(&mut self) {
277		self.bump.reset();
278	}
279
280	pub fn bump(&self) -> &BumpAlloc {
281		&self.bump
282	}
283}
284
285impl Default for QueryArena {
286	fn default() -> Self {
287		Self::new()
288	}
289}
290
291#[cfg(test)]
292mod tests {
293	use reifydb_type::storage::{DataBitVec, DataVec};
294
295	use super::*;
296
297	mod bump_vec {
298		use super::*;
299
300		#[test]
301		fn test_push_and_get() {
302			let bump = BumpAlloc::new();
303			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
304			v.push(10);
305			v.push(20);
306			v.push(30);
307
308			assert_eq!(DataVec::len(&v), 3);
309			assert_eq!(DataVec::get(&v, 0), Some(&10));
310			assert_eq!(DataVec::get(&v, 1), Some(&20));
311			assert_eq!(DataVec::get(&v, 2), Some(&30));
312			assert_eq!(DataVec::get(&v, 3), None);
313		}
314
315		#[test]
316		fn test_extend_from_slice() {
317			let bump = BumpAlloc::new();
318			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(8, &bump);
319			v.push(1);
320			DataVec::extend_from_slice(&mut v, &[2, 3, 4]);
321			assert_eq!(v.as_slice(), &[1, 2, 3, 4]);
322		}
323
324		#[test]
325		fn test_spawn() {
326			let bump = BumpAlloc::new();
327			let v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
328			let v2 = DataVec::<i32>::spawn(&v, 8);
329			assert_eq!(DataVec::len(&v2), 0);
330			assert!(DataVec::capacity(&v2) >= 8);
331		}
332
333		#[test]
334		fn test_clone() {
335			let bump = BumpAlloc::new();
336			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
337			v.push(1);
338			v.push(2);
339
340			let v2 = v.clone();
341			assert_eq!(v.as_slice(), v2.as_slice());
342		}
343
344		#[test]
345		fn test_clear() {
346			let bump = BumpAlloc::new();
347			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
348			v.push(1);
349			v.push(2);
350			DataVec::clear(&mut v);
351			assert_eq!(DataVec::len(&v), 0);
352		}
353
354		#[test]
355		fn test_deref() {
356			let bump = BumpAlloc::new();
357			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
358			v.push(10);
359			v.push(20);
360			let slice: &[i32] = &*v;
361			assert_eq!(slice, &[10, 20]);
362		}
363
364		#[test]
365		fn test_eq() {
366			let bump = BumpAlloc::new();
367			let mut v1: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
368			v1.push(1);
369			v1.push(2);
370
371			let mut v2: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
372			v2.push(1);
373			v2.push(2);
374
375			assert_eq!(v1, v2);
376
377			v2.push(3);
378			assert_ne!(v1, v2);
379		}
380
381		#[test]
382		fn test_take() {
383			let bump = BumpAlloc::new();
384			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, &bump);
385			v.push(10);
386			v.push(20);
387			v.push(30);
388
389			let taken = DataVec::take(&v, 2);
390			assert_eq!(taken.as_slice(), &[10, 20]);
391
392			let taken_all = DataVec::take(&v, 5);
393			assert_eq!(taken_all.as_slice(), &[10, 20, 30]);
394		}
395	}
396
397	mod bump_bitvec {
398		use super::*;
399
400		#[test]
401		fn test_push_and_get() {
402			let bump = BumpAlloc::new();
403			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
404			bv.push(true);
405			bv.push(false);
406			bv.push(true);
407
408			assert_eq!(DataBitVec::len(&bv), 3);
409			assert!(DataBitVec::get(&bv, 0));
410			assert!(!DataBitVec::get(&bv, 1));
411			assert!(DataBitVec::get(&bv, 2));
412		}
413
414		#[test]
415		fn test_set() {
416			let bump = BumpAlloc::new();
417			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
418			bv.push(false);
419			bv.push(false);
420			bv.push(false);
421
422			DataBitVec::set(&mut bv, 1, true);
423			assert!(!DataBitVec::get(&bv, 0));
424			assert!(DataBitVec::get(&bv, 1));
425			assert!(!DataBitVec::get(&bv, 2));
426
427			DataBitVec::set(&mut bv, 1, false);
428			assert!(!DataBitVec::get(&bv, 1));
429		}
430
431		#[test]
432		fn test_cross_byte_boundary() {
433			let bump = BumpAlloc::new();
434			let mut bv = BumpBitVec::with_capacity_in(16, &bump);
435			for i in 0..17 {
436				bv.push(i % 3 == 0);
437			}
438
439			assert_eq!(DataBitVec::len(&bv), 17);
440			for i in 0..17 {
441				assert_eq!(DataBitVec::get(&bv, i), i % 3 == 0, "mismatch at bit {}", i);
442			}
443		}
444
445		#[test]
446		fn test_count_ones() {
447			let bump = BumpAlloc::new();
448			let mut bv = BumpBitVec::with_capacity_in(16, &bump);
449			bv.push(true);
450			bv.push(false);
451			bv.push(true);
452			bv.push(false);
453			bv.push(true);
454
455			assert_eq!(DataBitVec::count_ones(&bv), 3);
456			assert_eq!(DataBitVec::count_zeros(&bv), 2);
457		}
458
459		#[test]
460		fn test_count_ones_cross_byte() {
461			let bump = BumpAlloc::new();
462			let mut bv = BumpBitVec::with_capacity_in(16, &bump);
463			for i in 0..17 {
464				bv.push(i % 3 == 0);
465			}
466			let expected = (0..17).filter(|&i| i % 3 == 0).count();
467			assert_eq!(DataBitVec::count_ones(&bv), expected);
468		}
469
470		#[test]
471		fn test_extend_from() {
472			let bump = BumpAlloc::new();
473			let mut bv1 = BumpBitVec::with_capacity_in(8, &bump);
474			bv1.push(true);
475			bv1.push(false);
476
477			let mut bv2 = BumpBitVec::with_capacity_in(8, &bump);
478			bv2.push(false);
479			bv2.push(true);
480
481			DataBitVec::extend_from(&mut bv1, &bv2);
482			assert_eq!(DataBitVec::len(&bv1), 4);
483			assert!(DataBitVec::get(&bv1, 0));
484			assert!(!DataBitVec::get(&bv1, 1));
485			assert!(!DataBitVec::get(&bv1, 2));
486			assert!(DataBitVec::get(&bv1, 3));
487		}
488
489		#[test]
490		fn test_clone() {
491			let bump = BumpAlloc::new();
492			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
493			bv.push(true);
494			bv.push(false);
495			bv.push(true);
496
497			let bv2 = bv.clone();
498			assert_eq!(DataBitVec::len(&bv2), 3);
499			assert!(DataBitVec::get(&bv2, 0));
500			assert!(!DataBitVec::get(&bv2, 1));
501			assert!(DataBitVec::get(&bv2, 2));
502		}
503
504		#[test]
505		fn test_clear() {
506			let bump = BumpAlloc::new();
507			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
508			bv.push(true);
509			bv.push(false);
510			DataBitVec::clear(&mut bv);
511			assert_eq!(DataBitVec::len(&bv), 0);
512		}
513
514		#[test]
515		fn test_spawn() {
516			let bump = BumpAlloc::new();
517			let bv = BumpBitVec::with_capacity_in(8, &bump);
518			let bv2 = DataBitVec::spawn(&bv, 16);
519			assert_eq!(DataBitVec::len(&bv2), 0);
520			assert!(DataBitVec::capacity(&bv2) >= 16);
521		}
522
523		#[test]
524		fn test_iter() {
525			let bump = BumpAlloc::new();
526			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
527			bv.push(true);
528			bv.push(false);
529			bv.push(true);
530			bv.push(false);
531
532			let collected: Vec<bool> = DataBitVec::iter(&bv).collect();
533			assert_eq!(collected, vec![true, false, true, false]);
534		}
535
536		#[test]
537		fn test_eq() {
538			let bump = BumpAlloc::new();
539			let mut bv1 = BumpBitVec::with_capacity_in(8, &bump);
540			bv1.push(true);
541			bv1.push(false);
542
543			let mut bv2 = BumpBitVec::with_capacity_in(8, &bump);
544			bv2.push(true);
545			bv2.push(false);
546
547			assert_eq!(bv1, bv2);
548
549			bv2.push(true);
550			assert_ne!(bv1, bv2);
551		}
552
553		#[test]
554		fn test_take() {
555			let bump = BumpAlloc::new();
556			let mut bv = BumpBitVec::with_capacity_in(8, &bump);
557			bv.push(true);
558			bv.push(false);
559			bv.push(true);
560
561			let taken = DataBitVec::take(&bv, 2);
562			assert_eq!(DataBitVec::len(&taken), 2);
563			assert!(DataBitVec::get(&taken, 0));
564			assert!(!DataBitVec::get(&taken, 1));
565
566			let taken_all = DataBitVec::take(&bv, 5);
567			assert_eq!(DataBitVec::len(&taken_all), 3);
568		}
569	}
570
571	mod bump_storage {
572		use reifydb_type::value::container::{bool::BoolContainer, number::NumberContainer};
573
574		use super::*;
575
576		#[test]
577		fn test_number_container_with_bump() {
578			let bump_alloc = BumpAlloc::new();
579			let data = BumpVec::with_capacity_in(4, &bump_alloc);
580			let mut container: NumberContainer<i32, Bump<'_>> = NumberContainer::from_parts(data);
581
582			container.push(42);
583			container.push(99);
584
585			assert_eq!(container.len(), 2);
586			assert_eq!(container.get(0), Some(&42));
587			assert_eq!(container.get(1), Some(&99));
588		}
589
590		#[test]
591		fn test_bool_container_with_bump() {
592			let bump_alloc = BumpAlloc::new();
593			let data = BumpBitVec::with_capacity_in(4, &bump_alloc);
594			let mut container: BoolContainer<Bump<'_>> = BoolContainer::from_parts(data);
595
596			container.push(true);
597			container.push(false);
598			container.push_default();
599
600			assert_eq!(container.len(), 3);
601			assert_eq!(container.get(0), Some(true));
602			assert_eq!(container.get(1), Some(false));
603			// push_default on a bare container pushes false;
604			// nullability is tracked by the Option wrapper at the ColumnBuffer level.
605			assert_eq!(container.get(2), Some(false));
606		}
607
608		#[test]
609		fn test_query_arena_reset() {
610			let mut arena = QueryArena::new();
611			let bump = arena.bump();
612			let mut v: BumpVec<i32> = BumpVec::with_capacity_in(4, bump);
613			v.push(1);
614			v.push(2);
615			assert_eq!(DataVec::len(&v), 2);
616
617			// After reset, the arena memory is reclaimed
618			drop(v);
619			arena.reset();
620
621			let bump = arena.bump();
622			let v2: BumpVec<i32> = BumpVec::with_capacity_in(4, bump);
623			assert_eq!(DataVec::len(&v2), 0);
624		}
625	}
626}
reifydb_engine/arena/mod.rs

reifydb_engine/arena/
mod.rs
