/*
 * 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.
 */

mod procfs;
pub use self::procfs::*;

/// Abstracted measured result for benchmark.
pub trait Measure: Sized {
    /// Expected return value from the function being measured.
    /// Usually just `()`.
    type FuncOutput;

    /// Measure a function.
    fn measure(func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String>;

    /// Merge with another measure result.
    fn merge(self, _rhs: Self) -> Self;

    /// Return true if another measurement is needed.
    fn need_more(&self) -> bool;

    /// Convert to human-readable result.
    fn to_string(&self) -> String;
}

/// Measure the best wall clock time quickly.
pub struct WallClock {
    pub best: f64,
    pub total: f64,
    pub count: usize,
}

/// Measure bytes. Run only once.
pub struct Bytes(pub u64);

/// Measure both.
///
/// A will be measured in an inner loop. B will be measured in the outer loop.
/// `A::FuncOutput` must be `()`.
pub struct Both<A, B>(Result<A, String>, Result<B, String>);

impl Measure for WallClock {
    type FuncOutput = ();

    fn measure(mut func: impl FnMut()) -> Result<Self, String> {
        use std::time::SystemTime;
        let now = SystemTime::now();
        func();
        let elapsed = now.elapsed().unwrap();
        let seconds = elapsed.as_secs() as f64 + elapsed.subsec_nanos() as f64 * 1e-9;
        Ok(Self {
            best: seconds,
            total: seconds,
            count: 1,
        })
    }

    fn merge(self, rhs: Self) -> Self {
        Self {
            best: self.best.min(rhs.best),
            total: self.total + rhs.total,
            count: self.count + rhs.count,
        }
    }

    fn need_more(&self) -> bool {
        // Run for 5 seconds, or 40 times at most.
        self.count < 40 && self.total < 5.0
    }

    fn to_string(&self) -> String {
        let value = self.best;
        if value <= 1.0 {
            format!("{:7.3} ms", value * 1000.0)
        } else {
            format!("{:7.3} s ", value)
        }
    }
}

impl Measure for Bytes {
    type FuncOutput = u64;

    fn measure(mut func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String> {
        Ok(Self(func()))
    }

    fn merge(self, _rhs: Self) -> Self {
        Self(self.0)
    }

    fn need_more(&self) -> bool {
        false
    }

    fn to_string(&self) -> String {
        let value = self.0;
        if value < 1_000 {
            format!("{:7} B ", value)
        } else if value < 1_000_000 {
            format!("{:7.3} KB", (value as f64) / 1000.0)
        } else {
            format!("{:7.3} MB", (value as f64) / 1000_000.0)
        }
    }
}

impl<D: Default, A: Measure<FuncOutput = ()>, B: Measure<FuncOutput = D>> Measure for Both<A, B> {
    type FuncOutput = B::FuncOutput;

    fn measure(mut func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String> {
        let mut a = None;
        let mut a_run = false;
        let mut b_run = false;
        let mut b = B::measure(|| {
            b_run = true;
            let mut v = None;
            a = Some(A::measure(|| {
                a_run = true;
                v = Some(func());
            }));
            v.unwrap_or_default()
        });

        if !b_run {
            // B failed without running the function. Try just running A.
            a = Some(A::measure(|| {
                func();
            }));
        } else if !a_run {
            // A failed without running the function.
            // B's result is meaningless. Re-run B.
            b = B::measure(|| func());
        }

        Ok(Self(a.unwrap(), b))
    }

    fn merge(self, rhs: Self) -> Self {
        let (a1, b1) = self.into();
        let (a2, b2) = rhs.into();
        let a = match (a1, a2) {
            (Ok(a1), Ok(a2)) => Ok(a1.merge(a2)),
            (Ok(a1), Err(_)) => Ok(a1),
            (Err(_), Ok(a2)) => Ok(a2),
            (Err(a1), Err(_)) => Err(a1),
        };
        let b = match (b1, b2) {
            (Ok(b1), Ok(b2)) => Ok(b1.merge(b2)),
            (Ok(b1), Err(_)) => Ok(b1),
            (Err(_), Ok(b2)) => Ok(b2),
            (Err(b1), Err(_)) => Err(b1),
        };
        Self(a, b)
    }

    fn need_more(&self) -> bool {
        match (&self.0, &self.1) {
            (Ok(a), Ok(b)) => a.need_more() || b.need_more(),
            _ => false,
        }
    }

    fn to_string(&self) -> String {
        let a = match &self.0 {
            Ok(a) => a.to_string(),
            Err(a) => a.clone(),
        };
        let b = match &self.1 {
            Ok(b) => b.to_string(),
            Err(b) => b.clone(),
        };
        format!("{} {}", a, b)
    }
}

impl<A, B> Into<(Result<A, String>, Result<B, String>)> for Both<A, B> {
    fn into(self) -> (Result<A, String>, Result<B, String>) {
        (self.0, self.1)
    }
}

// For ad-hoc messages
impl Measure for String {
    type FuncOutput = String;

    fn measure(mut func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String> {
        Ok(func())
    }

    fn merge(self, _rhs: Self) -> Self {
        self
    }

    fn need_more(&self) -> bool {
        false
    }

    fn to_string(&self) -> String {
        self.clone()
    }
}

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

    struct Dummy();
    struct DummyError();
    impl Measure for Dummy {
        type FuncOutput = ();

        fn measure(mut func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String> {
            func();
            Ok(Self())
        }

        fn merge(self, _rhs: Self) -> Self {
            Self()
        }

        fn need_more(&self) -> bool {
            false
        }

        fn to_string(&self) -> String {
            "dummy".to_string()
        }
    }

    impl Measure for DummyError {
        type FuncOutput = ();

        fn measure(mut _func: impl FnMut() -> Self::FuncOutput) -> Result<Self, String> {
            Err("unsupported".to_string())
        }

        fn merge(self, _rhs: Self) -> Self {
            Self()
        }

        fn need_more(&self) -> bool {
            false
        }

        fn to_string(&self) -> String {
            unreachable!()
        }
    }

    #[test]
    fn test_bytes() {
        let measured = Bytes::measure(|| 10).unwrap();
        assert_eq!(measured.to_string(), "     10 B ");
    }

    #[test]
    fn test_string() {
        let measured = String::measure(|| "abc def".to_string()).unwrap();
        assert_eq!(measured, "abc def");
    }

    #[test]
    fn test_both() {
        let measured = Both::<Dummy, Bytes>::measure(|| 10).unwrap();
        assert!(!measured.need_more());
        assert_eq!(measured.to_string(), "dummy      10 B ");
    }

    #[test]
    fn test_chained_both() {
        let measured = Both::<Both<Dummy, Dummy>, Both<Dummy, Bytes>>::measure(|| 10).unwrap();
        assert!(!measured.need_more());
        assert_eq!(measured.to_string(), "dummy dummy dummy      10 B ");
    }

    #[test]
    fn test_both_errors() {
        let measured = Both::<Dummy, DummyError>::measure(|| ()).unwrap();
        assert!(!measured.need_more());
        assert_eq!(measured.to_string(), "dummy unsupported");

        let measured = Both::<DummyError, Dummy>::measure(|| ()).unwrap();
        assert!(!measured.need_more());
        assert_eq!(measured.to_string(), "unsupported dummy");

        let measured = Both::<DummyError, DummyError>::measure(|| ()).unwrap();
        assert!(!measured.need_more());
        assert_eq!(measured.to_string(), "unsupported unsupported");
    }
}


/*
 * 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.
 */

//! Measurement based on procfs (/proc)

use super::Bytes;
use super::Measure;

/// Measure IO.
pub struct IO {
    rchar: u64,
    wchar: u64,
}

#[derive(Debug)]
struct IOSnapshot {
    rchar: u64,
    wchar: u64,
    rchar_overhead: u64,
}

fn read_io() -> Result<IOSnapshot, String> {
    let io_str = std::fs::read_to_string("/proc/self/io").map_err(|_| "(no data)".to_string())?;
    let mut rchar: u64 = 0;
    let mut wchar: u64 = 0;
    const RCHAR_PREFIX: &str = "rchar: ";
    const WCHAR_PREFIX: &str = "wchar: ";
    for line in io_str.lines() {
        if line.starts_with(RCHAR_PREFIX) {
            rchar += line[RCHAR_PREFIX.len()..]
                .parse::<u64>()
                .map_err(|_| "unexpected rchar".to_string())?;
        } else if line.starts_with(WCHAR_PREFIX) {
            wchar += line[WCHAR_PREFIX.len()..]
                .parse::<u64>()
                .map_err(|_| "unexpected wchar".to_string())?;
        }
    }
    // Reading io has side effect on rchar. Record it.
    let rchar_overhead = io_str.len() as u64;
    Ok(IOSnapshot {
        rchar,
        wchar,
        rchar_overhead,
    })
}

impl Measure for IO {
    type FuncOutput = ();

    fn measure(mut func: impl FnMut()) -> Result<Self, String> {
        let before = read_io()?;
        func();
        let after = read_io()?;
        let rchar = after.rchar - before.rchar - before.rchar_overhead;
        let wchar = after.wchar - before.wchar;
        Ok(Self { rchar, wchar })
    }

    fn merge(self, rhs: Self) -> Self {
        Self {
            rchar: self.rchar.max(rhs.rchar),
            wchar: self.wchar.max(rhs.wchar),
        }
    }

    fn need_more(&self) -> bool {
        false
    }

    fn to_string(&self) -> String {
        format!(
            "{}/{}",
            Bytes(self.rchar).to_string(),
            Bytes(self.wchar).to_string()
        )
    }
}

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

    #[test]
    fn test_io() {
        if let Ok(io) = IO::measure(|| {}) {
            // The test runner can run things in multi-thread and breaks the measurement here :/
            if io.rchar == 0 && io.wchar == 0 {
                assert_eq!(io.to_string(), "      0 B /      0 B ");
            }
        }
    }
}


/*
 * Portions 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.
 */

/* Copyright (c) 2018 Carl Lerche
 *
 * Permission is hereby granted, free of charge, to any
 * person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the
 * Software without restriction, including without
 * limitation the rights to use, copy, modify, merge,
 * publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software
 * is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice
 * shall be included in all copies or substantial portions
 * of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
 * ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
 * TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
 * PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
 * SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
 * IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */
#![feature(test)]

extern crate test;

use minibytes::Bytes;
use test::Bencher;

#[bench]
fn deref_unique(b: &mut Bencher) {
    let buf = Bytes::from(vec![0; 1024]);

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&buf[..]);
        }
    })
}

#[bench]
fn deref_shared(b: &mut Bencher) {
    let buf = Bytes::from(vec![0; 1024]);
    let _b2 = buf.clone();

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&buf[..]);
        }
    })
}

#[bench]
fn deref_static(b: &mut Bencher) {
    let buf = Bytes::from_static(b"hello world");

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&buf[..]);
        }
    })
}

#[bench]
fn clone_static(b: &mut Bencher) {
    let bytes =
        Bytes::from_static("hello world 1234567890 and have a good byte 0987654321".as_bytes());

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&bytes.clone());
        }
    })
}

#[bench]
fn clone_shared(b: &mut Bencher) {
    let bytes = Bytes::from(b"hello world 1234567890 and have a good byte 0987654321".to_vec());

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&bytes.clone());
        }
    })
}

#[bench]
fn clone_arc_vec(b: &mut Bencher) {
    use std::sync::Arc;
    let bytes = Arc::new(b"hello world 1234567890 and have a good byte 0987654321".to_vec());

    b.iter(|| {
        for _ in 0..1024 {
            test::black_box(&bytes.clone());
        }
    })
}

#[bench]
fn from_long_slice(b: &mut Bencher) {
    let data = [0u8; 128];
    b.bytes = data.len() as u64;
    b.iter(|| {
        let buf = Bytes::copy_from_slice(&data[..]);
        test::black_box(buf);
    })
}

#[bench]
fn slice_empty(b: &mut Bencher) {
    b.iter(|| {
        let b = Bytes::from(vec![17; 1024]).clone();
        for i in 0..1000 {
            test::black_box(b.slice(i % 100..i % 100));
        }
    })
}

#[bench]
fn slice_short_from_arc(b: &mut Bencher) {
    b.iter(|| {
        // `clone` is to convert to ARC
        let b = Bytes::from(vec![17; 1024]).clone();
        for i in 0..1000 {
            test::black_box(b.slice(1..2 + i % 10));
        }
    })
}


/*
 * 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::any::Any;
use std::ops::Range;
use std::ops::RangeBounds;
use std::sync::Arc;

pub type Bytes = AbstractBytes<[u8]>;
pub trait BytesOwner: AsRef<[u8]> + Send + Sync + 'static {}

/// Immutable bytes with zero-copy slicing and cloning.
pub struct AbstractBytes<T: ?Sized> {
    pub(crate) ptr: *const u8,
    pub(crate) len: usize,

    // Actual owner of the bytes. None for static buffers.
    pub(crate) owner: Option<Arc<dyn AbstractOwner<T>>>,
}

/// The actual storage owning the bytes.
pub trait AbstractOwner<T: ?Sized>: AsRef<T> + Send + Sync + 'static {
    fn as_any_mut(&mut self) -> &mut dyn Any;
}

impl<T: BytesOwner> AbstractOwner<[u8]> for T {
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

// AbstractOwner<T> is Send + Sync and AbstractBytes<T> is immutable.
unsafe impl<T: ?Sized> Send for AbstractBytes<T> {}
unsafe impl<T: ?Sized> Sync for AbstractBytes<T> {}

// #[derive(Clone)] does not work well with type parameters.
// Therefore implement Clone manually.
impl<T: ?Sized> Clone for AbstractBytes<T> {
    fn clone(&self) -> Self {
        Self {
            ptr: self.ptr,
            len: self.len,
            owner: self.owner.clone(),
        }
    }
}

// Core implementation of Bytes.
impl<T> AbstractBytes<T>
where
    T: SliceLike + ?Sized,
    T::Owned: AbstractOwner<T>,
{
    /// Returns a slice of self for the provided range.
    /// This operation is `O(1)`.
    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
        use std::ops::Bound;
        let start = match range.start_bound() {
            Bound::Included(&n) => n,
            Bound::Excluded(&n) => n + 1,
            Bound::Unbounded => 0,
        };
        let end = match range.end_bound() {
            Bound::Included(&n) => n + 1,
            Bound::Excluded(&n) => n,
            Bound::Unbounded => self.len,
        };
        assert!(start <= end, "invalid slice {}..{}", start, end);
        assert!(end <= self.len, "{} exceeds Bytes length {}", end, self.len);
        if start == end {
            Self::new()
        } else {
            T::check_slice_bytes(self.as_bytes(), start, end);
            Self {
                ptr: unsafe { self.ptr.add(start) },
                len: end - start,
                owner: self.owner.clone(),
            }
        }
    }

    /// Attempt to convert `slice` to a zero-copy slice of this `Bytes`.
    /// Copy the `slice` if zero-copy cannot be done.
    ///
    /// This is similar to `bytes::Bytes::slice_ref` from `bytes 0.5.4`,
    /// but does not panic.
    pub fn slice_to_bytes(&self, slice: &T) -> Self {
        match self.range_of_slice(slice) {
            Some(range) => self.slice(range),
            None => Self::copy_from_slice(slice),
        }
    }

    /// Return a range `x` so that `self[x]` matches `slice` exactly
    /// (not only content, but also internal pointer addresses).
    ///
    /// Returns `None` if `slice` is outside the memory range of this
    /// `Bytes`.
    ///
    /// This operation is `O(1)`.
    pub fn range_of_slice(&self, slice: &T) -> Option<Range<usize>> {
        let slice_start = slice.as_ptr() as usize;
        let slice_end = slice_start + slice.len();
        let bytes_start = self.ptr as usize;
        let bytes_end = bytes_start + self.len;
        if slice_start >= bytes_start && slice_end <= bytes_end {
            let start = slice_start - bytes_start;
            Some(start..start + slice.len())
        } else {
            None
        }
    }

    /// Creates an empty `Bytes`.
    #[inline]
    pub fn new() -> Self {
        let empty = T::EMPTY.as_bytes();
        Self {
            ptr: empty.as_ptr(),
            len: empty.len(),
            owner: None,
        }
    }

    /// Creates `Bytes` from a [`BytesOwner`] (for example, `Vec<u8>`).
    pub fn from_owner(value: impl AbstractOwner<T>) -> Self {
        let slice: &T = value.as_ref();
        let bytes = slice.as_bytes();
        Self {
            ptr: bytes.as_ptr(),
            len: bytes.len(),
            owner: Some(Arc::new(value)),
        }
    }

    /// Creates `Bytes` instance from slice, by copying it.
    pub fn copy_from_slice(data: &T) -> Self {
        Self::from_owner(data.to_owned())
    }

    #[inline]
    pub(crate) fn as_bytes(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
    }

    /// Attempt to downcast to an exclusive mut reference.
    ///
    /// Returns None if the type mismatches, or the internal reference count is
    /// not 0.
    pub fn downcast_mut<A: Any>(&mut self) -> Option<&mut A> {
        let arc_owner = match self.owner.as_mut() {
            None => return None,
            Some(owner) => owner,
        };
        let owner = match Arc::get_mut(arc_owner) {
            None => return None,
            Some(owner) => owner,
        };
        let any = owner.as_any_mut();
        any.downcast_mut()
    }
}

impl Bytes {
    #[inline]
    pub(crate) fn as_slice(&self) -> &[u8] {
        self.as_bytes()
    }

    /// Creates `Bytes` from a static slice.
    pub const fn from_static(slice: &'static [u8]) -> Self {
        Self {
            ptr: slice.as_ptr(),
            len: slice.len(),
            owner: None,
        }
    }

    /// Convert to `Vec<u8>`, in a zero-copy way if possible.
    pub fn into_vec(mut self) -> Vec<u8> {
        let len = self.len();
        match self.downcast_mut::<Vec<u8>>() {
            Some(ref mut owner) if owner.len() == len => {
                let mut result: Vec<u8> = Vec::new();
                std::mem::swap(&mut result, owner);
                result
            }
            Some(_) | None => self.as_slice().to_vec(),
        }
    }
}

pub trait SliceLike: 'static {
    type Owned;
    const EMPTY: &'static Self;

    fn as_bytes(&self) -> &[u8];
    fn to_owned(&self) -> Self::Owned;

    #[inline]
    fn check_slice_bytes(bytes: &[u8], start: usize, end: usize) {
        let _ = (bytes, start, end);
    }

    #[inline]
    fn len(&self) -> usize {
        self.as_bytes().len()
    }

    #[inline]
    fn as_ptr(&self) -> *const u8 {
        self.as_bytes().as_ptr()
    }
}

impl SliceLike for [u8] {
    type Owned = Vec<u8>;
    const EMPTY: &'static Self = b"";

    #[inline]
    fn as_bytes(&self) -> &[u8] {
        self
    }
    #[inline]
    fn to_owned(&self) -> Self::Owned {
        self.to_vec()
    }
}


/*
 * 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.
 */

//! Implement common traits for [`Bytes`] and [`Text`].

use std::ascii::escape_default;
use std::borrow;
use std::cmp;
use std::fmt;
use std::hash;
use std::ops;

use crate::Bytes;
use crate::BytesOwner;
use crate::Text;
use crate::TextOwner;

impl<T: BytesOwner> From<T> for Bytes {
    fn from(value: T) -> Self {
        Self::from_owner(value)
    }
}

impl From<&'static [u8]> for Bytes {
    fn from(value: &'static [u8]) -> Self {
        Self::from_static(value)
    }
}

impl From<&'static str> for Bytes {
    fn from(value: &'static str) -> Self {
        Self::from_static(value.as_bytes())
    }
}

impl AsRef<[u8]> for Bytes {
    #[inline]
    fn as_ref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl ops::Deref for Bytes {
    type Target = [u8];
    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl hash::Hash for Bytes {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state);
    }
}

impl borrow::Borrow<[u8]> for Bytes {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}

impl Default for Bytes {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: AsRef<[u8]>> PartialEq<T> for Bytes {
    fn eq(&self, other: &T) -> bool {
        self.as_slice() == other.as_ref()
    }
}

impl Eq for Bytes {}

impl<T: AsRef<[u8]>> PartialOrd<T> for Bytes {
    fn partial_cmp(&self, other: &T) -> Option<cmp::Ordering> {
        self.as_slice().partial_cmp(other.as_ref())
    }
}

impl Ord for Bytes {
    fn cmp(&self, other: &Bytes) -> cmp::Ordering {
        self.as_slice().cmp(other.as_slice())
    }
}

impl fmt::Debug for Bytes {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Use `[u8]::escape_ascii` when inherent_ascii_escape is stabilized.
        f.write_str("b\"")?;
        for &byte in self.as_slice() {
            fmt::Display::fmt(&escape_default(byte), f)?;
        }
        f.write_str("\"")?;
        Ok(())
    }
}

impl<T: TextOwner> From<T> for Text {
    fn from(value: T) -> Self {
        Self::from_owner(value)
    }
}

impl From<&'static str> for Text {
    fn from(value: &'static str) -> Self {
        Self::from_static(value)
    }
}

impl AsRef<str> for Text {
    #[inline]
    fn as_ref(&self) -> &str {
        self.as_slice()
    }
}

impl ops::Deref for Text {
    type Target = str;
    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl hash::Hash for Text {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state);
    }
}

impl borrow::Borrow<str> for Text {
    fn borrow(&self) -> &str {
        self.as_slice()
    }
}

impl Default for Text {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: AsRef<str>> PartialEq<T> for Text {
    fn eq(&self, other: &T) -> bool {
        self.as_slice() == other.as_ref()
    }
}

impl Eq for Text {}

impl<T: AsRef<str>> PartialOrd<T> for Text {
    fn partial_cmp(&self, other: &T) -> Option<cmp::Ordering> {
        self.as_slice().partial_cmp(other.as_ref())
    }
}

impl Ord for Text {
    fn cmp(&self, other: &Text) -> cmp::Ordering {
        self.as_slice().cmp(other.as_slice())
    }
}

impl fmt::Debug for Text {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(self.as_slice(), f)
    }
}

impl fmt::Display for Text {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(self.as_slice(), f)
    }
}


/*
 * 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.
 */

//! # minibytes
//!
//! This create provides the [`Bytes`] type. It is similar to `&[u8]`: cloning
//! or slicing are zero-copy. Unlike `&[u8]`, `Bytes` does not have lifetime.
//! This is done by maintaining the life cycle of the underlying storage using
//! reference count.
//!
//! Aside from supporting `Vec<u8>` as the underlying storage, [`Bytes`] also
//! supports [`memmap2::Mmap`]. Libraries can implement [`BytesOwner`] for other
//! types to further extend storage support.

mod bytes;
mod impls;
mod owners;
mod serde;
mod text;

#[cfg(test)]
mod tests;

pub use text::Text;
pub use text::TextOwner;

pub use crate::bytes::Bytes;
pub use crate::bytes::BytesOwner;


/*
 * 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.
 */

//! Implement [`BytesOwner`] and [`TextOwner`] for common types.

use crate::BytesOwner;
use crate::TextOwner;

impl BytesOwner for Vec<u8> {}
impl BytesOwner for Box<[u8]> {}
impl BytesOwner for String {}
#[cfg(feature = "frommmap")]
impl BytesOwner for memmap2::Mmap {}
#[cfg(feature = "frombytes")]
impl BytesOwner for bytes::Bytes {}

impl TextOwner for String {}


/*
 * 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::fmt;

use serde::de;
use serde::Deserialize;
use serde::Deserializer;
use serde::Serialize;
use serde::Serializer;

use crate::Bytes;

impl Serialize for Bytes {
    #[inline]
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_bytes(&self)
    }
}

struct BytesVisitor;

impl<'de> de::Visitor<'de> for BytesVisitor {
    type Value = Bytes;

    fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("byte slice")
    }

    fn visit_borrowed_bytes<E: de::Error>(self, v: &'de [u8]) -> Result<Self::Value, E> {
        Ok(Bytes::copy_from_slice(v))
    }

    fn visit_borrowed_str<E: de::Error>(self, v: &'de str) -> Result<Self::Value, E> {
        Ok(Bytes::copy_from_slice(v.as_bytes()))
    }

    fn visit_byte_buf<E>(self, v: Vec<u8>) -> Result<Self::Value, E> {
        Ok(Bytes::from(v))
    }

    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E> {
        Ok(Bytes::copy_from_slice(v))
    }
}

impl<'de> Deserialize<'de> for Bytes {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        deserializer.deserialize_bytes(BytesVisitor)
    }
}


/*
 * 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 quickcheck::quickcheck;

use crate::Bytes;
use crate::Text;

quickcheck! {
    fn test_shallow_clone(v: Vec<u8>) -> bool {
        let a: Bytes = v.into();
        let b: Bytes = a.clone();
        a == b && a.as_ptr() == b.as_ptr()
    }

    fn test_shallow_slice(v: Vec<u8>) -> bool {
        let a: Bytes = v.into();
        let b: Bytes = a.slice(..a.len() / 2);
        b == &a[..b.len()] && (b.is_empty() || a.as_ptr() == b.as_ptr())
    }

    fn test_range_of_slice(v: Vec<u8>) -> bool {
        let a: Bytes = v.into();
        let range1 = a.len() / 3.. a.len() * 2 / 3;
        let slice = a.slice(range1.clone());
        if slice.is_empty() {
            true
        } else {
            let range2 = a.range_of_slice(&slice).unwrap();
            range1 == range2
        }
    }

    fn test_text_shallow_clone(v: String) -> bool {
        let a: Text = v.into();
        let b: Text = a.clone();
        a == b && a.as_ptr() == b.as_ptr()
    }
}

static SAMPLE_TEXT: &str = "è¿™æ˜¯æµ‹è¯•ç”¨çš„æ–‡å—";

#[test]
fn test_text_slice_valid() {
    let a: Text = SAMPLE_TEXT.into();
    let b = a.slice(3..12);
    let s: &str = b.as_ref();
    let c = a.slice_to_bytes(s);
    let d = b.slice_to_bytes(s);
    assert_eq!(b.as_ptr(), c.as_ptr());
    assert_eq!(b.as_ptr(), d.as_ptr());
}

#[test]
fn test_text_to_string() {
    let a: Text = SAMPLE_TEXT.into();
    assert_eq!(a.to_string(), SAMPLE_TEXT.to_string());
}

#[test]
#[should_panic]
fn test_text_slice_invalid() {
    let a: Text = SAMPLE_TEXT.into();
    let _b = a.slice(3..11); // invalid utf-8 boundary
}

#[test]
fn test_downcast_mut() {
    let v = b"abcd".to_vec();
    let mut b = Bytes::from(v);
    assert!(b.downcast_mut::<Vec<u8>>().is_some());
    assert!(b.downcast_mut::<String>().is_none());
    let mut c = b.clone();
    assert!(b.downcast_mut::<Vec<u8>>().is_none());
    assert!(c.downcast_mut::<Vec<u8>>().is_none());
}

#[test]
fn test_into_vec() {
    let v = b"abcd".to_vec();
    let ptr1 = &v[0] as *const u8;
    let b = Bytes::from(v);
    let v = b.into_vec(); // zero-copy
    let ptr2 = &v[0] as *const u8;
    assert_eq!(ptr1, ptr2);

    let b = Bytes::from(v);
    let _c = b.clone();
    let v = b.into_vec(); // not zero-copy because refcount > 1
    let ptr3 = &v[0] as *const u8;
    assert_ne!(ptr1, ptr3);

    let b = Bytes::from(v);
    let c = b.slice(1..3);
    drop(b);
    assert_eq!(c.into_vec(), b"bc");
}

#[test]
fn test_bytes_debug_format() {
    let v = b"printable\t\r\n\'\"\\\x00\x01\x02printable".to_vec();
    let b = Bytes::from(v);
    let escaped = format!("{:?}", b);
    let expected = r#"b"printable\t\r\n\'\"\\\x00\x01\x02printable""#;
    assert_eq!(escaped, expected);
}


/*
 * 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::any::Any;

use super::bytes::AbstractBytes;
use super::bytes::AbstractOwner;
use super::bytes::SliceLike;

pub type Text = AbstractBytes<str>;
pub trait TextOwner: AsRef<str> + Send + Sync + 'static {}

impl<T: TextOwner> AbstractOwner<str> for T {
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

impl Text {
    /// Creates `Text` from a static str.
    pub const fn from_static(slice: &'static str) -> Self {
        Self {
            ptr: slice.as_ptr(),
            len: slice.len(),
            owner: None,
        }
    }

    #[inline]
    pub(crate) fn as_slice(&self) -> &str {
        let bytes = self.as_bytes();
        // bytes was validated as utf-8.
        unsafe { std::str::from_utf8_unchecked(bytes) }
    }
}

impl SliceLike for str {
    type Owned = String;
    const EMPTY: &'static Self = "";

    #[inline]
    fn check_slice_bytes(bytes: &[u8], start: usize, end: usize) {
        // called by AbstractBytes::slice, bytes was validated as utf-8.
        let s = unsafe { std::str::from_utf8_unchecked(bytes) };
        // check whether the slicing is valid.
        let _ = s[start..end];
    }
    #[inline]
    fn as_bytes(&self) -> &[u8] {
        self.as_bytes()
    }
    #[inline]
    fn to_owned(&self) -> Self::Owned {
        self.to_string()
    }
}


/*
 * 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 mpatch;


/*
 * 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::os::raw::c_char;
use std::os::raw::c_void;
use std::ptr;

use libc::ssize_t;
use mpatch_sys::*;

unsafe extern "C" fn get_next_link(deltas: *mut c_void, index: ssize_t) -> *mut mpatch_flist {
    let deltas = (deltas as *const Vec<&[u8]>).as_ref().unwrap();
    if index < 0 || index as usize >= deltas.len() {
        return ptr::null_mut();
    }

    let delta: &[u8] = deltas[index as usize];

    let mut res: *mut mpatch_flist = ptr::null_mut();
    if mpatch_decode(
        delta.as_ptr() as *const c_char,
        delta.len() as isize,
        &mut res,
    ) < 0
    {
        return ptr::null_mut();
    }

    return res;
}

pub fn get_full_text(base_text: &[u8], deltas: &Vec<&[u8]>) -> Result<Vec<u8>, &'static str> {
    // If there are no deltas, just return the full text portion
    if deltas.len() == 0 {
        return Ok(base_text.to_vec());
    }

    unsafe {
        let patch: *mut mpatch_flist = mpatch_fold(
            deltas as *const Vec<&[u8]> as *mut c_void,
            Some(get_next_link),
            0,
            deltas.len() as isize,
        );
        if patch.is_null() {
            return Err("mpatch failed to process the deltas");
        }

        let outlen = mpatch_calcsize(base_text.len() as isize, patch);
        if outlen < 0 {
            mpatch_lfree(patch);
            return Err("mpatch failed to calculate size");
        }

        let outlen = outlen as usize;
        let mut result: Vec<u8> = Vec::with_capacity(outlen);
        result.set_len(outlen);
        if mpatch_apply(
            result.as_mut_ptr() as *mut c_char,
            base_text.as_ptr() as *const c_char,
            base_text.len() as ssize_t,
            patch,
        ) < 0
        {
            mpatch_lfree(patch);
            return Err("mpatch failed to apply patches");
        }

        mpatch_lfree(patch);
        return Ok(result);
    }
}

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

    #[test]
    fn no_deltas() {
        let base_text = b"hello";
        let full_text = get_full_text(&base_text[..], &vec![]).unwrap();
        assert_eq!(base_text, full_text.as_slice());
    }

    #[test]
    fn no_deltas_empty_base() {
        let base_text = b"";
        let full_text = get_full_text(&base_text[..], &vec![]).unwrap();
        assert_eq!(base_text, full_text.as_slice());
    }

    #[test]
    fn test_apply_delta() {
        let base_text = b"My data";
        let deltas: Vec<&[u8]> =
            vec![b"\x00\x00\x00\x03\x00\x00\x00\x03\x00\x00\x00\x0Adeltafied "];

        let full_text = get_full_text(&base_text[..], &deltas).unwrap();
        assert_eq!(b"My deltafied data", full_text[..].as_ref());
    }

    #[test]
    fn test_apply_deltas() {
        let base_text = b"My data";
        let deltas: Vec<&[u8]> = vec![
            b"\x00\x00\x00\x03\x00\x00\x00\x03\x00\x00\x00\x0Adeltafied ",
            b"\x00\x00\x00\x03\x00\x00\x00\x0D\x00\x00\x00\x10still deltafied ",
        ];

        let full_text = get_full_text(&base_text[..], &deltas).unwrap();
        assert_eq!(b"My still deltafied data", full_text[..].as_ref());
    }

    #[test]
    fn test_apply_invalid_deltas() {
        let base_text = b"My data";

        // Short delta
        let deltas: Vec<&[u8]> = vec![b"\x00\x03"];

        let full_text = get_full_text(&base_text[..], &deltas);
        assert!(full_text.is_err());

        // Short data
        let deltas: Vec<&[u8]> = vec![
            b"\x00\x00\x00\x03\x00\x00\x00\x03\x00\x00\x00\x0Adeltafied ",
            b"\x00\x00\x00\x03\x00\x00\x00\x03\x00\x00\x00\x0Adelta",
        ];

        let full_text = get_full_text(&base_text[..], &deltas);
        assert!(full_text.is_err());

        // Delta doesn't match base_text
        let deltas: Vec<&[u8]> =
            vec![b"\x00\x00\x00\xFF\x00\x00\x01\x00\x00\x00\x00\x0Adeltafied "];

        let full_text = get_full_text(&base_text[..], &deltas);
        assert!(full_text.is_err());
    }
}


/*
 * 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.
 */

fn main() {
    cc::Build::new()
        .file("../../edenscm/mpatch.c")
        .include("../../../../")
        .compile("mpatch");
}


/*
 * 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.
 */

// Include the mpatch bindings
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use std::os::raw::c_char;
use std::os::raw::c_int;
use std::os::raw::c_void;

#[repr(C)]
pub struct mpatch_flist {
    pub base: *mut mpatch_frag,
    pub head: *mut mpatch_frag,
    pub tail: *mut mpatch_frag,
}

#[repr(C)]
pub struct mpatch_frag {
    pub start: c_int,
    pub end: c_int,
    pub len: c_int,
    pub data: *const c_char,
}

extern "C" {
    pub fn mpatch_decode(bin: *const c_char, len: isize, res: *mut *mut mpatch_flist) -> c_int;

    pub fn mpatch_calcsize(len: isize, l: *mut mpatch_flist) -> isize;

    pub fn mpatch_lfree(a: *mut mpatch_flist);

    pub fn mpatch_apply(
        buf: *mut c_char,
        orig: *const c_char,
        len: isize,
        l: *mut mpatch_flist,
    ) -> c_int;

    pub fn mpatch_fold(
        bins: *mut c_void,
        get_next_item: Option<
            unsafe extern "C" fn(arg1: *mut c_void, arg2: isize) -> *mut mpatch_flist,
        >,
        start: isize,
        end: isize,
    ) -> *mut mpatch_flist;
}


/*
 * 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.
 */

//! The mutationstore is a store for recording records of commit mutations for
//! commits that are not in the local repository.
//!
//! It uses an indexedlog to store the data.  Each mutation entry corresponds to
//! the information of the mutation that led to the creation of a particular commit,
//! which is recorded as the successor in the entry.
//!
//! Entries can come from three possible places:
//!
//! * Commit metadata for a commit not available locally
//!
//! * Obsmarkers for repos that have been migrated from evolution tracking
//!
//! * Synthetic for entries created synthetically, e.g. by a pullcreatemarkers
//!   implementation.
//!
//! The other commits referred to in an entry must predate the successor commit.
//! For entries that originated from commits, this is ensured, as the successor
//! commit hash includes the other commit hashes.  For other entry types, it is
//! an error to refer to later commits, and any entry that causes a cycle will
//! be ignored.

#![allow(clippy::redundant_closure)]

use std::collections::HashMap;
use std::collections::HashSet;
use std::io::Cursor;
use std::path::Path;
use std::sync::Arc;

use anyhow::Result;
use bitflags::bitflags;
use dag::namedag::MemNameDag;
use dag::nameset::hints::Flags;
use dag::ops::DagAddHeads;
use dag::DagAlgorithm;
use dag::Set;
use dag::VertexName;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use indexedlog::log::IndexDef;
use indexedlog::log::IndexOutput;
use indexedlog::log::Log;
use indexedlog::log::{self as ilog};
use indexedlog::DefaultOpenOptions;
use indexedlog::OpenWithRepair;
pub use indexedlog::Repair;
use types::mutation::MutationEntry;
use types::node::Node;
use vlqencoding::VLQDecodeAt;

pub struct MutationStore {
    log: Log,
    pending: Vec<MutationEntry>,
}

bitflags! {
    pub struct DagFlags: u8 {
        /// Include successors.
        const SUCCESSORS = 0b1;

        /// Include predecessors.
        const PREDECESSORS = 0b10;
    }
}

const INDEX_PRED: usize = 0;
const INDEX_SUCC: usize = 1;
const INDEX_SPLIT: usize = 2;

impl DefaultOpenOptions<ilog::OpenOptions> for MutationStore {
    fn default_open_options() -> ilog::OpenOptions {
        const NODE_LEN: usize = Node::len();
        const SUCC_START: usize = 1usize;
        const PRED_COUNT_START: usize = SUCC_START + NODE_LEN;
        let succ_index = |_data: &[u8]| {
            vec![IndexOutput::Reference(
                SUCC_START as u64..PRED_COUNT_START as u64,
            )]
        };
        let pred_index = |data: &[u8]| {
            let (pred_count, pred_start) = data
                .read_vlq_at(PRED_COUNT_START)
                .map(|(pred_count, vlq_size)| (pred_count, PRED_COUNT_START + vlq_size))
                .unwrap_or((0, 0));
            (0..pred_count)
                .map(|i| pred_start + NODE_LEN * i)
                .map(|i: usize| IndexOutput::Reference(i as u64..i as u64 + NODE_LEN as u64))
                .collect()
        };
        let split_index = |data: &[u8]| {
            let (split_count, split_start) = data
                .read_vlq_at(PRED_COUNT_START)
                .and_then(|(pred_count, vlq1_size): (usize, usize)| {
                    data.read_vlq_at(PRED_COUNT_START + vlq1_size + NODE_LEN * pred_count)
                        .map(|(split_count, vlq2_size)| {
                            (
                                split_count,
                                PRED_COUNT_START + vlq1_size + NODE_LEN * pred_count + vlq2_size,
                            )
                        })
                })
                .unwrap_or((0, 0));
            (0..split_count)
                .map(|i| split_start + NODE_LEN * i)
                .map(|i: usize| IndexOutput::Reference(i as u64..i as u64 + NODE_LEN as u64))
                .collect()
        };

        // Allow some lag to make the indexing more efficient.  Set to 10KB, which is roughly
        // 100 records.
        let lag_threshold = 10000;
        ilog::OpenOptions::new().create(true).index_defs(vec![
            IndexDef::new("pred", pred_index).lag_threshold(lag_threshold),
            IndexDef::new("succ", succ_index).lag_threshold(lag_threshold),
            IndexDef::new("split", split_index).lag_threshold(lag_threshold),
        ])
    }
}

impl MutationStore {
    pub fn open(path: impl AsRef<Path>) -> Result<MutationStore> {
        let log = Self::default_open_options().open_with_repair(path.as_ref())?;
        let pending = Vec::new();
        Ok(MutationStore { log, pending })
    }

    /// Add an entry. Consider adding automatic entries based on this entry.
    /// See `flush` for automatic entries.
    pub fn add(&mut self, entry: &MutationEntry) -> Result<()> {
        self.add_raw(entry)?;
        self.pending.push(entry.clone());
        Ok(())
    }

    /// Add an entry. Do not consider adding automatic entries.
    pub fn add_raw(&mut self, entry: &MutationEntry) -> Result<()> {
        let mut buf = Vec::with_capacity(types::mutation::DEFAULT_ENTRY_SIZE);
        entry.serialize(&mut buf)?;
        self.log.append(buf.as_slice())?;
        Ok(())
    }

    pub async fn flush(&mut self) -> Result<()> {
        // If P -> Q, X -> Y are being added, and there is an existing chain P
        // -> ... -> X, add a Q -> Y marker automatically.
        // Note: P must not equal to X or Y.
        //
        // See also D7121487.

        // Prepare for calculation.
        let mut pred_map = HashMap::with_capacity(self.pending.len()); // node -> index
        let mut pred_nodes = Vec::with_capacity(self.pending.len());
        for (i, entry) in self.pending.iter().enumerate() {
            let pred = entry.preds[0];
            pred_map.insert(pred, i);
            pred_nodes.push(pred);
        }
        let pred_set =
            Set::from_static_names(pred_nodes.iter().map(|p| VertexName::copy_from(p.as_ref())));
        let dag = self
            .get_dag_advanced(pred_nodes, DagFlags::SUCCESSORS)
            .await?;
        let mut new_entries = Vec::new();

        // Scan through "X"s.
        for entry in &self.pending {
            let x = entry.preds[0];
            // Find all "P"s, as in P -> ... -> X, and X -> Y.
            let x_set = VertexName::copy_from(x.as_ref()).into();
            // "dag" is locally built and should be non-blocking.
            let x_ancestors = match dag.ancestors(x_set).await {
                Ok(set) => set,
                Err(_) => continue, // might have cycles
            };
            let ps = x_ancestors & pred_set.clone();
            let mut iter = ps.iter().await?;
            while let Some(p) = iter.next().await {
                let p = Node::from_slice(p?.as_ref())?;
                let y = entry.succ;
                if p == x || p == y {
                    continue;
                }
                let q = self.pending[pred_map[&p]].succ;
                if q == x || q == y || q == p {
                    continue;
                }
                // Copy P -> X to Q -> Y.
                let copy_entry = match self.get(x)? {
                    Some(entry) => entry,
                    _ => continue,
                };
                let op = if copy_entry.op.ends_with("-copy") {
                    copy_entry.op.clone()
                } else {
                    format!("{}-copy", &copy_entry.op)
                };
                // The new entry will be the one returned by `get(y)`.
                // It overrides the "x -> y" entry.
                let new_entry = MutationEntry {
                    succ: y,
                    preds: vec![x, q],
                    op,
                    ..copy_entry
                };
                new_entries.push(new_entry);
            }
        }

        let mut buf = Vec::with_capacity(types::mutation::DEFAULT_ENTRY_SIZE);
        for entry in new_entries {
            buf.clear();
            entry.serialize(&mut buf)?;
            self.log.append(buf.as_slice())?;
        }

        self.log.flush()?;
        self.pending.clear();
        Ok(())
    }

    fn try_clone(&self) -> Result<Self> {
        Ok(Self {
            log: self.log.try_clone()?,
            pending: self.pending.clone(),
        })
    }

    /// Calculate the "obsolete" set, a subset of `draft` with visible successors.
    /// A vertex that is in `public` or `draft` is considered visible.
    pub async fn calculate_obsolete(&self, public: Set, draft: Set) -> Result<Set> {
        let visible = public | draft.clone();
        let this = Arc::new(self.try_clone()?);
        let hints = draft.hints().clone();
        hints.update_flags_with(|f| f - Flags::ANCESTORS - Flags::FULL);

        // Vertex -> Node.
        let draft_nodes = draft
            .iter()
            .await?
            .filter_map(|v| async { v.ok().and_then(|v| Node::from_slice(v.as_ref()).ok()) })
            .collect::<Vec<Node>>()
            .await;

        // Obtain the obsolete graph about draft successors.
        let obsdag = this
            .get_dag_advanced(draft_nodes, DagFlags::SUCCESSORS)
            .await
            .map_err(|e| dag::errors::BackendError::Other(e))?;

        // Filter out invisible successors.
        let obsvisible = {
            let mut obsall = obsdag.all().await?;
            // In a non-lazy graph the following code is good enough to calculate
            // obsvisible: obsdag.ancestors(obsall & visible).await?
            //
            // However, in a lazy graph obsdag.all() might have too many names
            // outside the main graph and cause excessive server-side lookups.
            // So we manually ignore names not in the local graph to avoid the
            // slow path.
            if let Some(visible_id_convert) = visible.id_convert() {
                let obsnames: Vec<VertexName> = { obsall.iter().await?.try_collect().await? };
                let obsnames: Vec<VertexName> = {
                    let contains = visible_id_convert
                        .contains_vertex_name_locally(&obsnames)
                        .await?;
                    obsnames
                        .into_iter()
                        .zip(contains)
                        .filter_map(|(v, c)| if c { Some(v) } else { None })
                        .collect()
                };
                obsall = Set::from_static_names(obsnames);
            }
            obsdag.ancestors(obsall & visible).await?
        };

        // Heads of `obsvisible` are not obsoleted. Other part (parent) of
        // `obsvisible` are obsoleted.
        let obsoleted = obsdag.parents(obsvisible).await?;

        // Filter out unknown nodes.
        let obsoleted = draft & obsoleted;

        // Flatten the set for performance.
        Ok(obsoleted.flatten().await?)
    }

    pub fn get_successors_sets(&self, node: Node) -> Result<Vec<Vec<Node>>> {
        let mut successors_sets = Vec::new();
        for entry in self.log.lookup(INDEX_PRED, &node)? {
            let mutation_entry = MutationEntry::deserialize(&mut Cursor::new(entry?))?;
            let mut successors = Vec::new();
            successors.extend(&mutation_entry.split);
            successors.push(mutation_entry.succ);
            successors_sets.push(successors);
        }
        Ok(successors_sets)
    }

    pub fn get_predecessors(&self, node: Node) -> Result<Vec<Node>> {
        let mut lookup = self
            .log
            .lookup(INDEX_SUCC, &node)?
            .chain(self.log.lookup(INDEX_SPLIT, &node)?);
        let predecessors = if let Some(entry) = lookup.next() {
            let mutation_entry = MutationEntry::deserialize(&mut Cursor::new(entry?))?;
            mutation_entry.preds.clone()
        } else {
            vec![]
        };
        Ok(predecessors)
    }

    pub fn get_split_head(&self, node: Node) -> Result<Option<MutationEntry>> {
        let mutation_entry = match self.log.lookup(INDEX_SPLIT, &node)?.next() {
            Some(entry) => Some(MutationEntry::deserialize(&mut Cursor::new(entry?))?),
            None => None,
        };
        Ok(mutation_entry)
    }

    pub fn get(&self, succ: Node) -> Result<Option<MutationEntry>> {
        let mutation_entry = match self.log.lookup(INDEX_SUCC, &succ)?.next() {
            Some(entry) => Some(MutationEntry::deserialize(&mut Cursor::new(entry?))?),
            None => None,
        };
        Ok(mutation_entry)
    }

    /// Return a connected component that includes `nodes` and represents
    /// commit replacement relations.  The returned graph supports graph
    /// operations like common ancestors, heads, roots, etc. Parents in the
    /// graph are predecessors.
    pub async fn get_dag(&self, nodes: Vec<Node>) -> Result<MemNameDag> {
        self.get_dag_advanced(nodes, DagFlags::SUCCESSORS | DagFlags::PREDECESSORS)
            .await
    }

    /// Advanced version of `get_dag`. Specify whether to include successors or
    /// predecessors explicitly.
    pub async fn get_dag_advanced(&self, nodes: Vec<Node>, flags: DagFlags) -> Result<MemNameDag> {
        // Raw parent map. Might contain cycles.
        let mut parent_map = HashMap::<Node, Vec<Node>>::new();
        let mut non_heads = HashSet::<Node>::new();
        let mut add_parent = |parent: &Node, child: &Node| {
            let parents = parent_map.entry(child.clone()).or_default();
            if !parents.contains(parent) {
                parents.push(parent.clone());
                non_heads.insert(parent.clone());
            }
        };

        // Visit nodes. Fill parent map.
        let mut to_visit = nodes;
        let mut connected = HashSet::new();
        while let Some(node) = to_visit.pop() {
            if !connected.insert(node.clone()) {
                continue;
            }
            if flags.contains(DagFlags::SUCCESSORS) {
                for entry in self.log.lookup(INDEX_PRED, &node)? {
                    let entry = MutationEntry::deserialize(&mut Cursor::new(entry?))?;
                    add_parent(&node, &entry.succ);
                    to_visit.push(entry.succ);
                }
            }
            if flags.contains(DagFlags::PREDECESSORS) {
                for entry in self.log.lookup(INDEX_SUCC, &node)? {
                    let entry = MutationEntry::deserialize(&mut Cursor::new(entry?))?;
                    for pred in entry.preds {
                        add_parent(&pred, &node);
                        to_visit.push(pred);
                    }
                }
            }
        }

        // Construct parent_func.
        let parent_func = move |node: VertexName| -> dag::Result<Vec<VertexName>> {
            match parent_map.get(&Node::from_slice(node.as_ref()).unwrap()) {
                None => Ok(Vec::new()),
                Some(parents) => Ok(parents
                    .iter()
                    .map(|n| VertexName::copy_from(n.as_ref()))
                    .collect()),
            }
        };
        let parent_func = dag::utils::break_parent_func_cycle(parent_func);

        // Calculate heads. This makes multiple things more efficient:
        // `add_heads`, `break_parent_func_cycle`, and the resulting graph.
        let mut heads: Vec<Node> = connected.difference(&non_heads).cloned().collect();
        heads.sort_unstable();
        let heads: Vec<VertexName> = heads
            .into_iter()
            .map(|n| VertexName::copy_from(n.as_ref()))
            .collect();

        // Construct the graph.
        let mut dag = MemNameDag::new();
        let parents: Box<dyn Fn(VertexName) -> dag::Result<Vec<VertexName>> + Send + Sync> =
            Box::new(parent_func);

        // Inserting to a memory DAG from a fully known parent function is non-blocking.
        dag.add_heads(&parents, &heads.into()).await?;
        Ok(dag)
    }
}

#[cfg(test)]
mod tests {
    use dag::nonblocking::non_blocking_result;
    use dag::nonblocking::non_blocking_result as r;
    use dag::DagAlgorithm;
    use rand::SeedableRng;
    use rand_chacha::ChaChaRng;
    use tempdir::TempDir;

    use super::*;

    #[test]
    fn test_basic_store() {
        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let dir = TempDir::new("mutationstore").unwrap();
        let nodes = Node::random_distinct(&mut rng, 20);

        {
            let mut ms = MutationStore::open(dir.path()).expect("can open the store");
            ms.add(&MutationEntry {
                succ: nodes[1],
                preds: vec![nodes[0], nodes[2], nodes[3]],
                split: vec![],
                op: "fold".into(),
                user: "test".into(),
                time: 123456789,
                tz: -7200,
                extra: vec![(
                    Box::from(&b"note"[..]),
                    Box::from(&b"note about folding"[..]),
                )],
            })
            .expect("can add to the store");
            ms.add(&MutationEntry {
                succ: nodes[4],
                preds: vec![nodes[0]],
                split: vec![nodes[5], nodes[6]],
                op: "split".into(),
                user: "test".into(),
                time: 123456789,
                tz: -7200,
                extra: vec![],
            })
            .expect("can add to the store");

            r(ms.flush()).expect("can flush the store");
        }
        {
            let ms = MutationStore::open(dir.path()).expect("can re-open the store");
            let mut expected_successors_sets =
                vec![vec![nodes[1]], vec![nodes[5], nodes[6], nodes[4]]];
            expected_successors_sets.sort_unstable();
            let mut successors_sets = ms
                .get_successors_sets(nodes[0])
                .expect("can get successors sets");
            successors_sets.sort_unstable();
            assert_eq!(successors_sets, expected_successors_sets);
            assert_eq!(
                ms.get_successors_sets(nodes[3])
                    .expect("can get successors sets"),
                vec![vec![nodes[1]]]
            );
            assert_eq!(
                ms.get_successors_sets(nodes[1])
                    .expect("can get successors sets"),
                Vec::<Vec<Node>>::new()
            );
            assert_eq!(
                ms.get_predecessors(nodes[1]).expect("can get predecessors"),
                vec![nodes[0], nodes[2], nodes[3]]
            );
            assert_eq!(
                ms.get_predecessors(nodes[4]).expect("can get predecessors"),
                vec![nodes[0]]
            );
            assert_eq!(
                ms.get_predecessors(nodes[5]).expect("can get predecessors"),
                vec![nodes[0]]
            );
            assert_eq!(
                ms.get_predecessors(nodes[3]).expect("can get predecessors"),
                vec![]
            );
            assert_eq!(
                ms.get_split_head(nodes[5])
                    .expect("can get split head")
                    .unwrap()
                    .succ,
                nodes[4],
            );
            assert_eq!(
                &ms.get(nodes[4])
                    .expect("can get mutation entry")
                    .unwrap()
                    .user,
                "test",
            );
            assert_eq!(
                &ms.get(nodes[1])
                    .expect("can get mutation entry")
                    .unwrap()
                    .extra[0]
                    .1,
                &Box::from(&b"note about folding"[..])
            );
        }
    }

    #[test]
    fn test_dag() -> Result<()> {
        let dir = TempDir::new("mutationstore")?;
        let mut ms = MutationStore::open(dir.path())?;
        let parents = drawdag::parse(
            r#"
             D E Z
             |\| |
             B C Y
             |/  |
             A   X
             "#,
        );
        // str (length 1) -> Node
        let n = |s: &str| -> Node { Node::from_slice(s.repeat(Node::len()).as_bytes()).unwrap() };
        let mut iter = parents.iter().collect::<Vec<_>>();
        iter.sort();
        for (name, parents) in iter {
            let node = n(name);
            for parent in parents {
                let parent = n(parent);
                ms.add(&MutationEntry {
                    succ: node,
                    preds: vec![parent],
                    split: vec![],
                    op: "rewrite".into(),
                    user: "test".into(),
                    time: 123456789,
                    tz: -7200,
                    extra: vec![],
                })?;
            }
        }

        let dag = r(ms.get_dag(vec![n("B")]))?;
        assert_eq!(r(non_blocking_result(dag.all())?.count())?, 5); // A B C D E
        assert_eq!(
            dag::render::render_namedag(&dag, |v| Some(format!("({})", v.as_ref()[0] as char)))?,
            r#"
            o  4545454545454545454545454545454545454545 (E)
            â”‚
            â”‚ o  4444444444444444444444444444444444444444 (D)
            â•â”€â”¤
            o â”‚  4343434343434343434343434343434343434343 (C)
            â”‚ â”‚
            â”‚ o  4242424242424242424242424242424242424242 (B)
            â”œâ”€â•¯
            o  4141414141414141414141414141414141414141 (A)"#
        );
        Ok(())
    }

    #[test]
    fn test_dag_cycle() -> Result<()> {
        let dir = TempDir::new("mutationstore")?;
        let mut ms = MutationStore::open(dir.path())?;

        for (pred, succ) in [("A", "B"), ("B", "C"), ("C", "A")] {
            add(&mut ms, pred, succ)?;
        }
        r(ms.flush())?;

        // Nothing - cycles without a head is not rendered.
        assert_eq!(render(&ms, "A")?, "\n");
        assert_eq!(render(&ms, "B")?, "\n");
        assert_eq!(render(&ms, "C")?, "\n");

        // Add a "head" to "revive" the graph.
        add(&mut ms, "C", "D")?;
        r(ms.flush())?;
        assert_eq!(
            render(&ms, "D")?,
            r#"
            o  4444444444444444444444444444444444444444 (D)
            â”‚
            o  4343434343434343434343434343434343434343 (C)
            â”‚
            o  4242424242424242424242424242424242424242 (B)
            â”‚
            o  4141414141414141414141414141414141414141 (A)"#
        );

        Ok(())
    }

    #[test]
    fn test_copy_entries() -> Result<()> {
        let dir = TempDir::new("mutationstore")?;
        let mut ms = MutationStore::open(dir.path())?;

        for (pred, succ) in [("P", "E"), ("E", "X")] {
            add(&mut ms, pred, succ)?;
        }
        r(ms.flush())?;

        for (pred, succ) in [("P", "Q"), ("X", "Y")] {
            add(&mut ms, pred, succ)?;
        }

        // Before flush, Q -> Y is not connected.
        assert_eq!(
            render(&ms, "P")?,
            r#"
            o  5959595959595959595959595959595959595959 (Y)
            â”‚
            o  5858585858585858585858585858585858585858 (X)
            â”‚
            o  4545454545454545454545454545454545454545 (E)
            â”‚
            â”‚ o  5151515151515151515151515151515151515151 (Q)
            â”œâ”€â•¯
            o  5050505050505050505050505050505050505050 (P)"#
        );

        // After flush, Q -> Y is auto created.
        r(ms.flush())?;
        assert_eq!(
            render(&ms, "P")?,
            r#"
            o    5959595959595959595959595959595959595959 (Y)
            â”œâ”€â•®
            â”‚ o  5151515151515151515151515151515151515151 (Q)
            â”‚ â”‚
            o â”‚  5858585858585858585858585858585858585858 (X)
            â”‚ â”‚
            o â”‚  4545454545454545454545454545454545454545 (E)
            â”œâ”€â•¯
            o  5050505050505050505050505050505050505050 (P)"#
        );

        Ok(())
    }

    #[test]
    fn test_calculate_obsolete() -> Result<()> {
        let dir = TempDir::new("mutationstore")?;
        let mut ms = MutationStore::open(dir.path())?;

        // C   F  # C -> F
        // |   |
        // B D E  # B -> D -> E
        //  \|/   # X -> Y
        // X Y Z  # Y, Z are public; X, B, E, C are draft

        for (pred, succ) in [("B", "D"), ("D", "E"), ("C", "F"), ("X", "Y")] {
            add(&mut ms, pred, succ)?;
        }
        let public = Set::from_static_names(vec![v("Y"), v("Z")]);
        let draft = Set::from_static_names(vec![v("B"), v("E"), v("C"), v("X")]);
        let obsolete = non_blocking_result(ms.calculate_obsolete(public, draft))?;

        // B is obsoleted. It has a visible successor E (draft).
        assert!(r(obsolete.contains(&v("B")))?);

        // X is obsoleted. It has a visible successor Y (public).
        assert!(r(obsolete.contains(&v("X")))?);

        // C does not have a visible successor.
        assert!(!r(obsolete.contains(&v("C")))?);

        // D is not visible.
        assert!(!r(obsolete.contains(&v("D")))?);

        // A is not a draft.
        assert!(!r(obsolete.contains(&v("A")))?);

        // The set evaluated.
        // (0x42 == b'B', 0x58 == b'X')
        assert_eq!(format!("{:.2?}", &obsolete), "<static [42, 58]>");

        // E does not have a successor.
        assert!(!r(obsolete.contains(&v("E")))?);

        Ok(())
    }

    /// Create a node from a single-char string.
    fn n(s: impl ToString) -> Node {
        Node::from_slice(s.to_string().repeat(Node::len()).as_bytes()).unwrap()
    }

    /// Create a vertex from a single-char string.
    fn v(s: impl ToString) -> VertexName {
        n(s).as_ref().to_vec().into()
    }

    /// Add (test) edges to the mutation store.
    fn add(ms: &mut MutationStore, pred: &str, succ: &str) -> Result<()> {
        ms.add(&MutationEntry {
            preds: vec![n(pred)],
            succ: n(succ),
            split: vec![],
            op: "rewrite".into(),
            user: "test".into(),
            time: 1,
            tz: -7200,
            extra: vec![],
        })
    }

    /// Render the mutation store for the given nodes.
    fn render(ms: &MutationStore, s: &str) -> Result<String> {
        let dag = r(ms.get_dag(s.chars().map(n).collect::<Vec<Node>>()))?;
        dag::render::render_namedag(&dag, |v| Some(format!("({})", v.as_ref()[0] as char)))
    }
}


/*
 * 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.
 */

//! nodemap - A store for node-to-node mappings, with bidirectional indexes.

pub mod nodemap;
pub mod nodeset;

pub use indexedlog::Repair;

pub use crate::nodemap::NodeMap;
pub use crate::nodeset::NodeSet;


/*
 * 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::Range;
use std::path::Path;

use anyhow::Result;
use indexedlog::log;
use indexedlog::log::IndexOutput;
use indexedlog::log::Log;
use indexedlog::DefaultOpenOptions;
use indexedlog::OpenWithRepair;
use thiserror::Error;
use types::errors::KeyError;
use types::node::Node;

#[derive(Debug, Error)]
#[error("Node Map Error: {0:?}")]
struct NodeMapError(String);

impl From<NodeMapError> for KeyError {
    fn from(err: NodeMapError) -> Self {
        KeyError::new(err.into())
    }
}

/// A persistent bidirectional mapping between two Nodes
///
/// [NodeMap] is implemented on top of [indexedlog::log::Log] to store a mapping between two kinds
/// of nodes.
pub struct NodeMap {
    log: Log,
}

impl DefaultOpenOptions<log::OpenOptions> for NodeMap {
    fn default_open_options() -> log::OpenOptions {
        let first_index = |_data: &[u8]| vec![IndexOutput::Reference(0..20)];
        let second_index = |_data: &[u8]| vec![IndexOutput::Reference(20..40)];
        log::OpenOptions::new()
            .create(true)
            .index("first", first_index)
            .index("second", second_index)
    }
}

impl NodeMap {
    pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
        Ok(NodeMap {
            log: Self::default_open_options().open_with_repair(dir.as_ref())?,
        })
    }

    pub fn flush(&mut self) -> Result<()> {
        self.log.flush()?;
        Ok(())
    }

    pub fn add(&mut self, first: &Node, second: &Node) -> Result<()> {
        let mut buf = Vec::with_capacity(40);
        buf.extend_from_slice(first.as_ref());
        buf.extend_from_slice(second.as_ref());
        self.log.append(buf).map_err(|e| e.into())
    }

    pub fn lookup_by_first(&self, first: &Node) -> Result<Option<Node>> {
        self.lookup(first, 0, 20..40)
    }

    pub fn lookup_by_second(&self, second: &Node) -> Result<Option<Node>> {
        self.lookup(second, 1, 0..20)
    }

    fn lookup(&self, key: &Node, index_id: usize, range: Range<usize>) -> Result<Option<Node>> {
        let mut lookup_iter = self.log.lookup(index_id, key)?;
        Ok(match lookup_iter.next() {
            Some(result) => Some(Node::from_slice(&result?[range])?),
            None => None,
        })
    }

    pub fn iter<'a>(&'a self) -> Result<Box<dyn Iterator<Item = Result<(Node, Node)>> + 'a>> {
        let iter = self.log.iter().map(move |entry| match entry {
            Ok(data) => {
                let mut first = self.log.index_func(0, &data)?;
                if first.len() != 1 {
                    return Err(NodeMapError(format!(
                        "invalid index 1 keys in {:?}",
                        self.log.dir
                    ))
                    .into());
                }
                let first = first.pop().unwrap();
                let mut second = self.log.index_func(1, &data)?;
                if second.len() != 1 {
                    return Err(NodeMapError(format!(
                        "invalid index 2 keys in {:?}",
                        self.log.dir
                    ))
                    .into());
                }
                let second = second.pop().unwrap();

                Ok((Node::from_slice(&first)?, Node::from_slice(&second)?))
            }
            Err(e) => Err(e.into()),
        });
        Ok(Box::new(iter))
    }
}

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

    use super::*;

    quickcheck! {
        fn test_roundtrip(pairs: Vec<(Node, Node)>) -> bool {
            let mut pairs = pairs;
            if pairs.len() == 0 {
                return true;
            }

            let dir = TempDir::new().unwrap();
            let mut map = NodeMap::open(dir).unwrap();
            let last = pairs.pop().unwrap();
            for (first, second) in pairs.iter() {
                map.add(&first, &second).unwrap();
            }

            for (first, second) in pairs.iter() {
                if first != &map.lookup_by_second(second).unwrap().unwrap() {
                    return false;
                }
                if second != &map.lookup_by_first(first).unwrap().unwrap() {
                    return false;
                }
            }

            for value in vec![last.0, last.1].iter() {
                if !map.lookup_by_first(value).unwrap().is_none() {
                    return false;
                }
                if !map.lookup_by_second(value).unwrap().is_none() {
                    return false;
                }

            }

            let actual_pairs = map.iter().unwrap().collect::<Result<Vec<_>>>().unwrap();
            actual_pairs == pairs
        }
    }
}


/*
 * 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::path::Path;

use anyhow::Result;
use indexedlog::log;
use indexedlog::log::IndexOutput;
use indexedlog::log::Log;
use indexedlog::DefaultOpenOptions;
use indexedlog::OpenWithRepair;
use thiserror::Error;
use types::errors::KeyError;
use types::node::Node;

#[derive(Debug, Error)]
#[error("Node Set Error: {0:?}")]
struct NodeSetError(String);

impl From<NodeSetError> for KeyError {
    fn from(err: NodeSetError) -> Self {
        KeyError::new(err.into())
    }
}

/// A persistent set of Nodes.
///
/// [NodeSet] is implemented on top of [indexedlog::log::Log] to store
/// a set of nodes. Its insertion and lookup complexity are `O(log N)`.
pub struct NodeSet {
    log: Log,
}

impl DefaultOpenOptions<log::OpenOptions> for NodeSet {
    fn default_open_options() -> log::OpenOptions {
        let node_index = |_data: &[u8]| vec![IndexOutput::Reference(0..Node::len() as u64)];
        log::OpenOptions::new()
            .create(true)
            .index("node", node_index)
    }
}

impl NodeSet {
    const INDEX_NODE: usize = 0;

    pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
        Ok(NodeSet {
            log: Self::default_open_options().open_with_repair(dir.as_ref())?,
        })
    }

    pub fn flush(&mut self) -> Result<()> {
        self.log.flush()?;
        Ok(())
    }

    pub fn add(&mut self, node: &Node) -> Result<()> {
        if !self.contains(node)? {
            self.log.append(node.as_ref())?;
        }
        Ok(())
    }

    pub fn contains(&self, node: &Node) -> Result<bool> {
        let mut lookup_iter = self.log.lookup(Self::INDEX_NODE, node.as_ref())?;
        Ok(lookup_iter.next().is_some())
    }

    pub fn iter<'a>(&'a self) -> impl Iterator<Item = Result<Node>> + 'a {
        self.log
            .iter()
            .map(|slice| Node::from_slice(slice?).map_err(Into::into))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use quickcheck::quickcheck;
    use tempfile::TempDir;

    use super::*;

    quickcheck! {
        fn test_compare_with_hashset(nodes: HashSet<Node>) -> bool {
            let dir = TempDir::new().unwrap();
            let mut set = NodeSet::open(dir).unwrap();

            for node in &nodes {
                assert!(!set.contains(node).unwrap());
                set.add(node).unwrap();
                assert!(set.contains(node).unwrap());
            }

            let nodes2: HashSet<Node> = set.iter().map(|node| node.unwrap()).collect();
            nodes2 == nodes
        }
    }
}


/*
 * 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.
 */

//! Turn an async function to a sync non-blocking function.

use std::future::Future;
use std::io;
use std::task::Context;
use std::task::Poll;
use std::task::RawWaker;
use std::task::RawWakerVTable;
use std::task::Waker;

/// Attempt to resolve a `future` without blocking.
/// Return `WouldBlock` error if the future will block.
/// Return the resolved value otherwise.
pub fn non_blocking<F, R>(future: F) -> io::Result<R>
where
    F: Future<Output = R>,
{
    let waker = waker();
    let mut cx = Context::from_waker(&waker);
    let mut future = Box::pin(future);
    match future.as_mut().poll(&mut cx) {
        Poll::Ready(result) => Ok(result),
        Poll::Pending => Err(io::ErrorKind::WouldBlock.into()),
    }
}

/// Similar to `non_blocking`, but unwraps a level of `Result`.
pub fn non_blocking_result<F, T, E>(future: F) -> Result<T, E>
where
    F: Future<Output = Result<T, E>>,
    E: From<io::Error>,
{
    non_blocking(future)?
}

fn waker() -> Waker {
    let raw_waker = clone(std::ptr::null());
    unsafe { Waker::from_raw(raw_waker) }
}

fn vtable() -> &'static RawWakerVTable {
    &RawWakerVTable::new(clone, noop, noop, noop)
}

fn clone(data: *const ()) -> RawWaker {
    RawWaker::new(data, vtable())
}

fn noop(_: *const ()) {}

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

    #[test]
    fn test_non_blocking_ok() {
        async fn f() -> usize {
            g().await + 4
        }

        async fn g() -> usize {
            3
        }

        async fn h() -> io::Result<usize> {
            Ok(5)
        }

        assert_eq!(non_blocking(async { f().await }).unwrap(), 7);
        assert_eq!(non_blocking_result(h()).unwrap(), 5);
    }

    #[test]
    fn test_non_blocking_err() {
        let (sender, receiver) = futures::channel::oneshot::channel::<usize>();
        assert_eq!(
            non_blocking(async { receiver.await }).unwrap_err().kind(),
            io::ErrorKind::WouldBlock
        );
        drop(sender);
    }
}


/*
 * 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.
 */

//! # pathhistory
//!
//! This crate provides file or directory history algorithms that does not
//! depend on per-path indexes, is better than scanning commits one by one,
//! and is friendly for lazy stores.
//!
//! The basic idea is to use the segment struct from the `dag` crate and
//! (aggressively) skip large chunks of commits without visiting all commits
//! one by one.
//!
//! This might miss some "change then revert" commits but practically it
//! might be good enough.
//!
//! See `PathHistory` for the main structure.

mod pathhistory;
mod pathops;
mod renametrace;
mod utils;

#[cfg(test)]
mod tests;

pub use crate::pathhistory::PathHistory;
pub use crate::renametrace::RenameTracer;

#[cfg(test)]
dev_logger::init!();


/*
 * 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::VecDeque;
use std::sync::Arc;

use anyhow::bail;
use anyhow::Result;
use dag::ops::DagAlgorithm;
use dag::ops::IdConvert;
use dag::ops::IdDagAlgorithm;
use dag::Id;
use dag::IdSegment;
use dag::IdSet;
use dag::IdSpan;
use dag::Set;
use dag::Vertex;
use manifest_tree::TreeStore;
use storemodel::ReadRootTreeIds;
use types::HgId;
use types::RepoPathBuf;

use crate::pathops::CompiledPaths;
use crate::pathops::ContentId;
use crate::utils;

/// State for answering history questions about paths.
pub struct PathHistory {
    // Input
    //
    /// Resolve commit ids to trees in batch.
    root_tree_reader: Arc<dyn ReadRootTreeIds + Send + Sync>,

    /// Resolve and prefetch trees in batch.
    tree_store: Arc<dyn TreeStore + Send + Sync>,

    // Derived from input
    //
    /// Commit graph algorithms. Derived from `set`.
    id_dag: Arc<dyn IdDagAlgorithm + Send + Sync>,

    /// IdSet derived from `set`.
    id_set: IdSet,

    /// Resolve commit ids to hashes in batch. Derived from `set`.
    id_map: Arc<dyn IdConvert + Send + Sync>,

    /// Root commits. Sorted in ASC order. Derived from `set`.
    roots: Vec<Id>,

    /// Operations to resolve content ids of `paths` starting
    /// from a root tree. Derived from `paths`.
    compiled_paths: CompiledPaths,

    // States
    //
    /// Segments to visit. They are sorted in DESC (output) order.
    queue: VecDeque<IdSegment>,

    /// Cached commit states. Include the commit hash, and content id.
    commit_map: HashMap<Id, CommitState>,

    /// Partial output. Buffer of commits with at least one path changed.
    output_buffer: VecDeque<Vertex>,

    /// Commits that are decided to be skipped.
    skipped: IdSet,

    /// Statistics for reasoning about performance.
    stats: PathHistoryStats,
}

#[derive(Debug, Default)]
pub struct PathHistoryStats {
    pub commit_count: usize,
    pub split_count: usize,
}

/// Partial information of a commit.
/// Track what's partially known.
/// Used to check paths in two commits are different.
#[derive(Debug)]
struct CommitState {
    /// Content of this commit for comparison.
    content_id: ContentId,

    /// Commit id in the IdDag.
    id: Id,

    /// Commit hash.
    vertex: Vertex,
}

impl PathHistory {
    /// Prepare context to search changes of `paths` in commit `set`.
    ///
    /// The search is lazy. Call `next()` to get the next commit.
    /// Commits are emitted in topo order, latest first.
    pub async fn new(
        set: Set,
        paths: Vec<RepoPathBuf>,
        root_tree_reader: Arc<dyn ReadRootTreeIds + Send + Sync>,
        tree_store: Arc<dyn TreeStore + Send + Sync>,
    ) -> Result<Self> {
        Self::new_internal(set, paths, root_tree_reader, tree_store, false).await
    }

    /// Internal version of `PathHistory::new` method.
    ///
    /// This creation method adds an additional `ignore_file_content` parameter.
    /// When enabled, file content and file type (e.g. executable) will be ignored,
    /// this is used for tracing renames.
    pub(crate) async fn new_internal(
        set: Set,
        paths: Vec<RepoPathBuf>,
        root_tree_reader: Arc<dyn ReadRootTreeIds + Send + Sync>,
        tree_store: Arc<dyn TreeStore + Send + Sync>,
        ignore_file_content: bool,
    ) -> Result<Self> {
        tracing::debug!("PathHistory::new(set={:.12?}, paths={:?})", &set, &paths);
        let (id_set, id_map) = set_to_id_set_and_id_map(&set)?;
        let (roots, id_dag) = if let Some(dag) = set.dag() {
            // "Real" roots that have no parents. That is:
            // roots(set + parents(set)) & set
            let id_dag = dag.id_dag_snapshot()?;
            let parents = id_dag.parents(id_set.clone())?;
            let union = parents.union(&id_set);
            let roots = id_dag.roots(union)?;
            let roots = roots.intersection(&id_set);
            (roots, id_dag)
        } else {
            bail!("PathHistory requires Dag associated with `set`");
        };
        let roots: Vec<Id> = roots.iter_asc().collect();
        tracing::trace!(" roots: {:?}", &roots);
        let queue = id_dag.id_set_to_id_segments(&id_set)?;
        let compiled_paths =
            CompiledPaths::compile(paths).with_ignore_file_content(ignore_file_content);
        let mut path_history = Self {
            id_dag,
            id_map,
            id_set,
            root_tree_reader,
            tree_store,
            roots,
            compiled_paths,
            queue,
            commit_map: Default::default(),
            output_buffer: Default::default(),
            skipped: Default::default(),
            stats: Default::default(),
        };

        // Prefetch commits in segments, and roots.
        // This takes longer in a repo with lots of merges (ex. 0.2+s in linux.git).
        // In a more linear repo the overhead is smaller (ex. <0.1s).
        // PERF: It might be interesting to skip ancestors of root trees that do not
        // have the paths.
        let segments = path_history.queue.clone();
        let mut commit_ids = path_history.segments_to_prefetch_commit_ids(&mut segments.iter());
        commit_ids.extend_from_slice(&path_history.roots);
        path_history.populate_commit_states(&commit_ids).await?;

        Ok(path_history)
    }

    /// Obtain statistics for performance analysis.
    pub fn stats(&self) -> &PathHistoryStats {
        &self.stats
    }

    /// Find the next vertex that changes at least one of the specified paths.
    /// Output in DESC topo order defined by the `IdDag`.
    pub async fn next(&mut self) -> Result<Option<Vertex>> {
        loop {
            if let Some(v) = self.output_buffer.pop_front() {
                return Ok(Some(v));
            } else if self.is_done() {
                return Ok(None);
            } else {
                self.step().await?
            }
        }
    }

    /// Whether the history operation is completed.
    fn is_done(&self) -> bool {
        self.queue.is_empty()
    }

    /// Attempt to make some progress for `next()`.
    async fn step(&mut self) -> Result<()> {
        let seg = match self.queue.pop_front() {
            None => return Ok(()),
            Some(seg) => seg,
        };

        let mut action = Action::Undecided;
        let mut different_parents: Vec<Id> = Vec::new();
        let mut same_parents: Vec<Id> = Vec::new();
        let mut outside_same_parents: Vec<Id> = Vec::new();
        let span = IdSpan::new(seg.low, seg.high);

        if self.skipped.contains(span) {
            action = Action::Skip;
        } else {
            let head_content_id = self.commit_id_to_content_id(seg.high);
            for &id in &seg.parents {
                if self.commit_id_to_content_id(id) == head_content_id {
                    if !self.id_set.contains(id) {
                        // Parent is outside the requested set. Do not use its content
                        // to skip segments.
                        outside_same_parents.push(id);
                    } else {
                        same_parents.push(id)
                    }
                } else {
                    different_parents.push(id)
                }
            }

            // If there are "same parents" (content id is the same with
            // the segment head), ignore (::different_parents - ::same_parents)
            //
            //               head
            //              /    \
            //             :      :
            //   same_parent      different_parent <- segment parents
            //
            // The idea is that "same_parents" are "picked" by merges to reach
            // "head", "different_parents" are "abandoned" by the merges, and
            // not interesting.
            //
            // This is an important optimization for mergy repos.
            if action.is_undecided() && !same_parents.is_empty() {
                action = Action::OutputRoots;
                // Mark "::different_parents - ::same_parents" for skipping.
                // This also avoids showing commits that got reverted in the
                // current segment, which can cause confusion.
                if !different_parents.is_empty() {
                    let ancestors =
                        |ids: &[Id]| self.id_dag.ancestors(IdSet::from_spans(ids.to_vec()));
                    let diff_ancestors = ancestors(&different_parents)?;
                    let same_ancestors = ancestors(&same_parents)?;
                    let new_skip = diff_ancestors.difference(&same_ancestors);
                    if !new_skip.is_empty() {
                        tracing::trace!("   skip += [{:?}]", &new_skip);
                        self.skipped = self.skipped.union(&new_skip);
                    }
                }
            }

            // Single commit. Need to make final decision now.
            if action.is_undecided() && seg.low == seg.high {
                let changed = if seg.parents.is_empty() {
                    // A global root (roots(dag.all())). Check if it contains any of the "paths".
                    !head_content_id.is_empty()
                } else {
                    // Not a global root. Check all parents are different.
                    different_parents.len() == seg.parents.len()
                };
                if changed {
                    action = Action::OutputHigh;
                } else {
                    action = Action::OutputRoots;
                }
            }
        }

        // Take action.
        tracing::debug!("  {:11?} {:?}", &action, &seg);
        match action {
            Action::Skip => {}
            Action::OutputRoots => self.output_roots(seg),
            Action::OutputHigh => self.push_output(seg.high),
            Action::Undecided => {
                let interesting_parents = different_parents.into_iter().chain(outside_same_parents);
                self.split_segment(seg, IdSet::from_spans(interesting_parents))
                    .await?
            }
        }
        Ok(())
    }

    /// Output roots in the given segment.
    fn output_roots(&mut self, seg: IdSegment) {
        while let Some(&root) = self.roots.last() {
            if root < seg.low {
                break;
            }
            let _ = self.roots.pop();
            if root <= seg.high {
                let root_content_id = self.commit_id_to_content_id(root);
                if !root_content_id.is_empty() {
                    self.push_output(root);
                }
            }
        }
    }

    /// Split a segment into smaller pieces so `step()` might make progress.
    ///
    /// If a segment does not have parent (or grand parent) in
    /// `interesting_ancestors`, the segment will be skipped.
    /// `interesting_ancestors` is usually initially `different_parents`
    /// calculated in `step`.
    async fn split_segment(
        &mut self,
        seg: IdSegment,
        mut interesting_ancestors: IdSet,
    ) -> Result<()> {
        self.stats.split_count += 1;
        let mut new_segs = Vec::new();
        if seg.level > 0 {
            // Split a high level segment into lower level segments.
            let sub_segs = self.id_dag.id_segment_to_lower_level_id_segments(&seg)?;
            for seg in sub_segs.into_iter().rev() {
                if !seg.has_root
                    && seg
                        .parents
                        .iter()
                        .all(|&p| !interesting_ancestors.contains(p))
                {
                    // Skip the segment. It does not contain "interesting" parents.
                    continue;
                } else {
                    interesting_ancestors =
                        interesting_ancestors.union(&IdSet::from(IdSpan::new(seg.low, seg.high)));
                    new_segs.push(seg);
                }
            }
        } else {
            // Split a flat segment into two halves.
            let mid = self.pick_mid_id(seg.low, seg.high);
            new_segs.push(IdSegment {
                high: mid,
                low: seg.low,
                parents: seg.parents,
                has_root: seg.has_root,
                level: 0,
            });
            assert!(mid < seg.high);
            new_segs.push(IdSegment {
                high: seg.high,
                low: mid + 1,
                parents: vec![mid],
                has_root: false,
                level: 0,
            });
        }

        let commit_ids = self.segments_to_prefetch_commit_ids(&mut new_segs.iter());
        self.populate_commit_states(&commit_ids).await?;
        tracing::trace!("   split to {} subsegs", new_segs.len());
        for seg in new_segs {
            self.queue.push_front(seg);
        }
        Ok(())
    }

    /// Resolve commit id to content id for easier comparison.
    /// Callsite is responsible for prefeching the commit ids.
    /// Practically, prefetching is done by code changing `self.queue`:
    /// `new` and `split_segment`.
    fn commit_id_to_content_id(&mut self, id: Id) -> ContentId {
        self.commit_map[&id].content_id
    }

    /// Given segments. Find commit ids to prefetch (high + parents).
    fn segments_to_prefetch_commit_ids(
        &mut self,
        segments: &mut dyn Iterator<Item = &IdSegment>,
    ) -> Vec<Id> {
        let mut commit_ids = Vec::new();
        for seg in segments {
            commit_ids.push(seg.high);
            commit_ids.extend_from_slice(&seg.parents);
        }
        commit_ids.sort_unstable();
        commit_ids.dedup();
        commit_ids
    }

    /// Push the commit to output. The commit is confirmed changed from its
    /// parent, or is a root with paths.
    fn push_output(&mut self, id: Id) {
        if !self.skipped.contains(id) {
            let state = &self.commit_map[&id];
            tracing::trace!("   output: {:?} ({:?})", id, &state.vertex);
            self.output_buffer.push_back(state.vertex.clone());
        }
    }

    /// Pick the bisect mid point between low and high.
    /// If low == high - 1, return low.
    fn pick_mid_id(&self, low: Id, high: Id) -> Id {
        let mid = utils::pick_mid(low.0, high.0).unwrap_or(low.0);
        Id(mid)
    }

    /// Insert initial commit states to `commit_map` if they do not already exist.
    async fn populate_commit_states(&mut self, ids: &[Id]) -> Result<()> {
        let missing_ids = ids
            .iter()
            .cloned()
            .filter(|i| !self.commit_map.contains_key(i))
            .collect::<Vec<_>>();
        if !missing_ids.is_empty() {
            for state in self.prepare_commit_states(&missing_ids).await? {
                self.stats.commit_count += 1;
                self.commit_map.insert(state.id, state);
            }
        }
        Ok(())
    }

    /// Prepare initial `CommitState`s for the given indexes.
    /// The initial state includes commit id, hash, root tree id.
    /// Root trees are prefetched.
    async fn prepare_commit_states(&mut self, ids: &[Id]) -> Result<Vec<CommitState>> {
        // Index => Commit id => Commit hash (vertex)
        // (about 160 commits per millisecond)
        tracing::debug!("  prepare_commit_states for {} ids", ids.len());
        let vertexes: Vec<Vertex> = self
            .id_map
            .vertex_name_batch(ids)
            .await?
            .into_iter()
            .collect::<std::result::Result<Vec<Vertex>, _>>()?;

        // Commit hash => Root tree id
        // (about 26 commits per millisecond)
        let commit_hgids: Vec<HgId> = vertexes
            .iter()
            .map(|v| HgId::from_slice(v.as_ref()))
            .collect::<std::result::Result<Vec<HgId>, _>>()?;
        let commit_to_tree_id_map: HashMap<HgId, HgId> = self
            .root_tree_reader
            .read_root_tree_ids(commit_hgids.clone())
            .await?
            .into_iter()
            .collect();
        let root_tree_ids: Vec<HgId> = commit_hgids
            .iter()
            .map(|i| commit_to_tree_id_map[i])
            .collect();

        // Root tree id => Content Id.
        // (about 14 commits per millisecond, or slower for deeper paths)
        let content_ids = self.resolve_root_tree_to_content_id(root_tree_ids).await?;

        // Construct the commit states
        let states = ids
            .iter()
            .zip(vertexes)
            .zip(content_ids)
            .map(|((id, vertex), content_id)| CommitState {
                id: *id,
                vertex,
                content_id,
            })
            .collect();

        Ok(states)
    }

    /// Resolve trees recursively. Turn them into `ContentId` for easier comparison.
    async fn resolve_root_tree_to_content_id(
        &mut self,
        root_tree_ids: Vec<HgId>,
    ) -> Result<Vec<ContentId>> {
        let tree_store = self.tree_store.clone();
        self.compiled_paths.execute(root_tree_ids, tree_store).await
    }
}

/// Used by `step`.
#[derive(Debug, Copy, Clone)]
enum Action {
    /// Not yet decided. Needs to split the segment for better decision.
    Undecided,
    /// Skip the segment. Output nothing in the segment (not even roots).
    Skip,
    /// Skip the segment. Output roots in the segment.
    OutputRoots,
    /// Skip the segment. Output `high`. The segment length should be 1.
    OutputHigh,
}

impl Action {
    fn is_undecided(&self) -> bool {
        match self {
            Action::Undecided => true,
            _ => false,
        }
    }
}

fn set_to_id_set_and_id_map(set: &Set) -> Result<(IdSet, Arc<dyn IdConvert + Send + Sync>)> {
    match set.to_id_set_and_id_map_in_o1() {
        Some(v) => Ok(v),
        None => bail!(
            "PathHistory requires {:?} to convert to IdSet and IdMap in O(1)",
            set
        ),
    }
}