From 1b1127b6f3a254af84113d28c01058bd36fc60d6 Mon Sep 17 00:00:00 2001 From: Simon LUCIDO Date: Fri, 24 Mar 2023 22:00:28 +0100 Subject: [PATCH 1/3] feat: add more memory types Signed-off-by: Simon LUCIDO --- src/address_allocator.rs | 85 +++++++++++++++++--------- src/allocation_engine/interval_tree.rs | 72 ++++++++++++---------- src/allocation_engine/mod.rs | 3 +- src/lib.rs | 6 +- 4 files changed, 100 insertions(+), 66 deletions(-) diff --git a/src/address_allocator.rs b/src/address_allocator.rs index d2d5833..289e169 100644 --- a/src/address_allocator.rs +++ b/src/address_allocator.rs @@ -8,7 +8,7 @@ //! This module implements an allocation strategies for memory slots in an //! address space (for example MMIO and PIO). -use crate::allocation_engine::IntervalTree; +use crate::allocation_engine::{IntervalTree, NodeState}; use crate::{AllocPolicy, Constraint, Error, RangeInclusive, Result}; // Internal representation of AddressAllocator. Contains the managed address @@ -61,9 +61,10 @@ impl AddressAllocator { size: u64, alignment: u64, policy: AllocPolicy, + node_state: NodeState, ) -> Result { let constraint = Constraint::new(size, alignment, policy)?; - self.interval_tree.allocate(constraint) + self.interval_tree.allocate(constraint, node_state) } /// Deletes the specified memory slot or returns `ResourceNotAvailable` if @@ -81,21 +82,26 @@ mod tests { fn test_regression_exact_match_length_check() { let mut pool = AddressAllocator::new(0x0, 0x2000).unwrap(); let res = pool - .allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x1000)) + .allocate( + 0x1000, + 0x1000, + AllocPolicy::ExactMatch(0x1000), + NodeState::Ram, + ) .unwrap(); assert_eq!( - pool.allocate(0x0, 0x1000, AllocPolicy::FirstMatch) + pool.allocate(0x0, 0x1000, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap_err(), Error::InvalidSize(0x0) ); assert_eq!( - pool.allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x3)) + pool.allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x3), NodeState::Ram) .unwrap_err(), Error::UnalignedAddress ); assert_eq!(res, RangeInclusive::new(0x1000, 0x1FFF).unwrap()); let res = pool - .allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x0)) + .allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x0), NodeState::Ram) .unwrap(); assert_eq!(res, RangeInclusive::new(0x0, 0x0FFF).unwrap()); } @@ -120,7 +126,7 @@ mod tests { fn test_allocate_fails_alignment_zero() { let mut pool = AddressAllocator::new(0x1000, 0x10000).unwrap(); assert_eq!( - pool.allocate(0x100, 0, AllocPolicy::FirstMatch) + pool.allocate(0x100, 0, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap_err(), Error::InvalidAlignment ); @@ -130,7 +136,7 @@ mod tests { fn test_allocate_fails_alignment_non_power_of_two() { let mut pool = AddressAllocator::new(0x1000, 0x10000).unwrap(); assert_eq!( - pool.allocate(0x100, 200, AllocPolicy::FirstMatch) + pool.allocate(0x100, 200, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap_err(), Error::InvalidAlignment ); @@ -140,16 +146,17 @@ mod tests { fn test_allocate_fails_not_enough_space() { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); assert_eq!( - pool.allocate(0x800, 0x100, AllocPolicy::LastMatch).unwrap(), + pool.allocate(0x800, 0x100, AllocPolicy::LastMatch, NodeState::Ram) + .unwrap(), RangeInclusive::new(0x1800, 0x1FFF).unwrap() ); assert_eq!( - pool.allocate(0x900, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x900, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap_err(), Error::ResourceNotAvailable ); assert_eq!( - pool.allocate(0x400, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x400, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1000, 0x13FF).unwrap() ); @@ -159,17 +166,18 @@ mod tests { fn test_allocate_with_alignment_first_ok() { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); assert_eq!( - pool.allocate(0x110, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x110, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1000, 0x110F).unwrap() ); assert_eq!( - pool.allocate(0x100, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x100, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1200, 0x12FF).unwrap() ); assert_eq!( - pool.allocate(0x10, 0x100, AllocPolicy::FirstMatch).unwrap(), + pool.allocate(0x10, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) + .unwrap(), RangeInclusive::new(0x1300, 0x130F).unwrap() ); } @@ -179,19 +187,19 @@ mod tests { let mut pool_reverse = AddressAllocator::new(0x1000, 0x10000).unwrap(); assert_eq!( pool_reverse - .allocate(0x110, 0x100, AllocPolicy::LastMatch) + .allocate(0x110, 0x100, AllocPolicy::LastMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x10E00, 0x10F0F).unwrap() ); assert_eq!( pool_reverse - .allocate(0x100, 0x100, AllocPolicy::LastMatch) + .allocate(0x100, 0x100, AllocPolicy::LastMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x10D00, 0x10DFF).unwrap() ); assert_eq!( pool_reverse - .allocate(0x10, 0x100, AllocPolicy::LastMatch) + .allocate(0x10, 0x100, AllocPolicy::LastMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x10C00, 0x10C0F).unwrap() ); @@ -202,27 +210,42 @@ mod tests { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); // First range is [0x1000:0x17FF] assert_eq!( - pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1000, 0x17FF).unwrap() ); // Second range is [0x1A00:0x1BFF] assert_eq!( - pool.allocate(0x200, 0x100, AllocPolicy::ExactMatch(0x1A00)) - .unwrap(), + pool.allocate( + 0x200, + 0x100, + AllocPolicy::ExactMatch(0x1A00), + NodeState::Ram + ) + .unwrap(), RangeInclusive::new(0x1A00, 0x1BFF).unwrap() ); // There is 0x200 between the first 2 ranges. // We ask for an available address but the range is too big assert_eq!( - pool.allocate(0x800, 0x100, AllocPolicy::ExactMatch(0x1800)) - .unwrap_err(), + pool.allocate( + 0x800, + 0x100, + AllocPolicy::ExactMatch(0x1800), + NodeState::Ram + ) + .unwrap_err(), Error::ResourceNotAvailable ); // We ask for an available address, with a small enough range assert_eq!( - pool.allocate(0x100, 0x100, AllocPolicy::ExactMatch(0x1800)) - .unwrap(), + pool.allocate( + 0x100, + 0x100, + AllocPolicy::ExactMatch(0x1800), + NodeState::Ram + ) + .unwrap(), RangeInclusive::new(0x1800, 0x18FF).unwrap() ); } @@ -231,14 +254,14 @@ mod tests { fn test_tree_allocate_address_free_and_realloc() { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); assert_eq!( - pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1000, 0x17FF).unwrap() ); let _ = pool.free(&RangeInclusive::new(0x1000, 0x17FF).unwrap()); assert_eq!( - pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap(), RangeInclusive::new(0x1000, 0x17FF).unwrap() ); @@ -248,11 +271,13 @@ mod tests { fn test_allow_range_size_one_left() { let mut pool = AddressAllocator::new(1, 1000).unwrap(); assert_eq!( - pool.allocate(10, 2, AllocPolicy::FirstMatch).unwrap(), + pool.allocate(10, 2, AllocPolicy::FirstMatch, NodeState::Ram) + .unwrap(), RangeInclusive::new(2, 11).unwrap() ); assert_eq!( - pool.allocate(1, 1, AllocPolicy::FirstMatch).unwrap(), + pool.allocate(1, 1, AllocPolicy::FirstMatch, NodeState::Ram) + .unwrap(), RangeInclusive::new(1, 1).unwrap() ); } @@ -262,7 +287,7 @@ mod tests { let mut pool = AddressAllocator::new(0x0, 0x1000).unwrap(); //First allocation fails assert_eq!( - pool.allocate(0x2000, 0x100, AllocPolicy::FirstMatch) + pool.allocate(0x2000, 0x100, AllocPolicy::FirstMatch, NodeState::Ram) .unwrap_err(), Error::ResourceNotAvailable ); @@ -274,7 +299,7 @@ mod tests { ); // Now we try an allocation that should succeed. assert_eq!( - pool.allocate(0x4FE, 0x100, AllocPolicy::ExactMatch(0x500)) + pool.allocate(0x4FE, 0x100, AllocPolicy::ExactMatch(0x500), NodeState::Ram) .unwrap(), RangeInclusive::new(0x500, 0x9FD).unwrap() ); diff --git a/src/allocation_engine/interval_tree.rs b/src/allocation_engine/interval_tree.rs index 2797afc..8e28464 100644 --- a/src/allocation_engine/interval_tree.rs +++ b/src/allocation_engine/interval_tree.rs @@ -44,8 +44,14 @@ pub fn align_up(address: u64, alignment: u64) -> Result { pub enum NodeState { /// Node is free. Free, - /// Node is allocated. - Allocated, + /// Reserved, and should be allocated (x86's real mode for example). + ReservedAllocated, + /// Reserved, and should not be allocated (x86's upper memory, for example). + ReservedUnallocated, + /// RAM + Ram, + /// MMIO + Mmio, } impl NodeState { @@ -279,22 +285,22 @@ impl InnerNode { /// find an existing node with the state `NodeState::Free` that satisfies /// all constraints of an allocation request. The recursion is safe as we /// have in place a maximum height for the tree. - fn mark_as_allocated(&mut self, key: &RangeInclusive) -> Result<()> { + fn mark_as_allocated(&mut self, key: &RangeInclusive, node_state: NodeState) -> Result<()> { match self.key.cmp(key) { Ordering::Equal => { if self.node_state != NodeState::Free { return Err(Error::InvalidStateTransition(self.key, self.node_state)); } - self.node_state = NodeState::Allocated; + self.node_state = node_state; Ok(()) } Ordering::Less => match self.right.as_mut() { None => Err(Error::ResourceNotAvailable), - Some(node) => node.mark_as_allocated(key), + Some(node) => node.mark_as_allocated(key, node_state), }, Ordering::Greater => match self.left.as_mut() { None => Err(Error::ResourceNotAvailable), - Some(node) => node.mark_as_allocated(key), + Some(node) => node.mark_as_allocated(key, node_state), }, } } @@ -518,10 +524,10 @@ impl IntervalTree { Ok(()) } - fn mark_as_allocated(&mut self, key: &RangeInclusive) -> Result<()> { + fn mark_as_allocated(&mut self, key: &RangeInclusive, node_state: NodeState) -> Result<()> { match self.root.as_mut() { None => (), - Some(node) => node.mark_as_allocated(key)?, + Some(node) => node.mark_as_allocated(key, node_state)?, }; Ok(()) } @@ -543,7 +549,11 @@ impl IntervalTree { /// such that the RangeInclusive representing the desired memory slot will appear as /// an node with the state `NodeState::Allocated` while the leftovers of /// the previous node will be present in the tree as free nodes. - pub fn allocate(&mut self, constraint: Constraint) -> Result { + pub fn allocate( + &mut self, + constraint: Constraint, + node_state: NodeState, + ) -> Result { // Return ResourceNotAvailable if we can not get a reference to the // root node. let root = self.root.as_ref().ok_or(Error::ResourceNotAvailable)?; @@ -562,7 +572,7 @@ impl IntervalTree { // Allocate a resource from the node, no need to split the candidate node. if node_key.start() == result.start() && node_key.len() == constraint.size { - self.mark_as_allocated(&node_key)?; + self.mark_as_allocated(&node_key, node_state)?; return Ok(node_key); } @@ -583,7 +593,7 @@ impl IntervalTree { )?; } - self.insert(result, NodeState::Allocated)?; + self.insert(result, node_state)?; if result.end() < node_key.end() { self.insert( RangeInclusive::new( @@ -675,7 +685,7 @@ mod tests { #[test] fn test_is_free() { - let mut ns = NodeState::Allocated; + let mut ns = NodeState::Ram; assert!(!ns.is_free()); ns = NodeState::Free; assert!(ns.is_free()); @@ -685,7 +695,7 @@ mod tests { fn test_search() { let mut tree = Box::new(InnerNode::new( RangeInclusive::new(0x100, 0x110).unwrap(), - NodeState::Allocated, + NodeState::Ram, )); let left_child = InnerNode::new(RangeInclusive::new(0x90, 0x99).unwrap(), NodeState::Free); @@ -712,7 +722,7 @@ mod tests { fn test_search_superset() { let mut tree = Box::new(InnerNode::new( RangeInclusive::new(0x100, 0x110).unwrap(), - NodeState::Allocated, + NodeState::Ram, )); let right_child = InnerNode::new(RangeInclusive::new(0x200, 0x2FF).unwrap(), NodeState::Free); @@ -780,7 +790,7 @@ mod tests { fn test_tree_insert_balanced() { let mut tree = Box::new(InnerNode::new( RangeInclusive::new(0x300, 0x310).unwrap(), - NodeState::Allocated, + NodeState::Ram, )); tree = tree .insert(RangeInclusive::new(0x100, 0x110).unwrap(), NodeState::Free) @@ -798,7 +808,7 @@ mod tests { assert!(is_balanced(Some(tree))); tree = Box::new(InnerNode::new( RangeInclusive::new(0x300, 0x310).unwrap(), - NodeState::Allocated, + NodeState::Ram, )); tree = tree .insert(RangeInclusive::new(0x100, 0x110).unwrap(), NodeState::Free) @@ -849,7 +859,7 @@ mod tests { fn test_tree_insert_intersect_negative() { let mut tree = Box::new(InnerNode::new( RangeInclusive::new(0x100, 0x200).unwrap(), - NodeState::Allocated, + NodeState::Ram, )); tree = tree .insert(RangeInclusive::new(0x201, 0x2FF).unwrap(), NodeState::Free) @@ -882,17 +892,15 @@ mod tests { #[test] fn test_tree_insert_duplicate_negative() { let range = RangeInclusive::new(0x100, 0x200).unwrap(); - let tree = Box::new(InnerNode::new(range, NodeState::Allocated)); + let tree = Box::new(InnerNode::new(range, NodeState::Ram)); let res = tree.insert(range, NodeState::Free); assert_eq!(res.unwrap_err(), Error::Overlap(range, range)); } #[test] fn test_tree_stack_overflow_negative() { - let mut inner_node = InnerNode::new( - RangeInclusive::new(0x100, 0x200).unwrap(), - NodeState::Allocated, - ); + let mut inner_node = + InnerNode::new(RangeInclusive::new(0x100, 0x200).unwrap(), NodeState::Ram); inner_node.height = 50; let tree = Box::new(inner_node); let res = tree.insert(RangeInclusive::new(0x100, 0x200).unwrap(), NodeState::Free); @@ -902,24 +910,24 @@ mod tests { #[test] fn test_tree_mark_as_allocated_invalid_transition() { let range = RangeInclusive::new(0x100, 0x110).unwrap(); - let mut tree = Box::new(InnerNode::new(range, NodeState::Allocated)); + let mut tree = Box::new(InnerNode::new(range, NodeState::Ram)); assert_eq!( - tree.mark_as_allocated(&range).unwrap_err(), - Error::InvalidStateTransition(range, NodeState::Allocated) + tree.mark_as_allocated(&range, tree.node_state).unwrap_err(), + Error::InvalidStateTransition(range, NodeState::Ram) ); } #[test] fn test_tree_mark_as_allocated_resource_not_available() { let range = RangeInclusive::new(0x100, 0x110).unwrap(); - let mut tree = Box::new(InnerNode::new(range, NodeState::Allocated)); + let mut tree = Box::new(InnerNode::new(range, NodeState::Ram)); assert_eq!( - tree.mark_as_allocated(&RangeInclusive::new(0x111, 0x112).unwrap()) + tree.mark_as_allocated(&RangeInclusive::new(0x111, 0x112).unwrap(), tree.node_state) .unwrap_err(), Error::ResourceNotAvailable ); assert_eq!( - tree.mark_as_allocated(&RangeInclusive::new(0x90, 0x92).unwrap()) + tree.mark_as_allocated(&RangeInclusive::new(0x90, 0x92).unwrap(), tree.node_state) .unwrap_err(), Error::ResourceNotAvailable ); @@ -929,12 +937,12 @@ mod tests { fn test_tree_mark_as_allocated() { let range = RangeInclusive::new(0x100, 0x110).unwrap(); let range2 = RangeInclusive::new(0x200, 0x2FF).unwrap(); - let mut tree = Box::new(InnerNode::new(range, NodeState::Allocated)); + let mut tree = Box::new(InnerNode::new(range, NodeState::Ram)); tree = tree.insert(range2, NodeState::Free).unwrap(); - assert!(tree.mark_as_allocated(&range2).is_ok()); + assert!(tree.mark_as_allocated(&range2, tree.node_state).is_ok()); assert_eq!( *tree.search(&range2).unwrap(), - InnerNode::new(range2, NodeState::Allocated) + InnerNode::new(range2, NodeState::Ram) ); } @@ -996,7 +1004,7 @@ mod tests { AllocPolicy::ExactMatch(0x8000000000000000), ) .unwrap(); - let res = tree.allocate(constraint); + let res = tree.allocate(constraint, NodeState::Ram); assert_eq!(res.unwrap_err(), Error::Overflow); } } diff --git a/src/allocation_engine/mod.rs b/src/allocation_engine/mod.rs index 43b2c64..575956d 100644 --- a/src/allocation_engine/mod.rs +++ b/src/allocation_engine/mod.rs @@ -4,4 +4,5 @@ mod interval_tree; -pub(crate) use interval_tree::{IntervalTree, NodeState}; +pub(crate) use interval_tree::IntervalTree; +pub use interval_tree::NodeState; diff --git a/src/lib.rs b/src/lib.rs index 4127872..ff43deb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,7 +15,7 @@ //! ```rust //! use std::collections::HashMap; //! use std::process::id; -//! use vm_allocator::{AddressAllocator, AllocPolicy, Error, IdAllocator, RangeInclusive, Result}; +//! use vm_allocator::{AddressAllocator, AllocPolicy, NodeState, Error, IdAllocator, RangeInclusive, Result}; //! //! const FIRST_ADDR_PAST_32BITS: u64 = 1 << 32; //! const MEM_32BIT_GAP_SIZE: u64 = 768 << 20; @@ -48,7 +48,7 @@ //! // For each device we reserve one page of addresses. //! let mmio_range = //! self.mmio_allocator -//! .allocate(PAGE_SIZE, PAGE_SIZE, AllocPolicy::FirstMatch)?; +//! .allocate(PAGE_SIZE, PAGE_SIZE, AllocPolicy::FirstMatch, NodeState::Mmio)?; //! let slot = DeviceSlot { //! id: self.id_allocator.allocate_id()?, //! mmio_range, @@ -82,7 +82,7 @@ mod id_allocator; use std::{cmp::max, cmp::min, result}; use thiserror::Error; -use crate::allocation_engine::NodeState; +pub use crate::allocation_engine::NodeState; pub use crate::{address_allocator::AddressAllocator, id_allocator::IdAllocator}; /// Default alignment that can be used for creating a `Constraint`. From 5555cfc26b1c61c5eab07ce27cb90f5f07b8c085 Mon Sep 17 00:00:00 2001 From: Simon LUCIDO Date: Fri, 24 Mar 2023 22:55:09 +0100 Subject: [PATCH 2/3] feat: add method to get all flatten nodes Signed-off-by: Simon LUCIDO --- src/address_allocator.rs | 7 ++++++- src/allocation_engine/interval_tree.rs | 23 ++++++++++++++++++++++- src/allocation_engine/mod.rs | 1 + 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/src/address_allocator.rs b/src/address_allocator.rs index 289e169..92ca21e 100644 --- a/src/address_allocator.rs +++ b/src/address_allocator.rs @@ -8,7 +8,7 @@ //! This module implements an allocation strategies for memory slots in an //! address space (for example MMIO and PIO). -use crate::allocation_engine::{IntervalTree, NodeState}; +use crate::allocation_engine::{InnerNode, IntervalTree, NodeState}; use crate::{AllocPolicy, Constraint, Error, RangeInclusive, Result}; // Internal representation of AddressAllocator. Contains the managed address @@ -72,6 +72,11 @@ impl AddressAllocator { pub fn free(&mut self, key: &RangeInclusive) -> Result<()> { self.interval_tree.free(key) } + + /// Get all the allocated memory slots. + pub fn allocated_slots(&self) -> Vec<&InnerNode> { + self.interval_tree.allocated_slots() + } } #[cfg(test)] diff --git a/src/allocation_engine/interval_tree.rs b/src/allocation_engine/interval_tree.rs index 8e28464..295c09b 100644 --- a/src/allocation_engine/interval_tree.rs +++ b/src/allocation_engine/interval_tree.rs @@ -62,7 +62,7 @@ impl NodeState { /// Internal tree node to implement interval tree. #[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord)] -pub(crate) struct InnerNode { +pub struct InnerNode { /// Interval handled by this node. key: RangeInclusive, /// NodeState, can be Free or Allocated. @@ -87,6 +87,21 @@ impl InnerNode { } } + /// Returns all the allocated nodes in the tree. + fn allocated_nodes(&self) -> Vec<&InnerNode> { + let mut allocated_nodes = Vec::new(); + if self.node_state != NodeState::Free { + allocated_nodes.push(self); + } + if let Some(left) = &self.left { + allocated_nodes.extend(left.allocated_nodes()); + } + if let Some(right) = &self.right { + allocated_nodes.extend(right.allocated_nodes()); + } + allocated_nodes + } + /// Returns a readonly reference to the node associated with the `key` or /// None if the searched key does not exist in the tree. fn search(&self, key: &RangeInclusive) -> Option<&InnerNode> { @@ -659,6 +674,12 @@ impl IntervalTree { self.insert(range, NodeState::Free)?; Ok(()) } + + pub fn allocated_slots(&self) -> Vec<&InnerNode> { + self.root + .as_ref() + .map_or(vec![], |root| root.allocated_nodes()) + } } #[cfg(test)] diff --git a/src/allocation_engine/mod.rs b/src/allocation_engine/mod.rs index 575956d..060c72c 100644 --- a/src/allocation_engine/mod.rs +++ b/src/allocation_engine/mod.rs @@ -4,5 +4,6 @@ mod interval_tree; +pub use interval_tree::InnerNode; pub(crate) use interval_tree::IntervalTree; pub use interval_tree::NodeState; From 2d5c1bf9275fae7fc794f4502e795f516805ea72 Mon Sep 17 00:00:00 2001 From: Simon LUCIDO Date: Fri, 24 Mar 2023 23:00:35 +0100 Subject: [PATCH 3/3] feat: add public getters for inner node Signed-off-by: Simon LUCIDO --- src/allocation_engine/interval_tree.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/allocation_engine/interval_tree.rs b/src/allocation_engine/interval_tree.rs index 295c09b..6cc1eff 100644 --- a/src/allocation_engine/interval_tree.rs +++ b/src/allocation_engine/interval_tree.rs @@ -88,7 +88,7 @@ impl InnerNode { } /// Returns all the allocated nodes in the tree. - fn allocated_nodes(&self) -> Vec<&InnerNode> { + pub fn allocated_nodes(&self) -> Vec<&InnerNode> { let mut allocated_nodes = Vec::new(); if self.node_state != NodeState::Free { allocated_nodes.push(self); @@ -102,6 +102,15 @@ impl InnerNode { allocated_nodes } + /// Returns the key of the node. + pub fn key(&self) -> &RangeInclusive { + &self.key + } + + pub fn node_state(&self) -> NodeState { + self.node_state + } + /// Returns a readonly reference to the node associated with the `key` or /// None if the searched key does not exist in the tree. fn search(&self, key: &RangeInclusive) -> Option<&InnerNode> {