use crate::bisect_multi::{bisect_multi_bytes, BisectStatus};
use std::collections::HashMap;
pub const SIGNATURE: &[u8] = b"Bazaar Graph Index 1\n";
pub const OPTION_NODE_REFS: &[u8] = b"node_ref_lists=";
pub const OPTION_KEY_ELEMENTS: &[u8] = b"key_elements=";
pub const OPTION_LEN: &[u8] = b"len=";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexNode {
pub key: Vec<Vec<u8>>,
pub absent: bool,
pub references: Vec<Vec<Vec<Vec<u8>>>>,
pub value: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IndexError {
UnknownReference(Vec<Vec<u8>>),
LengthMismatch { expected: usize, actual: usize },
BadSignature,
BadOptions,
BadLineData,
BadReferenceOffset(Vec<u8>),
BadKey(IndexKey),
BadValue(String),
DuplicateKey(IndexKey),
BadIndexData,
Other(String),
}
impl std::fmt::Display for IndexError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
IndexError::UnknownReference(k) => {
write!(f, "reference to unknown key: {:?}", k)
}
IndexError::LengthMismatch { expected, actual } => write!(
f,
"mismatched output length and expected length: {} {}",
actual, expected
),
IndexError::BadSignature => write!(f, "bad index format signature"),
IndexError::BadOptions => write!(f, "bad index options"),
IndexError::BadLineData => write!(f, "bad index line data"),
IndexError::BadReferenceOffset(s) => {
write!(f, "bad reference offset: {:?}", s)
}
IndexError::BadKey(k) => write!(f, "bad index key: {:?}", k),
IndexError::BadValue(msg) => write!(f, "bad index value: {}", msg),
IndexError::DuplicateKey(k) => {
write!(f, "duplicate index key: {:?}", k)
}
IndexError::BadIndexData => write!(f, "bad index data"),
IndexError::Other(msg) => write!(f, "{}", msg),
}
}
}
impl std::error::Error for IndexError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexHeader {
pub node_ref_lists: usize,
pub key_length: usize,
pub key_count: usize,
pub header_end: usize,
}
pub fn parse_header(data: &[u8]) -> Result<IndexHeader, IndexError> {
if !data.starts_with(SIGNATURE) {
return Err(IndexError::BadSignature);
}
let after_sig = &data[SIGNATURE.len()..];
let mut option_lines: [&[u8]; 3] = [b"", b"", b""];
let mut offset = 0usize;
for slot in option_lines.iter_mut() {
let nl = after_sig[offset..]
.iter()
.position(|&b| b == b'\n')
.ok_or(IndexError::BadOptions)?;
*slot = &after_sig[offset..offset + nl];
offset += nl + 1;
}
let node_ref_lists = parse_option(option_lines[0], OPTION_NODE_REFS)?;
let key_length = parse_option(option_lines[1], OPTION_KEY_ELEMENTS)?;
let key_count = parse_option(option_lines[2], OPTION_LEN)?;
let header_end =
SIGNATURE.len() + option_lines[0].len() + option_lines[1].len() + option_lines[2].len() + 3;
Ok(IndexHeader {
node_ref_lists,
key_length,
key_count,
header_end,
})
}
fn parse_option(line: &[u8], prefix: &[u8]) -> Result<usize, IndexError> {
if !line.starts_with(prefix) {
return Err(IndexError::BadOptions);
}
std::str::from_utf8(&line[prefix.len()..])
.ok()
.and_then(|s| s.parse::<usize>().ok())
.ok_or(IndexError::BadOptions)
}
pub type IndexKey = Vec<Vec<u8>>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RawNode {
pub key: IndexKey,
pub absent: bool,
pub ref_offsets: Vec<Vec<u64>>,
pub value: Vec<u8>,
}
pub type ParsedNode = (IndexKey, Vec<u8>, Vec<Vec<u64>>);
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct ParsedLines {
pub first_key: Option<IndexKey>,
pub last_key: Option<IndexKey>,
pub nodes: Vec<ParsedNode>,
pub keys_by_offset: Vec<(u64, RawNode)>,
pub trailers: usize,
}
pub fn parse_lines(
lines: &[&[u8]],
start_pos: u64,
key_length: usize,
) -> Result<ParsedLines, IndexError> {
let expected_elements = 3 + key_length;
let mut out = ParsedLines::default();
let mut pos = start_pos;
for line in lines {
if line.is_empty() {
out.trailers += 1;
continue;
}
let elements: Vec<&[u8]> = line.split(|&b| b == b'\x00').collect();
if elements.len() != expected_elements {
return Err(IndexError::BadLineData);
}
let key: Vec<Vec<u8>> = elements[..key_length].iter().map(|e| e.to_vec()).collect();
if out.first_key.is_none() {
out.first_key = Some(key.clone());
}
out.last_key = Some(key.clone());
let absent_field = elements[elements.len() - 3];
let references_field = elements[elements.len() - 2];
let value_field = elements[elements.len() - 1];
let absent = absent_field == b"a";
let mut ref_lists: Vec<Vec<u64>> = Vec::new();
for ref_string in references_field.split(|&b| b == b'\t') {
let mut list = Vec::new();
for piece in ref_string.split(|&b| b == b'\r') {
if piece.is_empty() {
continue;
}
let parsed = std::str::from_utf8(piece)
.ok()
.and_then(|s| s.parse::<u64>().ok())
.ok_or_else(|| IndexError::BadReferenceOffset(piece.to_vec()))?;
list.push(parsed);
}
ref_lists.push(list);
}
let raw = RawNode {
key: key.clone(),
absent,
ref_offsets: ref_lists.clone(),
value: value_field.to_vec(),
};
out.keys_by_offset.push((pos, raw));
pos += line.len() as u64 + 1;
if absent {
continue;
}
out.nodes.push((key, value_field.to_vec(), ref_lists));
}
Ok(out)
}
pub fn serialize_graph_index(
nodes: &[IndexNode],
reference_lists: usize,
key_elements: usize,
) -> Result<Vec<u8>, IndexError> {
let mut sorted: Vec<&IndexNode> = nodes.iter().collect();
sorted.sort_by(|a, b| a.key.cmp(&b.key));
let mut header = Vec::new();
header.extend_from_slice(SIGNATURE);
header.extend_from_slice(OPTION_NODE_REFS);
header.extend_from_slice(reference_lists.to_string().as_bytes());
header.push(b'\n');
header.extend_from_slice(OPTION_KEY_ELEMENTS);
header.extend_from_slice(key_elements.to_string().as_bytes());
header.push(b'\n');
header.extend_from_slice(OPTION_LEN);
let key_count = sorted.iter().filter(|n| !n.absent).count();
header.extend_from_slice(key_count.to_string().as_bytes());
header.push(b'\n');
let prefix_length = header.len();
let (digits, addresses, expected_bytes) = if reference_lists > 0 {
let mut key_offset_info: Vec<(usize, usize)> = Vec::with_capacity(sorted.len());
let mut non_ref_bytes = prefix_length;
let mut total_references = 0usize;
for (idx, node) in sorted.iter().enumerate() {
key_offset_info.push((idx, non_ref_bytes));
for element in &node.key {
non_ref_bytes += element.len();
}
if key_elements > 1 {
non_ref_bytes += key_elements - 1;
}
non_ref_bytes += node.value.len() + 3 + 1;
if node.absent {
non_ref_bytes += 1;
} else {
non_ref_bytes += reference_lists - 1;
for ref_list in &node.references {
total_references += ref_list.len();
if !ref_list.is_empty() {
non_ref_bytes += ref_list.len() - 1;
}
}
}
}
let mut digits = 1usize;
let mut possible_total = non_ref_bytes + total_references * digits;
while 10usize.pow(digits as u32) < possible_total {
digits += 1;
possible_total = non_ref_bytes + total_references * digits;
}
let expected = possible_total + 1;
let mut addresses: HashMap<Vec<Vec<u8>>, usize> = HashMap::new();
let mut total_refs_so_far = 0usize;
for (idx, non_ref_so_far) in &key_offset_info {
let node = sorted[*idx];
addresses.insert(
node.key.clone(),
non_ref_so_far + total_refs_so_far * digits,
);
if !node.absent {
for ref_list in &node.references {
total_refs_so_far += ref_list.len();
}
}
}
(digits, addresses, Some(expected))
} else {
(0, HashMap::new(), None)
};
let mut out = header;
for node in &sorted {
let mut flattened = Vec::new();
for (i, ref_list) in node.references.iter().enumerate() {
if i > 0 {
flattened.push(b'\t');
}
for (j, reference) in ref_list.iter().enumerate() {
if j > 0 {
flattened.push(b'\r');
}
let addr = addresses
.get(reference)
.ok_or_else(|| IndexError::UnknownReference(reference.clone()))?;
let formatted = format!("{:0>width$}", addr, width = digits);
flattened.extend_from_slice(formatted.as_bytes());
}
}
for (i, element) in node.key.iter().enumerate() {
if i > 0 {
out.push(b'\x00');
}
out.extend_from_slice(element);
}
out.push(b'\x00');
if node.absent {
out.push(b'a');
}
out.push(b'\x00');
out.extend_from_slice(&flattened);
out.push(b'\x00');
out.extend_from_slice(&node.value);
out.push(b'\n');
}
out.push(b'\n');
if let Some(expected) = expected_bytes {
if out.len() != expected {
return Err(IndexError::LengthMismatch {
expected,
actual: out.len(),
});
}
}
Ok(out)
}
pub trait IndexTransport {
fn get_bytes(&self, path: &str) -> Result<Vec<u8>, IndexError>;
fn abspath(&self, path: &str) -> String {
path.to_string()
}
fn recommended_page_size(&self) -> u64 {
64 * 1024
}
fn readv(
&self,
path: &str,
ranges: &[(u64, u64)],
_adjust_for_latency: bool,
_upper_limit: u64,
) -> Result<Vec<(u64, Vec<u8>)>, IndexError> {
let data = self.get_bytes(path)?;
let mut out = Vec::with_capacity(ranges.len());
for &(offset, length) in ranges {
let end = (offset + length) as usize;
if end > data.len() {
return Err(IndexError::Other(format!(
"readv past end of {} at offset {}+{}",
path, offset, length
)));
}
out.push((offset, data[offset as usize..end].to_vec()));
}
Ok(out)
}
}
impl IndexError {
fn missing_trailer() -> Self {
IndexError::Other("BadIndexData: missing trailer".to_string())
}
}
pub type RefList = Vec<IndexKey>;
pub type RefLists = Vec<RefList>;
pub type NodeBody = (Vec<u8>, RefLists);
pub type IndexEntry = (IndexKey, Vec<u8>, RefLists);
pub type KeyPrefix = Vec<Option<Vec<u8>>>;
#[inline]
pub fn key_matches_prefix(prefix: &KeyPrefix, key: &[Vec<u8>]) -> bool {
if prefix.len() != key.len() {
return false;
}
prefix.iter().zip(key.iter()).all(|(p, k)| match p {
Some(p) => p == k,
None => true,
})
}
#[inline]
pub fn key_matches_any_prefix(prefixes: &[KeyPrefix], key: &[Vec<u8>]) -> bool {
prefixes.iter().any(|p| key_matches_prefix(p, key))
}
#[inline]
fn is_key_disallowed(b: u8) -> bool {
matches!(b, b'\t' | b'\n' | 0x0b | 0x0c | b'\r' | 0 | b' ')
}
pub fn key_is_valid(key: &[Vec<u8>], key_length: usize) -> bool {
if key.len() != key_length {
return false;
}
key.iter()
.all(|element| !element.is_empty() && !element.iter().any(|&b| is_key_disallowed(b)))
}
pub fn value_is_valid(value: &[u8]) -> bool {
!value.iter().any(|&b| b == b'\n' || b == 0)
}
#[derive(Debug, Default, Clone)]
pub struct ParsedRangeMap {
byte_map: Vec<(u64, u64)>,
key_map: Vec<(Option<IndexKey>, Option<IndexKey>)>,
}
impl ParsedRangeMap {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.byte_map.len()
}
pub fn is_empty(&self) -> bool {
self.byte_map.is_empty()
}
pub fn byte_range(&self, index: usize) -> Option<(u64, u64)> {
self.byte_map.get(index).copied()
}
pub fn key_range(&self, index: usize) -> Option<(Option<IndexKey>, Option<IndexKey>)> {
self.key_map.get(index).cloned()
}
pub fn byte_index(&self, offset: u64) -> isize {
self.byte_map
.iter()
.rposition(|r| r.0 <= offset)
.map(|i| i as isize)
.unwrap_or(-1)
}
pub fn key_index(&self, key: &Option<IndexKey>) -> isize {
self.key_map
.iter()
.rposition(|r| r.0 <= *key)
.map(|i| i as isize)
.unwrap_or(-1)
}
pub fn is_parsed(&self, offset: u64) -> bool {
let index = self.byte_index(offset);
if index < 0 {
return false;
}
let (start, end) = self.byte_map[index as usize];
offset >= start && offset < end
}
pub fn mark_parsed(
&mut self,
start: u64,
start_key: Option<IndexKey>,
end: u64,
end_key: Option<IndexKey>,
) {
let index = self.byte_index(start);
let new_value = (start, end);
let new_key = (start_key, end_key);
if index < 0 {
self.byte_map.insert(0, new_value);
self.key_map.insert(0, new_key);
return;
}
let index = index as usize;
let next = index + 1;
if next < self.byte_map.len()
&& self.byte_map[index].1 == start
&& self.byte_map[next].0 == end
{
self.byte_map[index] = (self.byte_map[index].0, self.byte_map[next].1);
self.key_map[index] = (self.key_map[index].0.clone(), self.key_map[next].1.clone());
self.byte_map.remove(next);
self.key_map.remove(next);
} else if self.byte_map[index].1 == start {
self.byte_map[index] = (self.byte_map[index].0, end);
self.key_map[index] = (self.key_map[index].0.clone(), new_key.1);
} else if next < self.byte_map.len() && self.byte_map[next].0 == end {
self.byte_map[next] = (start, self.byte_map[next].1);
self.key_map[next] = (new_key.0, self.key_map[next].1.clone());
} else {
self.byte_map.insert(next, new_value);
self.key_map.insert(next, new_key);
}
}
}
pub fn parse_full(data: &[u8]) -> Result<(IndexHeader, HashMap<IndexKey, NodeBody>), IndexError> {
let header = parse_header(data)?;
let body = &data[header.header_end..];
let mut segments: Vec<&[u8]> = body.split(|&b| b == b'\n').collect();
segments.pop();
let parsed = parse_lines(&segments, header.header_end as u64, header.key_length)?;
if parsed.trailers != 1 {
return Err(IndexError::missing_trailer());
}
let mut offset_to_key: HashMap<u64, IndexKey> = HashMap::new();
for (offset, raw_node) in &parsed.keys_by_offset {
offset_to_key.insert(*offset, raw_node.key.clone());
}
let mut nodes: HashMap<IndexKey, NodeBody> = HashMap::new();
let node_ref_lists = header.node_ref_lists;
for (_, raw_node) in parsed.keys_by_offset.into_iter() {
if raw_node.absent {
continue;
}
let resolved = if node_ref_lists == 0 {
Vec::new()
} else {
let mut out: Vec<Vec<IndexKey>> = Vec::with_capacity(node_ref_lists);
for ref_list in &raw_node.ref_offsets {
let mut list: Vec<IndexKey> = Vec::with_capacity(ref_list.len());
for off in ref_list {
let k = offset_to_key.get(off).ok_or_else(|| {
IndexError::Other(format!("unresolved reference offset {}", off))
})?;
list.push(k.clone());
}
out.push(list);
}
out
};
nodes.insert(raw_node.key, (raw_node.value, resolved));
}
Ok((header, nodes))
}
pub type BisectNodeBody = (Vec<u8>, Vec<Vec<u64>>);
pub struct GraphIndex<T: IndexTransport> {
transport: T,
name: String,
base_offset: u64,
size: Option<u64>,
nodes: Option<HashMap<IndexKey, NodeBody>>,
header: Option<IndexHeader>,
bisect_nodes: Option<HashMap<IndexKey, BisectNodeBody>>,
keys_by_offset: HashMap<u64, RawNode>,
range_map: ParsedRangeMap,
bytes_read: u64,
}
impl<T: IndexTransport> GraphIndex<T> {
pub fn new(transport: T, name: impl Into<String>, base_offset: u64) -> Self {
Self::with_size(transport, name, base_offset, None)
}
pub fn with_size(
transport: T,
name: impl Into<String>,
base_offset: u64,
size: Option<u64>,
) -> Self {
Self {
transport,
name: name.into(),
base_offset,
size,
nodes: None,
header: None,
bisect_nodes: None,
keys_by_offset: HashMap::new(),
range_map: ParsedRangeMap::new(),
bytes_read: 0,
}
}
pub fn size(&self) -> Option<u64> {
self.size
}
pub fn bytes_read(&self) -> u64 {
self.bytes_read
}
pub fn is_buffered_already(&self) -> bool {
self.nodes.is_some()
}
pub fn header(&self) -> Option<&IndexHeader> {
self.header.as_ref()
}
pub fn nodes_iter(&self) -> impl Iterator<Item = (&IndexKey, &NodeBody)> {
self.nodes.iter().flat_map(|m| m.iter())
}
pub fn ensure_header_parsed(&mut self) -> Result<(), IndexError> {
if self.header.is_some() {
return Ok(());
}
self.read_and_parse(vec![(0, 200)])?;
Ok(())
}
pub fn read_and_parse_for_test(
&mut self,
readv_ranges: Vec<(u64, u64)>,
) -> Result<(), IndexError> {
self.read_and_parse(readv_ranges)
}
pub fn key_count_or_zero(&self) -> usize {
self.header.as_ref().map(|h| h.key_count).unwrap_or(0)
}
pub fn range_map(&self) -> &ParsedRangeMap {
&self.range_map
}
pub fn bisect_nodes(&self) -> Option<&HashMap<IndexKey, BisectNodeBody>> {
self.bisect_nodes.as_ref()
}
pub fn keys_by_offset(&self) -> &HashMap<u64, RawNode> {
&self.keys_by_offset
}
pub fn buffer_all(&mut self) -> Result<(), IndexError> {
if self.nodes.is_some() {
return Ok(());
}
let raw = self.transport.get_bytes(&self.name)?;
let data = if self.base_offset == 0 {
raw
} else {
raw[self.base_offset as usize..].to_vec()
};
let (header, nodes) = parse_full(&data)?;
self.nodes = Some(nodes);
self.header = Some(header);
Ok(())
}
pub fn key_count(&mut self) -> Result<usize, IndexError> {
if let Some(h) = &self.header {
return Ok(h.key_count);
}
if self.size.is_some() {
self.ensure_header_parsed()?;
} else {
self.buffer_all()?;
}
Ok(self
.header
.as_ref()
.expect("header set by ensure_header_parsed/buffer_all")
.key_count)
}
pub fn node_ref_lists(&mut self) -> Result<usize, IndexError> {
if let Some(h) = &self.header {
return Ok(h.node_ref_lists);
}
if self.size.is_some() {
self.ensure_header_parsed()?;
} else {
self.buffer_all()?;
}
Ok(self.header.as_ref().expect("header set").node_ref_lists)
}
pub fn key_length(&mut self) -> Result<usize, IndexError> {
if let Some(h) = &self.header {
return Ok(h.key_length);
}
if self.size.is_some() {
self.ensure_header_parsed()?;
} else {
self.buffer_all()?;
}
Ok(self.header.as_ref().expect("header set").key_length)
}
pub fn iter_all_entries(&mut self) -> Result<Vec<IndexEntry>, IndexError> {
self.buffer_all()?;
let nodes = self.nodes.as_ref().expect("buffer_all populated nodes");
Ok(nodes
.iter()
.map(|(k, (v, r))| (k.clone(), v.clone(), r.clone()))
.collect())
}
pub fn iter_entries(&mut self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError> {
if keys.is_empty() {
return Ok(Vec::new());
}
let mut deduped: Vec<IndexKey> = Vec::with_capacity(keys.len());
let mut seen: std::collections::HashSet<&IndexKey> = std::collections::HashSet::new();
for k in keys {
if seen.insert(k) {
deduped.push(k.clone());
}
}
if !self.should_bisect_for(deduped.len())? {
self.buffer_all()?;
let nodes = self.nodes.as_ref().expect("buffer_all populated nodes");
let mut out = Vec::with_capacity(deduped.len());
for k in &deduped {
if let Some((v, r)) = nodes.get(k) {
out.push((k.clone(), v.clone(), r.clone()));
}
}
return Ok(out);
}
self.iter_entries_via_bisect(deduped)
}
fn should_bisect_for(&mut self, requested_count: usize) -> Result<bool, IndexError> {
if self.size().is_none() || self.is_buffered_already() {
return Ok(false);
}
if self.key_count_or_zero() == 0 {
self.ensure_header_parsed()?;
}
if self.is_buffered_already() {
return Ok(false);
}
Ok(requested_count * 20 <= self.key_count_or_zero())
}
fn iter_entries_via_bisect(
&mut self,
keys: Vec<IndexKey>,
) -> Result<Vec<IndexEntry>, IndexError> {
let size = self.size().unwrap_or(0) as usize;
let stashed_err: std::cell::RefCell<Option<IndexError>> = std::cell::RefCell::new(None);
let self_cell = std::cell::RefCell::new(self);
let found: Vec<(IndexKey, (Vec<u8>, Vec<Vec<IndexKey>>))> = bisect_multi_bytes(
|probes| {
if stashed_err.borrow().is_some() {
return probes
.into_iter()
.map(|p| (p, BisectStatus::Absent))
.collect();
}
let probe_input: Vec<(u64, IndexKey)> = probes
.iter()
.map(|(loc, key)| (*loc as u64, key.clone()))
.collect();
let results = match self_cell
.borrow_mut()
.lookup_keys_via_location(&probe_input)
{
Ok(rs) => rs,
Err(e) => {
*stashed_err.borrow_mut() = Some(e);
return probes
.into_iter()
.map(|p| (p, BisectStatus::Absent))
.collect();
}
};
probes
.into_iter()
.zip(results)
.map(|(probe, (_, lookup))| {
let status = match lookup {
LookupResult::Missing => BisectStatus::Absent,
LookupResult::Direction(d) if d < 0 => BisectStatus::Earlier,
LookupResult::Direction(_) => BisectStatus::Later,
LookupResult::Found { value, refs } => {
BisectStatus::Found((value, refs))
}
};
(probe, status)
})
.collect()
},
size,
keys,
);
if let Some(e) = stashed_err.into_inner() {
return Err(e);
}
Ok(found
.into_iter()
.map(|(key, (value, refs))| (key, value, refs))
.collect())
}
pub fn iter_entries_prefix(
&mut self,
prefixes: &[KeyPrefix],
) -> Result<Vec<IndexEntry>, IndexError> {
self.buffer_all()?;
let key_length = self.header.as_ref().expect("header").key_length;
for p in prefixes {
if p.len() != key_length || !matches!(p.first(), Some(Some(_))) {
return Err(IndexError::BadKey(
p.iter().map(|e| e.clone().unwrap_or_default()).collect(),
));
}
}
let nodes = self.nodes.as_ref().expect("buffer_all populated nodes");
if key_length == 1 {
return self.iter_entries(
&prefixes
.iter()
.map(|p| {
p.iter()
.map(|e| e.clone().expect("validated above"))
.collect::<IndexKey>()
})
.collect::<Vec<_>>(),
);
}
let mut out = Vec::new();
let mut emitted: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
for prefix in prefixes {
for (k, (v, r)) in nodes.iter() {
if k.len() != key_length {
continue;
}
let matches = prefix
.iter()
.zip(k.iter())
.all(|(p_elem, k_elem)| match p_elem {
Some(p) => p == k_elem,
None => true,
});
if matches && emitted.insert(k.clone()) {
out.push((k.clone(), v.clone(), r.clone()));
}
}
}
Ok(out)
}
pub fn external_references(
&mut self,
ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
self.buffer_all()?;
let header = self.header.as_ref().expect("header");
if ref_list_num + 1 > header.node_ref_lists {
return Err(IndexError::Other(format!(
"No ref list {}, index has {} ref lists",
ref_list_num, header.node_ref_lists
)));
}
let nodes = self.nodes.as_ref().expect("nodes");
let mut refs = std::collections::HashSet::new();
for (_k, (_v, ref_lists)) in nodes.iter() {
let list = &ref_lists[ref_list_num];
for r in list {
if !nodes.contains_key(r) {
refs.insert(r.clone());
}
}
}
Ok(refs)
}
pub fn validate(&mut self) -> Result<(), IndexError> {
self.buffer_all()?;
Ok(())
}
pub fn resolve_references(
&self,
references: &[Vec<u64>],
) -> Result<Vec<Vec<IndexKey>>, IndexError> {
let mut out = Vec::with_capacity(references.len());
for ref_list in references {
let mut resolved = Vec::with_capacity(ref_list.len());
for off in ref_list {
let raw = self.keys_by_offset.get(off).ok_or_else(|| {
IndexError::Other(format!("unresolved reference offset {}", off))
})?;
resolved.push(raw.key.clone());
}
out.push(resolved);
}
Ok(out)
}
fn parse_header_from_bytes<'a>(
&mut self,
bytes: &'a [u8],
) -> Result<(u64, &'a [u8]), IndexError> {
let header = parse_header(bytes)?;
self.range_map.mark_parsed(
0,
Some(Vec::new()),
header.header_end as u64,
Some(Vec::new()),
);
let header_end = header.header_end as u64;
self.header = Some(header);
self.bisect_nodes = Some(HashMap::new());
Ok((header_end, &bytes[header_end as usize..]))
}
fn parse_segment(
&mut self,
offset: u64,
data: &[u8],
end: u64,
index: isize,
) -> Result<(u64, bool), IndexError> {
let lower_end = self
.range_map
.byte_range(index as usize)
.ok_or_else(|| IndexError::Other("parse_segment: index out of range".into()))?
.1;
let mut trim_start: Option<u64>;
let start_adjacent;
if offset < lower_end {
trim_start = Some(lower_end - offset);
start_adjacent = true;
} else if offset == lower_end {
trim_start = None;
start_adjacent = true;
} else {
trim_start = None;
start_adjacent = false;
}
let size = self.size.unwrap_or(0);
let mut trim_end: Option<u64>;
let end_adjacent;
let last_segment;
if Some(end) == self.size {
trim_end = None;
end_adjacent = true;
last_segment = true;
} else if (index as usize) + 1 == self.range_map.len() {
trim_end = None;
end_adjacent = false;
last_segment = true;
} else {
let (higher_start, higher_end) = self
.range_map
.byte_range((index as usize) + 1)
.expect("higher region present");
if end == higher_start {
trim_end = None;
end_adjacent = true;
last_segment = true;
} else if end > higher_start {
trim_end = Some(higher_start - offset);
end_adjacent = true;
last_segment = end < higher_end;
} else {
trim_end = None;
end_adjacent = false;
last_segment = true;
}
}
let _ = size;
if !start_adjacent {
let start_idx = trim_start.map(|s| s as usize).unwrap_or(0);
let pos = data[start_idx..]
.iter()
.position(|&b| b == b'\n')
.ok_or_else(|| IndexError::Other("no \\n was present".into()))?;
trim_start = Some((start_idx + pos + 1) as u64);
}
if !end_adjacent {
let end_idx = trim_end.map(|e| e as usize).unwrap_or(data.len());
let pos = data[..end_idx]
.iter()
.rposition(|&b| b == b'\n')
.ok_or_else(|| IndexError::Other("no \\n was present".into()))?;
trim_end = Some((pos + 1) as u64);
}
let ts = trim_start.map(|t| t as usize).unwrap_or(0);
let te = trim_end.map(|t| t as usize).unwrap_or(data.len());
let trimmed = &data[ts..te];
if trimmed.is_empty() {
return Err(IndexError::Other(format!(
"read unneeded data [{}:{}] from [{}:{}]",
ts,
te,
offset,
offset + data.len() as u64
)));
}
let parse_offset = if ts != 0 { offset + ts as u64 } else { offset };
let mut segments: Vec<&[u8]> = trimmed.split(|&b| b == b'\n').collect();
segments.pop();
let key_length = self.header.as_ref().expect("header parsed").key_length;
let parsed = parse_lines(&segments, parse_offset, key_length)?;
let bisect_nodes = self
.bisect_nodes
.as_mut()
.expect("bisect_nodes initialised by parse_header_from_bytes");
for (key, value, ref_offsets) in parsed.nodes {
bisect_nodes.insert(key, (value, ref_offsets));
}
for (off, raw) in &parsed.keys_by_offset {
self.keys_by_offset.insert(*off, raw.clone());
}
self.range_map.mark_parsed(
parse_offset,
parsed.first_key,
parse_offset + trimmed.len() as u64,
parsed.last_key,
);
Ok((parse_offset + trimmed.len() as u64, last_segment))
}
fn parse_region(&mut self, offset: u64, data: &[u8]) -> Result<(), IndexError> {
let end = offset + data.len() as u64;
let mut high_parsed = offset;
loop {
let index = self.range_map.byte_index(high_parsed);
let cur_end = self
.range_map
.byte_range(index as usize)
.map(|r| r.1)
.unwrap_or(0);
if end < cur_end {
return Ok(());
}
let (next_high, last_segment) = self.parse_segment(offset, data, end, index)?;
high_parsed = next_high;
if last_segment {
return Ok(());
}
}
}
fn read_and_parse(&mut self, mut readv_ranges: Vec<(u64, u64)>) -> Result<(), IndexError> {
if readv_ranges.is_empty() {
return Ok(());
}
let size = self
.size
.ok_or_else(|| IndexError::Other("read_and_parse called without a size".into()))?;
if self.nodes.is_none() && self.bytes_read * 2 >= size {
self.buffer_all()?;
return Ok(());
}
if self.base_offset != 0 {
for r in &mut readv_ranges {
r.0 += self.base_offset;
}
}
let upper = size + self.base_offset;
let readv_data = self
.transport
.readv(&self.name, &readv_ranges, true, upper)?;
for (raw_offset, raw_data) in readv_data {
let signed_offset = raw_offset as i64 - self.base_offset as i64;
let (mut offset, mut data) = if signed_offset < 0 {
let drop = (-signed_offset) as usize;
if drop >= raw_data.len() {
self.bytes_read += raw_data.len() as u64;
continue;
}
(0u64, raw_data[drop..].to_vec())
} else {
(signed_offset as u64, raw_data)
};
self.bytes_read += data.len() as u64;
if data.len() as u64 == size && offset == 0 {
self.buffer_all_from_bytes(data)?;
return Ok(());
}
if self.bisect_nodes.is_none() {
if offset != 0 {
return Err(IndexError::Other(
"header chunk did not start at offset 0".into(),
));
}
let (header_end, body) = self.parse_header_from_bytes(&data)?;
let body_vec = body.to_vec();
offset = header_end;
data = body_vec;
}
self.parse_region(offset, &data)?;
}
Ok(())
}
fn buffer_all_from_bytes(&mut self, data: Vec<u8>) -> Result<(), IndexError> {
if self.nodes.is_some() {
return Ok(());
}
let (header, nodes) = parse_full(&data)?;
self.nodes = Some(nodes);
self.header = Some(header);
Ok(())
}
pub fn lookup_keys_via_location(
&mut self,
location_keys: &[(u64, IndexKey)],
) -> Result<Vec<((u64, IndexKey), LookupResult)>, IndexError> {
let size = self
.size
.ok_or_else(|| IndexError::Other("lookup_keys_via_location requires a size".into()))?;
let mut readv_ranges: Vec<(u64, u64)> = Vec::new();
for (location, key) in location_keys {
if let Some(bn) = &self.bisect_nodes {
if bn.contains_key(key) {
continue;
}
}
let key_idx = self.range_map.key_index(&Some(key.clone()));
if !self.range_map.is_empty() && key_idx >= 0 {
let (key_start, key_end) = self
.range_map
.key_range(key_idx as usize)
.expect("idx in range");
let (_, byte_end) = self
.range_map
.byte_range(key_idx as usize)
.expect("idx in range");
if key_start.as_ref().map(|k| k <= key).unwrap_or(true)
&& (key_end.as_ref().map(|k| k >= key).unwrap_or(false) || byte_end == size)
{
continue;
}
}
let byte_idx = self.range_map.byte_index(*location);
if !self.range_map.is_empty() && byte_idx >= 0 {
let (byte_start, byte_end) = self
.range_map
.byte_range(byte_idx as usize)
.expect("idx in range");
if byte_start <= *location && byte_end > *location {
continue;
}
}
let mut length = 800u64;
if location + length > size {
length = size - location;
}
if length > 0 {
readv_ranges.push((*location, length));
}
}
if self.bisect_nodes.is_none() {
readv_ranges.push((0, 200));
}
self.read_and_parse(readv_ranges)?;
let mut result: Vec<((u64, IndexKey), LookupResult)> = Vec::new();
if let Some(nodes) = &self.nodes {
for (location, key) in location_keys {
if !nodes.contains_key(key) {
result.push(((*location, key.clone()), LookupResult::Missing));
} else {
let (value, refs) = nodes.get(key).unwrap();
result.push((
(*location, key.clone()),
LookupResult::Found {
value: value.clone(),
refs: refs.clone(),
},
));
}
}
return Ok(result);
}
let mut pending_references: Vec<(u64, IndexKey)> = Vec::new();
let mut pending_locations: std::collections::HashSet<u64> =
std::collections::HashSet::new();
let bisect_nodes_view = self
.bisect_nodes
.as_ref()
.expect("bisect_nodes initialised");
let header_ref_lists = self.header.as_ref().expect("header").node_ref_lists;
for (location, key) in location_keys {
if bisect_nodes_view.contains_key(key) {
let (value, refs) = bisect_nodes_view.get(key).unwrap();
if header_ref_lists > 0 {
let mut wanted: Vec<u64> = Vec::new();
for ref_list in refs {
for r in ref_list {
if !self.keys_by_offset.contains_key(r) {
wanted.push(*r);
}
}
}
if !wanted.is_empty() {
pending_locations.extend(wanted);
pending_references.push((*location, key.clone()));
continue;
}
let resolved = self.resolve_references(refs)?;
result.push((
(*location, key.clone()),
LookupResult::Found {
value: value.clone(),
refs: resolved,
},
));
} else {
result.push((
(*location, key.clone()),
LookupResult::Found {
value: value.clone(),
refs: Vec::new(),
},
));
}
continue;
}
let key_idx = self.range_map.key_index(&Some(key.clone()));
if key_idx >= 0 {
let (key_start, key_end) = self
.range_map
.key_range(key_idx as usize)
.expect("idx in range");
let (_, byte_end) = self
.range_map
.byte_range(key_idx as usize)
.expect("idx in range");
if key_start.as_ref().map(|k| k <= key).unwrap_or(true)
&& (key_end.as_ref().map(|k| k >= key).unwrap_or(false) || byte_end == size)
{
result.push(((*location, key.clone()), LookupResult::Missing));
continue;
}
}
let byte_idx = self.range_map.byte_index(*location);
let (probed_key_start, _) = self
.range_map
.key_range(byte_idx.max(0) as usize)
.unwrap_or((None, None));
let direction = if probed_key_start.as_ref().map(|k| key < k).unwrap_or(false) {
LookupResult::Direction(-1)
} else {
LookupResult::Direction(1)
};
result.push(((*location, key.clone()), direction));
}
let mut more_ranges: Vec<(u64, u64)> = Vec::new();
for location in &pending_locations {
let mut length = 800u64;
if location + length > size {
length = size - location;
}
if length > 0 {
more_ranges.push((*location, length));
}
}
self.read_and_parse(more_ranges)?;
if let Some(nodes) = &self.nodes {
for (location, key) in pending_references {
let (value, refs) = nodes.get(&key).expect("nodes contains pending key");
result.push((
(location, key.clone()),
LookupResult::Found {
value: value.clone(),
refs: refs.clone(),
},
));
}
return Ok(result);
}
let bisect_nodes_view = self
.bisect_nodes
.as_ref()
.expect("bisect_nodes initialised");
let pending_clone: Vec<(u64, IndexKey)> = pending_references.clone();
for (location, key) in pending_clone {
let (value, refs) = bisect_nodes_view
.get(&key)
.expect("bisect_nodes contains pending key");
let value = value.clone();
let refs = refs.clone();
let resolved = self.resolve_references(&refs)?;
result.push((
(location, key),
LookupResult::Found {
value,
refs: resolved,
},
));
}
Ok(result)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LookupResult {
Found {
value: Vec<u8>,
refs: Vec<Vec<IndexKey>>,
},
Missing,
Direction(i32),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BuilderNode {
pub absent: bool,
pub references: Vec<Vec<IndexKey>>,
pub value: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct GraphIndexBuilder {
reference_lists: usize,
key_length: usize,
nodes: HashMap<IndexKey, BuilderNode>,
absent_keys: std::collections::HashSet<IndexKey>,
optimize_for_size: bool,
combine_backing_indices: bool,
}
impl GraphIndexBuilder {
pub fn new(reference_lists: usize, key_elements: usize) -> Self {
Self {
reference_lists,
key_length: key_elements,
nodes: HashMap::new(),
absent_keys: std::collections::HashSet::new(),
optimize_for_size: false,
combine_backing_indices: true,
}
}
pub fn reference_lists(&self) -> usize {
self.reference_lists
}
pub fn key_length(&self) -> usize {
self.key_length
}
pub fn optimize_for_size(&self) -> bool {
self.optimize_for_size
}
pub fn combine_backing_indices(&self) -> bool {
self.combine_backing_indices
}
pub fn set_optimize(&mut self, for_size: Option<bool>, combine_backing_indices: Option<bool>) {
if let Some(v) = for_size {
self.optimize_for_size = v;
}
if let Some(v) = combine_backing_indices {
self.combine_backing_indices = v;
}
}
pub fn nodes(&self) -> &HashMap<IndexKey, BuilderNode> {
&self.nodes
}
pub fn contains_present(&self, key: &IndexKey) -> bool {
self.nodes.get(key).map(|n| !n.absent).unwrap_or(false)
}
pub fn check_key(&self, key: &IndexKey) -> Result<(), IndexError> {
if !key_is_valid(key, self.key_length) {
return Err(IndexError::BadKey(key.clone()));
}
Ok(())
}
pub fn check_key_ref_value(
&self,
key: &IndexKey,
references: &[Vec<IndexKey>],
value: &[u8],
) -> Result<(Vec<Vec<IndexKey>>, Vec<IndexKey>), IndexError> {
self.check_key(key)?;
if !value_is_valid(value) {
return Err(IndexError::BadValue(format!(
"value {:?} contains \\n or \\0",
value
)));
}
if references.len() != self.reference_lists {
return Err(IndexError::BadValue(format!(
"expected {} reference lists, got {}",
self.reference_lists,
references.len()
)));
}
let mut absent = Vec::new();
let mut node_refs = Vec::with_capacity(references.len());
for ref_list in references {
let mut tupled: Vec<IndexKey> = Vec::with_capacity(ref_list.len());
for r in ref_list {
if !self.nodes.contains_key(r) {
self.check_key(r)?;
absent.push(r.clone());
}
tupled.push(r.clone());
}
node_refs.push(tupled);
}
Ok((node_refs, absent))
}
pub fn add_node(
&mut self,
key: IndexKey,
value: Vec<u8>,
references: Vec<Vec<IndexKey>>,
) -> Result<(), IndexError> {
let (node_refs, absent_refs) = self.check_key_ref_value(&key, &references, &value)?;
if let Some(existing) = self.nodes.get(&key) {
if !existing.absent {
return Err(IndexError::DuplicateKey(key));
}
}
for r in &absent_refs {
self.nodes.entry(r.clone()).or_insert_with(|| BuilderNode {
absent: true,
references: Vec::new(),
value: Vec::new(),
});
self.absent_keys.insert(r.clone());
}
self.absent_keys.remove(&key);
self.nodes.insert(
key,
BuilderNode {
absent: false,
references: node_refs,
value,
},
);
Ok(())
}
pub fn key_count(&self) -> usize {
self.nodes.len() - self.absent_keys.len()
}
pub fn iter_all_entries(
&self,
) -> impl Iterator<Item = (IndexKey, Vec<u8>, Vec<Vec<IndexKey>>)> + '_ {
self.nodes.iter().filter_map(|(k, n)| {
if n.absent {
None
} else {
Some((k.clone(), n.value.clone(), n.references.clone()))
}
})
}
pub fn iter_entries<'a, I>(
&'a self,
keys: I,
) -> impl Iterator<Item = (IndexKey, Vec<u8>, Vec<Vec<IndexKey>>)> + 'a
where
I: IntoIterator<Item = IndexKey> + 'a,
{
keys.into_iter().filter_map(move |k| {
let n = self.nodes.get(&k)?;
if n.absent {
None
} else {
Some((k, n.value.clone(), n.references.clone()))
}
})
}
pub fn iter_entries_prefix(
&self,
prefixes: &[KeyPrefix],
) -> Result<Vec<IndexEntry>, IndexError> {
for p in prefixes {
if p.len() != self.key_length {
return Err(IndexError::BadKey(
p.iter().map(|e| e.clone().unwrap_or_default()).collect(),
));
}
if matches!(p.first(), Some(None)) {
return Err(IndexError::BadKey(Vec::new()));
}
}
let mut out = Vec::new();
let mut emitted: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
for prefix in prefixes {
for (k, n) in self.nodes.iter() {
if n.absent {
continue;
}
if k.len() != self.key_length {
continue;
}
let matches = prefix
.iter()
.zip(k.iter())
.all(|(p_elem, k_elem)| match p_elem {
Some(p) => p == k_elem,
None => true,
});
if matches && emitted.insert(k.clone()) {
out.push((k.clone(), n.value.clone(), n.references.clone()));
}
}
}
Ok(out)
}
pub fn external_references(&self) -> std::collections::HashSet<IndexKey> {
let mut refs = std::collections::HashSet::new();
if self.reference_lists < 2 {
return refs;
}
let mut keys: std::collections::HashSet<&IndexKey> = std::collections::HashSet::new();
for (k, n) in &self.nodes {
if n.absent {
continue;
}
keys.insert(k);
if let Some(list) = n.references.get(1) {
for r in list {
refs.insert(r.clone());
}
}
}
refs.retain(|r| !keys.contains(r));
refs
}
pub fn finish(&self) -> Result<Vec<u8>, IndexError> {
let mut nodes: Vec<IndexNode> = Vec::with_capacity(self.nodes.len());
for (key, node) in &self.nodes {
nodes.push(IndexNode {
key: key.clone(),
absent: node.absent,
references: node.references.clone(),
value: node.value.clone(),
});
}
serialize_graph_index(&nodes, self.reference_lists, self.key_length)
}
pub fn find_ancestry(
&self,
keys: &[IndexKey],
ref_list_num: usize,
) -> Result<
(
HashMap<IndexKey, Vec<IndexKey>>,
std::collections::HashSet<IndexKey>,
),
IndexError,
> {
let mut pending: std::collections::HashSet<IndexKey> = keys.iter().cloned().collect();
let mut parent_map: HashMap<IndexKey, Vec<IndexKey>> = HashMap::new();
let mut missing: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
while !pending.is_empty() {
let mut next_pending: std::collections::HashSet<IndexKey> =
std::collections::HashSet::new();
let snapshot: Vec<IndexKey> = pending.iter().cloned().collect();
for (k, _v, refs) in self.iter_entries(snapshot) {
let parent_keys = refs.get(ref_list_num).cloned().unwrap_or_default();
for p in &parent_keys {
if !parent_map.contains_key(p) {
next_pending.insert(p.clone());
}
}
parent_map.insert(k, parent_keys);
}
for k in pending.iter() {
if !parent_map.contains_key(k) {
missing.insert(k.clone());
}
}
pending = next_pending;
}
Ok((parent_map, missing))
}
}
pub trait IndexLike {
fn key_count(&self) -> Result<usize, IndexError>;
fn node_ref_lists(&self) -> Result<usize, IndexError>;
fn iter_all(&self) -> Result<Vec<IndexEntry>, IndexError>;
fn iter(&self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError>;
fn iter_prefix(&self, prefixes: &[KeyPrefix]) -> Result<Vec<IndexEntry>, IndexError>;
fn external_refs(
&self,
ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError>;
fn validate(&self) -> Result<(), IndexError> {
let _ = self.iter_all()?;
Ok(())
}
fn clear_cache(&self) {}
fn find_ancestors(
&self,
search_keys: &[IndexKey],
ref_list_num: usize,
parent_map: &mut HashMap<IndexKey, Vec<IndexKey>>,
missing_keys: &mut std::collections::HashSet<IndexKey>,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
let entries = self.iter(search_keys)?;
let mut found: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let mut new_search: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
for (key, _value, refs) in entries {
let parents: Vec<IndexKey> = refs.get(ref_list_num).cloned().unwrap_or_default();
for p in &parents {
if !parent_map.contains_key(p) {
new_search.insert(p.clone());
}
}
found.insert(key.clone());
parent_map.insert(key, parents);
}
for k in search_keys {
if !found.contains(k) {
missing_keys.insert(k.clone());
}
}
new_search.retain(|k| !parent_map.contains_key(k));
Ok(new_search)
}
}
impl IndexLike for GraphIndexBuilder {
fn key_count(&self) -> Result<usize, IndexError> {
Ok(self.key_count())
}
fn node_ref_lists(&self) -> Result<usize, IndexError> {
Ok(self.reference_lists)
}
fn iter_all(&self) -> Result<Vec<IndexEntry>, IndexError> {
Ok(self.iter_all_entries().collect())
}
fn iter(&self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError> {
Ok(self.iter_entries(keys.iter().cloned()).collect())
}
fn iter_prefix(&self, prefixes: &[KeyPrefix]) -> Result<Vec<IndexEntry>, IndexError> {
self.iter_entries_prefix(prefixes)
}
fn external_refs(
&self,
ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
if ref_list_num + 1 > self.reference_lists {
return Err(IndexError::Other(format!(
"No ref list {}, index has {} ref lists",
ref_list_num, self.reference_lists
)));
}
if ref_list_num != 1 {
return Ok(std::collections::HashSet::new());
}
Ok(self.external_references())
}
}
impl<T: IndexTransport> IndexLike for std::sync::Mutex<GraphIndex<T>> {
fn key_count(&self) -> Result<usize, IndexError> {
self.lock().unwrap().key_count()
}
fn node_ref_lists(&self) -> Result<usize, IndexError> {
self.lock().unwrap().node_ref_lists()
}
fn iter_all(&self) -> Result<Vec<IndexEntry>, IndexError> {
self.lock().unwrap().iter_all_entries()
}
fn iter(&self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError> {
self.lock().unwrap().iter_entries(keys)
}
fn iter_prefix(&self, prefixes: &[KeyPrefix]) -> Result<Vec<IndexEntry>, IndexError> {
self.lock().unwrap().iter_entries_prefix(prefixes)
}
fn external_refs(
&self,
ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
self.lock().unwrap().external_references(ref_list_num)
}
fn validate(&self) -> Result<(), IndexError> {
self.lock().unwrap().validate()
}
}
pub struct CombinedGraphIndex {
indices: Vec<Box<dyn IndexLike + Send + Sync>>,
}
impl CombinedGraphIndex {
pub fn new() -> Self {
Self {
indices: Vec::new(),
}
}
pub fn from_indices(indices: Vec<Box<dyn IndexLike + Send + Sync>>) -> Self {
Self { indices }
}
pub fn push(&mut self, index: Box<dyn IndexLike + Send + Sync>) {
self.indices.push(index);
}
pub fn insert(&mut self, pos: usize, index: Box<dyn IndexLike + Send + Sync>) {
self.indices.insert(pos, index);
}
pub fn len(&self) -> usize {
self.indices.len()
}
pub fn is_empty(&self) -> bool {
self.indices.is_empty()
}
pub fn indices(&self) -> &[Box<dyn IndexLike + Send + Sync>] {
&self.indices
}
pub fn move_to_front(&mut self, hits: &[usize]) {
if hits.is_empty() {
return;
}
let mut hit_set: std::collections::HashSet<usize> = hits.iter().copied().collect();
let mut new_order: Vec<Box<dyn IndexLike + Send + Sync>> =
Vec::with_capacity(self.indices.len());
for &h in hits {
if h < self.indices.len() {
hit_set.remove(&h);
}
}
let mut taken: HashMap<usize, Box<dyn IndexLike + Send + Sync>> = HashMap::new();
for (i, idx) in std::mem::take(&mut self.indices).into_iter().enumerate() {
taken.insert(i, idx);
}
for &h in hits {
if let Some(idx) = taken.remove(&h) {
new_order.push(idx);
}
}
let mut leftover: Vec<(usize, Box<dyn IndexLike + Send + Sync>)> =
taken.into_iter().collect();
leftover.sort_by_key(|(i, _)| *i);
for (_, idx) in leftover {
new_order.push(idx);
}
self.indices = new_order;
}
}
impl Default for CombinedGraphIndex {
fn default() -> Self {
Self::new()
}
}
impl CombinedGraphIndex {
pub fn iter_entries_with_hits(
&mut self,
keys: &[IndexKey],
) -> Result<(Vec<IndexEntry>, Vec<usize>), IndexError> {
let mut remaining: std::collections::HashSet<IndexKey> = keys.iter().cloned().collect();
let mut out = Vec::new();
let mut hits: Vec<usize> = Vec::new();
for (i, idx) in self.indices.iter().enumerate() {
if remaining.is_empty() {
break;
}
let snapshot: Vec<IndexKey> = remaining.iter().cloned().collect();
let entries = idx.iter(&snapshot)?;
let mut hit = false;
for entry in entries {
if remaining.remove(&entry.0) {
out.push(entry);
hit = true;
}
}
if hit {
hits.push(i);
}
}
Ok((out, hits))
}
pub fn iter_entries_prefix_with_hits(
&mut self,
prefixes: &[KeyPrefix],
) -> Result<(Vec<IndexEntry>, Vec<usize>), IndexError> {
let mut seen: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let mut out = Vec::new();
let mut hits: Vec<usize> = Vec::new();
for (i, idx) in self.indices.iter().enumerate() {
let entries = idx.iter_prefix(prefixes)?;
let mut hit = false;
for entry in entries {
if seen.insert(entry.0.clone()) {
out.push(entry);
hit = true;
}
}
if hit {
hits.push(i);
}
}
Ok((out, hits))
}
pub fn find_ancestry(
&self,
keys: &[IndexKey],
ref_list_num: usize,
) -> Result<
(
HashMap<IndexKey, Vec<IndexKey>>,
std::collections::HashSet<IndexKey>,
),
IndexError,
> {
let mut parent_map: HashMap<IndexKey, Vec<IndexKey>> = HashMap::new();
let mut missing_keys: std::collections::HashSet<IndexKey> =
std::collections::HashSet::new();
let mut keys_to_lookup: std::collections::HashSet<IndexKey> =
keys.iter().cloned().collect();
while !keys_to_lookup.is_empty() {
let mut all_index_missing: Option<std::collections::HashSet<IndexKey>> = None;
let mut current = keys_to_lookup.clone();
for idx in &self.indices {
let mut index_missing: std::collections::HashSet<IndexKey> =
std::collections::HashSet::new();
let snapshot: Vec<IndexKey> = current.iter().cloned().collect();
let mut search_keys = snapshot;
while !search_keys.is_empty() {
let new_search = idx.find_ancestors(
&search_keys,
ref_list_num,
&mut parent_map,
&mut index_missing,
)?;
search_keys = new_search.into_iter().collect();
}
match all_index_missing.as_mut() {
None => {
all_index_missing = Some(index_missing.clone());
}
Some(prev) => {
prev.retain(|k| index_missing.contains(k));
}
}
current = index_missing;
if current.is_empty() {
break;
}
}
match all_index_missing {
None => {
missing_keys.extend(current);
break;
}
Some(s) => {
missing_keys.extend(s.iter().cloned());
keys_to_lookup = current.difference(&s).cloned().collect();
}
}
}
Ok((parent_map, missing_keys))
}
pub fn get_parent_map(
&self,
keys: &[IndexKey],
null_revision: &IndexKey,
) -> Result<HashMap<IndexKey, Vec<IndexKey>>, IndexError> {
let mut search_keys: Vec<IndexKey> = keys.to_vec();
let mut found_parents: HashMap<IndexKey, Vec<IndexKey>> = HashMap::new();
if let Some(pos) = search_keys.iter().position(|k| k == null_revision) {
search_keys.remove(pos);
found_parents.insert(null_revision.clone(), Vec::new());
}
for (key, _value, refs) in self.iter(&search_keys)? {
let parents = refs.first().cloned().unwrap_or_default();
if parents.is_empty() {
found_parents.insert(key, vec![null_revision.clone()]);
} else {
found_parents.insert(key, parents);
}
}
Ok(found_parents)
}
}
impl IndexLike for CombinedGraphIndex {
fn key_count(&self) -> Result<usize, IndexError> {
let mut total = 0;
for idx in &self.indices {
total += idx.key_count()?;
}
Ok(total)
}
fn node_ref_lists(&self) -> Result<usize, IndexError> {
if let Some(first) = self.indices.first() {
first.node_ref_lists()
} else {
Ok(0)
}
}
fn iter_all(&self) -> Result<Vec<IndexEntry>, IndexError> {
let mut seen: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let mut out = Vec::new();
for idx in &self.indices {
for entry in idx.iter_all()? {
if seen.insert(entry.0.clone()) {
out.push(entry);
}
}
}
Ok(out)
}
fn iter(&self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError> {
let mut remaining: std::collections::HashSet<IndexKey> = keys.iter().cloned().collect();
let mut out = Vec::new();
for idx in &self.indices {
if remaining.is_empty() {
break;
}
let snapshot: Vec<IndexKey> = remaining.iter().cloned().collect();
for entry in idx.iter(&snapshot)? {
if remaining.remove(&entry.0) {
out.push(entry);
}
}
}
Ok(out)
}
fn iter_prefix(&self, prefixes: &[KeyPrefix]) -> Result<Vec<IndexEntry>, IndexError> {
let mut seen: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let mut out = Vec::new();
for idx in &self.indices {
for entry in idx.iter_prefix(prefixes)? {
if seen.insert(entry.0.clone()) {
out.push(entry);
}
}
}
Ok(out)
}
fn external_refs(
&self,
ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
let mut refs = std::collections::HashSet::new();
for idx in &self.indices {
refs.extend(idx.external_refs(ref_list_num)?);
}
Ok(refs)
}
fn validate(&self) -> Result<(), IndexError> {
for idx in &self.indices {
idx.validate()?;
}
Ok(())
}
fn clear_cache(&self) {
for idx in &self.indices {
idx.clear_cache();
}
}
}
pub struct GraphIndexPrefixAdapter<I: IndexLike> {
inner: I,
prefix: IndexKey,
prefix_len: usize,
prefix_query: KeyPrefix,
}
impl<I: IndexLike> GraphIndexPrefixAdapter<I> {
pub fn new(inner: I, prefix: IndexKey, missing_key_length: usize) -> Self {
let prefix_len = prefix.len();
let mut prefix_query: KeyPrefix = prefix.iter().cloned().map(Some).collect();
for _ in 0..missing_key_length {
prefix_query.push(None);
}
Self {
inner,
prefix,
prefix_len,
prefix_query,
}
}
fn extend_key(&self, key: &IndexKey) -> IndexKey {
let mut full = self.prefix.clone();
full.extend(key.iter().cloned());
full
}
fn strip_entry(&self, entry: IndexEntry) -> Result<IndexEntry, IndexError> {
let (key, value, refs) = entry;
if key.len() < self.prefix_len {
return Err(IndexError::BadIndexData);
}
for (a, b) in self.prefix.iter().zip(key.iter()) {
if a != b {
return Err(IndexError::BadIndexData);
}
}
let stripped_key: IndexKey = key[self.prefix_len..].to_vec();
let mut stripped_refs: Vec<Vec<IndexKey>> = Vec::with_capacity(refs.len());
for ref_list in refs {
let mut new_list: Vec<IndexKey> = Vec::with_capacity(ref_list.len());
for ref_key in ref_list {
if ref_key.len() < self.prefix_len {
return Err(IndexError::BadIndexData);
}
for (a, b) in self.prefix.iter().zip(ref_key.iter()) {
if a != b {
return Err(IndexError::BadIndexData);
}
}
new_list.push(ref_key[self.prefix_len..].to_vec());
}
stripped_refs.push(new_list);
}
Ok((stripped_key, value, stripped_refs))
}
}
impl<I: IndexLike> IndexLike for GraphIndexPrefixAdapter<I> {
fn key_count(&self) -> Result<usize, IndexError> {
Ok(self.iter_all()?.len())
}
fn node_ref_lists(&self) -> Result<usize, IndexError> {
self.inner.node_ref_lists()
}
fn iter_all(&self) -> Result<Vec<IndexEntry>, IndexError> {
let inner_entries = self.inner.iter_prefix(&[self.prefix_query.clone()])?;
let mut out = Vec::with_capacity(inner_entries.len());
for e in inner_entries {
out.push(self.strip_entry(e)?);
}
Ok(out)
}
fn iter(&self, keys: &[IndexKey]) -> Result<Vec<IndexEntry>, IndexError> {
let extended: Vec<IndexKey> = keys.iter().map(|k| self.extend_key(k)).collect();
let inner_entries = self.inner.iter(&extended)?;
let mut out = Vec::with_capacity(inner_entries.len());
for e in inner_entries {
out.push(self.strip_entry(e)?);
}
Ok(out)
}
fn iter_prefix(&self, prefixes: &[KeyPrefix]) -> Result<Vec<IndexEntry>, IndexError> {
let extended: Vec<KeyPrefix> = prefixes
.iter()
.map(|p| {
let mut full: KeyPrefix = self.prefix.iter().cloned().map(Some).collect();
full.extend(p.iter().cloned());
full
})
.collect();
let inner_entries = self.inner.iter_prefix(&extended)?;
let mut out = Vec::with_capacity(inner_entries.len());
for e in inner_entries {
out.push(self.strip_entry(e)?);
}
Ok(out)
}
fn external_refs(
&self,
_ref_list_num: usize,
) -> Result<std::collections::HashSet<IndexKey>, IndexError> {
Ok(std::collections::HashSet::new())
}
fn validate(&self) -> Result<(), IndexError> {
self.inner.validate()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn key(parts: &[&[u8]]) -> Vec<Vec<u8>> {
parts.iter().map(|p| p.to_vec()).collect()
}
fn node(k: &[&[u8]], absent: bool, refs: Vec<Vec<Vec<Vec<u8>>>>, value: &[u8]) -> IndexNode {
IndexNode {
key: key(k),
absent,
references: refs,
value: value.to_vec(),
}
}
#[test]
fn serialize_empty_index_no_refs() {
let out = serialize_graph_index(&[], 0, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=1\nlen=0\n\n".to_vec()
);
}
#[test]
fn serialize_single_node_no_refs() {
let nodes = vec![node(&[b"a"], false, vec![], b"val")];
let out = serialize_graph_index(&nodes, 0, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=1\nlen=1\na\x00\x00\x00val\n\n"
.to_vec()
);
}
#[test]
fn serialize_with_reference_back_to_earlier_key() {
let nodes = vec![
node(&[b"a"], false, vec![vec![]], b"val1"),
node(&[b"b"], false, vec![vec![key(&[b"a"])]], b"val2"),
];
let out = serialize_graph_index(&nodes, 1, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=1\nkey_elements=1\nlen=2\na\x00\x00\x00val1\nb\x00\x0059\x00val2\n\n"
.to_vec()
);
}
#[test]
fn serialize_absent_node_has_no_tab_between_ref_lists() {
let nodes = vec![
node(
&[b"a"],
false,
vec![vec![key(&[b"missing"])], vec![]],
b"value",
),
node(&[b"missing"], true, vec![], b""),
];
let out = serialize_graph_index(&nodes, 2, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=1\nlen=1\na\x00\x0072\t\x00value\nmissing\x00a\x00\x00\n\n"
.to_vec()
);
}
#[test]
fn serialize_multi_element_key() {
let nodes = vec![node(&[b"x", b"y"], false, vec![], b"v")];
let out = serialize_graph_index(&nodes, 0, 2).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=2\nlen=1\nx\x00y\x00\x00\x00v\n\n"
.to_vec()
);
}
#[test]
fn serialize_reports_unknown_reference() {
let nodes = vec![node(&[b"a"], false, vec![vec![key(&[b"missing"])]], b"v")];
let err = serialize_graph_index(&nodes, 1, 1).unwrap_err();
assert_eq!(err, IndexError::UnknownReference(key(&[b"missing"])));
}
#[test]
fn serialize_pads_reference_offsets_to_matching_width() {
let mut nodes: Vec<IndexNode> = Vec::new();
for i in 0..20 {
let k = format!("key{:03}", i).into_bytes();
let refs = if i == 0 {
vec![vec![]]
} else {
vec![vec![key(&[&format!("key{:03}", i - 1).into_bytes()])]]
};
nodes.push(node(
&[k.as_slice()],
false,
refs,
format!("value{:03}", i).as_bytes(),
));
}
let out = serialize_graph_index(&nodes, 1, 1).unwrap();
assert_eq!(out.len(), 478);
assert!(out
.windows(b"key001\x00\x00060\x00value001\n".len())
.any(|w| w == b"key001\x00\x00060\x00value001\n"));
}
#[test]
fn parse_header_minimal_index() {
let data = b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=1\nlen=0\n\n";
let h = parse_header(data).unwrap();
assert_eq!(h.node_ref_lists, 0);
assert_eq!(h.key_length, 1);
assert_eq!(h.key_count, 0);
assert_eq!(h.header_end, 59);
}
#[test]
fn parse_header_non_zero_values() {
let data = b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=3\nlen=42\n";
let h = parse_header(data).unwrap();
assert_eq!(h.node_ref_lists, 2);
assert_eq!(h.key_length, 3);
assert_eq!(h.key_count, 42);
}
#[test]
fn parse_header_rejects_bad_signature() {
assert_eq!(
parse_header(b"not an index\n"),
Err(IndexError::BadSignature)
);
}
#[test]
fn parse_header_rejects_missing_option() {
let data = b"Bazaar Graph Index 1\nwrong_option=1\nkey_elements=1\nlen=0\n\n";
assert_eq!(parse_header(data), Err(IndexError::BadOptions));
}
#[test]
fn parse_header_rejects_non_decimal_option() {
let data = b"Bazaar Graph Index 1\nnode_ref_lists=abc\nkey_elements=1\nlen=0\n\n";
assert_eq!(parse_header(data), Err(IndexError::BadOptions));
}
#[test]
fn parse_lines_single_node_no_refs() {
let line: &[u8] = b"a\x00\x00\x00val";
let lines = vec![line, b""];
let parsed = parse_lines(&lines, 100, 1).unwrap();
assert_eq!(parsed.first_key, Some(key(&[b"a"])));
assert_eq!(parsed.last_key, Some(key(&[b"a"])));
assert_eq!(parsed.trailers, 1);
assert_eq!(parsed.nodes.len(), 1);
let (k, v, refs) = &parsed.nodes[0];
assert_eq!(k, &key(&[b"a"]));
assert_eq!(v, b"val");
assert_eq!(refs, &vec![Vec::<u64>::new()]);
}
#[test]
fn parse_lines_tracks_offsets() {
let line_a: &[u8] = b"a\x00\x00\x00val1";
let line_b: &[u8] = b"b\x00\x0000\x00val2";
let lines = vec![line_a, line_b];
let parsed = parse_lines(&lines, 0, 1).unwrap();
assert_eq!(parsed.keys_by_offset.len(), 2);
assert_eq!(parsed.keys_by_offset[0].0, 0);
assert_eq!(parsed.keys_by_offset[1].0, line_a.len() as u64 + 1);
}
#[test]
fn parse_lines_absent_node_not_in_output_but_in_offset_map() {
let line: &[u8] = b"ghost\x00a\x00\x00";
let parsed = parse_lines(&[line], 50, 1).unwrap();
assert!(parsed.nodes.is_empty());
assert_eq!(parsed.keys_by_offset.len(), 1);
assert!(parsed.keys_by_offset[0].1.absent);
assert_eq!(parsed.keys_by_offset[0].1.key, key(&[b"ghost"]));
}
#[test]
fn parse_lines_references() {
let line: &[u8] = b"k\x00\x00100\r200\t300\x00val";
let parsed = parse_lines(&[line], 0, 1).unwrap();
let refs = &parsed.nodes[0].2;
assert_eq!(refs.len(), 2);
assert_eq!(refs[0], vec![100u64, 200]);
assert_eq!(refs[1], vec![300u64]);
}
#[test]
fn parse_lines_bad_field_count_errors() {
let line: &[u8] = b"k\x00\x00val"; assert_eq!(parse_lines(&[line], 0, 1), Err(IndexError::BadLineData));
}
#[test]
fn parse_lines_bad_reference_offset_errors() {
let line: &[u8] = b"k\x00\x00notnumeric\x00val";
assert!(matches!(
parse_lines(&[line], 0, 1),
Err(IndexError::BadReferenceOffset(_))
));
}
#[test]
fn round_trip_serialize_then_parse() {
let nodes = vec![
node(&[b"a"], false, vec![vec![]], b"val1"),
node(&[b"b"], false, vec![vec![key(&[b"a"])]], b"val2"),
];
let bytes = serialize_graph_index(&nodes, 1, 1).unwrap();
let header = parse_header(&bytes).unwrap();
assert_eq!(header.node_ref_lists, 1);
assert_eq!(header.key_length, 1);
assert_eq!(header.key_count, 2);
let body = &bytes[header.header_end..];
let body_lines: Vec<&[u8]> = body.split(|&b| b == b'\n').collect();
let parsed = parse_lines(&body_lines, header.header_end as u64, 1).unwrap();
assert_eq!(parsed.nodes.len(), 2);
assert_eq!(parsed.nodes[0].0, key(&[b"a"]));
assert_eq!(parsed.nodes[1].0, key(&[b"b"]));
assert_eq!(parsed.nodes[1].2, vec![vec![header.header_end as u64]]);
assert!(parsed.trailers >= 1);
}
#[test]
fn serialize_empty_index_two_element_keys() {
let out = serialize_graph_index(&[], 0, 2).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=2\nlen=0\n\n".to_vec()
);
}
#[test]
fn serialize_empty_index_one_reference_list() {
let out = serialize_graph_index(&[], 1, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=1\nkey_elements=1\nlen=0\n\n".to_vec()
);
}
#[test]
fn serialize_empty_index_two_reference_lists() {
let out = serialize_graph_index(&[], 2, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=1\nlen=0\n\n".to_vec()
);
}
#[test]
fn serialize_empty_value_node() {
let nodes = vec![node(&[b"akey"], false, vec![], b"")];
let out = serialize_graph_index(&nodes, 0, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=1\nlen=1\nakey\x00\x00\x00\n\n"
.to_vec()
);
}
#[test]
fn serialize_sorts_three_nodes_byte_exact() {
let nodes = vec![
node(&[b"2002"], false, vec![], b"data"),
node(&[b"2000"], false, vec![], b"data"),
node(&[b"2001"], false, vec![], b"data"),
];
let out = serialize_graph_index(&nodes, 0, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=1\nlen=3\n\
2000\x00\x00\x00data\n\
2001\x00\x00\x00data\n\
2002\x00\x00\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_sorts_two_element_keys_lexicographically() {
let mut nodes = Vec::new();
for first in &[b"2002", b"2000", b"2001"] {
for second in &[b"2002", b"2000", b"2001"] {
nodes.push(node(&[*first, *second], false, vec![], b"data"));
}
}
let out = serialize_graph_index(&nodes, 0, 2).unwrap();
let expected: Vec<u8> = [
b"Bazaar Graph Index 1\nnode_ref_lists=0\nkey_elements=2\nlen=9\n".as_slice(),
b"2000\x002000\x00\x00\x00data\n",
b"2000\x002001\x00\x00\x00data\n",
b"2000\x002002\x00\x00\x00data\n",
b"2001\x002000\x00\x00\x00data\n",
b"2001\x002001\x00\x00\x00data\n",
b"2001\x002002\x00\x00\x00data\n",
b"2002\x002000\x00\x00\x00data\n",
b"2002\x002001\x00\x00\x00data\n",
b"2002\x002002\x00\x00\x00data\n",
b"\n",
]
.concat();
assert_eq!(out, expected);
}
#[test]
fn serialize_single_node_with_empty_ref_list_of_one() {
let nodes = vec![node(&[b"key"], false, vec![vec![]], b"data")];
let out = serialize_graph_index(&nodes, 1, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=1\nkey_elements=1\nlen=1\nkey\x00\x00\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_single_node_with_empty_ref_lists_of_two() {
let nodes = vec![node(&[b"key"], false, vec![vec![], vec![]], b"data")];
let out = serialize_graph_index(&nodes, 2, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=1\nlen=1\nkey\x00\x00\t\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_ref_list_with_two_element_keys() {
let nodes = vec![node(&[b"key", b"key2"], false, vec![vec![]], b"data")];
let out = serialize_graph_index(&nodes, 1, 2).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=1\nkey_elements=2\nlen=1\nkey\x00key2\x00\x00\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_cr_delimits_multiple_refs_in_one_list() {
let nodes = vec![
node(&[b"reference"], false, vec![vec![]], b"data"),
node(&[b"reference2"], false, vec![vec![]], b"data"),
node(
&[b"key"],
false,
vec![vec![key(&[b"reference"]), key(&[b"reference2"])]],
b"data",
),
];
let out = serialize_graph_index(&nodes, 1, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=1\nkey_elements=1\nlen=3\n\
key\x00\x00077\r094\x00data\n\
reference\x00\x00\x00data\n\
reference2\x00\x00\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_tab_delimits_multiple_reference_lists() {
let nodes = vec![
node(&[b"keference"], false, vec![vec![], vec![]], b"data"),
node(
&[b"rey"],
false,
vec![vec![key(&[b"keference"])], vec![key(&[b"keference"])]],
b"data",
),
];
let out = serialize_graph_index(&nodes, 2, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=1\nlen=2\n\
keference\x00\x00\t\x00data\n\
rey\x00\x0059\t59\x00data\n\n"
.to_vec()
);
}
#[test]
fn serialize_absent_record_has_no_reference_overhead() {
let nodes = vec![
node(&[b"aail"], true, vec![], b""),
node(
&[b"parent"],
false,
vec![vec![key(&[b"aail"]), key(&[b"zther"])], vec![]],
b"",
),
node(&[b"zther"], true, vec![], b""),
];
let out = serialize_graph_index(&nodes, 2, 1).unwrap();
assert_eq!(
out,
b"Bazaar Graph Index 1\nnode_ref_lists=2\nkey_elements=1\nlen=1\n\
aail\x00a\x00\x00\n\
parent\x00\x0059\r84\t\x00\n\
zther\x00a\x00\x00\n\n"
.to_vec()
);
}
#[test]
fn serialize_sorts_nodes_by_key() {
let nodes = vec![
node(&[b"c"], false, vec![], b"3"),
node(&[b"a"], false, vec![], b"1"),
node(&[b"b"], false, vec![], b"2"),
];
let out = serialize_graph_index(&nodes, 0, 1).unwrap();
let body_start = out
.windows(b"len=3\n".len())
.position(|w| w == b"len=3\n")
.unwrap()
+ b"len=3\n".len();
let body = &out[body_start..];
assert!(body.starts_with(b"a\x00\x00\x001"));
assert!(body.windows(5).any(|w| w == b"b\x00\x00\x002"));
assert!(body.windows(5).any(|w| w == b"c\x00\x00\x003"));
}
struct MemTransport {
files: std::collections::HashMap<String, Vec<u8>>,
}
impl MemTransport {
fn new() -> Self {
Self {
files: std::collections::HashMap::new(),
}
}
fn put(&mut self, path: &str, bytes: Vec<u8>) {
self.files.insert(path.to_string(), bytes);
}
}
impl IndexTransport for MemTransport {
fn get_bytes(&self, path: &str) -> Result<Vec<u8>, IndexError> {
self.files
.get(path)
.cloned()
.ok_or_else(|| IndexError::Other(format!("NoSuchFile: {}", path)))
}
}
fn build_index(nodes: &[IndexNode], reference_lists: usize, key_elements: usize) -> Vec<u8> {
serialize_graph_index(nodes, reference_lists, key_elements).unwrap()
}
#[test]
fn graph_index_buffer_all_no_refs() {
let bytes = build_index(
&[
node(&[b"a"], false, vec![], b"v1"),
node(&[b"b"], false, vec![], b"v2"),
],
0,
1,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
assert_eq!(idx.key_count().unwrap(), 2);
assert_eq!(idx.node_ref_lists().unwrap(), 0);
let mut entries = idx.iter_all_entries().unwrap();
entries.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(
entries,
vec![
(key(&[b"a"]), b"v1".to_vec(), vec![]),
(key(&[b"b"]), b"v2".to_vec(), vec![]),
]
);
}
#[test]
fn graph_index_resolves_references() {
let bytes = build_index(
&[
node(&[b"a"], false, vec![vec![]], b"v1"),
node(&[b"b"], false, vec![vec![key(&[b"a"])]], b"v2"),
],
1,
1,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let mut entries = idx.iter_all_entries().unwrap();
entries.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(
entries,
vec![
(key(&[b"a"]), b"v1".to_vec(), vec![vec![]]),
(key(&[b"b"]), b"v2".to_vec(), vec![vec![key(&[b"a"])]],),
]
);
}
#[test]
fn graph_index_iter_entries_filters_to_requested_keys() {
let bytes = build_index(
&[
node(&[b"a"], false, vec![], b"v1"),
node(&[b"b"], false, vec![], b"v2"),
node(&[b"c"], false, vec![], b"v3"),
],
0,
1,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let mut entries = idx
.iter_entries(&[key(&[b"a"]), key(&[b"missing"]), key(&[b"c"])])
.unwrap();
entries.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(
entries,
vec![
(key(&[b"a"]), b"v1".to_vec(), vec![]),
(key(&[b"c"]), b"v3".to_vec(), vec![]),
]
);
}
#[test]
fn graph_index_iter_entries_dedupes_repeated_keys() {
let bytes = build_index(&[node(&[b"a"], false, vec![], b"v1")], 0, 1);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let entries = idx.iter_entries(&[key(&[b"a"]), key(&[b"a"])]).unwrap();
assert_eq!(entries, vec![(key(&[b"a"]), b"v1".to_vec(), vec![])]);
}
#[test]
fn graph_index_external_references() {
let bytes = build_index(
&[
node(&[b"a"], false, vec![vec![key(&[b"missing"])]], b"v1"),
node(&[b"missing"], true, vec![], b""),
node(&[b"b"], false, vec![vec![key(&[b"a"])]], b"v2"),
],
1,
1,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let externals = idx.external_references(0).unwrap();
let expected: std::collections::HashSet<IndexKey> =
vec![key(&[b"missing"])].into_iter().collect();
assert_eq!(externals, expected);
}
#[test]
fn graph_index_external_references_rejects_invalid_ref_list() {
let bytes = build_index(&[node(&[b"a"], false, vec![], b"v1")], 0, 1);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let err = idx.external_references(0).unwrap_err();
assert_eq!(
err,
IndexError::Other("No ref list 0, index has 0 ref lists".to_string())
);
}
#[test]
fn graph_index_iter_entries_prefix_one_element() {
let bytes = build_index(
&[
node(&[b"a"], false, vec![], b"v1"),
node(&[b"b"], false, vec![], b"v2"),
],
0,
1,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let entries = idx
.iter_entries_prefix(&[vec![Some(b"a".to_vec())]])
.unwrap();
assert_eq!(entries, vec![(key(&[b"a"]), b"v1".to_vec(), vec![])]);
}
#[test]
fn graph_index_iter_entries_prefix_multi_element() {
let bytes = build_index(
&[
node(&[b"foo", b"bar"], false, vec![], b"v1"),
node(&[b"foo", b"baz"], false, vec![], b"v2"),
node(&[b"qux", b"bar"], false, vec![], b"v3"),
],
0,
2,
);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let mut entries = idx
.iter_entries_prefix(&[vec![Some(b"foo".to_vec()), None]])
.unwrap();
entries.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(
entries,
vec![
(key(&[b"foo", b"bar"]), b"v1".to_vec(), vec![]),
(key(&[b"foo", b"baz"]), b"v2".to_vec(), vec![]),
]
);
}
#[test]
fn graph_index_iter_entries_prefix_rejects_none_first_element() {
let bytes = build_index(&[node(&[b"a"], false, vec![], b"v1")], 0, 1);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let err = idx.iter_entries_prefix(&[vec![None]]).unwrap_err();
assert_eq!(err, IndexError::BadKey(vec![Vec::new()]));
}
#[test]
fn graph_index_validate_ok_for_well_formed_index() {
let bytes = build_index(&[node(&[b"a"], false, vec![], b"v")], 0, 1);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
idx.validate().unwrap();
}
#[test]
fn graph_index_buffer_all_idempotent() {
let bytes = build_index(&[node(&[b"a"], false, vec![], b"v")], 0, 1);
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
idx.buffer_all().unwrap();
idx.buffer_all().unwrap();
assert_eq!(idx.key_count().unwrap(), 1);
}
#[test]
fn graph_index_missing_trailer_is_error() {
let mut bytes = build_index(&[node(&[b"a"], false, vec![], b"v")], 0, 1);
assert_eq!(bytes.last(), Some(&b'\n'));
bytes.pop();
let mut t = MemTransport::new();
t.put("idx", bytes);
let mut idx = GraphIndex::new(t, "idx", 0);
let err = idx.buffer_all().unwrap_err();
assert_eq!(
err,
IndexError::Other("BadIndexData: missing trailer".to_string())
);
}
#[test]
fn graph_index_respects_base_offset() {
let inner = build_index(&[node(&[b"a"], false, vec![], b"v")], 0, 1);
let mut wrapped = b"junk-before-header".to_vec();
let prefix_len = wrapped.len() as u64;
wrapped.extend_from_slice(&inner);
let mut t = MemTransport::new();
t.put("idx", wrapped);
let mut idx = GraphIndex::new(t, "idx", prefix_len);
assert_eq!(idx.key_count().unwrap(), 1);
let entries = idx.iter_all_entries().unwrap();
assert_eq!(entries, vec![(key(&[b"a"]), b"v".to_vec(), vec![])]);
}
#[test]
fn key_is_valid_accepts_clean_bytes() {
assert!(key_is_valid(&[b"foo".to_vec()], 1));
assert!(key_is_valid(&[b"foo".to_vec(), b"bar".to_vec()], 2));
}
#[test]
fn key_is_valid_rejects_wrong_length() {
assert!(!key_is_valid(&[b"foo".to_vec()], 2));
assert!(!key_is_valid(&[b"a".to_vec(), b"b".to_vec()], 1));
}
#[test]
fn key_is_valid_rejects_empty_element() {
assert!(!key_is_valid(&[b"".to_vec()], 1));
assert!(!key_is_valid(&[b"a".to_vec(), b"".to_vec()], 2));
}
#[test]
fn key_is_valid_rejects_separator_bytes() {
for &bad in [b'\t', b'\n', 0x0b, 0x0c, b'\r', 0, b' '].iter() {
let elem = vec![b'a', bad];
assert!(
!key_is_valid(&[elem.clone()], 1),
"byte 0x{:02x} should disqualify",
bad
);
}
}
#[test]
fn value_is_valid_accepts_arbitrary_bytes() {
assert!(value_is_valid(b"any value"));
assert!(value_is_valid(b""));
assert!(value_is_valid(b"with\ttab and CR\r is fine"));
}
#[test]
fn value_is_valid_rejects_newline_or_null() {
assert!(!value_is_valid(b"has\nnewline"));
assert!(!value_is_valid(b"has\0null"));
}
#[test]
fn builder_add_node_and_finish_roundtrip() {
let mut b = GraphIndexBuilder::new(0, 1);
b.add_node(key(&[b"a"]), b"v1".to_vec(), vec![]).unwrap();
b.add_node(key(&[b"b"]), b"v2".to_vec(), vec![]).unwrap();
assert_eq!(b.key_count(), 2);
let bytes = b.finish().unwrap();
let (header, parsed) = parse_full(&bytes).unwrap();
assert_eq!(header.key_count, 2);
assert_eq!(parsed.get(&key(&[b"a"])), Some(&(b"v1".to_vec(), vec![])));
}
#[test]
fn builder_rejects_duplicate_key() {
let mut b = GraphIndexBuilder::new(0, 1);
b.add_node(key(&[b"a"]), b"v1".to_vec(), vec![]).unwrap();
let err = b
.add_node(key(&[b"a"]), b"v2".to_vec(), vec![])
.unwrap_err();
assert!(matches!(err, IndexError::DuplicateKey(_)));
}
#[test]
fn builder_rejects_bad_key() {
let mut b = GraphIndexBuilder::new(0, 1);
let err = b
.add_node(key(&[b"with space"]), b"v".to_vec(), vec![])
.unwrap_err();
assert!(matches!(err, IndexError::BadKey(_)));
}
#[test]
fn builder_records_absent_references() {
let mut b = GraphIndexBuilder::new(1, 1);
b.add_node(key(&[b"a"]), b"v".to_vec(), vec![vec![key(&[b"missing"])]])
.unwrap();
assert_eq!(b.key_count(), 1);
assert!(b.nodes().contains_key(&key(&[b"missing"])));
assert!(b.nodes().get(&key(&[b"missing"])).unwrap().absent);
}
#[test]
fn builder_external_references_returns_unresolved_second_refs() {
let mut b = GraphIndexBuilder::new(2, 1);
b.add_node(
key(&[b"a"]),
b"v".to_vec(),
vec![vec![], vec![key(&[b"parent1"]), key(&[b"a"])]],
)
.unwrap();
let refs = b.external_references();
assert!(refs.contains(&key(&[b"parent1"])));
assert!(!refs.contains(&key(&[b"a"])));
}
#[test]
fn combined_iter_dedups_keys() {
let mut b1 = GraphIndexBuilder::new(0, 1);
b1.add_node(key(&[b"a"]), b"v-from-1".to_vec(), vec![])
.unwrap();
let mut b2 = GraphIndexBuilder::new(0, 1);
b2.add_node(key(&[b"a"]), b"v-from-2".to_vec(), vec![])
.unwrap();
b2.add_node(key(&[b"b"]), b"vb".to_vec(), vec![]).unwrap();
let combined = CombinedGraphIndex::from_indices(vec![Box::new(b1), Box::new(b2)]);
let mut all = combined.iter_all().unwrap();
all.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(all.len(), 2);
assert_eq!(all[0].1, b"v-from-1".to_vec());
}
#[test]
fn prefix_adapter_strips_keys_and_refs() {
let mut b = GraphIndexBuilder::new(1, 2);
b.add_node(
key(&[b"prefix", b"k1"]),
b"v1".to_vec(),
vec![vec![key(&[b"prefix", b"k2"])]],
)
.unwrap();
b.add_node(key(&[b"prefix", b"k2"]), b"v2".to_vec(), vec![vec![]])
.unwrap();
let adapter = GraphIndexPrefixAdapter::new(b, key(&[b"prefix"]), 1);
let mut entries = adapter.iter_all().unwrap();
entries.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(entries[0].0, key(&[b"k1"]));
assert_eq!(entries[0].2, vec![vec![key(&[b"k2"])]]);
}
#[test]
fn parsed_range_map_starts_empty() {
let m = ParsedRangeMap::new();
assert!(m.is_empty());
assert_eq!(m.len(), 0);
assert_eq!(m.byte_index(0), -1);
assert_eq!(m.key_index(&None), -1);
assert!(!m.is_parsed(0));
}
#[test]
fn parsed_range_map_first_insert() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 100, Some(key(&[b"k"])));
assert_eq!(m.len(), 1);
assert_eq!(m.byte_range(0), Some((0, 100)));
assert_eq!(m.key_range(0), Some((None, Some(key(&[b"k"])))));
}
#[test]
fn parsed_range_map_byte_index_matches_python_doctest() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, Some(key(&[b"a"])), 10, Some(key(&[b"b"])));
m.mark_parsed(11, Some(key(&[b"c"])), 12, Some(key(&[b"d"])));
assert_eq!(m.byte_index(0), 0);
assert_eq!(m.byte_index(10), 0);
assert_eq!(m.byte_index(11), 1);
assert_eq!(m.byte_index(12), 1);
}
#[test]
fn parsed_range_map_extend_lower_region() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 50, Some(key(&[b"k1"])));
m.mark_parsed(50, Some(key(&[b"k1"])), 100, Some(key(&[b"k2"])));
assert_eq!(m.len(), 1);
assert_eq!(m.byte_range(0), Some((0, 100)));
assert_eq!(m.key_range(0), Some((None, Some(key(&[b"k2"])))));
}
#[test]
fn parsed_range_map_extend_higher_region() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 30, None);
m.mark_parsed(60, Some(key(&[b"k2"])), 100, Some(key(&[b"k3"])));
m.mark_parsed(30, Some(key(&[b"k1"])), 60, Some(key(&[b"k2"])));
assert_eq!(m.len(), 1);
assert_eq!(m.byte_range(0), Some((0, 100)));
}
#[test]
fn parsed_range_map_combine_two_regions() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 50, Some(key(&[b"k1"])));
m.mark_parsed(60, Some(key(&[b"k2"])), 100, Some(key(&[b"k3"])));
assert_eq!(m.len(), 2);
m.mark_parsed(50, Some(key(&[b"k1"])), 60, Some(key(&[b"k2"])));
assert_eq!(m.len(), 1);
assert_eq!(m.byte_range(0), Some((0, 100)));
}
#[test]
fn parsed_range_map_disjoint_new_region() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 50, Some(key(&[b"k1"])));
m.mark_parsed(200, Some(key(&[b"k5"])), 300, Some(key(&[b"k6"])));
assert_eq!(m.len(), 2);
assert_eq!(m.byte_range(0), Some((0, 50)));
assert_eq!(m.byte_range(1), Some((200, 300)));
}
#[test]
fn parsed_range_map_is_parsed_inside_only() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(10, Some(key(&[b"a"])), 20, Some(key(&[b"b"])));
assert!(!m.is_parsed(9));
assert!(m.is_parsed(10));
assert!(m.is_parsed(15));
assert!(!m.is_parsed(20));
assert!(!m.is_parsed(100));
}
#[test]
fn parsed_range_map_key_index() {
let mut m = ParsedRangeMap::new();
m.mark_parsed(0, None, 50, Some(key(&[b"a"])));
m.mark_parsed(60, Some(key(&[b"b"])), 100, Some(key(&[b"c"])));
assert_eq!(m.key_index(&None), 0);
assert_eq!(m.key_index(&Some(key(&[b"a"]))), 0);
assert_eq!(m.key_index(&Some(key(&[b"b"]))), 1);
assert_eq!(m.key_index(&Some(key(&[b"e"]))), 1);
}
#[test]
fn key_matches_prefix_exact_and_wildcard() {
let p: KeyPrefix = vec![Some(b"a".to_vec()), None, Some(b"c".to_vec())];
assert!(key_matches_prefix(&p, &key(&[b"a", b"foo", b"c"])));
assert!(key_matches_prefix(&p, &key(&[b"a", b"bar", b"c"])));
assert!(!key_matches_prefix(&p, &key(&[b"a", b"foo", b"d"])));
assert!(!key_matches_prefix(&p, &key(&[b"x", b"foo", b"c"])));
}
#[test]
fn key_matches_prefix_rejects_length_mismatch() {
let p: KeyPrefix = vec![Some(b"a".to_vec())];
assert!(!key_matches_prefix(&p, &key(&[b"a", b"b"])));
assert!(!key_matches_prefix(&p, &key(&[])));
}
#[test]
fn key_matches_any_prefix_unions() {
let p1: KeyPrefix = vec![Some(b"a".to_vec()), None];
let p2: KeyPrefix = vec![Some(b"b".to_vec()), None];
assert!(key_matches_any_prefix(
&[p1.clone(), p2.clone()],
&key(&[b"a", b"x"])
));
assert!(key_matches_any_prefix(
&[p1.clone(), p2.clone()],
&key(&[b"b", b"y"])
));
assert!(!key_matches_any_prefix(&[p1, p2], &key(&[b"c", b"z"])));
}
#[test]
fn find_ancestors_single_index_walks_frontier() {
let mut b = GraphIndexBuilder::new(1, 1);
b.add_node(
key(&[b"key-1"]),
b"value".to_vec(),
vec![vec![key(&[b"key-2"])]],
)
.unwrap();
b.add_node(key(&[b"key-2"]), b"value".to_vec(), vec![vec![]])
.unwrap();
let mut parent_map: HashMap<IndexKey, Vec<IndexKey>> = HashMap::new();
let mut missing: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let search = b
.find_ancestors(&[key(&[b"key-1"])], 0, &mut parent_map, &mut missing)
.unwrap();
assert_eq!(
parent_map.get(&key(&[b"key-1"])),
Some(&vec![key(&[b"key-2"])])
);
assert!(missing.is_empty());
let expected: std::collections::HashSet<IndexKey> =
vec![key(&[b"key-2"])].into_iter().collect();
assert_eq!(search, expected);
let search2: Vec<IndexKey> = search.into_iter().collect();
let search = b
.find_ancestors(&search2, 0, &mut parent_map, &mut missing)
.unwrap();
assert_eq!(parent_map.get(&key(&[b"key-2"])), Some(&vec![]));
assert!(missing.is_empty());
assert!(search.is_empty());
}
#[test]
fn find_ancestors_records_missing_keys() {
let mut b = GraphIndexBuilder::new(1, 1);
b.add_node(
key(&[b"key-1"]),
b"value".to_vec(),
vec![vec![key(&[b"key-2"])]],
)
.unwrap();
b.add_node(key(&[b"key-2"]), b"value".to_vec(), vec![vec![]])
.unwrap();
let mut parent_map: HashMap<IndexKey, Vec<IndexKey>> = HashMap::new();
let mut missing: std::collections::HashSet<IndexKey> = std::collections::HashSet::new();
let search = b
.find_ancestors(
&[key(&[b"key-2"]), key(&[b"key-3"])],
0,
&mut parent_map,
&mut missing,
)
.unwrap();
assert_eq!(parent_map.get(&key(&[b"key-2"])), Some(&vec![]));
let expected: std::collections::HashSet<IndexKey> =
vec![key(&[b"key-3"])].into_iter().collect();
assert_eq!(missing, expected);
assert!(search.is_empty());
}
#[test]
fn combined_find_ancestry_across_indexes() {
let mut b1 = GraphIndexBuilder::new(1, 1);
b1.add_node(key(&[b"key-1"]), b"value".to_vec(), vec![vec![]])
.unwrap();
b1.add_node(
key(&[b"key-2"]),
b"value".to_vec(),
vec![vec![key(&[b"key-1"])]],
)
.unwrap();
let mut b2 = GraphIndexBuilder::new(1, 1);
b2.add_node(
key(&[b"key-3"]),
b"value".to_vec(),
vec![vec![key(&[b"key-2"])]],
)
.unwrap();
b2.add_node(
key(&[b"key-4"]),
b"value".to_vec(),
vec![vec![key(&[b"key-3"])]],
)
.unwrap();
let combined = CombinedGraphIndex::from_indices(vec![Box::new(b1), Box::new(b2)]);
let (pm, missing) = combined.find_ancestry(&[key(&[b"key-1"])], 0).unwrap();
assert_eq!(pm.get(&key(&[b"key-1"])), Some(&vec![]));
assert_eq!(pm.len(), 1);
assert!(missing.is_empty());
let (pm, missing) = combined.find_ancestry(&[key(&[b"key-3"])], 0).unwrap();
assert_eq!(pm.get(&key(&[b"key-1"])), Some(&vec![]));
assert_eq!(pm.get(&key(&[b"key-2"])), Some(&vec![key(&[b"key-1"])]));
assert_eq!(pm.get(&key(&[b"key-3"])), Some(&vec![key(&[b"key-2"])]));
assert!(missing.is_empty());
}
#[test]
fn combined_find_ancestry_missing_and_no_indexes() {
let combined = CombinedGraphIndex::new();
let (pm, missing) = combined.find_ancestry(&[key(&[b"key-1"])], 0).unwrap();
assert!(pm.is_empty());
let expected: std::collections::HashSet<IndexKey> =
vec![key(&[b"key-1"])].into_iter().collect();
assert_eq!(missing, expected);
let mut b = GraphIndexBuilder::new(1, 1);
b.add_node(
key(&[b"key-1"]),
b"value".to_vec(),
vec![vec![key(&[b"ghost"])]],
)
.unwrap();
let combined = CombinedGraphIndex::from_indices(vec![Box::new(b)]);
let (pm, missing) = combined.find_ancestry(&[key(&[b"key-1"])], 0).unwrap();
assert_eq!(pm.get(&key(&[b"key-1"])), Some(&vec![key(&[b"ghost"])]));
let expected: std::collections::HashSet<IndexKey> =
vec![key(&[b"ghost"])].into_iter().collect();
assert_eq!(missing, expected);
}
#[test]
fn combined_get_parent_map_fills_null_for_roots() {
let null: IndexKey = vec![b"null:".to_vec()];
let mut b = GraphIndexBuilder::new(1, 1);
b.add_node(
key(&[b"key-1"]),
b"v".to_vec(),
vec![vec![key(&[b"key-2"])]],
)
.unwrap();
b.add_node(key(&[b"key-2"]), b"v".to_vec(), vec![vec![]])
.unwrap();
let combined = CombinedGraphIndex::from_indices(vec![Box::new(b)]);
let pm = combined
.get_parent_map(&[key(&[b"key-1"]), key(&[b"key-2"]), null.clone()], &null)
.unwrap();
assert_eq!(pm.get(&key(&[b"key-1"])), Some(&vec![key(&[b"key-2"])]));
assert_eq!(pm.get(&key(&[b"key-2"])), Some(&vec![null.clone()]));
assert_eq!(pm.get(&null), Some(&vec![]));
}
#[test]
fn combined_iter_with_hits_then_move_to_front() {
let mk = |k: &[u8]| {
let mut b = GraphIndexBuilder::new(0, 1);
b.add_node(key(&[k]), b"v".to_vec(), vec![]).unwrap();
Box::new(b) as Box<dyn IndexLike + Send + Sync>
};
let mut combined =
CombinedGraphIndex::from_indices(vec![mk(b"k0"), mk(b"k1"), mk(b"k2"), mk(b"k3")]);
let (entries, hits) = combined
.iter_entries_with_hits(&[key(&[b"k3"]), key(&[b"k1"])])
.unwrap();
assert_eq!(entries.len(), 2);
assert_eq!(hits, vec![1, 3]);
combined.move_to_front(&hits);
let key_at = |c: &CombinedGraphIndex, i: usize| {
c.indices()[i].iter(&[]).ok();
c.indices()[i]
.iter_all()
.unwrap()
.into_iter()
.next()
.unwrap()
.0
};
assert_eq!(key_at(&combined, 0), key(&[b"k1"]));
assert_eq!(key_at(&combined, 1), key(&[b"k3"]));
assert_eq!(key_at(&combined, 2), key(&[b"k0"]));
assert_eq!(key_at(&combined, 3), key(&[b"k2"]));
}
}