use std::fmt::Debug; use std::marker::PhantomData; use super::node_ptr; use super::node_ptr::NodePtr; use crate::EpochPin; use crate::Value; use crate::algorithm::lock_and_version::AtomicLockAndVersion; use crate::algorithm::lock_and_version::ConcurrentUpdateError; use crate::allocator::ArtAllocator; use crate::allocator::OutOfMemoryError; pub struct NodeRef<'e, V> { ptr: NodePtr, phantom: PhantomData<&'e EpochPin<'e>>, } impl<'e, V> Debug for NodeRef<'e, V> { fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { write!(fmt, "{:?}", self.ptr) } } impl<'e, V: Value> NodeRef<'e, V> { pub(crate) fn from_root_ptr(root_ptr: NodePtr) -> NodeRef<'e, V> { NodeRef { ptr: root_ptr, phantom: PhantomData, } } pub(crate) fn read_lock_or_restart( &self, ) -> Result, ConcurrentUpdateError> { let version = self.lockword().read_lock_or_restart()?; Ok(ReadLockedNodeRef { ptr: self.ptr, version, phantom: self.phantom, }) } pub(crate) fn write_lock_or_restart( &self, ) -> Result, ConcurrentUpdateError> { self.lockword().write_lock_or_restart()?; Ok(WriteLockedNodeRef { ptr: self.ptr, phantom: self.phantom, }) } fn lockword(&self) -> &AtomicLockAndVersion { self.ptr.lockword() } } /// A reference to a node that has been optimistically read-locked. The functions re-check /// the version after each read. pub struct ReadLockedNodeRef<'e, V> { ptr: NodePtr, version: u64, phantom: PhantomData<&'e EpochPin<'e>>, } impl<'e, V: Value> ReadLockedNodeRef<'e, V> { pub(crate) fn is_leaf(&self) -> bool { self.ptr.is_leaf() } pub(crate) fn is_full(&self) -> bool { self.ptr.is_full() } pub(crate) fn get_prefix(&self) -> &[u8] { self.ptr.get_prefix() } /// Note: because we're only holding a read lock, the prefix can change concurrently. /// You must be prepared to restart, if read_unlock() returns error later. /// /// Returns the length of the prefix, or None if it's not a match pub(crate) fn prefix_matches(&self, key: &[u8]) -> Option { self.ptr.prefix_matches(key) } pub(crate) fn find_child_or_restart( &self, key_byte: u8, ) -> Result>, ConcurrentUpdateError> { let child_or_value = self.ptr.find_child(key_byte); self.ptr.lockword().check_or_restart(self.version)?; match child_or_value { None => Ok(None), Some(child_ptr) => Ok(Some(NodeRef { ptr: child_ptr, phantom: self.phantom, })), } } pub(crate) fn find_next_child_or_restart( &self, min_key_byte: u8, ) -> Result)>, ConcurrentUpdateError> { let child_or_value = self.ptr.find_next_child(min_key_byte); self.ptr.lockword().check_or_restart(self.version)?; match child_or_value { None => Ok(None), Some((k, child_ptr)) => Ok(Some(( k, NodeRef { ptr: child_ptr, phantom: self.phantom, }, ))), } } pub(crate) fn get_leaf_value_ptr(&self) -> Result<*const V, ConcurrentUpdateError> { let result = self.ptr.get_leaf_value(); self.ptr.lockword().check_or_restart(self.version)?; // Extend the lifetime. let result = std::ptr::from_ref(result); Ok(result) } pub(crate) fn upgrade_to_write_lock_or_restart( self, ) -> Result, ConcurrentUpdateError> { self.ptr .lockword() .upgrade_to_write_lock_or_restart(self.version)?; Ok(WriteLockedNodeRef { ptr: self.ptr, phantom: self.phantom, }) } pub(crate) fn read_unlock_or_restart(self) -> Result<(), ConcurrentUpdateError> { self.ptr.lockword().check_or_restart(self.version)?; Ok(()) } pub(crate) fn check_or_restart(&self) -> Result<(), ConcurrentUpdateError> { self.ptr.lockword().check_or_restart(self.version)?; Ok(()) } } /// A reference to a node that has been optimistically read-locked. The functions re-check /// the version after each read. pub struct WriteLockedNodeRef<'e, V> { ptr: NodePtr, phantom: PhantomData<&'e EpochPin<'e>>, } impl<'e, V: Value> WriteLockedNodeRef<'e, V> { pub(crate) fn can_shrink(&self) -> bool { self.ptr.can_shrink() } pub(crate) fn num_children(&self) -> usize { self.ptr.num_children() } pub(crate) fn write_unlock(mut self) { self.ptr.lockword().write_unlock(); self.ptr = NodePtr::null(); } pub(crate) fn write_unlock_obsolete(mut self) { self.ptr.lockword().write_unlock_obsolete(); self.ptr = NodePtr::null(); } pub(crate) fn get_prefix(&self) -> &[u8] { self.ptr.get_prefix() } pub(crate) fn truncate_prefix(&mut self, new_prefix_len: usize) { self.ptr.truncate_prefix(new_prefix_len) } pub(crate) fn prepend_prefix(&mut self, prefix: &[u8], prefix_byte: u8) { self.ptr.prepend_prefix(prefix, prefix_byte) } pub(crate) fn insert_child(&mut self, key_byte: u8, child: NodePtr) { self.ptr.insert_child(key_byte, child) } pub(crate) fn get_leaf_value_mut(&mut self) -> &mut V { self.ptr.get_leaf_value_mut() } pub(crate) fn grow<'a, A>( &self, allocator: &'a A, ) -> Result, OutOfMemoryError> where A: ArtAllocator, { let new_node = self.ptr.grow(allocator)?; Ok(NewNodeRef { ptr: new_node, allocator, extra_nodes: Vec::new(), }) } pub(crate) fn shrink<'a, A>( &self, allocator: &'a A, ) -> Result, OutOfMemoryError> where A: ArtAllocator, { let new_node = self.ptr.shrink(allocator)?; Ok(NewNodeRef { ptr: new_node, allocator, extra_nodes: Vec::new(), }) } pub(crate) fn as_ptr(&self) -> NodePtr { self.ptr } pub(crate) fn replace_child(&mut self, key_byte: u8, replacement: NodePtr) { self.ptr.replace_child(key_byte, replacement); } pub(crate) fn delete_child(&mut self, key_byte: u8) { self.ptr.delete_child(key_byte); } pub(crate) fn find_remaining_child(&self) -> (u8, NodeRef<'e, V>) { assert_eq!(self.num_children(), 1); let child_or_value = self.ptr.find_next_child(0); match child_or_value { None => panic!("could not find only child in node"), Some((k, child_ptr)) => ( k, NodeRef { ptr: child_ptr, phantom: self.phantom, }, ), } } } impl<'e, V> Drop for WriteLockedNodeRef<'e, V> { fn drop(&mut self) { if !self.ptr.is_null() { self.ptr.lockword().write_unlock(); } } } pub(crate) struct NewNodeRef<'a, V, A> where V: Value, A: ArtAllocator, { ptr: NodePtr, allocator: &'a A, extra_nodes: Vec>, } impl<'a, V, A> NewNodeRef<'a, V, A> where V: Value, A: ArtAllocator, { pub(crate) fn insert_old_child(&mut self, key_byte: u8, child: &WriteLockedNodeRef) { self.ptr.insert_child(key_byte, child.as_ptr()) } pub(crate) fn into_ptr(mut self) -> NodePtr { let ptr = self.ptr; self.ptr = NodePtr::null(); ptr } pub(crate) fn insert_new_child(&mut self, key_byte: u8, child: NewNodeRef<'a, V, A>) { let child_ptr = child.into_ptr(); self.ptr.insert_child(key_byte, child_ptr); self.extra_nodes.push(child_ptr); } } impl<'a, V, A> Drop for NewNodeRef<'a, V, A> where V: Value, A: ArtAllocator, { /// This drop implementation deallocates the newly allocated node, if into_ptr() was not called. fn drop(&mut self) { if !self.ptr.is_null() { self.ptr.deallocate(self.allocator); for p in self.extra_nodes.iter() { p.deallocate(self.allocator); } } } } pub(crate) fn new_internal<'a, V, A>( prefix: &[u8], allocator: &'a A, ) -> Result, OutOfMemoryError> where V: Value, A: ArtAllocator, { Ok(NewNodeRef { ptr: node_ptr::new_internal(prefix, allocator)?, allocator, extra_nodes: Vec::new(), }) } pub(crate) fn new_leaf<'a, V, A>( prefix: &[u8], value: V, allocator: &'a A, ) -> Result, OutOfMemoryError> where V: Value, A: ArtAllocator, { Ok(NewNodeRef { ptr: node_ptr::new_leaf(prefix, value, allocator)?, allocator, extra_nodes: Vec::new(), }) }