virtio-spec-rs/src/virtq/alloc.rs at db19647ed811ee8ceb30eb01d6c2c9b3548739f0 · rust-osdev/virtio-spec-rs · GitHub

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
use ::alloc::alloc::handle_alloc_error;
#[cfg(feature = "nightly")]
use ::alloc::{
    alloc::{AllocError, Allocator, Global},
    boxed::Box,
};
#[cfg(not(feature = "nightly"))]
use allocator_api2::{
    alloc::{AllocError, Allocator, Global},
    boxed::Box,
};

use super::*;

impl Avail {
    pub fn new(queue_size: u16, has_event_idx: bool) -> Box<Self> {
        Self::new_in(queue_size, has_event_idx, Global)
    }

    pub fn try_new(queue_size: u16, has_event_idx: bool) -> Result<Box<Self>, AllocError> {
        Self::try_new_in(queue_size, has_event_idx, Global)
    }

    pub fn new_in<A: Allocator>(queue_size: u16, has_event_idx: bool, alloc: A) -> Box<Self, A> {
        Self::try_new_in(queue_size, has_event_idx, alloc)
            .unwrap_or_else(|_| handle_alloc_error(Self::layout(queue_size, has_event_idx)))
    }

    pub fn try_new_in<A: Allocator>(
        queue_size: u16,
        has_event_idx: bool,
        alloc: A,
    ) -> Result<Box<Self, A>, AllocError> {
        let layout = Self::layout(queue_size, has_event_idx);

        let mem = alloc.allocate_zeroed(layout)?;
        let mem = NonNull::slice_from_raw_parts(mem.cast(), layout.size());
        let raw = Self::from_ptr(mem).unwrap();
        let boxed = unsafe { Box::from_raw_in(raw.as_ptr(), alloc) };

        debug_assert_eq!(Layout::for_value(&*boxed), layout);
        debug_assert_eq!(boxed.ring(has_event_idx).len(), queue_size.into());
        debug_assert_eq!(boxed.used_event(has_event_idx).is_some(), has_event_idx);

        Ok(boxed)
    }
}

impl Used {
    pub fn new(queue_size: u16, has_event_idx: bool) -> Box<Self> {
        Self::new_in(queue_size, has_event_idx, Global)
    }

    pub fn try_new(queue_size: u16, has_event_idx: bool) -> Result<Box<Self>, AllocError> {
        Self::try_new_in(queue_size, has_event_idx, Global)
    }

    pub fn new_in<A: Allocator>(queue_size: u16, has_event_idx: bool, alloc: A) -> Box<Self, A> {
        Self::try_new_in(queue_size, has_event_idx, alloc)
            .unwrap_or_else(|_| handle_alloc_error(Self::layout(queue_size, has_event_idx)))
    }

    pub fn try_new_in<A: Allocator>(
        queue_size: u16,
        has_event_idx: bool,
        alloc: A,
    ) -> Result<Box<Self, A>, AllocError> {
        let layout = Self::layout(queue_size, has_event_idx);

        let mem = alloc.allocate_zeroed(layout)?;
        let mem = NonNull::slice_from_raw_parts(mem.cast(), layout.size());
        let raw = Self::from_ptr(mem, has_event_idx).unwrap();
        let boxed = unsafe { Box::from_raw_in(raw.as_ptr(), alloc) };

        debug_assert_eq!(Layout::for_value(&*boxed), layout);
        debug_assert_eq!(boxed.ring().len(), queue_size.into());
        debug_assert_eq!(boxed.avail_event().is_some(), has_event_idx);

        Ok(boxed)
    }
}

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

    #[test]
    fn avail_layout() {
        for queue_size in [255, 256, 257] {
            for has_event_idx in [false, true] {
                Avail::new(queue_size, has_event_idx);
            }
        }
    }

    #[test]
    fn used_layout() {
        for queue_size in [255, 256, 257] {
            for has_event_idx in [false, true] {
                Used::new(queue_size, has_event_idx);
            }
        }
    }
}