/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

mod bytearrayobject;
mod bytes;
mod bytesobject;
mod cell;
pub mod convert;
pub mod de;
pub mod error;
mod extract;
mod io;
mod none;
mod path;
mod pybuf;
mod pyset;
pub mod ser;
mod str;

#[cfg(test)]
mod tests;

pub use bytes::Bytes;
// Re-export
pub use cpython;

pub use crate::bytearrayobject::boxed_slice_to_pyobj;
pub use crate::bytearrayobject::vec_to_pyobj;
pub use crate::bytesobject::allocate_pybytes;
pub use crate::cell::PyCell;
pub use crate::error::format_py_error;
pub use crate::error::AnyhowResultExt;
pub use crate::error::PyErr;
pub use crate::error::ResultPyErrExt;
pub use crate::extract::ExtractInner;
pub use crate::extract::ExtractInnerRef;
pub use crate::io::wrap_pyio;
pub use crate::io::wrap_rust_write;
pub use crate::io::PyRustWrite;
pub use crate::io::WrappedIO;
pub use crate::none::PyNone;
pub use crate::path::Error;
pub use crate::path::PyPath;
pub use crate::path::PyPathBuf;
pub use crate::pybuf::SimplePyBuf;
pub use crate::pyset::pyset_add;
pub use crate::pyset::pyset_new;
pub use crate::str::Str;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::exc;
use cpython::FromPyObject;
use cpython::PyErr;
use cpython::PyObject;
use cpython::PyResult;
use cpython::Python;
use cpython::ToPyObject;

#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Default, Hash, Ord)]
pub struct PyNone;

impl ToPyObject for PyNone {
    type ObjectType = PyObject;

    #[inline]
    fn to_py_object(&self, py: Python) -> PyObject {
        py.None()
    }
}

impl FromPyObject<'_> for PyNone {
    fn extract(py: Python, obj: &PyObject) -> PyResult<Self> {
        if *obj == py.None() {
            Ok(PyNone)
        } else {
            Err(PyErr::new::<exc::TypeError, _>(
                py,
                format!("Expected None but received {}", obj.get_type(py).name(py)),
            ))
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::borrow::Borrow;
use std::fmt;
use std::mem;
use std::ops::Deref;
use std::path::Path;
use std::path::PathBuf;

use anyhow::Result;
use cpython::*;
use thiserror::Error;
use types::PathComponentBuf;
use types::RepoPath;
use types::RepoPathBuf;

#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Default, Hash, Ord)]
pub struct PyPathBuf(String);

#[derive(Debug, Eq, PartialEq, PartialOrd, Hash, Ord)]
pub struct PyPath(str);

#[derive(Error, Debug)]
pub enum Error {
    #[error("{0:?} is not a valid UTF-8 path")]
    NonUTF8Path(PathBuf),
}

impl PyPathBuf {
    pub fn as_pypath(&self) -> &PyPath {
        self
    }

    pub fn as_path(&self) -> &Path {
        self.0.as_ref()
    }

    pub fn to_path_buf(&self) -> PathBuf {
        Path::new(&self.0).to_path_buf()
    }

    pub fn to_repo_path_buf(self) -> Result<RepoPathBuf> {
        Ok(RepoPathBuf::from_string(self.0)?)
    }

    pub fn to_repo_path<'a>(&'a self) -> Result<&'a RepoPath> {
        Ok(RepoPath::from_str(&self.0)?)
    }

    pub fn into_utf8_bytes(self) -> Vec<u8> {
        self.0.into()
    }

    pub fn as_utf8_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

    pub fn from_string(s: String) -> Self {
        Self(s)
    }

    pub fn from_utf8_bytes(utf8_bytes: Vec<u8>) -> Result<Self> {
        Ok(Self(String::from_utf8(utf8_bytes)?))
    }
}

impl ToPyObject for PyPathBuf {
    type ObjectType = PyUnicode;

    #[inline]
    fn to_py_object(&self, py: Python) -> Self::ObjectType {
        self.0.to_py_object(py)
    }
}

impl<'source> FromPyObject<'source> for PyPathBuf {
    fn extract(py: Python, obj: &'source PyObject) -> PyResult<Self> {
        {
            let s = obj.cast_as::<PyUnicode>(py)?.data(py);
            Ok(Self(s.to_string(py)?.into()))
        }
    }
}

impl TryFrom<PathBuf> for PyPathBuf {
    type Error = anyhow::Error;

    fn try_from(path: PathBuf) -> Result<Self> {
        path.as_path().try_into()
    }
}

impl<'a> TryFrom<&'a Path> for PyPathBuf {
    type Error = anyhow::Error;

    fn try_from(path: &'a Path) -> Result<Self> {
        Ok(Self(
            path.to_str()
                .ok_or_else(|| Error::NonUTF8Path(path.to_path_buf()))?
                .into(),
        ))
    }
}

impl From<String> for PyPathBuf {
    fn from(s: String) -> PyPathBuf {
        Self(s)
    }
}

impl<'a> From<&'a RepoPath> for PyPathBuf {
    fn from(repo_path: &'a RepoPath) -> PyPathBuf {
        PyPathBuf(repo_path.as_str().to_owned())
    }
}

impl From<RepoPathBuf> for PyPathBuf {
    fn from(repo_path_buf: RepoPathBuf) -> PyPathBuf {
        PyPathBuf(repo_path_buf.into_string())
    }
}

impl From<PathComponentBuf> for PyPathBuf {
    fn from(path_component_buf: PathComponentBuf) -> PyPathBuf {
        PyPathBuf(path_component_buf.into_string())
    }
}

impl fmt::Display for PyPathBuf {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&*self.0, formatter)
    }
}

impl From<PyPathBuf> for String {
    fn from(path: PyPathBuf) -> String {
        path.0
    }
}

impl Deref for PyPathBuf {
    type Target = PyPath;
    fn deref(&self) -> &Self::Target {
        unsafe { mem::transmute(&*self.0) }
    }
}

impl AsRef<PyPath> for PyPathBuf {
    fn as_ref(&self) -> &PyPath {
        self
    }
}

impl PyPath {
    pub fn from_str(s: &str) -> &Self {
        unsafe { mem::transmute(s) }
    }

    pub fn as_path(&self) -> &Path {
        self.0.as_ref()
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }

    pub fn as_utf8_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

    pub fn to_path_buf(&self) -> PathBuf {
        Path::new(&self.0).to_path_buf()
    }

    pub fn to_repo_path(&self) -> Result<&RepoPath> {
        Ok(RepoPath::from_str(&self.0)?)
    }

    pub fn to_repo_path_buf(&self) -> Result<RepoPathBuf> {
        Ok(RepoPathBuf::from_string(self.0.to_string())?)
    }

    pub fn into_utf8_bytes(&self) -> Vec<u8> {
        self.0.into()
    }
}

impl ToOwned for PyPath {
    type Owned = PyPathBuf;

    fn to_owned(&self) -> Self::Owned {
        PyPathBuf(self.0.to_string())
    }
}

impl Borrow<PyPath> for PyPathBuf {
    fn borrow(&self) -> &PyPath {
        self
    }
}

impl RefFromPyObject for PyPath {
    fn with_extracted<F, R>(py: Python, obj: &PyObject, f: F) -> PyResult<R>
    where
        F: FnOnce(&PyPath) -> R,
    {
        {
            let s = obj.cast_as::<PyUnicode>(py)?.to_string(py)?;
            Ok(f(PyPath::from_str(s.as_ref())))
        }
    }
}

impl AsRef<PyPath> for PyPath {
    fn as_ref(&self) -> &PyPath {
        self
    }
}

impl AsRef<Path> for PyPath {
    fn as_ref(&self) -> &Path {
        self.as_path()
    }
}

impl fmt::Display for PyPath {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, formatter)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! A simple `Py_buffer` wrapper that allows zero-copy reading of Python
//! owned memory.

// The objects in memory have a relationship like:
//
// ```text
//     SimplePyBuf<T>   |        Raw Data     Python object
//     +-----------+    |        +-------+    +-----------+
//     | Py_buffer |    |        |       | <-- owns -- _  |
//     | +-------+ |    |        |       |    +-----------+
//     | | *buf -- points to --> |       |
//     | |  len  | |    |        |       |
//     | +-------+ |    |        +-------+
//     +-----------+    |
//                      |
//      Rust-managed    |   Python-managed
// ```
//
// Notes:
// - Raw data is owned by (embedded in, or pointed by) the Python object.
//   Raw data gets freed when the Python object is destructed.
// - Py_buffer is not a Python object but a Python-defined C struct allowing
//   native code to access "Raw data" directly. When constructing Py_buffer
//   from a Python object, the refcount of that Python object increases.
//   The refcount decreases when Py_buffer gets destructed via PyBuffer_Release.
// - Py_buffer is used to expose the raw pointer and length.
// - Memory alignment is up to the actual implementation of "Python object".
//   For a mmap buffer, the libc mmap function guarantees that.

use std::marker::PhantomData;
use std::mem;
use std::slice;

use cpython::PyObject;
use cpython::Python;
use python3_sys as cpy;

pub struct SimplePyBuf<T>(cpy::Py_buffer, PhantomData<T>);

// Since the buffer is read-only and Python cannot move the raw buffer (because
// we own the Py_buffer struct). It's safe to share and use SimplePyBuf in other
// threads.
unsafe impl<T> Send for SimplePyBuf<T> {}
unsafe impl<T> Sync for SimplePyBuf<T> {}

unsafe fn is_safe_type(obj: &PyObject) -> bool {
    if cpy::PyByteArray_Check(obj.as_ptr()) != 0 {
        return true;
    }
    if cpy::PyBytes_Check(obj.as_ptr()) != 0 {
        return true;
    }
    {
        if cpy::PyMemoryView_Check(obj.as_ptr()) != 0 {
            return true;
        }
    }
    return false;
}

impl<T: Copy> SimplePyBuf<T> {
    pub fn new(py: Python<'_>, obj: &PyObject) -> Self {
        // Note about GC on obj:
        //
        // Practically, obj here is some low-level, non-container ones like
        // bytes or memoryview that does not support GC (i.e. do not have
        // Py_TPFLAGS_HAVE_GC set).  refcount is the only way to release them.
        // So no need to pay extra attention on them - SimplePyBuf will get
        // refcount right and that's enough.
        //
        // Otherwise (obj is a container type that does support GC), whoever
        // owns this SimplePyBuf in the Rust world needs to do one of the
        // following:
        //   - implement tp_traverse in its Python class
        //   - call PyObject_GC_UnTrack to let GC ignore obj

        // Note about buffer mutability:
        //
        // The code here wants to access the buffer without taking Python GIL.
        // Therefore `obj` should be a read-only object. That is true for Python
        // bytes or buffer(some_other_immutable_obj). For now, explicitly
        // allow those two types. Beware that `PyBuffer_Check` won't guarnatee
        // its inner object is also immutable.
        unsafe {
            if !is_safe_type(obj) {
                let ty = obj.get_type(py);
                panic!("potentially unsafe type for SimplePyBuf: {}", ty.name(py));
            }

            let mut buf = mem::zeroed::<SimplePyBuf<T>>();
            let r = cpy::PyObject_GetBuffer(obj.as_ptr(), &mut buf.0, cpy::PyBUF_SIMPLE);
            if r == -1 {
                panic!("failed to get Py_buffer");
            }
            buf
        }
    }
}

impl<T> AsRef<[T]> for SimplePyBuf<T> {
    #[inline]
    fn as_ref(&self) -> &[T] {
        let len = self.0.len as usize / mem::size_of::<T>();
        unsafe { slice::from_raw_parts(self.0.buf as *const T, len) }
    }
}

impl<T> Drop for SimplePyBuf<T> {
    fn drop(&mut self) {
        let _gil = Python::acquire_gil();
        unsafe { cpy::PyBuffer_Release(&mut self.0) }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::ptr;

use python3_sys as ffi;

pub fn pyset_new(py: cpython::Python<'_>) -> cpython::PyResult<cpython::PyObject> {
    unsafe {
        let inner = ffi::PySet_New(ptr::null_mut());
        if inner.is_null() {
            return Err(cpython::PyErr::new::<cpython::exc::RuntimeError, _>(
                py,
                "Could not allocate set",
            ));
        }
        Ok(cpython::PyObject::from_owned_ptr(py, inner))
    }
}

pub fn pyset_add<V: cpython::ToPyObject>(
    py: cpython::Python<'_>,
    set: &mut cpython::PyObject,
    value: V,
) -> cpython::PyResult<()> {
    value.with_borrowed_ptr(py, |value| unsafe {
        let result = ffi::PySet_Add(set.as_ptr(), value);
        if result != -1 {
            Ok(())
        } else {
            Err(cpython::PyErr::fetch(py))
        }
    })
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::marker::PhantomData;

use cpython::PyBytes;
use cpython::PyDict;
use cpython::PyObject;
use cpython::PyResult;
use cpython::PyTuple;
use cpython::Python;
use cpython::PythonObject;
use cpython::ToPyObject;
use serde::ser;
use serde::Serialize;

use crate::PyErr as Error;

/// Serialize into Python object.
pub fn to_object<T>(py: Python, value: &T) -> PyResult<PyObject>
where
    T: Serialize,
{
    let serializer = Serializer { py };
    value.serialize(&serializer).map_err(Into::into)
}

// ------- Serializer Types -------

type Result<T> = std::result::Result<T, Error>;

struct Serializer<'a> {
    py: Python<'a>,
}

trait PyCollectItems {
    fn collect_items(py: Python, items: &[PyObject]) -> Result<PyObject>;
}

trait PyBuildVariant {
    fn build_variant(&self, py: Python, obj: PyObject) -> Result<PyObject>;
}

struct PyItems<'a, C, V> {
    py: Python<'a>,
    items: Vec<PyObject>,
    collector: PhantomData<C>,
    variant: V,
}

impl<'a, C, V> PyItems<'a, C, V> {
    fn new(py: Python<'a>, variant: V) -> Self {
        Self {
            py,
            items: Default::default(),
            collector: PhantomData,
            variant,
        }
    }
}

impl<'a, C, V> PyItems<'a, C, V>
where
    C: PyCollectItems,
    V: PyBuildVariant,
{
    fn collect(&self) -> Result<PyObject> {
        let obj = C::collect_items(self.py, &self.items)?;
        self.variant.build_variant(self.py, obj)
    }
}

impl PyBuildVariant for &'static str {
    fn build_variant(&self, py: Python, obj: PyObject) -> Result<PyObject> {
        enum_variant(py, self, obj)
    }
}

impl PyBuildVariant for () {
    fn build_variant(&self, _py: Python, obj: PyObject) -> Result<PyObject> {
        Ok(obj)
    }
}

struct BuildList;
struct BuildTuple;
struct BuildDict;

impl PyCollectItems for BuildList {
    fn collect_items(py: Python, items: &[PyObject]) -> Result<PyObject> {
        Ok(items.to_py_object(py).into_object())
    }
}

impl PyCollectItems for BuildTuple {
    fn collect_items(py: Python, items: &[PyObject]) -> Result<PyObject> {
        Ok(PyTuple::new(py, items).into_object())
    }
}

impl PyCollectItems for BuildDict {
    fn collect_items(py: Python, items: &[PyObject]) -> Result<PyObject> {
        let dict = PyDict::new(py);
        for chunk in items.chunks(2) {
            if let [key, value] = chunk {
                dict.set_item(py, key, value)?;
            }
        }
        Ok(dict.into_object())
    }
}

// ------- Serde APIs -------

impl<'a> Serializer<'a> {
    fn serialize<T: ToPyObject>(&self, obj: T) -> Result<PyObject> {
        Ok(obj.into_py_object(self.py).into_object())
    }

    fn to_object<T: Serialize + ?Sized>(py: Python, value: &T) -> Result<PyObject> {
        let serializer = Serializer { py };
        value.serialize(&serializer)
    }
}

impl<'a, 'b> ser::Serializer for &'a Serializer<'b> {
    type Ok = PyObject;
    type Error = Error;

    type SerializeSeq = PyItems<'b, BuildList, ()>;
    type SerializeTuple = PyItems<'b, BuildTuple, ()>;
    type SerializeTupleStruct = PyItems<'b, BuildTuple, ()>;
    type SerializeTupleVariant = PyItems<'b, BuildTuple, &'static str>;
    type SerializeMap = PyItems<'b, BuildDict, ()>;
    type SerializeStruct = PyItems<'b, BuildDict, ()>;
    type SerializeStructVariant = PyItems<'b, BuildDict, &'static str>;

    fn serialize_bool(self, v: bool) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_i8(self, v: i8) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_i16(self, v: i16) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_i32(self, v: i32) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_i64(self, v: i64) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_u8(self, v: u8) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_u16(self, v: u16) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_u32(self, v: u32) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_u64(self, v: u64) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_f32(self, v: f32) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_f64(self, v: f64) -> Result<PyObject> {
        self.serialize(v)
    }

    fn serialize_char(self, v: char) -> Result<PyObject> {
        self.serialize_str(&v.to_string())
    }

    fn serialize_str(self, v: &str) -> Result<PyObject> {
        Ok(v.to_py_object(self.py).into_object())
    }

    fn serialize_bytes(self, v: &[u8]) -> Result<PyObject> {
        Ok(PyBytes::new(self.py, v).into_object())
    }

    fn serialize_none(self) -> Result<PyObject> {
        Ok(self.py.None())
    }

    fn serialize_some<T>(self, value: &T) -> Result<PyObject>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(self)
    }

    fn serialize_unit(self) -> Result<PyObject> {
        self.serialize_none()
    }

    fn serialize_unit_struct(self, _name: &'static str) -> Result<PyObject> {
        self.serialize_unit()
    }

    fn serialize_unit_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        variant: &'static str,
    ) -> Result<PyObject> {
        self.serialize_str(variant)
    }

    fn serialize_newtype_struct<T>(self, _name: &'static str, value: &T) -> Result<PyObject>
    where
        T: ?Sized + Serialize,
    {
        value.serialize(self)
    }

    fn serialize_newtype_variant<T>(
        self,
        _name: &'static str,
        _variant_index: u32,
        variant: &'static str,
        value: &T,
    ) -> Result<PyObject>
    where
        T: ?Sized + Serialize,
    {
        let value = value.serialize(self)?;
        enum_variant(self.py, variant, value)
    }

    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq> {
        Ok(Self::SerializeSeq::new(self.py, ()))
    }

    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple> {
        Ok(Self::SerializeTuple::new(self.py, ()))
    }

    fn serialize_tuple_struct(
        self,
        _name: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleStruct> {
        Ok(Self::SerializeTupleStruct::new(self.py, ()))
    }

    fn serialize_tuple_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleVariant> {
        Ok(Self::SerializeTupleVariant::new(self.py, variant))
    }

    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap> {
        Ok(Self::SerializeMap::new(self.py, ()))
    }

    fn serialize_struct(self, _name: &'static str, _len: usize) -> Result<Self::SerializeStruct> {
        Ok(Self::SerializeStruct::new(self.py, ()))
    }

    fn serialize_struct_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeStructVariant> {
        Ok(Self::SerializeStructVariant::new(self.py, variant))
    }
}

macro_rules! impl_seq {
    ($trait: ty,  $name: tt) => {
        impl<C, V> $trait for PyItems<'_, C, V>
        where
            C: PyCollectItems,
            V: PyBuildVariant,
        {
            type Ok = PyObject;
            type Error = Error;

            fn $name<T>(&mut self, value: &T) -> Result<()>
            where
                T: ?Sized + Serialize,
            {
                let obj = Serializer::to_object(self.py, value)?;
                self.items.push(obj);
                Ok(())
            }

            fn end(self) -> Result<PyObject> {
                self.collect()
            }
        }
    };
}

impl_seq!(ser::SerializeSeq, serialize_element);
impl_seq!(ser::SerializeTuple, serialize_element);
impl_seq!(ser::SerializeTupleStruct, serialize_field);
impl_seq!(ser::SerializeTupleVariant, serialize_field);

impl<V> ser::SerializeMap for PyItems<'_, BuildDict, V>
where
    V: PyBuildVariant,
{
    type Ok = PyObject;
    type Error = Error;

    fn serialize_key<T>(&mut self, key: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        let obj = Serializer::to_object(self.py, key)?;
        self.items.push(obj);
        Ok(())
    }

    fn serialize_value<T>(&mut self, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        let obj = Serializer::to_object(self.py, value)?;
        self.items.push(obj);
        Ok(())
    }

    fn end(self) -> Result<PyObject> {
        self.collect()
    }
}

impl<V> ser::SerializeStruct for PyItems<'_, BuildDict, V>
where
    V: PyBuildVariant,
{
    type Ok = PyObject;
    type Error = Error;

    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        let key = Serializer::to_object(self.py, key)?;
        let value = Serializer::to_object(self.py, value)?;
        self.items.push(key);
        self.items.push(value);
        Ok(())
    }

    fn end(self) -> Result<PyObject> {
        self.collect()
    }
}

impl<V> ser::SerializeStructVariant for PyItems<'_, BuildDict, V>
where
    V: PyBuildVariant,
{
    type Ok = PyObject;
    type Error = Error;

    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<()>
    where
        T: ?Sized + Serialize,
    {
        ser::SerializeStruct::serialize_field(self, key, value)
    }

    fn end(self) -> Result<PyObject> {
        self.collect()
    }
}

// ------- Utilities -------

/// Convert `value` to `{key: value}` to represent an enum variant.
fn enum_variant(py: Python, key: &'static str, value: PyObject) -> Result<PyObject> {
    let dict = PyDict::new(py);
    dict.set_item(py, key, value)?;
    Ok(dict.into_object())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::*;

use crate::Bytes;

#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Default, Hash, Ord)]
pub struct Str(crate::Bytes);

impl ToPyObject for Str {
    type ObjectType = PyUnicode;

    #[inline]
    fn to_py_object(&self, py: Python) -> Self::ObjectType {
        PyUnicode::new(py, std::str::from_utf8(self.0.as_ref()).unwrap())
    }
}

impl From<Bytes> for Str {
    fn from(b: Bytes) -> Str {
        Str(b)
    }
}

impl From<String> for Str {
    fn from(s: String) -> Str {
        Str(Bytes::from(s))
    }
}

impl From<Vec<u8>> for Str {
    fn from(s: Vec<u8>) -> Str {
        Str(Bytes::from(s))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::fmt::Debug;

use cpython::PyClone;
use cpython::Python;
use serde::Deserialize;
use serde::Serialize;
use serde_bytes::ByteBuf;
use serde_cbor::Value;

#[test]
fn test_serde_basic_types() {
    check_serde_round_trip(&42);
    check_serde_round_trip(&Some(true));
    check_serde_round_trip(&false);
    check_serde_round_trip(&"abc".to_string());
    check_serde_round_trip(&b"abc".to_vec());
    check_serde_round_trip(&(1, (), Some(false), 3));
    check_serde_round_trip(&vec![1, 2, 3]);
    check_serde_round_trip(&());
    check_serde_round_trip(&ByteBuf::from(b"abc".to_vec()));
}

#[test]
fn test_serde_basic_structs() {
    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    struct S;

    check_serde_round_trip(&S);

    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    struct A {
        i: i64,
        s: String,
        u: S,
        e: (),
        t: (bool, (u64, f32)),
        v: Vec<Option<bool>>,
        m: HashMap<u64, String>,
        #[serde(with = "serde_bytes")]
        b: Vec<u8>,
    }

    let a = A {
        i: i64::MIN,
        s: "foo".to_string(),
        u: S,
        e: (),
        t: (true, (u64::MAX, -2.0)),
        v: vec![Some(true), None, Some(false)],
        m: example_hashmap(),
        b: b"abcdef".to_vec(),
    };

    check_serde_round_trip(&a);
}

#[test]
fn test_serde_nested_structs() {
    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    struct A(String, usize);

    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    struct B(Option<Box<A>>, Option<Box<A>>);

    let b = B(None, Some(Box::new(A("abc".to_string(), 42))));
    check_serde_round_trip(&b);
}

#[test]
fn test_serde_enums() {
    #[derive(PartialEq, Debug, Serialize, Deserialize)]
    enum E {
        A,
        B(bool, bool),
        C { x: u8, y: Option<i8> },
        D(Option<Box<E>>),
    }

    let values = [
        E::A,
        E::B(true, false),
        E::C { x: 42, y: None },
        E::D(None),
        E::D(Some(Box::new(E::B(false, true)))),
    ];

    for value in &values {
        check_serde_round_trip(value);
    }
}

#[test]
fn test_serde_deserialize_string_as_bytes() {
    #[derive(Serialize)]
    struct S(String);

    #[derive(Deserialize)]
    struct B(ByteBuf);

    let gil = Python::acquire_gil();
    let py = gil.python();
    let s = S("abc-æ–‡å—".to_string());
    let obj = crate::ser::to_object(py, &s).unwrap();

    let b: B = crate::de::from_object(py, obj.clone_ref(py)).unwrap();
    assert_eq!(b.0.as_ref(), s.0.as_bytes());
}

fn example_hashmap() -> HashMap<u64, String> {
    let mut m = HashMap::new();
    for i in 1..10 {
        m.insert(i, i.to_string());
    }
    m
}

fn check_serde_round_trip<S>(value: &S)
where
    S: Serialize + PartialEq + Debug,
    for<'de> S: Deserialize<'de>,
{
    let gil = Python::acquire_gil();
    let py = gil.python();
    let obj = crate::ser::to_object(py, &value).unwrap();

    let other: S = crate::de::from_object(py, obj.clone_ref(py)).unwrap();
    assert_eq!(value, &other);

    // Try deserializing into a dynamic type.
    // This exercises the `deserialize_any` code path.
    let dynamic_value: Value = crate::de::from_object(py, obj).unwrap();
    let another: S = serde_cbor::value::from_value(dynamic_value).unwrap();
    assert_eq!(value, &another);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use dag::iddagstore::IdDagStore;
use dag::idmap::IdMap;
use dag::idmap::IdMapAssignHead;
use dag::namedag::NameDag;
use dag::ops::DagAlgorithm;
use dag::ops::DagPersistent;
use dag::ops::Persist;
use dag::Group;
use dag::Id;
use dag::IdDag;
use dag::IdSet;
use dag::Set;
use dag::VertexListWithOptions;
use dag::VertexName;
use minibench::bench;
use minibench::elapsed;
use nonblocking::non_blocking_result as nbr;
use tempfile::tempdir;

type ParentsFunc<'a> = Box<dyn Fn(VertexName) -> dag::Result<Vec<VertexName>> + Send + Sync + 'a>;

pub fn main() {
    let dag_dir = tempdir().unwrap();

    bench_with_iddag(|| IdDag::open(&dag_dir.path()).unwrap());
    bench_with_iddag(|| IdDag::new_in_process());

    bench_many_heads_namedag();
}

fn bench_with_iddag<S: IdDagStore + Persist>(get_empty_iddag: impl Fn() -> IdDag<S>) {
    println!("benchmarking {}", std::any::type_name::<S>());
    let parents = bindag::parse_bindag(bindag::MOZILLA);

    let head_name = VertexName::copy_from(format!("{}", parents.len() - 1).as_bytes());
    let parents_by_name: ParentsFunc =
        Box::new(|name: VertexName| -> dag::Result<Vec<VertexName>> {
            let i = String::from_utf8(name.as_ref().to_vec())
                .unwrap()
                .parse::<usize>()
                .unwrap();
            Ok(parents[i]
                .iter()
                .map(|p| format!("{}", p).as_bytes().to_vec().into())
                .collect())
        });

    let id_map_dir = tempdir().unwrap();
    let mut id_map = IdMap::open(id_map_dir.path()).unwrap();
    let mut covered_ids = IdSet::empty();
    let reserved_ids = IdSet::empty();
    let outcome = nbr(id_map.assign_head(
        head_name.clone(),
        &parents_by_name,
        Group::MASTER,
        &mut covered_ids,
        &reserved_ids,
    ))
    .unwrap();

    bench("building segments", || {
        let mut dag = get_empty_iddag();
        elapsed(|| {
            dag.build_segments_from_prepared_flat_segments(&outcome)
                .unwrap();
        })
    });

    // Write segments to filesystem.
    let mut dag = get_empty_iddag();
    {
        dag.build_segments_from_prepared_flat_segments(&outcome)
            .unwrap();
        let locked = dag.lock().unwrap();
        dag.persist(&locked).unwrap();
    }

    let sample_two_ids: Vec<IdSet> = (0..parents.len() as u64)
        .step_by(10079)
        .flat_map(|i| {
            (1..parents.len() as u64)
                .step_by(7919)
                .map(move |j| (Id(i), Id(j)).into())
        })
        .collect(); // 2679 samples
    let sample_one_ids: Vec<IdSet> = (0..parents.len() as u64)
        .step_by(153)
        .map(|i| IdSet::from(Id(i)))
        .collect();
    let sample_sets: Vec<IdSet> = (0..parents.len() as u64)
        .step_by(10079)
        .flat_map(|i| {
            ((i + 7919)..parents.len() as u64)
                .step_by(7919)
                .map(move |j| (Id(i)..=Id(j)).into())
        })
        .collect(); // 1471 samples

    bench("ancestors", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.ancestors(set.clone()).unwrap();
            }
        })
    });

    bench("children (spans)", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.children(set.clone()).unwrap();
            }
        })
    });

    bench("children (1 id)", || {
        elapsed(|| {
            for set in &sample_one_ids {
                dag.children(set.clone()).unwrap();
            }
        })
    });

    bench("common_ancestors (spans)", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.common_ancestors(set.clone()).unwrap();
            }
        })
    });

    bench("descendants (small subset)", || {
        elapsed(|| {
            // "descendants" is extremely slow. Therefore only test a very
            // small subset.
            for set in sample_sets.iter().skip(500).take(2) {
                dag.descendants(set.clone()).unwrap();
            }
        })
    });

    bench("gca_one (2 ids)", || {
        elapsed(|| {
            for set in &sample_two_ids {
                dag.gca_one(set.clone()).unwrap();
            }
        })
    });

    bench("gca_one (spans)", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.gca_one(set.clone()).unwrap();
            }
        })
    });

    bench("gca_all (2 ids)", || {
        elapsed(|| {
            for set in &sample_two_ids {
                dag.gca_all(set.clone()).unwrap();
            }
        })
    });

    bench("gca_all (spans)", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.gca_all(set.clone()).unwrap();
            }
        })
    });

    bench("heads", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.heads(set.clone()).unwrap();
            }
        })
    });

    bench("heads_ancestors", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.heads_ancestors(set.clone()).unwrap();
            }
        })
    });

    bench("is_ancestor", || {
        elapsed(|| {
            for set in &sample_two_ids {
                let ids: Vec<_> = set.iter_desc().collect();
                dag.is_ancestor(ids[0], ids[1]).unwrap();
            }
        })
    });

    bench("parents", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.parents(set.clone()).unwrap();
            }
        })
    });

    bench("parent_ids", || {
        elapsed(|| {
            for set in &sample_two_ids {
                for id in set.iter_desc() {
                    dag.parent_ids(id).unwrap();
                }
            }
        })
    });

    bench("range (2 ids)", || {
        elapsed(|| {
            for set in &sample_two_ids {
                let ids: Vec<_> = set.iter_desc().collect();
                dag.range(ids[0].into(), ids[1].into()).unwrap();
            }
        })
    });

    bench("range (spans)", || {
        elapsed(|| {
            let mut iter = sample_sets.iter();
            if let (Some(set1), Some(set2)) = (iter.next(), iter.next_back()) {
                dag.range(set1.clone(), set2.clone()).unwrap();
            }
        })
    });

    bench("roots", || {
        elapsed(|| {
            for set in &sample_sets {
                dag.roots(set.clone()).unwrap();
            }
        })
    });
}

fn bench_many_heads_namedag() {
    println!("benchmarking NameDag with many heads");
    // Create a graph with M linear vertexes in the master branch, and M
    // child for every vertex in the master branch.
    //
    // VertexName are just strings of Ids (0, 1, ..., N0, N1, ...).
    const M: usize = 8192;
    let parent_func: ParentsFunc = Box::new(|v: VertexName| -> dag::Result<Vec<VertexName>> {
        let is_non_master = v.as_ref().starts_with(b"N");
        let idx: usize = if is_non_master {
            std::str::from_utf8(&v.as_ref()[1..])
                .unwrap()
                .parse()
                .unwrap()
        } else {
            std::str::from_utf8(v.as_ref()).unwrap().parse().unwrap()
        };
        let parents = if is_non_master {
            vec![VertexName::copy_from(&v.as_ref()[1..])]
        } else if idx > 0 {
            vec![VertexName::copy_from(format!("{}", idx - 1).as_bytes())]
        } else {
            vec![]
        };
        Ok(parents)
    });
    let non_master_heads: Vec<VertexName> = (0..M)
        .map(|i| VertexName::copy_from(format!("N{}", i).as_bytes()))
        .collect::<Vec<_>>();
    let master_heads: Vec<VertexName> =
        vec![VertexName::copy_from(format!("{}", M - 1).as_bytes())];
    let heads = VertexListWithOptions::from(master_heads)
        .with_highest_group(Group::MASTER)
        .chain(non_master_heads);
    let dag_dir = tempdir().unwrap();
    let mut dag = NameDag::open(&dag_dir.path()).unwrap();
    nbr(dag.add_heads_and_flush(&parent_func, &heads)).unwrap();

    let to_set = |v: &str| -> Set {
        nbr(dag.sort(&Set::from_static_names(vec![VertexName::copy_from(
            v.as_bytes(),
        )])))
        .unwrap()
    };
    let head_root_pairs: Vec<(Set, Set)> = (0..M)
        .map(|i| {
            let head = to_set(&format!("N{}", i));
            let root = to_set(&format!("{}", i));
            (head, root)
        })
        .collect();
    bench("range (master::draft)", || {
        elapsed(|| {
            for (head, root) in &head_root_pairs {
                nbr(dag.range(root.clone(), head.clone())).unwrap();
            }
        })
    });

    let heads = nbr(dag.heads(nbr(dag.all()).unwrap())).unwrap();
    let root_list: Vec<Set> = ((M - 64)..M).map(|i| to_set(&format!("N{}", i))).collect();
    bench("range (recent_draft::drafts)", || {
        elapsed(|| {
            for root in &root_list {
                nbr(dag.range(root.clone(), heads.clone())).unwrap();
            }
        })
    });
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use dag::idmap::IdMap;
use dag::idmap::IdMapAssignHead;
use dag::Group;
use dag::IdDag;
use dag::IdSet;
use dag::InProcessIdDag;
use dag::VertexName;
use minibench::bench;
use minibench::elapsed;
use nonblocking::non_blocking_result as nbr;
use tempfile::tempdir;

type ParentsFunc<'a> = Box<dyn Fn(VertexName) -> dag::Result<Vec<VertexName>> + Send + Sync + 'a>;

pub fn main() {
    println!(
        "benchmarking {} serde",
        std::any::type_name::<InProcessIdDag>()
    );
    let parents = bindag::parse_bindag(bindag::MOZILLA);

    let head_name = VertexName::copy_from(format!("{}", parents.len() - 1).as_bytes());
    let parents_by_name: ParentsFunc =
        Box::new(|name: VertexName| -> dag::Result<Vec<VertexName>> {
            let i = String::from_utf8(name.as_ref().to_vec())
                .unwrap()
                .parse::<usize>()
                .unwrap();
            Ok(parents[i]
                .iter()
                .map(|p| format!("{}", p).as_bytes().to_vec().into())
                .collect())
        });

    let id_map_dir = tempdir().unwrap();
    let mut id_map = IdMap::open(id_map_dir.path()).unwrap();
    let mut covered_ids = IdSet::empty();
    let reserved_ids = IdSet::empty();
    let outcome = nbr(id_map.assign_head(
        head_name.clone(),
        &parents_by_name,
        Group::MASTER,
        &mut covered_ids,
        &reserved_ids,
    ))
    .unwrap();
    let mut iddag = IdDag::new_in_process();
    iddag
        .build_segments_from_prepared_flat_segments(&outcome)
        .unwrap();

    let mut blob = Vec::new();
    bench("serializing inprocess iddag with mincode", || {
        elapsed(|| {
            blob = mincode::serialize(&iddag).unwrap();
        })
    });

    println!("mincode serialized blob has {} bytes", blob.len());

    bench("deserializing inprocess iddag with mincode", || {
        elapsed(|| {
            let _new_iddag: InProcessIdDag = mincode::deserialize(&blob).unwrap();
        })
    });
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

mod dag_ops;
mod inprocess_iddag_serde;
mod segment_sizes;
mod spanset;

fn main() {
    dag_ops::main();
    inprocess_iddag_serde::main();
    spanset::main();
    segment_sizes::main();
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use dag::idmap::IdMap;
use dag::idmap::IdMapAssignHead;
use dag::ops::Persist;
use dag::Group;
use dag::Id;
use dag::IdDag;
use dag::IdSet;
use dag::VertexName;
use minibench::bench;
use minibench::elapsed;
use nonblocking::non_blocking_result as nbr;
use tempfile::tempdir;

type ParentsFunc<'a> = Box<dyn Fn(VertexName) -> dag::Result<Vec<VertexName>> + Send + Sync + 'a>;

pub fn main() {
    let parents = bindag::parse_bindag(bindag::MOZILLA);

    let head_name = VertexName::copy_from(format!("{}", parents.len() - 1).as_bytes());
    let parents_by_name = |name: VertexName| -> dag::Result<Vec<VertexName>> {
        let i = String::from_utf8(name.as_ref().to_vec())
            .unwrap()
            .parse::<usize>()
            .unwrap();
        Ok(parents[i]
            .iter()
            .map(|p| format!("{}", p).as_bytes().to_vec().into())
            .collect())
    };

    let id_map_dir = tempdir().unwrap();
    let mut id_map = IdMap::open(id_map_dir.path()).unwrap();
    let mut covered_ids = IdSet::empty();
    let reserved_ids = IdSet::empty();
    let prepared_segments = nbr(id_map.assign_head(
        head_name,
        &(Box::new(parents_by_name) as ParentsFunc),
        Group::MASTER,
        &mut covered_ids,
        &reserved_ids,
    ))
    .unwrap();

    // Test the size, and generation speed, and ancestor calcuation speed
    // with some different segment sizes.
    for segment_size in [4, 8, 10, 12, 14, 16, 18, 20, 22, 24, 32, 64, 128] {
        let dag_dir = tempdir().unwrap();
        let mut dag = IdDag::open(&dag_dir.path()).unwrap();
        dag.set_new_segment_size(segment_size);
        let segment_len = dag
            .build_segments_from_prepared_flat_segments(&prepared_segments)
            .unwrap();
        {
            let locked = dag.lock().unwrap();
            dag.persist(&locked).unwrap();
        }

        let log_len = dag_dir.path().join("log").metadata().unwrap().len();
        eprintln!("segments: {}  log len: {}", segment_len, log_len);

        bench(
            format!("ancestor calcuation segment_size={}", segment_size),
            || {
                let dag = IdDag::open(&dag_dir.path()).unwrap();
                elapsed(|| {
                    for i in (0..parents.len() as u64).step_by(10079) {
                        for j in (1..parents.len() as u64).step_by(2351) {
                            dag.gca_one((Id(i), Id(j)).into()).unwrap();
                        }
                    }
                })
            },
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

/// GCA based on linear scan.
///
/// Ported from Mercurial's C code `find_gca_candidates()`.
///
/// The algorithm was written by Bryan O'Sullivan on 2013-04-16.
/// He provided both a Python [1] and a C version [2]. Since then, the only real
/// logic change is removing an unnecessary "if" branch by Mads Kiilerich on
/// 2014-02-24 [4].
///
/// The C implementation is quite competitive among linear algorithms on
/// performance. It is cache-efficient, has fast paths to exit early, and
/// takes up to 62 (if bitmask is u64) revs at once. Other implementations
/// might just take 2 revs at most. For example, the older implemenation
/// by Olivia Mackall in 2006 [3] takes 2 revs explicitly.
///
/// Changes in this Rust implementation:
/// - Change `bitmask` from `u64` to `u8` for smaller memory footage, at the
///   cost of losing support for more than 6 revs.
/// - Support octopus merge.
///
/// [1]: https://www.mercurial-scm.org/repo/hg/rev/2f7186400a07
/// [2]: https://www.mercurial-scm.org/repo/hg/rev/5bae936764bb
/// [3]: https://www.mercurial-scm.org/repo/hg/rev/b1db258e875c
/// [4]: https://www.mercurial-scm.org/repo/hg/rev/4add43865a9b
pub fn gca(dag: &[impl AsRef<[usize]>], revs: &[usize]) -> Vec<usize> {
    type BitMask = u8;
    let revcount = revs.len();
    assert!(revcount < 7);
    if revcount == 0 {
        return Vec::new();
    }

    let allseen: BitMask = (1 << revcount) - 1;
    let poison: BitMask = 1 << revcount;
    let maxrev = revs.iter().max().cloned().unwrap();
    let mut interesting = revcount;
    let mut gca = Vec::new();
    let mut seen: Vec<BitMask> = vec![0; maxrev + 1];

    for (i, &rev) in revs.iter().enumerate() {
        seen[rev] = 1 << i;
    }

    for v in (0..=maxrev).rev() {
        if interesting == 0 {
            break;
        }
        let mut sv = seen[v];
        if sv == 0 {
            continue;
        }
        if sv < poison {
            interesting -= 1;
            if sv == allseen {
                gca.push(v);
                sv |= poison;
                if revs.iter().any(|&r| r == v) {
                    break;
                }
            }
        }
        for &p in dag[v].as_ref() {
            let sp = seen[p];
            if sv < poison {
                if sp == 0 {
                    seen[p] = sv;
                    interesting += 1
                } else if sp != sv {
                    seen[p] |= sv
                }
            } else {
                if sp != 0 && sp < poison {
                    interesting -= 1
                }
                seen[p] = sv
            }
        }
    }

    gca
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::ops::Deref;
use std::ops::Range;

use vlqencoding::VLQDecode;

mod gca;
pub mod octopus;
mod range;
mod test_context;

pub use gca::gca;
pub use range::range;
pub use test_context::GeneralTestContext;
pub use test_context::OctopusTestContext;
pub use test_context::TestContext;

pub static MOZILLA: &[u8] = include_bytes!("mozilla-central.bindag");
pub static GIT: &[u8] = include_bytes!("git.bindag");

/// "smallvec" optimization
#[derive(Clone, Copy)]
pub struct ParentRevs([usize; 2]);

impl ParentRevs {
    const NONE: usize = usize::max_value();
}

impl From<Vec<usize>> for ParentRevs {
    fn from(revs: Vec<usize>) -> Self {
        assert!(revs.len() <= 2);
        match revs.len() {
            0 => Self([Self::NONE, Self::NONE]),
            1 => Self([revs[0], Self::NONE]),
            2 => Self([revs[0], revs[1]]),
            n => panic!("unsupported len: {}", n),
        }
    }
}

impl AsRef<[usize]> for ParentRevs {
    fn as_ref(&self) -> &[usize] {
        if self.0[0] == Self::NONE {
            &self.0[0..0]
        } else if self.0[1] == Self::NONE {
            &self.0[0..1]
        } else {
            &self.0[..]
        }
    }
}

impl Deref for ParentRevs {
    type Target = [usize];

    fn deref(&self) -> &[usize] {
        self.as_ref()
    }
}

pub fn parse_bindag(bindag: &[u8]) -> Vec<ParentRevs> {
    let mut parents = Vec::new();
    let mut cur = std::io::Cursor::new(bindag);
    let mut read_next = move || -> Result<usize, _> { cur.read_vlq() };

    while let Ok(i) = read_next() {
        let next_id = parents.len();
        match i {
            0 => {
                // no parents
                parents.push(vec![].into());
            }
            1 => {
                // 1 specified parent
                let p1 = next_id - read_next().unwrap() - 1;
                parents.push(vec![p1].into());
            }
            2 => {
                // 2 specified parents
                let p1 = next_id - read_next().unwrap() - 1;
                let p2 = next_id - read_next().unwrap() - 1;
                parents.push(vec![p1, p2].into());
            }
            3 => {
                // 2 parents, p2 specified
                let p1 = next_id - 1;
                let p2 = next_id - read_next().unwrap() - 1;
                parents.push(vec![p1, p2].into());
            }
            4 => {
                // 2 parents, p1 specified
                let p1 = next_id - read_next().unwrap() - 1;
                let p2 = next_id - 1;
                parents.push(vec![p1, p2].into());
            }
            _ => {
                // n commits
                for _ in 0..(i - 4) {
                    let p1 = parents.len() - 1;
                    parents.push(vec![p1].into());
                }
            }
        }
    }

    parents
}

/// Slice a graph. Remove unrefered edges.
pub fn slice_parents(parents: Vec<ParentRevs>, range: Range<usize>) -> Vec<ParentRevs> {
    let start: usize = range.start;
    let end: usize = range.end;
    if start == 0 && end >= parents.len() {
        return parents;
    }

    let mut result = Vec::with_capacity(end - start);
    for i in &parents[range] {
        let new_parents: Vec<usize> = i
            .as_ref()
            .iter()
            .filter_map(|&p| {
                if p < start || p >= end {
                    None
                } else {
                    Some(p - start)
                }
            })
            .collect();
        result.push(new_parents.into())
    }
    result
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::VecDeque;

/// Range based on linear scan.
///
/// Ported from Mercurial's C code `reachableroots2()`.
///
/// The C implementation was added by Laurent Charignon on 2015-08-06 [1].
/// It was based on the a Python implementation added by Bryan O'Sullivan on
/// 2012-06-01 [2], which is similar to an older implementation by Eric Hopper
/// on 2005-10-07 [3], but faster and shorter.
///
/// The C code was then revised by others. The most significant change was
/// switching the "contains" check of "roots" and "reachable" from Python sets
/// to bits in the pure C "revstates" array for easier error handling and
/// better performance, by Yuya Nishihara on 2015-08-13 [4] [5].
///
/// Improvements in this Rust implementation:
/// - Use `VecDeque` for `tovisit` (roughly O(len(result)) -> O(len(heads))).
/// - Truncate `revstates` (O(len(changelog)) -> O(max_head - min_root)).
/// - Add `reachable.is_empty()` fast path that existed in the Python code.
/// - Support octopus merge.
///
/// [1]: https://www.mercurial-scm.org/repo/hg/rev/ff89383a97db
/// [2]: https://www.mercurial-scm.org/repo/hg/rev/b6efeb27e733
/// [3]: https://www.mercurial-scm.org/repo/hg/rev/518da3c3b6ce
/// [4]: https://www.mercurial-scm.org/repo/hg/rev/b68c9d232db6
/// [5]: https://www.mercurial-scm.org/repo/hg/rev/b3ad349d0e50
pub fn range(dag: &[impl AsRef<[usize]>], roots: &[usize], heads: &[usize]) -> Vec<usize> {
    if roots.is_empty() || heads.is_empty() {
        return Vec::new();
    }
    let min_root = *roots.iter().min().unwrap();
    let max_head = *heads.iter().max().unwrap();
    let len = max_head.max(min_root) - min_root + 1;
    let mut reachable = Vec::with_capacity(len);
    let mut tovisit = VecDeque::new();
    let mut revstates = vec![0u8; len];

    const RS_SEEN: u8 = 1;
    const RS_ROOT: u8 = 2;
    const RS_REACHABLE: u8 = 4;

    for &rev in roots {
        if rev <= max_head {
            revstates[rev - min_root] |= RS_ROOT;
        }
    }

    for &rev in heads {
        if rev >= min_root && revstates[rev - min_root] & RS_SEEN == 0 {
            tovisit.push_back(rev);
            revstates[rev - min_root] |= RS_SEEN;
        }
    }

    // Visit the tovisit list and find the reachable roots
    while let Some(rev) = tovisit.pop_front() {
        // Add the node to reachable if it is a root
        if revstates[rev - min_root] & RS_ROOT != 0 {
            revstates[rev - min_root] |= RS_REACHABLE;
            reachable.push(rev);
        }

        // Add its parents to the list of nodes to visit
        for &p in dag[rev].as_ref() {
            if p >= min_root && revstates[p - min_root] & RS_SEEN == 0 {
                tovisit.push_back(p);
                revstates[p - min_root] |= RS_SEEN;
            }
        }
    }

    if reachable.is_empty() {
        return Vec::new();
    }

    // Find all the nodes in between the roots we found and the heads
    // and add them to the reachable set
    for rev in min_root..=max_head {
        if revstates[rev - min_root] & RS_SEEN == 0 {
            continue;
        }
        if dag[rev]
            .as_ref()
            .iter()
            .any(|&p| p >= min_root && revstates[p - min_root] & RS_REACHABLE != 0)
            && revstates[rev - min_root] & RS_REACHABLE == 0
        {
            revstates[rev - min_root] |= RS_REACHABLE;
            reachable.push(rev);
        }
    }

    reachable
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::ops::Range;

use dag::nameset::SyncNameSetQuery;
use dag::ops::DagAlgorithm;
use dag::ops::DagPersistent;
use dag::ops::IdConvert;
use dag::Group;
use dag::Id;
use dag::IdSet;
use dag::NameDag;
use dag::OnDiskIdDag;
use dag::VertexListWithOptions;
use dag::VertexName;
use nonblocking::non_blocking_result;
use tempfile::TempDir;

use crate::parse_bindag;
use crate::ParentRevs;

/// Context for testing purpose.
/// Contains the parsed bindag and NameDag from the dag crate.
pub struct GeneralTestContext<T> {
    /// Plain DAG parsed from bindag
    pub parents: Vec<T>,

    /// Complex DAG, with ids re-assigned
    pub dag: NameDag,

    /// Simple IdMap. NameDag Id -> Plain DAG id
    pub idmap: HashMap<Id, usize>,

    /// Plain DAG id -> NameDag Id.
    pub rev_idmap: Vec<Id>,

    /// Temporary dir containing the NameDag daga.
    pub dir: TempDir,
}

pub type TestContext = GeneralTestContext<ParentRevs>;
pub type OctopusTestContext = GeneralTestContext<Vec<usize>>;

impl TestContext {
    pub fn from_bin(bin: &[u8]) -> Self {
        Self::from_bin_sliced(bin, 0..usize::max_value())
    }

    pub fn from_bin_sliced(bin: &[u8], range: Range<usize>) -> Self {
        // Prepare the plain DAG (parents)
        let parents = parse_bindag(bin);
        let parents = crate::slice_parents(parents, range);
        Self::from_parents(parents)
    }
}

impl<T: AsRef<[usize]> + Send + Sync + Clone + 'static> GeneralTestContext<T> {
    pub fn from_parents(parents: Vec<T>) -> Self {
        // Prepare NameDag
        let parents_by_name = {
            let parents = parents.clone();
            move |name: VertexName| -> dag::Result<Vec<VertexName>> {
                let i = String::from_utf8(name.as_ref().to_vec())
                    .unwrap()
                    .parse::<usize>()
                    .unwrap();
                Ok(parents[i]
                    .as_ref()
                    .iter()
                    .map(|p| format!("{}", p).as_bytes().to_vec().into())
                    .collect())
            }
        };
        // Pick heads from 0..n
        let get_heads = |n: usize| -> Vec<VertexName> {
            let mut heads: HashSet<usize> = (0..n).collect();
            for ps in parents.iter().take(n) {
                for p in ps.as_ref().iter() {
                    heads.remove(&p);
                }
            }
            let mut names: Vec<VertexName> = Vec::new();
            for h in heads {
                names.push(VertexName::copy_from(format!("{}", h).as_bytes()));
            }
            names
        };
        let head_names: Vec<VertexName> = get_heads(parents.len());
        let master_names: Vec<VertexName> = get_heads(parents.len() / 2);

        let dir = tempfile::tempdir().unwrap();
        let mut dag = NameDag::open(dir.path()).unwrap();
        let parents_map: Box<dyn Fn(VertexName) -> dag::Result<Vec<VertexName>> + Send + Sync> =
            Box::new(parents_by_name);
        let heads = VertexListWithOptions::from(master_names)
            .with_highest_group(Group::MASTER)
            .chain(head_names);
        non_blocking_result(dag.add_heads_and_flush(&parents_map, &heads)).unwrap();

        // Prepare idmap
        let idmap: HashMap<Id, usize> = non_blocking_result(dag.all())
            .unwrap()
            .iter()
            .unwrap()
            .map(|name| {
                let name = name.unwrap();
                let dag_id: Id = non_blocking_result(dag.vertex_id(name.clone())).unwrap();
                let plain_id: usize = std::str::from_utf8(name.as_ref()).unwrap().parse().unwrap();
                (dag_id, plain_id)
            })
            .collect();
        let mut rev_idmap = vec![Id(0); idmap.len()];
        for (&k, &v) in idmap.iter() {
            rev_idmap[v] = k;
        }
        assert_eq!(rev_idmap.len(), parents.len());

        Self {
            parents,
            dag,
            idmap,
            rev_idmap,
            dir,
        }
    }
}

impl<T> GeneralTestContext<T> {
    /// Limit the size of `parents`.
    pub fn truncate(mut self, size: usize) -> Self {
        if size < self.parents.len() {
            self.parents.truncate(size);
        }
        self
    }

    /// Convert a IdSet (used by IdDag) to plain revs (used by `parents`).
    pub fn to_plain_revs(&self, set: &IdSet) -> Vec<usize> {
        set.iter_desc().map(|i| self.idmap[&i]).collect()
    }

    /// Get the IdDag reference.
    pub fn id_dag(&self) -> &OnDiskIdDag {
        self.dag.dag()
    }

    /// Make `revs` in the range of parents.
    pub fn clamp_revs(&self, revs: &[impl Into<usize> + Clone]) -> Vec<usize> {
        revs.iter()
            .cloned()
            .map(|i| Into::<usize>::into(i) % self.parents.len())
            .collect()
    }

    /// Convert `usize` plain revs to IdSet used by IdDag.
    pub fn to_dag_revs(&self, revs: &[usize]) -> IdSet {
        let iter = revs.iter().map(|&i| self.rev_idmap[i].clone());
        IdSet::from_spans(iter)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeMap;

use serde::Deserialize;
use serde::Serialize;

use crate::id::Id;
use crate::segment::PreparedFlatSegments;

#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Serialize, Deserialize)]
pub struct CloneData<Name> {
    pub flat_segments: PreparedFlatSegments,
    pub idmap: BTreeMap<Id, Name>,
}

impl<Name> CloneData<Name> {
    pub fn convert_vertex<T, F: Fn(Name) -> T>(self, f: F) -> CloneData<T> {
        let idmap = self.idmap.into_iter().map(|(k, v)| (k, f(v))).collect();
        CloneData {
            flat_segments: self.flat_segments,
            idmap,
        }
    }
}

#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Arbitrary;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Gen;

#[cfg(any(test, feature = "for-tests"))]
impl<Name> Arbitrary for CloneData<Name>
where
    Name: Arbitrary,
{
    fn arbitrary(g: &mut Gen) -> Self {
        let flat_segments = PreparedFlatSegments {
            segments: Arbitrary::arbitrary(g),
        };
        CloneData {
            flat_segments,
            idmap: Arbitrary::arbitrary(g),
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! # id
//!
//! Defines types around [`Id`].

use std::fmt;
use std::io;
use std::ops;

#[cfg(feature = "serialize-abomonation")]
use abomonation_derive::Abomonation;
pub use minibytes::Bytes;
use serde::Deserialize;
use serde::Serialize;

/// An integer [`Id`] representing a node in the graph.
/// [`Id`]s are topologically sorted.
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(Serialize, Deserialize)]
#[cfg_attr(feature = "serialize-abomonation", derive(Abomonation))]
pub struct Id(pub u64);

/// Name of a vertex in the graph.
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct VertexName(pub Bytes);

impl AsRef<[u8]> for VertexName {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl VertexName {
    pub fn to_hex(&self) -> String {
        const HEX_CHARS: &[u8] = b"0123456789abcdef";
        let mut v = Vec::with_capacity(self.0.len() * 2);
        for &byte in self.as_ref() {
            v.push(HEX_CHARS[(byte >> 4) as usize]);
            v.push(HEX_CHARS[(byte & 0xf) as usize]);
        }
        unsafe { String::from_utf8_unchecked(v) }
    }

    /// Convert from hex.
    ///
    /// If `len(hex)` is an odd number, hex + '0' will be used.
    pub fn from_hex(hex: &[u8]) -> io::Result<Self> {
        let mut bytes = vec![0u8; (hex.len() + 1) / 2];
        for (i, byte) in hex.iter().enumerate() {
            let value = match byte {
                b'0'..=b'9' => byte - b'0',
                b'a'..=b'f' => byte - b'a' + 10,
                b'A'..=b'F' => byte - b'A' + 10,
                _ => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        format!("{:?} is not a hex character", *byte as char),
                    ));
                }
            };
            if i & 1 == 0 {
                bytes[i / 2] |= value << 4;
            } else {
                bytes[i / 2] |= value;
            }
        }
        Ok(VertexName(Bytes::from(bytes)))
    }

    pub fn copy_from(value: &[u8]) -> Self {
        Self(value.to_vec().into())
    }
}

impl<T> From<T> for VertexName
where
    Bytes: From<T>,
{
    fn from(value: T) -> Self {
        Self(Bytes::from(value))
    }
}

impl fmt::Debug for VertexName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.0.len() >= 4 || !looks_like_ascii_identifier(self.as_ref()) {
            // Use hex format for long names (ex. binary commit hashes).
            let hex = self.to_hex();
            // Truncate to specified width (ex. '{:#.12}').
            if let Some(width) = f.precision() {
                let truncated = hex.get(..width).unwrap_or(&hex);
                f.write_str(truncated)
            } else {
                f.write_str(&hex)
            }
        } else {
            // Do not use hex if it looks like an ASCII identifier.
            match std::str::from_utf8(self.as_ref()) {
                Ok(s) => write!(f, "{}", s),
                Err(_) => write!(f, "{}", self.to_hex()),
            }
        }
    }
}

fn looks_like_ascii_identifier(bytes: &[u8]) -> bool {
    let mut iter = bytes.iter().copied();
    if !(iter.next().unwrap_or(b'\0') as char).is_ascii_alphabetic() {
        return false;
    }
    iter.all(|b| b.is_ascii_alphanumeric())
}

/// An integer that separates distinct groups of [`Id`]s.
///
/// This can be seen as a way to pre-allocate consecutive integers
/// for one group to make segments less fragmented.
///
/// `(Group, Id)` are also topologically sorted.
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(Serialize, Deserialize)]
pub struct Group(pub usize);

impl Group {
    /// The "master" group. `ancestors(master)`.
    /// - Expected to have most of the commits in a repo.
    /// - Expected to be free from fragmentation. In other words,
    ///   `ancestors(master)` can be represented in a single Span.
    pub const MASTER: Self = Self(0);

    /// The "non-master" group.
    /// - Anything not in `ancestors(master)`. For example, public release
    ///   branches, local feature branches.
    /// - Expected to have multiple heads. In other words, is fragmented.
    /// - Expected to be sparse referred. For example, the "visible heads"
    ///   will refer to a bounded subset in this group.
    pub const NON_MASTER: Self = Self(1);

    pub const ALL: [Self; 2] = [Self::MASTER, Self::NON_MASTER];

    pub const COUNT: usize = Self::ALL.len();

    // 1 byte for Group so it's easier to remove everything in a group.
    pub const BITS: u32 = 8;
    pub const BYTES: usize = 1;

    /// The first [`Id`] in this group.
    pub const fn min_id(self) -> Id {
        Id((self.0 as u64) << (64 - Self::BITS))
    }

    /// The maximum [`Id`] in this group.
    pub const fn max_id(self) -> Id {
        Id(self.min_id().0 + ((1u64 << (64 - Self::BITS)) - 1))
    }

    /// Convert to array.
    pub const fn bytes(self) -> [u8; 1] {
        [self.0 as u8]
    }

    /// Convert to hex array.
    pub fn hex_bytes(self) -> [u8; 2] {
        if self.0 < 10 {
            [b'0', b'0' + (self.0 as u8)]
        } else {
            unreachable!()
        }
    }
}

impl Id {
    /// The [`Group`] of an Id.
    pub fn group(self) -> Group {
        let group = (self.0 >> (64 - Group::BITS)) as usize;
        debug_assert!(group < Group::COUNT);
        Group(group)
    }

    /// Similar to `self..=other`.
    pub fn to(self, other: Id) -> IdIter {
        IdIter {
            current: self,
            end: other,
        }
    }

    /// Convert to a byte array. Useful for indexedlog range query.
    pub fn to_bytearray(self) -> [u8; 8] {
        // The field can be used for index range query. So it has to be BE.
        unsafe { std::mem::transmute(self.0.to_be()) }
    }

    /// Similar to `to_bytearray`, but insert a `prefix` at the head.
    /// Useful for segment queries where `level` is the `prefix`.
    pub fn to_prefixed_bytearray(self, prefix: u8) -> [u8; 9] {
        let a = self.to_bytearray();
        [prefix, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]]
    }

    pub const MAX: Self = Group::ALL[Group::COUNT - 1].max_id();
    pub const MIN: Self = Group::ALL[0].min_id();
}

impl fmt::Display for Id {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let group = self.group();
        if group == Group::NON_MASTER {
            write!(f, "N")?;
        }
        write!(f, "{}", self.0 - group.min_id().0)
    }
}

impl fmt::Debug for Id {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for Group {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Group::MASTER => write!(f, "Group Master"),
            Group::NON_MASTER => write!(f, "Group Non-Master"),
            _ => write!(f, "Group {}", self.0),
        }
    }
}

impl ops::Add<u64> for Id {
    type Output = Id;

    fn add(self, other: u64) -> Self {
        Self(self.0 + other)
    }
}

impl ops::Sub<u64> for Id {
    type Output = Id;

    fn sub(self, other: u64) -> Self {
        Self(self.0 - other)
    }
}

// Consider replacing this with iter::Step once it's stable.
pub struct IdIter {
    current: Id,
    end: Id,
}

impl Iterator for IdIter {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.current > self.end {
            None
        } else {
            let result = self.current;
            self.current = self.current + 1;
            Some(result)
        }
    }
}

#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Arbitrary;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Gen;

#[cfg(any(test, feature = "for-tests"))]
impl Arbitrary for Id {
    fn arbitrary(g: &mut Gen) -> Self {
        let group = Group((u32::arbitrary(g) & 1) as usize);
        group.min_id() + u64::arbitrary(g) % (group.max_id().0 - group.min_id().0)
    }
}

// For convenience.
impl std::cmp::PartialEq<u64> for Id {
    fn eq(&self, other: &u64) -> bool {
        self.0 == *other
    }
}

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

    use super::*;

    #[test]
    fn test_vertex_from_hex_odd() {
        let vertex = VertexName::from_hex(b"a").unwrap();
        let vertex2 = VertexName::from_hex(b"a0").unwrap();
        assert_eq!(vertex, vertex2);
        assert_eq!(vertex.to_hex(), "a0");
    }

    quickcheck! {
        fn test_vertex_hex_roundtrip(slice: Vec<u8>) -> bool {
            let vertex = VertexName::from(slice);
            let hex = vertex.to_hex();
            let vertex2 = VertexName::from_hex(hex.as_bytes()).unwrap();
            vertex2 == vertex
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! A subset of types that are related to wire protocols and used by the dag crate.

pub mod clone;
pub mod id;
pub mod location;
pub mod segment;

pub use clone::CloneData;
pub use id::Bytes;
pub use id::Group;
pub use id::Id;
pub use id::IdIter;
pub use id::VertexName;
pub use location::Location;
pub use segment::FlatSegment;
pub use segment::PreparedFlatSegments;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::future::Future;

use serde::Deserialize;
use serde::Serialize;

/// The position of one node in the graph relative their known descendant.
/// This structure is relevant when the IdMap is lazy. Assuming that all parents of merges
/// and all heads are known then any node can be represented as their first `descendant`
/// and the distance to that descendant.
///
/// Example:
/// 0 - a - b - c
/// In this example our initial commit is `0`, then we have `a` the first commit, `b` second,
/// `c` third.
/// {
///   descendant: c,
///   distance: 1,
/// }
/// => [b]
#[derive(Copy, Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[derive(Serialize, Deserialize)]
pub struct Location<Name> {
    pub descendant: Name,
    pub distance: u64,
}

impl<Name> Location<Name> {
    pub fn new(descendant: Name, distance: u64) -> Self {
        Self {
            descendant,
            distance,
        }
    }

    // What's up with all these map-like functions?
    // The most common operation for the location is to transform it between various
    // types of graphs, crate::Id, HgId, HgChangesetId, ChangesetId etc.
    // Had some fun adding with these functions. "Experimenting".

    pub fn map_descendant<T, F>(self, f: F) -> Location<T>
    where
        F: FnOnce(Name) -> T,
    {
        let new_name = f(self.descendant);
        Location::new(new_name, self.distance)
    }

    pub fn try_map_descendant<T, E, F>(self, f: F) -> Result<Location<T>, E>
    where
        F: FnOnce(Name) -> Result<T, E>,
    {
        let new_name = f(self.descendant)?;
        Ok(Location::new(new_name, self.distance))
    }

    pub async fn then_descendant<T, Fut, F>(self, f: F) -> Location<T>
    where
        F: FnOnce(Name) -> Fut,
        Fut: Future<Output = T>,
    {
        let new_name = f(self.descendant).await;
        Location::new(new_name, self.distance)
    }

    pub async fn and_then_descendant<T, E, Fut, F>(self, f: F) -> Result<Location<T>, E>
    where
        F: FnOnce(Name) -> Fut,
        Fut: Future<Output = Result<T, E>>,
    {
        let new_name = f(self.descendant).await?;
        Ok(Location::new(new_name, self.distance))
    }

    pub fn with_descendant<T>(self, descendant: T) -> Location<T> {
        Location::new(descendant, self.distance)
    }
}

#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Arbitrary;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck::Gen;

#[cfg(any(test, feature = "for-tests"))]
impl<Name> Arbitrary for Location<Name>
where
    Name: Arbitrary,
{
    fn arbitrary(g: &mut Gen) -> Self {
        Location {
            descendant: Name::arbitrary(g),
            distance: u64::arbitrary(g),
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use serde::Deserialize;
use serde::Serialize;

use crate::id::Id;

/// Base segment.
///
/// Intermediate structure between processing a Dag and constructing high level segments.
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Serialize, Deserialize, Ord, PartialOrd)]
pub struct FlatSegment {
    pub low: Id,
    pub high: Id,
    pub parents: Vec<Id>,
}

#[cfg(any(test, feature = "for-tests"))]
impl quickcheck::Arbitrary for FlatSegment {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            low: Id::arbitrary(g),
            high: Id::arbitrary(g),
            parents: Vec::arbitrary(g),
        }
    }
}

use std::collections::BTreeSet;

/// These segments can be used directly in the build process of the IdDag.
/// They produced by `IdMap::assign_head` and `IdDag::all_flat_segments`.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
#[derive(Serialize, Deserialize)]
pub struct PreparedFlatSegments {
    /// New flat segments.
    pub segments: BTreeSet<FlatSegment>,
}

impl PreparedFlatSegments {
    pub fn vertex_count(&self) -> u64 {
        let mut count = 0;
        for segment in &self.segments {
            count += segment.high.0 - segment.low.0 + 1;
        }
        count
    }

    pub fn segment_count(&self) -> usize {
        self.segments.len()
    }

    /// Return set of all (unique) parents + head + roots of flat segments.
    ///
    /// Used by the pull fast path to provide necessary "anchor" vertexes
    /// ("universally known", and ones needed by the client to make decisions)
    /// in the IdMap.
    ///
    /// Might return some extra `Id`s that are not part of parents, heads, or
    /// roots. They are useful for the client to verify the graph is the same
    /// as the server.
    ///
    /// The size of the returned `Id`s is about `O(segments)`.
    pub fn parents_head_and_roots(&self) -> BTreeSet<Id> {
        self.segments
            .iter()
            .map(|seg| {
                // `seg.high` is either a head, or a parent referred by another seg
                // `seg.low` is either a room, or something unnecessary for lazy protocol,
                // but speeds up graph shape verification (see `check_isomorphic_graph`).
                // `parents` are either "universally known", essential for lazy protocol,
                // or something necessary for the pull protocol to re-map the IdMap.
                [seg.high, seg.low]
                    .into_iter()
                    .chain(seg.parents.clone().into_iter())
            })
            .flatten()
            .collect()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::TestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    static ref CONTEXT: TestContext =
        TestContext::from_bin(bindag::GIT).truncate(u16::max_value() as usize);
}

fuzz_target!(|input: Vec<u16>| {
    // gca with > 3 revs is less interesting to this test.
    let revs = CONTEXT.clamp_revs(&input[..input.len().min(3)]);
    tests::test_gca(&CONTEXT, revs);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::octopus;
use bindag::OctopusTestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    static ref CONTEXT: OctopusTestContext =
        OctopusTestContext::from_parents(octopus::cross_octopus());
}

fuzz_target!(|input: Vec<u8>| {
    let revs = CONTEXT.clamp_revs(&input[..input.len().min(5)]);
    tests::test_gca(&CONTEXT, revs);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::TestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    static ref CONTEXT: TestContext = {
        let start = 60146;
        TestContext::from_bin_sliced(bindag::GIT, start..start + 256)
    };
}

fuzz_target!(|input: Vec<u8>| {
    // plain gca supports 6 revs at maximum.
    let revs = CONTEXT.clamp_revs(&input[..input.len().min(6)]);
    tests::test_gca(&CONTEXT, revs);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::TestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    // The complete DAG is too large for `range` operation to run reasonably fast.
    // Therefore take a subset of it.
    static ref CONTEXT: TestContext = TestContext::from_bin_sliced(bindag::GIT, 49040..60415);
}

fuzz_target!(|input: (Vec<u16>, Vec<u16>)| {
    let roots = CONTEXT.clamp_revs(&input.0);
    let heads = CONTEXT.clamp_revs(&input.1);
    tests::test_range(&CONTEXT, roots, heads);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::TestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    // Pick a subset of the DAG (with many merges).
    static ref CONTEXT: TestContext = TestContext::from_bin_sliced(bindag::GIT, 56666..60415);
}

fuzz_target!(|input: (Vec<u16>, Vec<u16>)| {
    let roots = CONTEXT.clamp_revs(&input.0);
    let heads = CONTEXT.clamp_revs(&input.1);
    tests::test_range(&CONTEXT, roots, heads);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::octopus;
use bindag::OctopusTestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    static ref CONTEXT: OctopusTestContext =
        OctopusTestContext::from_parents(octopus::cross_octopus());
}

fuzz_target!(|input: (Vec<u8>, Vec<u8>)| {
    let roots = CONTEXT.clamp_revs(&input.0);
    let heads = CONTEXT.clamp_revs(&input.1);
    tests::test_range(&CONTEXT, roots, heads);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![no_main]

use bindag::TestContext;
use lazy_static::lazy_static;
use libfuzzer_sys::fuzz_target;

mod tests;

lazy_static! {
    static ref CONTEXT: TestContext = {
        let start = 60146;
        TestContext::from_bin_sliced(bindag::GIT, start..start + 256)
    };
}

fuzz_target!(|input: (Vec<u8>, Vec<u8>)| {
    let roots = CONTEXT.clamp_revs(&input.0);
    let heads = CONTEXT.clamp_revs(&input.1);
    tests::test_range(&CONTEXT, roots, heads);
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use bindag::GeneralTestContext;

#[allow(dead_code)]
pub fn test_gca<T: AsRef<[usize]>>(context: &GeneralTestContext<T>, mut plain_revs: Vec<usize>) {
    plain_revs.sort_unstable();
    plain_revs.dedup();
    let iddag_revs = context.to_dag_revs(&plain_revs);
    let mut plain_gca = bindag::gca(&context.parents, &plain_revs);
    let mut iddag_gca = context.to_plain_revs(&context.id_dag().gca_all(iddag_revs).unwrap());

    plain_gca.sort_unstable();
    iddag_gca.sort_unstable();
    assert_eq!(plain_gca, iddag_gca, "gca({:?})", &plain_revs);
}

#[allow(dead_code)]
pub fn test_range<T: AsRef<[usize]>>(
    context: &GeneralTestContext<T>,
    plain_roots: Vec<usize>,
    plain_heads: Vec<usize>,
) {
    let plain_roots = context.clamp_revs(&plain_roots);
    let plain_heads = context.clamp_revs(&plain_heads);
    // let plain_roots: Vec<_> = roots.iter().map(|i| context.idmap[i]).collect();
    // let plain_heads: Vec<_> = heads.iter().map(|i| context.idmap[i]).collect();

    let mut plain_range = bindag::range(&context.parents, &plain_roots, &plain_heads);
    plain_range.sort_unstable();

    let iddag_roots = context.to_dag_revs(&plain_roots);
    let iddag_heads = context.to_dag_revs(&plain_heads);
    let mut iddag_range = context.to_plain_revs(
        &context
            .id_dag()
            .range(iddag_roots, iddag_heads)
            .expect("IdDag::range"),
    );

    iddag_range.sort_unstable();
    assert_eq!(
        plain_range, iddag_range,
        "range({:?}::{:?})",
        &plain_roots, &plain_heads
    );
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

pub trait MapDagError<T> {
    fn context(self, message: &'static str) -> dag::Result<T>;
    fn with_context(self, func: impl Fn() -> String) -> dag::Result<T>;
}

impl<T> MapDagError<T> for Result<T, git2::Error> {
    fn context(self, message: &'static str) -> dag::Result<T> {
        anyhow::Context::context(self, message)
            .map_err(|e| dag::errors::BackendError::Other(e).into())
    }
    fn with_context(self, func: impl Fn() -> String) -> dag::Result<T> {
        anyhow::Context::with_context(self, func)
            .map_err(|e| dag::errors::BackendError::Other(e).into())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BTreeMap;
use std::ops::Deref;
use std::ops::DerefMut;
use std::path::Path;

use dag::ops::DagAlgorithm;
use dag::ops::DagPersistent;
use dag::Dag;
use dag::Group;
use dag::Set;
use dag::Vertex;
use dag::VertexListWithOptions;
use nonblocking::non_blocking_result;
use parking_lot::Mutex;

use crate::errors::MapDagError;

/// `GitDag` maintains segmented changelog as an index on the Git commit graph.
///
/// This struct provides a "read-only" view for the commit graph. To read other
/// parts of the git repo, or make changes to the Git commit graph, use a
/// separate `git2::Repository` object.
///
/// The `dag` part is append-only. It might include vertexes no longer referred
/// by the git repo. Use `ancestors(git_heads())` to get commits referred by
/// the git repo, and use `&` to filter them.
pub struct GitDag {
    dag: Dag,
    heads: Set,
    references: BTreeMap<String, Vertex>,
}

impl GitDag {
    /// `open` a Git repo at `git_dir`. Build index at `dag_dir`, with specified `main_branch`.
    pub fn open(git_dir: &Path, dag_dir: &Path, main_branch: &str) -> dag::Result<Self> {
        let git_repo = git2::Repository::open(git_dir)
            .with_context(|| format!("opening git repo at {}", git_dir.display()))?;
        Self::open_git_repo(&git_repo, dag_dir, main_branch)
    }

    /// For an git repo, build index at `dag_dir` with specified `main_branch`.
    pub fn open_git_repo(
        git_repo: &git2::Repository,
        dag_dir: &Path,
        main_branch: &str,
    ) -> dag::Result<Self> {
        let dag = Dag::open(dag_dir)?;
        Ok(sync_from_git(dag, git_repo, main_branch)?)
    }

    /// Get "snapshotted" references.
    pub fn git_references(&self) -> &BTreeMap<String, Vertex> {
        &self.references
    }

    /// Get "snapshotted" heads.
    pub fn git_heads(&self) -> Set {
        self.heads.clone()
    }
}

impl Deref for GitDag {
    type Target = Dag;

    fn deref(&self) -> &Self::Target {
        &self.dag
    }
}

impl DerefMut for GitDag {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.dag
    }
}

/// Read references from git, build segments for new heads.
///
/// Useful when the git repo is changed by other processes or threads.
fn sync_from_git(
    mut dag: Dag,
    git_repo: &git2::Repository,
    main_branch: &str,
) -> dag::Result<GitDag> {
    let mut master_heads = Vec::new();
    let mut non_master_heads = Vec::new();
    let mut references = BTreeMap::new();

    let git_refs = git_repo.references().context("listing git references")?;
    for git_ref in git_refs {
        let git_ref = git_ref.context("resolving git reference")?;
        let commit = match git_ref.peel_to_commit() {
            Err(e) => {
                tracing::warn!(
                    "git ref {} cannot resolve to commit: {}",
                    String::from_utf8_lossy(git_ref.name_bytes()),
                    e
                );
                // Ignore this error. Some git references (ex. tags) can point
                // to trees instead of commits.
                continue;
            }
            Ok(c) => c,
        };
        let oid = commit.id();
        let vertex = Vertex::copy_from(oid.as_bytes());
        if let Some(name) = git_ref.name() {
            references.insert(name.to_string(), vertex.clone());
        }
        if git_ref.name() == Some(main_branch) {
            master_heads.push(vertex);
        } else {
            non_master_heads.push(vertex);
        }
    }

    struct ForceSend<T>(T);

    // See https://github.com/rust-lang/git2-rs/issues/194, libgit2 can be
    // accessed by a different thread.
    unsafe impl<T> Send for ForceSend<T> {}

    let git_repo = ForceSend(git_repo);
    let git_repo = Mutex::new(git_repo);

    let parent_func = move |v: Vertex| -> dag::Result<Vec<Vertex>> {
        tracing::trace!("visiting git commit {:?}", &v);
        let oid = git2::Oid::from_bytes(v.as_ref())
            .with_context(|| format!("converting to git oid for {:?}", &v))?;
        let commit = git_repo
            .lock()
            .0
            .find_commit(oid)
            .with_context(|| format!("resolving {:?} to git commit", &v))?;
        Ok(commit
            .parent_ids()
            .map(|id| Vertex::copy_from(id.as_bytes()))
            .collect())
    };
    let parents: Box<dyn Fn(Vertex) -> dag::Result<Vec<Vertex>> + Send + Sync> =
        Box::new(parent_func);
    let heads = VertexListWithOptions::from(master_heads.clone())
        .with_highest_group(Group::MASTER)
        .chain(non_master_heads.clone());
    non_blocking_result(dag.add_heads_and_flush(&parents, &heads))?;

    let possible_heads =
        Set::from_static_names(master_heads.into_iter().chain(non_master_heads.into_iter()));
    let heads = non_blocking_result(dag.heads_ancestors(possible_heads))?;

    Ok(GitDag {
        dag,
        heads,
        references,
    })
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

pub mod errors;
mod gitdag;

pub use git2;

pub use self::gitdag::GitDag;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::cmp::Ordering;
use std::collections::VecDeque;

use Ordering::Equal;
use Ordering::Greater;
use Ordering::Less;

// Based on rust stdlib slice::binary_search_by:
// https://github.com/rust-lang/rust/blob/dbad8c93680710e80c54cbaf8416821f7a5750c8/library/core/src/slice/mod.rs#L1721

pub trait BinarySearchBy<T> {
    /// Binary searches this sorted slice with a comparator function.
    ///
    /// The comparator function should implement an order consistent
    /// with the sort order of the underlying slice, returning an
    /// order code that indicates whether its argument is `Less`,
    /// `Equal` or `Greater` the desired target.
    ///
    /// If the value is found then [`Result::Ok`] is returned, containing the
    /// index of the matching element. If there are multiple matches, then any
    /// one of the matches could be returned. If the value is not found then
    /// [`Result::Err`] is returned, containing the index where a matching
    /// element could be inserted while maintaining sorted order.
    fn bsearch_by<F>(&self, f: F) -> Result<usize, usize>
    where
        F: FnMut(&T) -> Ordering;
}

impl<T> BinarySearchBy<T> for VecDeque<T> {
    fn bsearch_by<F>(&self, mut f: F) -> Result<usize, usize>
    where
        F: FnMut(&T) -> Ordering,
    {
        let s = self;
        let mut size = s.len();
        if size == 0 {
            return Err(0);
        }
        let mut base = 0usize;
        while size > 1 {
            let half = size / 2;
            let mid = base + half;
            // mid is always in [0, size), that means mid is >= 0 and < size.
            // mid >= 0: by definition
            // mid < size: mid = size / 2 + size / 4 + size / 8 ...
            let cmp = f(s.get(mid).unwrap());
            base = if cmp == Greater { base } else { mid };
            size -= half;
        }
        // base is always in [0, size) because base <= mid.
        let cmp = f(s.get(base).unwrap());
        if cmp == Equal {
            Ok(base)
        } else {
            Err(base + (cmp == Less) as usize)
        }
    }
}