featherdb_storage/

freelist.rs

//! Free page list management

use crate::BufferPool;
use bitvec::prelude::*;
use featherdb_core::{constants, Error, PageId, Result};
use std::sync::atomic::{AtomicU64, Ordering};

/// Bitmap-based free page tracking
pub struct FreeList {
    /// Root page of the bitmap (stored on disk)
    root_page: PageId,
    /// Cached bitmap for fast allocation
    bitmap: BitVec<u8, Lsb0>,
    /// Hint for next free page (speeds up allocation)
    next_free_hint: AtomicU64,
    /// Whether the freelist has been loaded from disk
    loaded: bool,
}

impl FreeList {
    /// Create a new free list
    pub fn new(root_page: PageId) -> Self {
        FreeList {
            root_page,
            bitmap: BitVec::new(),
            next_free_hint: AtomicU64::new(constants::FIRST_DATA_PAGE),
            loaded: false,
        }
    }

    /// Load the free list from disk
    fn load(&mut self, buffer_pool: &BufferPool) -> Result<()> {
        if self.loaded {
            return Ok(());
        }

        // For now, we'll use a simple approach: the bitmap is stored directly
        // in the freelist page. For larger databases, this would need to span
        // multiple pages.

        let guard = buffer_pool.get_page(self.root_page)?;
        let page = guard.read();

        // The page data after the header contains the bitmap
        let data = page.as_bytes();
        let bitmap_start = constants::PAGE_HEADER_SIZE;
        let bitmap_data = &data[bitmap_start..];

        self.bitmap = BitVec::from_slice(bitmap_data);
        self.loaded = true;

        Ok(())
    }

    /// Save the free list to disk
    fn save(&self, buffer_pool: &BufferPool) -> Result<()> {
        let guard = buffer_pool.get_page_for_write(self.root_page)?;
        let mut page = guard.write();

        let data = page.as_bytes_mut();
        let bitmap_start = constants::PAGE_HEADER_SIZE;

        // Write bitmap to page
        let bitmap_bytes = self.bitmap.as_raw_slice();
        let copy_len = bitmap_bytes.len().min(data.len() - bitmap_start);
        data[bitmap_start..bitmap_start + copy_len].copy_from_slice(&bitmap_bytes[..copy_len]);

        Ok(())
    }

    /// Allocate a new page
    pub fn allocate(&mut self, buffer_pool: &BufferPool) -> Result<PageId> {
        self.load(buffer_pool)?;

        let hint = self.next_free_hint.load(Ordering::Relaxed) as usize;

        // Search from hint
        for i in hint..self.bitmap.len() {
            if !self.bitmap[i] {
                self.bitmap.set(i, true);
                self.next_free_hint.store((i + 1) as u64, Ordering::Relaxed);
                self.save(buffer_pool)?;
                return Ok(PageId(i as u64));
            }
        }

        // Wrap around
        for i in constants::FIRST_DATA_PAGE as usize..hint {
            if !self.bitmap[i] {
                self.bitmap.set(i, true);
                self.next_free_hint.store((i + 1) as u64, Ordering::Relaxed);
                self.save(buffer_pool)?;
                return Ok(PageId(i as u64));
            }
        }

        // Need to extend the bitmap (allocate new page from file)
        let new_page = PageId(self.bitmap.len() as u64);

        // Extend bitmap
        while self.bitmap.len() <= new_page.0 as usize {
            self.bitmap.push(false);
        }
        self.bitmap.set(new_page.0 as usize, true);
        self.next_free_hint.store(new_page.0 + 1, Ordering::Relaxed);

        self.save(buffer_pool)?;
        Ok(new_page)
    }

    /// Free a page
    pub fn free(&mut self, page_id: PageId, buffer_pool: &BufferPool) -> Result<()> {
        self.load(buffer_pool)?;

        let idx = page_id.0 as usize;

        // Extend bitmap if necessary
        while self.bitmap.len() <= idx {
            self.bitmap.push(false);
        }

        if !self.bitmap[idx] {
            return Err(Error::DoubleFree(page_id));
        }

        self.bitmap.set(idx, false);

        // Update hint if this is earlier
        let hint = self.next_free_hint.load(Ordering::Relaxed) as usize;
        if idx < hint {
            self.next_free_hint.store(idx as u64, Ordering::Relaxed);
        }

        self.save(buffer_pool)?;
        Ok(())
    }

    /// Check if a page is allocated
    pub fn is_allocated(&mut self, page_id: PageId, buffer_pool: &BufferPool) -> Result<bool> {
        self.load(buffer_pool)?;

        let idx = page_id.0 as usize;
        if idx >= self.bitmap.len() {
            return Ok(false);
        }

        Ok(self.bitmap[idx])
    }

    /// Get the number of allocated pages
    pub fn allocated_count(&mut self, buffer_pool: &BufferPool) -> Result<usize> {
        self.load(buffer_pool)?;
        Ok(self.bitmap.count_ones())
    }

    /// Get the number of free pages
    pub fn free_count(&mut self, buffer_pool: &BufferPool) -> Result<usize> {
        self.load(buffer_pool)?;
        Ok(self.bitmap.count_zeros())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FileManager;
    use featherdb_core::Config;
    use std::sync::Arc;
    use tempfile::TempDir;

    fn create_test_freelist() -> (TempDir, Arc<BufferPool>, FreeList) {
        let tmp = TempDir::new().unwrap();
        let path = tmp.path().join("test.db");
        let config = Config::new(&path);

        let fm = Arc::new(FileManager::open(&config).unwrap());
        fm.init_if_needed(&config).unwrap();

        let bp = Arc::new(BufferPool::new(100, fm));
        let fl = FreeList::new(PageId::from(constants::FREELIST_ROOT_PAGE));

        (tmp, bp, fl)
    }

    #[test]
    fn test_freelist_allocate() {
        let (_tmp, bp, mut fl) = create_test_freelist();

        let page1 = fl.allocate(&bp).unwrap();
        let page2 = fl.allocate(&bp).unwrap();

        assert_ne!(page1, page2);
        assert!(fl.is_allocated(page1, &bp).unwrap());
        assert!(fl.is_allocated(page2, &bp).unwrap());
    }

    #[test]
    fn test_freelist_free() {
        let (_tmp, bp, mut fl) = create_test_freelist();

        let page = fl.allocate(&bp).unwrap();
        assert!(fl.is_allocated(page, &bp).unwrap());

        fl.free(page, &bp).unwrap();
        assert!(!fl.is_allocated(page, &bp).unwrap());
    }

    #[test]
    fn test_freelist_reuse() {
        let (_tmp, bp, mut fl) = create_test_freelist();

        let page1 = fl.allocate(&bp).unwrap();
        fl.free(page1, &bp).unwrap();

        let page2 = fl.allocate(&bp).unwrap();
        // Should reuse the freed page
        assert_eq!(page1, page2);
    }
}