diff --git a/Cargo.toml b/Cargo.toml index 0cebe72..fb23b09 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] -name = "crate-template" +name = "vm-allocator" version = "0.1.0" -authors = [TODO] +authors = ["Liu Jiang "] license = "Apache-2.0 OR BSD-3-Clause" edition = "2018" diff --git a/src/interval_tree.rs b/src/interval_tree.rs new file mode 100644 index 0000000..033776c --- /dev/null +++ b/src/interval_tree.rs @@ -0,0 +1,1144 @@ +// Copyright (C) 2019 Alibaba Cloud. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! An interval tree implementation specialized for VMM resource management. +//! +//! It's not designed as a generic interval tree, but specialized for VMM resource management. +//! In addition to the normal get/insert/delete/update operations, it also implements allocate/free. +//! +//! # Examples +//! ```rust +//! extern crate vm_allocator; +//! use vm_allocator::{IntervalTree, Range, Constraint, NodeState}; +//! +//! // Create an interval tree and add available resources. +//! let mut tree = IntervalTree::::new(); +//! tree.insert(Range::new(0x100u32, 0x100u32), None); +//! tree.insert(Range::new(0x200u16, 0x2ffu16), None); +//! +//! // Allocate a range with constraints. +//! let mut constraint = Constraint::new(8u64); +//! constraint.min = 0x211; +//! constraint.max = 0x21f; +//! constraint.align = 0x8; +//! +//! let key = tree.allocate(&constraint); +//! assert_eq!(key, Some(Range::new(0x218u64, 0x21fu64))); +//! let val = tree.get(&Range::new(0x218u64, 0x21fu64)); +//! assert_eq!(val, Some(NodeState::Allocated)); +//! +//! // Associate data with the allocated range and mark the range as occupied. +//! // Note: caller needs to protect from concurrent access between allocate() and the first call +//! // to update() to mark range as occupied. +//! let old = tree.update(&Range::new(0x218u32, 0x21fu32), 2); +//! assert_eq!(old, None); +//! let old = tree.update(&Range::new(0x218u32, 0x21fu32), 3); +//! assert_eq!(old, Some(2)); +//! let val = tree.get(&Range::new(0x218u32, 0x21fu32)); +//! assert_eq!(val, Some(NodeState::Valued(&3))); +//! +//! // Free allocated resource. +//! let old = tree.free(key.as_ref().unwrap()); +//! assert_eq!(old, Some(3)); +//! +//! ``` + +use crate::{AllocPolicy, Constraint}; +use std::cmp::{max, min, Ordering}; + +/// A closed interval range [min, max]. +#[allow(missing_docs)] +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Hash)] +pub struct Range { + pub min: u64, + pub max: u64, +} + +impl Range { + /// Create a new Range object. + /// + /// Note: panic if encounters invalid parameters. + pub fn new(min: T, max: T) -> Self + where + u64: From, + { + let umin = u64::from(min); + let umax = u64::from(max); + if umin > umax || (umin == 0 && umax == std::u64::MAX) { + panic!("interval_tree: Range({}, {}) is invalid", umin, umax); + } + Range { + min: umin, + max: umax, + } + } + + /// Create a new Range object. + /// + /// Note: panic if encounters invalid parameters. + pub fn with_size(base: T, size: T) -> Self + where + u64: From, + { + let umin = u64::from(base); + let umax = u64::from(size).checked_add(umin).unwrap(); + if umin > umax || (umin == 0 && umax == std::u64::MAX) { + panic!("interval_tree: Range({}, {}) is invalid", umin, umax); + } + Range { + min: umin, + max: umax, + } + } + + /// Create a new Range object containing just a point. + pub fn new_point(value: T) -> Self + where + u64: From, + { + let val = u64::from(value); + Range { min: val, max: val } + } + + /// Get length of the range. + pub fn len(&self) -> u64 { + self.max - self.min + 1 + } + + /// Check whether the range is empty. + pub fn is_empty(&self) -> bool { + false + } + + /// Check whether two Range objects intersect with each other. + pub fn intersect(&self, other: &Range) -> bool { + max(self.min, other.min) <= min(self.max, other.max) + } + + /// Check whether the key is fully covered. + pub fn contain(&self, other: &Range) -> bool { + self.min <= other.min && self.max >= other.max + } + + /// Create a new Range object with min aligned to `align`. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::Range; + /// + /// let a = Range::new(2u32, 6u32); + /// assert_eq!(a.align_to(0), Some(Range::new(2u32, 6u32))); + /// assert_eq!(a.align_to(1), Some(Range::new(2u16, 6u16))); + /// assert_eq!(a.align_to(2), Some(Range::new(2u64, 6u64))); + /// assert_eq!(a.align_to(4), Some(Range::new(4u8, 6u8))); + /// assert_eq!(a.align_to(8), None); + /// assert_eq!(a.align_to(3), None); + /// let b = Range::new(2u8, 2u8); + /// assert_eq!(b.align_to(2), Some(Range::new(2u8, 2u8))); + /// ``` + pub fn align_to(&self, align: u64) -> Option { + match align { + 0 | 1 => Some(*self), + _ => { + if align & (align - 1) != 0 { + return None; + } + if let Some(min) = self.min.checked_add(align - 1).map(|v| v & !(align - 1)) { + if min <= self.max { + return Some(Range::new(min, self.max)); + } + } + None + } + } + } +} + +impl Ord for Range { + fn cmp(&self, other: &Self) -> Ordering { + match self.min.cmp(&other.min) { + Ordering::Equal => self.max.cmp(&other.max), + res => res, + } + } +} + +/// Node state for interval tree nodes. +/// +/// Valid state transition: +/// - None -> Free: IntervalTree::insert() +/// - None -> Valued: IntervalTree::insert() +/// - Free -> Allocated: IntervalTree::allocate() +/// - Allocated -> Valued(T): IntervalTree::update() +/// - Valued -> Valued(T): IntervalTree::update() +/// - Allocated -> Free: IntervalTree::free() +/// - Valued(T) -> Free: IntervalTree::free() +/// - * -> None: IntervalTree::delete() +#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord)] +pub enum NodeState { + /// Node is free + Free, + /// Node is allocated but without associated data + Allocated, + /// Node is allocated with associated data. + Valued(T), +} + +impl NodeState { + fn take(&mut self) -> Self { + std::mem::replace(self, NodeState::::Free) + } + + fn replace(&mut self, value: NodeState) -> Self { + std::mem::replace(self, value) + } + + fn as_ref(&self) -> NodeState<&T> { + match self { + NodeState::::Valued(ref x) => NodeState::<&T>::Valued(x), + NodeState::::Allocated => NodeState::<&T>::Allocated, + NodeState::::Free => NodeState::<&T>::Free, + } + } + + fn as_mut(&mut self) -> NodeState<&mut T> { + match self { + NodeState::::Valued(ref mut x) => NodeState::<&mut T>::Valued(x), + NodeState::::Allocated => NodeState::<&mut T>::Allocated, + NodeState::::Free => NodeState::<&mut T>::Free, + } + } + + fn is_free(&self) -> bool { + if let NodeState::::Free = self { + true + } else { + false + } + } +} + +impl Into> for NodeState { + fn into(self) -> Option { + match self { + NodeState::::Free | NodeState::::Allocated => None, + NodeState::::Valued(data) => Some(data), + } + } +} + +/// Internal tree node to implement interval tree. +#[derive(Debug)] +struct InnerNode { + /// Interval handled by this node. + key: Range, + /// Optional contained data, None if the node is free. + data: NodeState, + /// Optional left child of current node. + left: Option>, + /// Optional right child of current node. + right: Option>, + /// Cached height of the node. + height: u32, + /// Cached maximum valued covered by this node. + max_key: u64, +} + +impl InnerNode { + fn new(key: Range, data: NodeState) -> Self { + InnerNode { + key, + data, + left: None, + right: None, + height: 1, + max_key: key.max, + } + } +} + +/// Newtype for interval tree nodes. +#[derive(Debug)] +struct Node(Box>); + +impl Node { + fn new(key: Range, data: Option) -> Self { + let value = if let Some(t) = data { + NodeState::Valued(t) + } else { + NodeState::Free + }; + Node(Box::new(InnerNode::new(key, value))) + } + + /// Returns a readonly reference to the node associated with the `key` or None if not found. + fn search(&self, key: &Range) -> Option<&Self> { + match self.0.key.cmp(key) { + Ordering::Equal => Some(self), + Ordering::Less => self.0.right.as_ref().and_then(|node| node.search(key)), + Ordering::Greater => self.0.left.as_ref().and_then(|node| node.search(key)), + } + } + + /// Returns a shared reference to the node covers full range of the `key`. + fn search_superset(&self, key: &Range) -> Option<&Self> { + if self.0.key.contain(key) { + Some(self) + } else if key.max < self.0.key.min && self.0.left.is_some() { + // Safe to unwrap() because we have just checked it. + self.0.left.as_ref().unwrap().search_superset(key) + } else if key.min > self.0.key.max && self.0.right.is_some() { + // Safe to unwrap() because we have just checked it. + self.0.right.as_ref().unwrap().search_superset(key) + } else { + None + } + } + + /// Returns a mutable reference to the node covers full range of the `key`. + fn search_superset_mut(&mut self, key: &Range) -> Option<&mut Self> { + if self.0.key.contain(key) { + Some(self) + } else if key.max < self.0.key.min && self.0.left.is_some() { + // Safe to unwrap() because we have just checked it. + self.0.left.as_mut().unwrap().search_superset_mut(key) + } else if key.min > self.0.key.max && self.0.right.is_some() { + // Safe to unwrap() because we have just checked it. + self.0.right.as_mut().unwrap().search_superset_mut(key) + } else { + None + } + } + + /// Insert a new (key, data) pair into the subtree. + /// + /// Note: it will panic if the new key intersects with existing nodes. + fn insert(mut self, key: Range, data: Option) -> Self { + match self.0.key.cmp(&key) { + Ordering::Equal => { + panic!("interval_tree: key {:?} exists", key); + } + Ordering::Less => { + if self.0.key.intersect(&key) { + panic!( + "interval_tree: key {:?} intersects with existing {:?}", + key, self.0.key + ); + } + match self.0.right { + None => self.0.right = Some(Node::new(key, data)), + Some(_) => self.0.right = self.0.right.take().map(|n| n.insert(key, data)), + } + } + Ordering::Greater => { + if self.0.key.intersect(&key) { + panic!( + "interval_tree: key {:?} intersects with existing {:?}", + key, self.0.key + ); + } + match self.0.left { + None => self.0.left = Some(Node::new(key, data)), + Some(_) => self.0.left = self.0.left.take().map(|n| n.insert(key, data)), + } + } + } + self.updated_node() + } + + /// Update an existing entry and return the old value. + fn update(&mut self, key: &Range, data: NodeState) -> Option { + match self.0.key.cmp(&key) { + Ordering::Equal => { + match (self.0.data.as_ref(), data.as_ref()) { + (NodeState::<&T>::Free, NodeState::<&T>::Free) + | (NodeState::<&T>::Free, NodeState::<&T>::Valued(_)) + | (NodeState::<&T>::Allocated, NodeState::<&T>::Free) + | (NodeState::<&T>::Allocated, NodeState::<&T>::Allocated) + | (NodeState::<&T>::Valued(_), NodeState::<&T>::Free) + | (NodeState::<&T>::Valued(_), NodeState::<&T>::Allocated) => { + panic!("try to update unallocated interval tree node"); + } + _ => {} + } + self.0.data.replace(data).into() + } + Ordering::Less => match self.0.right.as_mut() { + None => None, + Some(node) => node.update(key, data), + }, + Ordering::Greater => match self.0.left.as_mut() { + None => None, + Some(node) => node.update(key, data), + }, + } + } + + /// Delete `key` from the subtree. + /// + /// Note: it doesn't return whether the key exists in the subtree, so caller need to ensure the + /// logic. + fn delete(mut self, key: &Range) -> (Option, Option) { + match self.0.key.cmp(&key) { + Ordering::Equal => { + let data = self.0.data.take(); + return (data.into(), self.delete_root()); + } + Ordering::Less => { + if let Some(node) = self.0.right.take() { + let (data, right) = node.delete(key); + self.0.right = right; + return (data, Some(self.updated_node())); + } + } + Ordering::Greater => { + if let Some(node) = self.0.left.take() { + let (data, left) = node.delete(key); + self.0.left = left; + return (data, Some(self.updated_node())); + } + } + } + (None, Some(self)) + } + + /// Rotate the node if necessary to keep balance. + fn rotate(self) -> Self { + let l = height(&self.0.left); + let r = height(&self.0.right); + match (l as i32) - (r as i32) { + 1 | 0 | -1 => self, + 2 => self.rotate_left_successor(), + -2 => self.rotate_right_successor(), + _ => unreachable!(), + } + } + + /// Perform a single left rotation on this node. + fn rotate_left(mut self) -> Self { + let mut new_root = self.0.right.take().expect("Node is broken"); + self.0.right = new_root.0.left.take(); + self.update_cached_info(); + new_root.0.left = Some(self); + new_root.update_cached_info(); + new_root + } + + /// Perform a single right rotation on this node. + fn rotate_right(mut self) -> Self { + let mut new_root = self.0.left.take().expect("Node is broken"); + self.0.left = new_root.0.right.take(); + self.update_cached_info(); + new_root.0.right = Some(self); + new_root.update_cached_info(); + new_root + } + + /// Performs a rotation when the left successor is too high. + fn rotate_left_successor(mut self) -> Self { + let left = self.0.left.take().expect("Node is broken"); + if height(&left.0.left) < height(&left.0.right) { + let rotated = left.rotate_left(); + self.0.left = Some(rotated); + self.update_cached_info(); + } else { + self.0.left = Some(left); + } + self.rotate_right() + } + + /// Performs a rotation when the right successor is too high. + fn rotate_right_successor(mut self) -> Self { + let right = self.0.right.take().expect("Node is broken"); + if height(&right.0.left) > height(&right.0.right) { + let rotated = right.rotate_right(); + self.0.right = Some(rotated); + self.update_cached_info(); + } else { + self.0.right = Some(right); + } + self.rotate_left() + } + + fn delete_root(mut self) -> Option { + match (self.0.left.take(), self.0.right.take()) { + (None, None) => None, + (Some(l), None) => Some(l), + (None, Some(r)) => Some(r), + (Some(l), Some(r)) => Some(Self::combine_subtrees(l, r)), + } + } + + /// Find the minimal key below the tree and returns a new optional tree where the minimal + /// value has been removed and the (optional) minimal node as tuple (min_node, remaining) + fn get_new_root(mut self) -> (Self, Option) { + match self.0.left.take() { + None => { + let remaining = self.0.right.take(); + (self, remaining) + } + Some(left) => { + let (min_node, left) = left.get_new_root(); + self.0.left = left; + (min_node, Some(self.updated_node())) + } + } + } + + fn combine_subtrees(l: Self, r: Self) -> Self { + let (mut new_root, remaining) = r.get_new_root(); + new_root.0.left = Some(l); + new_root.0.right = remaining; + new_root.updated_node() + } + + fn find_candidate(&self, constraint: &Constraint) -> Option<&Self> { + match constraint.policy { + AllocPolicy::FirstMatch => self.first_match(constraint), + AllocPolicy::Default => self.first_match(constraint), + } + } + + fn first_match(&self, constraint: &Constraint) -> Option<&Self> { + let mut candidate = if self.0.left.is_some() { + self.0.left.as_ref().unwrap().first_match(constraint) + } else { + None + }; + + if candidate.is_none() && self.check_constraint(constraint) { + candidate = Some(self); + } + if candidate.is_none() && self.0.right.is_some() { + candidate = self.0.right.as_ref().unwrap().first_match(constraint); + } + candidate + } + + fn check_constraint(&self, constraint: &Constraint) -> bool { + if self.0.data.is_free() { + let min = std::cmp::max(self.0.key.min, constraint.min); + let max = std::cmp::min(self.0.key.max, constraint.max); + if min <= max { + let key = Range::new(min, max); + if constraint.align == 0 || constraint.align == 1 { + return key.len() >= constraint.size; + } + return match key.align_to(constraint.align) { + None => false, + Some(aligned_key) => aligned_key.len() >= constraint.size, + }; + } + } + false + } + + /// Update cached information of the node. + /// Please make sure that the cached values of both children are up to date. + fn update_cached_info(&mut self) { + self.0.height = max(height(&self.0.left), height(&self.0.right)) + 1; + self.0.max_key = max( + max_key(&self.0.left), + max(max_key(&self.0.right), self.0.key.max), + ); + } + + /// Update the sub-tree to keep balance. + fn updated_node(mut self) -> Self { + self.update_cached_info(); + self.rotate() + } +} + +/// Compute height of the optional sub-tree. +fn height(node: &Option>) -> u32 { + node.as_ref().map_or(0, |n| n.0.height) +} + +/// Compute maximum key value covered by the optional sub-tree. +fn max_key(node: &Option>) -> u64 { + node.as_ref().map_or(0, |n| n.0.max_key) +} + +/// An interval tree implementation specialized for VMM resource management. +#[derive(Debug, Default)] +pub struct IntervalTree { + root: Option>, +} + +impl IntervalTree { + /// Construct a new empty IntervalTree. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// + /// let tree = vm_allocator::IntervalTree::::new(); + /// ``` + pub fn new() -> Self { + IntervalTree { root: None } + } + + /// Check whether the interval tree is empty. + pub fn is_empty(&self) -> bool { + self.root.is_none() + } + + /// Get the data item associated with the key, or return None if no match found. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = vm_allocator::IntervalTree::::new(); + /// assert!(tree.is_empty()); + /// assert_eq!(tree.get(&Range::new(0x101u64, 0x101u64)), None); + /// tree.insert(Range::new(0x100u64, 0x100u64), Some(1)); + /// tree.insert(Range::new(0x200u64, 0x2ffu64), None); + /// assert!(!tree.is_empty()); + /// assert_eq!(tree.get(&Range::new(0x100u64, 0x100u64)), Some(NodeState::Valued(&1))); + /// assert_eq!(tree.get(&Range::new(0x200u64, 0x2ffu64)), Some(NodeState::Free)); + /// assert_eq!(tree.get(&Range::new(0x101u64, 0x101u64)), None); + /// assert_eq!(tree.get(&Range::new(0x100u64, 0x101u64)), None); + /// ``` + pub fn get(&self, key: &Range) -> Option> { + match self.root { + None => None, + Some(ref node) => node.search(key).map(|n| n.0.data.as_ref()), + } + } + + /// Get a shared reference to the node fully covering the entire key range. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1)); + /// tree.insert(Range::new(0x200u32, 0x2ffu32), None); + /// assert_eq!(tree.get_superset(&Range::new(0x100u32, 0x100u32)), + /// Some((&Range::new(0x100u32, 0x100u32), NodeState::Valued(&1)))); + /// assert_eq!(tree.get_superset(&Range::new(0x210u32, 0x210u32)), + /// Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))); + /// assert_eq!(tree.get_superset(&Range::new(0x2ffu32, 0x2ffu32)), + /// Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))); + /// ``` + pub fn get_superset(&self, key: &Range) -> Option<(&Range, NodeState<&T>)> { + match self.root { + None => None, + Some(ref node) => node + .search_superset(key) + .map(|n| (&n.0.key, n.0.data.as_ref())), + } + } + + /// Get a mutable reference to the node fully covering the entire key range. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1)); + /// tree.insert(Range::new(0x200u32, 0x2ffu32), None); + /// assert_eq!(tree.get_superset_mut(&Range::new(0x100u32, 0x100u32)), + /// Some((&Range::new(0x100u32, 0x100u32), NodeState::Valued(&mut 1)))); + /// assert_eq!(tree.get_superset_mut(&Range::new(0x210u32, 0x210u32)), + /// Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))); + /// assert_eq!(tree.get_superset_mut(&Range::new(0x2ffu32, 0x2ffu32)), + /// Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))); + /// ``` + pub fn get_superset_mut(&mut self, key: &Range) -> Option<(&Range, NodeState<&mut T>)> { + match self.root { + None => None, + Some(ref mut node) => node + .search_superset_mut(key) + .map(|n| (&n.0.key, n.0.data.as_mut())), + } + } + + /// Get a shared reference to the value associated with the id. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u16, 0x100u16), Some(1)); + /// tree.insert(Range::new(0x200u16, 0x2ffu16), None); + /// assert_eq!(tree.get_by_id(0x100u16), Some(&1)); + /// assert_eq!(tree.get_by_id(0x210u32), None); + /// assert_eq!(tree.get_by_id(0x2ffu64), None); + /// ``` + pub fn get_by_id(&self, id: U) -> Option<&T> + where + u64: From, + { + match self.root { + None => None, + Some(ref node) => { + let key = Range::new_point(id); + match node.search_superset(&key) { + Some(node) => node.0.data.as_ref().into(), + None => None, + } + } + } + } + + /// Get a mutable reference to the value associated with the id. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u16, 0x100u16), Some(1)); + /// tree.insert(Range::new(0x200u16, 0x2ffu16), None); + /// assert_eq!(tree.get_by_id_mut(0x100u16), Some(&mut 1)); + /// assert_eq!(tree.get_by_id_mut(0x210u32), None); + /// assert_eq!(tree.get_by_id_mut(0x2ffu64), None); + /// ``` + pub fn get_by_id_mut(&mut self, id: U) -> Option<&mut T> + where + u64: From, + { + match self.root { + None => None, + Some(ref mut node) => { + let key = Range::new_point(id); + match node.search_superset_mut(&key) { + Some(node) => node.0.data.as_mut().into(), + None => None, + } + } + } + } + + /// Insert the (key, data) pair into the interval tree, panic if intersects with existing nodes. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, NodeState}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1)); + /// tree.insert(Range::new(0x200u32, 0x2ffu32), None); + /// assert_eq!(tree.get(&Range::new(0x100u64, 0x100u64)), Some(NodeState::Valued(&1))); + /// assert_eq!(tree.get(&Range::new(0x200u64, 0x2ffu64)), Some(NodeState::Free)); + /// ``` + pub fn insert(&mut self, key: Range, data: Option) { + match self.root.take() { + None => self.root = Some(Node::new(key, data)), + Some(node) => self.root = Some(node.insert(key, data)), + } + } + + /// Update an existing entry and return the old value. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, Constraint}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u64, 0x100u64), None); + /// tree.insert(Range::new(0x200u64, 0x2ffu64), None); + /// + /// let constraint = Constraint::new(2u32); + /// let key = tree.allocate(&constraint); + /// assert_eq!(key, Some(Range::new(0x200u64, 0x201u64))); + /// let old = tree.update(&Range::new(0x200u64, 0x201u64), 2); + /// assert_eq!(old, None); + /// let old = tree.update(&Range::new(0x200u64, 0x201u64), 3); + /// assert_eq!(old, Some(2)); + /// ``` + pub fn update(&mut self, key: &Range, data: T) -> Option { + match self.root.as_mut() { + None => None, + Some(node) => node.update(key, NodeState::::Valued(data)), + } + } + + /// Remove the `key` from the tree and return the associated data. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u64, 0x100u64), Some(1)); + /// tree.insert(Range::new(0x200u64, 0x2ffu64), None); + /// let old = tree.delete(&Range::new(0x100u64, 0x100u64)); + /// assert_eq!(old, Some(1)); + /// let old = tree.delete(&Range::new(0x200u64, 0x2ffu64)); + /// assert_eq!(old, None); + /// ``` + pub fn delete(&mut self, key: &Range) -> Option { + match self.root.take() { + Some(node) => { + let (data, root) = node.delete(key); + self.root = root; + data + } + None => None, + } + } + + /// Allocate a resource range according the allocation constraints. + /// + /// # Examples + /// ```rust + /// extern crate vm_allocator; + /// use vm_allocator::{IntervalTree, Range, Constraint}; + /// + /// let mut tree = IntervalTree::::new(); + /// tree.insert(Range::new(0x100u64, 0x100u64), None); + /// tree.insert(Range::new(0x200u64, 0x2ffu64), None); + /// + /// let constraint = Constraint::new(2u8); + /// let key = tree.allocate(&constraint); + /// assert_eq!(key, Some(Range::new(0x200u64, 0x201u64))); + /// tree.update(&Range::new(0x200u64, 0x201u64), 2); + /// ``` + pub fn allocate(&mut self, constraint: &Constraint) -> Option { + if constraint.size == 0 { + return None; + } + let candidate = match self.root.as_mut() { + None => None, + Some(node) => node.find_candidate(constraint), + }; + + match candidate { + None => None, + Some(node) => { + let node_key = node.0.key; + let range = Range::new( + max(node_key.min, constraint.min), + min(node_key.max, constraint.max), + ); + // Safe to unwrap because candidate satisfy the constraints. + let aligned_key = range.align_to(constraint.align).unwrap(); + let result = Range::new(aligned_key.min, aligned_key.min + constraint.size - 1); + + // Allocate a resource from the node, no need to split the candidate node. + if node_key.min == aligned_key.min && node_key.len() == constraint.size { + self.root + .as_mut() + .unwrap() + .update(&node_key, NodeState::::Allocated); + return Some(node_key); + } + + // Split the candidate node. + // TODO: following algorithm is not optimal in preference of simplicity. + self.delete(&node_key); + if aligned_key.min > node_key.min { + self.insert(Range::new(node_key.min, aligned_key.min - 1), None); + } + self.insert(result, None); + if result.max < node_key.max { + self.insert(Range::new(result.max + 1, node_key.max), None); + } + + self.root + .as_mut() + .unwrap() + .update(&result, NodeState::::Allocated); + Some(result) + } + } + } + + /// Free an allocated range and return the associated data. + pub fn free(&mut self, key: &Range) -> Option { + let result = self.delete(key); + let mut range = *key; + + // Try to merge with adjacent free nodes. + if range.min > 0 { + if let Some((r, v)) = self.get_superset(&Range::new(range.min - 1, range.min - 1)) { + if v.is_free() { + range.min = r.min; + } + } + } + if range.max < std::u64::MAX { + if let Some((r, v)) = self.get_superset(&Range::new(range.max + 1, range.max + 1)) { + if v.is_free() { + range.max = r.max; + } + } + } + + if range.min < key.min { + self.delete(&Range::new(range.min, key.min - 1)); + } + if range.max > key.max { + self.delete(&Range::new(key.max + 1, range.max)); + } + self.insert(range, None); + + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[should_panic] + fn test_new_range() { + let _ = Range::new(2u8, 1u8); + } + + #[test] + #[should_panic] + fn test_new_range_overflow() { + let _ = Range::new(0u64, std::u64::MAX); + } + + #[test] + fn test_range_intersect() { + let a = Range::new(1u8, 4u8); + let b = Range::new(4u16, 6u16); + let c = Range::new(2u32, 3u32); + let d = Range::new(4u64, 4u64); + let e = Range::new(5u32, 6u32); + + assert!(a.intersect(&b)); + assert!(b.intersect(&a)); + assert!(a.intersect(&c)); + assert!(c.intersect(&a)); + assert!(a.intersect(&d)); + assert!(d.intersect(&a)); + assert!(!a.intersect(&e)); + assert!(!e.intersect(&a)); + + assert_eq!(a.len(), 4); + assert_eq!(d.len(), 1); + } + + #[test] + fn test_range_contain() { + let a = Range::new(2u8, 6u8); + assert!(a.contain(&Range::new(2u8, 3u8))); + assert!(a.contain(&Range::new(3u8, 4u8))); + assert!(a.contain(&Range::new(5u8, 5u8))); + assert!(a.contain(&Range::new(5u8, 6u8))); + assert!(a.contain(&Range::new(6u8, 6u8))); + assert!(!a.contain(&Range::new(1u8, 1u8))); + assert!(!a.contain(&Range::new(1u8, 2u8))); + assert!(!a.contain(&Range::new(1u8, 3u8))); + assert!(!a.contain(&Range::new(1u8, 7u8))); + assert!(!a.contain(&Range::new(7u8, 8u8))); + assert!(!a.contain(&Range::new(6u8, 7u8))); + assert!(!a.contain(&Range::new(7u8, 8u8))); + } + + #[test] + fn test_range_align_to() { + let a = Range::new(2u32, 6); + assert_eq!(a.align_to(0), Some(Range::new(2u64, 6u64))); + assert_eq!(a.align_to(1), Some(Range::new(2u8, 6u8))); + assert_eq!(a.align_to(2), Some(Range::new(2u16, 6u16))); + assert_eq!(a.align_to(4), Some(Range::new(4u32, 6u32))); + assert_eq!(a.align_to(8), None); + assert_eq!(a.align_to(3), None); + + let a = Range::new(0xFFFF_FFFF_FFFF_FFFDu64, 0xFFFF_FFFF_FFFF_FFFFu64); + assert_eq!( + a.align_to(2), + Some(Range::new(0xFFFF_FFFF_FFFF_FFFEu64, 0xFFFF_FFFF_FFFF_FFFF)) + ); + assert_eq!(a.align_to(4), None); + } + + #[test] + fn test_range_ord() { + let a = Range::new(1u32, 4u32); + let b = Range::new(1u32, 4u32); + let c = Range::new(1u32, 3u32); + let d = Range::new(1u32, 5u32); + let e = Range::new(2u32, 2u32); + + assert_eq!(a, b); + assert_eq!(b, a); + assert!(a > c); + assert!(c < a); + assert!(a < d); + assert!(d > a); + assert!(a < e); + assert!(e > a); + } + + #[should_panic] + #[test] + fn test_tree_insert_equal() { + let mut tree = IntervalTree::::new(); + tree.insert(Range::new(0x100u16, 0x200), Some(1)); + tree.insert(Range::new(0x100u32, 0x200), None); + } + + #[should_panic] + #[test] + fn test_tree_insert_intersect() { + let mut tree = IntervalTree::::new(); + tree.insert(Range::new(0x100, 0x200u32), Some(1)); + tree.insert(Range::new(0x200, 0x2ffu64), None); + } + + #[test] + fn test_tree_get_superset() { + let mut tree = IntervalTree::::new(); + tree.insert(Range::new(0x100u64, 0x100u64), Some(1)); + tree.insert(Range::new(0x200, 0x2ffu64), None); + assert_eq!( + tree.get_superset(&Range::new(0x100u32, 0x100)), + Some((&Range::new(0x100, 0x100u32), NodeState::Valued(&1))) + ); + assert_eq!( + tree.get_superset(&Range::new(0x200u16, 0x200)), + Some((&Range::new(0x200, 0x2ffu64), NodeState::Free)) + ); + assert_eq!( + tree.get_superset(&Range::new(0x200u32, 0x2ff)), + Some((&Range::new(0x200, 0x2ffu16), NodeState::Free)) + ); + assert_eq!( + tree.get_superset(&Range::new(0x210u32, 0x210)), + Some((&Range::new(0x200, 0x2ffu32), NodeState::Free)) + ); + assert_eq!( + tree.get_superset(&Range::new(0x2ffu32, 0x2ff)), + Some((&Range::new(0x200, 0x2ffu32), NodeState::Free)) + ); + assert_eq!(tree.get_superset(&Range::new(0x2ffu32, 0x300)), None); + assert_eq!(tree.get_superset(&Range::new(0x300u32, 0x300)), None); + assert_eq!(tree.get_superset(&Range::new(0x1ffu32, 0x300)), None); + } + + #[test] + fn test_tree_update() { + let mut tree = IntervalTree::::new(); + tree.insert(Range::new(0x100u32, 0x100u32), None); + tree.insert(Range::new(0x200u32, 0x2ffu32), None); + + let constraint = Constraint::new(2u32); + let key = tree.allocate(&constraint); + assert_eq!(key, Some(Range::new(0x200u32, 0x201u32))); + let old = tree.update(&Range::new(0x200u32, 0x201u32), 2); + assert_eq!(old, None); + let old = tree.update(&Range::new(0x200u32, 0x201u32), 3); + assert_eq!(old, Some(2)); + let old = tree.update(&Range::new(0x200u32, 0x200u32), 4); + assert_eq!(old, None); + let old = tree.update(&Range::new(0x200u32, 0x203u32), 5); + assert_eq!(old, None); + + tree.delete(&Range::new(0x200u32, 0x201u32)); + let old = tree.update(&Range::new(0x200u32, 0x201u32), 2); + assert_eq!(old, None); + } + + #[test] + fn test_tree_delete() { + let mut tree = IntervalTree::::new(); + assert_eq!(tree.get(&Range::new(0x101u32, 0x101u32)), None); + assert!(tree.is_empty()); + tree.insert(Range::new(0x100u32, 0x100u32), Some(1)); + tree.insert(Range::new(0x200u32, 0x2ffu32), None); + assert!(!tree.is_empty()); + assert_eq!( + tree.get(&Range::new(0x100u32, 0x100u32)), + Some(NodeState::Valued(&1)) + ); + assert_eq!( + tree.get(&Range::new(0x200u32, 0x2ffu32)), + Some(NodeState::Free) + ); + assert_eq!(tree.get(&Range::new(0x101u32, 0x101u32)), None); + + let old = tree.delete(&Range::new(0x100u32, 0x100u32)); + assert_eq!(old, Some(1)); + let old = tree.delete(&Range::new(0x200u32, 0x2ffu32)); + assert_eq!(old, None); + + assert!(tree.is_empty()); + assert_eq!(tree.get(&Range::new(0x100u32, 0x100u32)), None); + assert_eq!(tree.get(&Range::new(0x200u32, 0x2ffu32)), None); + } + + #[test] + fn test_allocate_free() { + let mut tree = IntervalTree::::new(); + let mut constraint = Constraint::new(1u8); + + assert_eq!(tree.allocate(&constraint), None); + tree.insert(Range::new(0x100u16, 0x100u16), None); + tree.insert(Range::new(0x200u16, 0x2ffu16), None); + + let key = tree.allocate(&constraint); + assert_eq!(key, Some(Range::new(0x100u16, 0x100u16))); + let old = tree.update(&Range::new(0x100u16, 0x100u16), 2); + assert_eq!(old, None); + let val = tree.get(&Range::new(0x100u16, 0x100u16)); + assert_eq!(val, Some(NodeState::Valued(&2))); + + constraint.min = 0x100; + constraint.max = 0x100; + assert_eq!(tree.allocate(&constraint), None); + + constraint.min = 0x201; + constraint.max = 0x300; + constraint.align = 0x8; + constraint.size = 0x10; + assert_eq!( + tree.allocate(&constraint), + Some(Range::new(0x208u16, 0x217u16)) + ); + + // Free the node when it's still in 'Allocated' state. + let old = tree.free(&Range::new(0x208u16, 0x217u16)); + assert_eq!(old, None); + + // Reallocate the freed resource. + assert_eq!( + tree.allocate(&constraint), + Some(Range::new(0x208u16, 0x217u16)) + ); + + constraint.size = 0x100; + assert_eq!(tree.allocate(&constraint), None); + + // Verify that allocating a bigger range with smaller allocated range fails. + constraint.min = 0x200; + constraint.max = 0x2ff; + constraint.align = 0x8; + constraint.size = 0x100; + assert_eq!(tree.allocate(&constraint), None); + + // Free the node when it's in 'Valued' state. + tree.update(&Range::new(0x208u16, 0x217u16), 0x10); + assert_eq!(tree.allocate(&constraint), None); + let old = tree.free(&Range::new(0x208u16, 0x217u16)); + assert_eq!(old, Some(0x10)); + + // Reallocate the freed resource, verify that adjacent free nodes have been merged. + assert_eq!( + tree.allocate(&constraint), + Some(Range::new(0x200u32, 0x2ffu32)) + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 845a298..7069964 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,82 @@ +// Copyright (C) 2019 Alibaba Cloud. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Generic algorithms for VMM resource management. + #![deny(missing_docs)] -//! Dummy crate needs high-level documentation. -/// Dummy public function needs documentation. -pub fn it_works() { - assert!(true); + +mod interval_tree; +pub use interval_tree::{IntervalTree, NodeState, Range}; + +/// Policy for resource allocation. +#[derive(Copy, Clone, Debug)] +pub enum AllocPolicy { + /// Default allocation policy. + Default, + /// Allocate from the first matched entry. + FirstMatch, +} + +/// Struct to describe resource allocation constraints. +#[derive(Copy, Clone, Debug)] +pub struct Constraint { + /// Size to allocate. + pub size: u64, + /// Lower boundary for the allocated resource. + pub min: u64, + /// Upper boundary for the allocated resource. + pub max: u64, + /// Alignment for the allocated resource. + pub align: u64, + /// Resource allocation policy. + pub policy: AllocPolicy, +} + +impl Constraint { + /// Create a new constraint object with default settings. + pub fn new(size: T) -> Self + where + u64: From, + { + Constraint { + size: u64::from(size), + min: 0, + max: std::u64::MAX, + align: 1, + policy: AllocPolicy::Default, + } + } + + /// Set the min constraint. + pub fn min(mut self, min: T) -> Self + where + u64: From, + { + self.min = u64::from(min); + self + } + + /// Set the max constraint. + pub fn max(mut self, max: T) -> Self + where + u64: From, + { + self.max = u64::from(max); + self + } + + /// Set the alignment constraint. + pub fn align(mut self, align: T) -> Self + where + u64: From, + { + self.align = u64::from(align); + self + } + + /// Set the allocation policy. + pub fn policy(mut self, policy: AllocPolicy) -> Self { + self.policy = policy; + self + } }