diff --git a/rust/radix204/src/errors.rs b/rust/radix204/src/errors.rs --- a/rust/radix204/src/errors.rs +++ b/rust/radix204/src/errors.rs @@ -22,5 +22,8 @@ description("incomplete key") display("key (id={}) is incomplete", key_id.to_usize()) } + AmbiguousPrefix { + description("ambiguous prefix") + } } } diff --git a/rust/radix204/src/radix.rs b/rust/radix204/src/radix.rs --- a/rust/radix204/src/radix.rs +++ b/rust/radix204/src/radix.rs @@ -65,6 +65,24 @@ } } + /// Return an iterator of non-none (not 0) indexes. + #[inline] + pub fn keys(&self, vec: &mut RB) -> Result> { + let pos = self.0.to_usize(); + let end = pos + 16; + if end > vec.len() { + return Err(ErrorKind::OffsetOverflow(pos).into()); + } + Ok( + vec[pos..end] + .iter() + .enumerate() + .filter(|&(_, v)| *v != 0) + .map(|(i, _)| i as u8) + .collect(), + ) + } + /// Rewrite an index entry to point to another `RadixNode`. #[inline] pub fn write_radix(&self, vec: &mut RB, index: u8, node: &RadixNode) -> Result<()> { @@ -132,6 +150,52 @@ }) } + pub fn lookup_prefix(&self, base16seq: &[u8]) -> Result> { + borrow_shared(&self.1, |rbuf| { + let mut radix = RadixNode::new(self.2); + for &b in base16seq { + debug_assert!(b < 16); + match radix.follow(rbuf, b)? { + FollowResult::Radix(next_radix) => { + radix = next_radix; + } + FollowResult::Missing => { + return Ok(None); + } + FollowResult::KeyId(id) => { + let key = self.read_key(id)?; + let matched = key.to_base16iter().zip(base16seq.iter()).all(|(b1, &b2)| { + b1 == b2 + }); + return if key.len() * 2 >= base16seq.len() && matched { + Ok(Some(id)) + } else { + Ok(None) + }; + } + } + } + loop { + let keys = radix.keys(rbuf)?; + match keys.len() { + 0 => return Ok(None), + 1 => { + match radix.follow(rbuf, keys[0])? { + FollowResult::Radix(next_radix) => { + radix = next_radix; + } + FollowResult::Missing => { + return Ok(None); + } + FollowResult::KeyId(id) => return Ok(Some(id)), + } + } + _ => return Err(ErrorKind::AmbiguousPrefix.into()), + } + } + }) + } + /// Insert when `KeyId` is known pub fn insert(&self, new_key_id: KeyId) -> Result<()> { let new_key = self.read_key(new_key_id)?; @@ -294,6 +358,39 @@ ); } + #[test] + fn test_lookup_prefix() { + let (kbuf, _rbuf, radix) = prepare_kbuf_radix(); + assert_eq!(radix.lookup_prefix(&[0x1]).unwrap(), None); + let key = b"0123456789abcdefghij".into(); + let key_id = append_key(&kbuf, &key); + radix.insert(key_id).expect("insert"); + let base16: Vec = key.to_base16iter() + .chain([0u8; 10].iter().cloned()) + .collect(); + for i in 0..base16.len() { + let r = radix.lookup_prefix(&base16[0..i]).unwrap(); + if i <= 40 { + assert_eq!(r, Some(key_id)); + } else { + assert_eq!(r, None); + } + } + + let key = b"0123456789bbcdefghij".into(); + radix.insert(append_key(&kbuf, &key)).expect("insert"); + for i in 0..base16.len() { + let r = radix.lookup_prefix(&base16[0..i]); + if i <= 21 { + assert_eq!(r.unwrap_err().description(), "ambiguous prefix") + } else if i <= 40 { + assert_eq!(r.unwrap(), Some(key_id)); + } else { + assert_eq!(r.unwrap(), None); + } + } + } + #[bench] fn bench_50k_insert(b: &mut Bencher) { const COUNT: u32 = 51200;