weavatrix-rust 2.0.1

use crate::engine::RepositoryState;
use crate::operations::arg_str;
use blazingly_json::Value;
use std::collections::{BTreeSet, VecDeque};
use weavatrix_graph::{Direction, EdgeIndex, GraphView, NodeIndex};

pub(super) fn resolve_seeds(
    state: &RepositoryState,
    args: &Value,
) -> Result<Vec<NodeIndex>, String> {
    let mut seeds = Vec::new();
    for key in ["seed_symbols", "seed_files"] {
        if let Some(values) = args.get(key).and_then(Value::as_array) {
            for value in values.iter().filter_map(Value::as_str) {
                seeds.push(state.resolve_node(value)?);
            }
        }
    }
    if seeds.is_empty() {
        let question = arg_str(args, "question")?.to_ascii_lowercase();
        seeds.extend(
            state
                .graph()
                .nodes()
                .iter()
                .enumerate()
                .filter(|(_, node)| node.label.to_ascii_lowercase().contains(&question))
                .take(12)
                .map(|(index, _)| NodeIndex::new(u32::try_from(index).unwrap_or(u32::MAX))),
        );
    }
    if seeds.is_empty() {
        Err("query did not resolve any graph seed".to_owned())
    } else {
        Ok(seeds)
    }
}

/// Structural hubs contain or are imported by half the repository, so walking
/// through them turns any local question into "the whole repository". They
/// stay reachable as endpoints but never conduct a traversal.
fn is_hub(state: &RepositoryState, index: NodeIndex) -> bool {
    state.graph().node_at(index).is_some_and(|node| {
        matches!(
            node.kind,
            weavatrix_graph::NodeKind::Repository | weavatrix_graph::NodeKind::Package
        )
    })
}

pub(in crate::operations) fn traverse(
    state: &RepositoryState,
    seeds: Vec<NodeIndex>,
    max_depth: usize,
    max_nodes: usize,
    direction: Direction,
    dfs: bool,
    relations: Option<&BTreeSet<String>>,
) -> (Vec<(NodeIndex, usize)>, BTreeSet<EdgeIndex>) {
    let mut seen = BTreeSet::new();
    let mut order = Vec::new();
    let mut edges = BTreeSet::new();
    let mut queue = VecDeque::new();
    for seed in seeds {
        if seen.insert(seed) {
            order.push((seed, 0));
            queue.push_back((seed, 0));
        }
    }
    while let Some((node, depth)) = if dfs {
        queue.pop_back()
    } else {
        queue.pop_front()
    } {
        if depth >= max_depth || seen.len() >= max_nodes || (depth > 0 && is_hub(state, node)) {
            continue;
        }
        let adjacent = match direction {
            Direction::Outgoing => state.graph().outgoing_edges(node).collect::<Vec<_>>(),
            Direction::Incoming => state.graph().incoming_edges(node).collect::<Vec<_>>(),
            Direction::Both => state
                .graph()
                .outgoing_edges(node)
                .chain(state.graph().incoming_edges(node))
                .collect::<Vec<_>>(),
        };
        for edge in adjacent {
            if relations.is_some_and(|allowed| {
                state
                    .graph()
                    .edge_at(edge)
                    .is_none_or(|edge| !allowed.contains(edge.kind.as_str()))
            }) {
                continue;
            }
            let Some(endpoints) = state.graph().edge_endpoints(edge) else {
                continue;
            };
            edges.insert(edge);
            let neighbor = if endpoints.source() == node {
                endpoints.target()
            } else {
                endpoints.source()
            };
            if seen.len() < max_nodes && seen.insert(neighbor) {
                order.push((neighbor, depth + 1));
                queue.push_back((neighbor, depth + 1));
            }
        }
    }
    (order, edges)
}