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
119
120
121
122
123
124
125

// SPDX-License-Identifier: MPL-2.0

//! Options for allocating frames

use super::{Frame, FrameVec, Segment};
use crate::{mm::page::allocator, prelude::*, Error};

/// Options for allocating physical memory pages (or frames).
///
/// All allocated frames are safe to use in the sense that they are
/// not _typed memory_. We define typed memory as the memory that
/// may store Rust objects or affect Rust memory safety, e.g.,
/// the code and data segments of the OS kernel, the stack and heap
/// allocated for the OS kernel.
pub struct FrameAllocOptions {
    nframes: usize,
    is_contiguous: bool,
    uninit: bool,
}

impl FrameAllocOptions {
    /// Creates new options for allocating the specified number of frames.
    pub fn new(nframes: usize) -> Self {
        Self {
            nframes,
            is_contiguous: false,
            uninit: false,
        }
    }

    /// Sets whether the allocated frames should be contiguous.
    ///
    /// The default value is `false`.
    pub fn is_contiguous(&mut self, is_contiguous: bool) -> &mut Self {
        self.is_contiguous = is_contiguous;
        self
    }

    /// Sets whether the allocated frames should be uninitialized.
    ///
    /// If `uninit` is set as `false`, the frame will be zeroed once allocated.
    /// If `uninit` is set as `true`, the frame will **NOT** be zeroed and should *NOT* be read before writing.
    ///
    /// The default value is false.
    pub fn uninit(&mut self, uninit: bool) -> &mut Self {
        self.uninit = uninit;
        self
    }

    /// Allocates a collection of page frames according to the given options.
    pub fn alloc(&self) -> Result<FrameVec> {
        let frames = if self.is_contiguous {
            allocator::alloc(self.nframes).ok_or(Error::NoMemory)?
        } else {
            let mut frame_list = Vec::new();
            for _ in 0..self.nframes {
                frame_list.push(allocator::alloc_single().ok_or(Error::NoMemory)?);
            }
            FrameVec(frame_list)
        };
        if !self.uninit {
            for frame in frames.iter() {
                frame.writer().fill(0);
            }
        }

        Ok(frames)
    }

    /// Allocates a single page frame according to the given options.
    pub fn alloc_single(&self) -> Result<Frame> {
        if self.nframes != 1 {
            return Err(Error::InvalidArgs);
        }

        let frame = allocator::alloc_single().ok_or(Error::NoMemory)?;
        if !self.uninit {
            frame.writer().fill(0);
        }

        Ok(frame)
    }

    /// Allocates a contiguous range of page frames according to the given options.
    ///
    /// The returned [`Segment`] contains at least one page frame.
    pub fn alloc_contiguous(&self) -> Result<Segment> {
        // It's no use to checking `self.is_contiguous` here.
        if self.nframes == 0 {
            return Err(Error::InvalidArgs);
        }

        let segment = allocator::alloc_contiguous(self.nframes).ok_or(Error::NoMemory)?;
        if !self.uninit {
            segment.writer().fill(0);
        }

        Ok(segment)
    }
}

#[cfg(ktest)]
#[ktest]
fn test_alloc_dealloc() {
    // Here we allocate and deallocate frames in random orders to test the allocator.
    // We expect the test to fail if the underlying implementation panics.
    let single_options = FrameAllocOptions::new(1);
    let multi_options = FrameAllocOptions::new(10);
    let mut contiguous_options = FrameAllocOptions::new(10);
    contiguous_options.is_contiguous(true);
    let mut remember_vec = Vec::new();
    for i in 0..10 {
        for i in 0..10 {
            let single_frame = single_options.alloc_single().unwrap();
            if i % 3 == 0 {
                remember_vec.push(single_frame);
            }
        }
        let contiguous_segment = contiguous_options.alloc_contiguous().unwrap();
        drop(contiguous_segment);
        let multi_frames = multi_options.alloc().unwrap();
        remember_vec.extend(multi_frames.into_iter());
        remember_vec.pop();
    }
}