Skip to content

Commit

Permalink
Add dict_get_bound
Browse files Browse the repository at this point in the history
  • Loading branch information
Rexagon committed Aug 7, 2023
1 parent fd98f6b commit e021969
Show file tree
Hide file tree
Showing 3 changed files with 250 additions and 3 deletions.
152 changes: 152 additions & 0 deletions src/dict/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,158 @@ pub fn dict_get_owned(
})
}

/// Finds the specified dict bound and returns a key and a value corresponding to the key.
pub fn dict_find_bound<'a: 'b, 'b>(
root: &'a Option<Cell>,
mut key_bit_len: u16,
bound: DictBound,
signed: bool,
) -> Result<Option<(CellBuilder, CellSlice<'b>)>, Error> {
let mut data = match root {
Some(data) => ok!(data.as_slice()),
None => return Ok(None),
};

let mut direction = None;
let mut key = CellBuilder::new();

// Try to find the required leaf
loop {
// Read the key part written in the current edge
let prefix = ok!(read_label(&mut data, key_bit_len));
#[allow(clippy::needless_borrow)]
if !prefix.is_data_empty() {
ok!(key.store_slice_data(&prefix));
}

match key_bit_len.checked_sub(prefix.remaining_bits()) {
Some(0) => break,
Some(remaining) => {
key_bit_len = remaining;
if data.remaining_refs() < 2 {
return Err(Error::CellUnderflow);
}
}
None => return Err(Error::CellUnderflow),
}

if key_bit_len < 1 {
return Err(Error::CellUnderflow);
}
key_bit_len -= 1;

let child_index = match direction {
// Compute direction by the first part
None => {
let mut child_index = *direction.insert(bound.into_bit());
// Invert first bit for signed keys if starting from the empty part
if signed && prefix.is_data_empty() {
child_index = !child_index;
}
child_index
}
// Use the same direction for all remaining parts
Some(direction) => direction,
};

ok!(key.store_bit(child_index));

// Load next child based on the next bit
data = ok!(data.cell().get_reference_as_slice(child_index as u8));
}

// Return the last slice as data
Ok(Some((key, data)))
}

/// Finds the specified dict bound and returns a key and cell slice parts corresponding to the key.
pub fn dict_find_bound_owned(
root: &Option<Cell>,
mut key_bit_len: u16,
bound: DictBound,
signed: bool,
) -> Result<Option<(CellBuilder, CellSliceParts)>, Error> {
let Some(root) = root else {
return Ok(None);
};
let mut data = ok!(root.as_slice());
let mut prev = None;

let mut direction = None;
let mut key = CellBuilder::new();

// Try to find the required leaf
loop {
// Read the key part written in the current edge
let prefix = ok!(read_label(&mut data, key_bit_len));
#[allow(clippy::needless_borrow)]
if !prefix.is_data_empty() {
ok!(key.store_slice_data(&prefix));
}

match key_bit_len.checked_sub(prefix.remaining_bits()) {
Some(0) => break,
Some(remaining) => {
key_bit_len = remaining;
if data.remaining_refs() < 2 {
return Err(Error::CellUnderflow);
}
}
None => return Err(Error::CellUnderflow),
}

let child_index = match direction {
// Compute direction by the first part
None => {
let mut child_index = *direction.insert(bound.into_bit());
// Invert first bit for signed keys if starting from the empty part
if signed && prefix.is_data_empty() {
child_index = !child_index;
}
child_index
}
// Use the same direction for all remaining parts
Some(direction) => direction,
};

ok!(key.store_bit(child_index));

// Load next child based on the next bit
prev = Some((data.cell(), child_index));
data = ok!(data.cell().get_reference_as_slice(child_index as u8));
}

// Build cell slice parts
let slice = match prev {
Some((prev, child_index)) => {
let cell = match prev.reference_cloned(child_index as u8) {
Some(cell) => cell,
None => return Err(Error::CellUnderflow),
};
(cell, data.range())
}
None => (root.clone(), data.range()),
};

// Return the last slice as data
Ok(Some((key, slice)))
}

/// Dictionary bound.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub enum DictBound {
/// The lowest dictionary key.
Min,
/// The largest dictionary key.
Max,
}

impl DictBound {
fn into_bit(self) -> bool {
self == Self::Max
}
}

/// Loads a non-empty dictionary from the root cell.
pub fn dict_load_from_root(
slice: &mut CellSlice<'_>,
Expand Down
48 changes: 46 additions & 2 deletions src/dict/raw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@ use crate::error::Error;
use crate::util::{unlikely, IterStatus};

use super::{
dict_get, dict_get_owned, dict_insert, dict_load_from_root, dict_remove_owned, read_label,
SetMode,
dict_find_bound, dict_find_bound_owned, dict_get, dict_get_owned, dict_insert,
dict_load_from_root, dict_remove_owned, read_label, DictBound, SetMode,
};

/// Dictionary with fixed length keys (where `N` is a number of bits in each key).
Expand Down Expand Up @@ -137,6 +137,50 @@ impl<const N: u16> RawDict<N> {
dict_get_owned(&self.0, N, key)
}

/// Returns the lowest key and a value corresponding to the key.
pub fn get_min(&self, signed: bool) -> Result<Option<(CellBuilder, CellSlice<'_>)>, Error> {
dict_find_bound(&self.0, N, DictBound::Min, signed)
}

/// Returns the largest key and a value corresponding to the key.
pub fn get_max(&self, signed: bool) -> Result<Option<(CellBuilder, CellSlice<'_>)>, Error> {
dict_find_bound(&self.0, N, DictBound::Max, signed)
}

/// Finds the specified dict bound and returns a key and a value corresponding to the key.
pub fn get_bound(
&self,
bound: DictBound,
signed: bool,
) -> Result<Option<(CellBuilder, CellSlice<'_>)>, Error> {
dict_find_bound(&self.0, N, bound, signed)
}

/// Returns the lowest key and cell slice parts corresponding to the key.
pub fn get_min_owned(
&self,
signed: bool,
) -> Result<Option<(CellBuilder, CellSliceParts)>, Error> {
dict_find_bound_owned(&self.0, N, DictBound::Min, signed)
}

/// Returns the largest key and cell slice parts corresponding to the key.
pub fn get_max_owned(
&self,
signed: bool,
) -> Result<Option<(CellBuilder, CellSliceParts)>, Error> {
dict_find_bound_owned(&self.0, N, DictBound::Max, signed)
}

/// Finds the specified dict bound and returns a key and cell slice parts corresponding to the key.
pub fn get_bound_owned(
&self,
bound: DictBound,
signed: bool,
) -> Result<Option<(CellBuilder, CellSliceParts)>, Error> {
dict_find_bound_owned(&self.0, N, bound, signed)
}

/// Returns `true` if the dictionary contains a value for the specified key.
pub fn contains_key(&self, key: CellSlice<'_>) -> Result<bool, Error> {
Ok(ok!(dict_get(&self.0, N, key)).is_some())
Expand Down
53 changes: 52 additions & 1 deletion src/dict/typed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::dict::dict_remove_owned;
use crate::error::Error;
use crate::util::*;

use super::raw::*;
use super::{dict_find_bound, raw::*, DictBound};
use super::{dict_get, dict_insert, dict_load_from_root, serialize_entry, DictKey, SetMode};

/// Typed dictionary with fixed length keys.
Expand Down Expand Up @@ -301,6 +301,43 @@ where
{
Values::new(&self.root, K::BITS)
}

/// Returns the lowest key and a value corresponding to the key.
pub fn get_min<'a>(&'a self, signed: bool) -> Result<Option<(K, V)>, Error>
where
V: Load<'a>,
{
Ok(match ok!(self.get_bound_raw(DictBound::Min, signed)) {
Some((key, ref mut value)) => Some((key, ok!(V::load_from(value)))),
None => None,
})
}

/// Returns the lowest key and a value corresponding to the key.
pub fn get_max<'a>(&'a self, signed: bool) -> Result<Option<(K, V)>, Error>
where
V: Load<'a>,
{
Ok(match ok!(self.get_bound_raw(DictBound::Max, signed)) {
Some((key, ref mut value)) => Some((key, ok!(V::load_from(value)))),
None => None,
})
}

/// Finds the specified dict bound and returns a key and a raw value corresponding to the key.
pub fn get_bound_raw(
&self,
bound: DictBound,
signed: bool,
) -> Result<Option<(K, CellSlice<'_>)>, Error> {
let Some((key, value)) = ok!(dict_find_bound(&self.root, K::BITS, bound, signed)) else {
return Ok(None);
};
match K::from_raw_data(key.raw_data()) {
Some(key) => Ok(Some((key, value))),
None => Err(Error::CellUnderflow),
}
}
}

impl<K, V> Dict<K, V>
Expand Down Expand Up @@ -708,6 +745,20 @@ mod tests {
}
}

#[test]
fn dict_bounds() {
let mut dict = Dict::<i32, bool>::new();
for i in -10..=10 {
dict.set(i, i < 0).unwrap();
}

assert_eq!(dict.get_min(false).unwrap(), Some((0, false)));
assert_eq!(dict.get_max(false).unwrap(), Some((-1, true)));

assert_eq!(dict.get_min(true).unwrap(), Some((-10, true)));
assert_eq!(dict.get_max(true).unwrap(), Some((10, false)));
}

#[test]
fn dict_replace() {
let mut dict = Dict::<u32, bool>::new();
Expand Down

0 comments on commit e021969

Please sign in to comment.