This is an archive of the discontinued Mercurial Phabricator instance.

Differential D1517

radixbuf: make key reader function more flexible
AbandonedPublic

Authored by quark on Nov 27 2017, 1:30 PM.
Tags
None
Subscribers
jsgf

Details

Reviewers
jsgf
Group Reviewers
Restricted Project
Summary

Previously the key reader function (convert KeyId to Key) can take a
&'a T and return a &'a [u8]. That was inspired by C void *, looked
flexible - in theory T could have some RefCell and do whatever it wants.
However, in practise, passing the radix buffer itself to the key reader
function is often desired and there is no easy way to do that cleanly.

This patch replaces the key reader arguments to an enum that support
different dedicated types of key readers - zero-copy from key buffer,
construct something from the radix buffer, or from both buffers.

This adds flexibility like embedding other data directly into the radix
buffer with an optional external key buffer as a dependence.

Test Plan

cargo test --lib

Diff Detail

Repository
rFBHGX Facebook Mercurial Extensions
Lint
Lint Skipped
Unit
Unit Tests Skipped

Event Timeline

quark created this revision.Nov 27 2017, 1:30 PM
Herald added a reviewer: Restricted Project. · View Herald TranscriptNov 27 2017, 1:30 PM
jsgf requested changes to this revision.Dec 7 2017, 12:19 PM
jsgf added a subscriber: jsgf.
jsgf added inline comments.
rust/radixbuf/src/key.rs
70

_like_

96

_like_

rust/radixbuf/src/radix.rs
138–141

I don't really understand this. It looks like something that could be expressed with a function and closures. Unless you really want the early return from $block?

180

This is too big to inline. The compiler can always inline it with LTO if it really wants to.

(Ditto other large functions with #[inline].)

233

Does rustfmt make a mess of this?

234

WAY to big to #[inline].

This revision now requires changes to proceed.Dec 7 2017, 12:19 PM
quark added inline comments.Dec 7 2017, 3:18 PM
rust/radixbuf/src/radix.rs
138–141

I'd be happy to rewrite if there is a cleaner way to get rid of this. What I want is not holding &radix_buf reference when running $block since $block may want &mut radix_buf.

180

Good idea.

233

It folds line 235 to 237 into a single line exceeding 100 chars.

quark abandoned this revision.Dec 14 2017, 8:51 PM

I have changed the childmap to use separate key buffer so this becomes no longer necessary.

Revision Contents
Changeset List

Diff 3879

rust/indexes/src/nodemap.rs

// Copyright 2017 Facebook, Inc. // Copyright 2017 Facebook, Inc.
// //
// This software may be used and distributed according to the terms of the // This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version. // GNU General Public License version 2 or any later version.
use errors::{Result, ResultExt, ErrorKind}; 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::key::KeyId;
use radixbuf::errors as rerrors; use radixbuf::errors as rerrors;
/// An index for node to rev lookups. /// An index for node to rev lookups.
/// ///
/// The index depends entirely on an append-only changelog.i source /// The index depends entirely on an append-only changelog.i source
/// of truth. It does not support in-memory overrides, which could be /// of truth. It does not support in-memory overrides, which could be
/// implemented at a higher level. /// implemented at a higher level.
if next_rev > end_rev { if next_rev > end_rev {
// next_rev cannot be larger than what changelog has. // next_rev cannot be larger than what changelog has.
return Err(ErrorKind::IndexCorrupted.into()); return Err(ErrorKind::IndexCorrupted.into());
} else if next_rev > 0 { } else if next_rev > 0 {
// Sanity check: if the last node stored in the index does not match the changelog, // 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. // the index is broken and needs rebuilt. That could happen if strip happens.
let rev: KeyId = (next_rev - 1).into(); 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 let Ok(Some(id)) = radix_lookup_unchecked(&main_index, 1, &node) {
if id != rev { if id != rev {
return Err(ErrorKind::IndexCorrupted.into()); return Err(ErrorKind::IndexCorrupted.into());
} }
} else { } else {
return Err(ErrorKind::IndexCorrupted.into()); return Err(ErrorKind::IndexCorrupted.into());
} }
} }
// Build side_index for the revisions not in the main index // Build side_index for the revisions not in the main index
let mut side_index = vec![0u32; 16]; 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 { Ok(NodeRevMap {
changelog, changelog,
main_index, main_index,
side_index, side_index,
}) })
} }
/// hex prefix -> node /// hex prefix -> node
pub fn hex_prefix_to_node<T: AsRef<[u8]>>(&self, hex_prefix: T) -> Result<Option<&[u8]>> { pub fn hex_prefix_to_node<T: AsRef<[u8]>>(&self, hex_prefix: T) -> Result<Option<&[u8]>> {
let bin_prefix = match hex_to_bin_base16(hex_prefix) { let bin_prefix = match hex_to_bin_base16(hex_prefix) {
Some(v) => v, Some(v) => v,
None => return Ok(None), None => return Ok(None),
}; };
let iter = bin_prefix.iter().cloned(); let iter = bin_prefix.iter().cloned();
let cl = &self.changelog; let key_reader = KeyReader::KeyBufferZeroCopy(rev_to_node, self.changelog.as_ref());
let main_res = radix_prefix_lookup(&self.main_index, 1, iter.clone(), rev_to_node, cl)?; 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, rev_to_node, cl)?; let side_res = radix_prefix_lookup(&self.side_index, 0, iter, &key_reader)?;
match (main_res, side_res) { match (main_res, side_res) {
(Some(_), Some(_)) => { (Some(_), Some(_)) => {
let err: rerrors::Error = rerrors::ErrorKind::AmbiguousPrefix.into(); let err: rerrors::Error = rerrors::ErrorKind::AmbiguousPrefix.into();
Err(err.into()) Err(err.into())
} }
(Some(rev), None) | (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), _ => Ok(None),
} }
} }
/// node -> rev /// node -> rev
pub fn node_to_rev<T: AsRef<[u8]>>(&self, node: T) -> Result<Option<u32>> { pub fn node_to_rev<T: AsRef<[u8]>>(&self, node: T) -> Result<Option<u32>> {
let cl = &self.changelog; let key_reader = KeyReader::KeyBufferZeroCopy(rev_to_node, self.changelog.as_ref());
if let Some(rev) = radix_lookup(&self.main_index, 1, &node, rev_to_node, cl)? { if let Some(rev) = radix_lookup(&self.main_index, 1, &node, &key_reader)? {
Ok(Some(rev.into())) 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())) Ok(Some(rev.into()))
} else { } else {
Ok(None) Ok(None)
} }
} }
/// How many revisions the side index has. /// How many revisions the side index has.
pub fn lag(&self) -> u32 { pub fn lag(&self) -> u32 {
let next_rev = self.main_index.as_ref()[0].to_le(); let next_rev = self.main_index.as_ref()[0].to_le();
let end_rev = changelog_next_rev(&self.changelog); let end_rev = changelog_next_rev(&self.changelog);
end_rev - next_rev end_rev - next_rev
} }
/// Incrementally build the main index based on the existing one. /// Incrementally build the main index based on the existing one.
pub fn build_incrementally(&self) -> Result<Vec<u32>> { pub fn build_incrementally(&self) -> Result<Vec<u32>> {
// Copy the main index since we need to modify it. // Copy the main index since we need to modify it.
let mut index = self.main_index.as_ref().to_vec(); let mut index = self.main_index.as_ref().to_vec();
let end_rev = changelog_next_rev(&self.changelog); let end_rev = changelog_next_rev(&self.changelog);
let next_rev = index[0].to_le(); 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(); index[0] = end_rev.to_le();
Ok(index) Ok(index)
} }
} }
/// Return the minimal revision number the changelog does not have. /// Return the minimal revision number the changelog does not have.
fn changelog_next_rev<T: AsRef<[u8]>>(changelog: &T) -> u32 { fn changelog_next_rev<T: AsRef<[u8]>>(changelog: &T) -> u32 {
let changelog = changelog.as_ref(); let changelog = changelog.as_ref();
let rev = changelog.len() as u64 / 64; let rev = changelog.len() as u64 / 64;
if rev > (1u64 << 31) { if rev > (1u64 << 31) {
panic!("rev exceeds 32 bit integers") panic!("rev exceeds 32 bit integers")
} }
rev as u32 rev as u32
} }
/// Read the given range of revisions (from `start_rev` (inclusive) to /// Read the given range of revisions (from `start_rev` (inclusive) to
/// `end_rev` (exclusive)) from changelog. Insert them to a "forked" /// `end_rev` (exclusive)) from changelog. Insert them to a "forked"
/// index and return that index. /// index and return that index.
fn build<T>(changelog: &T, index: &mut Vec<u32>, radix_offset: u32, start_rev: u32, end_rev: u32) fn build(changelog: &[u8], index: &mut Vec<u32>, offset: u32, start_rev: u32, end_rev: u32)
-> Result<()> -> Result<()> {
where
T: AsRef<[u8]>,
{
// Reserve the approximate size needed for the index - 28 bytes for each revision. // 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. // 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); 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 { 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)?; res.chain_err(|| ErrorKind::IndexCorrupted)?;
} }
Ok(()) Ok(())
} }
/// Helper method similar to `radixbuf::key::FixedKey::read`, but takes a revision number instead. /// Helper method similar to `radixbuf::key::FixedKey::read`, but takes a revision number instead.
fn rev_to_node<K: AsRef<[u8]>>(changelog: &K, rev: KeyId) -> rerrors::Result<&[u8]> { #[inline]
let buf = changelog.as_ref(); fn rev_to_node(changelog: &[u8], rev: KeyId) -> rerrors::Result<&[u8]> {
let buf = changelog;
let rev_usize: usize = rev.into(); let rev_usize: usize = rev.into();
let start_pos = rev_usize * 64 + 32; let start_pos = rev_usize * 64 + 32;
let end_pos = start_pos + 20; let end_pos = start_pos + 20;
if buf.len() < end_pos { if buf.len() < end_pos {
Err(rerrors::ErrorKind::InvalidKeyId(rev).into()) Err(rerrors::ErrorKind::InvalidKeyId(rev).into())
} else { } else {
Ok(&buf[start_pos..end_pos]) Ok(&buf[start_pos..end_pos])
} }

rust/radixbuf/src/key.rs

/// Keys with fixed 20 bytes length. /// Keys with fixed 20 bytes length.
pub struct FixedKey; pub struct FixedKey;
/// Keys with variant-length. Serialized as a VLQ-encoded length, followed by the actual key. /// Keys with variant-length. Serialized as a VLQ-encoded length, followed by the actual key.
pub struct VariantKey; pub struct VariantKey;
impl FixedKey { impl FixedKey {
#[inline] #[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]> {
jsgfUnsubmitted
Not Done

_like_

jsgf: _like_
let key_buf = key_buf.as_ref();
let start_pos: usize = key_id.into(); let start_pos: usize = key_id.into();
// LANG: Consider making 20 a type parameter once supported. // LANG: Consider making 20 a type parameter once supported.
let end_pos = start_pos + 20; let end_pos = start_pos + 20;
if key_buf.len() < end_pos { if key_buf.len() < end_pos {
Err(ErrorKind::InvalidKeyId(key_id).into()) Err(ErrorKind::InvalidKeyId(key_id).into())
} else { } else {
Ok(&key_buf[start_pos..end_pos]) Ok(&key_buf[start_pos..end_pos])
} }
let key_buf = key_buf.as_mut(); let key_buf = key_buf.as_mut();
key_buf[pos..pos + 20].copy_from_slice(key); key_buf[pos..pos + 20].copy_from_slice(key);
pos.into() pos.into()
} }
} }
impl VariantKey { impl VariantKey {
#[inline] #[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]> {
jsgfUnsubmitted
Not Done

_like_

jsgf: _like_
let key_buf = key_buf.as_ref(); let key_buf = key_buf.as_ref();
let mut reader = Cursor::new(key_buf); let mut reader = Cursor::new(key_buf);
reader.seek(SeekFrom::Start(key_id.into()))?; reader.seek(SeekFrom::Start(key_id.into()))?;
let key_len: usize = reader.read_vlq()?; let key_len: usize = reader.read_vlq()?;
let start_pos = reader.seek(SeekFrom::Current(0))? as usize; let start_pos = reader.seek(SeekFrom::Current(0))? as usize;
let end_pos = start_pos + key_len; let end_pos = start_pos + key_len;
if key_buf.len() < end_pos { if key_buf.len() < end_pos {

rust/radixbuf/src/radix.rs

} }
vec[pos] = value.to_le(); vec[pos] = value.to_le();
Ok(()) Ok(())
} }
} }
// Public APIs // 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<Box<[u8]>>),
// 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<Box<[u8]>>, &'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 {
jsgfUnsubmitted
Not Done

I don't really understand this. It looks like something that could be expressed with a function and closures. Unless you really want the early return from $block?

jsgf: I don't really understand this. It //looks// like something that could be expressed with a…
quarkAuthorUnsubmitted
Not Done

I'd be happy to rewrite if there is a cleaner way to get rid of this. What I want is not holding &radix_buf reference when running $block since $block may want &mut radix_buf.

quark: I'd be happy to rewrite if there is a cleaner way to get rid of this. What I want is not…
&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`. /// Lookup a given `Key`. Return an optional potentially matched `KeyId`.
/// The caller is responsible to check whether `KeyId` matches the given `Key` or not. /// The caller is responsible to check whether `KeyId` matches the given `Key` or not.
#[inline] #[inline]
pub fn radix_lookup_unchecked<R, K>(radix_buf: &R, offset: u32, key: &K) -> Result<Option<KeyId>> pub fn radix_lookup_unchecked<R, K>(radix_buf: &R, offset: u32, key: &K) -> Result<Option<KeyId>>
where where
R: AsRef<[u32]>, R: AsRef<[u32]>,
K: AsRef<[u8]>, K: AsRef<[u8]>,
{ {
let (key_id, _) = RadixOffset::new(offset).follow( let (key_id, _) = RadixOffset::new(offset).follow(
radix_buf, radix_buf,
key.to_base16iter(), key.to_base16iter(),
)?; )?;
Ok(key_id) 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`. /// Lookup a given `Key`. Return a verified `KeyId` or `None`.
#[cfg_attr(rustfmt, rustfmt_skip)]
#[inline] #[inline]
jsgfUnsubmitted
Not Done

This is too big to inline. The compiler can always inline it with LTO if it really wants to.

(Ditto other large functions with #[inline].)

jsgf: This is too big to inline. The compiler can always inline it with LTO if it really wants to.
quarkAuthorUnsubmitted
Not Done

Good idea.

quark: Good idea.
pub fn radix_lookup<R, K, KR, KA>( pub fn radix_lookup<R, K>(radix_buf: &R, offset: u32, key: &K, key_reader: &KeyReader)
radix_buf: &R, offset: u32, key: &K, key_reader: KR, key_reader_arg: &KA)
-> Result<Option<KeyId>> -> Result<Option<KeyId>>
where where
R: AsRef<[u32]>, R: AsRef<[u32]>,
K: AsRef<[u8]>, K: AsRef<[u8]>,
KR: Fn(&KA, KeyId) -> Result<&[u8]>,
{ {
let key_id = radix_lookup_unchecked(radix_buf, offset, key)?; let key_id = radix_lookup_unchecked(radix_buf, offset, key)?;
if let Some(id) = key_id { if let Some(id) = key_id {
let existing_key = key_reader(key_reader_arg, id).unwrap(); read_key!(key_reader, id, existing_key, radix_buf, {
if existing_key != key.as_ref() { if existing_key != key.as_ref() {
return Ok(None); return Ok(None);
} }
});
} }
Ok(key_id) Ok(key_id)
} }
/// Lookup a unique `KeyId` given a prefix of a binary base16 sequence. /// Lookup a unique `KeyId` given a prefix of a binary base16 sequence.
#[cfg_attr(rustfmt, rustfmt_skip)]
#[inline] #[inline]
pub fn radix_prefix_lookup<R, P, KR, KA>( pub fn radix_prefix_lookup<R, P>(radix_buf: &R, offset: u32, prefix: P, key_reader: &KeyReader)
radix_buf: &R, offset: u32, prefix: P, key_reader: KR, key_reader_arg: &KA)
-> Result<Option<KeyId>> -> Result<Option<KeyId>>
where where
R: AsRef<[u32]>, R: AsRef<[u32]>,
P: Iterator<Item = u8> + Clone, P: Iterator<Item = u8> + Clone,
KR: Fn(&KA, KeyId) -> Result<&[u8]>,
{ {
let root = RadixOffset::new(offset); let root = RadixOffset::new(offset);
let (key_id, (_radix, i, _b)) = root.follow(radix_buf, prefix.clone())?; let (key_id, (_radix, i, _b)) = root.follow(radix_buf, prefix.clone())?;
if let Some(id) = key_id { if let Some(id) = key_id {
let key = key_reader(key_reader_arg, id)?; read_key!(key_reader, id, key, radix_buf, {
let iter = key.to_base16iter(); let iter = key.to_base16iter();
let matched = iter.clone().skip(i).zip(prefix.clone().skip(i)).all(|(b1, b2)| b1 == b2); 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() { if !matched || iter.count() < prefix.count() {
return Ok(None); return Ok(None);
} }
});
} }
Ok(key_id) Ok(key_id)
} }
/// Insert a `Key` to `KeyId` relationship into the radix tree. /// Insert a `Key` to `KeyId` relationship into the radix tree.
#[cfg_attr(rustfmt, rustfmt_skip)]
#[inline] #[inline]
pub fn radix_insert<R, KR, KA>( pub fn radix_insert<R>(radix_buf: &mut R, offset: u32, key_id: KeyId, key_reader: &KeyReader)
radix_buf: &mut R, offset: u32, key_id: KeyId, key_reader: KR, key_reader_arg: &KA)
-> Result<()> -> Result<()>
where where
R: Resize<u32> + AsRef<[u32]> + AsMut<[u32]>, R: Resize<u32> + AsRef<[u32]> + AsMut<[u32]>,
KR: Fn(&KA, KeyId) -> Result<&[u8]>,
{ {
let new_key = key_reader(key_reader_arg, key_id)?; read_key!(key_reader, key_id, new_key, radix_buf, {
radix_insert_with_key( radix_insert_with_key(radix_buf, offset, key_id, &new_key, key_reader)
radix_buf, })
offset,
key_id,
&new_key,
key_reader,
key_reader_arg,
)
} }
#[cfg_attr(rustfmt, rustfmt_skip)] #[cfg_attr(rustfmt, rustfmt_skip)]
jsgfUnsubmitted
Not Done

Does rustfmt make a mess of this?

jsgf: Does rustfmt make a mess of this?
quarkAuthorUnsubmitted
Not Done

It folds line 235 to 237 into a single line exceeding 100 chars.

quark: It folds line 235 to 237 into a single line exceeding 100 chars.
#[inline] #[inline]
jsgfUnsubmitted
Not Done

WAY to big to #[inline].

jsgf: **WAY** to big to `#[inline]`.
pub fn radix_insert_with_key<R, K, KR, KA>( pub fn radix_insert_with_key<R, K>(
radix_buf: &mut R, offset: u32, id: KeyId, new_key: &K, key_reader: KR, key_reader_arg: &KA) radix_buf: &mut R, offset: u32, id: KeyId, new_key: &K, key_reader: &KeyReader)
-> Result<()> -> Result<()>
where where
R: Resize<u32> + AsRef<[u32]> + AsMut<[u32]>, R: Resize<u32> + AsRef<[u32]> + AsMut<[u32]>,
K: AsRef<[u8]>, K: AsRef<[u8]>,
KR: Fn(&KA, KeyId) -> Result<&[u8]>,
{ {
let new_key_id = id; let new_key_id = id;
let root = RadixOffset::new(offset); let root = RadixOffset::new(offset);
let (old_key_id, (radix, i, b)) = root.follow(radix_buf, new_key.to_base16iter())?; let (old_key_id, (radix, i, b)) = root.follow(radix_buf, new_key.to_base16iter())?;
match old_key_id { match old_key_id {
Some(old_key_id) => { Some(old_key_id) => {
// No need to re-insert a same key // No need to re-insert a same key
if old_key_id == new_key_id { if old_key_id == new_key_id {
return Ok(()); return Ok(());
} }
// Need to do a leaf split // 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 // Find common prefix starting from the next base16 integer
let mut common_len = 0; let mut common_len = 0;
let old_iter = old_key.to_base16iter().skip(i + 1); let old_iter = old_key.to_base16iter().skip(i + 1);
let new_iter = new_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) { for (b1, b2) in old_iter.zip(new_iter) {
if b1 == b2 { if b1 == b2 {
common_len += 1; common_len += 1;
} else { } else {
// Got a chain of radix nodes to write back // Got a chain of radix nodes to write back
let mut node = RadixOffset::create(radix_buf)?; let mut node = RadixOffset::create(radix_buf)?;
node.write_key_id(radix_buf, b1, old_key_id)?; node.write_key_id(radix_buf, b1, old_key_id)?;
node.write_key_id(radix_buf, b2, new_key_id)?; node.write_key_id(radix_buf, b2, new_key_id)?;
let new_iter = new_key.to_base16iter().skip(i + 1); let new_iter = new_key.to_base16iter().skip(i + 1);
for k in new_iter.take(common_len).rev() { for k in new_iter.take(common_len).rev() {
let new_node = RadixOffset::create(radix_buf)?; let new_node = RadixOffset::create(radix_buf)?;
new_node.write_radix(radix_buf, k, node)?; new_node.write_radix(radix_buf, k, node)?;
node = new_node; node = new_node;
} }
return radix.write_radix(radix_buf, b, node); return radix.write_radix(radix_buf, b, node);
} }
} }
// new_key is a prefix of old_key, or vice-versa // new_key is a prefix of old_key, or vice-versa
if old_key.len() > new_key.as_ref().len() { if old_key.len() > new_key.as_ref().len() {
Err(ErrorKind::PrefixConflict(new_key_id, old_key_id).into()) bail!(ErrorKind::PrefixConflict(new_key_id, old_key_id))
} else { } else {
Err(ErrorKind::PrefixConflict(old_key_id, new_key_id).into()) bail!(ErrorKind::PrefixConflict(old_key_id, new_key_id))
} }
})
} }
None => radix.write_key_id(radix_buf, b, new_key_id), None => radix.write_key_id(radix_buf, b, new_key_id),
} }
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::collections::HashSet; use std::collections::HashSet;
use std::mem::transmute; use std::mem::transmute;
use key::{VariantKey, FixedKey}; use key::{VariantKey, FixedKey};
use rand::{ChaChaRng, Rng}; use rand::{ChaChaRng, Rng};
use test::Bencher; use test::Bencher;
#[test] #[test]
fn test_errors() { fn test_errors() {
let mut key_buf = vec![0u8; 10]; let mut key_buf = vec![0u8; 10];
let mut radix_buf = vec![0u32; 15]; let mut radix_buf = vec![0u32; 15];
{
// KeyId exceeds format limit // KeyId exceeds format limit
let key = [0u8; 20]; let key = [0u8; 20];
let key_id = (1u32 << 31).into(); let key_id = (1u32 << 31).into();
let r = radix_insert_with_key(&mut radix_buf, 0, key_id, &key, FixedKey::read, &key_buf); 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"); assert_eq!(r.unwrap_err().description(), "offset overflow");
// KeyId exceeds key buffer length // KeyId exceeds key buffer length
let key_id = 30u32.into(); let key_id = 30u32.into();
let r = radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id, &key_reader);
assert_eq!(r.unwrap_err().description(), "invalid key id"); assert_eq!(r.unwrap_err().description(), "invalid key id");
let r = radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id, &key_reader);
let t = format!("{}", r.unwrap_err()); let t = format!("{}", r.unwrap_err());
assert_eq!(t, "KeyId(30) cannot be resolved"); assert_eq!(t, "KeyId(30) cannot be resolved");
// Radix root node offset exceeds radix buffer length // 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); 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"); assert_eq!(format!("{}", r.unwrap_err()), "offset 16 is out of range");
// Radix node offset out of range during a lookup // Radix node offset out of range during a lookup
let prefix = [0xf].iter().cloned(); let prefix = [0xf].iter().cloned();
let r = radix_prefix_lookup(&radix_buf, 0, prefix, FixedKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader);
assert_eq!(format!("{}", r.unwrap_err()), "offset 15 is out of range"); assert_eq!(format!("{}", r.unwrap_err()), "offset 15 is out of range");
// Base16 sequence overflow // Base16 sequence overflow
let prefix = [21].iter().cloned(); let prefix = [21].iter().cloned();
let r = radix_prefix_lookup(&radix_buf, 0, prefix.clone(), FixedKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, prefix.clone(), &key_reader);
assert_eq!(r.unwrap_err().description(), "invalid base16 value"); assert_eq!(r.unwrap_err().description(), "invalid base16 value");
let r = radix_prefix_lookup(&radix_buf, 0, prefix, FixedKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader);
assert_eq!(format!("{}", r.unwrap_err()), "21 is not a base16 value"); 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_id1 = VariantKey::append(&mut key_buf, &b"ab");
let key_id2 = 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 // But not okay if `KeyId` are different
let r = radix_insert(&mut radix_buf, 0, key_id2, VariantKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id2, &key_reader);
assert_eq!(r.unwrap_err().description(), "key prefix conflict"); 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_id4 = VariantKey::append(&mut key_buf, &b"a");
let key_id5 = VariantKey::append(&mut key_buf, &b"abc"); let key_id5 = VariantKey::append(&mut key_buf, &b"abc");
let r = radix_insert(&mut radix_buf, 0, key_id4, VariantKey::read, &key_buf); {
// 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!( assert_eq!(
format!("{}", r.unwrap_err()), format!("{}", r.unwrap_err()),
format!("{:?} cannot be a prefix of {:?}", key_id4, key_id1) format!("{:?} cannot be a prefix of {:?}", key_id4, key_id1)
); );
let r = radix_insert(&mut radix_buf, 0, key_id5, VariantKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id5, &key_reader);
assert_eq!( assert_eq!(
format!("{}", r.unwrap_err()), format!("{}", r.unwrap_err()),
format!("{:?} cannot be a prefix of {:?}", key_id1, key_id5) format!("{:?} cannot be a prefix of {:?}", key_id1, key_id5)
); );
}
{
// Enforce a leaf split of key_id1 // Enforce a leaf split of key_id1
let key_id3 = VariantKey::append(&mut key_buf, &b"ac"); let key_id3 = VariantKey::append(&mut key_buf, &b"ac");
radix_insert(&mut radix_buf, 0, key_id3, VariantKey::read, &key_buf).expect("insert"); 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 // Still impossible to cause key prefix conflicts
let r = radix_insert(&mut radix_buf, 0, key_id4, VariantKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id4, &key_reader);
assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); assert_eq!(r.unwrap_err().description(), "ambiguous prefix");
let r = radix_insert(&mut radix_buf, 0, key_id5, VariantKey::read, &key_buf); let r = radix_insert(&mut radix_buf, 0, key_id5, &key_reader);
assert_eq!(r.unwrap_err().description(), "key prefix conflict"); assert_eq!(r.unwrap_err().description(), "key prefix conflict");
} }
}
#[test] #[test]
fn test_prefix_lookup() { fn test_prefix_lookup() {
let mut key_buf: Vec<u8> = vec![]; let mut key_buf: Vec<u8> = vec![];
let mut radix_buf = vec![0u32; 16]; let mut radix_buf = vec![0u32; 16];
let query = b"01abc".to_base16iter(); let query = b"01abc".to_base16iter();
// With a single key // With a single key
let key1 = b"01ab"; let key1 = b"01ab";
let key1_id = VariantKey::append(&mut key_buf, &key1); let key1_id = VariantKey::append(&mut key_buf, &key1);
radix_insert(&mut radix_buf, 0, key1_id, VariantKey::read, &key_buf).expect("insert"); {
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() { for i in 0..query.len() {
let prefix = query.clone().take(i); let prefix = query.clone().take(i);
let r = radix_prefix_lookup(&radix_buf, 0, prefix, VariantKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader);
if i == 0 { if i == 0 {
// This is sub-optimal. But see the NOTE in RadixOffset::follow. // This is sub-optimal. But see the NOTE in RadixOffset::follow.
assert_eq!(r.unwrap_err().description(), "ambiguous prefix"); assert_eq!(r.unwrap_err().description(), "ambiguous prefix");
} else if i <= key1.len() * 2 { } else if i <= key1.len() * 2 {
assert_eq!(r.unwrap(), Some(key1_id)); assert_eq!(r.unwrap(), Some(key1_id));
} else { } else {
assert_eq!(r.unwrap(), None); assert_eq!(r.unwrap(), None);
} }
} }
}
// With another key // With another key
let key2 = b"01bbc"; let key2 = b"01bbc";
let key2_id = VariantKey::append(&mut key_buf, &key2); let key2_id = VariantKey::append(&mut key_buf, &key2);
radix_insert(&mut radix_buf, 0, key2_id, VariantKey::read, &key_buf).expect("insert"); {
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() { for i in 0..query.len() {
let prefix = query.clone().take(i); let prefix = query.clone().take(i);
let r = radix_prefix_lookup(&radix_buf, 0, prefix, VariantKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, prefix, &key_reader);
if i <= 5 { if i <= 5 {
assert_eq!(r.unwrap_err().description(), "ambiguous prefix") assert_eq!(r.unwrap_err().description(), "ambiguous prefix")
} else if i <= key1.len() * 2 { } else if i <= key1.len() * 2 {
assert_eq!(r.unwrap(), Some(key1_id)); assert_eq!(r.unwrap(), Some(key1_id));
} else { } else {
assert_eq!(r.unwrap(), None); assert_eq!(r.unwrap(), None);
} }
} }
let query = b"1".to_base16iter(); let query = b"1".to_base16iter();
let r = radix_prefix_lookup(&radix_buf, 0, query, VariantKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, query, &key_reader);
assert_eq!(r.unwrap(), None); assert_eq!(r.unwrap(), None);
let query = b"01b".to_base16iter(); let query = b"01b".to_base16iter();
let r = radix_prefix_lookup(&radix_buf, 0, query, VariantKey::read, &key_buf); let r = radix_prefix_lookup(&radix_buf, 0, query, &key_reader);
assert_eq!(r.unwrap(), Some(key2_id)); 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())
);
// 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()));
// 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! { quickcheck! {
fn test_compare_with_stdset_sparse(std_set: HashSet<u64>) -> bool { fn test_compare_with_stdset_sparse(std_set: HashSet<u64>) -> bool {
let std_set: HashSet<[u8; 10]> = std_set.iter().map(|&x| { let std_set: HashSet<[u8; 10]> = std_set.iter().map(|&x| {
let mut buf = [0u8; 10]; let mut buf = [0u8; 10];
let slice: [u8; 8] = unsafe { transmute(x) }; let slice: [u8; 8] = unsafe { transmute(x) };
buf[0..8].copy_from_slice(&slice); buf[0..8].copy_from_slice(&slice);
buf buf
// Compare with `HashSet`. // Compare with `HashSet`.
fn check_with_stdset(std_set: HashSet<[u8; 10]>) -> bool { fn check_with_stdset(std_set: HashSet<[u8; 10]>) -> bool {
let mut key_buf = Vec::<u8>::with_capacity(std_set.len() * 11); let mut key_buf = Vec::<u8>::with_capacity(std_set.len() * 11);
let mut radix_buf = vec![0u32; 16]; let mut radix_buf = vec![0u32; 16];
// Insert to radix tree // Insert to radix tree
for key in &std_set { for key in &std_set {
let key_id = VariantKey::append(&mut key_buf, key); 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 // Test key existence
let key_reader = KeyReader::KeyBufferZeroCopy(VariantKey::read, &key_buf);
std_set.iter().all(|key| { 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() r.unwrap().is_some()
}) })
} }
const COUNT: usize = 51200; const COUNT: usize = 51200;
#[bench] #[bench]
fn bench_insert(b: &mut Bencher) { fn bench_insert(b: &mut Bencher) {
radix_lookup_unchecked(&radix_buf, 0, &key).expect("lookup"); radix_lookup_unchecked(&radix_buf, 0, &key).expect("lookup");
}) })
} }
#[bench] #[bench]
fn bench_lookups(b: &mut Bencher) { fn bench_lookups(b: &mut Bencher) {
let key_buf = randomized_key_buf(COUNT); let key_buf = randomized_key_buf(COUNT);
let radix_buf = batch_insert_radix_buf(&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 { b.iter(|| for i in 0..COUNT {
let key_id = (i as u32 * 20).into(); let key_id = (i as u32 * 20).into();
let key = FixedKey::read(&key_buf, key_id).unwrap(); 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");
}) })
} }
fn randomized_key_buf(count: usize) -> Vec<u8> { fn randomized_key_buf(count: usize) -> Vec<u8> {
let mut key_buf = vec![0u8; count * 20]; let mut key_buf = vec![0u8; count * 20];
// Using an unseeded rng so benchmarks are more stable across multiple runs // Using an unseeded rng so benchmarks are more stable across multiple runs
ChaChaRng::new_unseeded().fill_bytes(key_buf.as_mut()); ChaChaRng::new_unseeded().fill_bytes(key_buf.as_mut());
key_buf key_buf
} }
fn batch_insert_radix_buf(key_buf: &Vec<u8>, count: usize) -> Vec<u32> { fn batch_insert_radix_buf(key_buf: &Vec<u8>, count: usize) -> Vec<u32> {
let mut radix_buf = vec![0u32; 16]; let mut radix_buf = vec![0u32; 16];
let key_reader = KeyReader::KeyBufferZeroCopy(FixedKey::read, key_buf);
for i in 0..count { for i in 0..count {
let key_id: KeyId = ((i * 20) as u32).into(); 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 radix_buf
} }
} }