diff --git a/rust/indexes/src/nodemap.rs b/rust/indexes/src/nodemap.rs --- a/rust/indexes/src/nodemap.rs +++ b/rust/indexes/src/nodemap.rs @@ -4,7 +4,8 @@ // GNU General Public License version 2 or any later version. use errors::{Result, ResultExt, ErrorKind}; -use radixbuf::radix::{radix_lookup, radix_prefix_lookup, radix_lookup_unchecked, radix_insert}; +use radixbuf::radix::{radix_lookup, radix_prefix_lookup, radix_lookup_unchecked, + radix_insert_with_key, KeyReader}; use radixbuf::key::KeyId; use radixbuf::errors as rerrors; @@ -87,7 +88,7 @@ // Sanity check: if the last node stored in the index does not match the changelog, // the index is broken and needs rebuilt. That could happen if strip happens. let rev: KeyId = (next_rev - 1).into(); - let node = rev_to_node(&changelog, rev)?; + let node = rev_to_node(changelog.as_ref(), rev)?; if let Ok(Some(id)) = radix_lookup_unchecked(&main_index, 1, &node) { if id != rev { return Err(ErrorKind::IndexCorrupted.into()); @@ -97,9 +98,10 @@ } } + // Build side_index for the revisions not in the main index let mut side_index = vec![0u32; 16]; - build(&changelog, &mut side_index, 0, next_rev, end_rev)?; + build(&changelog.as_ref(), &mut side_index, 0, next_rev, end_rev)?; Ok(NodeRevMap { changelog, @@ -115,26 +117,26 @@ None => return Ok(None), }; let iter = bin_prefix.iter().cloned(); - let cl = &self.changelog; - let main_res = radix_prefix_lookup(&self.main_index, 1, iter.clone(), rev_to_node, cl)?; - let side_res = radix_prefix_lookup(&self.side_index, 0, iter, rev_to_node, cl)?; + let key_reader = KeyReader::KeyBufferZeroCopy(rev_to_node, self.changelog.as_ref()); + let main_res = radix_prefix_lookup(&self.main_index, 1, iter.clone(), &key_reader)?; + let side_res = radix_prefix_lookup(&self.side_index, 0, iter, &key_reader)?; match (main_res, side_res) { (Some(_), Some(_)) => { let err: rerrors::Error = rerrors::ErrorKind::AmbiguousPrefix.into(); Err(err.into()) } (Some(rev), None) | - (None, Some(rev)) => Ok(Some(rev_to_node(&self.changelog, rev)?)), + (None, Some(rev)) => Ok(Some(rev_to_node(self.changelog.as_ref(), rev)?)), _ => Ok(None), } } /// node -> rev pub fn node_to_rev>(&self, node: T) -> Result> { - let cl = &self.changelog; - if let Some(rev) = radix_lookup(&self.main_index, 1, &node, rev_to_node, cl)? { + let key_reader = KeyReader::KeyBufferZeroCopy(rev_to_node, self.changelog.as_ref()); + if let Some(rev) = radix_lookup(&self.main_index, 1, &node, &key_reader)? { Ok(Some(rev.into())) - } else if let Some(rev) = radix_lookup(&self.side_index, 0, &node, rev_to_node, cl)? { + } else if let Some(rev) = radix_lookup(&self.side_index, 0, &node, &key_reader)? { Ok(Some(rev.into())) } else { Ok(None) @@ -154,7 +156,7 @@ let mut index = self.main_index.as_ref().to_vec(); let end_rev = changelog_next_rev(&self.changelog); let next_rev = index[0].to_le(); - build(&self.changelog, &mut index, 1, next_rev, end_rev)?; + build(self.changelog.as_ref(), &mut index, 1, next_rev, end_rev)?; index[0] = end_rev.to_le(); Ok(index) } @@ -173,24 +175,24 @@ /// Read the given range of revisions (from `start_rev` (inclusive) to /// `end_rev` (exclusive)) from changelog. Insert them to a "forked" /// index and return that index. -fn build(changelog: &T, index: &mut Vec, radix_offset: u32, start_rev: u32, end_rev: u32) - -> Result<()> -where - T: AsRef<[u8]>, -{ +fn build(changelog: &[u8], index: &mut Vec, offset: u32, start_rev: u32, end_rev: u32) + -> Result<()> { // Reserve the approximate size needed for the index - 28 bytes for each revision. // See https://phab.mercurial-scm.org/D1291 for a table of number of revisions and index sizes. index.reserve(7 * (end_rev - start_rev) as usize); + let key_reader = KeyReader::KeyBufferZeroCopy(rev_to_node, changelog); for i in start_rev..end_rev { - let res = radix_insert(index, radix_offset, i.into(), rev_to_node, changelog); + let node = rev_to_node(changelog, i.into())?; + let res = radix_insert_with_key(index, offset, i.into(), &node, &key_reader); res.chain_err(|| ErrorKind::IndexCorrupted)?; } Ok(()) } /// Helper method similar to `radixbuf::key::FixedKey::read`, but takes a revision number instead. -fn rev_to_node>(changelog: &K, rev: KeyId) -> rerrors::Result<&[u8]> { - let buf = changelog.as_ref(); +#[inline] +fn rev_to_node(changelog: &[u8], rev: KeyId) -> rerrors::Result<&[u8]> { + let buf = changelog; let rev_usize: usize = rev.into(); let start_pos = rev_usize * 64 + 32; let end_pos = start_pos + 20; diff --git a/rust/radixbuf/src/key.rs b/rust/radixbuf/src/key.rs --- a/rust/radixbuf/src/key.rs +++ b/rust/radixbuf/src/key.rs @@ -67,8 +67,7 @@ impl FixedKey { #[inline] - pub fn read<'a, K: AsRef<[u8]>>(key_buf: &'a K, key_id: KeyId) -> Result<&'a [u8]> { - let key_buf = key_buf.as_ref(); + pub fn read(key_buf: &[u8], key_id: KeyId) -> Result<&[u8]> { let start_pos: usize = key_id.into(); // LANG: Consider making 20 a type parameter once supported. let end_pos = start_pos + 20; @@ -94,7 +93,7 @@ impl VariantKey { #[inline] - pub fn read<'a, K: AsRef<[u8]>>(key_buf: &'a K, key_id: KeyId) -> Result<&'a [u8]> { + pub fn read(key_buf: &[u8], key_id: KeyId) -> Result<&[u8]> { let key_buf = key_buf.as_ref(); let mut reader = Cursor::new(key_buf); reader.seek(SeekFrom::Start(key_id.into()))?; diff --git a/rust/radixbuf/src/radix.rs b/rust/radixbuf/src/radix.rs --- a/rust/radixbuf/src/radix.rs +++ b/rust/radixbuf/src/radix.rs @@ -120,6 +120,47 @@ // Public APIs +/// Defines how to get a key from `KeyId`. +pub enum KeyReader<'a> { + // Read from a key buffer. Return zero-copy slice from key buffer. + KeyBufferZeroCopy(fn(&[u8], KeyId) -> Result<&[u8]>, &'a [u8]), + + // Read from the radix buffer. Return allocated slice. + RadixBufferAllocated(fn(&[u32], KeyId) -> Result>), + + // Read from both buffers. Return zero-copy slice from the key buffer. + BothBufferZeroCopy(fn(&'a [u8], &[u32], KeyId) -> Result<&'a [u8]>, &'a [u8]), + + // Read from both buffers. Return allocated slice. Slowest but most flexible. + BothBufferAllocated(fn(&[u8], &[u32], KeyId) -> Result>, &'a [u8]), +} + +// Use macro so "$radix" won't be borrowed for the entire "match" block. +macro_rules! read_key { + ($reader: expr, $id: expr, $key: ident, $radix: tt, $block: tt) => { + match $reader { + &KeyReader::KeyBufferZeroCopy(ref f, ref arg) => { + let $key = f(arg, $id)?; + $block + } + &KeyReader::RadixBufferAllocated(ref f) => { + let key = f($radix.as_ref(), $id)?; + let $key = &*key; + $block + } + &KeyReader::BothBufferZeroCopy(ref f, ref arg) => { + let $key = f(arg, $radix.as_ref(), $id)?; + $block + } + &KeyReader::BothBufferAllocated(ref f, ref arg) => { + let key = f(arg, $radix.as_ref(), $id)?; + let $key = &*key; + $block + } + } + } +} + /// Lookup a given `Key`. Return an optional potentially matched `KeyId`. /// The caller is responsible to check whether `KeyId` matches the given `Key` or not. #[inline] @@ -135,83 +176,68 @@ Ok(key_id) } -// unfortunately rustfmt makes the parameter list longer than 100 chars so it's disabled for now. - /// Lookup a given `Key`. Return a verified `KeyId` or `None`. -#[cfg_attr(rustfmt, rustfmt_skip)] #[inline] -pub fn radix_lookup( - radix_buf: &R, offset: u32, key: &K, key_reader: KR, key_reader_arg: &KA) +pub fn radix_lookup(radix_buf: &R, offset: u32, key: &K, key_reader: &KeyReader) -> Result> where R: AsRef<[u32]>, K: AsRef<[u8]>, - KR: Fn(&KA, KeyId) -> Result<&[u8]>, { let key_id = radix_lookup_unchecked(radix_buf, offset, key)?; if let Some(id) = key_id { - let existing_key = key_reader(key_reader_arg, id).unwrap(); - if existing_key != key.as_ref() { - return Ok(None); - } + read_key!(key_reader, id, existing_key, radix_buf, { + if existing_key != key.as_ref() { + return Ok(None); + } + }); } Ok(key_id) } /// Lookup a unique `KeyId` given a prefix of a binary base16 sequence. -#[cfg_attr(rustfmt, rustfmt_skip)] #[inline] -pub fn radix_prefix_lookup( - radix_buf: &R, offset: u32, prefix: P, key_reader: KR, key_reader_arg: &KA) +pub fn radix_prefix_lookup(radix_buf: &R, offset: u32, prefix: P, key_reader: &KeyReader) -> Result> where R: AsRef<[u32]>, P: Iterator + Clone, - KR: Fn(&KA, KeyId) -> Result<&[u8]>, { let root = RadixOffset::new(offset); let (key_id, (_radix, i, _b)) = root.follow(radix_buf, prefix.clone())?; if let Some(id) = key_id { - let key = key_reader(key_reader_arg, id)?; - let iter = key.to_base16iter(); - let matched = iter.clone().skip(i).zip(prefix.clone().skip(i)).all(|(b1, b2)| b1 == b2); - if !matched || iter.count() < prefix.count() { - return Ok(None); - } + read_key!(key_reader, id, key, radix_buf, { + let iter = key.to_base16iter(); + let mut zipped = iter.clone().skip(i).zip(prefix.clone().skip(i)); + let matched = zipped.all(|(b1, b2)| b1 == b2); + if !matched || iter.count() < prefix.count() { + return Ok(None); + } + }); } Ok(key_id) } /// Insert a `Key` to `KeyId` relationship into the radix tree. -#[cfg_attr(rustfmt, rustfmt_skip)] #[inline] -pub fn radix_insert( - radix_buf: &mut R, offset: u32, key_id: KeyId, key_reader: KR, key_reader_arg: &KA) +pub fn radix_insert(radix_buf: &mut R, offset: u32, key_id: KeyId, key_reader: &KeyReader) -> Result<()> where R: Resize + AsRef<[u32]> + AsMut<[u32]>, - KR: Fn(&KA, KeyId) -> Result<&[u8]>, { - let new_key = key_reader(key_reader_arg, key_id)?; - radix_insert_with_key( - radix_buf, - offset, - key_id, - &new_key, - key_reader, - key_reader_arg, - ) + read_key!(key_reader, key_id, new_key, radix_buf, { + radix_insert_with_key(radix_buf, offset, key_id, &new_key, key_reader) + }) } #[cfg_attr(rustfmt, rustfmt_skip)] #[inline] -pub fn radix_insert_with_key( - radix_buf: &mut R, offset: u32, id: KeyId, new_key: &K, key_reader: KR, key_reader_arg: &KA) +pub fn radix_insert_with_key( + radix_buf: &mut R, offset: u32, id: KeyId, new_key: &K, key_reader: &KeyReader) -> Result<()> where R: Resize + AsRef<[u32]> + AsMut<[u32]>, K: AsRef<[u8]>, - KR: Fn(&KA, KeyId) -> Result<&[u8]>, { let new_key_id = id; let root = RadixOffset::new(offset); @@ -224,36 +250,36 @@ } // Need to do a leaf split - let old_key = key_reader(key_reader_arg, old_key_id)?; + read_key!(key_reader, old_key_id, old_key, radix_buf, { + // Find common prefix starting from the next base16 integer + let mut common_len = 0; + let old_iter = old_key.to_base16iter().skip(i + 1); + let new_iter = new_key.to_base16iter().skip(i + 1); + for (b1, b2) in old_iter.zip(new_iter) { + if b1 == b2 { + common_len += 1; + } else { + // Got a chain of radix nodes to write back + let mut node = RadixOffset::create(radix_buf)?; + node.write_key_id(radix_buf, b1, old_key_id)?; + node.write_key_id(radix_buf, b2, new_key_id)?; + let new_iter = new_key.to_base16iter().skip(i + 1); + for k in new_iter.take(common_len).rev() { + let new_node = RadixOffset::create(radix_buf)?; + new_node.write_radix(radix_buf, k, node)?; + node = new_node; + } + return radix.write_radix(radix_buf, b, node); + } + } - // Find common prefix starting from the next base16 integer - let mut common_len = 0; - let old_iter = old_key.to_base16iter().skip(i + 1); - let new_iter = new_key.to_base16iter().skip(i + 1); - for (b1, b2) in old_iter.zip(new_iter) { - if b1 == b2 { - common_len += 1; + // new_key is a prefix of old_key, or vice-versa + if old_key.len() > new_key.as_ref().len() { + bail!(ErrorKind::PrefixConflict(new_key_id, old_key_id)) } else { - // Got a chain of radix nodes to write back - let mut node = RadixOffset::create(radix_buf)?; - node.write_key_id(radix_buf, b1, old_key_id)?; - node.write_key_id(radix_buf, b2, new_key_id)?; - let new_iter = new_key.to_base16iter().skip(i + 1); - for k in new_iter.take(common_len).rev() { - let new_node = RadixOffset::create(radix_buf)?; - new_node.write_radix(radix_buf, k, node)?; - node = new_node; - } - return radix.write_radix(radix_buf, b, node); + bail!(ErrorKind::PrefixConflict(old_key_id, new_key_id)) } - } - - // new_key is a prefix of old_key, or vice-versa - if old_key.len() > new_key.as_ref().len() { - Err(ErrorKind::PrefixConflict(new_key_id, old_key_id).into()) - } else { - Err(ErrorKind::PrefixConflict(old_key_id, new_key_id).into()) - } + }) } None => radix.write_key_id(radix_buf, b, new_key_id), } @@ -273,69 +299,81 @@ let mut key_buf = vec![0u8; 10]; let mut radix_buf = vec![0u32; 15]; - // KeyId exceeds format limit - let key = [0u8; 20]; - let key_id = (1u32 << 31).into(); - let r = radix_insert_with_key(&mut radix_buf, 0, key_id, &key, FixedKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "offset overflow"); + { + // KeyId exceeds format limit + let key = [0u8; 20]; + let key_id = (1u32 << 31).into(); + let key_reader = KeyReader::KeyBufferZeroCopy(FixedKey::read, &key_buf); + let r = radix_insert_with_key(&mut radix_buf, 0, key_id, &key, &key_reader); + assert_eq!(r.unwrap_err().description(), "offset overflow"); - // KeyId exceeds key buffer length - let key_id = 30u32.into(); - let r = radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "invalid key id"); - let r = radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, &key_buf); - let t = format!("{}", r.unwrap_err()); - assert_eq!(t, "KeyId(30) cannot be resolved"); + // KeyId exceeds key buffer length + let key_id = 30u32.into(); + let r = radix_insert(&mut radix_buf, 0, key_id, &key_reader); + assert_eq!(r.unwrap_err().description(), "invalid key id"); + let r = radix_insert(&mut radix_buf, 0, key_id, &key_reader); + let t = format!("{}", r.unwrap_err()); + assert_eq!(t, "KeyId(30) cannot be resolved"); - // Radix root node offset exceeds radix buffer length - let r = radix_insert_with_key(&mut radix_buf, 16, key_id, &key, FixedKey::read, &key_buf); - assert_eq!(format!("{}", r.unwrap_err()), "offset 16 is out of range"); + // Radix root node offset exceeds radix buffer length + let r = radix_insert_with_key(&mut radix_buf, 16, key_id, &key, &key_reader); + assert_eq!(format!("{}", r.unwrap_err()), "offset 16 is out of range"); - // Radix node offset out of range during a lookup - let prefix = [0xf].iter().cloned(); - let r = radix_prefix_lookup(&radix_buf, 0, prefix, FixedKey::read, &key_buf); - assert_eq!(format!("{}", r.unwrap_err()), "offset 15 is out of range"); + // Radix node offset out of range during a lookup + let prefix = [0xf].iter().cloned(); + let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader); + assert_eq!(format!("{}", r.unwrap_err()), "offset 15 is out of range"); - // Base16 sequence overflow - let prefix = [21].iter().cloned(); - let r = radix_prefix_lookup(&radix_buf, 0, prefix.clone(), FixedKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "invalid base16 value"); - let r = radix_prefix_lookup(&radix_buf, 0, prefix, FixedKey::read, &key_buf); - assert_eq!(format!("{}", r.unwrap_err()), "21 is not a base16 value"); + // Base16 sequence overflow + let prefix = [21].iter().cloned(); + let r = radix_prefix_lookup(&radix_buf, 0, prefix.clone(), &key_reader); + assert_eq!(r.unwrap_err().description(), "invalid base16 value"); + let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader); + assert_eq!(format!("{}", r.unwrap_err()), "21 is not a base16 value"); + } - // Inserting a same key with a same `KeyId` is okay let key_id1 = VariantKey::append(&mut key_buf, &b"ab"); let key_id2 = VariantKey::append(&mut key_buf, &b"ab"); - radix_insert(&mut radix_buf, 0, key_id1, VariantKey::read, &key_buf).expect("insert"); - radix_insert(&mut radix_buf, 0, key_id1, VariantKey::read, &key_buf).expect("insert"); + { + // Inserting a same key with a same `KeyId` is okay + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + radix_insert(&mut radix_buf, 0, key_id1, &key_reader).expect("insert"); + radix_insert(&mut radix_buf, 0, key_id1, &key_reader).expect("insert"); - // But not okay if `KeyId` are different - let r = radix_insert(&mut radix_buf, 0, key_id2, VariantKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "key prefix conflict"); + // But not okay if `KeyId` are different + let r = radix_insert(&mut radix_buf, 0, key_id2, &key_reader); + assert_eq!(r.unwrap_err().description(), "key prefix conflict"); + } - // A key cannot be a prefix of another key let key_id4 = VariantKey::append(&mut key_buf, &b"a"); let key_id5 = VariantKey::append(&mut key_buf, &b"abc"); - let r = radix_insert(&mut radix_buf, 0, key_id4, VariantKey::read, &key_buf); - assert_eq!( - format!("{}", r.unwrap_err()), - format!("{:?} cannot be a prefix of {:?}", key_id4, key_id1) - ); - let r = radix_insert(&mut radix_buf, 0, key_id5, VariantKey::read, &key_buf); - assert_eq!( - format!("{}", r.unwrap_err()), - format!("{:?} cannot be a prefix of {:?}", key_id1, key_id5) - ); + { + // A key cannot be a prefix of another key + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + let r = radix_insert(&mut radix_buf, 0, key_id4, &key_reader); + assert_eq!( + format!("{}", r.unwrap_err()), + format!("{:?} cannot be a prefix of {:?}", key_id4, key_id1) + ); + let r = radix_insert(&mut radix_buf, 0, key_id5, &key_reader); + assert_eq!( + format!("{}", r.unwrap_err()), + format!("{:?} cannot be a prefix of {:?}", key_id1, key_id5) + ); + } - // Enforce a leaf split of key_id1 - let key_id3 = VariantKey::append(&mut key_buf, &b"ac"); - radix_insert(&mut radix_buf, 0, key_id3, VariantKey::read, &key_buf).expect("insert"); + { + // Enforce a leaf split of key_id1 + let key_id3 = VariantKey::append(&mut key_buf, &b"ac"); + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + radix_insert(&mut radix_buf, 0, key_id3, &key_reader).expect("insert"); - // Still impossible to cause key prefix conflicts - let r = radix_insert(&mut radix_buf, 0, key_id4, VariantKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); - let r = radix_insert(&mut radix_buf, 0, key_id5, VariantKey::read, &key_buf); - assert_eq!(r.unwrap_err().description(), "key prefix conflict"); + // Still impossible to cause key prefix conflicts + let r = radix_insert(&mut radix_buf, 0, key_id4, &key_reader); + assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); + let r = radix_insert(&mut radix_buf, 0, key_id5, &key_reader); + assert_eq!(r.unwrap_err().description(), "key prefix conflict"); + } } #[test] @@ -348,43 +386,111 @@ // With a single key let key1 = b"01ab"; let key1_id = VariantKey::append(&mut key_buf, &key1); - radix_insert(&mut radix_buf, 0, key1_id, VariantKey::read, &key_buf).expect("insert"); - for i in 0..query.len() { - let prefix = query.clone().take(i); - let r = radix_prefix_lookup(&radix_buf, 0, prefix, VariantKey::read, &key_buf); - if i == 0 { - // This is sub-optimal. But see the NOTE in RadixOffset::follow. - assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); - } else if i <= key1.len() * 2 { - assert_eq!(r.unwrap(), Some(key1_id)); - } else { - assert_eq!(r.unwrap(), None); + { + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + radix_insert(&mut radix_buf, 0, key1_id, &key_reader).expect("insert"); + for i in 0..query.len() { + let prefix = query.clone().take(i); + let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader); + if i == 0 { + // This is sub-optimal. But see the NOTE in RadixOffset::follow. + assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); + } else if i <= key1.len() * 2 { + assert_eq!(r.unwrap(), Some(key1_id)); + } else { + assert_eq!(r.unwrap(), None); + } } } // With another key let key2 = b"01bbc"; let key2_id = VariantKey::append(&mut key_buf, &key2); - radix_insert(&mut radix_buf, 0, key2_id, VariantKey::read, &key_buf).expect("insert"); - for i in 0..query.len() { - let prefix = query.clone().take(i); - let r = radix_prefix_lookup(&radix_buf, 0, prefix, VariantKey::read, &key_buf); - if i <= 5 { - assert_eq!(r.unwrap_err().description(), "ambiguous prefix") - } else if i <= key1.len() * 2 { - assert_eq!(r.unwrap(), Some(key1_id)); - } else { - assert_eq!(r.unwrap(), None); + { + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + radix_insert(&mut radix_buf, 0, key2_id, &key_reader).expect("insert"); + for i in 0..query.len() { + let prefix = query.clone().take(i); + let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader); + if i <= 5 { + assert_eq!(r.unwrap_err().description(), "ambiguous prefix") + } else if i <= key1.len() * 2 { + assert_eq!(r.unwrap(), Some(key1_id)); + } else { + assert_eq!(r.unwrap(), None); + } } + + let query = b"1".to_base16iter(); + let r = radix_prefix_lookup(&radix_buf, 0, query, &key_reader); + assert_eq!(r.unwrap(), None); + + let query = b"01b".to_base16iter(); + let r = radix_prefix_lookup(&radix_buf, 0, query, &key_reader); + assert_eq!(r.unwrap(), Some(key2_id)); } + } + + #[test] + fn test_key_reader() { + // KeyReader::KeyBufferZeroCopy is tested in other places. Test other readers. + + // RadixBufferAllocated + let key_reader = KeyReader::RadixBufferAllocated(|radix_buf, key_id| { + // One level of redirection: radix_buf[key_id] + let pos: usize = key_id.into(); + let v = radix_buf[pos]; + let buf: [u8; 4] = unsafe { transmute(v.to_le()) }; + Ok(Vec::from(&buf[..]).into_boxed_slice()) + }); + let mut radix_buf = vec![0u32; 16]; + radix_buf.push(0xa); + radix_buf.push(0xb); + radix_insert(&mut radix_buf, 0, 16u32.into(), &key_reader).expect("insert"); + assert_eq!(radix_buf.len(), 18); + radix_insert(&mut radix_buf, 0, 17u32.into(), &key_reader).expect("insert"); + assert_eq!(radix_buf.len(), 34); + assert_eq!( + radix_lookup_unchecked(&radix_buf, 0, &[0xb, 0, 0, 0]).unwrap(), + Some(17u32.into()) + ); - let query = b"1".to_base16iter(); - let r = radix_prefix_lookup(&radix_buf, 0, query, VariantKey::read, &key_buf); - assert_eq!(r.unwrap(), None); + // BothBufferZeroCopy + let key_reader = KeyReader::BothBufferZeroCopy( + |key_buf, radix_buf, key_id| { + // key_buf[radix_buf[key_id]..radix_buf[key_id+1]] + let pos: usize = key_id.into(); + let start = radix_buf[pos] as usize; + let len = radix_buf[pos + 1] as usize; + Ok(&key_buf[start..start + len]) + }, + b"0123456789abcdef", + ); + let mut radix_buf = vec![0u32; 16]; + radix_buf.extend_from_slice(&[1, 3, 7, 5]); + radix_insert(&mut radix_buf, 0, 16u32.into(), &key_reader).expect("insert"); + radix_insert(&mut radix_buf, 0, 18u32.into(), &key_reader).expect("insert"); + assert_eq!(radix_buf.len(), 36); + let r = radix_lookup(&radix_buf, 0, b"123", &key_reader).unwrap(); + assert_eq!(r, Some(16u32.into())); + let r = radix_lookup(&radix_buf, 0, b"789ab", &key_reader).unwrap(); + assert_eq!(r, Some(18u32.into())); - let query = b"01b".to_base16iter(); - let r = radix_prefix_lookup(&radix_buf, 0, query, VariantKey::read, &key_buf); - assert_eq!(r.unwrap(), Some(key2_id)); + // BothBufferAllocated + let key_reader = KeyReader::BothBufferAllocated( + |key_buf, radix_buf, key_id| { + // key_buf[radix_buf[key_id]..radix_buf[key_id+1]] + let pos: usize = key_id.into(); + let start = radix_buf[pos] as usize; + let len = radix_buf[pos + 1] as usize; + Ok(key_buf[start..start + len].to_vec().into_boxed_slice()) + }, + b"0123456789abcdef", + ); + let r = radix_lookup(&radix_buf, 0, b"123", &key_reader).unwrap(); + assert_eq!(r, Some(16u32.into())); + let r = radix_lookup(&radix_buf, 0, b"789ab", &key_reader).unwrap(); + assert_eq!(r, Some(18u32.into())); } quickcheck! { @@ -417,12 +523,14 @@ // Insert to radix tree for key in &std_set { let key_id = VariantKey::append(&mut key_buf, key); - radix_insert(&mut radix_buf, 0, key_id, VariantKey::read, &key_buf).expect("insert"); + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); + radix_insert(&mut radix_buf, 0, key_id, &key_reader).expect("insert"); } // Test key existence + let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf); std_set.iter().all(|key| { - let r = radix_lookup(&radix_buf, 0, key, VariantKey::read, &key_buf); + let r = radix_lookup(&radix_buf, 0, key, &key_reader); r.unwrap().is_some() }) } @@ -450,10 +558,11 @@ fn bench_lookups(b: &mut Bencher) { let key_buf = randomized_key_buf(COUNT); let radix_buf = batch_insert_radix_buf(&key_buf, COUNT); + let key_reader = KeyReader::KeyBufferZeroCopy(FixedKey::read, &key_buf); b.iter(|| for i in 0..COUNT { let key_id = (i as u32 * 20).into(); let key = FixedKey::read(&key_buf, key_id).unwrap(); - radix_lookup(&radix_buf, 0, &key, FixedKey::read, &key_buf).expect("lookup"); + radix_lookup(&radix_buf, 0, &key, &key_reader).expect("lookup"); }) } @@ -466,9 +575,10 @@ fn batch_insert_radix_buf(key_buf: &Vec, count: usize) -> Vec { let mut radix_buf = vec![0u32; 16]; + let key_reader = KeyReader::KeyBufferZeroCopy(FixedKey::read, key_buf); for i in 0..count { let key_id: KeyId = ((i * 20) as u32).into(); - radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, key_buf).expect("insert"); + radix_insert(&mut radix_buf, 0, key_id, &key_reader).expect("insert"); } radix_buf }