iridium-db 0.4.0

use std::collections::HashMap;

pub const PAGE_SIZE_BYTES: usize = 8 * 1024;

pub type PageId = u64;

#[derive(Debug)]
pub enum BufferPoolError {
    InvalidCapacity,
    NotFound(PageId),
    AlreadyPinned(PageId),
    NotPinned(PageId),
}

pub type Result<T> = std::result::Result<T, BufferPoolError>;

#[derive(Clone)]
pub struct PageSnapshot {
    pub page_id: PageId,
    pub data: Vec<u8>,
}

struct PageFrame {
    data: Box<[u8; PAGE_SIZE_BYTES]>,
    pin_count: u32,
    last_access: u64,
    second_last_access: u64,
    access_count: u32,
}

pub struct BufferPool {
    capacity_pages: usize,
    clock: u64,
    pages: HashMap<PageId, PageFrame>,
}

impl BufferPool {
    pub fn new(capacity_pages: usize) -> Self {
        let capacity_pages = capacity_pages.max(1);
        Self {
            capacity_pages,
            clock: 0,
            pages: HashMap::new(),
        }
    }

    pub fn contains(&self, page_id: PageId) -> bool {
        self.pages.contains_key(&page_id)
    }

    pub fn read_page(&mut self, page_id: PageId) -> PageSnapshot {
        self.access(page_id);
        let frame = self
            .pages
            .get(&page_id)
            .expect("page must exist after access");
        PageSnapshot {
            page_id,
            data: frame.data.to_vec(),
        }
    }

    pub fn pin_page(&mut self, page_id: PageId) -> Result<()> {
        self.access(page_id);
        let frame = self
            .pages
            .get_mut(&page_id)
            .ok_or(BufferPoolError::NotFound(page_id))?;
        frame.pin_count = frame.pin_count.saturating_add(1);
        Ok(())
    }

    pub fn unpin_page(&mut self, page_id: PageId) -> Result<()> {
        let frame = self
            .pages
            .get_mut(&page_id)
            .ok_or(BufferPoolError::NotFound(page_id))?;
        if frame.pin_count == 0 {
            return Err(BufferPoolError::NotPinned(page_id));
        }
        frame.pin_count -= 1;
        Ok(())
    }

    pub fn prefetch(&mut self, page_ids: &[PageId]) {
        for &page_id in page_ids {
            self.access(page_id);
        }
    }

    fn access(&mut self, page_id: PageId) {
        self.clock = self.clock.saturating_add(1);
        if !self.pages.contains_key(&page_id) {
            self.evict_if_needed();
            self.pages.insert(
                page_id,
                PageFrame {
                    data: Box::new([0u8; PAGE_SIZE_BYTES]),
                    pin_count: 0,
                    last_access: 0,
                    second_last_access: 0,
                    access_count: 0,
                },
            );
        }
        if let Some(frame) = self.pages.get_mut(&page_id) {
            frame.access_count = frame.access_count.saturating_add(1);
            frame.second_last_access = frame.last_access;
            frame.last_access = self.clock;
        }
    }

    fn evict_if_needed(&mut self) {
        if self.pages.len() < self.capacity_pages {
            return;
        }

        let mut candidate: Option<PageId> = None;
        let mut candidate_score: u64 = u64::MAX;
        let mut candidate_has_second = true;

        for (page_id, frame) in &self.pages {
            if frame.pin_count > 0 {
                continue;
            }

            let has_second = frame.access_count >= 2 && frame.second_last_access > 0;
            let score = if has_second {
                frame.second_last_access
            } else {
                0
            };

            match candidate {
                None => {
                    candidate = Some(*page_id);
                    candidate_score = score;
                    candidate_has_second = has_second;
                }
                Some(_) => {
                    if candidate_has_second && !has_second {
                        candidate = Some(*page_id);
                        candidate_score = score;
                        candidate_has_second = has_second;
                        continue;
                    }
                    if !candidate_has_second && has_second {
                        continue;
                    }
                    if score < candidate_score {
                        candidate = Some(*page_id);
                        candidate_score = score;
                        candidate_has_second = has_second;
                    }
                }
            }
        }

        if let Some(page_id) = candidate {
            self.pages.remove(&page_id);
        }
    }
}

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

    #[test]
    fn lru2_evicts_single_access_first() {
        let mut pool = BufferPool::new(2);

        pool.read_page(1);
        pool.read_page(1);
        pool.read_page(2);

        pool.read_page(3);

        assert!(pool.contains(1));
        assert!(!pool.contains(2));
        assert!(pool.contains(3));
    }

    #[test]
    fn pin_prevents_eviction() {
        let mut pool = BufferPool::new(2);

        pool.read_page(10);
        pool.read_page(20);
        pool.pin_page(10).unwrap();

        pool.read_page(30);

        assert!(pool.contains(10));
        assert!(!pool.contains(20));
        assert!(pool.contains(30));

        pool.unpin_page(10).unwrap();
    }

    #[test]
    fn page_size_is_fixed() {
        let mut pool = BufferPool::new(1);
        let snapshot = pool.read_page(42);
        assert_eq!(snapshot.data.len(), PAGE_SIZE_BYTES);
    }
}