use crate::model::{
artifacts::ArtifactChecksumRecord,
inventory::{
InventoryRecord, InventoryRecordError, InventoryTargetRecord, MAX_INVENTORY_TARGETS,
},
};
use std::collections::{BTreeSet, VecDeque};
use thiserror::Error;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SelectionExpansion {
Exact,
DirectChildren,
Descendants,
}
#[derive(Clone, Debug)]
pub struct SelectionRequest {
pub canister_ids: Vec<String>,
pub expansion: SelectionExpansion,
}
#[derive(Clone, Debug)]
pub struct SelectionView<'a> {
pub inventory: ArtifactChecksumRecord,
pub targets: Vec<&'a InventoryTargetRecord>,
}
pub fn select<'a>(
inventory: &'a InventoryRecord,
request: &SelectionRequest,
) -> Result<SelectionView<'a>, SelectionError> {
if request.canister_ids.is_empty() {
return Err(SelectionError::EmptySelection);
}
if request.canister_ids.len() > MAX_INVENTORY_TARGETS {
return Err(SelectionError::TooManySelectors);
}
let mut selected = BTreeSet::new();
let mut queue = VecDeque::new();
for id in &request.canister_ids {
let target = inventory.target(id)?;
if !selected.insert(target.canister_id()) {
return Err(SelectionError::DuplicateSelector(
target.canister_id().into(),
));
}
queue.push_back(target.canister_id());
}
if request.expansion != SelectionExpansion::Exact {
while let Some(parent) = queue.pop_front() {
for child in inventory
.targets()
.iter()
.filter(|target| target.parent_canister_id() == Some(parent))
{
if selected.insert(child.canister_id())
&& request.expansion == SelectionExpansion::Descendants
{
queue.push_back(child.canister_id());
}
}
}
}
Ok(SelectionView {
inventory: inventory.digest(),
targets: inventory
.targets()
.iter()
.filter(|target| selected.contains(target.canister_id()))
.collect(),
})
}
#[derive(Debug, Error)]
pub enum SelectionError {
#[error("selection contains no targets")]
EmptySelection,
#[error("selection exceeds {MAX_INVENTORY_TARGETS} selectors")]
TooManySelectors,
#[error("duplicate physical selector {0}")]
DuplicateSelector(String),
#[error(transparent)]
Inventory(#[from] InventoryRecordError),
}
#[cfg(test)]
mod tests;