use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use anyhow::{Result, anyhow};
use znippy_common::GUNNAR_GRAPH_MODULE;
use znippy_common::arrow::array::{
Array, Int64Array, Int64Builder, ListBuilder, StringArray, StringBuilder, UInt32Array,
UInt32Builder,
};
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;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommitNode {
pub oid: String,
pub parents: Vec<String>,
pub tree: Option<String>,
pub committer_time: Option<i64>,
pub generation: u32,
}
pub fn graph_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("oid", DataType::Utf8, false),
Field::new(
"parents",
DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
false,
),
Field::new("tree", DataType::Utf8, true),
Field::new("committer_time", DataType::Int64, true),
Field::new("generation", DataType::UInt32, false),
]))
}
pub fn assign_generations(mut nodes: Vec<CommitNode>) -> Vec<CommitNode> {
let index: HashMap<&str, usize> = nodes
.iter()
.enumerate()
.map(|(i, n)| (n.oid.as_str(), i))
.collect();
let n = nodes.len();
let mut parent_ids: Vec<Vec<usize>> = Vec::with_capacity(n);
let mut children: Vec<Vec<usize>> = vec![Vec::new(); n];
let mut indeg: Vec<usize> = vec![0; n];
for (i, node) in nodes.iter().enumerate() {
let mut ps = Vec::new();
for p in &node.parents {
if let Some(&pi) = index.get(p.as_str()) {
if pi != i {
ps.push(pi);
}
}
}
ps.sort_unstable();
ps.dedup();
indeg[i] = ps.len();
for &pi in &ps {
children[pi].push(i);
}
parent_ids.push(ps);
}
let mut generation = vec![1u32; n];
let mut queue: Vec<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
let mut head = 0usize;
let mut settled = 0usize;
while head < queue.len() {
let i = queue[head];
head += 1;
settled += 1;
let g = generation[i];
for &c in &children[i] {
generation[c] = generation[c].max(g.saturating_add(1));
indeg[c] -= 1;
if indeg[c] == 0 {
queue.push(c);
}
}
}
debug_assert!(settled == n || n == 0);
for (i, node) in nodes.iter_mut().enumerate() {
node.generation = generation[i];
}
nodes.sort_by(|a, b| a.generation.cmp(&b.generation).then_with(|| a.oid.cmp(&b.oid)));
nodes
}
pub fn build_graph_batch(nodes: &[CommitNode]) -> Result<RecordBatch> {
let n = nodes.len();
let mut oid_b = StringBuilder::with_capacity(n, n * 64);
let mut parents_b = ListBuilder::new(StringBuilder::new());
let mut tree_b = StringBuilder::with_capacity(n, n * 64);
let mut time_b = Int64Builder::with_capacity(n);
let mut gen_b = UInt32Builder::with_capacity(n);
for node in nodes {
oid_b.append_value(&node.oid);
for p in &node.parents {
parents_b.values().append_value(p);
}
parents_b.append(true);
match &node.tree {
Some(t) => tree_b.append_value(t),
None => tree_b.append_null(),
}
match node.committer_time {
Some(t) => time_b.append_value(t),
None => time_b.append_null(),
}
gen_b.append_value(node.generation);
}
RecordBatch::try_new(
graph_schema(),
vec![
Arc::new(oid_b.finish()),
Arc::new(parents_b.finish()),
Arc::new(tree_b.finish()),
Arc::new(time_b.finish()),
Arc::new(gen_b.finish()),
],
)
.map_err(|e| anyhow!("commit-graph batch: {e}"))
}
pub fn decode_graph(bytes: &[u8]) -> Result<Vec<CommitNode>> {
let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
.map_err(|e| anyhow!("commit-graph reader: {e}"))?;
let mut out = Vec::new();
for batch in reader {
let batch = batch.map_err(|e| anyhow!("commit-graph batch read: {e}"))?;
let oids = col::<StringArray>(&batch, "oid")?;
let parents = batch
.column_by_name("parents")
.ok_or_else(|| anyhow!("commit-graph: no `parents` column"))?
.as_any()
.downcast_ref::<znippy_common::arrow::array::ListArray>()
.ok_or_else(|| anyhow!("commit-graph: `parents` is not a list"))?;
let trees = col::<StringArray>(&batch, "tree")?;
let times = col::<Int64Array>(&batch, "committer_time")?;
let gens = col::<UInt32Array>(&batch, "generation")?;
for i in 0..batch.num_rows() {
let plist = parents.value(i);
let plist = plist
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| anyhow!("commit-graph: parent items are not strings"))?;
out.push(CommitNode {
oid: oids.value(i).to_string(),
parents: (0..plist.len()).map(|j| plist.value(j).to_string()).collect(),
tree: (!trees.is_null(i)).then(|| trees.value(i).to_string()),
committer_time: (!times.is_null(i)).then(|| times.value(i)),
generation: gens.value(i),
});
}
}
Ok(out)
}
pub fn read_graph(archive: &Path) -> Result<Option<Vec<CommitNode>>> {
match read_reserved_section_bytes(archive, GUNNAR_GRAPH_MODULE)? {
Some(b) => Ok(Some(decode_graph(&b)?)),
None => Ok(None),
}
}
fn col<'a, T: Array + 'static>(batch: &'a RecordBatch, name: &str) -> Result<&'a T> {
batch
.column_by_name(name)
.ok_or_else(|| anyhow!("commit-graph: no `{name}` column"))?
.as_any()
.downcast_ref::<T>()
.ok_or_else(|| anyhow!("commit-graph: `{name}` has an unexpected type"))
}
#[cfg(test)]
mod tests {
use super::*;
fn node(oid: &str, parents: &[&str]) -> CommitNode {
CommitNode {
oid: oid.to_string(),
parents: parents.iter().map(|s| s.to_string()).collect(),
tree: Some(format!("t{}", &oid[1..])),
committer_time: Some(1_700_000_000),
generation: 0,
}
}
fn hexid(c: char) -> String {
std::iter::repeat_n(c, 40).collect()
}
#[test]
fn generations_are_one_plus_the_max_parent() {
let a = hexid('a');
let b = hexid('b');
let c = hexid('c');
let d = hexid('d');
let nodes = assign_generations(vec![
node(&d, &[&b, &c]),
node(&b, &[&a]),
node(&c, &[&a]),
node(&a, &[]),
]);
let g: HashMap<&str, u32> = nodes.iter().map(|n| (n.oid.as_str(), n.generation)).collect();
assert_eq!(g[a.as_str()], 1, "root");
assert_eq!(g[b.as_str()], 2);
assert_eq!(g[c.as_str()], 2);
assert_eq!(g[d.as_str()], 3, "merge is 1 + max(2,2)");
let gens: Vec<u32> = nodes.iter().map(|n| n.generation).collect();
assert!(gens.windows(2).all(|w| w[0] <= w[1]), "not topologically ordered: {gens:?}");
}
#[test]
fn a_parent_outside_the_archive_does_not_inflate_the_generation() {
let a = hexid('a');
let missing = hexid('f');
let nodes = assign_generations(vec![node(&a, &[&missing])]);
assert_eq!(nodes[0].generation, 1, "a commit whose parent is not in the archive is a root here");
}
#[test]
fn deep_chain_does_not_blow_the_stack() {
let ids: Vec<String> = (0..50_000u32).map(|i| format!("{i:040x}")).collect();
let mut nodes = Vec::with_capacity(ids.len());
for (i, id) in ids.iter().enumerate() {
let parents: Vec<&str> = if i == 0 { vec![] } else { vec![ids[i - 1].as_str()] };
nodes.push(node(id, &parents));
}
let out = assign_generations(nodes);
let max = out.iter().map(|n| n.generation).max().unwrap();
assert_eq!(max, 50_000);
}
#[test]
fn batch_roundtrips_through_arrow_ipc() {
let a = hexid('a');
let b = hexid('b');
let nodes = assign_generations(vec![node(&b, &[&a]), node(&a, &[])]);
let batch = build_graph_batch(&nodes).unwrap();
let mut buf = Vec::new();
{
let mut w = znippy_common::arrow::ipc::writer::StreamWriter::try_new(
&mut buf,
&graph_schema(),
)
.unwrap();
w.write(&batch).unwrap();
w.finish().unwrap();
}
let back = decode_graph(&buf).unwrap();
assert_eq!(back, nodes);
assert_eq!(back[1].parents, vec![a.clone()]);
assert_eq!(back[1].generation, 2);
}
}