prax_query/
memory.rs

1//! Memory optimization utilities for prax-query.
2//!
3//! This module provides memory-efficient alternatives and pooling mechanisms
4//! to reduce allocation counts and peak memory usage.
5//!
6//! # Optimization Strategies
7//!
8//! 1. **Object Pooling**: Reuse allocated objects instead of allocating new ones
9//! 2. **Compact Types**: Smaller type representations for common cases
10//! 3. **Inline Storage**: Store small data inline to avoid heap allocation
11//! 4. **String Deduplication**: Share identical strings via interning
12//!
13//! # Example
14//!
15//! ```rust
16//! use prax_query::memory::{StringPool, CompactFilter};
17//! use std::sync::Arc;
18//!
19//! // Reuse strings from a pool
20//! let pool = StringPool::new();
21//! let s1 = pool.intern("id");
22//! let s2 = pool.intern("id");
23//! // s1 and s2 point to the same allocation
24//!
25//! // Use compact filters for simple cases
26//! let filter = CompactFilter::eq_int(Arc::from("id"), 42);
27//! ```
28
29use parking_lot::Mutex;
30use std::collections::HashSet;
31use std::sync::Arc;
32use tracing::info;
33
34// ============================================================================
35// String Pool
36// ============================================================================
37
38/// A thread-safe pool for interning strings.
39///
40/// Reduces memory usage by ensuring only one copy of each unique string exists.
41#[derive(Debug, Default)]
42pub struct StringPool {
43    strings: Mutex<HashSet<Arc<str>>>,
44}
45
46impl StringPool {
47    /// Create a new empty string pool.
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// Create a pool with pre-allocated capacity.
53    pub fn with_capacity(capacity: usize) -> Self {
54        info!(capacity, "StringPool initialized");
55        Self {
56            strings: Mutex::new(HashSet::with_capacity(capacity)),
57        }
58    }
59
60    /// Intern a string, returning a shared reference.
61    pub fn intern(&self, s: &str) -> Arc<str> {
62        let mut strings = self.strings.lock();
63
64        // Check if already interned
65        if let Some(existing) = strings.get(s) {
66            return Arc::clone(existing);
67        }
68
69        // Intern new string
70        let arc: Arc<str> = Arc::from(s);
71        strings.insert(Arc::clone(&arc));
72        arc
73    }
74
75    /// Get the number of interned strings.
76    pub fn len(&self) -> usize {
77        self.strings.lock().len()
78    }
79
80    /// Check if the pool is empty.
81    pub fn is_empty(&self) -> bool {
82        self.strings.lock().is_empty()
83    }
84
85    /// Clear all interned strings.
86    pub fn clear(&self) {
87        self.strings.lock().clear();
88    }
89
90    /// Get memory statistics.
91    pub fn stats(&self) -> PoolStats {
92        let strings = self.strings.lock();
93        let count = strings.len();
94        let total_bytes: usize = strings.iter().map(|s| s.len()).sum();
95        PoolStats {
96            count,
97            total_bytes,
98            avg_bytes: if count > 0 { total_bytes / count } else { 0 },
99        }
100    }
101}
102
103/// Statistics for a pool.
104#[derive(Debug, Clone, Default)]
105pub struct PoolStats {
106    /// Number of items in the pool.
107    pub count: usize,
108    /// Total bytes used.
109    pub total_bytes: usize,
110    /// Average bytes per item.
111    pub avg_bytes: usize,
112}
113
114// ============================================================================
115// Compact Filter Types
116// ============================================================================
117
118/// A compact filter representation for simple equality filters.
119///
120/// Uses less memory than the full `Filter` enum for common cases:
121/// - Equality on integer fields: 16 bytes vs ~40 bytes
122/// - Equality on string fields with short values: ~32 bytes vs ~56 bytes
123#[derive(Debug, Clone, PartialEq)]
124pub enum CompactFilter {
125    /// Equality on integer field.
126    EqInt {
127        /// Field name (interned).
128        field: Arc<str>,
129        /// Integer value.
130        value: i64,
131    },
132    /// Equality on boolean field.
133    EqBool {
134        /// Field name (interned).
135        field: Arc<str>,
136        /// Boolean value.
137        value: bool,
138    },
139    /// Equality on string field.
140    EqStr {
141        /// Field name (interned).
142        field: Arc<str>,
143        /// String value.
144        value: Arc<str>,
145    },
146    /// IS NULL check.
147    IsNull {
148        /// Field name (interned).
149        field: Arc<str>,
150    },
151    /// IS NOT NULL check.
152    IsNotNull {
153        /// Field name (interned).
154        field: Arc<str>,
155    },
156    /// Greater than on integer field.
157    GtInt {
158        /// Field name (interned).
159        field: Arc<str>,
160        /// Integer value.
161        value: i64,
162    },
163    /// Less than on integer field.
164    LtInt {
165        /// Field name (interned).
166        field: Arc<str>,
167        /// Integer value.
168        value: i64,
169    },
170    /// AND of two compact filters.
171    And(Box<CompactFilter>, Box<CompactFilter>),
172    /// OR of two compact filters.
173    Or(Box<CompactFilter>, Box<CompactFilter>),
174}
175
176impl CompactFilter {
177    /// Create an equality filter on an integer field.
178    #[inline]
179    pub fn eq_int(field: impl Into<Arc<str>>, value: i64) -> Self {
180        Self::EqInt {
181            field: field.into(),
182            value,
183        }
184    }
185
186    /// Create an equality filter on a boolean field.
187    #[inline]
188    pub fn eq_bool(field: impl Into<Arc<str>>, value: bool) -> Self {
189        Self::EqBool {
190            field: field.into(),
191            value,
192        }
193    }
194
195    /// Create an equality filter on a string field.
196    #[inline]
197    pub fn eq_str(field: impl Into<Arc<str>>, value: impl Into<Arc<str>>) -> Self {
198        Self::EqStr {
199            field: field.into(),
200            value: value.into(),
201        }
202    }
203
204    /// Create an IS NULL filter.
205    #[inline]
206    pub fn is_null(field: impl Into<Arc<str>>) -> Self {
207        Self::IsNull {
208            field: field.into(),
209        }
210    }
211
212    /// Create an IS NOT NULL filter.
213    #[inline]
214    pub fn is_not_null(field: impl Into<Arc<str>>) -> Self {
215        Self::IsNotNull {
216            field: field.into(),
217        }
218    }
219
220    /// Create a greater-than filter on an integer field.
221    #[inline]
222    pub fn gt_int(field: impl Into<Arc<str>>, value: i64) -> Self {
223        Self::GtInt {
224            field: field.into(),
225            value,
226        }
227    }
228
229    /// Create a less-than filter on an integer field.
230    #[inline]
231    pub fn lt_int(field: impl Into<Arc<str>>, value: i64) -> Self {
232        Self::LtInt {
233            field: field.into(),
234            value,
235        }
236    }
237
238    /// Combine with another filter using AND.
239    #[inline]
240    pub fn and(self, other: Self) -> Self {
241        Self::And(Box::new(self), Box::new(other))
242    }
243
244    /// Combine with another filter using OR.
245    #[inline]
246    pub fn or(self, other: Self) -> Self {
247        Self::Or(Box::new(self), Box::new(other))
248    }
249
250    /// Convert to SQL condition string for PostgreSQL.
251    pub fn to_sql_postgres(&self, param_offset: &mut usize) -> String {
252        match self {
253            Self::EqInt { field, .. } => {
254                *param_offset += 1;
255                format!("{} = ${}", field, *param_offset)
256            }
257            Self::EqBool { field, .. } => {
258                *param_offset += 1;
259                format!("{} = ${}", field, *param_offset)
260            }
261            Self::EqStr { field, .. } => {
262                *param_offset += 1;
263                format!("{} = ${}", field, *param_offset)
264            }
265            Self::IsNull { field } => format!("{} IS NULL", field),
266            Self::IsNotNull { field } => format!("{} IS NOT NULL", field),
267            Self::GtInt { field, .. } => {
268                *param_offset += 1;
269                format!("{} > ${}", field, *param_offset)
270            }
271            Self::LtInt { field, .. } => {
272                *param_offset += 1;
273                format!("{} < ${}", field, *param_offset)
274            }
275            Self::And(left, right) => {
276                let left_sql = left.to_sql_postgres(param_offset);
277                let right_sql = right.to_sql_postgres(param_offset);
278                format!("({} AND {})", left_sql, right_sql)
279            }
280            Self::Or(left, right) => {
281                let left_sql = left.to_sql_postgres(param_offset);
282                let right_sql = right.to_sql_postgres(param_offset);
283                format!("({} OR {})", left_sql, right_sql)
284            }
285        }
286    }
287
288    /// Get the approximate size in bytes.
289    pub fn size_bytes(&self) -> usize {
290        match self {
291            Self::EqInt { .. } | Self::GtInt { .. } | Self::LtInt { .. } => 24, // Arc<str> + i64
292            Self::EqBool { .. } => 17,                                          // Arc<str> + bool
293            Self::EqStr { field, value } => 16 + field.len() + value.len(),
294            Self::IsNull { .. } | Self::IsNotNull { .. } => 16, // Arc<str>
295            Self::And(l, r) | Self::Or(l, r) => 16 + l.size_bytes() + r.size_bytes(),
296        }
297    }
298}
299
300// ============================================================================
301// Reusable Buffer Pool
302// ============================================================================
303
304/// A pool of reusable String buffers.
305///
306/// Reduces allocation by reusing String buffers for SQL generation.
307#[derive(Debug, Default)]
308pub struct BufferPool {
309    buffers: Mutex<Vec<String>>,
310    default_capacity: usize,
311}
312
313impl BufferPool {
314    /// Create a new buffer pool.
315    pub fn new() -> Self {
316        Self {
317            buffers: Mutex::new(Vec::new()),
318            default_capacity: 256,
319        }
320    }
321
322    /// Create a pool with custom default buffer capacity.
323    pub fn with_capacity(default_capacity: usize) -> Self {
324        info!(default_capacity, "BufferPool initialized");
325        Self {
326            buffers: Mutex::new(Vec::new()),
327            default_capacity,
328        }
329    }
330
331    /// Get a buffer from the pool or create a new one.
332    pub fn get(&self) -> PooledBuffer<'_> {
333        let buffer = self
334            .buffers
335            .lock()
336            .pop()
337            .unwrap_or_else(|| String::with_capacity(self.default_capacity));
338        PooledBuffer { buffer, pool: self }
339    }
340
341    /// Return a buffer to the pool.
342    fn return_buffer(&self, mut buffer: String) {
343        buffer.clear();
344        // Only keep reasonably sized buffers to avoid memory bloat
345        if buffer.capacity() <= 4096 {
346            self.buffers.lock().push(buffer);
347        }
348    }
349
350    /// Get the number of available buffers.
351    pub fn available(&self) -> usize {
352        self.buffers.lock().len()
353    }
354
355    /// Clear all pooled buffers.
356    pub fn clear(&self) {
357        self.buffers.lock().clear();
358    }
359}
360
361/// A buffer borrowed from a pool.
362///
363/// Automatically returns to the pool when dropped.
364pub struct PooledBuffer<'a> {
365    buffer: String,
366    pool: &'a BufferPool,
367}
368
369impl<'a> PooledBuffer<'a> {
370    /// Get mutable access to the buffer.
371    pub fn as_mut_str(&mut self) -> &mut String {
372        &mut self.buffer
373    }
374
375    /// Take ownership of the buffer (does not return to pool).
376    pub fn take(mut self) -> String {
377        std::mem::take(&mut self.buffer)
378    }
379}
380
381impl<'a> std::ops::Deref for PooledBuffer<'a> {
382    type Target = String;
383
384    fn deref(&self) -> &Self::Target {
385        &self.buffer
386    }
387}
388
389impl<'a> std::ops::DerefMut for PooledBuffer<'a> {
390    fn deref_mut(&mut self) -> &mut Self::Target {
391        &mut self.buffer
392    }
393}
394
395impl<'a> Drop for PooledBuffer<'a> {
396    fn drop(&mut self) {
397        if !self.buffer.is_empty() || self.buffer.capacity() > 0 {
398            let buffer = std::mem::take(&mut self.buffer);
399            self.pool.return_buffer(buffer);
400        }
401    }
402}
403
404// ============================================================================
405// Memory Usage Tracking
406// ============================================================================
407
408/// Track memory usage for debugging and optimization.
409#[derive(Debug, Clone, Default)]
410pub struct MemoryStats {
411    /// Number of allocations.
412    pub allocations: u64,
413    /// Number of deallocations.
414    pub deallocations: u64,
415    /// Current bytes allocated.
416    pub current_bytes: usize,
417    /// Peak bytes allocated.
418    pub peak_bytes: usize,
419    /// Total bytes allocated (lifetime).
420    pub total_bytes: usize,
421}
422
423impl MemoryStats {
424    /// Create new empty stats.
425    pub fn new() -> Self {
426        Self::default()
427    }
428
429    /// Record an allocation.
430    pub fn record_alloc(&mut self, bytes: usize) {
431        self.allocations += 1;
432        self.current_bytes += bytes;
433        self.total_bytes += bytes;
434        if self.current_bytes > self.peak_bytes {
435            self.peak_bytes = self.current_bytes;
436        }
437    }
438
439    /// Record a deallocation.
440    pub fn record_dealloc(&mut self, bytes: usize) {
441        self.deallocations += 1;
442        self.current_bytes = self.current_bytes.saturating_sub(bytes);
443    }
444
445    /// Get the net allocation count.
446    pub fn net_allocations(&self) -> i64 {
447        self.allocations as i64 - self.deallocations as i64
448    }
449}
450
451// ============================================================================
452// Global Pools (Lazy Initialized)
453// ============================================================================
454
455/// Global string pool for common field names.
456pub static GLOBAL_STRING_POOL: std::sync::LazyLock<StringPool> = std::sync::LazyLock::new(|| {
457    let pool = StringPool::with_capacity(128);
458    // Pre-populate with common field names
459    for name in COMMON_FIELD_NAMES {
460        pool.intern(name);
461    }
462    pool
463});
464
465/// Global buffer pool for SQL generation.
466pub static GLOBAL_BUFFER_POOL: std::sync::LazyLock<BufferPool> =
467    std::sync::LazyLock::new(|| BufferPool::with_capacity(256));
468
469/// Common field names to pre-populate the string pool.
470const COMMON_FIELD_NAMES: &[&str] = &[
471    "id",
472    "uuid",
473    "name",
474    "email",
475    "username",
476    "password",
477    "title",
478    "description",
479    "content",
480    "body",
481    "status",
482    "type",
483    "role",
484    "active",
485    "enabled",
486    "deleted",
487    "verified",
488    "published",
489    "count",
490    "score",
491    "priority",
492    "order",
493    "position",
494    "age",
495    "amount",
496    "price",
497    "quantity",
498    "user_id",
499    "post_id",
500    "comment_id",
501    "category_id",
502    "parent_id",
503    "author_id",
504    "owner_id",
505    "created_at",
506    "updated_at",
507    "deleted_at",
508    "published_at",
509    "expires_at",
510    "starts_at",
511    "ends_at",
512    "last_login_at",
513    "verified_at",
514    "slug",
515    "url",
516    "path",
517    "key",
518    "value",
519    "token",
520    "code",
521    "version",
522];
523
524/// Intern a string using the global pool.
525#[inline]
526pub fn intern(s: &str) -> Arc<str> {
527    GLOBAL_STRING_POOL.intern(s)
528}
529
530/// Get a buffer from the global pool.
531#[inline]
532pub fn get_buffer() -> PooledBuffer<'static> {
533    GLOBAL_BUFFER_POOL.get()
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539
540    #[test]
541    fn test_string_pool_interning() {
542        let pool = StringPool::new();
543
544        let s1 = pool.intern("hello");
545        let s2 = pool.intern("hello");
546
547        // Should be the same Arc
548        assert!(Arc::ptr_eq(&s1, &s2));
549        assert_eq!(pool.len(), 1);
550    }
551
552    #[test]
553    fn test_string_pool_different_strings() {
554        let pool = StringPool::new();
555
556        let s1 = pool.intern("hello");
557        let s2 = pool.intern("world");
558
559        assert!(!Arc::ptr_eq(&s1, &s2));
560        assert_eq!(pool.len(), 2);
561    }
562
563    #[test]
564    fn test_compact_filter_eq_int() {
565        let filter = CompactFilter::eq_int(Arc::from("id"), 42);
566        let mut offset = 0;
567        let sql = filter.to_sql_postgres(&mut offset);
568        assert_eq!(sql, "id = $1");
569        assert_eq!(offset, 1);
570    }
571
572    #[test]
573    fn test_compact_filter_and() {
574        let filter = CompactFilter::eq_int(Arc::from("id"), 42)
575            .and(CompactFilter::eq_bool(Arc::from("active"), true));
576        let mut offset = 0;
577        let sql = filter.to_sql_postgres(&mut offset);
578        assert_eq!(sql, "(id = $1 AND active = $2)");
579        assert_eq!(offset, 2);
580    }
581
582    #[test]
583    fn test_compact_filter_is_null() {
584        let filter = CompactFilter::is_null(Arc::from("deleted_at"));
585        let mut offset = 0;
586        let sql = filter.to_sql_postgres(&mut offset);
587        assert_eq!(sql, "deleted_at IS NULL");
588        assert_eq!(offset, 0); // No params for IS NULL
589    }
590
591    #[test]
592    fn test_buffer_pool() {
593        let pool = BufferPool::new();
594
595        {
596            let mut buffer = pool.get();
597            buffer.push_str("hello");
598            assert_eq!(&*buffer, "hello");
599        } // Buffer returned to pool here
600
601        assert_eq!(pool.available(), 1);
602
603        {
604            let buffer = pool.get();
605            assert!(buffer.is_empty()); // Buffer was cleared
606        }
607    }
608
609    #[test]
610    fn test_global_intern() {
611        let s1 = intern("id");
612        let s2 = intern("id");
613        assert!(Arc::ptr_eq(&s1, &s2));
614    }
615
616    #[test]
617    fn test_memory_stats() {
618        let mut stats = MemoryStats::new();
619
620        stats.record_alloc(100);
621        stats.record_alloc(200);
622        assert_eq!(stats.current_bytes, 300);
623        assert_eq!(stats.peak_bytes, 300);
624
625        stats.record_dealloc(100);
626        assert_eq!(stats.current_bytes, 200);
627        assert_eq!(stats.peak_bytes, 300); // Peak unchanged
628    }
629}
prax_query/memory.rs

prax_query/
memory.rs
