armature-core 0.9.0

//! Per-Request Arena Allocator
//!
//! This module provides a thread-local arena allocator for batch allocation
//! and deallocation of request-scoped data. All allocations made during
//! request processing are freed at once when the request completes, reducing
//! allocator pressure and improving cache locality.
//!
//! ## Performance Benefits
//!
//! - **Batch deallocation**: All allocations freed in O(1) time
//! - **Cache locality**: Consecutive allocations are adjacent in memory
//! - **Reduced fragmentation**: No per-allocation bookkeeping overhead
//! - **Thread-local**: No synchronization overhead for allocation
//!
//! ## Usage
//!
//! ```rust,ignore
//! use armature_core::arena::{with_arena, ArenaString, ArenaVec};
//!
//! // Process a request with arena allocation
//! with_arena(|arena| {
//!     // Allocate strings in the arena
//!     let method = ArenaString::from_str(arena, "GET");
//!     let path = ArenaString::from_str(arena, "/api/users/123");
//!
//!     // Allocate vectors in the arena
//!     let mut headers = ArenaVec::new_in(arena);
//!     headers.push(("Content-Type", "application/json"));
//!
//!     // All allocations are freed when this closure returns
//! });
//! ```
//!
//! ## Arena-Backed Request
//!
//! For maximum performance, use `ArenaRequest` which allocates all
//! request data from the arena:
//!
//! ```rust,ignore
//! use armature_core::arena::{with_arena, ArenaRequest};
//!
//! with_arena(|arena| {
//!     let request = ArenaRequest::new(arena, "GET", "/api/users");
//!     // Process request...
//! }); // All request data freed at once
//! ```

use bumpalo::Bump;
use bytes::Bytes;
use std::cell::RefCell;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::ops::Deref;

// ============================================================================
// Thread-Local Arena Pool
// ============================================================================

/// Default arena size (64 KB - fits most requests)
const DEFAULT_ARENA_SIZE: usize = 64 * 1024;

/// Maximum arena size before reset (1 MB)
const MAX_ARENA_SIZE: usize = 1024 * 1024;

thread_local! {
    /// Thread-local arena for per-request allocations.
    /// Reused across requests to avoid repeated arena creation.
    static ARENA: RefCell<Bump> = RefCell::new(Bump::with_capacity(DEFAULT_ARENA_SIZE));
}

/// Execute a function with access to the thread-local arena.
///
/// The arena is reset after the function completes, freeing all allocations.
/// This is the primary entry point for arena-backed request processing.
///
/// # Example
///
/// ```rust,ignore
/// use armature_core::arena::with_arena;
///
/// let result = with_arena(|arena| {
///     // Allocate data in the arena
///     let s = arena.alloc_str("hello");
///     s.len()
/// });
/// // Arena is reset here - all allocations freed
/// ```
#[inline]
pub fn with_arena<F, R>(f: F) -> R
where
    F: FnOnce(&Bump) -> R,
{
    ARENA.with(|arena| {
        let arena = arena.borrow();

        // Note: We don't reset here - let the caller do it explicitly
        // for better control over when deallocations happen
        f(&arena)
    })
}

/// Execute a function with mutable access to the thread-local arena.
///
/// The arena is reset after the function completes.
#[inline]
pub fn with_arena_mut<F, R>(f: F) -> R
where
    F: FnOnce(&Bump) -> R,
{
    ARENA.with(|arena| {
        let arena = arena.borrow();

        f(&arena)
    })
}

/// Reset the thread-local arena, freeing all allocations.
///
/// Call this after processing a request to reclaim memory.
/// If the arena has grown too large, it will be shrunk.
#[inline]
pub fn reset_arena() {
    ARENA.with(|arena| {
        let mut arena = arena.borrow_mut();

        // If arena grew too large, recreate it with default size
        if arena.allocated_bytes() > MAX_ARENA_SIZE {
            *arena = Bump::with_capacity(DEFAULT_ARENA_SIZE);
        } else {
            arena.reset();
        }
    });
}

/// Get the current arena allocation size (for diagnostics).
#[inline]
pub fn arena_allocated_bytes() -> usize {
    ARENA.with(|arena| arena.borrow().allocated_bytes())
}

// ============================================================================
// Arena-Backed String
// ============================================================================

/// An arena-allocated string slice.
///
/// This is a lightweight reference to a string stored in the arena.
/// It does not own the memory - the arena does.
///
/// # Performance
///
/// - No allocation overhead per string
/// - No deallocation overhead per string
/// - Deref to `&str` for easy use
#[derive(Clone, Copy)]
pub struct ArenaStr<'a> {
    inner: &'a str,
}

impl<'a> ArenaStr<'a> {
    /// Create a new arena string from a string slice.
    ///
    /// The string is copied into the arena.
    #[inline]
    pub fn from_str(arena: &'a Bump, s: &str) -> Self {
        Self {
            inner: arena.alloc_str(s),
        }
    }

    /// Create an empty arena string.
    #[inline]
    pub const fn empty() -> Self {
        Self { inner: "" }
    }

    /// Get the string as a `&str`.
    #[inline]
    pub fn as_str(&self) -> &'a str {
        self.inner
    }

    /// Get the length of the string.
    #[inline]
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Check if the string is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Convert to owned String (allocates on heap).
    ///
    /// Note: Prefer using `Display::to_string()` if `Display` is in scope.
    #[inline]
    pub fn into_string(&self) -> String {
        self.inner.to_string()
    }
}

impl<'a> Deref for ArenaStr<'a> {
    type Target = str;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.inner
    }
}

impl<'a> AsRef<str> for ArenaStr<'a> {
    #[inline]
    fn as_ref(&self) -> &str {
        self.inner
    }
}

impl<'a> fmt::Debug for ArenaStr<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self.inner)
    }
}

impl<'a> fmt::Display for ArenaStr<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.inner)
    }
}

impl<'a> PartialEq for ArenaStr<'a> {
    fn eq(&self, other: &Self) -> bool {
        self.inner == other.inner
    }
}

impl<'a> Eq for ArenaStr<'a> {}

impl<'a> PartialEq<str> for ArenaStr<'a> {
    fn eq(&self, other: &str) -> bool {
        self.inner == other
    }
}

impl<'a> PartialEq<&str> for ArenaStr<'a> {
    fn eq(&self, other: &&str) -> bool {
        self.inner == *other
    }
}

impl<'a> PartialEq<String> for ArenaStr<'a> {
    fn eq(&self, other: &String) -> bool {
        self.inner == other.as_str()
    }
}

impl<'a> Hash for ArenaStr<'a> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.inner.hash(state);
    }
}

// ============================================================================
// Arena-Backed Vector
// ============================================================================

/// An arena-allocated vector.
///
/// Uses bumpalo's `Vec` under the hood for efficient arena allocation.
pub type ArenaVec<'a, T> = bumpalo::collections::Vec<'a, T>;

// ============================================================================
// Arena-Backed HashMap
// ============================================================================

/// A simple arena-backed hash map for headers and parameters.
///
/// This uses a vector of key-value pairs for small maps (typical for headers).
/// For maps with many entries, a standard HashMap may be more efficient.
pub struct ArenaMap<'a, K, V> {
    entries: ArenaVec<'a, (K, V)>,
}

impl<'a, K: PartialEq, V> ArenaMap<'a, K, V> {
    /// Create a new empty arena map.
    #[inline]
    pub fn new_in(arena: &'a Bump) -> Self {
        Self {
            entries: ArenaVec::new_in(arena),
        }
    }

    /// Create a new arena map with pre-allocated capacity.
    #[inline]
    pub fn with_capacity_in(arena: &'a Bump, capacity: usize) -> Self {
        Self {
            entries: ArenaVec::with_capacity_in(capacity, arena),
        }
    }

    /// Insert a key-value pair.
    ///
    /// If the key exists, the value is updated and the old value returned.
    #[inline]
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        for entry in self.entries.iter_mut() {
            if entry.0 == key {
                let old = std::mem::replace(&mut entry.1, value);
                return Some(old);
            }
        }
        self.entries.push((key, value));
        None
    }

    /// Get a reference to a value by key.
    #[inline]
    pub fn get(&self, key: &K) -> Option<&V> {
        self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v)
    }

    /// Get a mutable reference to a value by key.
    #[inline]
    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        self.entries
            .iter_mut()
            .find(|(k, _)| k == key)
            .map(|(_, v)| v)
    }

    /// Check if a key exists.
    #[inline]
    pub fn contains_key(&self, key: &K) -> bool {
        self.entries.iter().any(|(k, _)| k == key)
    }

    /// Get the number of entries.
    #[inline]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Check if the map is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Iterate over entries.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &(K, V)> {
        self.entries.iter()
    }

    /// Iterate over keys.
    #[inline]
    pub fn keys(&self) -> impl Iterator<Item = &K> {
        self.entries.iter().map(|(k, _)| k)
    }

    /// Iterate over values.
    #[inline]
    pub fn values(&self) -> impl Iterator<Item = &V> {
        self.entries.iter().map(|(_, v)| v)
    }
}

impl<'a, K: PartialEq + fmt::Debug, V: fmt::Debug> fmt::Debug for ArenaMap<'a, K, V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map()
            .entries(self.entries.iter().map(|(k, v)| (k, v)))
            .finish()
    }
}

// ============================================================================
// Arena-Backed HTTP Request
// ============================================================================

/// An arena-allocated HTTP request.
///
/// All string data (method, path, headers, etc.) is allocated from
/// the arena, enabling batch deallocation when the request completes.
///
/// # Performance
///
/// Compared to `HttpRequest` with heap allocations:
/// - **Allocation**: ~5x faster (single bump vs multiple malloc)
/// - **Deallocation**: ~10x faster (reset vs multiple free)
/// - **Cache locality**: Better due to contiguous allocation
///
/// # Example
///
/// ```rust,ignore
/// use armature_core::arena::{with_arena, ArenaRequest};
///
/// with_arena(|arena| {
///     let mut request = ArenaRequest::new(arena, "GET", "/api/users");
///     request.add_header(arena, "Content-Type", "application/json");
///     request.add_query_param(arena, "page", "1");
///
///     // Process request...
/// }); // All data freed at once
/// ```
pub struct ArenaRequest<'a> {
    /// HTTP method (GET, POST, etc.)
    pub method: ArenaStr<'a>,
    /// Request path (e.g., "/api/users/123")
    pub path: ArenaStr<'a>,
    /// Request headers
    pub headers: ArenaMap<'a, ArenaStr<'a>, ArenaStr<'a>>,
    /// Path parameters extracted from route (e.g., {"id": "123"})
    pub path_params: ArenaMap<'a, ArenaStr<'a>, ArenaStr<'a>>,
    /// Query parameters (e.g., {"page": "1", "limit": "10"})
    pub query_params: ArenaMap<'a, ArenaStr<'a>, ArenaStr<'a>>,
    /// Request body (not arena-allocated - typically from hyper)
    pub body: &'a [u8],
}

impl<'a> ArenaRequest<'a> {
    /// Create a new arena-backed request.
    #[inline]
    pub fn new(arena: &'a Bump, method: &str, path: &str) -> Self {
        Self {
            method: ArenaStr::from_str(arena, method),
            path: ArenaStr::from_str(arena, path),
            headers: ArenaMap::with_capacity_in(arena, 16), // Typical header count
            path_params: ArenaMap::with_capacity_in(arena, 4),
            query_params: ArenaMap::with_capacity_in(arena, 8),
            body: &[],
        }
    }

    /// Create a new arena-backed request with body reference.
    #[inline]
    pub fn with_body(arena: &'a Bump, method: &str, path: &str, body: &'a [u8]) -> Self {
        Self {
            method: ArenaStr::from_str(arena, method),
            path: ArenaStr::from_str(arena, path),
            headers: ArenaMap::with_capacity_in(arena, 16),
            path_params: ArenaMap::with_capacity_in(arena, 4),
            query_params: ArenaMap::with_capacity_in(arena, 8),
            body,
        }
    }

    /// Add a header to the request.
    #[inline]
    pub fn add_header(&mut self, arena: &'a Bump, name: &str, value: &str) {
        let name = ArenaStr::from_str(arena, name);
        let value = ArenaStr::from_str(arena, value);
        self.headers.insert(name, value);
    }

    /// Add a path parameter.
    #[inline]
    pub fn add_path_param(&mut self, arena: &'a Bump, name: &str, value: &str) {
        let name = ArenaStr::from_str(arena, name);
        let value = ArenaStr::from_str(arena, value);
        self.path_params.insert(name, value);
    }

    /// Add a query parameter.
    #[inline]
    pub fn add_query_param(&mut self, arena: &'a Bump, name: &str, value: &str) {
        let name = ArenaStr::from_str(arena, name);
        let value = ArenaStr::from_str(arena, value);
        self.query_params.insert(name, value);
    }

    /// Get a header value by name.
    #[inline]
    pub fn header(&self, name: &str) -> Option<&str> {
        // Linear search is fine for typical header counts (< 20)
        self.headers
            .iter()
            .find(|(k, _)| k.as_str().eq_ignore_ascii_case(name))
            .map(|(_, v)| v.as_str())
    }

    /// Get a path parameter by name.
    #[inline]
    pub fn param(&self, name: &str) -> Option<&str> {
        self.path_params
            .iter()
            .find(|(k, _)| k.as_str() == name)
            .map(|(_, v)| v.as_str())
    }

    /// Get a query parameter by name.
    #[inline]
    pub fn query(&self, name: &str) -> Option<&str> {
        self.query_params
            .iter()
            .find(|(k, _)| k.as_str() == name)
            .map(|(_, v)| v.as_str())
    }

    /// Parse body as JSON.
    #[inline]
    pub fn json<T: serde::de::DeserializeOwned>(&self) -> Result<T, crate::Error> {
        crate::json::from_slice(self.body).map_err(|e| crate::Error::Deserialization(e.to_string()))
    }

    /// Convert to standard HttpRequest (allocates on heap).
    ///
    /// Use this when you need to pass the request to code that
    /// expects `HttpRequest`.
    pub fn to_http_request(&self) -> crate::HttpRequest {
        let mut req = crate::HttpRequest::new(self.method.to_string(), self.path.to_string());

        for (k, v) in self.headers.iter() {
            // `as_str` on both halves: `ArenaStr` derefs to `str`, and the
            // header map takes anything that can become a value without an
            // intermediate `String`.
            req.headers.insert(k.as_str(), v.as_str());
        }

        // Names are interned: the arena request owns its names, and a request
        // param name has to be `&'static str`. The interner is hard-capped
        // (see [`crate::param_intern`]), so a name that came in off the wire
        // rather than from a route pattern cannot grow the process.
        req.path_params = self
            .path_params
            .iter()
            .map(|(k, v)| {
                (
                    crate::param_intern::intern(k.as_str()),
                    Bytes::copy_from_slice(v.as_str().as_bytes()),
                )
            })
            .collect();

        // The query lives in the target now, so it has to be re-encoded onto
        // the path rather than handed over as a map.
        if !self.query_params.is_empty() {
            let pairs: Vec<(&str, &str)> = self
                .query_params
                .iter()
                .map(|(k, v)| (k.as_str(), v.as_str()))
                .collect();
            if let Ok(query) = serde_urlencoded::to_string(&pairs) {
                req.path = crate::ByteStr::from(format!("{}?{}", self.path.as_str(), query));
            }
        }

        req.body = Bytes::copy_from_slice(self.body);
        req
    }
}

impl<'a> fmt::Debug for ArenaRequest<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ArenaRequest")
            .field("method", &self.method)
            .field("path", &self.path)
            .field("headers", &self.headers)
            .field("path_params", &self.path_params)
            .field("query_params", &self.query_params)
            .field("body_len", &self.body.len())
            .finish()
    }
}

// ============================================================================
// Conversion from Hyper Request
// ============================================================================

/// Convert a hyper request to an arena request.
///
/// This is the most efficient way to create a request from incoming HTTP data.
#[inline]
pub fn arena_request_from_hyper<'a>(
    arena: &'a Bump,
    method: &str,
    path: &str,
    headers: impl Iterator<Item = (&'a str, &'a str)>,
    body: &'a [u8],
) -> ArenaRequest<'a> {
    let mut request = ArenaRequest::with_body(arena, method, path, body);

    for (name, value) in headers {
        request.add_header(arena, name, value);
    }

    request
}

// ============================================================================
// Request Scope Guard
// ============================================================================

/// A scope guard that resets the arena when dropped.
///
/// Use this to ensure the arena is reset even if a panic occurs.
///
/// # Example
///
/// ```rust,ignore
/// use armature_core::arena::RequestScope;
///
/// {
///     let _scope = RequestScope::new();
///     // Allocate in arena...
/// } // Arena automatically reset here
/// ```
pub struct RequestScope {
    _private: (),
}

impl RequestScope {
    /// Create a new request scope.
    ///
    /// When this is dropped, the thread-local arena is reset.
    #[inline]
    pub fn new() -> Self {
        Self { _private: () }
    }
}

impl Default for RequestScope {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for RequestScope {
    #[inline]
    fn drop(&mut self) {
        reset_arena();
    }
}

// ============================================================================
// Statistics and Diagnostics
// ============================================================================

/// Arena allocation statistics.
#[derive(Debug, Clone, Copy)]
pub struct ArenaStats {
    /// Total bytes allocated in the arena.
    pub allocated_bytes: usize,
    /// Number of chunks allocated by the arena.
    pub chunk_count: usize,
}

/// Get current arena statistics.
#[inline]
pub fn arena_stats() -> ArenaStats {
    ARENA.with(|arena| {
        let mut arena = arena.borrow_mut();
        let chunk_count = arena.iter_allocated_chunks().count();
        ArenaStats {
            allocated_bytes: arena.allocated_bytes(),
            chunk_count,
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_arena_str() {
        with_arena(|arena| {
            let s = ArenaStr::from_str(arena, "hello world");
            assert_eq!(s.as_str(), "hello world");
            assert_eq!(s.len(), 11);
            assert!(!s.is_empty());
            assert_eq!(&*s, "hello world"); // Deref
        });
        reset_arena();
    }

    #[test]
    fn test_arena_str_empty() {
        let s = ArenaStr::empty();
        assert!(s.is_empty());
        assert_eq!(s.len(), 0);
    }

    #[test]
    fn test_arena_str_equality() {
        with_arena(|arena| {
            let s1 = ArenaStr::from_str(arena, "test");
            let s2 = ArenaStr::from_str(arena, "test");
            let s3 = ArenaStr::from_str(arena, "other");

            assert_eq!(s1, s2);
            assert_ne!(s1, s3);
            assert!(s1 == "test"); // Compare with &str
            assert!(s1 == "test"); // Compare with String
            assert_eq!(s1.as_str(), "test"); // Direct comparison
        });
        reset_arena();
    }

    #[test]
    fn test_arena_map() {
        with_arena(|arena| {
            let mut map = ArenaMap::<ArenaStr, ArenaStr>::new_in(arena);

            let key1 = ArenaStr::from_str(arena, "key1");
            let val1 = ArenaStr::from_str(arena, "value1");
            map.insert(key1, val1);

            let key2 = ArenaStr::from_str(arena, "key2");
            let val2 = ArenaStr::from_str(arena, "value2");
            map.insert(key2, val2);

            assert_eq!(map.len(), 2);
            assert!(map.contains_key(&ArenaStr::from_str(arena, "key1")));

            let lookup = ArenaStr::from_str(arena, "key1");
            assert_eq!(map.get(&lookup).map(|v| v.as_str()), Some("value1"));
        });
        reset_arena();
    }

    #[test]
    fn test_arena_request() {
        with_arena(|arena| {
            let mut request = ArenaRequest::new(arena, "POST", "/api/users");
            request.add_header(arena, "Content-Type", "application/json");
            request.add_header(arena, "Authorization", "Bearer token123");
            request.add_query_param(arena, "page", "1");
            request.add_path_param(arena, "id", "42");

            assert_eq!(request.method.as_str(), "POST");
            assert_eq!(request.path.as_str(), "/api/users");
            assert_eq!(request.header("Content-Type"), Some("application/json"));
            assert_eq!(request.header("content-type"), Some("application/json")); // Case-insensitive
            assert_eq!(request.query("page"), Some("1"));
            assert_eq!(request.param("id"), Some("42"));
        });
        reset_arena();
    }

    #[test]
    fn test_arena_request_to_http_request() {
        with_arena(|arena| {
            let mut request = ArenaRequest::new(arena, "GET", "/test");
            request.add_header(arena, "Accept", "application/json");

            let http_request = request.to_http_request();
            assert_eq!(http_request.method, "GET");
            assert_eq!(http_request.path, "/test");
            assert_eq!(http_request.headers.get("Accept"), Some("application/json"));
        });
        reset_arena();
    }

    #[test]
    fn test_request_scope() {
        // RequestScope ensures arena is reset on drop
        {
            let _scope = RequestScope::new();
            with_arena(|arena| {
                // Allocate some data
                let s1 = ArenaStr::from_str(arena, "test string that uses memory");
                let s2 = ArenaStr::from_str(arena, "another string");
                assert!(!s1.is_empty());
                assert!(!s2.is_empty());
            });
        }
        // After scope drop, arena is reset (can allocate again)
        with_arena(|arena| {
            let s = ArenaStr::from_str(arena, "new allocation after reset");
            assert!(!s.is_empty());
        });
        reset_arena();
    }

    #[test]
    fn test_arena_stats() {
        reset_arena();
        let stats = arena_stats();
        // Arena has at least one chunk even when empty
        assert!(stats.chunk_count >= 1);

        // Allocate significant data to ensure we see the difference
        with_arena(|arena| {
            // Allocate a large vector to ensure we see allocation changes
            let mut vec: ArenaVec<u8> = ArenaVec::with_capacity_in(1024, arena);
            for i in 0..255u8 {
                vec.push(i);
            }
            assert_eq!(vec.len(), 255);
        });

        // Stats should show allocations
        let stats_after = arena_stats();
        // At minimum, we should have some allocation
        assert!(stats_after.allocated_bytes > 0);
        reset_arena();
    }

    #[test]
    fn test_arena_vec() {
        with_arena(|arena| {
            let mut vec: ArenaVec<i32> = ArenaVec::new_in(arena);
            vec.push(1);
            vec.push(2);
            vec.push(3);

            assert_eq!(vec.len(), 3);
            assert_eq!(vec[0], 1);
            assert_eq!(vec[1], 2);
            assert_eq!(vec[2], 3);
        });
        reset_arena();
    }
}