#![allow(non_snake_case)]

use vcs_graph::bfs::BfsState;
use vcs_graph::graph::{Graph as RsGraph, GraphError};
use vcs_graph::known_graph::KnownGraph as RsKnownGraph;
use vcs_graph::{
    CachingParentsProvider as RsCachingParentsProvider, ChildMap, ParentMap, Parents,
    ParentsProvider, RevnoVec,
};

use pyo3::import_exception;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyTuple};
use pyo3::wrap_pyfunction;
use std::collections::{HashMap, HashSet};
use std::hash::Hash;

import_exception!(vcsgraph.errors, GraphCycleError);

struct PyNode(Py<PyAny>);

impl std::fmt::Debug for PyNode {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        Python::attach(|py| {
            let repr = self.0.bind(py).repr();
            if PyErr::occurred(py) {
                return Err(std::fmt::Error);
            }
            if let Ok(repr) = repr {
                write!(f, "{}", repr)
            } else {
                write!(f, "???")
            }
        })
    }
}

impl std::fmt::Display for PyNode {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        Python::attach(|py| {
            let repr = self.0.bind(py).repr();
            if PyErr::occurred(py) {
                return Err(std::fmt::Error);
            }
            if let Ok(repr) = repr {
                write!(f, "{}", repr)
            } else {
                write!(f, "???")
            }
        })
    }
}

impl Clone for PyNode {
    fn clone(&self) -> PyNode {
        Python::attach(|py| PyNode(self.0.clone_ref(py)))
    }
}

impl From<Py<PyAny>> for PyNode {
    fn from(obj: Py<PyAny>) -> PyNode {
        PyNode(obj)
    }
}

impl<'a> From<Bound<'a, PyAny>> for PyNode {
    fn from(obj: Bound<'a, PyAny>) -> PyNode {
        PyNode(obj.unbind())
    }
}

impl<'py> FromPyObject<'_, 'py> for PyNode {
    type Error = PyErr;

    fn extract(obj: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
        Ok(PyNode(obj.to_owned().unbind()))
    }
}

impl<'py> IntoPyObject<'py> for PyNode {
    type Target = PyAny;

    type Output = Bound<'py, Self::Target>;

    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        Ok(self.0.clone_ref(py).into_bound(py))
    }
}

impl Hash for PyNode {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        Python::attach(|py| match self.0.bind(py).hash() {
            Err(err) => err.restore(py),
            Ok(hash) => state.write_isize(hash),
        });
    }
}

impl PartialEq for PyNode {
    fn eq(&self, other: &PyNode) -> bool {
        Python::attach(|py| match self.0.bind(py).eq(other.0.bind(py)) {
            Err(err) => {
                err.restore(py);
                false
            }
            Ok(b) => b,
        })
    }
}

impl std::cmp::Eq for PyNode {}

impl PartialOrd for PyNode {
    fn partial_cmp(&self, other: &PyNode) -> Option<std::cmp::Ordering> {
        self.cmp(other).into()
    }
}

impl Ord for PyNode {
    fn cmp(&self, other: &PyNode) -> std::cmp::Ordering {
        Python::attach(|py| match self.0.bind(py).lt(other.0.bind(py)) {
            Err(err) => {
                err.restore(py);
                std::cmp::Ordering::Equal
            }
            Ok(b) => {
                if b {
                    std::cmp::Ordering::Less
                } else {
                    match self.0.bind(py).gt(other.0.bind(py)) {
                        Err(err) => {
                            err.restore(py);
                            std::cmp::Ordering::Equal
                        }
                        Ok(b) => {
                            if b {
                                std::cmp::Ordering::Greater
                            } else {
                                std::cmp::Ordering::Equal
                            }
                        }
                    }
                }
            }
        })
    }
}

/// Given a map from child => parents, create a map of parent => children
#[pyfunction]
fn invert_parent_map(parent_map: ParentMap<PyNode>) -> ChildMap<PyNode> {
    vcs_graph::invert_parent_map::<PyNode>(&parent_map)
}

/// Collapse regions of the graph that are 'linear'.
///
/// For example::
///
///   A:[B], B:[C]
///
/// can be collapsed by removing B and getting::
///
///   A:[C]
///
/// :param parent_map: A dictionary mapping children to their parents
/// :return: Another dictionary with 'linear' chains collapsed
#[pyfunction]
fn collapse_linear_regions(parent_map: ParentMap<PyNode>) -> PyResult<ParentMap<PyNode>> {
    Ok(vcs_graph::collapse_linear_regions::<PyNode>(&parent_map))
}

/// A parents provider for Graph objects, backed by a dict.
///
/// Mirrors `vcsgraph.graph.DictParentsProvider`: takes a mapping of
/// `{key: parents_tuple}` and serves `get_parent_map` lookups from it.
/// The dict is consulted on every call so mutations after construction
/// are visible — matching the original Python implementation.
#[pyclass(name = "DictParentsProvider", dict)]
struct PyDictParentsProvider {
    ancestry: Py<PyAny>,
}

impl PyDictParentsProvider {
    fn lookup(&self, py: Python<'_>, keys: &HashSet<PyNode>) -> ParentMap<PyNode> {
        let ancestry = self.ancestry.bind(py);
        let mut pm = ParentMap::new();
        for k in keys {
            let key_obj = k.0.bind(py);
            let v_obj = match ancestry.get_item(key_obj) {
                Ok(v) => v,
                Err(_) => continue,
            };
            let Ok(iter) = v_obj.try_iter() else {
                continue;
            };
            let mut parents: Vec<PyNode> = Vec::new();
            let mut ok = true;
            for p in iter {
                match p {
                    Ok(p) => parents.push(PyNode::from(p.unbind())),
                    Err(_) => {
                        ok = false;
                        break;
                    }
                }
            }
            if ok {
                pm.insert(k.clone(), Parents::Known(parents));
            }
        }
        pm
    }
}

#[pymethods]
impl PyDictParentsProvider {
    #[new]
    fn new(ancestry: Py<PyAny>) -> PyResult<Self> {
        Ok(PyDictParentsProvider { ancestry })
    }

    /// The underlying mapping, preserved as the original Python object.
    #[getter]
    fn ancestry(&self, py: Python) -> Py<PyAny> {
        self.ancestry.clone_ref(py)
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let r = self.ancestry.bind(py).repr()?;
        Ok(format!("DictParentsProvider({})", r))
    }

    fn get_parent_map<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let hs: HashSet<PyNode> = nodes.into_iter().collect();
        let pm = self.lookup(py, &hs);
        parent_map_to_pydict(py, pm)
    }
}

impl ParentsProvider<PyNode> for PyDictParentsProvider {
    fn get_parent_map(&self, keys: &HashSet<PyNode>) -> ParentMap<PyNode> {
        Python::attach(|py| self.lookup(py, keys))
    }
}

/// A parents provider which stacks multiple child providers. Each child is
/// queried in order, matching `vcsgraph.graph.StackedParentsProvider`.
///
/// If a child exposes `get_cached_parent_map`, that fast path is queried
/// first for all children before any full `get_parent_map` call — just
/// like the Python version.
///
/// The provider iterable is stored as-is and re-iterated on every
/// `get_parent_map` call, so live mutations (e.g. appending a fallback to
/// the list passed in) take effect on subsequent lookups.
#[pyclass(name = "StackedParentsProvider")]
struct PyStackedParentsProvider {
    parent_providers: Py<PyAny>,
}

#[pymethods]
impl PyStackedParentsProvider {
    #[new]
    fn new(parent_providers: Py<PyAny>) -> PyResult<Self> {
        Ok(PyStackedParentsProvider { parent_providers })
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let r = self.parent_providers.bind(py).repr()?;
        Ok(format!("StackedParentsProvider({})", r))
    }

    fn get_parent_map<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let found = PyDict::new(py);
        let remaining = pyo3::types::PySet::empty(py)?;
        for k in keys.bind(py).try_iter()? {
            remaining.add(k?.unbind())?;
        }

        // First pass: any provider that implements get_cached_parent_map
        // gets queried cheaply before we hit the slow path.
        for pp in self.parent_providers.bind(py).try_iter()? {
            if remaining.is_empty() {
                break;
            }
            let pp = pp?;
            let get_cached = pp.getattr("get_cached_parent_map");
            let Ok(get_cached) = get_cached else {
                continue;
            };
            if get_cached.is_none() {
                continue;
            }
            let new_found = get_cached.call1((remaining.clone(),))?;
            let new_found_dict = new_found.cast::<PyDict>()?;
            for (k, v) in new_found_dict.iter() {
                found.set_item(&k, v)?;
                remaining.discard(&k)?;
            }
        }
        if remaining.is_empty() {
            return Ok(found);
        }

        // Second pass: full get_parent_map calls, in order.
        for pp in self.parent_providers.bind(py).try_iter()? {
            if remaining.is_empty() {
                break;
            }
            let pp = pp?;
            let new_found = match pp.call_method1("get_parent_map", (remaining.clone(),)) {
                Ok(r) => r,
                Err(err) => {
                    // Match Python's behaviour of catching UnsupportedOperation
                    // and moving on.
                    if err.get_type(py).is(&py
                        .import("vcsgraph.errors")?
                        .getattr("UnsupportedOperation")?)
                    {
                        continue;
                    }
                    return Err(err);
                }
            };
            let new_found_dict = new_found.cast::<PyDict>()?;
            for (k, v) in new_found_dict.iter() {
                found.set_item(&k, v)?;
                remaining.discard(&k)?;
            }
        }
        Ok(found)
    }
}

/// A parents provider that wraps any `get_parent_map`-like callable.
///
/// Mirrors `vcsgraph.graph.CallableToParentsProviderAdapter`.
#[pyclass(name = "CallableToParentsProviderAdapter")]
struct PyCallableToParentsProviderAdapter {
    callable: Py<PyAny>,
}

#[pymethods]
impl PyCallableToParentsProviderAdapter {
    #[new]
    fn new(callable: Py<PyAny>) -> Self {
        PyCallableToParentsProviderAdapter { callable }
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let r = self.callable.bind(py).repr()?;
        Ok(format!("CallableToParentsProviderAdapter({})", r))
    }

    fn get_parent_map(&self, py: Python, keys: Py<PyAny>) -> PyResult<Py<PyAny>> {
        Ok(self.callable.bind(py).call1((keys,))?.unbind())
    }
}

#[pyclass]
struct TopoSorter {
    sorter: vcs_graph::tsort::TopoSorter<PyNode>,
}

#[pymethods]
impl TopoSorter {
    #[new]
    fn new(py: Python, graph: Py<PyAny>) -> PyResult<TopoSorter> {
        let iter = if graph.bind(py).is_instance_of::<PyDict>() {
            graph
                .cast_bound::<PyDict>(py)?
                .call_method0("items")?
                .try_iter()?
        } else {
            graph.bind(py).try_iter()?
        };
        let graph = iter
            .map(|k| k?.extract::<(Py<PyAny>, Vec<Py<PyAny>>)>())
            .map(|k| k.map(|(k, vs)| (PyNode::from(k), vs.into_iter().map(PyNode::from).collect())))
            .collect::<PyResult<Vec<(PyNode, Vec<PyNode>)>>>()?;

        let sorter = vcs_graph::tsort::TopoSorter::<PyNode>::new(graph.into_iter());
        Ok(TopoSorter { sorter })
    }

    fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyAny>>> {
        match self.sorter.next() {
            None => Ok(None),
            Some(Ok(node)) => Ok(Some(node.into_pyobject(py)?.unbind())),
            Some(Err(vcs_graph::Error::Cycle(e))) => Err(GraphCycleError::new_err(e)),
            Some(Err(e)) => panic!("Unexpected error: {:?}", e),
        }
    }

    fn __iter__(slf: PyRefMut<Self>) -> PyRefMut<Self> {
        slf
    }

    fn iter_topo_order(slf: PyRefMut<Self>) -> PyRefMut<Self> {
        slf
    }

    fn sorted(&mut self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
        let mut ret = Vec::new();
        while let Some(node) = self.__next__(py)? {
            ret.push(node);
        }
        Ok(ret)
    }
}

fn revno_vec_to_py(py: Python, revno: RevnoVec) -> Py<PyAny> {
    PyTuple::new(py, revno.into_iter().map(|v| v.into_pyobject(py).unwrap()))
        .unwrap()
        .into_any()
        .unbind()
}

#[pyclass]
struct MergeSorter {
    sorter: vcs_graph::tsort::MergeSorter<PyNode>,
}

fn branch_tip_is_null(py: Python, branch_tip: Py<PyAny>) -> bool {
    if let Ok(branch_tip) = branch_tip.extract::<&[u8]>(py) {
        branch_tip == b"null:"
    } else if let Ok((branch_tip,)) = branch_tip.extract::<(Vec<u8>,)>(py) {
        branch_tip.as_slice() == b"null:"
    } else {
        false
    }
}

#[pymethods]
impl MergeSorter {
    #[new]
    #[pyo3(signature = (graph, branch_tip=None, mainline_revisions=None, generate_revno=false))]
    fn new(
        py: Python,
        graph: Py<PyAny>,
        mut branch_tip: Option<Py<PyAny>>,
        mainline_revisions: Option<Py<PyAny>>,
        generate_revno: Option<bool>,
    ) -> PyResult<MergeSorter> {
        let iter = if graph.bind(py).is_instance_of::<PyDict>() {
            graph
                .cast_bound::<PyDict>(py)?
                .call_method0("items")?
                .try_iter()?
        } else {
            graph.bind(py).try_iter()?
        };
        let graph = iter
            .map(|k| k?.extract::<(Py<PyAny>, Vec<Py<PyAny>>)>())
            .map(|k| k.map(|(k, vs)| (PyNode::from(k), vs.into_iter().map(PyNode::from).collect())))
            .collect::<PyResult<HashMap<PyNode, Vec<PyNode>>>>()?;

        let mainline_revisions = if let Some(mainline_revisions) = mainline_revisions {
            let mainline_revisions = mainline_revisions
                .bind(py)
                .try_iter()?
                .map(|k| {
                    let item = k?;
                    Ok(item.extract::<Py<PyAny>>()?)
                })
                .collect::<PyResult<Vec<Py<PyAny>>>>()?;
            Some(mainline_revisions.into_iter().map(PyNode::from).collect())
        } else {
            None
        };

        // The null: revision doesn't exist in the graph, so don't attempt to remove it
        if let Some(ref mut tip_obj) = branch_tip {
            if branch_tip_is_null(py, tip_obj.clone_ref(py)) {
                branch_tip = None;
            }
        }

        let sorter = vcs_graph::tsort::MergeSorter::<PyNode>::new(
            graph,
            branch_tip.map(PyNode::from),
            mainline_revisions,
            generate_revno.unwrap_or(false),
        );
        Ok(MergeSorter { sorter })
    }

    fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyAny>>> {
        match self.sorter.next() {
            None => Ok(None),
            Some(Ok((sequence_number, node, merge_depth, None, end_of_merge))) => Ok(Some(
                (
                    sequence_number,
                    node.into_pyobject(py)?.unbind(),
                    merge_depth,
                    end_of_merge,
                )
                    .into_pyobject(py)?
                    .unbind()
                    .into(),
            )),

            Some(Ok((sequence_number, node, merge_depth, Some(revno), end_of_merge))) => Ok(Some(
                (
                    sequence_number,
                    node.into_pyobject(py)?.unbind(),
                    merge_depth,
                    revno_vec_to_py(py, revno),
                    end_of_merge,
                )
                    .into_pyobject(py)?
                    .unbind()
                    .into(),
            )),
            Some(Err(vcs_graph::Error::Cycle(e))) => Err(GraphCycleError::new_err(e)),
            Some(Err(e)) => panic!("Unexpected error: {:?}", e),
        }
    }

    fn __iter__(slf: PyRefMut<Self>) -> PyRefMut<Self> {
        slf
    }

    fn iter_topo_order(slf: PyRefMut<Self>) -> PyRefMut<Self> {
        slf
    }

    fn sorted<'a>(&mut self, py: Python<'a>) -> PyResult<Bound<'a, PyList>> {
        let ret = PyList::empty(py);
        loop {
            let item = self.__next__(py)?;
            if let Some(item) = item {
                ret.append(item)?;
            } else {
                break;
            }
        }
        Ok(ret)
    }
}

/// Topological sort a graph which groups merges.
///
/// :param graph: sequence of pairs of node->parents_list.
/// :param branch_tip: the tip of the branch to graph. Revisions not
///                    reachable from branch_tip are not included in the
///                    output.
/// :param mainline_revisions: If not None this forces a mainline to be
///                            used rather than synthesised from the graph.
///                            This must be a valid path through some part
///                            of the graph. If the mainline does not cover all
///                            the revisions, output stops at the start of the
///                            old revision listed in the mainline revisions
///                            list.
///                            The order for this parameter is oldest-first.
/// :param generate_revno: Optional parameter controlling the generation of
///     revision number sequences in the output. See the output description of
///     the MergeSorter docstring for details.
/// :result: See the MergeSorter docstring for details.
///
/// Node identifiers can be any hashable object, and are typically strings.
#[pyfunction]
#[pyo3(signature = (graph, branch_tip=None, mainline_revisions=None, generate_revno=false))]
fn merge_sort(
    py: Python,
    graph: Py<PyAny>,
    branch_tip: Option<Py<PyAny>>,
    mainline_revisions: Option<Py<PyAny>>,
    generate_revno: Option<bool>,
) -> PyResult<Bound<PyList>> {
    let mut sorter = MergeSorter::new(py, graph, branch_tip, mainline_revisions, generate_revno)?;
    sorter.sorted(py)
}

const NULL_REVISION: &[u8] = b"null:";

fn is_null_revision(py: Python, obj: &Py<PyAny>) -> bool {
    if let Ok(b) = obj.extract::<&[u8]>(py) {
        b == NULL_REVISION
    } else {
        false
    }
}

#[pyclass(name = "KnownGraph")]
struct PyKnownGraph {
    inner: RsKnownGraph<PyNode>,
}

fn extract_parent_map(py: Python, parent_map: Py<PyAny>) -> PyResult<Vec<(PyNode, Vec<PyNode>)>> {
    let iter = if parent_map.bind(py).is_instance_of::<PyDict>() {
        parent_map
            .cast_bound::<PyDict>(py)?
            .call_method0("items")?
            .try_iter()?
    } else {
        parent_map.bind(py).try_iter()?
    };
    iter.map(|k| k?.extract::<(Py<PyAny>, Vec<Py<PyAny>>)>())
        .map(|k| k.map(|(k, vs)| (PyNode::from(k), vs.into_iter().map(PyNode::from).collect())))
        .collect()
}

#[pymethods]
impl PyKnownGraph {
    #[new]
    #[pyo3(signature = (parent_map, do_cache=true))]
    fn new(py: Python, parent_map: Py<PyAny>, do_cache: Option<bool>) -> PyResult<Self> {
        let pm = extract_parent_map(py, parent_map)?;
        Ok(PyKnownGraph {
            inner: RsKnownGraph::new(pm, do_cache.unwrap_or(true)),
        })
    }

    /// Add a new node to the graph. If `key` was a ghost, it is filled in.
    fn add_node(
        &mut self,
        py: Python,
        key: Py<PyAny>,
        parent_keys: Vec<Py<PyAny>>,
    ) -> PyResult<()> {
        let key = PyNode::from(key);
        let parents: Vec<PyNode> = parent_keys.into_iter().map(PyNode::from).collect();
        self.inner
            .add_node(key.clone(), parents)
            .map_err(|e| match e {
                vcs_graph::Error::ParentMismatch {
                    key,
                    expected,
                    actual,
                } => {
                    let key_repr = format!("{:?}", key);
                    let expected_py: Vec<Py<PyAny>> =
                        expected.into_iter().map(|n| n.0.clone_ref(py)).collect();
                    let actual_py: Vec<Py<PyAny>> =
                        actual.into_iter().map(|n| n.0.clone_ref(py)).collect();
                    pyo3::exceptions::PyValueError::new_err(format!(
                        "Parent key mismatch, existing node {} has parents of {:?} not {:?}",
                        key_repr, expected_py, actual_py
                    ))
                }
                other => pyo3::exceptions::PyValueError::new_err(format!("{:?}", other)),
            })
    }

    /// Return the parent keys for `key`. Returns `None` for ghosts; raises
    /// `KeyError` if `key` is not in the graph.
    fn get_parent_keys(&self, py: Python, key: Py<PyAny>) -> PyResult<Option<Vec<Py<PyAny>>>> {
        let node = PyNode::from(key);
        if !self.inner.contains(&node) {
            return Err(pyo3::exceptions::PyKeyError::new_err(format!("{:?}", node)));
        }
        Ok(self
            .inner
            .get_parent_keys(&node)
            .map(|ps| ps.iter().map(|n| n.0.clone_ref(py)).collect()))
    }

    /// Return the child keys for `key`. Raises `KeyError` if `key` is not in
    /// the graph.
    fn get_child_keys(&self, py: Python, key: Py<PyAny>) -> PyResult<Vec<Py<PyAny>>> {
        let node = PyNode::from(key);
        match self.inner.get_child_keys(&node) {
            Some(cs) => Ok(cs.iter().map(|n| n.0.clone_ref(py)).collect()),
            None => Err(pyo3::exceptions::PyKeyError::new_err(format!("{:?}", node))),
        }
    }

    /// Return the heads from amongst `keys`.
    fn heads<'py>(&mut self, py: Python<'py>, keys: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
        let mut candidates: Vec<Py<PyAny>> = Vec::new();
        let mut had_null = false;
        for k in keys.bind(py).try_iter()? {
            let item: Py<PyAny> = k?.extract()?;
            if is_null_revision(py, &item) {
                had_null = true;
            } else {
                candidates.push(item);
            }
        }
        if candidates.is_empty() && had_null {
            // NULL_REVISION is only a head if it's the only entry.
            let fs = py.import("builtins")?.getattr("frozenset")?;
            let null = pyo3::types::PyBytes::new(py, NULL_REVISION);
            return fs.call1((vec![null.into_any()],));
        }
        let nodes: Vec<PyNode> = candidates.into_iter().map(PyNode::from).collect();
        let heads = self.inner.heads(nodes);
        let py_set: Vec<Py<PyAny>> = heads.into_iter().map(|n| n.0).collect();
        let fs = py.import("builtins")?.getattr("frozenset")?;
        fs.call1((py_set,))
    }

    /// Return the nodes in topological order (parents first).
    fn topo_sort(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
        match self.inner.topo_sort() {
            Ok(v) => Ok(v.into_iter().map(|n| n.0.clone_ref(py)).collect()),
            Err(vcs_graph::Error::Cycle(_)) => {
                Err(GraphCycleError::new_err(("cycle in known graph",)))
            }
            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{:?}", e))),
        }
    }

    /// Return a reverse topological ordering grouped by prefix.
    ///
    /// Mirrors the Python implementation: keys that are bytes use their first
    /// byte as the prefix bucket; single-element keys use an empty prefix.
    fn gc_sort(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
        let prefix_of = |k: &PyNode| -> Vec<u8> {
            Python::attach(|py| {
                let bound = k.0.bind(py);
                if let Ok(b) = bound.extract::<&[u8]>() {
                    if b.len() == 1 {
                        Vec::new()
                    } else {
                        vec![b[0]]
                    }
                } else if let Ok(s) = bound.extract::<&str>() {
                    if s.len() == 1 {
                        Vec::new()
                    } else {
                        vec![s.as_bytes()[0]]
                    }
                } else if let Ok(first) = bound.get_item(0) {
                    if let Ok(b) = first.extract::<&[u8]>() {
                        b.to_vec()
                    } else if let Ok(s) = first.extract::<&str>() {
                        s.as_bytes().to_vec()
                    } else {
                        Vec::new()
                    }
                } else {
                    Vec::new()
                }
            })
        };
        let v = self.inner.gc_sort(prefix_of);
        Ok(v.into_iter().map(|n| n.0.clone_ref(py)).collect())
    }

    /// Compute the merge-sorted graph output starting at `tip_key`.
    ///
    /// If `tip_key` is `None`, `b"null:"`, or `(b"null:",)`, returns an empty
    /// list (matches the Python null-tip semantics).
    fn merge_sort(
        &self,
        py: Python,
        tip_key: Py<PyAny>,
    ) -> PyResult<Vec<Py<PyKnownGraphMergeSortNode>>> {
        if tip_key.is_none(py) || branch_tip_is_null(py, tip_key.clone_ref(py)) {
            return Ok(Vec::new());
        }
        let tip = PyNode::from(tip_key);
        if !self.inner.contains(&tip) {
            return Ok(Vec::new());
        }
        let result = self.inner.merge_sort(tip).map_err(|e| match e {
            vcs_graph::Error::Cycle(_) => GraphCycleError::new_err(("cycle in known graph",)),
            other => pyo3::exceptions::PyValueError::new_err(format!("{:?}", other)),
        })?;
        result
            .into_iter()
            .map(|n| {
                let revno = revno_vec_to_py(py, n.revno);
                Py::new(
                    py,
                    PyKnownGraphMergeSortNode {
                        key: n.key.0,
                        merge_depth: n.merge_depth,
                        revno,
                        end_of_merge: n.end_of_merge,
                    },
                )
            })
            .collect()
    }

    /// Return a mapping-like view of all nodes in the graph, keyed by node
    /// key. Each value is a `_KnownGraphNode` exposing live `key`,
    /// `parent_keys`, `child_keys`, and `gdfo` attributes.
    #[getter]
    fn _nodes(slf: Py<Self>) -> PyKnownGraphNodesView {
        PyKnownGraphNodesView { graph: slf }
    }
}

#[pyclass(name = "_KnownGraphNodesView")]
struct PyKnownGraphNodesView {
    graph: Py<PyKnownGraph>,
}

#[pymethods]
impl PyKnownGraphNodesView {
    fn __getitem__(&self, py: Python, key: Py<PyAny>) -> PyResult<PyKnownGraphNode> {
        let g = self.graph.borrow(py);
        let node = PyNode::from(key.clone_ref(py));
        if !g.inner.contains(&node) {
            return Err(pyo3::exceptions::PyKeyError::new_err(format!("{:?}", node)));
        }
        Ok(PyKnownGraphNode {
            graph: self.graph.clone_ref(py),
            key,
        })
    }

    fn __contains__(&self, py: Python, key: Py<PyAny>) -> bool {
        self.graph.borrow(py).inner.contains(&PyNode::from(key))
    }

    fn __len__(&self, py: Python) -> usize {
        self.graph.borrow(py).inner.len()
    }

    fn __iter__(&self, py: Python) -> PyResult<Py<PyAny>> {
        let keys: Vec<Py<PyAny>> = self
            .graph
            .borrow(py)
            .inner
            .keys()
            .map(|n| n.0.clone_ref(py))
            .collect();
        Ok(pyo3::types::PyList::new(py, keys)?
            .call_method0("__iter__")?
            .unbind())
    }

    fn keys(&self, py: Python) -> Vec<Py<PyAny>> {
        self.graph
            .borrow(py)
            .inner
            .keys()
            .map(|n| n.0.clone_ref(py))
            .collect()
    }

    fn values(&self, py: Python) -> Vec<PyKnownGraphNode> {
        self.graph
            .borrow(py)
            .inner
            .keys()
            .map(|n| PyKnownGraphNode {
                graph: self.graph.clone_ref(py),
                key: n.0.clone_ref(py),
            })
            .collect()
    }
}

#[pyclass(name = "_KnownGraphNode")]
struct PyKnownGraphNode {
    graph: Py<PyKnownGraph>,
    key: Py<PyAny>,
}

#[pymethods]
impl PyKnownGraphNode {
    #[getter]
    fn key(&self, py: Python) -> Py<PyAny> {
        self.key.clone_ref(py)
    }

    #[getter]
    fn gdfo(&self, py: Python) -> Option<u64> {
        self.graph
            .borrow(py)
            .inner
            .gdfo(&PyNode::from(self.key.clone_ref(py)))
    }

    #[getter]
    fn parent_keys(&self, py: Python) -> Option<Vec<Py<PyAny>>> {
        self.graph
            .borrow(py)
            .inner
            .get_parent_keys(&PyNode::from(self.key.clone_ref(py)))
            .map(|ps| ps.iter().map(|n| n.0.clone_ref(py)).collect())
    }

    #[getter]
    fn child_keys(&self, py: Python) -> Vec<Py<PyAny>> {
        self.graph
            .borrow(py)
            .inner
            .get_child_keys(&PyNode::from(self.key.clone_ref(py)))
            .map(|cs| cs.iter().map(|n| n.0.clone_ref(py)).collect())
            .unwrap_or_default()
    }

    fn __repr__(&self, py: Python) -> String {
        format!(
            "_KnownGraphNode({:?}  gdfo:{:?} par:{:?} child:{:?})",
            self.key
                .bind(py)
                .repr()
                .map(|r| r.to_string())
                .unwrap_or_default(),
            self.gdfo(py),
            self.parent_keys(py),
            self.child_keys(py),
        )
    }
}

#[pyclass(name = "_MergeSortNode")]
struct PyKnownGraphMergeSortNode {
    #[pyo3(get, set)]
    key: Py<PyAny>,
    #[pyo3(get, set)]
    merge_depth: usize,
    #[pyo3(get, set)]
    revno: Py<PyAny>,
    #[pyo3(get, set)]
    end_of_merge: bool,
}

/// Sort a collection of PyNodes by hash for use as a deterministic cache
/// key. Collisions only hurt cache hit rate, not correctness, because two
/// different sets can't produce identical sorted Vecs within a process.
fn sort_pynodes_by_hash(items: impl IntoIterator<Item = PyNode>) -> Vec<PyNode> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    let mut v: Vec<PyNode> = items.into_iter().collect();
    v.sort_by_key(|n| {
        let mut h = DefaultHasher::new();
        n.hash(&mut h);
        h.finish()
    });
    v
}

/// Lazy iterator yielding the lefthand ancestry of a starting key.
///
/// Mirrors the Python `Graph.iter_lefthand_ancestry` generator semantics:
/// each `__next__` call walks one step down the left-most parent chain.
/// Callers can break out of the iteration before the walk reaches a ghost
/// sentinel, matching Python's early-exit behaviour.
#[pyclass(name = "_LefthandAncestryIterator")]
struct PyLefthandAncestryIterator {
    provider: Py<PyAny>,
    next_key: Option<Py<PyAny>>,
    stop_keys: Py<PyAny>,
    graph_py: Py<PyAny>,
}

#[pymethods]
impl PyLefthandAncestryIterator {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&mut self, py: Python) -> PyResult<Py<PyAny>> {
        let key = match self.next_key.take() {
            Some(k) => k,
            None => return Err(pyo3::exceptions::PyStopIteration::new_err(())),
        };
        // Honour stop_keys (a Python container supporting `in`).
        if self.stop_keys.bind(py).contains(&key)? {
            return Err(pyo3::exceptions::PyStopIteration::new_err(()));
        }
        // Fetch the left-most parent for the next iteration. Python's
        // generator raises RevisionNotPresent if the key is missing from
        // the provider; mirror that exactly.
        let key_list = pyo3::types::PyList::new(py, [key.clone_ref(py)])?;
        let parent_map = self
            .provider
            .bind(py)
            .call_method1("get_parent_map", (key_list,))?;
        let parents = parent_map.get_item(key.clone_ref(py));
        match parents {
            Ok(ps) => {
                let ps_iter: Vec<Py<PyAny>> = {
                    let mut out = Vec::new();
                    for item in ps.try_iter()? {
                        out.push(item?.unbind());
                    }
                    out
                };
                if ps_iter.is_empty() {
                    self.next_key = None;
                } else {
                    self.next_key = Some(ps_iter.into_iter().next().unwrap());
                }
                Ok(key)
            }
            Err(_) => Err(RevisionNotPresent::new_err((
                key,
                self.graph_py.clone_ref(py),
            ))),
        }
    }
}

/// Adapter that wraps a graph whose keys are 1-tuples of ids. Each method
/// takes ids, wraps them in tuples, delegates to the underlying graph,
/// then unwraps the tuples in the response. Mirrors
/// `vcsgraph.graph.GraphThunkIdsToKeys`.
#[pyclass(name = "GraphThunkIdsToKeys")]
struct PyGraphThunkIdsToKeys {
    graph: Py<PyAny>,
}

#[pymethods]
impl PyGraphThunkIdsToKeys {
    #[new]
    fn new(graph: Py<PyAny>) -> Self {
        PyGraphThunkIdsToKeys { graph }
    }

    fn topo_sort(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
        let result = self.graph.bind(py).call_method0("topo_sort")?;
        let mut out = Vec::new();
        for item in result.try_iter()? {
            let tup = item?;
            out.push(tup.get_item(0)?.unbind());
        }
        Ok(out)
    }

    fn heads<'py>(
        &self,
        py: Python<'py>,
        ids: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let mut as_keys: Vec<Py<PyAny>> = Vec::new();
        for item in ids.bind(py).try_iter()? {
            let item = item?;
            let tup = pyo3::types::PyTuple::new(py, [item.unbind()])?;
            as_keys.push(tup.into_any().unbind());
        }
        let head_keys = self
            .graph
            .bind(py)
            .call_method1("heads", (pyo3::types::PyList::new(py, as_keys)?,))?;
        let mut out_items: Vec<Py<PyAny>> = Vec::new();
        for h in head_keys.try_iter()? {
            let h = h?;
            out_items.push(h.get_item(0)?.unbind());
        }
        pyo3::types::PySet::new(py, out_items)
    }

    fn merge_sort(&self, py: Python, tip_revision: Py<PyAny>) -> PyResult<Py<PyAny>> {
        let tip_tuple = pyo3::types::PyTuple::new(py, [tip_revision])?;
        let nodes = self
            .graph
            .bind(py)
            .call_method1("merge_sort", (tip_tuple,))?;
        for item in nodes.try_iter()? {
            let item = item?;
            let key = item.getattr("key")?;
            let unwrapped = key.get_item(0)?;
            item.setattr("key", unwrapped)?;
        }
        Ok(nodes.unbind())
    }

    fn add_node(&self, py: Python, revision: Py<PyAny>, parents: Py<PyAny>) -> PyResult<()> {
        let rev_tuple = pyo3::types::PyTuple::new(py, [revision])?;
        let mut parent_tuples: Vec<Py<PyAny>> = Vec::new();
        for p in parents.bind(py).try_iter()? {
            let p = p?;
            let tup = pyo3::types::PyTuple::new(py, [p.unbind()])?;
            parent_tuples.push(tup.into_any().unbind());
        }
        self.graph.bind(py).call_method1(
            "add_node",
            (rev_tuple, pyo3::types::PyList::new(py, parent_tuples)?),
        )?;
        Ok(())
    }
}

/// A cache of results for graph heads() calls.
///
/// The cache key is the unordered set of input keys; the value is the
/// `.heads()` result for that set. Every call returns a fresh mutable set
/// so callers can modify the result without affecting later lookups —
/// matches `vcsgraph.graph.HeadsCache`.
#[pyclass(name = "HeadsCache")]
struct PyHeadsCache {
    graph: Py<PyAny>,
    cache: std::sync::Mutex<std::collections::HashMap<Vec<PyNode>, Vec<PyNode>>>,
}

#[pymethods]
impl PyHeadsCache {
    #[new]
    fn new(graph: Py<PyAny>) -> Self {
        PyHeadsCache {
            graph,
            cache: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// The underlying graph. Exposed so Python callers can reach
    /// `cache.graph` just like the pure-Python version allowed.
    #[getter]
    fn graph(&self, py: Python) -> Py<PyAny> {
        self.graph.clone_ref(py)
    }

    /// Return the heads of `keys`. The result is a mutable Python set;
    /// callers may modify it without affecting future lookups.
    fn heads<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let key = sort_pynodes_by_hash(nodes);
        {
            let cache = self.cache.lock().unwrap();
            if let Some(cached) = cache.get(&key) {
                return pyo3::types::PySet::new(py, cached.iter().map(|n| n.0.clone_ref(py)));
            }
        }
        // Cache miss: delegate to the wrapped graph.
        let result: Vec<PyNode> = {
            let heads = self.graph.bind(py).call_method1(
                "heads",
                (pyo3::types::PyList::new(
                    py,
                    key.iter().map(|n| n.0.clone_ref(py)),
                )?,),
            )?;
            let mut out = Vec::new();
            for item in heads.try_iter()? {
                out.push(PyNode::from(item?));
            }
            out
        };
        let mut cache = self.cache.lock().unwrap();
        cache.insert(key, result.clone());
        pyo3::types::PySet::new(py, result.into_iter().map(|n| n.0))
    }
}

/// A cache of `heads()` results that returns frozen sets.
///
/// Same as [`PyHeadsCache`] but the results are immutable. Matches
/// `vcsgraph.graph.FrozenHeadsCache`.
#[pyclass(name = "FrozenHeadsCache")]
struct PyFrozenHeadsCache {
    graph: Py<PyAny>,
    cache: std::sync::Mutex<std::collections::HashMap<Vec<PyNode>, Vec<PyNode>>>,
}

#[pymethods]
impl PyFrozenHeadsCache {
    #[new]
    fn new(graph: Py<PyAny>) -> Self {
        PyFrozenHeadsCache {
            graph,
            cache: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    #[getter]
    fn graph(&self, py: Python) -> Py<PyAny> {
        self.graph.clone_ref(py)
    }

    /// Return the heads of `keys` as a frozenset.
    fn heads(&self, py: Python, keys: Py<PyAny>) -> PyResult<Py<PyAny>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let key = sort_pynodes_by_hash(nodes);
        {
            let cache = self.cache.lock().unwrap();
            if let Some(cached) = cache.get(&key) {
                let fs = py.import("builtins")?.getattr("frozenset")?;
                let items: Vec<Py<PyAny>> = cached.iter().map(|n| n.0.clone_ref(py)).collect();
                return Ok(fs.call1((items,))?.into_any().unbind());
            }
        }
        let result: Vec<PyNode> = {
            let heads = self.graph.bind(py).call_method1(
                "heads",
                (pyo3::types::PyList::new(
                    py,
                    key.iter().map(|n| n.0.clone_ref(py)),
                )?,),
            )?;
            let mut out = Vec::new();
            for item in heads.try_iter()? {
                out.push(PyNode::from(item?));
            }
            out
        };
        let mut cache = self.cache.lock().unwrap();
        cache.insert(key, result.clone());
        let fs = py.import("builtins")?.getattr("frozenset")?;
        let items: Vec<Py<PyAny>> = result.into_iter().map(|n| n.0).collect();
        Ok(fs.call1((items,))?.into_any().unbind())
    }

    /// Store a precomputed `(keys, heads)` pair directly in the cache.
    fn cache(&self, py: Python, keys: Py<PyAny>, heads: Py<PyAny>) -> PyResult<()> {
        let key_nodes = extract_iter_pynodes(py, &keys)?;
        let head_nodes = extract_iter_pynodes(py, &heads)?;
        let key = sort_pynodes_by_hash(key_nodes);
        self.cache.lock().unwrap().insert(key, head_nodes);
        Ok(())
    }
}

/// Python binding for [`CachingParentsProvider`]. Wraps an inner Python
/// parents provider (any object with a `get_parent_map(keys)` method) and
/// caches every lookup so repeated queries don't re-hit the inner provider.
///
/// Mirrors `vcsgraph.graph.CachingParentsProvider`.
#[pyclass(name = "CachingParentsProvider")]
struct PyCachingParentsProvider {
    inner: RsCachingParentsProvider<PyNode, PyParentsProviderAdapter>,
    real_provider_py: Py<PyAny>,
    /// Live Python set exposed as `.missing_keys`. Callers in breezy mutate
    /// it directly (e.g. `provider.missing_keys.clear()`), so it must be the
    /// same object across reads. We sync this with the Rust state around
    /// every `get_parent_map` call.
    missing_keys_py: Py<pyo3::types::PySet>,
}

impl PyCachingParentsProvider {
    /// Push Python set contents into Rust state.
    fn sync_missing_to_rust(&self, py: Python<'_>) -> PyResult<()> {
        let set = self.missing_keys_py.bind(py);
        let mut keys: HashSet<PyNode> = HashSet::new();
        for item in set.try_iter()? {
            keys.insert(PyNode::from(item?));
        }
        self.inner.set_missing_keys(keys);
        Ok(())
    }

    /// Pull Rust state into the Python set.
    fn sync_missing_to_py(&self, py: Python<'_>) -> PyResult<()> {
        let set = self.missing_keys_py.bind(py);
        set.clear();
        for k in self.inner.missing_keys() {
            set.add(k.0)?;
        }
        Ok(())
    }
}

#[pymethods]
impl PyCachingParentsProvider {
    /// Construct a caching wrapper around `parent_provider`, which must be
    /// either a Python object with a `get_parent_map` method or `None` when
    /// `get_parent_map` is supplied as a callable.
    #[new]
    #[pyo3(signature = (parent_provider=None, get_parent_map=None))]
    fn new(
        py: Python,
        parent_provider: Option<Py<PyAny>>,
        get_parent_map: Option<Py<PyAny>>,
    ) -> PyResult<Self> {
        let provider_obj = match (parent_provider, get_parent_map) {
            (Some(p), None) => p,
            (None, Some(cb)) => {
                // Wrap the callable into a tiny Python shim that exposes
                // a `get_parent_map` attribute, so the adapter can call it
                // uniformly.
                let builtins = py.import("builtins")?;
                let type_ = builtins.getattr("type")?;
                let ns = pyo3::types::PyDict::new(py);
                ns.set_item("get_parent_map", cb)?;
                let shim = type_.call1(("_CPPShim", (builtins.getattr("object")?,), ns))?;
                shim.call0()?.into_any().unbind()
            }
            (Some(_), Some(_)) => {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "Pass parent_provider OR get_parent_map, not both",
                ))
            }
            (None, None) => {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "Either parent_provider or get_parent_map must be supplied",
                ))
            }
        };
        let adapter = PyParentsProviderAdapter::new(provider_obj.clone_ref(py));
        let missing_keys_py = pyo3::types::PySet::empty(py)?.unbind();
        Ok(PyCachingParentsProvider {
            inner: RsCachingParentsProvider::new(adapter),
            real_provider_py: provider_obj,
            missing_keys_py,
        })
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let r = self.real_provider_py.bind(py).repr()?;
        Ok(format!("CachingParentsProvider({})", r))
    }

    #[pyo3(signature = (cache_misses=true))]
    fn enable_cache(&self, py: Python<'_>, cache_misses: bool) -> PyResult<()> {
        self.inner
            .enable_cache(cache_misses)
            .map_err(pyo3::exceptions::PyAssertionError::new_err)?;
        self.missing_keys_py.bind(py).clear();
        Ok(())
    }

    fn disable_cache(&self, py: Python<'_>) {
        self.inner.disable_cache();
        self.missing_keys_py.bind(py).clear();
    }

    /// Return a snapshot of the cache as a Python dict, or `None` if
    /// disabled.
    fn get_cached_map<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyDict>>> {
        match self.inner.get_cached_map() {
            None => Ok(None),
            Some(map) => Ok(Some(cache_map_to_pydict(py, map)?)),
        }
    }

    /// Backward-compatible access to the raw `_cache` attribute. Returns a
    /// dict snapshot when the cache is enabled, `None` otherwise. The
    /// Python tests read this directly, so mirror Python's behaviour of
    /// exposing an empty dict rather than `None` for a freshly-enabled
    /// cache.
    #[getter]
    fn _cache<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyDict>>> {
        self.get_cached_map(py)
    }

    fn get_cached_parent_map<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let hs: HashSet<PyNode> = nodes.into_iter().collect();
        let pm = self.inner.get_cached_parent_map(&hs);
        parent_map_to_pydict(py, pm)
    }

    fn get_parent_map<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let hs: HashSet<PyNode> = nodes.into_iter().collect();
        self.sync_missing_to_rust(py)?;
        let pm = self.inner.get_parent_map(&hs);
        if let Some(err) = self.inner.inner().take_error() {
            self.sync_missing_to_py(py)?;
            return Err(err);
        }
        self.sync_missing_to_py(py)?;
        parent_map_to_pydict(py, pm)
    }

    fn note_missing_key(&self, py: Python<'_>, key: Py<PyAny>) -> PyResult<()> {
        let node = PyNode::from(key);
        self.inner.note_missing_key(node.clone());
        // Mirror to the exposed Python set so callers see the new entry.
        if self.inner.missing_keys().contains(&node) {
            self.missing_keys_py.bind(py).add(node.0.clone_ref(py))?;
        }
        Ok(())
    }

    #[getter]
    fn missing_keys<'py>(&self, py: Python<'py>) -> Bound<'py, pyo3::types::PySet> {
        self.missing_keys_py.bind(py).clone()
    }
}

/// Adapter letting a Python parents provider satisfy Rust's `ParentsProvider`
/// trait. Holds a raw `Py<PyAny>` and dispatches `get_parent_map(keys)` via
/// the GIL.
///
/// The Python provider's `get_parent_map` must accept an iterable of keys
/// and return a dict-like `{key: parents_tuple}` (missing keys are treated
/// as ghosts). If the Python call raises, the error is captured in
/// `last_error` and an empty map is returned; callers should drain
/// `last_error` after a Rust-level operation completes and surface the error
/// back to Python.
struct PyParentsProviderAdapter {
    provider: Py<PyAny>,
    last_error: std::sync::Mutex<Option<PyErr>>,
}

impl PyParentsProviderAdapter {
    fn new(provider: Py<PyAny>) -> Self {
        Self {
            provider,
            last_error: std::sync::Mutex::new(None),
        }
    }

    fn take_error(&self) -> Option<PyErr> {
        self.last_error.lock().unwrap().take()
    }

    fn set_error(&self, err: PyErr) {
        let mut slot = self.last_error.lock().unwrap();
        if slot.is_none() {
            *slot = Some(err);
        }
    }

    fn has_error(&self) -> bool {
        self.last_error.lock().unwrap().is_some()
    }
}

impl ParentsProvider<PyNode> for PyParentsProviderAdapter {
    fn get_parent_map(&self, keys: &HashSet<PyNode>) -> ParentMap<PyNode> {
        if self.has_error() {
            return ParentMap::new();
        }
        Python::attach(|py| {
            let key_list = pyo3::types::PyList::empty(py);
            for k in keys {
                if let Err(err) = key_list.append(k.0.bind(py)) {
                    self.set_error(err);
                    return ParentMap::new();
                }
            }
            let result = match self
                .provider
                .bind(py)
                .call_method1("get_parent_map", (key_list,))
            {
                Ok(r) => r,
                Err(err) => {
                    self.set_error(err);
                    return ParentMap::new();
                }
            };
            match result.extract::<ParentMap<PyNode>>() {
                Ok(pm) => pm,
                Err(err) => {
                    self.set_error(err);
                    ParentMap::new()
                }
            }
        })
    }
}

#[pyclass(name = "_RustGraph")]
struct PyGraph {
    inner: RsGraph<PyNode, PyParentsProviderAdapter>,
    provider_py: Py<PyAny>,
}

impl PyGraph {
    fn check_error(&self) -> PyResult<()> {
        if let Some(err) = self.inner.parents_provider().take_error() {
            return Err(err);
        }
        Ok(())
    }
}

fn extract_iter_pynodes(py: Python, obj: &Py<PyAny>) -> PyResult<Vec<PyNode>> {
    let mut out = Vec::new();
    for item in obj.bind(py).try_iter()? {
        out.push(PyNode::from(item?));
    }
    Ok(out)
}

fn parent_map_to_pydict<'py>(
    py: Python<'py>,
    pm: ParentMap<PyNode>,
) -> PyResult<Bound<'py, PyDict>> {
    let d = PyDict::new(py);
    for (k, v) in pm {
        d.set_item(k.0, v.into_pyobject(py)?)?;
    }
    Ok(d)
}

fn cache_map_to_pydict<'py>(
    py: Python<'py>,
    map: rustc_hash::FxHashMap<PyNode, Parents<PyNode>>,
) -> PyResult<Bound<'py, PyDict>> {
    let d = PyDict::new(py);
    for (k, v) in map {
        d.set_item(k.0, v.into_pyobject(py)?)?;
    }
    Ok(d)
}

#[pymethods]
impl PyGraph {
    #[new]
    fn new(py: Python, parents_provider: Py<PyAny>) -> PyResult<Self> {
        let adapter = PyParentsProviderAdapter::new(parents_provider.clone_ref(py));
        Ok(PyGraph {
            inner: RsGraph::new(adapter),
            provider_py: parents_provider,
        })
    }

    /// Return the wrapped parents provider as given at construction time.
    #[getter]
    fn parents_provider(&self, py: Python) -> Py<PyAny> {
        self.provider_py.clone_ref(py)
    }

    fn get_parent_map<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let pm = self.inner.get_parent_map(nodes);
        if let Some(err) = self.inner.parents_provider().take_error() {
            return Err(err);
        }
        parent_map_to_pydict(py, pm)
    }

    fn get_child_map<'py>(&self, py: Python<'py>, keys: Py<PyAny>) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let cm = self.inner.get_child_map(nodes);
        if let Some(err) = self.inner.parents_provider().take_error() {
            return Err(err);
        }
        let d = PyDict::new(py);
        for (parent, children) in cm {
            let list: Vec<Py<PyAny>> = children.into_iter().map(|n| n.0).collect();
            d.set_item(parent.0, list)?;
        }
        Ok(d)
    }

    fn iter_topo_order(&self, py: Python, revisions: Py<PyAny>) -> PyResult<Vec<Py<PyAny>>> {
        let nodes = extract_iter_pynodes(py, &revisions)?;
        let res = self
            .inner
            .iter_topo_order(nodes)
            .map(|v| v.into_iter().map(|n| n.0).collect())
            .map_err(|e| match e {
                vcs_graph::Error::Cycle(_) => {
                    GraphCycleError::new_err(("cycle in graph while iter_topo_order",))
                }
                other => pyo3::exceptions::PyValueError::new_err(format!("{:?}", other)),
            });
        self.check_error()?;
        res
    }

    /// Return a lazy iterator over the lefthand ancestry of `start_key`.
    ///
    /// The iterator yields revisions one at a time, walking left-most
    /// parents, and raises RevisionNotPresent if a key in the walk is
    /// missing from the provider. `stop_keys` is any container supporting
    /// the `in` operator; iteration stops when the current key is in it.
    #[pyo3(signature = (start_key, stop_keys=None))]
    fn iter_lefthand_ancestry(
        slf: PyRef<'_, Self>,
        py: Python,
        start_key: Py<PyAny>,
        stop_keys: Option<Py<PyAny>>,
    ) -> PyResult<PyLefthandAncestryIterator> {
        let stop_keys = match stop_keys {
            Some(s) => s,
            None => pyo3::types::PyTuple::empty(py).into_any().unbind(),
        };
        Ok(PyLefthandAncestryIterator {
            provider: slf.provider_py.clone_ref(py),
            next_key: Some(start_key),
            stop_keys,
            graph_py: slf.into_pyobject(py)?.into_any().unbind(),
        })
    }

    /// Iterate ancestry, yielding `(key, parents_tuple_or_None)` pairs.
    fn iter_ancestry<'py>(
        &self,
        py: Python<'py>,
        revision_ids: Py<PyAny>,
    ) -> PyResult<Vec<Bound<'py, pyo3::types::PyTuple>>> {
        let nodes = extract_iter_pynodes(py, &revision_ids)?;
        let pairs = self.inner.iter_ancestry(nodes);
        self.check_error()?;
        pairs
            .into_iter()
            .map(|(k, parents)| {
                let parents_obj: Py<PyAny> = match parents {
                    Parents::Known(ps) => {
                        let list: Vec<Py<PyAny>> = ps.into_iter().map(|n| n.0).collect();
                        pyo3::types::PyTuple::new(py, list)?.into_any().unbind()
                    }
                    Parents::Ghost => py.None(),
                };
                pyo3::types::PyTuple::new(py, [k.0, parents_obj])
            })
            .collect()
    }

    fn find_distance_to_null(
        &self,
        py: Python,
        target_revision_id: Py<PyAny>,
        known_revision_ids: Py<PyAny>,
    ) -> PyResult<i64> {
        let target = PyNode::from(target_revision_id);
        let mut known: Vec<(PyNode, i64)> = Vec::new();
        for item in known_revision_ids.bind(py).try_iter()? {
            let (k, d): (Py<PyAny>, i64) = item?.extract()?;
            known.push((PyNode::from(k), d));
        }
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let res = self
            .inner
            .find_distance_to_null(target, known, null)
            .map_err(graph_error_to_py);
        self.check_error()?;
        res
    }

    fn find_lefthand_distances<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, PyDict>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let result = self.inner.find_lefthand_distances(nodes, null);
        self.check_error()?;
        let d = PyDict::new(py);
        for (k, dist) in result {
            d.set_item(k.0, dist)?;
        }
        Ok(d)
    }

    /// Return the heads from amongst keys.
    ///
    /// This is done by searching the ancestries of each key.  Any key that is
    /// reachable from another key is not returned; all the others are.
    ///
    /// This operation scales with the relative depth between any two keys. If
    /// any two keys are completely disconnected all ancestry of both sides
    /// will be retrieved.
    ///
    /// :param keys: An iterable of keys.
    /// :return: A set of the heads. Note that as a set there is no ordering
    ///     information. Callers will need to filter their input to create
    ///     order if they need it.
    fn heads<'py>(
        &self,
        py: Python<'py>,
        keys: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let nodes = extract_iter_pynodes(py, &keys)?;
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let result = self.inner.heads_with_null(nodes, &null);
        self.check_error()?;
        pynodes_to_pyset(py, result)
    }

    /// Determine the lowest common ancestors of the provided revisions.
    ///
    /// A lowest common ancestor is a common ancestor none of whose
    /// descendants are common ancestors.  In graphs, unlike trees, there may
    /// be multiple lowest common ancestors.
    ///
    /// This algorithm has two phases.  Phase 1 identifies border ancestors,
    /// and phase 2 filters border ancestors to determine lowest common
    /// ancestors.
    ///
    /// In phase 1, border ancestors are identified, using a breadth-first
    /// search starting at the bottom of the graph.  Searches are stopped
    /// whenever a node or one of its descendants is determined to be common
    ///
    /// In phase 2, the border ancestors are filtered to find the least
    /// common ancestors.  This is done by searching the ancestries of each
    /// border ancestor.
    ///
    /// Phase 2 is perfomed on the principle that a border ancestor that is
    /// not an ancestor of any other border ancestor is a least common
    /// ancestor.
    ///
    /// Searches are stopped when they find a node that is determined to be a
    /// common ancestor of all border ancestors, because this shows that it
    /// cannot be a descendant of any border ancestor.
    ///
    /// The scaling of this operation should be proportional to:
    ///
    /// 1. The number of uncommon ancestors
    /// 2. The number of border ancestors
    /// 3. The length of the shortest path between a border ancestor and an
    ///    ancestor of all border ancestors.
    fn find_lca<'py>(
        &self,
        py: Python<'py>,
        revisions: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let nodes = extract_iter_pynodes(py, &revisions)?;
        // Match Python's `_find_border_ancestors` precondition: None is
        // not a valid revision id. Raise InvalidRevisionId before running
        // the algorithm so callers see the same error they used to.
        for n in &nodes {
            if n.0.is_none(py) {
                return Err(InvalidRevisionId::new_err((
                    py.None(),
                    self.provider_py.clone_ref(py),
                )));
            }
        }
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let result = self.inner.find_lca(nodes, &null);
        self.check_error()?;
        pynodes_to_pyset(py, result)
    }

    /// Determine whether a revision is an ancestor of another.
    ///
    /// We answer this using heads() as heads() has the logic to perform the
    /// smallest number of parent lookups to determine the ancestral
    /// relationship between N revisions.
    fn is_ancestor(
        &self,
        py: Python,
        candidate_ancestor: Py<PyAny>,
        candidate_descendant: Py<PyAny>,
    ) -> PyResult<bool> {
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let res = self.inner.is_ancestor(
            PyNode::from(candidate_ancestor),
            PyNode::from(candidate_descendant),
            &null,
        );
        self.check_error()?;
        Ok(res)
    }

    /// Determine whether a revision is between two others.
    ///
    /// returns true if and only if:
    /// lower_bound_revid <= revid <= upper_bound_revid
    #[pyo3(signature = (revid, lower_bound_revid, upper_bound_revid))]
    fn is_between(
        &self,
        py: Python,
        revid: Py<PyAny>,
        lower_bound_revid: Option<Py<PyAny>>,
        upper_bound_revid: Option<Py<PyAny>>,
    ) -> PyResult<bool> {
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let res = self.inner.is_between(
            PyNode::from(revid),
            lower_bound_revid.map(PyNode::from),
            upper_bound_revid.map(PyNode::from),
            &null,
        );
        self.check_error()?;
        Ok(res)
    }

    /// Find the order that each revision was merged into tip.
    ///
    /// This basically just walks backwards with a stack, and walks left-first
    /// until it finds a node to stop.
    fn find_merge_order(
        &self,
        py: Python,
        tip_revision_id: Py<PyAny>,
        lca_revision_ids: Py<PyAny>,
    ) -> PyResult<Vec<Py<PyAny>>> {
        let tip = PyNode::from(tip_revision_id);
        let lcas = extract_iter_pynodes(py, &lca_revision_ids)?;
        let result = self.inner.find_merge_order(tip, lcas);
        self.check_error()?;
        Ok(result.into_iter().map(|n| n.0).collect())
    }

    /// Find the unique ancestors for a revision versus others.
    ///
    /// This returns the ancestry of unique_revision, excluding all revisions
    /// in the ancestry of common_revisions. If unique_revision is in the
    /// ancestry, then the empty set will be returned.
    ///
    /// :param unique_revision: The revision_id whose ancestry we are
    ///     interested in.
    ///     (XXX: Would this API be better if we allowed multiple revisions on
    ///     to be searched here?)
    /// :param common_revisions: Revision_ids of ancestries to exclude.
    /// :return: A set of revisions in the ancestry of unique_revision
    ///
    /// Algorithm description:
    ///
    /// 1) Walk backwards from the unique node and all common nodes.
    /// 2) When a node is seen by both sides, stop searching it in the unique
    ///    walker, include it in the common walker.
    /// 3) Stop searching when there are no nodes left for the unique walker.
    ///    At this point, you have a maximal set of unique nodes. Some of
    ///    them may actually be common, and you haven't reached them yet.
    /// 4) Start new searchers for the unique nodes, seeded with the
    ///    information you have so far.
    /// 5) Continue searching, stopping the common searches when the search
    ///    tip is an ancestor of all unique nodes.
    /// 6) Aggregate together unique searchers when they are searching the
    ///    same tips. When all unique searchers are searching the same node,
    ///    stop move it to a single 'all_unique_searcher'.
    /// 7) The 'all_unique_searcher' represents the very 'tip' of searching.
    ///    Most of the time this produces very little important information.
    ///    So don't step it as quickly as the other searchers.
    /// 8) Search is done when all common searchers have completed.
    fn find_unique_ancestors<'py>(
        &self,
        py: Python<'py>,
        unique_revision: Py<PyAny>,
        common_revisions: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let unique = PyNode::from(unique_revision);
        let commons = extract_iter_pynodes(py, &common_revisions)?;
        let result = self.inner.find_unique_ancestors(unique, commons);
        self.check_error()?;
        pynodes_to_pyset(py, result)
    }

    /// Determine the graph difference between two revisions.
    fn find_difference<'py>(
        &self,
        py: Python<'py>,
        left_revision: Py<PyAny>,
        right_revision: Py<PyAny>,
    ) -> PyResult<(
        Bound<'py, pyo3::types::PySet>,
        Bound<'py, pyo3::types::PySet>,
    )> {
        let left = PyNode::from(left_revision);
        let right = PyNode::from(right_revision);
        let (l, r) = self.inner.find_difference(left, right);
        self.check_error()?;
        Ok((pynodes_to_pyset(py, l)?, pynodes_to_pyset(py, r)?))
    }

    /// Find a unique LCA.
    ///
    /// Find lowest common ancestors.  If there is no unique common
    /// ancestor, find the lowest common ancestors of those ancestors.
    ///
    /// Iteration stops when a unique lowest common ancestor is found.
    /// The graph origin is necessarily a unique lowest common ancestor.
    ///
    /// Note that None is not an acceptable substitute for NULL_REVISION.
    /// in the input for this method.
    ///
    /// :param count_steps: If True, the return value will be a tuple of
    ///     (unique_lca, steps) where steps is the number of times that
    ///     find_lca was run.  If False, only unique_lca is returned.
    #[pyo3(signature = (left_revision, right_revision, count_steps=false))]
    fn find_unique_lca(
        &self,
        py: Python,
        left_revision: Py<PyAny>,
        right_revision: Py<PyAny>,
        count_steps: bool,
    ) -> PyResult<Py<PyAny>> {
        let left = PyNode::from(left_revision.clone_ref(py));
        let right = PyNode::from(right_revision.clone_ref(py));
        let null_bytes = pyo3::types::PyBytes::new(py, NULL_REVISION);
        let null = PyNode::from(null_bytes.into_any().unbind());
        let res = self.inner.find_unique_lca(left, right, &null);
        self.check_error()?;
        match res {
            Some((key, steps)) => {
                if count_steps {
                    Ok(pyo3::types::PyTuple::new(
                        py,
                        [key.0, steps.into_pyobject(py)?.into_any().unbind()],
                    )?
                    .into_any()
                    .unbind())
                } else {
                    Ok(key.0)
                }
            }
            None => Err(NoCommonAncestor::new_err((left_revision, right_revision))),
        }
    }

    /// Find the first lefthand ancestor of tip_key that merged merged_key.
    ///
    /// We do this by first finding the descendants of merged_key, then
    /// walking through the lefthand ancestry of tip_key until we find a key
    /// that doesn't descend from merged_key.  Its child is the key that
    /// merged merged_key.
    ///
    /// :return: The first lefthand ancestor of tip_key to merge merged_key.
    ///     merged_key if it is a lefthand ancestor of tip_key.
    ///     None if no ancestor of tip_key merged merged_key.
    fn find_lefthand_merger(
        &self,
        merged_key: Py<PyAny>,
        tip_key: Py<PyAny>,
    ) -> PyResult<Option<Py<PyAny>>> {
        let merged = PyNode::from(merged_key);
        let tip = PyNode::from(tip_key);
        let res = self.inner.find_lefthand_merger(merged, tip).map(|n| n.0);
        self.check_error()?;
        Ok(res)
    }

    /// Find descendants of `old_key` that are ancestors of `new_key`.
    fn find_descendants<'py>(
        &self,
        py: Python<'py>,
        old_key: Py<PyAny>,
        new_key: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let old = PyNode::from(old_key);
        let new = PyNode::from(new_key);
        let result = self.inner.find_descendants(old, new);
        self.check_error()?;
        pynodes_to_pyset(py, result)
    }

    /// Remove revisions which are children of other ones in the set.
    ///
    /// This doesn't do any graph searching, it just checks the immediate
    /// parent_map to find if there are any children which can be removed.
    ///
    /// :param revisions: A set of revision_ids
    /// :param parent_map: A mapping `{key: parents}` — the value may be
    ///     a list, tuple, or None (for ghosts).
    /// :return: A set of revision_ids with the children removed
    fn _remove_simple_descendants<'py>(
        &self,
        py: Python<'py>,
        revisions: Py<PyAny>,
        parent_map: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let revs: std::collections::HashSet<PyNode> = {
            let mut s = std::collections::HashSet::new();
            for item in revisions.bind(py).try_iter()? {
                s.insert(PyNode::from(item?));
            }
            s
        };
        // Walk parent_map.items() Python-side so we accept any dict-like
        // mapping, tuples or lists as the value.
        let items = parent_map.bind(py).call_method0("items")?;
        let mut simple = revs.clone();
        for item in items.try_iter()? {
            let item = item?;
            let key: Py<PyAny> = item.get_item(0)?.unbind();
            let parents = item.get_item(1)?;
            if parents.is_none() {
                continue;
            }
            for parent in parents.try_iter()? {
                let parent = parent?;
                let parent_node = PyNode::from(parent.unbind());
                if revs.contains(&parent_node) {
                    simple.remove(&PyNode::from(key.clone_ref(py)));
                    break;
                }
            }
        }
        let items: Vec<Py<PyAny>> = simple.into_iter().map(|n| n.0).collect();
        pyo3::types::PySet::new(py, items)
    }

    /// Find ancestors of `new_key` that may be descendants of `old_key`.
    fn _find_descendant_ancestors<'py>(
        &self,
        py: Python<'py>,
        old_key: Py<PyAny>,
        new_key: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let old = PyNode::from(old_key);
        let new = PyNode::from(new_key);
        let result = self.inner.find_descendant_ancestors(old, new);
        self.check_error()?;
        pynodes_to_pyset(py, result)
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let r = self.provider_py.bind(py).repr()?;
        Ok(format!("Graph({})", r))
    }
}

import_exception!(vcsgraph.errors, GhostRevisionsHaveNoRevno);
import_exception!(vcsgraph.errors, InvalidRevisionId);
import_exception!(vcsgraph.errors, NoCommonAncestor);
import_exception!(vcsgraph.errors, RevisionNotPresent);

fn graph_error_to_py(e: GraphError<PyNode>) -> PyErr {
    match e {
        GraphError::GhostRevision { target, ghost } => Python::attach(|py| {
            GhostRevisionsHaveNoRevno::new_err((target.0.clone_ref(py), ghost.0.clone_ref(py)))
        }),
        GraphError::RevisionNotPresent(key) => {
            Python::attach(|py| RevisionNotPresent::new_err((key.0.clone_ref(py),)))
        }
        GraphError::Cycle(_) => GraphCycleError::new_err(("cycle in graph",)),
    }
}

/// Helper: convert a set of PyNodes into a Python `set()`.
fn pynodes_to_pyset<'py>(
    py: Python<'py>,
    set: rustc_hash::FxHashSet<PyNode>,
) -> PyResult<Bound<'py, pyo3::types::PySet>> {
    let items: Vec<Py<PyAny>> = set.into_iter().map(|n| n.0).collect();
    pyo3::types::PySet::new(py, items)
}

/// Python binding for [`BfsState`]. Owns its own adapter over a Python
/// parents provider and holds the Rust BFS state as sibling fields.
///
/// `dict=true` gives the class a `__dict__`, which lets the existing Python
/// callers in `graph.py` stash a `_label` attribute on the instance for
/// debug logging. This can be removed once Phase 4 (heads / border
/// ancestors) no longer relies on it.
#[pyclass(name = "_BreadthFirstSearcher", dict)]
struct PyBFSearcher {
    state: BfsState<PyNode>,
    adapter: PyParentsProviderAdapter,
}

#[pymethods]
impl PyBFSearcher {
    #[new]
    fn new(py: Python, revisions: Py<PyAny>, parents_provider: Py<PyAny>) -> PyResult<Self> {
        let mut revs: Vec<PyNode> = Vec::new();
        for item in revisions.bind(py).try_iter()? {
            revs.push(PyNode::from(item?));
        }
        Ok(PyBFSearcher {
            state: BfsState::new(revs),
            adapter: PyParentsProviderAdapter::new(parents_provider),
        })
    }

    /// `seen` attribute — matches the Python API. Returns a set snapshot.
    #[getter]
    fn seen<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        pynodes_to_pyset(py, self.state.seen.clone())
    }

    /// `_next_query` attribute — the current search frontier.
    ///
    /// Read-only snapshot. Existing callers in `graph.py` iterate and
    /// truth-test this attribute but do not mutate it.
    #[getter]
    fn _next_query<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        pynodes_to_pyset(py, self.state.next_query().clone())
    }

    /// `_iterations` attribute — number of advance steps performed.
    #[getter]
    fn _iterations(&self) -> usize {
        self.state.iterations()
    }

    /// `_current_present` attribute — revisions present in the most recent
    /// advance. Mirrors the Python API; empty before the first advance.
    #[getter]
    fn _current_present<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        pynodes_to_pyset(py, self.state.current_present().clone())
    }

    /// `_current_ghosts` attribute — ghost revisions encountered in the most
    /// recent advance. Mirrors the Python API; empty before the first advance.
    #[getter]
    fn _current_ghosts<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        pynodes_to_pyset(py, self.state.current_ghosts().clone())
    }

    /// `_current_parents` attribute — parent map of present revisions from
    /// the most recent advance. Mirrors the Python API; empty before the
    /// first advance.
    #[getter]
    fn _current_parents<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
        let d = PyDict::new(py);
        for (rev, parents) in self.state.current_parents() {
            let tup = pyo3::types::PyTuple::new(py, parents.iter().map(|n| n.0.clone_ref(py)))?;
            d.set_item(rev.0.clone_ref(py), tup)?;
        }
        Ok(d)
    }

    /// Python `step` returns the next query set, or `()` on StopIteration.
    fn step<'py>(&mut self, py: Python<'py>) -> PyResult<Py<PyAny>> {
        let next = self.state.next_set(&self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        match next {
            Some(set) => Ok(pynodes_to_pyset(py, set)?.into_any().unbind()),
            None => Ok(pyo3::types::PyTuple::empty(py).into_any().unbind()),
        }
    }

    fn __next__<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let next = self.state.next_set(&self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        match next {
            Some(set) => pynodes_to_pyset(py, set),
            None => Err(pyo3::exceptions::PyStopIteration::new_err(())),
        }
    }

    fn next<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        self.__next__(py)
    }

    fn next_with_ghosts<'py>(
        &mut self,
        py: Python<'py>,
    ) -> PyResult<(
        Bound<'py, pyo3::types::PySet>,
        Bound<'py, pyo3::types::PySet>,
    )> {
        let next = self.state.next_with_ghosts(&self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        match next {
            Some((present, ghosts)) => Ok((
                pynodes_to_pyset(py, present)?,
                pynodes_to_pyset(py, ghosts)?,
            )),
            None => Err(pyo3::exceptions::PyStopIteration::new_err(())),
        }
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn get_state<'py>(
        &mut self,
        py: Python<'py>,
    ) -> PyResult<(
        Bound<'py, pyo3::types::PySet>,
        Bound<'py, pyo3::types::PySet>,
        Bound<'py, pyo3::types::PySet>,
    )> {
        let (started, excludes, included) = self.state.get_state(&self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        Ok((
            pynodes_to_pyset(py, started)?,
            pynodes_to_pyset(py, excludes)?,
            pynodes_to_pyset(py, included)?,
        ))
    }

    fn find_seen_ancestors<'py>(
        &self,
        py: Python<'py>,
        revisions: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let mut revs: Vec<PyNode> = Vec::new();
        for item in revisions.bind(py).try_iter()? {
            revs.push(PyNode::from(item?));
        }
        let result = self.state.find_seen_ancestors(revs, &self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        pynodes_to_pyset(py, result)
    }

    fn stop_searching_any<'py>(
        &mut self,
        py: Python<'py>,
        revisions: Py<PyAny>,
    ) -> PyResult<Bound<'py, pyo3::types::PySet>> {
        let mut revs: Vec<PyNode> = Vec::new();
        for item in revisions.bind(py).try_iter()? {
            revs.push(PyNode::from(item?));
        }
        let stopped = self.state.stop_searching_any(revs);
        pynodes_to_pyset(py, stopped)
    }

    fn start_searching(&mut self, py: Python, revisions: Py<PyAny>) -> PyResult<Py<PyAny>> {
        let mut revs: Vec<PyNode> = Vec::new();
        for item in revisions.bind(py).try_iter()? {
            revs.push(PyNode::from(item?));
        }
        let res = self.state.start_searching(revs, &self.adapter);
        if let Some(err) = self.adapter.take_error() {
            return Err(err);
        }
        match res {
            Some((present, ghosts)) => {
                let pres = pynodes_to_pyset(py, present)?;
                let gh = pynodes_to_pyset(py, ghosts)?;
                Ok(
                    pyo3::types::PyTuple::new(py, [pres.into_any(), gh.into_any()])?
                        .into_any()
                        .unbind(),
                )
            }
            // In Next mode Python returns None (the function has no explicit
            // return), which matches our None here.
            None => Ok(py.None()),
        }
    }

    fn __repr__(&self, py: Python) -> PyResult<String> {
        let prefix = if self.state.iterations() > 0 {
            "searching"
        } else {
            "starting"
        };
        let seen_repr =
            pyo3::types::PyList::new(py, self.state.seen.iter().map(|n| n.0.clone_ref(py)))?
                .repr()?;
        let next_repr = {
            let next_keys: Vec<Py<PyAny>> = self
                .state
                .started_keys
                .iter()
                .map(|n| n.0.clone_ref(py))
                .collect();
            pyo3::types::PyList::new(py, next_keys)?.repr()?
        };
        Ok(format!(
            "_BreadthFirstSearcher(iterations={}, {}={}, seen={})",
            self.state.iterations(),
            prefix,
            next_repr,
            seen_repr
        ))
    }
}

#[pymodule]
fn _graph_rs(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(invert_parent_map))?;
    m.add_wrapped(wrap_pyfunction!(collapse_linear_regions))?;
    m.add_wrapped(wrap_pyfunction!(merge_sort))?;
    m.add_class::<TopoSorter>()?;
    m.add_class::<MergeSorter>()?;
    m.add_class::<PyKnownGraph>()?;
    m.add_class::<PyKnownGraphMergeSortNode>()?;
    m.add_class::<PyKnownGraphNode>()?;
    m.add_class::<PyKnownGraphNodesView>()?;
    m.add_class::<PyGraph>()?;
    m.add_class::<PyBFSearcher>()?;
    m.add_class::<PyCachingParentsProvider>()?;
    m.add_class::<PyHeadsCache>()?;
    m.add_class::<PyFrozenHeadsCache>()?;
    m.add_class::<PyGraphThunkIdsToKeys>()?;
    m.add_class::<PyLefthandAncestryIterator>()?;
    m.add_class::<PyDictParentsProvider>()?;
    m.add_class::<PyStackedParentsProvider>()?;
    m.add_class::<PyCallableToParentsProviderAdapter>()?;
    Ok(())
}
