prax_query/mem_optimize/
arena.rs

1//! Typed arena allocation for query builder chains.
2//!
3//! This module provides arena-based allocation for efficient query construction
4//! with minimal heap allocations.
5//!
6//! # Benefits
7//!
8//! - **Batch deallocation**: All allocations freed at once when scope ends
9//! - **Cache-friendly**: Contiguous memory allocation
10//! - **Fast allocation**: O(1) bump pointer allocation
11//! - **No fragmentation**: No individual deallocation overhead
12//!
13//! # Example
14//!
15//! ```rust
16//! use prax_query::mem_optimize::arena::QueryArena;
17//!
18//! let arena = QueryArena::new();
19//!
20//! // Build query within arena scope
21//! let sql = arena.scope(|scope| {
22//!     let filter = scope.eq("status", "active");
23//!     let filter = scope.and(vec![
24//!         filter,
25//!         scope.gt("age", 18),
26//!     ]);
27//!     scope.build_select("users", filter)
28//! });
29//!
30//! // Arena memory freed, but sql String is owned
31//! ```
32
33use bumpalo::Bump;
34use std::cell::Cell;
35use std::fmt::Write;
36
37use super::interning::InternedStr;
38
39// ============================================================================
40// Query Arena
41// ============================================================================
42
43/// Arena allocator for query building.
44///
45/// Provides fast allocation with batch deallocation when the scope ends.
46pub struct QueryArena {
47    bump: Bump,
48    stats: Cell<ArenaStats>,
49}
50
51impl QueryArena {
52    /// Create a new query arena with default capacity.
53    pub fn new() -> Self {
54        Self {
55            bump: Bump::new(),
56            stats: Cell::new(ArenaStats::default()),
57        }
58    }
59
60    /// Create an arena with specified initial capacity.
61    pub fn with_capacity(capacity: usize) -> Self {
62        Self {
63            bump: Bump::with_capacity(capacity),
64            stats: Cell::new(ArenaStats::default()),
65        }
66    }
67
68    /// Execute a closure with an arena scope.
69    ///
70    /// The scope provides allocation methods. All allocations are valid
71    /// within the closure and freed when it returns.
72    pub fn scope<F, R>(&self, f: F) -> R
73    where
74        F: FnOnce(&ArenaScope<'_>) -> R,
75    {
76        let scope = ArenaScope::new(&self.bump, &self.stats);
77        f(&scope)
78    }
79
80    /// Reset the arena for reuse.
81    ///
82    /// This is O(1) - just resets the bump pointer.
83    pub fn reset(&mut self) {
84        self.bump.reset();
85        self.stats.set(ArenaStats::default());
86    }
87
88    /// Get the number of bytes currently allocated.
89    pub fn allocated_bytes(&self) -> usize {
90        self.bump.allocated_bytes()
91    }
92
93    /// Get arena statistics.
94    pub fn stats(&self) -> ArenaStats {
95        self.stats.get()
96    }
97}
98
99impl Default for QueryArena {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105// ============================================================================
106// Arena Scope
107// ============================================================================
108
109/// A scope for allocating within an arena.
110///
111/// All allocations made through this scope are freed when the scope ends.
112pub struct ArenaScope<'a> {
113    bump: &'a Bump,
114    stats: &'a Cell<ArenaStats>,
115}
116
117impl<'a> ArenaScope<'a> {
118    fn new(bump: &'a Bump, stats: &'a Cell<ArenaStats>) -> Self {
119        Self { bump, stats }
120    }
121
122    fn record_alloc(&self, bytes: usize) {
123        let mut s = self.stats.get();
124        s.allocations += 1;
125        s.total_bytes += bytes;
126        self.stats.set(s);
127    }
128
129    /// Allocate a string in the arena.
130    #[inline]
131    pub fn alloc_str(&self, s: &str) -> &'a str {
132        self.record_alloc(s.len());
133        self.bump.alloc_str(s)
134    }
135
136    /// Allocate a slice in the arena.
137    #[inline]
138    pub fn alloc_slice<T: Copy>(&self, slice: &[T]) -> &'a [T] {
139        self.record_alloc(std::mem::size_of_val(slice));
140        self.bump.alloc_slice_copy(slice)
141    }
142
143    /// Allocate a slice from an iterator.
144    #[inline]
145    pub fn alloc_slice_iter<T, I>(&self, iter: I) -> &'a [T]
146    where
147        I: IntoIterator<Item = T>,
148        I::IntoIter: ExactSizeIterator,
149    {
150        let iter = iter.into_iter();
151        self.record_alloc(iter.len() * std::mem::size_of::<T>());
152        self.bump.alloc_slice_fill_iter(iter)
153    }
154
155    /// Allocate a single value in the arena.
156    #[inline]
157    pub fn alloc<T>(&self, value: T) -> &'a T {
158        self.record_alloc(std::mem::size_of::<T>());
159        self.bump.alloc(value)
160    }
161
162    /// Allocate a mutable value in the arena.
163    #[inline]
164    pub fn alloc_mut<T>(&self, value: T) -> &'a mut T {
165        self.record_alloc(std::mem::size_of::<T>());
166        self.bump.alloc(value)
167    }
168
169    // ========================================================================
170    // Filter Construction
171    // ========================================================================
172
173    /// Create an equality filter.
174    #[inline]
175    pub fn eq<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
176        ScopedFilter::Equals(self.alloc_str(field), value.into())
177    }
178
179    /// Create a not-equals filter.
180    #[inline]
181    pub fn ne<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
182        ScopedFilter::NotEquals(self.alloc_str(field), value.into())
183    }
184
185    /// Create a less-than filter.
186    #[inline]
187    pub fn lt<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
188        ScopedFilter::Lt(self.alloc_str(field), value.into())
189    }
190
191    /// Create a less-than-or-equal filter.
192    #[inline]
193    pub fn lte<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
194        ScopedFilter::Lte(self.alloc_str(field), value.into())
195    }
196
197    /// Create a greater-than filter.
198    #[inline]
199    pub fn gt<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
200        ScopedFilter::Gt(self.alloc_str(field), value.into())
201    }
202
203    /// Create a greater-than-or-equal filter.
204    #[inline]
205    pub fn gte<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
206        ScopedFilter::Gte(self.alloc_str(field), value.into())
207    }
208
209    /// Create an IN filter.
210    #[inline]
211    pub fn is_in(&self, field: &str, values: Vec<ScopedValue<'a>>) -> ScopedFilter<'a> {
212        ScopedFilter::In(self.alloc_str(field), self.alloc_slice_iter(values))
213    }
214
215    /// Create a NOT IN filter.
216    #[inline]
217    pub fn not_in(&self, field: &str, values: Vec<ScopedValue<'a>>) -> ScopedFilter<'a> {
218        ScopedFilter::NotIn(self.alloc_str(field), self.alloc_slice_iter(values))
219    }
220
221    /// Create a CONTAINS filter.
222    #[inline]
223    pub fn contains<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
224        ScopedFilter::Contains(self.alloc_str(field), value.into())
225    }
226
227    /// Create a STARTS WITH filter.
228    #[inline]
229    pub fn starts_with<V: Into<ScopedValue<'a>>>(
230        &self,
231        field: &str,
232        value: V,
233    ) -> ScopedFilter<'a> {
234        ScopedFilter::StartsWith(self.alloc_str(field), value.into())
235    }
236
237    /// Create an ENDS WITH filter.
238    #[inline]
239    pub fn ends_with<V: Into<ScopedValue<'a>>>(&self, field: &str, value: V) -> ScopedFilter<'a> {
240        ScopedFilter::EndsWith(self.alloc_str(field), value.into())
241    }
242
243    /// Create an IS NULL filter.
244    #[inline]
245    pub fn is_null(&self, field: &str) -> ScopedFilter<'a> {
246        ScopedFilter::IsNull(self.alloc_str(field))
247    }
248
249    /// Create an IS NOT NULL filter.
250    #[inline]
251    pub fn is_not_null(&self, field: &str) -> ScopedFilter<'a> {
252        ScopedFilter::IsNotNull(self.alloc_str(field))
253    }
254
255    /// Combine filters with AND.
256    #[inline]
257    pub fn and(&self, filters: Vec<ScopedFilter<'a>>) -> ScopedFilter<'a> {
258        // Filter out None filters
259        let filters: Vec<_> = filters
260            .into_iter()
261            .filter(|f| !matches!(f, ScopedFilter::None))
262            .collect();
263
264        match filters.len() {
265            0 => ScopedFilter::None,
266            1 => filters.into_iter().next().unwrap(),
267            _ => ScopedFilter::And(self.alloc_slice_iter(filters)),
268        }
269    }
270
271    /// Combine filters with OR.
272    #[inline]
273    pub fn or(&self, filters: Vec<ScopedFilter<'a>>) -> ScopedFilter<'a> {
274        let filters: Vec<_> = filters
275            .into_iter()
276            .filter(|f| !matches!(f, ScopedFilter::None))
277            .collect();
278
279        match filters.len() {
280            0 => ScopedFilter::None,
281            1 => filters.into_iter().next().unwrap(),
282            _ => ScopedFilter::Or(self.alloc_slice_iter(filters)),
283        }
284    }
285
286    /// Negate a filter.
287    #[inline]
288    pub fn not(&self, filter: ScopedFilter<'a>) -> ScopedFilter<'a> {
289        if matches!(filter, ScopedFilter::None) {
290            return ScopedFilter::None;
291        }
292        ScopedFilter::Not(self.alloc(filter))
293    }
294
295    // ========================================================================
296    // Query Building
297    // ========================================================================
298
299    /// Build a SELECT query string.
300    pub fn build_select(&self, table: &str, filter: ScopedFilter<'a>) -> String {
301        let mut sql = String::with_capacity(128);
302        sql.push_str("SELECT * FROM ");
303        sql.push_str(table);
304
305        if !matches!(filter, ScopedFilter::None) {
306            sql.push_str(" WHERE ");
307            filter.write_sql(&mut sql, &mut 1);
308        }
309
310        sql
311    }
312
313    /// Build a SELECT query with specific columns.
314    pub fn build_select_columns(
315        &self,
316        table: &str,
317        columns: &[&str],
318        filter: ScopedFilter<'a>,
319    ) -> String {
320        let mut sql = String::with_capacity(128);
321        sql.push_str("SELECT ");
322
323        for (i, col) in columns.iter().enumerate() {
324            if i > 0 {
325                sql.push_str(", ");
326            }
327            sql.push_str(col);
328        }
329
330        sql.push_str(" FROM ");
331        sql.push_str(table);
332
333        if !matches!(filter, ScopedFilter::None) {
334            sql.push_str(" WHERE ");
335            filter.write_sql(&mut sql, &mut 1);
336        }
337
338        sql
339    }
340
341    /// Build a complete query with all parts.
342    pub fn build_query(&self, query: &ScopedQuery<'a>) -> String {
343        let mut sql = String::with_capacity(256);
344
345        // SELECT
346        sql.push_str("SELECT ");
347        if query.columns.is_empty() {
348            sql.push('*');
349        } else {
350            for (i, col) in query.columns.iter().enumerate() {
351                if i > 0 {
352                    sql.push_str(", ");
353                }
354                sql.push_str(col);
355            }
356        }
357
358        // FROM
359        sql.push_str(" FROM ");
360        sql.push_str(query.table);
361
362        // WHERE
363        if !matches!(query.filter, ScopedFilter::None) {
364            sql.push_str(" WHERE ");
365            query.filter.write_sql(&mut sql, &mut 1);
366        }
367
368        // ORDER BY
369        if !query.order_by.is_empty() {
370            sql.push_str(" ORDER BY ");
371            for (i, (col, dir)) in query.order_by.iter().enumerate() {
372                if i > 0 {
373                    sql.push_str(", ");
374                }
375                sql.push_str(col);
376                sql.push(' ');
377                sql.push_str(dir);
378            }
379        }
380
381        // LIMIT
382        if let Some(limit) = query.limit {
383            write!(sql, " LIMIT {}", limit).unwrap();
384        }
385
386        // OFFSET
387        if let Some(offset) = query.offset {
388            write!(sql, " OFFSET {}", offset).unwrap();
389        }
390
391        sql
392    }
393
394    /// Create a new query builder.
395    pub fn query(&self, table: &str) -> ScopedQuery<'a> {
396        ScopedQuery {
397            table: self.alloc_str(table),
398            columns: &[],
399            filter: ScopedFilter::None,
400            order_by: &[],
401            limit: None,
402            offset: None,
403        }
404    }
405}
406
407// ============================================================================
408// Scoped Filter
409// ============================================================================
410
411/// A filter allocated within an arena scope.
412#[derive(Debug, Clone)]
413pub enum ScopedFilter<'a> {
414    /// No filter.
415    None,
416    /// Equality.
417    Equals(&'a str, ScopedValue<'a>),
418    /// Not equals.
419    NotEquals(&'a str, ScopedValue<'a>),
420    /// Less than.
421    Lt(&'a str, ScopedValue<'a>),
422    /// Less than or equal.
423    Lte(&'a str, ScopedValue<'a>),
424    /// Greater than.
425    Gt(&'a str, ScopedValue<'a>),
426    /// Greater than or equal.
427    Gte(&'a str, ScopedValue<'a>),
428    /// In list.
429    In(&'a str, &'a [ScopedValue<'a>]),
430    /// Not in list.
431    NotIn(&'a str, &'a [ScopedValue<'a>]),
432    /// Contains.
433    Contains(&'a str, ScopedValue<'a>),
434    /// Starts with.
435    StartsWith(&'a str, ScopedValue<'a>),
436    /// Ends with.
437    EndsWith(&'a str, ScopedValue<'a>),
438    /// Is null.
439    IsNull(&'a str),
440    /// Is not null.
441    IsNotNull(&'a str),
442    /// And.
443    And(&'a [ScopedFilter<'a>]),
444    /// Or.
445    Or(&'a [ScopedFilter<'a>]),
446    /// Not.
447    Not(&'a ScopedFilter<'a>),
448}
449
450impl<'a> ScopedFilter<'a> {
451    /// Write SQL to a string buffer.
452    pub fn write_sql(&self, buf: &mut String, param_idx: &mut usize) {
453        match self {
454            ScopedFilter::None => {}
455            ScopedFilter::Equals(field, _) => {
456                write!(buf, "{} = ${}", field, param_idx).unwrap();
457                *param_idx += 1;
458            }
459            ScopedFilter::NotEquals(field, _) => {
460                write!(buf, "{} != ${}", field, param_idx).unwrap();
461                *param_idx += 1;
462            }
463            ScopedFilter::Lt(field, _) => {
464                write!(buf, "{} < ${}", field, param_idx).unwrap();
465                *param_idx += 1;
466            }
467            ScopedFilter::Lte(field, _) => {
468                write!(buf, "{} <= ${}", field, param_idx).unwrap();
469                *param_idx += 1;
470            }
471            ScopedFilter::Gt(field, _) => {
472                write!(buf, "{} > ${}", field, param_idx).unwrap();
473                *param_idx += 1;
474            }
475            ScopedFilter::Gte(field, _) => {
476                write!(buf, "{} >= ${}", field, param_idx).unwrap();
477                *param_idx += 1;
478            }
479            ScopedFilter::In(field, values) => {
480                write!(buf, "{} IN (", field).unwrap();
481                for (i, _) in values.iter().enumerate() {
482                    if i > 0 {
483                        buf.push_str(", ");
484                    }
485                    write!(buf, "${}", param_idx).unwrap();
486                    *param_idx += 1;
487                }
488                buf.push(')');
489            }
490            ScopedFilter::NotIn(field, values) => {
491                write!(buf, "{} NOT IN (", field).unwrap();
492                for (i, _) in values.iter().enumerate() {
493                    if i > 0 {
494                        buf.push_str(", ");
495                    }
496                    write!(buf, "${}", param_idx).unwrap();
497                    *param_idx += 1;
498                }
499                buf.push(')');
500            }
501            ScopedFilter::Contains(field, _) => {
502                write!(buf, "{} LIKE ${}", field, param_idx).unwrap();
503                *param_idx += 1;
504            }
505            ScopedFilter::StartsWith(field, _) => {
506                write!(buf, "{} LIKE ${}", field, param_idx).unwrap();
507                *param_idx += 1;
508            }
509            ScopedFilter::EndsWith(field, _) => {
510                write!(buf, "{} LIKE ${}", field, param_idx).unwrap();
511                *param_idx += 1;
512            }
513            ScopedFilter::IsNull(field) => {
514                write!(buf, "{} IS NULL", field).unwrap();
515            }
516            ScopedFilter::IsNotNull(field) => {
517                write!(buf, "{} IS NOT NULL", field).unwrap();
518            }
519            ScopedFilter::And(filters) => {
520                buf.push('(');
521                for (i, filter) in filters.iter().enumerate() {
522                    if i > 0 {
523                        buf.push_str(" AND ");
524                    }
525                    filter.write_sql(buf, param_idx);
526                }
527                buf.push(')');
528            }
529            ScopedFilter::Or(filters) => {
530                buf.push('(');
531                for (i, filter) in filters.iter().enumerate() {
532                    if i > 0 {
533                        buf.push_str(" OR ");
534                    }
535                    filter.write_sql(buf, param_idx);
536                }
537                buf.push(')');
538            }
539            ScopedFilter::Not(filter) => {
540                buf.push_str("NOT (");
541                filter.write_sql(buf, param_idx);
542                buf.push(')');
543            }
544        }
545    }
546}
547
548// ============================================================================
549// Scoped Value
550// ============================================================================
551
552/// A value allocated within an arena scope.
553#[derive(Debug, Clone)]
554pub enum ScopedValue<'a> {
555    /// Null.
556    Null,
557    /// Boolean.
558    Bool(bool),
559    /// Integer.
560    Int(i64),
561    /// Float.
562    Float(f64),
563    /// String (borrowed from arena).
564    String(&'a str),
565    /// Interned string (shared reference).
566    Interned(InternedStr),
567}
568
569impl<'a> From<bool> for ScopedValue<'a> {
570    fn from(v: bool) -> Self {
571        ScopedValue::Bool(v)
572    }
573}
574
575impl<'a> From<i32> for ScopedValue<'a> {
576    fn from(v: i32) -> Self {
577        ScopedValue::Int(v as i64)
578    }
579}
580
581impl<'a> From<i64> for ScopedValue<'a> {
582    fn from(v: i64) -> Self {
583        ScopedValue::Int(v)
584    }
585}
586
587impl<'a> From<f64> for ScopedValue<'a> {
588    fn from(v: f64) -> Self {
589        ScopedValue::Float(v)
590    }
591}
592
593impl<'a> From<&'a str> for ScopedValue<'a> {
594    fn from(v: &'a str) -> Self {
595        ScopedValue::String(v)
596    }
597}
598
599impl<'a> From<InternedStr> for ScopedValue<'a> {
600    fn from(v: InternedStr) -> Self {
601        ScopedValue::Interned(v)
602    }
603}
604
605// ============================================================================
606// Scoped Query
607// ============================================================================
608
609/// A query being built within an arena scope.
610#[derive(Debug, Clone)]
611pub struct ScopedQuery<'a> {
612    /// Table name.
613    pub table: &'a str,
614    /// Columns to select.
615    pub columns: &'a [&'a str],
616    /// Filter.
617    pub filter: ScopedFilter<'a>,
618    /// Order by clauses.
619    pub order_by: &'a [(&'a str, &'a str)],
620    /// Limit.
621    pub limit: Option<usize>,
622    /// Offset.
623    pub offset: Option<usize>,
624}
625
626impl<'a> ScopedQuery<'a> {
627    /// Set columns to select.
628    pub fn select(mut self, columns: &'a [&'a str]) -> Self {
629        self.columns = columns;
630        self
631    }
632
633    /// Set filter.
634    pub fn filter(mut self, filter: ScopedFilter<'a>) -> Self {
635        self.filter = filter;
636        self
637    }
638
639    /// Set order by.
640    pub fn order_by(mut self, order_by: &'a [(&'a str, &'a str)]) -> Self {
641        self.order_by = order_by;
642        self
643    }
644
645    /// Set limit.
646    pub fn limit(mut self, limit: usize) -> Self {
647        self.limit = Some(limit);
648        self
649    }
650
651    /// Set offset.
652    pub fn offset(mut self, offset: usize) -> Self {
653        self.offset = Some(offset);
654        self
655    }
656}
657
658// ============================================================================
659// Statistics
660// ============================================================================
661
662/// Statistics for arena usage.
663#[derive(Debug, Clone, Copy, Default)]
664pub struct ArenaStats {
665    /// Number of allocations.
666    pub allocations: usize,
667    /// Total bytes allocated.
668    pub total_bytes: usize,
669}
670
671#[cfg(test)]
672mod tests {
673    use super::*;
674
675    #[test]
676    fn test_arena_basic_filter() {
677        let arena = QueryArena::new();
678
679        let sql = arena.scope(|scope| scope.build_select("users", scope.eq("id", 42)));
680
681        assert!(sql.contains("SELECT * FROM users"));
682        assert!(sql.contains("WHERE"));
683        assert!(sql.contains("id = $1"));
684    }
685
686    #[test]
687    fn test_arena_complex_filter() {
688        let arena = QueryArena::new();
689
690        let sql = arena.scope(|scope| {
691            let filter = scope.and(vec![
692                scope.eq("active", true),
693                scope.or(vec![scope.gt("age", 18), scope.is_not_null("verified_at")]),
694            ]);
695            scope.build_select("users", filter)
696        });
697
698        assert!(sql.contains("AND"));
699        assert!(sql.contains("OR"));
700    }
701
702    #[test]
703    fn test_arena_reset() {
704        let mut arena = QueryArena::with_capacity(1024);
705
706        // Use arena
707        let _ = arena.scope(|scope| scope.build_select("users", scope.eq("id", 1)));
708        let bytes1 = arena.allocated_bytes();
709
710        // Reset
711        arena.reset();
712
713        // Use again
714        let _ = arena.scope(|scope| scope.build_select("posts", scope.eq("id", 2)));
715        let bytes2 = arena.allocated_bytes();
716
717        // Should be similar (reusing memory)
718        assert!(bytes2 <= bytes1 * 2);
719    }
720
721    #[test]
722    fn test_arena_query_builder() {
723        let arena = QueryArena::new();
724
725        let sql = arena.scope(|scope| {
726            let query = scope
727                .query("users")
728                .filter(scope.eq("active", true))
729                .limit(10)
730                .offset(20);
731            scope.build_query(&query)
732        });
733
734        assert!(sql.contains("SELECT * FROM users"));
735        assert!(sql.contains("LIMIT 10"));
736        assert!(sql.contains("OFFSET 20"));
737    }
738
739    #[test]
740    fn test_arena_in_filter() {
741        let arena = QueryArena::new();
742
743        let sql = arena.scope(|scope| {
744            let filter = scope.is_in(
745                "status",
746                vec!["pending".into(), "processing".into(), "completed".into()],
747            );
748            scope.build_select("orders", filter)
749        });
750
751        assert!(sql.contains("IN"));
752        assert!(sql.contains("$1"));
753        assert!(sql.contains("$2"));
754        assert!(sql.contains("$3"));
755    }
756
757    #[test]
758    fn test_arena_stats() {
759        let arena = QueryArena::new();
760
761        arena.scope(|scope| {
762            let _ = scope.alloc_str("test string");
763            let _ = scope.alloc_str("another string");
764        });
765
766        let stats = arena.stats();
767        assert_eq!(stats.allocations, 2);
768    }
769}
770
prax_query/mem_optimize/arena.rs

prax_query/mem_optimize/
arena.rs
