1use std::collections::HashMap;
21use std::path::Path;
22use std::sync::Arc;
23
24use anyhow::{Result, anyhow};
25use znippy_common::GUNNAR_GRAPH_MODULE;
26use znippy_common::arrow::array::{
27 Array, Int64Array, Int64Builder, ListBuilder, StringArray, StringBuilder, UInt32Array,
28 UInt32Builder,
29};
30use znippy_common::arrow::datatypes::{DataType, Field, Schema};
31use znippy_common::arrow::ipc::reader::StreamReader;
32use znippy_common::arrow::record_batch::RecordBatch;
33use znippy_common::read_reserved_section_bytes;
34
35#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct CommitNode {
38 pub oid: String,
39 pub parents: Vec<String>,
40 pub tree: Option<String>,
41 pub committer_time: Option<i64>,
43 pub generation: u32,
45}
46
47pub fn graph_schema() -> Arc<Schema> {
48 Arc::new(Schema::new(vec![
49 Field::new("oid", DataType::Utf8, false),
50 Field::new(
51 "parents",
52 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
53 false,
54 ),
55 Field::new("tree", DataType::Utf8, true),
56 Field::new("committer_time", DataType::Int64, true),
57 Field::new("generation", DataType::UInt32, false),
58 ]))
59}
60
61pub fn assign_generations(mut nodes: Vec<CommitNode>) -> Vec<CommitNode> {
68 let index: HashMap<&str, usize> = nodes
69 .iter()
70 .enumerate()
71 .map(|(i, n)| (n.oid.as_str(), i))
72 .collect();
73
74 let n = nodes.len();
76 let mut parent_ids: Vec<Vec<usize>> = Vec::with_capacity(n);
77 let mut children: Vec<Vec<usize>> = vec![Vec::new(); n];
78 let mut indeg: Vec<usize> = vec![0; n];
79 for (i, node) in nodes.iter().enumerate() {
80 let mut ps = Vec::new();
81 for p in &node.parents {
82 if let Some(&pi) = index.get(p.as_str()) {
83 if pi != i {
84 ps.push(pi);
85 }
86 }
87 }
88 ps.sort_unstable();
89 ps.dedup();
90 indeg[i] = ps.len();
91 for &pi in &ps {
92 children[pi].push(i);
93 }
94 parent_ids.push(ps);
95 }
96
97 let mut generation = vec![1u32; n];
98 let mut queue: Vec<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
99 let mut head = 0usize;
100 let mut settled = 0usize;
101 while head < queue.len() {
102 let i = queue[head];
103 head += 1;
104 settled += 1;
105 let g = generation[i];
106 for &c in &children[i] {
107 generation[c] = generation[c].max(g.saturating_add(1));
108 indeg[c] -= 1;
109 if indeg[c] == 0 {
110 queue.push(c);
111 }
112 }
113 }
114 debug_assert!(settled == n || n == 0);
118
119 for (i, node) in nodes.iter_mut().enumerate() {
120 node.generation = generation[i];
121 }
122 nodes.sort_by(|a, b| a.generation.cmp(&b.generation).then_with(|| a.oid.cmp(&b.oid)));
123 nodes
124}
125
126pub fn build_graph_batch(nodes: &[CommitNode]) -> Result<RecordBatch> {
128 let n = nodes.len();
129 let mut oid_b = StringBuilder::with_capacity(n, n * 64);
130 let mut parents_b = ListBuilder::new(StringBuilder::new());
131 let mut tree_b = StringBuilder::with_capacity(n, n * 64);
132 let mut time_b = Int64Builder::with_capacity(n);
133 let mut gen_b = UInt32Builder::with_capacity(n);
134
135 for node in nodes {
136 oid_b.append_value(&node.oid);
137 for p in &node.parents {
138 parents_b.values().append_value(p);
139 }
140 parents_b.append(true);
141 match &node.tree {
142 Some(t) => tree_b.append_value(t),
143 None => tree_b.append_null(),
144 }
145 match node.committer_time {
146 Some(t) => time_b.append_value(t),
147 None => time_b.append_null(),
148 }
149 gen_b.append_value(node.generation);
150 }
151
152 RecordBatch::try_new(
153 graph_schema(),
154 vec![
155 Arc::new(oid_b.finish()),
156 Arc::new(parents_b.finish()),
157 Arc::new(tree_b.finish()),
158 Arc::new(time_b.finish()),
159 Arc::new(gen_b.finish()),
160 ],
161 )
162 .map_err(|e| anyhow!("commit-graph batch: {e}"))
163}
164
165pub fn encode_graph(nodes: &[CommitNode]) -> Result<Vec<u8>> {
170 use znippy_common::arrow::ipc::writer::StreamWriter;
171 let batch = build_graph_batch(nodes)?;
172 let mut buf = Vec::new();
173 {
174 let mut w = StreamWriter::try_new(&mut buf, &graph_schema())
175 .map_err(|e| anyhow!("commit-graph writer: {e}"))?;
176 w.write(&batch).map_err(|e| anyhow!("commit-graph write: {e}"))?;
177 w.finish().map_err(|e| anyhow!("commit-graph finish: {e}"))?;
178 }
179 Ok(buf)
180}
181
182pub fn decode_graph(bytes: &[u8]) -> Result<Vec<CommitNode>> {
184 let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
185 .map_err(|e| anyhow!("commit-graph reader: {e}"))?;
186 let mut out = Vec::new();
187 for batch in reader {
188 let batch = batch.map_err(|e| anyhow!("commit-graph batch read: {e}"))?;
189 let oids = col::<StringArray>(&batch, "oid")?;
190 let parents = batch
191 .column_by_name("parents")
192 .ok_or_else(|| anyhow!("commit-graph: no `parents` column"))?
193 .as_any()
194 .downcast_ref::<znippy_common::arrow::array::ListArray>()
195 .ok_or_else(|| anyhow!("commit-graph: `parents` is not a list"))?;
196 let trees = col::<StringArray>(&batch, "tree")?;
197 let times = col::<Int64Array>(&batch, "committer_time")?;
198 let gens = col::<UInt32Array>(&batch, "generation")?;
199
200 for i in 0..batch.num_rows() {
201 let plist = parents.value(i);
202 let plist = plist
203 .as_any()
204 .downcast_ref::<StringArray>()
205 .ok_or_else(|| anyhow!("commit-graph: parent items are not strings"))?;
206 out.push(CommitNode {
207 oid: oids.value(i).to_string(),
208 parents: (0..plist.len()).map(|j| plist.value(j).to_string()).collect(),
209 tree: (!trees.is_null(i)).then(|| trees.value(i).to_string()),
210 committer_time: (!times.is_null(i)).then(|| times.value(i)),
211 generation: gens.value(i),
212 });
213 }
214 }
215 Ok(out)
216}
217
218pub fn read_graph(archive: &Path) -> Result<Option<Vec<CommitNode>>> {
220 match read_reserved_section_bytes(archive, GUNNAR_GRAPH_MODULE)? {
221 Some(b) => Ok(Some(decode_graph(&b)?)),
222 None => Ok(None),
223 }
224}
225
226fn col<'a, T: Array + 'static>(batch: &'a RecordBatch, name: &str) -> Result<&'a T> {
227 batch
228 .column_by_name(name)
229 .ok_or_else(|| anyhow!("commit-graph: no `{name}` column"))?
230 .as_any()
231 .downcast_ref::<T>()
232 .ok_or_else(|| anyhow!("commit-graph: `{name}` has an unexpected type"))
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238
239 fn node(oid: &str, parents: &[&str]) -> CommitNode {
240 CommitNode {
241 oid: oid.to_string(),
242 parents: parents.iter().map(|s| s.to_string()).collect(),
243 tree: Some(format!("t{}", &oid[1..])),
244 committer_time: Some(1_700_000_000),
245 generation: 0,
246 }
247 }
248
249 fn hexid(c: char) -> String {
250 std::iter::repeat_n(c, 40).collect()
251 }
252
253 #[test]
254 fn generations_are_one_plus_the_max_parent() {
255 let a = hexid('a');
256 let b = hexid('b');
257 let c = hexid('c');
258 let d = hexid('d');
259 let nodes = assign_generations(vec![
261 node(&d, &[&b, &c]),
262 node(&b, &[&a]),
263 node(&c, &[&a]),
264 node(&a, &[]),
265 ]);
266 let g: HashMap<&str, u32> = nodes.iter().map(|n| (n.oid.as_str(), n.generation)).collect();
267 assert_eq!(g[a.as_str()], 1, "root");
268 assert_eq!(g[b.as_str()], 2);
269 assert_eq!(g[c.as_str()], 2);
270 assert_eq!(g[d.as_str()], 3, "merge is 1 + max(2,2)");
271
272 let gens: Vec<u32> = nodes.iter().map(|n| n.generation).collect();
274 assert!(gens.windows(2).all(|w| w[0] <= w[1]), "not topologically ordered: {gens:?}");
275 }
276
277 #[test]
278 fn a_parent_outside_the_archive_does_not_inflate_the_generation() {
279 let a = hexid('a');
280 let missing = hexid('f');
281 let nodes = assign_generations(vec![node(&a, &[&missing])]);
282 assert_eq!(nodes[0].generation, 1, "a commit whose parent is not in the archive is a root here");
283 }
284
285 #[test]
286 fn deep_chain_does_not_blow_the_stack() {
287 let ids: Vec<String> = (0..50_000u32).map(|i| format!("{i:040x}")).collect();
289 let mut nodes = Vec::with_capacity(ids.len());
290 for (i, id) in ids.iter().enumerate() {
291 let parents: Vec<&str> = if i == 0 { vec![] } else { vec![ids[i - 1].as_str()] };
292 nodes.push(node(id, &parents));
293 }
294 let out = assign_generations(nodes);
295 let max = out.iter().map(|n| n.generation).max().unwrap();
296 assert_eq!(max, 50_000);
297 }
298
299 #[test]
300 fn batch_roundtrips_through_arrow_ipc() {
301 let a = hexid('a');
302 let b = hexid('b');
303 let nodes = assign_generations(vec![node(&b, &[&a]), node(&a, &[])]);
304 let batch = build_graph_batch(&nodes).unwrap();
305
306 let mut buf = Vec::new();
307 {
308 let mut w = znippy_common::arrow::ipc::writer::StreamWriter::try_new(
309 &mut buf,
310 &graph_schema(),
311 )
312 .unwrap();
313 w.write(&batch).unwrap();
314 w.finish().unwrap();
315 }
316 let back = decode_graph(&buf).unwrap();
317 assert_eq!(back, nodes);
318 assert_eq!(back[1].parents, vec![a.clone()]);
319 assert_eq!(back[1].generation, 2);
320 }
321}