use std::collections::{HashSet, VecDeque};
use objects::{
object::{ChangeId, ContentHash, EntryType},
store::ObjectStore,
};
use serde::{Deserialize, Serialize};
use crate::{ProtocolError, Result};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ObjectId {
Hash(ContentHash),
ChangeId(ChangeId),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ObjectInfo {
pub id: ObjectId,
pub obj_type: ObjectType,
pub size: u64,
pub delta_base: Option<ContentHash>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct PlannedObject {
pub id: ObjectId,
pub obj_type: ObjectType,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ObjectType {
Blob,
Tree,
State,
Action,
}
#[derive(Debug, Clone, Default)]
pub struct StateClosureOptions {
pub depth: Option<u32>,
pub exclude_states: Vec<ChangeId>,
}
pub fn enumerate_state_closure(
store: &dyn ObjectStore,
state_id: ChangeId,
) -> Result<Vec<ObjectInfo>> {
enumerate_state_closure_with_options(store, state_id, StateClosureOptions::default())
}
pub fn enumerate_state_closure_with_options(
store: &dyn ObjectStore,
state_id: ChangeId,
options: StateClosureOptions,
) -> Result<Vec<ObjectInfo>> {
let (excluded_states, excluded_hashes) = collect_excluded(store, &options.exclude_states)?;
let mut out = Vec::new();
let mut seen_states: HashSet<ChangeId> = HashSet::new();
let mut seen_hashes: HashSet<ContentHash> = HashSet::new();
let mut queue: VecDeque<(ChangeId, u32)> = VecDeque::new();
queue.push_back((state_id, 0));
while let Some((id, depth)) = queue.pop_front() {
if excluded_states.contains(&id) {
continue;
}
if !seen_states.insert(id) {
continue;
}
let state = store
.get_state(&id)?
.ok_or_else(|| ProtocolError::ObjectNotFound(id.to_string()))?;
let state_bytes = rmp_serde::to_vec_named(&state)?;
out.push(ObjectInfo {
id: ObjectId::ChangeId(id),
obj_type: ObjectType::State,
size: state_bytes.len() as u64,
delta_base: None,
});
if options.depth.map(|max| depth < max).unwrap_or(true) {
for parent in &state.parents {
queue.push_back((*parent, depth + 1));
}
}
enumerate_tree_closure_filtered(
store,
state.tree,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
if let Some(provenance_root) = state.provenance {
enumerate_tree_closure_filtered(
store,
provenance_root,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
}
if let Some(context_root) = state.context {
enumerate_tree_closure_filtered(
store,
context_root,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
}
}
Ok(out)
}
pub fn enumerate_state_closure_plan(
store: &dyn ObjectStore,
state_id: ChangeId,
) -> Result<Vec<PlannedObject>> {
enumerate_state_closure_plan_with_options(store, state_id, StateClosureOptions::default())
}
pub fn enumerate_state_closure_plan_with_options(
store: &dyn ObjectStore,
state_id: ChangeId,
options: StateClosureOptions,
) -> Result<Vec<PlannedObject>> {
let (excluded_states, excluded_hashes) = collect_excluded(store, &options.exclude_states)?;
let mut out = Vec::new();
let mut seen_states: HashSet<ChangeId> = HashSet::new();
let mut seen_hashes: HashSet<ContentHash> = HashSet::new();
let mut queue: VecDeque<(ChangeId, u32)> = VecDeque::new();
queue.push_back((state_id, 0));
while let Some((id, depth)) = queue.pop_front() {
if excluded_states.contains(&id) {
continue;
}
if !seen_states.insert(id) {
continue;
}
let state = store
.get_state(&id)?
.ok_or_else(|| ProtocolError::ObjectNotFound(id.to_string()))?;
out.push(PlannedObject {
id: ObjectId::ChangeId(id),
obj_type: ObjectType::State,
});
if options.depth.map(|max| depth < max).unwrap_or(true) {
for parent in &state.parents {
queue.push_back((*parent, depth + 1));
}
}
enumerate_tree_plan_filtered(
store,
state.tree,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
if let Some(provenance_root) = state.provenance {
enumerate_tree_plan_filtered(
store,
provenance_root,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
}
if let Some(context_root) = state.context {
enumerate_tree_plan_filtered(
store,
context_root,
&excluded_hashes,
&mut seen_hashes,
&mut out,
)?;
}
}
Ok(out)
}
fn enumerate_tree_closure_filtered(
store: &dyn ObjectStore,
tree_hash: ContentHash,
excluded: &HashSet<ContentHash>,
seen: &mut HashSet<ContentHash>,
out: &mut Vec<ObjectInfo>,
) -> Result<()> {
if excluded.contains(&tree_hash) {
return Ok(());
}
if !seen.insert(tree_hash) {
return Ok(());
}
let tree = store
.get_tree(&tree_hash)?
.ok_or_else(|| ProtocolError::ObjectNotFound(tree_hash.to_hex()))?;
let tree_bytes = rmp_serde::to_vec_named(&tree)?;
out.push(ObjectInfo {
id: ObjectId::Hash(tree_hash),
obj_type: ObjectType::Tree,
size: tree_bytes.len() as u64,
delta_base: None,
});
for entry in tree.entries() {
match entry.entry_type {
EntryType::Blob => {
if excluded.contains(&entry.hash) {
continue;
}
if !seen.insert(entry.hash) {
continue;
}
let blob = store
.get_blob(&entry.hash)?
.ok_or_else(|| ProtocolError::ObjectNotFound(entry.hash.to_hex()))?;
out.push(ObjectInfo {
id: ObjectId::Hash(entry.hash),
obj_type: ObjectType::Blob,
size: blob.size() as u64,
delta_base: None,
});
}
EntryType::Tree => {
enumerate_tree_closure_filtered(store, entry.hash, excluded, seen, out)?;
}
EntryType::Symlink => {
if excluded.contains(&entry.hash) {
continue;
}
if !seen.insert(entry.hash) {
continue;
}
let blob = store
.get_blob(&entry.hash)?
.ok_or_else(|| ProtocolError::ObjectNotFound(entry.hash.to_hex()))?;
out.push(ObjectInfo {
id: ObjectId::Hash(entry.hash),
obj_type: ObjectType::Blob,
size: blob.size() as u64,
delta_base: None,
});
}
}
}
Ok(())
}
fn enumerate_tree_plan_filtered(
store: &dyn ObjectStore,
tree_hash: ContentHash,
excluded: &HashSet<ContentHash>,
seen: &mut HashSet<ContentHash>,
out: &mut Vec<PlannedObject>,
) -> Result<()> {
if excluded.contains(&tree_hash) {
return Ok(());
}
if !seen.insert(tree_hash) {
return Ok(());
}
let tree = store
.get_tree(&tree_hash)?
.ok_or_else(|| ProtocolError::ObjectNotFound(tree_hash.to_hex()))?;
out.push(PlannedObject {
id: ObjectId::Hash(tree_hash),
obj_type: ObjectType::Tree,
});
for entry in tree.entries() {
match entry.entry_type {
EntryType::Blob | EntryType::Symlink => {
if excluded.contains(&entry.hash) {
continue;
}
if !seen.insert(entry.hash) {
continue;
}
out.push(PlannedObject {
id: ObjectId::Hash(entry.hash),
obj_type: ObjectType::Blob,
});
}
EntryType::Tree => {
enumerate_tree_plan_filtered(store, entry.hash, excluded, seen, out)?;
}
}
}
Ok(())
}
fn collect_excluded(
store: &dyn ObjectStore,
roots: &[ChangeId],
) -> Result<(HashSet<ChangeId>, HashSet<ContentHash>)> {
if roots.is_empty() {
return Ok((HashSet::new(), HashSet::new()));
}
let mut excluded_states: HashSet<ChangeId> = HashSet::new();
let mut excluded_hashes: HashSet<ContentHash> = HashSet::new();
let mut queue: VecDeque<ChangeId> = VecDeque::new();
for id in roots {
queue.push_back(*id);
}
while let Some(id) = queue.pop_front() {
if !excluded_states.insert(id) {
continue;
}
let state = match store.get_state(&id)? {
Some(state) => state,
None => continue,
};
for parent in &state.parents {
queue.push_back(*parent);
}
collect_tree_hashes(store, state.tree, &mut excluded_hashes)?;
if let Some(provenance_root) = state.provenance {
collect_tree_hashes(store, provenance_root, &mut excluded_hashes)?;
}
if let Some(context_root) = state.context {
collect_tree_hashes(store, context_root, &mut excluded_hashes)?;
}
}
Ok((excluded_states, excluded_hashes))
}
fn collect_tree_hashes(
store: &dyn ObjectStore,
tree_hash: ContentHash,
excluded: &mut HashSet<ContentHash>,
) -> Result<()> {
if !excluded.insert(tree_hash) {
return Ok(());
}
let tree = match store.get_tree(&tree_hash)? {
Some(tree) => tree,
None => return Ok(()),
};
for entry in tree.entries() {
match entry.entry_type {
EntryType::Blob | EntryType::Symlink => {
excluded.insert(entry.hash);
}
EntryType::Tree => {
collect_tree_hashes(store, entry.hash, excluded)?;
}
}
}
Ok(())
}
pub fn is_ancestor(
store: &dyn ObjectStore,
ancestor: ChangeId,
descendant: ChangeId,
) -> Result<bool> {
if ancestor == descendant {
return Ok(true);
}
let mut seen: HashSet<ChangeId> = HashSet::new();
let mut queue: VecDeque<ChangeId> = VecDeque::new();
queue.push_back(descendant);
while let Some(id) = queue.pop_front() {
if !seen.insert(id) {
continue;
}
let state = match store.get_state(&id)? {
Some(s) => s,
None => return Ok(false),
};
for parent in state.parents {
if parent == ancestor {
return Ok(true);
}
queue.push_back(parent);
}
}
Ok(false)
}
#[cfg(test)]
mod tests {
use repo::Repository;
use tempfile::TempDir;
use super::{
ObjectId, StateClosureOptions, enumerate_state_closure_plan_with_options,
enumerate_state_closure_with_options,
};
#[test]
fn lean_closure_planner_matches_object_info_ids_and_types() {
let temp = TempDir::new().unwrap();
let repo = Repository::init_default(temp.path()).unwrap();
std::fs::create_dir_all(temp.path().join("src")).unwrap();
std::fs::write(temp.path().join("README.md"), "hello\n").unwrap();
std::fs::write(temp.path().join("src/lib.rs"), "pub fn hi() {}\n").unwrap();
let state = repo.snapshot(Some("seed".to_string()), None).unwrap();
let full = enumerate_state_closure_with_options(
repo.store(),
state.change_id,
StateClosureOptions::default(),
)
.unwrap();
let lean = enumerate_state_closure_plan_with_options(
repo.store(),
state.change_id,
StateClosureOptions::default(),
)
.unwrap();
let full_pairs = full
.into_iter()
.map(|info| (info.id, info.obj_type))
.collect::<std::collections::HashSet<_>>();
let lean_pairs = lean
.into_iter()
.map(|info| (info.id, info.obj_type))
.collect::<std::collections::HashSet<_>>();
assert_eq!(full_pairs, lean_pairs);
assert!(
full_pairs
.iter()
.any(|(id, _)| matches!(id, ObjectId::ChangeId(_)))
);
}
}