lance-csr 0.3.1

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! `shortestPath` / `allShortestPaths` as a breadth-first search over the adjacency.
//!
//! [`ShortestPathNode`] takes rows that already bind both end nodes and appends, per
//! matching path, the path's length and its node ids (`<path>__length`, `<path>__nodes`),
//! repeating the input row once per path. Rows with no path within the length range are
//! dropped, as MATCH requires; OPTIONAL MATCH keeps them with nulls through its left join.
//!
//! The adjacency is the engine's CSR when it has one (`use_csr=True`). Otherwise the node has
//! a second input, a scan of the relationship's two endpoint columns, and builds an adjacency
//! from it once per query, shared by all partitions.
//!
//! A source with one target searches from both ends when reverse adjacency is available,
//! expanding the smaller complete frontier. Sources with several targets share one
//! forward search. Predecessors retain every shortest path, including parallel edges.

use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::hash::{Hash, Hasher};
use std::sync::Arc;

use arrow::array::Int64Builder;
use arrow::array::{Array, ArrayRef, AsArray, Int64Array, ListBuilder, RecordBatch, UInt32Array};
use arrow::compute::{cast, take};
use arrow::datatypes::{DataType, Field, Int64Type, Schema, SchemaRef};
use datafusion::common::{DFSchema, DFSchemaRef};
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
};
use futures::lock::Mutex;
use futures::StreamExt;

use crate::csr_expand::{Adjacency, Neighbors};
use crate::error::{GraphError, Result};

type FastMap<K, V> = HashMap<K, V, ahash::RandomState>;

/// Which way the search follows edges.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd)]
pub enum SearchDirection {
    Outgoing,
    Incoming,
    Undirected,
}

/// The search parameters shared by the logical and physical nodes.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)]
pub struct ShortestPathSpec {
    /// Input columns holding the source and target node ids.
    pub source_column: String,
    pub target_column: String,
    pub direction: SearchDirection,
    pub min_hops: u32,
    pub max_hops: u32,
    /// Every shortest path, not just one.
    pub all: bool,
    /// Output column names.
    pub length_column: String,
    pub nodes_column: String,
    /// Endpoint columns of the edge scan (the second input), when there is no CSR.
    pub edge_source: String,
    pub edge_target: String,
}

/// Logical node; see the module docs.
#[derive(Debug, Clone)]
pub struct ShortestPathNode {
    input: LogicalPlan,
    /// Scan of the relationship's endpoints; `None` when `adjacency` is given.
    edges: Option<LogicalPlan>,
    adjacency: Option<Neighbors>,
    reverse_adjacency: Option<Neighbors>,
    spec: ShortestPathSpec,
    schema: DFSchemaRef,
}

impl ShortestPathNode {
    pub fn try_new(
        input: LogicalPlan,
        edges: Option<LogicalPlan>,
        adjacency: Option<Neighbors>,
        spec: ShortestPathSpec,
    ) -> Result<Self> {
        if edges.is_none() == adjacency.is_none() {
            return Err(plan_error(
                "shortestPath needs exactly one of a CSR or an edge scan",
            ));
        }
        let added = DFSchema::try_from(Schema::new(vec![
            Field::new(&spec.length_column, DataType::Int64, true),
            Field::new(
                &spec.nodes_column,
                DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
                true,
            ),
        ]))
        .map_err(|e| plan_error(&format!("shortestPath schema: {e}")))?;
        let schema = input
            .schema()
            .join(&added)
            .map_err(|e| plan_error(&format!("shortestPath schema: {e}")))?;
        Ok(Self {
            input,
            edges,
            adjacency,
            reverse_adjacency: None,
            spec,
            schema: Arc::new(schema),
        })
    }

    pub(crate) fn with_reverse_adjacency(mut self, adjacency: Option<Neighbors>) -> Self {
        self.reverse_adjacency = adjacency;
        self
    }
}

impl PartialEq for ShortestPathNode {
    fn eq(&self, other: &Self) -> bool {
        self.input == other.input
            && self.edges == other.edges
            && self.spec == other.spec
            && match (&self.reverse_adjacency, &other.reverse_adjacency) {
                (Some(a), Some(b)) => a.ptr_eq(b),
                (None, None) => true,
                _ => false,
            }
            && match (&self.adjacency, &other.adjacency) {
                (Some(a), Some(b)) => a.ptr_eq(b),
                (None, None) => true,
                _ => false,
            }
    }
}
impl Eq for ShortestPathNode {}
impl Hash for ShortestPathNode {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.input.hash(state);
        self.spec.hash(state);
    }
}
impl PartialOrd for ShortestPathNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.spec.partial_cmp(&other.spec)
    }
}

impl UserDefinedLogicalNodeCore for ShortestPathNode {
    fn name(&self) -> &str {
        "ShortestPath"
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        std::iter::once(&self.input)
            .chain(self.edges.as_ref())
            .collect()
    }

    fn schema(&self) -> &DFSchemaRef {
        &self.schema
    }

    fn expressions(&self) -> Vec<Expr> {
        vec![]
    }

    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = &self.spec;
        write!(
            f,
            "ShortestPath: {} -> {}, {:?}, hops {}..{}, {}, {}",
            s.source_column,
            s.target_column,
            s.direction,
            s.min_hops,
            s.max_hops,
            if s.all { "all" } else { "one" },
            if self.adjacency.is_some() {
                "csr"
            } else {
                "edge scan"
            }
        )
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<Expr>,
        mut inputs: Vec<LogicalPlan>,
    ) -> datafusion::common::Result<Self> {
        let edges = (inputs.len() > 1).then(|| inputs.pop().expect("two inputs"));
        let input = inputs.pop().expect("an input");
        Self::try_new(input, edges, self.adjacency.clone(), self.spec.clone())
            .map(|node| node.with_reverse_adjacency(self.reverse_adjacency.clone()))
            .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string()))
    }
}

/// Physical operator for [`ShortestPathNode`].
#[derive(Debug)]
pub struct ShortestPathExec {
    input: Arc<dyn ExecutionPlan>,
    /// Single-partition edge scan, when there is no CSR.
    edges: Option<Arc<dyn ExecutionPlan>>,
    /// The CSR, or the adjacency built from `edges` by the first batch that needs it.
    adjacency: Arc<Mutex<Option<Neighbors>>>,
    reverse_adjacency: Option<Neighbors>,
    spec: ShortestPathSpec,
    schema: SchemaRef,
    properties: Arc<PlanProperties>,
}

impl ShortestPathExec {
    fn new(
        input: Arc<dyn ExecutionPlan>,
        edges: Option<Arc<dyn ExecutionPlan>>,
        adjacency: Option<Neighbors>,
        reverse_adjacency: Option<Neighbors>,
        spec: ShortestPathSpec,
        schema: SchemaRef,
    ) -> Self {
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(schema.clone()),
            input.output_partitioning().clone(),
            EmissionType::Incremental,
            Boundedness::Bounded,
        ));
        Self {
            input,
            edges,
            adjacency: Arc::new(Mutex::new(adjacency)),
            reverse_adjacency,
            spec,
            schema,
            properties,
        }
    }
}

impl DisplayAs for ShortestPathExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
        let s = &self.spec;
        write!(
            f,
            "ShortestPathExec: {} -> {}, {:?}, hops {}..{}, {}",
            s.source_column,
            s.target_column,
            s.direction,
            s.min_hops,
            s.max_hops,
            if s.all { "all" } else { "one" }
        )
    }
}

impl ExecutionPlan for ShortestPathExec {
    fn name(&self) -> &str {
        "ShortestPathExec"
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.properties
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        std::iter::once(&self.input)
            .chain(self.edges.as_ref())
            .collect()
    }

    /// The edge scan must be one partition: the build reads it as a single stream. Without
    /// this, DataFusion's distribution rule drops a coalesce added at planning time and the
    /// build sees only partition 0's edges.
    fn required_input_distribution(&self) -> Vec<datafusion::physical_plan::Distribution> {
        use datafusion::physical_plan::Distribution;
        std::iter::once(Distribution::UnspecifiedDistribution)
            .chain(self.edges.as_ref().map(|_| Distribution::SinglePartition))
            .collect()
    }

    fn with_new_children(
        self: Arc<Self>,
        mut children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
        let edges = (children.len() > 1).then(|| children.pop().expect("two children"));
        let input = children.pop().expect("a child");
        let mut exec = Self::new(
            input,
            edges,
            None,
            self.reverse_adjacency.clone(),
            self.spec.clone(),
            self.schema.clone(),
        );
        exec.adjacency = self.adjacency.clone();
        Ok(Arc::new(exec))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> datafusion::common::Result<SendableRecordBatchStream> {
        let input = self.input.execute(partition, context.clone())?;
        let source_index = input.schema().index_of(&self.spec.source_column)?;
        let target_index = input.schema().index_of(&self.spec.target_column)?;
        let (cell, edges, spec, schema) = (
            self.adjacency.clone(),
            self.edges.clone(),
            self.spec.clone(),
            self.schema.clone(),
        );
        let out_schema = schema.clone();
        let reverse_adjacency = self.reverse_adjacency.clone();
        let stream = input.then(move |batch| {
            let reverse_adjacency = reverse_adjacency.clone();
            let (cell, edges, spec, schema, context) = (
                cell.clone(),
                edges.clone(),
                spec.clone(),
                schema.clone(),
                context.clone(),
            );
            async move {
                let adjacency = {
                    let mut slot = cell.lock().await;
                    if slot.is_none() {
                        *slot = Some(build_adjacency(edges, &spec, context).await?);
                    }
                    slot.clone().expect("just set")
                };
                search_batch(
                    &batch?,
                    source_index,
                    target_index,
                    &adjacency,
                    reverse_adjacency.as_ref().or_else(|| {
                        (spec.direction == SearchDirection::Undirected).then_some(&adjacency)
                    }),
                    &spec,
                    &schema,
                )
                .await
            }
        });
        Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream)))
    }
}

/// The adjacency for a query without a CSR: built once from the edge scan, in the
/// direction(s) the search follows.
async fn build_adjacency(
    edges: Option<Arc<dyn ExecutionPlan>>,
    spec: &ShortestPathSpec,
    context: Arc<TaskContext>,
) -> datafusion::common::Result<Neighbors> {
    let edges = edges.expect("an edge scan when there is no CSR");
    let read = |from: String, to: String| {
        let (edges, context) = (edges.clone(), context.clone());
        async move {
            let stream = edges.execute(0, context)?;
            Adjacency::build(stream, &from, &to)
                .await
                .map(Arc::new)
                .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))
        }
    };
    let (src, dst) = (spec.edge_source.clone(), spec.edge_target.clone());
    Ok(match spec.direction {
        SearchDirection::Outgoing => Neighbors::directed(read(src, dst).await?),
        SearchDirection::Incoming => Neighbors::directed(read(dst, src).await?),
        SearchDirection::Undirected => {
            Neighbors::undirected(read(src.clone(), dst.clone()).await?, read(dst, src).await?)
        }
    })
}

/// Breadth-first search from `source` until every node in `targets` is reached (or the
/// frontier empties, or `max_hops`). Returns each reached node's depth and, for every node
/// at depth d > 0, its predecessors at depth d - 1.
async fn search(
    adjacency: &Neighbors,
    source: i64,
    targets: &HashSet<i64>,
    max_hops: u32,
) -> (FastMap<i64, u32>, FastMap<i64, Vec<i64>>) {
    let mut depth: FastMap<i64, u32> = FastMap::default();
    let mut parents: FastMap<i64, Vec<i64>> = FastMap::default();
    depth.insert(source, 0);
    let mut remaining: HashSet<i64> = targets.iter().copied().filter(|t| *t != source).collect();
    let mut frontier = vec![source];
    let mut level = 0;
    while !frontier.is_empty() && !remaining.is_empty() && level < max_hops {
        level += 1;
        let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents).await;
        // Finish the level before stopping, so every shortest path is recorded.
        for n in &next {
            remaining.remove(n);
        }
        frontier = next;
    }
    (depth, parents)
}

/// Finish the whole level so parallel edges and alternative shortest paths survive.
async fn expand_level(
    adjacency: &Neighbors,
    frontier: &[i64],
    level: u32,
    depth: &mut FastMap<i64, u32>,
    parents: &mut FastMap<i64, Vec<i64>>,
) -> Vec<i64> {
    let mut next = Vec::new();
    let mut work = 0_usize;
    for &node in frontier {
        let previous_work = work;
        let mut position = crate::csr_expand::NeighborPosition::default();
        for candidate in adjacency.candidates_from(node, &mut position) {
            work += 1;
            if work.is_multiple_of(1024) {
                crate::csr_expand::cooperative_yield().await;
            }
            let Some(n) = candidate else {
                continue;
            };
            let n = n as i64;
            match depth.get(&n) {
                Some(&d) if d == level => parents.entry(n).or_default().push(node),
                Some(_) => {}
                None => {
                    depth.insert(n, level);
                    parents.insert(n, vec![node]);
                    next.push(n);
                }
            }
        }
        if work == previous_work {
            work += 1;
            if work.is_multiple_of(1024) {
                crate::csr_expand::cooperative_yield().await;
            }
        }
    }
    next
}

/// Search a single pair using complete layers from the smaller frontier. At the
/// first intersection every meeting node lies on the same depth cut of a shortest path.
pub(crate) async fn search_bidirectional(
    forward: &Neighbors,
    reverse: &Neighbors,
    source: i64,
    target: i64,
    max_hops: u32,
) -> (FastMap<i64, u32>, FastMap<i64, Vec<i64>>) {
    let mut depth = FastMap::default();
    let mut parents = FastMap::default();
    depth.insert(source, 0);
    if source == target {
        return (depth, parents);
    }
    let mut reverse_depth = FastMap::default();
    let mut reverse_parents: FastMap<i64, Vec<i64>> = FastMap::default();
    reverse_depth.insert(target, 0);
    let (mut frontier, mut reverse_frontier) = (vec![source], vec![target]);
    let (mut level, mut reverse_level) = (0, 0);
    while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops {
        if frontier.len() <= reverse_frontier.len() {
            level += 1;
            frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents).await;
        } else {
            reverse_level += 1;
            reverse_frontier = expand_level(
                reverse,
                &reverse_frontier,
                reverse_level,
                &mut reverse_depth,
                &mut reverse_parents,
            )
            .await;
        }
        let mut pending: Vec<_> = frontier
            .iter()
            .copied()
            .filter(|n| reverse_depth.get(n) == Some(&reverse_level))
            .collect();
        if pending.is_empty() {
            continue;
        }
        // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths.
        let mut visited = HashSet::new();
        let mut work = 0_usize;
        while let Some(node) = pending.pop() {
            work += 1;
            if work.is_multiple_of(1024) {
                crate::csr_expand::cooperative_yield().await;
            }
            if !visited.insert(node) {
                continue;
            }
            if let Some(next) = reverse_parents.get(&node) {
                for &n in next {
                    parents.entry(n).or_default().push(node);
                    pending.push(n);
                }
            }
        }
        depth.insert(target, level + reverse_level);
        break;
    }
    (depth, parents)
}

/// Shortest paths from the search source to `target`, as node ids from source to target:
/// all of them, or the first found.
async fn paths_to(target: i64, parents: &FastMap<i64, Vec<i64>>, all: bool) -> Vec<Vec<i64>> {
    let mut out = Vec::new();
    let mut stack = vec![vec![target]];
    let mut work = 0_usize;
    while let Some(path) = stack.pop() {
        work += 1;
        if work.is_multiple_of(1024) {
            crate::csr_expand::cooperative_yield().await;
        }
        let head = *path.last().expect("non-empty");
        match parents.get(&head) {
            None => {
                let mut path = path;
                path.reverse();
                out.push(path);
                if !all {
                    break;
                }
            }
            Some(ps) => {
                let take = if all { ps.len() } else { 1 };
                for &p in ps.iter().take(take) {
                    let mut next = path.clone();
                    next.push(p);
                    stack.push(next);
                }
            }
        }
    }
    out
}

async fn search_batch(
    batch: &RecordBatch,
    source_index: usize,
    target_index: usize,
    adjacency: &Neighbors,
    reverse_adjacency: Option<&Neighbors>,
    spec: &ShortestPathSpec,
    schema: &SchemaRef,
) -> datafusion::common::Result<RecordBatch> {
    let sources = cast(batch.column(source_index), &DataType::Int64)?;
    let sources = sources.as_primitive::<Int64Type>();
    let targets = cast(batch.column(target_index), &DataType::Int64)?;
    let targets = targets.as_primitive::<Int64Type>();
    // Targets per source, so each source is searched once for the whole batch.
    let mut by_source: FastMap<i64, HashSet<i64>> = FastMap::default();
    for i in 0..batch.num_rows() {
        if sources.is_valid(i) && targets.is_valid(i) {
            by_source
                .entry(sources.value(i))
                .or_default()
                .insert(targets.value(i));
        }
    }
    let mut searches = FastMap::default();
    for (&source, targets) in &by_source {
        let result = match reverse_adjacency.filter(|_| targets.len() == 1) {
            Some(reverse) => {
                search_bidirectional(
                    adjacency,
                    reverse,
                    source,
                    *targets.iter().next().unwrap(),
                    spec.max_hops,
                )
                .await
            }
            None => search(adjacency, source, targets, spec.max_hops).await,
        };
        searches.insert(source, result);
    }
    let mut rows = Vec::new();
    let mut lengths = Vec::new();
    let mut nodes = ListBuilder::new(Int64Builder::new());
    for i in 0..batch.num_rows() {
        if !(sources.is_valid(i) && targets.is_valid(i)) {
            continue;
        }
        let (s, t) = (sources.value(i), targets.value(i));
        let (depth, parents) = &searches[&s];
        let Some(&d) = depth.get(&t) else {
            continue;
        };
        if d < spec.min_hops || d > spec.max_hops {
            continue;
        }
        for path in paths_to(t, parents, spec.all).await {
            rows.push(i as u32);
            lengths.push(d as i64);
            nodes.values().append_slice(&path);
            nodes.append(true);
        }
    }
    let rows = UInt32Array::from(rows);
    let mut columns: Vec<ArrayRef> = batch
        .columns()
        .iter()
        .map(|c| take(c.as_ref(), &rows, None))
        .collect::<std::result::Result<_, _>>()?;
    columns.push(Arc::new(Int64Array::from(lengths)));
    columns.push(Arc::new(nodes.finish()));
    Ok(RecordBatch::try_new(schema.clone(), columns)?)
}

/// Plan [`ShortestPathNode`]s; installed by [`crate::csr_expand::CsrQueryPlanner`].
pub(crate) fn plan_shortest_path(
    node: &ShortestPathNode,
    physical_inputs: &[Arc<dyn ExecutionPlan>],
) -> Arc<dyn ExecutionPlan> {
    let edges = physical_inputs.get(1).map(|scan| {
        // One stream feeds the build: merge the scan's partitions.
        Arc::new(CoalescePartitionsExec::new(scan.clone())) as Arc<dyn ExecutionPlan>
    });
    Arc::new(ShortestPathExec::new(
        physical_inputs[0].clone(),
        edges,
        node.adjacency.clone(),
        node.reverse_adjacency.clone(),
        node.spec.clone(),
        Arc::new(node.schema.as_arrow().clone()),
    ))
}

fn plan_error(message: &str) -> GraphError {
    GraphError::PlanError {
        message: message.to_string(),
        location: snafu::Location::new(file!(), line!(), column!()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    async fn adjacencies(edges: &[(i64, i64)]) -> (Arc<Adjacency>, Arc<Adjacency>) {
        let schema = Arc::new(Schema::new(vec![
            Field::new("src", DataType::Int64, false),
            Field::new("dst", DataType::Int64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(Int64Array::from(
                    edges.iter().map(|e| e.0).collect::<Vec<_>>(),
                )),
                Arc::new(Int64Array::from(
                    edges.iter().map(|e| e.1).collect::<Vec<_>>(),
                )),
            ],
        )
        .unwrap();
        let stream = || {
            let batch = batch.clone();
            Box::pin(RecordBatchStreamAdapter::new(
                schema.clone(),
                futures::stream::once(async move { Ok(batch) }),
            )) as SendableRecordBatchStream
        };
        let out = Adjacency::build(stream(), "src", "dst").await.unwrap();
        let inc = Adjacency::build(stream(), "dst", "src").await.unwrap();
        (Arc::new(out), Arc::new(inc))
    }

    async fn neighbors(edges: &[(i64, i64)]) -> Neighbors {
        let (out, inc) = adjacencies(edges).await;
        Neighbors::undirected(out, inc)
    }

    #[tokio::test]
    async fn finds_every_shortest_path_and_only_shortest_ones() {
        // Two 2-hop routes 1-2-4 and 1-3-4, and a longer 1-5-6-4.
        let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]).await;
        let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10).await;
        assert_eq!(depth[&4], 2);
        let mut all = paths_to(4, &parents, true).await;
        all.sort();
        assert_eq!(all, vec![vec![1, 2, 4], vec![1, 3, 4]]);
        assert_eq!(paths_to(4, &parents, false).await.len(), 1);
        // Beyond the hop limit, nothing is reached.
        let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1).await;
        assert!(!depth.contains_key(&4));
    }

    #[tokio::test]
    async fn cancellation_without_edges_yields_while_scanning_frontier() {
        let (out, _) = adjacencies(&[]).await;
        let adjacency = Neighbors::directed(out);
        let frontier = vec![99; 100_000];
        let mut depth = FastMap::default();
        let mut parents = FastMap::default();
        let mut next = Box::pin(expand_level(
            &adjacency,
            &frontier,
            1,
            &mut depth,
            &mut parents,
        ));
        assert!(futures::poll!(next.as_mut()).is_pending());
        drop(next);
        assert!(depth.is_empty() && parents.is_empty());
        assert!(
            expand_level(&adjacency, &frontier, 1, &mut depth, &mut parents)
                .await
                .is_empty()
        );
    }

    #[tokio::test]
    async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() {
        // A wide first layer makes the target's side expand. Two meeting nodes,
        // including parallel edges on both sides of the cut, yield six paths.
        let (out, inc) = adjacencies(&[
            (1, 2),
            (1, 2),
            (1, 3),
            (1, 9),
            (2, 4),
            (2, 4),
            (3, 4),
            (3, 4),
            (4, 5),
            (5, 5),
            (5, 1),
        ])
        .await;
        let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc));
        let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3).await;
        assert_eq!(depth[&5], 3);
        let mut all = paths_to(5, &parents, true).await;
        all.sort();
        assert_eq!(
            all,
            vec![vec![1, 2, 4, 5]; 4]
                .into_iter()
                .chain(vec![vec![1, 3, 4, 5]; 2])
                .collect::<Vec<_>>()
        );
        let one = paths_to(5, &parents, false).await;
        assert_eq!(one.len(), 1);
        assert!(all.contains(&one[0]));
        assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2)
            .await
            .0
            .contains_key(&5));
    }

    #[tokio::test]
    async fn bidirectional_matches_forward_search_for_small_graphs() {
        // All directed graphs on three nodes, including disconnected pairs and cycles.
        let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)];
        for mask in 0..(1 << edges.len()) {
            let selected: Vec<_> = edges
                .iter()
                .enumerate()
                .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge))
                .collect();
            let (out, inc) = adjacencies(&selected).await;
            let directed = Neighbors::directed(out.clone());
            let reverse = Neighbors::directed(inc.clone());
            let undirected = Neighbors::undirected(out, inc);
            for (forward, reverse) in [(&directed, &reverse), (&undirected, &undirected)] {
                for source in 0..3 {
                    for target in 0..3 {
                        for max_hops in 0..=3 {
                            let (expected_depth, expected_parents) =
                                search(forward, source, &HashSet::from([target]), max_hops).await;
                            let (depth, parents) =
                                search_bidirectional(forward, reverse, source, target, max_hops)
                                    .await;
                            assert_eq!(
                                depth.get(&target),
                                expected_depth.get(&target),
                                "graph {mask}, {source}->{target}, max {max_hops}"
                            );
                            if depth.contains_key(&target) {
                                let mut expected = paths_to(target, &expected_parents, true).await;
                                let mut actual = paths_to(target, &parents, true).await;
                                expected.sort();
                                actual.sort();
                                assert_eq!(
                                    actual, expected,
                                    "graph {mask}, {source}->{target}, max {max_hops}"
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    #[tokio::test]
    async fn bidirectional_avoids_expanding_the_wide_middle_layer() {
        let mut edges = Vec::new();
        for branch in 1..=100 {
            edges.push((0, branch));
            for leaf in 0..100 {
                edges.push((branch, 101 + (branch - 1) * 100 + leaf));
            }
        }
        let (out, inc) = adjacencies(&edges).await;
        let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc));
        let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2).await;
        assert_eq!(paths_to(101, &parents, true).await, vec![vec![0, 1, 101]]);
        assert_eq!(depth.len(), 102); // Root, first layer and stitched target.
        assert_eq!(
            search(&forward, 0, &HashSet::from([101]), 2).await.0.len(),
            10_101
        );
    }
}