//! Arena-based memory allocator for AST nodes (safe version)
//!
//! This module provides an efficient arena allocator for AST nodes that:
//! - Reduces allocation overhead by pooling allocations
//! - Improves cache locality by keeping related nodes close
//! - Enables fast bulk deallocation when the arena is dropped
//! - Uses only safe Rust code

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

/// Memory pool for reusing allocations
pub struct Arena {
    /// Track total allocated items for statistics
    total_allocated: RefCell<usize>,
    /// Storage for allocated values (using Rc for shared ownership)
    storage: RefCell<Vec<Rc<dyn std::any::Any>>>,
}

impl Arena {
    /// Create a new arena allocator
    pub fn new() -> Self {
        Self {
            total_allocated: RefCell::new(0),
            storage: RefCell::new(Vec::new()),
        }
    }

    /// Allocate a value in the arena (returns Rc for shared ownership)
    pub fn alloc<T: 'static>(&self, value: T) -> Rc<T> {
        *self.total_allocated.borrow_mut() += 1;
        let rc = Rc::new(value);
        self.storage
            .borrow_mut()
            .push(rc.clone() as Rc<dyn std::any::Any>);
        rc
    }

    /// Get total items allocated
    pub fn total_allocated(&self) -> usize {
        *self.total_allocated.borrow()
    }

    /// Get number of items in storage
    pub fn num_items(&self) -> usize {
        self.storage.borrow().len()
    }

    /// Clear all allocations (for reuse)
    pub fn clear(&self) {
        self.storage.borrow_mut().clear();
        *self.total_allocated.borrow_mut() = 0;
    }
}

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

/// String interner for deduplicating strings in the AST
pub struct StringInterner {
    map: RefCell<HashMap<String, Rc<str>>>,
}

impl StringInterner {
    /// Create a new string interner
    pub fn new() -> Self {
        Self {
            map: RefCell::new(HashMap::new()),
        }
    }

    /// Intern a string, returning an Rc that can be cheaply cloned
    pub fn intern(&self, s: &str) -> Rc<str> {
        let mut map = self.map.borrow_mut();

        if let Some(interned) = map.get(s) {
            Rc::clone(interned)
        } else {
            let rc: Rc<str> = Rc::from(s);
            map.insert(s.to_string(), Rc::clone(&rc));
            rc
        }
    }

    /// Get statistics about the interner
    pub fn stats(&self) -> (usize, usize) {
        let map = self.map.borrow();
        let total_size: usize = map.values().map(|s| s.len()).sum();
        (map.len(), total_size)
    }

    /// Clear the interner
    pub fn clear(&self) {
        self.map.borrow_mut().clear();
    }
}

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

/// Fast bump allocator for sequential allocations
pub struct BumpAllocator<T> {
    /// Pre-allocated vector with capacity
    storage: RefCell<Vec<T>>,
    /// Track allocations
    count: RefCell<usize>,
}

impl<T> BumpAllocator<T> {
    /// Create a new bump allocator with initial capacity
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            storage: RefCell::new(Vec::with_capacity(capacity)),
            count: RefCell::new(0),
        }
    }

    /// Allocate a value, returning its index
    pub fn alloc(&self, value: T) -> usize {
        let mut storage = self.storage.borrow_mut();
        let index = storage.len();
        storage.push(value);
        *self.count.borrow_mut() += 1;
        index
    }

    /// Get a reference to an allocated value by index
    pub fn get(&self, index: usize) -> Option<T>
    where
        T: Clone,
    {
        self.storage.borrow().get(index).cloned()
    }

    /// Get the number of allocated items
    pub fn len(&self) -> usize {
        self.storage.borrow().len()
    }

    /// Check if allocator is empty
    pub fn is_empty(&self) -> bool {
        self.storage.borrow().is_empty()
    }

    /// Clear all allocations
    pub fn clear(&self) {
        self.storage.borrow_mut().clear();
        *self.count.borrow_mut() = 0;
    }
}

impl<T> Default for BumpAllocator<T> {
    fn default() -> Self {
        Self::with_capacity(128)
    }
}

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

    #[test]
    fn test_arena_basic() {
        let arena = Arena::new();

        let x = arena.alloc(42i32);
        assert_eq!(*x, 42);

        let y = arena.alloc("hello".to_string());
        assert_eq!(y.as_ref(), "hello");

        assert_eq!(arena.total_allocated(), 2);
    }

    #[test]
    fn test_string_interner() {
        let interner = StringInterner::new();

        let s1 = interner.intern("hello");
        let s2 = interner.intern("hello");
        let s3 = interner.intern("world");

        assert_eq!(&*s1, "hello");
        assert_eq!(&*s2, "hello");
        assert_eq!(&*s3, "world");

        // Check that identical strings share the same Rc
        assert!(Rc::ptr_eq(&s1, &s2));
        assert!(!Rc::ptr_eq(&s1, &s3));

        let (num_strings, _) = interner.stats();
        assert_eq!(num_strings, 2); // Only "hello" and "world"
    }

    #[test]
    fn test_bump_allocator() {
        let alloc = BumpAllocator::with_capacity(10);

        let idx1 = alloc.alloc(42i32);
        let idx2 = alloc.alloc(100i32);

        assert_eq!(idx1, 0);
        assert_eq!(idx2, 1);
        assert_eq!(alloc.get(idx1), Some(42));
        assert_eq!(alloc.get(idx2), Some(100));
        assert_eq!(alloc.len(), 2);
    }

    #[test]
    fn test_arena_many_allocations() {
        let arena = Arena::new();

        for i in 0..1000 {
            let _x = arena.alloc(i);
        }

        assert_eq!(arena.total_allocated(), 1000);
        assert_eq!(arena.num_items(), 1000);

        arena.clear();
        assert_eq!(arena.total_allocated(), 0);
        assert_eq!(arena.num_items(), 0);
    }
}
