use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use anyhow::{Result, anyhow};
use roaring::RoaringBitmap;
use znippy_common::GUNNAR_REACH_MODULE;
use znippy_common::arrow::array::{Array, BinaryArray, BinaryBuilder, StringArray, StringBuilder};
use znippy_common::arrow::datatypes::{DataType, Field, Schema};
use znippy_common::arrow::ipc::reader::StreamReader;
use znippy_common::arrow::record_batch::RecordBatch;
use znippy_common::read_reserved_section_bytes;
use crate::graph::CommitNode;
use crate::object::{GitObjectKind, tree_entries};
#[derive(Debug, Clone, Copy)]
pub struct ReachPolicy {
pub max_commits: usize,
}
impl Default for ReachPolicy {
fn default() -> Self {
Self { max_commits: 512 }
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ReachEntry {
pub commit: String,
pub bitmap: RoaringBitmap,
}
pub fn reach_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("commit_oid", DataType::Utf8, false),
Field::new("cardinality", DataType::UInt64, false),
Field::new("bitmap", DataType::Binary, false),
]))
}
pub struct ObjectFacts<'a> {
pub ordinal: &'a HashMap<String, u32>,
pub trees: &'a HashMap<String, Vec<u8>>,
pub oid_len: usize,
}
pub fn build_reach(
commits: &[CommitNode],
facts: &ObjectFacts<'_>,
policy: ReachPolicy,
) -> Vec<ReachEntry> {
let n = commits.len();
if n == 0 || policy.max_commits == 0 {
return Vec::new();
}
let pos: HashMap<&str, usize> = commits
.iter()
.enumerate()
.map(|(i, c)| (c.oid.as_str(), i))
.collect();
let selected = select_commits(commits, &pos, policy);
let keep: Vec<bool> = {
let mut k = vec![false; n];
for &i in &selected {
k[i] = true;
}
k
};
let mut pending: Vec<usize> = vec![0; n];
let mut parent_idx: Vec<Vec<usize>> = Vec::with_capacity(n);
for c in commits {
let mut ps: Vec<usize> = c.parents.iter().filter_map(|p| pos.get(p.as_str()).copied()).collect();
ps.sort_unstable();
ps.dedup();
for &p in &ps {
pending[p] += 1;
}
parent_idx.push(ps);
}
let mut tree_memo: HashMap<String, RoaringBitmap> = HashMap::new();
let mut live: HashMap<usize, RoaringBitmap> = HashMap::new();
for i in 0..n {
let c = &commits[i];
let mut bm = RoaringBitmap::new();
if let Some(&o) = facts.ordinal.get(c.oid.as_str()) {
bm.insert(o);
}
if let Some(t) = &c.tree {
if let Some(tc) = tree_closure(t, facts, &mut tree_memo) {
bm |= tc;
}
}
for &p in &parent_idx[i] {
if let Some(pb) = live.get(&p) {
bm |= pb;
}
pending[p] -= 1;
if pending[p] == 0 && !keep[p] {
live.remove(&p);
}
}
if pending[i] > 0 || keep[i] {
live.insert(i, bm);
}
}
let mut out: Vec<ReachEntry> = Vec::with_capacity(selected.len());
for &i in &selected {
let bitmap = live.remove(&i).unwrap_or_else(|| {
panic!(
"commit {} was selected for a reachability bitmap and its bitmap is not live at \
the end of the build; an empty one here would under-send a clone",
commits[i].oid
)
});
out.push(ReachEntry {
commit: commits[i].oid.clone(),
bitmap,
});
}
out
}
fn select_commits(
commits: &[CommitNode],
pos: &HashMap<&str, usize>,
policy: ReachPolicy,
) -> Vec<usize> {
let n = commits.len();
let mut has_child = vec![false; n];
for c in commits {
for p in &c.parents {
if let Some(&pi) = pos.get(p.as_str()) {
has_child[pi] = true;
}
}
}
let mut chosen: Vec<usize> = (0..n).filter(|&i| !has_child[i]).collect();
chosen.truncate(policy.max_commits);
if chosen.len() < policy.max_commits {
let room = policy.max_commits - chosen.len();
let rest: Vec<usize> = (0..n).filter(|&i| has_child[i]).collect();
if !rest.is_empty() {
let stride = rest.len().div_ceil(room).max(1);
for &i in rest.iter().step_by(stride).take(room) {
chosen.push(i);
}
}
}
chosen.sort_unstable();
chosen.dedup();
chosen
}
pub(crate) fn tree_closure<'m>(
root: &str,
facts: &ObjectFacts<'_>,
memo: &'m mut HashMap<String, RoaringBitmap>,
) -> Option<&'m RoaringBitmap> {
if memo.contains_key(root) {
return memo.get(root);
}
let mut stack: Vec<(String, bool)> = vec![(root.to_string(), false)];
let mut visiting: std::collections::HashSet<String> = std::collections::HashSet::new();
while let Some((oid, expanded)) = stack.pop() {
if memo.contains_key(&oid) {
visiting.remove(&oid);
continue;
}
let Some(payload) = facts.trees.get(&oid) else {
let mut bm = RoaringBitmap::new();
if let Some(&o) = facts.ordinal.get(oid.as_str()) {
bm.insert(o);
}
memo.insert(oid, bm);
continue;
};
let children: Vec<String> = tree_entries(payload, facts.oid_len)
.iter()
.map(|e| hex::encode(e.oid))
.collect();
if !expanded {
visiting.insert(oid.clone());
let unresolved: Vec<String> = children
.iter()
.filter(|c| !memo.contains_key(*c) && !visiting.contains(*c))
.cloned()
.collect();
if !unresolved.is_empty() {
stack.push((oid, true));
for c in unresolved {
stack.push((c, false));
}
continue;
}
}
let mut bm = RoaringBitmap::new();
if let Some(&o) = facts.ordinal.get(oid.as_str()) {
bm.insert(o);
}
for c in &children {
if let Some(cb) = memo.get(c) {
bm |= cb;
} else if let Some(&o) = facts.ordinal.get(c.as_str()) {
bm.insert(o);
}
}
memo.insert(oid.clone(), bm);
visiting.remove(&oid);
}
memo.get(root)
}
pub(crate) fn accumulate(
tip: &str,
by_commit: &HashMap<&str, &RoaringBitmap>,
graph: &HashMap<&str, &CommitNode>,
facts: &ObjectFacts<'_>,
acc: &mut RoaringBitmap,
) {
let mut unbitmapped: Vec<&str> = Vec::new();
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
let mut stack: Vec<&str> = vec![tip];
while let Some(oid) = stack.pop() {
if !seen.insert(oid) {
continue;
}
if let Some(bm) = by_commit.get(oid) {
*acc |= *bm;
continue;
}
let Some(node) = graph.get(oid) else {
if let Some(&o) = facts.ordinal.get(oid) {
acc.insert(o);
}
continue;
};
unbitmapped.push(oid);
for p in &node.parents {
stack.push(p.as_str());
}
}
for oid in unbitmapped {
if let Some(&o) = facts.ordinal.get(oid) {
acc.insert(o);
}
let Some(node) = graph.get(oid) else { continue };
if let Some(tree) = &node.tree {
accumulate_tree(tree.as_str(), facts, acc);
}
}
}
fn accumulate_tree(root: &str, facts: &ObjectFacts<'_>, acc: &mut RoaringBitmap) {
let mut stack: Vec<String> = vec![root.to_string()];
while let Some(oid) = stack.pop() {
let Some(&o) = facts.ordinal.get(oid.as_str()) else {
continue;
};
if !acc.insert(o) {
continue;
}
if let Some(payload) = facts.trees.get(oid.as_str()) {
for e in tree_entries(payload, facts.oid_len) {
stack.push(hex::encode(e.oid));
}
}
}
}
pub fn needs_payload(kind: GitObjectKind) -> bool {
matches!(kind, GitObjectKind::Commit | GitObjectKind::Tree)
}
pub fn build_reach_batch(entries: &[ReachEntry]) -> Result<RecordBatch> {
use znippy_common::arrow::array::UInt64Builder;
let n = entries.len();
let mut oid_b = StringBuilder::with_capacity(n, n * 64);
let mut card_b = UInt64Builder::with_capacity(n);
let mut bm_b = BinaryBuilder::with_capacity(n, n * 128);
for e in entries {
oid_b.append_value(&e.commit);
card_b.append_value(e.bitmap.len());
let mut buf = Vec::new();
e.bitmap
.serialize_into(&mut buf)
.map_err(|err| anyhow!("reach bitmap serialize: {err}"))?;
bm_b.append_value(&buf);
}
RecordBatch::try_new(
reach_schema(),
vec![Arc::new(oid_b.finish()), Arc::new(card_b.finish()), Arc::new(bm_b.finish())],
)
.map_err(|e| anyhow!("reach batch: {e}"))
}
pub fn decode_reach(bytes: &[u8]) -> Result<Vec<ReachEntry>> {
let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
.map_err(|e| anyhow!("reach reader: {e}"))?;
let mut out = Vec::new();
for batch in reader {
let batch = batch.map_err(|e| anyhow!("reach batch read: {e}"))?;
let oids = batch
.column_by_name("commit_oid")
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.ok_or_else(|| anyhow!("reach: no `commit_oid` column"))?;
let bms = batch
.column_by_name("bitmap")
.and_then(|c| c.as_any().downcast_ref::<BinaryArray>())
.ok_or_else(|| anyhow!("reach: no `bitmap` column"))?;
for i in 0..batch.num_rows() {
let bitmap = RoaringBitmap::deserialize_from(bms.value(i))
.map_err(|e| anyhow!("reach bitmap deserialize: {e}"))?;
out.push(ReachEntry { commit: oids.value(i).to_string(), bitmap });
}
}
Ok(out)
}
pub fn read_reach(archive: &Path) -> Result<Option<Vec<ReachEntry>>> {
match read_reserved_section_bytes(archive, GUNNAR_REACH_MODULE)? {
Some(b) => Ok(Some(decode_reach(&b)?)),
None => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::graph::assign_generations;
fn hexid(c: char) -> String {
std::iter::repeat_n(c, 40).collect()
}
#[allow(clippy::type_complexity)]
fn tiny_repo() -> (Vec<CommitNode>, HashMap<String, u32>, HashMap<String, Vec<u8>>) {
let blob1 = hexid('1');
let blob2 = hexid('2');
let tree1 = hexid('3');
let tree2 = hexid('4');
let c1 = hexid('5');
let c2 = hexid('6');
let mut trees: HashMap<String, Vec<u8>> = HashMap::new();
let mut t1 = Vec::new();
t1.extend_from_slice(b"100644 a\0");
t1.extend_from_slice(&hex::decode(&blob1).unwrap());
trees.insert(tree1.clone(), t1);
let mut t2 = Vec::new();
t2.extend_from_slice(b"100644 a\0");
t2.extend_from_slice(&hex::decode(&blob1).unwrap());
t2.extend_from_slice(b"100644 b\0");
t2.extend_from_slice(&hex::decode(&blob2).unwrap());
trees.insert(tree2.clone(), t2);
let ordinal: HashMap<String, u32> = [
(blob1, 0u32),
(blob2, 1),
(tree1, 2),
(tree2, 3),
(c1.clone(), 4),
(c2.clone(), 5),
]
.into_iter()
.collect();
let commits = assign_generations(vec![
CommitNode {
oid: c1.clone(),
parents: vec![],
tree: Some(hexid('3')),
committer_time: Some(1),
generation: 0,
},
CommitNode {
oid: c2.clone(),
parents: vec![c1.clone()],
tree: Some(hexid('4')),
committer_time: Some(2),
generation: 0,
},
]);
(commits, ordinal, trees)
}
#[test]
fn bitmap_contains_exactly_the_reachable_objects() {
let (commits, ordinal, trees) = tiny_repo();
let facts = ObjectFacts { ordinal: &ordinal, trees: &trees, oid_len: 20 };
let entries = build_reach(&commits, &facts, ReachPolicy { max_commits: 16 });
let by_commit: HashMap<&str, &RoaringBitmap> =
entries.iter().map(|e| (e.commit.as_str(), &e.bitmap)).collect();
let b1: Vec<u32> = by_commit[hexid('5').as_str()].iter().collect();
assert_eq!(b1, vec![0, 2, 4], "c1 must not reach blob2/tree2/c2");
let b2: Vec<u32> = by_commit[hexid('6').as_str()].iter().collect();
assert_eq!(b2, vec![0, 1, 2, 3, 4, 5]);
}
#[test]
fn want_minus_have_is_an_andnot() {
let (commits, ordinal, trees) = tiny_repo();
let facts = ObjectFacts { ordinal: &ordinal, trees: &trees, oid_len: 20 };
let entries = build_reach(&commits, &facts, ReachPolicy { max_commits: 16 });
let by: HashMap<&str, &RoaringBitmap> =
entries.iter().map(|e| (e.commit.as_str(), &e.bitmap)).collect();
let want = by[hexid('6').as_str()].clone();
let have = by[hexid('5').as_str()].clone();
let delta: Vec<u32> = (want - have).iter().collect();
assert_eq!(delta, vec![1, 3, 5]);
}
#[test]
fn selection_takes_tips_first_and_respects_the_cap() {
let (commits, ordinal, trees) = tiny_repo();
let facts = ObjectFacts { ordinal: &ordinal, trees: &trees, oid_len: 20 };
let entries = build_reach(&commits, &facts, ReachPolicy { max_commits: 1 });
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].commit, hexid('6'), "the tip, not the root");
}
#[test]
fn batch_roundtrips_through_arrow_ipc() {
let (commits, ordinal, trees) = tiny_repo();
let facts = ObjectFacts { ordinal: &ordinal, trees: &trees, oid_len: 20 };
let entries = build_reach(&commits, &facts, ReachPolicy::default());
let batch = build_reach_batch(&entries).unwrap();
let mut buf = Vec::new();
{
let mut w = znippy_common::arrow::ipc::writer::StreamWriter::try_new(
&mut buf,
&reach_schema(),
)
.unwrap();
w.write(&batch).unwrap();
w.finish().unwrap();
}
let back = decode_reach(&buf).unwrap();
assert_eq!(back, entries);
}
#[test]
fn a_corrupt_tree_cycle_terminates() {
let t = hexid('a');
let mut payload = Vec::new();
payload.extend_from_slice(b"40000 self\0");
payload.extend_from_slice(&hex::decode(&t).unwrap());
let trees: HashMap<String, Vec<u8>> = [(t.clone(), payload)].into_iter().collect();
let ordinal: HashMap<String, u32> = [(t.clone(), 0u32), (hexid('b'), 1)].into_iter().collect();
let facts = ObjectFacts { ordinal: &ordinal, trees: &trees, oid_len: 20 };
let commits = assign_generations(vec![CommitNode {
oid: hexid('b'),
parents: vec![],
tree: Some(t),
committer_time: None,
generation: 0,
}]);
let entries = build_reach(&commits, &facts, ReachPolicy::default());
assert_eq!(entries.len(), 1);
let bits: Vec<u32> = entries[0].bitmap.iter().collect();
assert_eq!(bits, vec![0, 1]);
}
}