Skip to main content

proka_kernel/memory/
frame_allocator.rs

1//! Bitmap-based physical frame allocator with deallocation support
2//! Copyright (C) RainSTR Studio 2025, All Rights Reserved.
3//!
4//! This module provides a frame allocator using the bitmap-allocator crate,
5//! supporting both allocation and deallocation of physical frames.
6
7extern crate alloc;
8
9use crate::config::PAGE_SIZE;
10use bitmap_allocator::BitAlloc;
11use limine::memory_map::EntryType;
12use limine::response::MemoryMapResponse;
13use spin::Mutex;
14use x86_64::structures::paging::{FrameAllocator, PhysFrame, Size4KiB};
15
16/// Bitmap frame allocator type - supports up to 16M frames (64 GiB)
17type BitAlloc16M = bitmap_allocator::BitAlloc16M;
18
19/// Global allocator instance
20///
21/// We use a `static` variable here because `BitmapFrameAllocator` contains
22/// `BitAlloc16M` which is a very large struct (~2.2 MB). Creating it on
23/// the stack would cause a stack overflow and kernel panic (triple fault).
24static FRAME_ALLOCATOR_INNER: Mutex<BitmapFrameAllocator> = Mutex::new(BitmapFrameAllocator {
25    alloc: BitAlloc16M::DEFAULT,
26    total_frames: 0,
27    used_frames: 0,
28});
29
30/// Frame statistics
31#[derive(Debug, Clone, Copy)]
32pub struct FrameStats {
33    /// Total number of frames in the system
34    pub total_frames: usize,
35    /// Number of free frames
36    pub free_frames: usize,
37    /// Number of used frames
38    pub used_frames: usize,
39    /// Total memory in bytes
40    pub total_memory: usize,
41    /// Free memory in bytes
42    pub free_memory: usize,
43    /// Used memory in bytes
44    pub used_memory: usize,
45}
46
47/// Bitmap-based frame allocator
48pub struct BitmapFrameAllocator {
49    /// The bitmap allocator
50    alloc: BitAlloc16M,
51    /// Total number of frames
52    total_frames: usize,
53    /// Number of used frames
54    used_frames: usize,
55}
56
57impl BitmapFrameAllocator {
58    /// Initialize the allocator from the memory map
59    ///
60    /// # Safety
61    /// This function is unsafe because the caller must guarantee that:
62    /// - The passed memory map is valid
63    /// - All frames marked as `USABLE` in it are really unused
64    pub unsafe fn init(&mut self, memory_map: &'static MemoryMapResponse) {
65        // Mark all usable frames as available in the bitmap
66        for region in memory_map.entries().iter() {
67            if region.entry_type == EntryType::USABLE {
68                let start = region.base as usize;
69                let end = (region.base + region.length) as usize;
70                let start_frame = start / PAGE_SIZE;
71                let end_frame = end.div_ceil(PAGE_SIZE);
72
73                self.alloc.insert(start_frame..end_frame);
74                self.total_frames += end_frame - start_frame;
75            }
76        }
77    }
78
79    /// Allocate a contiguous block of frames
80    pub fn allocate_contiguous(&mut self, count: usize) -> Option<PhysFrame> {
81        if count == 1 {
82            self.allocate_frame()
83        } else {
84            let frame_num = self.alloc.alloc_contiguous(None, count, 0)?;
85            self.used_frames += count;
86            Some(PhysFrame::containing_address(x86_64::PhysAddr::new(
87                (frame_num * PAGE_SIZE) as u64,
88            )))
89        }
90    }
91
92    /// Deallocate a frame
93    pub fn deallocate_frame(&mut self, frame: PhysFrame) {
94        let frame_num = frame.start_address().as_u64() as usize / PAGE_SIZE;
95        if self.alloc.dealloc(frame_num) {
96            self.used_frames -= 1;
97        }
98    }
99
100    /// Deallocate a contiguous block of frames
101    pub fn deallocate_contiguous(&mut self, frame: PhysFrame, count: usize) {
102        let frame_num = frame.start_address().as_u64() as usize / PAGE_SIZE;
103        if self.alloc.dealloc_contiguous(frame_num, count) {
104            self.used_frames -= count;
105        }
106    }
107
108    /// Get memory statistics
109    pub fn stats(&self) -> FrameStats {
110        FrameStats {
111            total_frames: self.total_frames,
112            free_frames: self.total_frames - self.used_frames,
113            used_frames: self.used_frames,
114            total_memory: self.total_frames * PAGE_SIZE,
115            free_memory: (self.total_frames - self.used_frames) * PAGE_SIZE,
116            used_memory: self.used_frames * PAGE_SIZE,
117        }
118    }
119
120    /// Check if a frame is allocated
121    pub fn is_allocated(&self, frame: PhysFrame) -> bool {
122        let frame_num = frame.start_address().as_u64() as usize / PAGE_SIZE;
123        !self.alloc.test(frame_num)
124    }
125
126    /// Get the number of free frames
127    pub fn free_frames(&self) -> usize {
128        self.total_frames - self.used_frames
129    }
130}
131
132unsafe impl FrameAllocator<Size4KiB> for BitmapFrameAllocator {
133    fn allocate_frame(&mut self) -> Option<PhysFrame> {
134        let frame_num = self.alloc.alloc()?;
135        self.used_frames += 1;
136        Some(PhysFrame::containing_address(x86_64::PhysAddr::new(
137            (frame_num * PAGE_SIZE) as u64,
138        )))
139    }
140}
141
142/// Global frame allocator with spinlock protection
143/// Wrapper around a static mutex to avoid stack overflow
144pub struct LockedFrameAllocator(&'static Mutex<BitmapFrameAllocator>);
145
146impl LockedFrameAllocator {
147    /// Initialize the global allocator from the memory map
148    ///
149    /// # Safety
150    /// This function is unsafe because the caller must guarantee that:
151    /// - The passed memory map is valid
152    /// - All frames marked as `USABLE` in it are really unused
153    /// - This is called only once during initialization
154    pub unsafe fn init(memory_map: &'static MemoryMapResponse) -> Self {
155        let mut allocator = FRAME_ALLOCATOR_INNER.lock();
156        if allocator.total_frames == 0 {
157            allocator.init(memory_map);
158        }
159        LockedFrameAllocator(&FRAME_ALLOCATOR_INNER)
160    }
161
162    /// Deallocate a frame
163    pub fn deallocate_frame(&self, frame: PhysFrame) {
164        self.0.lock().deallocate_frame(frame);
165    }
166
167    /// Deallocate a contiguous block of frames
168    pub fn deallocate_contiguous(&self, frame: PhysFrame, count: usize) {
169        self.0.lock().deallocate_contiguous(frame, count);
170    }
171
172    /// Get memory statistics
173    pub fn stats(&self) -> FrameStats {
174        self.0.lock().stats()
175    }
176
177    /// Check if a frame is allocated
178    pub fn is_allocated(&self, frame: PhysFrame) -> bool {
179        self.0.lock().is_allocated(frame)
180    }
181
182    /// Get the number of free frames
183    pub fn free_frames(&self) -> usize {
184        self.0.lock().free_frames()
185    }
186}
187
188unsafe impl FrameAllocator<Size4KiB> for LockedFrameAllocator {
189    fn allocate_frame(&mut self) -> Option<PhysFrame> {
190        self.0.lock().allocate_frame()
191    }
192}
193
194/// Format byte count to human-readable string
195pub fn format_bytes(bytes: usize) -> alloc::string::String {
196    const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
197    let mut size = bytes;
198    let mut unit_index = 0;
199
200    while size >= 1024 && unit_index < UNITS.len() - 1 {
201        size /= 1024;
202        unit_index += 1;
203    }
204
205    alloc::format!("{} {}", size, UNITS[unit_index])
206}