kernel: use rbtree in pmm
This commit is contained in:
parent
9dcc22c7de
commit
b069ee78a8
SSH key fingerprint:
SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
|
|
@ -7,6 +7,7 @@
|
|||
allocator_api,
|
||||
ptr_cast_slice,
|
||||
likely_unlikely,
|
||||
int_roundings,
|
||||
never_type
|
||||
)]
|
||||
#![cfg_attr(test, feature(custom_test_frameworks))]
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ extern "C" fn _start() -> ! {
|
|||
let leaf = kernel::x86_64::cpuid::Leaf8000008::get();
|
||||
serial_println!("max phy: {:#?}", leaf);
|
||||
|
||||
// kernel::serial_println!("PMM: {pmm:#?}");
|
||||
kernel::serial_println!("PMM: {pmm:#?}");
|
||||
|
||||
let fb = limine_requests::FRAMEBUFFER_REQUEST
|
||||
.framebuffers()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
use core::{alloc::Allocator, fmt::Debug};
|
||||
use core::{alloc::Allocator, cell::Cell, fmt::Debug, num::NonZeroUsize, ptr::NonNull};
|
||||
|
||||
use crate::{memory::page_tree::PageTree, serial_println, sync::OnceLock, x86_64::PAGE_SIZE};
|
||||
use bit_field::BitField;
|
||||
use rbtree::RBTree;
|
||||
|
||||
use crate::{serial_println, sync::OnceLock, x86_64::PAGE_SIZE};
|
||||
|
||||
pub static HHDM_BASE: OnceLock<u64> = OnceLock::new();
|
||||
|
||||
|
|
@ -12,6 +15,12 @@ pub trait VirtAddrTranslationExt {
|
|||
pub struct PhyAddr(pub u64);
|
||||
|
||||
impl PhyAddr {
|
||||
pub fn from_hhdm_virt(virt: impl Into<VirtAddr>) -> Option<Self> {
|
||||
let hhdm_base = unsafe { *HHDM_BASE.get_unchecked() };
|
||||
|
||||
virt.into().0.checked_sub(hhdm_base).map(PhyAddr)
|
||||
}
|
||||
|
||||
pub fn into_hhdm_virt(&self) -> VirtAddr {
|
||||
VirtAddr(self.0 + unsafe { crate::memory::HHDM_BASE.get().unwrap_unchecked() })
|
||||
}
|
||||
|
|
@ -21,6 +30,10 @@ impl PhyAddr {
|
|||
pub fn byte_add(&self, offset: usize) -> PhyAddr {
|
||||
PhyAddr(self.0 + offset as u64)
|
||||
}
|
||||
|
||||
pub fn page_index(&self) -> u64 {
|
||||
self.0.div_floor(PAGE_SIZE as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for PhyAddr {
|
||||
|
|
@ -29,6 +42,36 @@ impl Debug for PhyAddr {
|
|||
}
|
||||
}
|
||||
|
||||
impl<T> From<&T> for VirtAddr {
|
||||
fn from(ptr: &T) -> Self {
|
||||
VirtAddr(ptr as *const T as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<&mut T> for VirtAddr {
|
||||
fn from(ptr: &mut T) -> Self {
|
||||
VirtAddr(ptr as *mut T as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<*const T> for VirtAddr {
|
||||
fn from(ptr: *const T) -> Self {
|
||||
VirtAddr(ptr as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<*mut T> for VirtAddr {
|
||||
fn from(ptr: *mut T) -> Self {
|
||||
VirtAddr(ptr as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<NonNull<T>> for VirtAddr {
|
||||
fn from(ptr: NonNull<T>) -> Self {
|
||||
VirtAddr(ptr.as_ptr() as u64)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub struct VirtAddr(pub u64);
|
||||
|
||||
|
|
@ -69,7 +112,35 @@ pub struct PhysicalMemoryManager {
|
|||
}
|
||||
|
||||
pub struct PhysicalMemoryAllocator {
|
||||
tree: PageTree,
|
||||
tree: RBTree<PhysicalPageNode>,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Debug)]
|
||||
struct PageHeader {
|
||||
node: PhysicalPageNode,
|
||||
count: u64,
|
||||
}
|
||||
|
||||
impl PageHeader {
|
||||
fn phy(&self) -> Option<PhyAddr> {
|
||||
PhyAddr::from_hhdm_virt(self)
|
||||
}
|
||||
|
||||
fn new_from_page_idx_and_count(page_idx: usize, count: usize) -> NonNull<Self> {
|
||||
let phy = PhyAddr(page_idx as u64 * PAGE_SIZE as u64);
|
||||
let virt = phy.into_hhdm_virt();
|
||||
let ptr = virt.as_mut::<Self>();
|
||||
|
||||
unsafe {
|
||||
ptr.write(PageHeader {
|
||||
node: PhysicalPageNode::new_red(),
|
||||
count: count as u64,
|
||||
});
|
||||
}
|
||||
|
||||
NonNull::new(ptr).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for PhysicalMemoryAllocator {
|
||||
|
|
@ -78,8 +149,17 @@ impl Debug for PhysicalMemoryAllocator {
|
|||
.field_with("tree", |f| {
|
||||
let iter = self.tree.iter();
|
||||
write!(f, "[")?;
|
||||
for info in iter {
|
||||
writeln!(f, "({:?}..{:?}), ", info.phy, info.end())?;
|
||||
for key in iter {
|
||||
let header = unsafe { (&raw const *key).cast::<PageHeader>().read_volatile() };
|
||||
|
||||
if let Some(phy) = header.phy() {
|
||||
writeln!(
|
||||
f,
|
||||
"({:?}..{:?}), ",
|
||||
phy,
|
||||
phy.page_add(header.count as usize)
|
||||
)?;
|
||||
}
|
||||
}
|
||||
write!(f, "]")
|
||||
})
|
||||
|
|
@ -94,7 +174,7 @@ impl PhysicalMemoryAllocator {
|
|||
.filter(|region| region.region_type.is_usable());
|
||||
|
||||
let mut pmm = PhysicalMemoryAllocator {
|
||||
tree: PageTree::from_root(usize::MAX),
|
||||
tree: RBTree::new(),
|
||||
};
|
||||
|
||||
usable_regions.for_each(|region| {
|
||||
|
|
@ -113,7 +193,16 @@ impl PhysicalMemoryAllocator {
|
|||
page_index,
|
||||
count
|
||||
);
|
||||
self.tree.free_region(page_index, count);
|
||||
|
||||
let node_ptr =
|
||||
PageHeader::new_from_page_idx_and_count(page_index, count).cast::<PhysicalPageNode>();
|
||||
|
||||
let None = self.tree.insert_node(node_ptr) else {
|
||||
panic!(
|
||||
"Failed to insert node for page_index = {}, count = {}",
|
||||
page_index, count
|
||||
);
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -129,17 +218,14 @@ unsafe impl Allocator for PanicingAllocator {
|
|||
fn allocate(
|
||||
&self,
|
||||
layout: core::alloc::Layout,
|
||||
) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
|
||||
) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
|
||||
match layout.size() {
|
||||
0 => Ok(core::ptr::NonNull::slice_from_raw_parts(
|
||||
core::ptr::NonNull::dangling(),
|
||||
0,
|
||||
)),
|
||||
0 => Ok(NonNull::slice_from_raw_parts(NonNull::dangling(), 0)),
|
||||
_ => panic!("PanicingAllocator cannot allocate memory"),
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn deallocate(&self, _ptr: core::ptr::NonNull<u8>, _layout: core::alloc::Layout) {}
|
||||
unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: core::alloc::Layout) {}
|
||||
}
|
||||
|
||||
pub mod bump {
|
||||
|
|
@ -559,7 +645,7 @@ pub mod bump {
|
|||
}
|
||||
}
|
||||
|
||||
fn asdf(bump: &mut Bump<super::PanicingAllocator>) {
|
||||
fn _asdf(bump: &mut Bump<super::PanicingAllocator>) {
|
||||
let mut bump = bump.as_scope();
|
||||
let x = bump.alloc_with(|| 3u64);
|
||||
bump.scope(|bump| {
|
||||
|
|
@ -571,558 +657,165 @@ pub mod bump {
|
|||
}
|
||||
}
|
||||
|
||||
mod page_tree {
|
||||
use bit_field::BitField;
|
||||
|
||||
use crate::{
|
||||
memory::{PAGE_SIZE, PhyAddr},
|
||||
serial_println,
|
||||
};
|
||||
|
||||
use core::ops::{Index, IndexMut};
|
||||
|
||||
pub enum SearchResult<T> {
|
||||
Found(T),
|
||||
NotFound(T),
|
||||
}
|
||||
|
||||
/// A node for a red-black tree of free physical page chunks.
|
||||
///
|
||||
/// On amd64 platforms, the maximum physical address is 52 bits, the lower
|
||||
/// 12 of which are zero for page aligned addresses.
|
||||
/// Our Tree Node entry needs to store 3 page indices (left, right, parent) and a color bit.
|
||||
struct CompactPage(u128);
|
||||
/// Our Tree Node entry needs to store 3 page indices (parent, left, right) and a color bit.
|
||||
struct PhysicalPageNode(Cell<u128>);
|
||||
|
||||
impl CompactPage {
|
||||
fn parent(&self) -> usize {
|
||||
self.0.get_bits(0..40) as usize
|
||||
}
|
||||
fn set_parent(&mut self, parent: usize) {
|
||||
self.0.set_bits(0..40, parent as u128);
|
||||
}
|
||||
impl Debug for PhysicalPageNode {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
use rbtree::UnsafeNode;
|
||||
|
||||
fn left_child(&self) -> usize {
|
||||
self.0.get_bits(40..80) as usize
|
||||
}
|
||||
fn set_left_child(&mut self, left: usize) {
|
||||
self.0.set_bits(40..80, left as u128);
|
||||
}
|
||||
|
||||
fn right_child(&self) -> usize {
|
||||
self.0.get_bits(80..120) as usize
|
||||
}
|
||||
fn set_right_child(&mut self, right: usize) {
|
||||
self.0.set_bits(80..120, right as u128);
|
||||
}
|
||||
|
||||
fn color(&self) -> bool {
|
||||
self.0.get_bit(120)
|
||||
}
|
||||
fn set_color(&mut self, color: bool) {
|
||||
self.0.set_bit(120, color);
|
||||
}
|
||||
|
||||
fn data(&self) -> u8 {
|
||||
self.0.get_bits(121..128) as u8
|
||||
}
|
||||
fn set_data(&mut self, data: u8) {
|
||||
self.0.set_bits(121..128, data as u128);
|
||||
f.debug_struct("PhysicalPageNode")
|
||||
.field("parent", &self.parent())
|
||||
.field("left", &self.left())
|
||||
.field("right", &self.right())
|
||||
.field("color", &self.color())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
struct Page {
|
||||
left_idx: usize,
|
||||
right_idx: usize,
|
||||
parent: usize,
|
||||
count: isize,
|
||||
impl PhysicalPageNode {
|
||||
fn new_red() -> Self {
|
||||
PhysicalPageNode(Cell::new(1 << 120))
|
||||
}
|
||||
|
||||
impl Page {
|
||||
fn new_red(_idx: usize, count: usize) -> Self {
|
||||
Self {
|
||||
left_idx: usize::MAX,
|
||||
right_idx: usize::MAX,
|
||||
parent: usize::MAX,
|
||||
count: -(count as isize),
|
||||
fn bits(&self) -> u128 {
|
||||
self.0.get()
|
||||
}
|
||||
|
||||
fn parent_bits(&self) -> u64 {
|
||||
self.bits().get_bits(0..40) as u64
|
||||
}
|
||||
fn child(&self, left: bool) -> usize {
|
||||
if left { self.left_idx } else { self.right_idx }
|
||||
|
||||
fn set_parent_bits(&self, parent: u64) {
|
||||
self.0.update(|mut bits| {
|
||||
bits.set_bits(0..40, parent as u128);
|
||||
bits
|
||||
});
|
||||
}
|
||||
fn red(&self) -> bool {
|
||||
self.count.is_negative()
|
||||
|
||||
fn left_bits(&self) -> u64 {
|
||||
self.bits().get_bits(40..80) as u64
|
||||
}
|
||||
fn set_color(&mut self, red: bool) {
|
||||
if red {
|
||||
self.count = -(self.count.unsigned_abs() as isize);
|
||||
} else {
|
||||
self.count = self.count.unsigned_abs() as isize;
|
||||
|
||||
fn set_left_bits(&self, left: u64) {
|
||||
self.0.update(|mut bits| {
|
||||
bits.set_bits(40..80, left as u128);
|
||||
bits
|
||||
});
|
||||
}
|
||||
|
||||
fn right_bits(&self) -> u64 {
|
||||
self.bits().get_bits(80..120) as u64
|
||||
}
|
||||
fn count(&self) -> usize {
|
||||
self.count.unsigned_abs()
|
||||
|
||||
fn set_right_bits(&self, right: u64) {
|
||||
self.0.update(|mut bits| {
|
||||
bits.set_bits(80..120, right as u128);
|
||||
bits
|
||||
});
|
||||
}
|
||||
fn set_count(&mut self, count: usize) {
|
||||
let red = self.red();
|
||||
self.count = if red {
|
||||
-(count as isize)
|
||||
} else {
|
||||
count as isize
|
||||
};
|
||||
|
||||
fn color_bit(&self) -> bool {
|
||||
self.bits().get_bit(120)
|
||||
}
|
||||
|
||||
fn set_color_bit(&self, color: bool) {
|
||||
self.0.update(|mut bits| {
|
||||
bits.set_bit(120, color);
|
||||
bits
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[repr(transparent)]
|
||||
pub struct PageTree {
|
||||
root: usize,
|
||||
pages: Pages,
|
||||
}
|
||||
/// A key for a red-black tree of free physical page chunks.
|
||||
/// The key uses the address of the node as the key, which is unique and means
|
||||
/// the key does not require any additional bits.
|
||||
struct PhysicalPageNodeKey;
|
||||
|
||||
struct Pages;
|
||||
impl Eq for PhysicalPageNodeKey {}
|
||||
|
||||
impl Pages {
|
||||
fn get_disjoint_mut<const N: usize>(&mut self, indices: [usize; N]) -> [&mut Page; N] {
|
||||
indices.map(|idx| unsafe {
|
||||
PhyAddr(idx as u64 * PAGE_SIZE as u64)
|
||||
.into_hhdm_virt()
|
||||
.as_mut::<Page>()
|
||||
.as_mut_unchecked()
|
||||
})
|
||||
}
|
||||
|
||||
fn get_ptr(&self, idx: usize) -> *const Page {
|
||||
PhyAddr(idx as u64 * PAGE_SIZE as u64)
|
||||
.into_hhdm_virt()
|
||||
.as_ptr::<Page>()
|
||||
}
|
||||
|
||||
fn get_mut(&mut self, idx: usize) -> &mut Page {
|
||||
unsafe {
|
||||
PhyAddr(idx as u64 * PAGE_SIZE as u64)
|
||||
.into_hhdm_virt()
|
||||
.as_mut::<Page>()
|
||||
.as_mut_unchecked()
|
||||
}
|
||||
impl PartialEq for PhysicalPageNodeKey {
|
||||
fn eq(&self, _other: &Self) -> bool {
|
||||
core::ptr::eq(self, _other)
|
||||
}
|
||||
}
|
||||
|
||||
impl Index<usize> for Pages {
|
||||
type Output = Page;
|
||||
fn index(&self, idx: usize) -> &Self::Output {
|
||||
unsafe {
|
||||
PhyAddr(idx as u64 * PAGE_SIZE as u64)
|
||||
.into_hhdm_virt()
|
||||
.as_ptr::<Page>()
|
||||
.as_ref_unchecked()
|
||||
}
|
||||
impl PartialOrd for PhysicalPageNodeKey {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
impl IndexMut<usize> for Pages {
|
||||
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
|
||||
unsafe {
|
||||
PhyAddr(idx as u64 * PAGE_SIZE as u64)
|
||||
.into_hhdm_virt()
|
||||
.as_mut::<Page>()
|
||||
.as_mut_unchecked()
|
||||
}
|
||||
impl Ord for PhysicalPageNodeKey {
|
||||
fn cmp(&self, other: &Self) -> core::cmp::Ordering {
|
||||
(&raw const *self).cmp(&(&raw const *other))
|
||||
}
|
||||
}
|
||||
|
||||
impl PageTree {
|
||||
pub fn from_root(root: usize) -> Self {
|
||||
Self { root, pages: Pages }
|
||||
}
|
||||
pub fn free_region(&mut self, page_index: usize, count: usize) {
|
||||
let pre = match self.find(page_index) {
|
||||
SearchResult::Found(_) => panic!("Page index already exists in free list"),
|
||||
SearchResult::NotFound(Self::INVALID_IDX) => Self::INVALID_IDX,
|
||||
SearchResult::NotFound(idx) => {
|
||||
if idx > page_index {
|
||||
self.next_back_of(idx)
|
||||
} else {
|
||||
idx
|
||||
}
|
||||
}
|
||||
};
|
||||
unsafe impl rbtree::UnsafeNode for PhysicalPageNode {
|
||||
type Key = PhysicalPageNodeKey;
|
||||
|
||||
match pre {
|
||||
Self::INVALID_IDX => {
|
||||
self.insert(page_index, count);
|
||||
fn left(&self) -> Option<NonNull<Self>> {
|
||||
NonZeroUsize::new((self.left_bits() << 12) as usize).map(NonNull::with_exposed_provenance)
|
||||
}
|
||||
_ => {
|
||||
if self.pages[pre].count() + pre == page_index {
|
||||
serial_println!(
|
||||
"Found previous page: pre = {}, pre_count = {}",
|
||||
pre,
|
||||
self.pages[pre].count()
|
||||
|
||||
fn right(&self) -> Option<NonNull<Self>> {
|
||||
NonZeroUsize::new((self.right_bits() << 12) as usize).map(NonNull::with_exposed_provenance)
|
||||
}
|
||||
|
||||
fn parent(&self) -> Option<NonNull<Self>> {
|
||||
NonZeroUsize::new((self.parent_bits() << 12) as usize).map(NonNull::with_exposed_provenance)
|
||||
}
|
||||
|
||||
fn key(&self) -> &Self::Key {
|
||||
assert_eq!(
|
||||
core::mem::size_of::<Self::Key>(),
|
||||
0,
|
||||
"PhysicalPageNodeKey must be zero-sized"
|
||||
);
|
||||
|
||||
let new_count = self.pages[pre].count() + count;
|
||||
self.pages[pre].set_count(new_count);
|
||||
// SAFETY: The key is a zero-sized type, so we can transmute the
|
||||
// reference to the node to a reference to the key.
|
||||
unsafe { core::mem::transmute_copy::<&PhysicalPageNode, &PhysicalPageNodeKey>(&self) }
|
||||
}
|
||||
|
||||
fn color(&self) -> rbtree::Color {
|
||||
if self.color_bit() {
|
||||
rbtree::Color::Red
|
||||
} else {
|
||||
self.insert(page_index, count);
|
||||
}
|
||||
}
|
||||
rbtree::Color::Black
|
||||
}
|
||||
}
|
||||
|
||||
// rb-tree impl
|
||||
|
||||
const INVALID_IDX: usize = usize::MAX;
|
||||
fn rotate(&mut self, x: usize, left: bool) {
|
||||
let (y, b, x_hole, y_hole) = if left {
|
||||
let y = self.pages[x].right_idx;
|
||||
let b = self.pages[y].left_idx;
|
||||
|
||||
let [x_page, y_page] = self.pages.get_disjoint_mut([x, y]);
|
||||
let x_hole = &mut x_page.right_idx;
|
||||
let y_hole = &mut y_page.left_idx;
|
||||
|
||||
(y, b, x_hole, y_hole)
|
||||
} else {
|
||||
let y = self.pages[x].left_idx;
|
||||
let b = self.pages[y].right_idx;
|
||||
|
||||
let [x_page, y_page] = self.pages.get_disjoint_mut([x, y]);
|
||||
let x_hole = &mut x_page.left_idx;
|
||||
let y_hole = &mut y_page.right_idx;
|
||||
|
||||
(y, b, x_hole, y_hole)
|
||||
};
|
||||
|
||||
*x_hole = b;
|
||||
*y_hole = x;
|
||||
|
||||
if b != usize::MAX {
|
||||
self.pages[b].parent = x;
|
||||
fn set_left(&self, left: Option<NonNull<Self>>) {
|
||||
self.set_left_bits(left.map_or(0, |ptr| {
|
||||
PhyAddr::from_hhdm_virt(VirtAddr::from(ptr))
|
||||
.expect("PhysicalPageNode left pointer is not in HHDM")
|
||||
.page_index()
|
||||
}));
|
||||
}
|
||||
|
||||
let x_parent = self.pages[x].parent;
|
||||
self.pages[y].parent = x_parent;
|
||||
if self.pages[x].parent == usize::MAX {
|
||||
self.root = y;
|
||||
} else if x == self.pages[x_parent].left_idx {
|
||||
self.pages[x_parent].left_idx = y;
|
||||
} else {
|
||||
self.pages[x_parent].right_idx = y;
|
||||
fn set_right(&self, right: Option<NonNull<Self>>) {
|
||||
self.set_right_bits(right.map_or(0, |ptr| {
|
||||
PhyAddr::from_hhdm_virt(VirtAddr::from(ptr))
|
||||
.expect("PhysicalPageNode right pointer is not in HHDM")
|
||||
.page_index()
|
||||
}));
|
||||
}
|
||||
|
||||
self.pages[x].parent = y;
|
||||
fn set_parent(&self, parent: Option<NonNull<Self>>) {
|
||||
self.set_parent_bits(parent.map_or(0, |ptr| {
|
||||
PhyAddr::from_hhdm_virt(VirtAddr::from(ptr))
|
||||
.expect("PhysicalPageNode parentpointer is not in HHDM")
|
||||
.page_index()
|
||||
}));
|
||||
}
|
||||
|
||||
pub fn find(&self, idx: usize) -> SearchResult<usize> {
|
||||
let mut node = self.root;
|
||||
let mut parent = usize::MAX;
|
||||
while node != usize::MAX {
|
||||
parent = node;
|
||||
if idx < node {
|
||||
node = self.pages[node].left_idx;
|
||||
} else if idx > node {
|
||||
node = self.pages[node].right_idx;
|
||||
} else {
|
||||
return SearchResult::Found(node);
|
||||
}
|
||||
}
|
||||
|
||||
SearchResult::NotFound(parent)
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, idx: usize, count: usize) {
|
||||
let SearchResult::NotFound(parent) = self.find(idx) else {
|
||||
panic!("Page index already exists in free list");
|
||||
};
|
||||
|
||||
let page = &mut self.pages[idx];
|
||||
*page = Page::new_red(idx, count);
|
||||
|
||||
page.parent = parent;
|
||||
if parent == usize::MAX {
|
||||
self.root = idx;
|
||||
} else if idx < parent {
|
||||
self.set_left_child_of(parent, idx);
|
||||
} else {
|
||||
self.set_right_child_of(parent, idx);
|
||||
}
|
||||
|
||||
self.fixup_insert(idx);
|
||||
}
|
||||
|
||||
fn fixup_insert(&mut self, mut idx: usize) {
|
||||
// idx is red; fixup the tree while the parent is red
|
||||
loop {
|
||||
let parent = self.parent_of(idx);
|
||||
if parent == usize::MAX || !self.color_of(parent) {
|
||||
break;
|
||||
}
|
||||
|
||||
// gp is guaranteed to exist because the parent is red and the root is black
|
||||
let gp = self.pages[parent].parent;
|
||||
let (uncle, uncle_left) = if parent == self.pages[gp].left_idx {
|
||||
(self.pages[gp].right_idx, false)
|
||||
} else {
|
||||
(self.pages[gp].left_idx, true)
|
||||
};
|
||||
|
||||
if self.color_of(uncle) {
|
||||
self.pages[parent].set_color(false);
|
||||
self.pages[uncle].set_color(false);
|
||||
self.pages[gp].set_color(true);
|
||||
idx = gp;
|
||||
} else {
|
||||
// uncle is black
|
||||
if idx == self.pages[parent].child(uncle_left) {
|
||||
// idx is on the same side as uncle
|
||||
self.rotate(parent, !uncle_left);
|
||||
idx = parent;
|
||||
}
|
||||
|
||||
self.pages[parent].set_color(false);
|
||||
self.pages[gp].set_color(true);
|
||||
self.rotate(gp, uncle_left);
|
||||
}
|
||||
}
|
||||
|
||||
self.pages[self.root].set_color(false);
|
||||
}
|
||||
|
||||
/// replaces the node at `at_idx` with the node at `with_idx`.
|
||||
fn replace(&mut self, at_idx: usize, with_idx: usize) {
|
||||
let at_parent = self.parent_of(at_idx);
|
||||
|
||||
match at_parent {
|
||||
Self::INVALID_IDX => {}
|
||||
_ if at_idx == self.left_child_of(at_parent) => {
|
||||
self.set_left_child_of(at_parent, with_idx);
|
||||
}
|
||||
_ => {
|
||||
self.set_right_child_of(at_parent, with_idx);
|
||||
}
|
||||
}
|
||||
|
||||
if let Self::INVALID_IDX = with_idx {
|
||||
} else {
|
||||
self.set_parent_of(with_idx, at_parent);
|
||||
}
|
||||
}
|
||||
|
||||
/// returns the next node in the tree after `idx`, or usize::MAX if there is no next node.
|
||||
fn minimum_of(&self, mut idx: usize) -> usize {
|
||||
loop {
|
||||
let next = self.left_child_of(idx);
|
||||
if next == usize::MAX {
|
||||
break;
|
||||
}
|
||||
idx = next;
|
||||
}
|
||||
|
||||
idx
|
||||
}
|
||||
|
||||
fn maximum_of(&self, mut idx: usize) -> usize {
|
||||
loop {
|
||||
let next = self.right_child_of(idx);
|
||||
if next == usize::MAX {
|
||||
break;
|
||||
}
|
||||
idx = next;
|
||||
}
|
||||
|
||||
idx
|
||||
}
|
||||
|
||||
fn next_back_of(&self, mut idx: usize) -> usize {
|
||||
match self.left_child_of(idx) {
|
||||
Self::INVALID_IDX => {
|
||||
let mut p = self.parent_of(idx);
|
||||
while p != Self::INVALID_IDX && idx == self.left_child_of(p) {
|
||||
idx = p;
|
||||
p = self.parent_of(idx);
|
||||
}
|
||||
p
|
||||
}
|
||||
left => self.maximum_of(left),
|
||||
}
|
||||
}
|
||||
|
||||
fn next_of(&self, mut idx: usize) -> usize {
|
||||
match self.right_child_of(idx) {
|
||||
Self::INVALID_IDX => {
|
||||
let mut p = self.parent_of(idx);
|
||||
while p != Self::INVALID_IDX && idx == self.right_child_of(p) {
|
||||
idx = p;
|
||||
p = self.parent_of(idx);
|
||||
}
|
||||
p
|
||||
}
|
||||
right => self.minimum_of(right),
|
||||
}
|
||||
}
|
||||
|
||||
fn left_child_of(&self, idx: usize) -> usize {
|
||||
self.pages[idx].left_idx
|
||||
}
|
||||
fn right_child_of(&self, idx: usize) -> usize {
|
||||
self.pages[idx].right_idx
|
||||
}
|
||||
fn child_of(&self, idx: usize, left: bool) -> usize {
|
||||
if left {
|
||||
self.pages[idx].left_idx
|
||||
} else {
|
||||
self.pages[idx].right_idx
|
||||
}
|
||||
}
|
||||
fn children_of(&self, idx: usize) -> (usize, usize) {
|
||||
(self.pages[idx].left_idx, self.pages[idx].right_idx)
|
||||
}
|
||||
fn parent_of(&self, idx: usize) -> usize {
|
||||
self.pages[idx].parent
|
||||
}
|
||||
fn set_parent_of(&mut self, idx: usize, parent: usize) {
|
||||
if idx != Self::INVALID_IDX {
|
||||
self.pages[idx].parent = parent;
|
||||
}
|
||||
}
|
||||
fn set_color_of(&mut self, idx: usize, red: bool) {
|
||||
self.pages[idx].set_color(red);
|
||||
}
|
||||
fn color_of(&self, idx: usize) -> bool {
|
||||
// leaves are implicitly black
|
||||
idx != Self::INVALID_IDX && self.pages[idx].red()
|
||||
}
|
||||
fn set_left_child_of(&mut self, idx: usize, left: usize) {
|
||||
self.pages[idx].left_idx = left;
|
||||
}
|
||||
fn set_right_child_of(&mut self, idx: usize, right: usize) {
|
||||
self.pages[idx].right_idx = right;
|
||||
}
|
||||
|
||||
pub fn remove(&mut self, z: usize) {
|
||||
let y = match self.children_of(z) {
|
||||
(Self::INVALID_IDX, _) | (_, Self::INVALID_IDX) => z,
|
||||
_ => self.next_of(z),
|
||||
};
|
||||
|
||||
let x = match self.left_child_of(y) {
|
||||
Self::INVALID_IDX => self.right_child_of(y),
|
||||
_ => self.left_child_of(y),
|
||||
};
|
||||
|
||||
let p_y = self.parent_of(y);
|
||||
self.set_parent_of(x, p_y);
|
||||
|
||||
match p_y {
|
||||
Self::INVALID_IDX => {
|
||||
self.root = x;
|
||||
}
|
||||
_ if y == self.left_child_of(p_y) => {
|
||||
self.set_left_child_of(p_y, x);
|
||||
}
|
||||
_ => {
|
||||
self.set_right_child_of(p_y, x);
|
||||
}
|
||||
}
|
||||
|
||||
let color = self.color_of(y);
|
||||
if y != z {
|
||||
self.replace(z, y);
|
||||
}
|
||||
|
||||
if !color {
|
||||
self.fixup_remove(x);
|
||||
}
|
||||
}
|
||||
|
||||
fn fixup_remove(&mut self, mut x: usize) {
|
||||
loop {
|
||||
if x == self.root || self.color_of(x) {
|
||||
break;
|
||||
}
|
||||
|
||||
let p_x = self.parent_of(x);
|
||||
let (mut w, w_left) = if x == self.left_child_of(p_x) {
|
||||
(self.right_child_of(p_x), false)
|
||||
} else {
|
||||
(self.left_child_of(p_x), true)
|
||||
};
|
||||
|
||||
if self.color_of(w) {
|
||||
self.set_color_of(w, false);
|
||||
self.set_color_of(p_x, true);
|
||||
|
||||
self.rotate(p_x, !w_left);
|
||||
|
||||
w = self.child_of(p_x, w_left);
|
||||
}
|
||||
|
||||
// w is now black
|
||||
|
||||
match (
|
||||
self.color_of(self.left_child_of(w)),
|
||||
self.color_of(self.right_child_of(w)),
|
||||
) {
|
||||
(false, false) => {
|
||||
self.set_color_of(w, true);
|
||||
x = p_x;
|
||||
continue;
|
||||
}
|
||||
(left @ true, false) | (left @ false, true) => {
|
||||
if left != w_left {
|
||||
self.set_color_of(self.child_of(w, left), false);
|
||||
self.set_color_of(w, true);
|
||||
self.rotate(w, w_left);
|
||||
}
|
||||
|
||||
self.set_color_of(w, self.color_of(p_x));
|
||||
self.set_color_of(p_x, false);
|
||||
self.set_color_of(self.child_of(w, w_left), false);
|
||||
self.rotate(p_x, !w_left);
|
||||
x = self.root;
|
||||
}
|
||||
(true, true) => {
|
||||
x = self.root;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn iter(&self) -> PageTreeIter<'_> {
|
||||
PageTreeIter {
|
||||
tree: self,
|
||||
next: self.minimum_of(self.root),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PageTreeIter<'a> {
|
||||
tree: &'a PageTree,
|
||||
next: usize,
|
||||
}
|
||||
|
||||
pub struct PageInfo {
|
||||
pub phy: PhyAddr,
|
||||
pub count: usize,
|
||||
}
|
||||
|
||||
impl PageInfo {
|
||||
pub fn end(&self) -> PhyAddr {
|
||||
self.phy.page_add(self.count)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Iterator for PageTreeIter<'a> {
|
||||
type Item = PageInfo;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
if self.next == usize::MAX {
|
||||
return None;
|
||||
}
|
||||
|
||||
let page_info = PageInfo {
|
||||
phy: PhyAddr(self.next as u64 * PAGE_SIZE as u64),
|
||||
count: self.tree.pages[self.next].count(),
|
||||
};
|
||||
self.next = self.tree.next_of(self.next);
|
||||
|
||||
Some(page_info)
|
||||
}
|
||||
fn set_color(&self, color: rbtree::Color) {
|
||||
self.set_color_bit(color == rbtree::Color::Red);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue