1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

//! The `frame_allocator` module provides a frame allocator for the kernel.

use alloc::vec::Vec;

use lazy_static::lazy_static;

use crate::{
    constant::MEM_LIMIT,
    mem::{FrameNumber, PhysicalAddress},
    sync::Mutex,
};

/// The `FrameTracker` struct represents a frame in the physical memory.
/// It contains the frame number and is responsible for zeroing out the frame when it is created.
/// It deallocates the frame when it is dropped, which follows the RAII idiom.
pub struct FrameTracker {
    frame_number: FrameNumber,
}

impl FrameTracker {
    /// Initializes a frame with a specific [FrameNumber] and zeros out the frame.
    pub fn new(frame_number: FrameNumber) -> Self {
        for byte in frame_number.as_bytes_mut() {
            *byte = 0;
        }

        Self { frame_number }
    }

    pub fn frame_number(&self) -> FrameNumber {
        self.frame_number
    }
}

impl Drop for FrameTracker {
    fn drop(&mut self) {
        deallocate_frame(self.frame_number)
    }
}

trait FrameAllocator {
    fn new() -> Self;
    fn allocate(&mut self) -> Option<FrameNumber>;
    fn deallocate(&mut self, frame_number: FrameNumber);
}

pub struct StackFrameAllocator {
    frame_start: FrameNumber,
    frame_end: FrameNumber,
    deallocated_page: Vec<FrameNumber>,
}

impl StackFrameAllocator {
    fn init(&mut self, frame_start: FrameNumber, frame_end: FrameNumber) {
        self.frame_start = frame_start;
        self.frame_end = frame_end;
    }
}

impl FrameAllocator for StackFrameAllocator {
    fn new() -> Self {
        StackFrameAllocator {
            frame_start: FrameNumber::from(0),
            frame_end: FrameNumber::from(0),
            deallocated_page: Vec::new(),
        }
    }

    fn allocate(&mut self) -> Option<FrameNumber> {
        if let Some(frame_number) = self.deallocated_page.pop() {
            Some(frame_number)
        } else if self.frame_start == self.frame_end {
            None
        } else {
            let result = Some(self.frame_start);
            self.frame_start += 1;
            result
        }
    }

    fn deallocate(&mut self, frame_number: FrameNumber) {
        if self.frame_start <= frame_number || self.deallocated_page.contains(&frame_number) {
            panic!(
                "the frame {:#x} has not been allocated",
                usize::from(frame_number)
            )
        }
        self.deallocated_page.push(frame_number);
    }
}

lazy_static! {
    pub static ref FRAME_ALLOCATOR: Mutex<StackFrameAllocator> =
        Mutex::new(StackFrameAllocator::new());
}

/// Initializes a frame allocator that manages the physical address from `kernel_end` to
/// [MEM_LIMIT].
pub fn init() {
    extern "C" {
        fn kernel_end();
    }

    FRAME_ALLOCATOR.lock().init(
        PhysicalAddress::from(kernel_end as usize).ceil(),
        PhysicalAddress::from(MEM_LIMIT).floor(),
    );
}

/// Allocates a frame and returns a [FrameTracker] to track the allocated frame when succeeded.
pub fn allocate_frame() -> Option<FrameTracker> {
    FRAME_ALLOCATOR.lock().allocate().map(FrameTracker::new)
}

/// Deallocates the frame with a specific [FrameNumber].
pub fn deallocate_frame(frame_number: FrameNumber) {
    FRAME_ALLOCATOR.lock().deallocate(frame_number);
}