use std::collections::BTreeSet;

use bone_types::SketchEntityId;
use slotmap::SlotMap;

pub(crate) fn entity_closure<E, R>(
    entities: &SlotMap<SketchEntityId, E>,
    seed: SketchEntityId,
    references: impl Fn(&E) -> R + Copy,
) -> BTreeSet<SketchEntityId>
where
    R: IntoIterator<Item = SketchEntityId>,
{
    grow(entities, BTreeSet::from([seed]), references)
}

fn grow<E, R>(
    entities: &SlotMap<SketchEntityId, E>,
    closure: BTreeSet<SketchEntityId>,
    references: impl Fn(&E) -> R + Copy,
) -> BTreeSet<SketchEntityId>
where
    R: IntoIterator<Item = SketchEntityId>,
{
    let grown: BTreeSet<SketchEntityId> = entities
        .iter()
        .filter_map(|(id, entity)| {
            references(entity)
                .into_iter()
                .any(|r| closure.contains(&r))
                .then_some(id)
        })
        .chain(closure.iter().copied())
        .collect();
    if grown == closure {
        closure
    } else {
        grow(entities, grown, references)
    }
}
