use crate::error::RetrievalError;
use ankurah_proto as proto;
pub async fn expand_states(
mut states: Vec<proto::Attested<proto::EntityState>>,
additional_entity_ids: impl IntoIterator<Item = proto::EntityId>,
collection: &crate::storage::StorageCollectionWrapper,
) -> Result<Vec<proto::Attested<proto::EntityState>>, RetrievalError> {
let mut entity_map: std::collections::HashSet<_> = states.iter().map(|s| s.payload.entity_id).collect();
for entity_id in additional_entity_ids {
if !entity_map.contains(&entity_id) {
match collection.get_state(entity_id).await {
Ok(state) => {
states.push(state);
entity_map.insert(entity_id);
}
Err(RetrievalError::EntityNotFound(_)) => {
}
Err(e) => {
return Err(e);
}
}
}
}
Ok(states)
}