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

//! Unpacking changegroups.
//! See <https://www.mercurial-scm.org/repo/hg/file/@/mercurial/help/internals/changegroups.txt>

use std::cmp;
use std::mem;
use std::str::FromStr;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bytes_old::BytesMut;
use mercurial_types::MPath;
use mercurial_types::RevFlags;
use slog::Logger;
use tokio_io::codec::Decoder;

use super::CgDeltaChunk;
use super::Part;
use super::Section;
use crate::delta;
use crate::errors::ErrorKind;
use crate::utils::BytesExt;

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum CgVersion {
    Cg2Version,
    Cg3Version,
}

impl CgVersion {
    pub fn to_str(&self) -> &str {
        match self {
            CgVersion::Cg2Version => "02",
            CgVersion::Cg3Version => "03",
        }
    }
}

impl FromStr for CgVersion {
    type Err = Error;

    fn from_str(s: &str) -> Result<CgVersion> {
        match s {
            "02" => Ok(CgVersion::Cg2Version),
            "03" => Ok(CgVersion::Cg3Version),
            bad => Err(ErrorKind::CgDecode(format!(
                "Non supported Cg version in Part Header {}",
                bad
            ))
            .into()),
        }
    }
}

// See the chunk header definition below for the first 100 bytes. The last 4 is
// for the length field itself. _CHANGEGROUPV2_DELTA_HEADER = "20s20s20s20s20s"
const CHUNK_HEADER2_LEN: usize = 20 + 20 + 20 + 20 + 20 + 4;

// See the chunk header definition below for the first 102 bytes. The last 4 is
// for the length field itself. _CHANGEGROUPV3_DELTA_HEADER = ">20s20s20s20s20sH"
const CHUNK_HEADER3_LEN: usize = 20 + 20 + 20 + 20 + 20 + 2 + 4;

#[derive(Debug)]
pub struct CgUnpacker {
    #[allow(dead_code)]
    logger: Logger,
    state: State,
    version: CgVersion,
}

impl Part {
    /// Gets the section this part is from.
    ///
    /// # Panics
    ///
    /// When self does not contain a Section.
    pub fn section(&self) -> &Section {
        match *self {
            Part::CgChunk(ref section, _) => section,
            Part::SectionEnd(ref section) => section,
            _ => panic!("this Part does not contain a Section"),
        }
    }

    /// Gets the chunk inside this part.
    ///
    /// # Panics
    ///
    /// When self does not contain a CgDeltaChunk (it is a SectionEnd or End).
    pub fn chunk(&self) -> &CgDeltaChunk {
        match self {
            Part::CgChunk(_, chunk) => chunk,
            _ => panic!("this Part does not contain a CgDeltaChunk"),
        }
    }
}

impl Decoder for CgUnpacker {
    type Item = Part;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        match Self::decode_next(buf, self.state.take(), &self.version) {
            Err(e) => {
                self.state = State::Invalid;
                Err(e)
            }
            Ok((ret, state)) => {
                self.state = state;
                match ret {
                    None => Ok(None),
                    Some(v) => Ok(Some(v)),
                }
            }
        }
    }

    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        match self.decode(buf)? {
            None => {
                if !buf.is_empty() {
                    let len = buf.len();
                    let bytes = &buf[..cmp::min(len, 128)];
                    let msg = format!(
                        "incomplete changegroup: {} bytes remaining in \
                         buffer. State: {:?}, First 128 bytes: {:?}",
                        len, self.state, bytes,
                    );
                    bail!(ErrorKind::CgDecode(msg));
                }
                if self.state != State::End {
                    let msg = format!(
                        "incomplete changegroup: expected state End, found {:?}",
                        self.state
                    );
                    bail!(ErrorKind::CgDecode(msg));
                }
                Ok(None)
            }
            Some(v) => Ok(Some(v)),
        }
    }
}

impl CgUnpacker {
    pub fn new(logger: Logger, version: CgVersion) -> Self {
        CgUnpacker {
            logger,
            state: State::Changeset,
            version,
        }
    }

    fn chunk_header_len(version: &CgVersion) -> usize {
        match version {
            CgVersion::Cg2Version => CHUNK_HEADER2_LEN,
            CgVersion::Cg3Version => CHUNK_HEADER3_LEN,
        }
    }

    fn decode_next(
        buf: &mut BytesMut,
        state: State,
        version: &CgVersion,
    ) -> Result<(Option<Part>, State)> {
        match state {
            State::Changeset => match Self::decode_chunk(buf, version)? {
                None => Ok((None, State::Changeset)),
                Some(CgChunk::Empty) => {
                    Ok((Some(Part::SectionEnd(Section::Changeset)), State::Manifest))
                }
                Some(CgChunk::Delta(chunk)) => Ok((
                    Some(Part::CgChunk(Section::Changeset, chunk)),
                    State::Changeset,
                )),
            },
            State::Manifest => match Self::decode_chunk(buf, version)? {
                None => Ok((None, State::Manifest)),
                Some(CgChunk::Empty) => {
                    let next_state = match version {
                        CgVersion::Cg2Version => State::Filename,
                        CgVersion::Cg3Version => State::Treemanifest,
                    };
                    Ok((Some(Part::SectionEnd(Section::Manifest)), next_state))
                }
                Some(CgChunk::Delta(chunk)) => Ok((
                    Some(Part::CgChunk(Section::Manifest, chunk)),
                    State::Manifest,
                )),
            },
            State::Treemanifest => match Self::decode_chunk(buf, version)? {
                None => Ok((None, State::Treemanifest)),
                Some(CgChunk::Empty) => Ok((
                    Some(Part::SectionEnd(Section::Treemanifest)),
                    State::Filename,
                )),
                Some(CgChunk::Delta(_)) => {
                    Err(ErrorKind::CgDecode("Empty TreeManifest has expected".into()).into())
                }
            },
            State::Filename => {
                let filename = Self::decode_filename(buf)?;
                match filename {
                    DecodeRes::None => Ok((None, State::Filename)),
                    DecodeRes::Some(f) => Self::decode_filelog_chunk(buf, f, version),
                    DecodeRes::End => Ok((Some(Part::End), State::End)),
                }
            }
            State::Filelog(filename) => Self::decode_filelog_chunk(buf, filename, version),
            State::End => Ok((None, State::End)),
            State::Invalid => Err(ErrorKind::CgDecode("byte stream corrupt".into()).into()),
        }
    }

    fn decode_filelog_chunk(
        buf: &mut BytesMut,
        f: MPath,
        version: &CgVersion,
    ) -> Result<(Option<Part>, State)> {
        match Self::decode_chunk(buf, version)? {
            None => Ok((None, State::Filelog(f))),
            Some(CgChunk::Empty) => {
                Ok((Some(Part::SectionEnd(Section::Filelog(f))), State::Filename))
            }
            Some(CgChunk::Delta(chunk)) => Ok((
                Some(Part::CgChunk(Section::Filelog(f.clone()), chunk)),
                State::Filelog(f),
            )),
        }
    }

    fn decode_chunk(buf: &mut BytesMut, version: &CgVersion) -> Result<Option<CgChunk>> {
        if buf.len() < 4 {
            return Ok(None);
        }

        let chunk_len = buf.peek_i32();
        // Note that chunk_len includes the 4 bytes consumed by itself
        // TODO: chunk_len < 0 = error
        let chunk_len = chunk_len as usize;
        if chunk_len == 0 {
            let _ = buf.drain_i32();
            return Ok(Some(CgChunk::Empty));
        }
        if chunk_len < Self::chunk_header_len(version) {
            let msg = format!(
                "invalid chunk for version {:?}: length >= {} required, found {}",
                version,
                Self::chunk_header_len(version),
                chunk_len,
            );
            bail!(ErrorKind::CgDecode(msg));
        }

        if buf.len() < chunk_len {
            return Ok(None);
        }
        Self::decode_delta(buf, chunk_len, version)
    }

    fn decode_delta(
        buf: &mut BytesMut,
        chunk_len: usize,
        version: &CgVersion,
    ) -> Result<Option<CgChunk>> {
        let _ = buf.drain_i32();

        // A chunk header has:
        // ---
        // node: HgNodeHash (20 bytes)
        // p1: HgNodeHash (20 bytes)
        // p2: HgNodeHash (20 bytes) -- NULL_HASH if only 1 parent
        // base node: HgNodeHash (20 bytes) (new in changegroup2)
        // link node: HgNodeHash (20 bytes)
        // flags: unsigned short (2 bytes) -- (version 3 only)
        // ---

        let node = buf.drain_node();
        let p1 = buf.drain_node();
        let p2 = buf.drain_node();
        let base = buf.drain_node();
        let linknode = buf.drain_node();
        let flags = match version {
            CgVersion::Cg2Version => None,
            CgVersion::Cg3Version => {
                let bits = buf.drain_u16();
                let flags = RevFlags::from_bits(bits)
                    .ok_or_else(|| format_err!("unknown revlog flags: {}", bits))?;
                Some(flags)
            }
        };

        let delta = delta::decode_delta(buf.split_to(chunk_len - Self::chunk_header_len(version)))?;

        Ok(Some(CgChunk::Delta(CgDeltaChunk {
            node,
            p1,
            p2,
            base,
            linknode,
            delta,
            flags,
        })))
    }

    fn decode_filename(buf: &mut BytesMut) -> Result<DecodeRes<MPath>> {
        if buf.len() < 4 {
            return Ok(DecodeRes::None);
        }
        let filename_len = buf.peek_i32();
        // TODO: filename_len < 0 == error
        if filename_len == 0 {
            let _ = buf.split_to(4);
            return Ok(DecodeRes::End);
        }
        let filename_len = filename_len as usize;
        // filename_len includes the 4 bytes for the length field.
        if buf.len() < filename_len {
            return Ok(DecodeRes::None);
        }
        let _ = buf.split_to(4);
        let filename = buf.drain_path(filename_len - 4).with_context(|| {
            let msg = format!("invalid filename of length {}", filename_len);
            ErrorKind::CgDecode(msg)
        })?;
        Ok(DecodeRes::Some(filename))
    }
}

enum DecodeRes<T> {
    None,
    Some(T),
    End,
}

enum CgChunk {
    Delta(CgDeltaChunk),
    Empty,
}

#[derive(Debug, Eq, PartialEq)]
enum State {
    Changeset,
    Manifest,
    Treemanifest,
    Filename,
    Filelog(MPath),
    End,
    Invalid,
}

impl State {
    pub fn take(&mut self) -> Self {
        mem::replace(self, State::Invalid)
    }
}


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

// Codecs related to infinitepush also known as Commit Cloud.

use std::io::Cursor;

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use mercurial_mutation::HgMutationEntry;
use tokio_io::codec::Decoder;
use types::mutation::MutationEntry;
use vlqencoding::VLQDecode;
use vlqencoding::VLQEncode;

use crate::utils::BytesExt;

#[derive(Debug)]
pub struct InfinitepushBookmarksUnpacker {
    finished: bool,
    expected_len: Option<usize>,
}

impl InfinitepushBookmarksUnpacker {
    pub fn new() -> Self {
        Self {
            finished: false,
            expected_len: None,
        }
    }
}

impl Decoder for InfinitepushBookmarksUnpacker {
    type Item = Bytes;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        if self.finished {
            return Ok(None);
        }
        match self.expected_len {
            Some(toread) => {
                if buf.len() < toread {
                    Ok(None)
                } else {
                    self.finished = true;
                    Ok(Some(buf.split_to(toread).freeze()))
                }
            }
            None => {
                if buf.len() >= 4 {
                    self.expected_len = Some(buf.drain_u32() as usize);
                }
                Ok(None)
            }
        }
    }
}

#[derive(Debug)]
pub struct InfinitepushMutationUnpacker {}

impl InfinitepushMutationUnpacker {
    pub fn new() -> Self {
        Self {}
    }
}

const MUTATION_PART_VERSION: u8 = 1;

/// Decoder for infinitepush mutation entries
///
/// This decoder decodes all entries in one operation, so needs to wait for eof.
impl Decoder for InfinitepushMutationUnpacker {
    type Item = Vec<HgMutationEntry>;
    type Error = Error;

    fn decode(&mut self, _buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        Ok(None)
    }

    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        let mut entries = Vec::new();
        let mut cursor = Cursor::new(buf);
        let version = cursor.read_u8()?;
        if version != MUTATION_PART_VERSION {
            bail!("Unsupported infinitepush mutation part format: {}", version);
        }
        let count = cursor.read_vlq()?;
        entries.reserve_exact(count);
        for _ in 0..count {
            let entry = HgMutationEntry::deserialize(&mut cursor)?;
            entries.push(entry);
        }
        let size = cursor.position();
        cursor.into_inner().advance(size as usize);
        Ok(Some(entries))
    }
}

pub fn infinitepush_mutation_packer(entries: Vec<HgMutationEntry>) -> Result<Bytes> {
    let mut buf = Vec::with_capacity(entries.len() * types::mutation::DEFAULT_ENTRY_SIZE);
    buf.write_u8(MUTATION_PART_VERSION)?;
    buf.write_vlq(entries.len())?;
    for entry in entries {
        let entry: MutationEntry = entry.into();
        entry.serialize(&mut buf)?;
    }
    Ok(buf.into())
}


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

pub mod packer;

#[cfg(test)]
use quickcheck_arbitrary_derive::Arbitrary;

#[derive(Debug, Clone)]
#[cfg_attr(test, derive(Arbitrary))]
pub struct MetadataEntry {
    key: String,
    value: String,
}

impl MetadataEntry {
    pub fn new(key: impl Into<String>, value: impl Into<String>) -> Self {
        Self {
            key: key.into(),
            value: value.into(),
        }
    }
}


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

//! Packing obsmarkers to be sent after e.g. a pushrebase
//! Format documentation: <https://www.mercurial-scm.org/repo/hg/file/tip/mercurial/obsolete.py>

use anyhow::Error;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::BufMut;
use futures_old::stream::iter_result;
use futures_old::Stream;
use mercurial_types::HgChangesetId;
use mononoke_types::DateTime;

use super::MetadataEntry;
use crate::chunk::Chunk;

const VERSION: u8 = 1;

pub fn obsmarkers_packer_stream(
    pairs_stream: impl Stream<Item = (HgChangesetId, Vec<HgChangesetId>), Error = Error>,
    time: DateTime,
    metadata: Vec<MetadataEntry>,
) -> impl Stream<Item = Chunk, Error = Error> {
    let version_chunk = Chunk::new(vec![VERSION]);

    let chunks_stream = pairs_stream.and_then(move |(predecessor, successors)| {
        prepare_obsmarker_chunk(&predecessor, &successors, &time, &metadata)
    });

    iter_result(vec![version_chunk]).chain(chunks_stream)
}

fn prepare_obsmarker_chunk(
    predecessor: &HgChangesetId,
    successors: &Vec<HgChangesetId>,
    time: &DateTime,
    metadata: &Vec<MetadataEntry>,
) -> Result<Chunk> {
    // Reserve space, fill it with zeros before writing out our data.
    let mut v: Vec<u8> = vec![0; 19];

    // 0: length: uint32. We'll write it out later.

    // 4: seconds since epoch: f64
    BigEndian::write_f64(&mut v[4..], time.timestamp_secs() as f64);

    // 12: timezone offset in minutes: int16
    BigEndian::write_i16(&mut v[12..], time.tz_offset_minutes());

    // 14: flags: uint16. We don't have any.

    // 16: number of successors: uint8.
    let new_len = u8::try_from(successors.len())?;
    v[16] = new_len;

    // 17: number of parents: uint8. We set 3, for "no parent data stored".
    v[17] = 3;

    // 18: Number of metadata entries: uint8.
    let metadata_len = u8::try_from(metadata.len())?;
    v[18] = metadata_len;

    // predecessor changeset identifier (20 bytes)
    v.put_slice(predecessor.as_ref());

    // successors changeset identifiers (20 bytes, each)
    for succ in successors.iter() {
        v.put_slice(succ.as_ref());
    }

    let metadata_bytes: Vec<&[u8]> = metadata
        .iter()
        .flat_map(|entry| vec![entry.key.as_bytes(), entry.value.as_bytes()])
        .collect();

    // Metadata sizes, uint8 each
    for bytes in metadata_bytes.iter() {
        v.put_u8(u8::try_from(bytes.len())?);
    }

    // Metadata keys and values, variable size
    for bytes in metadata_bytes.iter() {
        v.put_slice(bytes);
    }

    // And now, we write out the final length into the first 4 bytes we reserved.
    let len = v.len();
    BigEndian::write_u32(&mut v[0..], u32::try_from(len)?);

    Chunk::new(v)
}

#[cfg(test)]
mod test {
    use anyhow::Error;
    use futures_ext::StreamExt;
    use futures_old::stream;
    use futures_old::Async;
    use futures_old::Poll;
    use mercurial_types_mocks::nodehash;
    use quickcheck::quickcheck;

    use super::*;

    fn long_string() -> String {
        String::from_utf8(vec![b'T'; u16::max_value() as usize]).unwrap()
    }

    fn size_matches(data: &[u8]) -> bool {
        (BigEndian::read_u32(data) as usize) == data.len()
    }

    fn metadata_matches(data: &[u8], metadata: &[MetadataEntry]) -> bool {
        let mut off = 0;

        let sizes_ok = metadata.iter().all(|e| {
            let key_size = data[off] as usize;
            off += 1;

            let value_size = data[off] as usize;
            off += 1;

            key_size == e.key.as_bytes().len() && value_size == e.value.as_bytes().len()
        });

        let strings_ok = metadata.iter().all(|e| {
            let key_size = e.key.as_bytes().len();
            let value_size = e.value.as_bytes().len();

            let key = String::from_utf8(data[off..off + key_size].to_vec());
            off += key_size;

            let value = String::from_utf8(data[off..off + value_size].to_vec());
            off += value_size;

            key.is_ok() && key.unwrap() == e.key && value.is_ok() && value.unwrap() == e.value
        });

        sizes_ok && strings_ok
    }

    fn successors_match(data: &[u8], successors: &[HgChangesetId]) -> bool {
        let mut off = 0;

        successors.iter().all(|succ| {
            let end = off + 20;
            let range = off..end;
            off = end;

            HgChangesetId::from_bytes(&data[range]).expect("not a changeset") == *succ
        })
    }

    fn content_matches(
        data: &[u8],
        predecessor: &HgChangesetId,
        successors: &Vec<HgChangesetId>,
        time: &DateTime,
        metadata: &Vec<MetadataEntry>,
    ) -> bool {
        BigEndian::read_f64(&data[4..]) == time.timestamp_secs() as f64
            && BigEndian::read_i16(&data[12..]) == time.tz_offset_minutes()
            && BigEndian::read_i16(&data[14..]) == 0
            && data[16] == (successors.len() as u8)
            && data[17] == 3
            && data[18] == (metadata.len() as u8)
            && HgChangesetId::from_bytes(&data[19..39]).expect("not a changeset") == *predecessor
            && successors_match(&data[39..], successors)
            && metadata_matches(&data[(39 + 20 * successors.len())..], metadata)
    }

    quickcheck! {
        fn test_prepare_no_metadata(predecessor: HgChangesetId, successors: Vec<HgChangesetId>, time: DateTime) -> bool {
            let chunk = prepare_obsmarker_chunk(&predecessor, &successors, &time, &vec![]);
            chunk.is_ok() && size_matches(&chunk.unwrap().into_bytes().unwrap())
        }

        fn test_prepare_metadata(predecessor: HgChangesetId, successors: Vec<HgChangesetId>, time: DateTime, metadata: Vec<MetadataEntry>) -> bool {
            let chunk = prepare_obsmarker_chunk(&predecessor, &successors, &time, &metadata);
            let max_size = u8::max_value() as usize;

            if metadata.len() > max_size || successors.len() > max_size {
                // NOTE: With the default quickcheck configuration, we won't exercise this. We
                // actually test this below, but this branch ensures the tests won't fail if
                // we use a larger quickcheck vector size (through QUICKCHECK_GENERATOR_SIZE).
                chunk.is_err()
            } else {
                chunk.is_ok() && size_matches(&chunk.unwrap().into_bytes().unwrap())
            }
        }

        fn test_roundtrip(predecessor: HgChangesetId, successors: Vec<HgChangesetId>, time: DateTime, metadata: Vec<MetadataEntry>) -> bool {
            let metadata = metadata.into_iter().take(4).collect(); // See above
            let chunk = prepare_obsmarker_chunk(&predecessor, &successors, &time, &metadata);
            chunk.is_ok() && content_matches(&chunk.unwrap().into_bytes().unwrap(), &predecessor, &successors, &time, &metadata)
        }
    }

    #[test]
    fn test_no_successors() {
        let time = DateTime::now();
        let successors = vec![];
        let metadata = vec![];
        let chunk = prepare_obsmarker_chunk(&nodehash::ONES_CSID, &successors, &time, &metadata);
        assert!(chunk.is_ok());
    }

    #[test]
    fn test_successor_count_overflow() {
        let time = DateTime::now();
        let successors = vec![nodehash::TWOS_CSID; u16::max_value() as usize];
        let metadata = vec![];
        let chunk = prepare_obsmarker_chunk(&nodehash::ONES_CSID, &successors, &time, &metadata);
        assert!(chunk.is_err());
    }

    #[test]
    fn test_metadata_count_overflow() {
        let entry = MetadataEntry::new("key", "value");
        let time = DateTime::now();
        let metadata = vec![entry; u16::max_value() as usize];
        let chunk = prepare_obsmarker_chunk(
            &nodehash::ONES_CSID,
            &vec![nodehash::TWOS_CSID],
            &time,
            &metadata,
        );
        assert!(chunk.is_err());
    }

    #[test]
    fn test_metadata_key_overflow() {
        let entry = MetadataEntry::new(long_string(), "value");
        let time = DateTime::now();
        let metadata = vec![entry];
        let chunk = prepare_obsmarker_chunk(
            &nodehash::ONES_CSID,
            &vec![nodehash::TWOS_CSID],
            &time,
            &metadata,
        );
        assert!(chunk.is_err());
    }

    #[test]
    fn test_metadata_value_overflow() {
        let entry = MetadataEntry::new("key", long_string());
        let time = DateTime::now();
        let metadata = vec![entry];
        let chunk = prepare_obsmarker_chunk(
            &nodehash::ONES_CSID,
            &vec![nodehash::TWOS_CSID],
            &time,
            &metadata,
        );
        assert!(chunk.is_err());
    }

    fn stream_for_pairs(
        pairs: Vec<(HgChangesetId, Vec<HgChangesetId>)>,
    ) -> impl Stream<Item = (HgChangesetId, Vec<HgChangesetId>), Error = Error> {
        stream::iter_ok(pairs).boxify()
    }

    fn extract_chunk(wrapped: Poll<Option<Chunk>, Error>) -> Result<Chunk> {
        if let Ok(Async::Ready(Some(chunk))) = wrapped {
            return Ok(chunk);
        }
        Err(Error::msg("no chunk"))
    }

    fn extract_vec(wrapped: Poll<Option<Chunk>, Error>) -> Result<Vec<u8>> {
        Ok(extract_chunk(wrapped)?.into_bytes()?.to_vec())
    }

    #[test]
    fn test_stream_emits_version() {
        let pairs_stream = stream_for_pairs(vec![
            (nodehash::ONES_CSID, vec![nodehash::TWOS_CSID]),
            (nodehash::THREES_CSID, vec![nodehash::FOURS_CSID]),
        ]);
        let time = DateTime::now();
        let meta = vec![];
        let mut packer = obsmarkers_packer_stream(pairs_stream, time, meta.clone());

        let r1 = extract_vec(packer.poll()).unwrap();
        assert_eq!(r1, vec![VERSION]);

        let r2 = extract_vec(packer.poll()).unwrap();
        assert!(size_matches(&r2));
        assert!(content_matches(
            &r2,
            &nodehash::ONES_CSID,
            &vec![nodehash::TWOS_CSID],
            &time,
            &meta
        ));

        let r3 = extract_vec(packer.poll()).unwrap();
        assert!(size_matches(&r3));
        assert!(content_matches(
            &r3,
            &nodehash::THREES_CSID,
            &vec![nodehash::FOURS_CSID],
            &time,
            &meta
        ));

        assert!(extract_vec(packer.poll()).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.
 */

use std::collections::HashMap;
use std::fmt;
use std::io::Write;

use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::WriteBytesExt;
use bytes::Bytes as BytesNew;
use bytes_old::Bytes;
use context::CoreContext;
use futures::compat::Future01CompatExt;
use futures::FutureExt;
use futures::TryFutureExt;
use futures_ext::BoxFuture;
use futures_ext::BoxStream;
use futures_ext::StreamExt;
use futures_old::stream::iter_ok;
use futures_old::stream::once;
use futures_old::Future;
use futures_old::Stream;
use futures_stats::TimedFutureExt;
use mercurial_mutation::HgMutationEntry;
use mercurial_types::Delta;
use mercurial_types::HgBlobNode;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;
use mercurial_types::RepoPath;
use mercurial_types::RevFlags;
use mercurial_types::NULL_HASH;
use mononoke_types::DateTime;
use phases::Phase;

use super::changegroup::packer::CgPacker;
use super::changegroup::unpacker::CgVersion;
use super::changegroup::CgDeltaChunk;
use super::changegroup::Part;
use super::changegroup::Section;
use super::chunk::Chunk;
use super::infinitepush::infinitepush_mutation_packer;
use super::obsmarkers::packer::obsmarkers_packer_stream;
use super::obsmarkers::MetadataEntry;
use super::wirepack;
use super::wirepack::packer::WirePackPacker;
use crate::errors::ErrorKind;
use crate::part_encode::PartEncodeBuilder;
use crate::part_header::PartHeaderType;
use crate::part_header::PartId;

pub type FilenodeEntry = (HgFileNodeId, HgChangesetId, HgBlobNode, Option<RevFlags>);

pub fn listkey_part<N, S, K, V>(namespace: N, items: S) -> Result<PartEncodeBuilder>
where
    N: Into<BytesNew>,
    S: Stream<Item = (K, V), Error = Error> + Send + 'static,
    K: AsRef<[u8]>,
    V: AsRef<[u8]>,
{
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::Listkeys)?;
    builder.add_mparam("namespace", namespace)?;
    // Ideally we'd use a size_hint here, but streams don't appear to have one.
    let payload = Vec::with_capacity(256);
    let fut = items
        .fold(payload, |mut payload, (key, value)| {
            payload.extend_from_slice(key.as_ref());
            payload.push(b'\t');
            payload.extend_from_slice(value.as_ref());
            payload.push(b'\n');
            Ok::<_, Error>(payload)
        })
        .map_err(|err| err.context(ErrorKind::ListkeyGeneration));

    builder.set_data_future(fut);

    Ok(builder)
}

pub fn phases_part<S>(ctx: CoreContext, phases_entries: S) -> Result<PartEncodeBuilder>
where
    S: Stream<Item = (HgChangesetId, Phase), Error = Error> + Send + 'static,
{
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::PhaseHeads)?;
    let mut scuba_logger = ctx.scuba().clone();
    let payload = Vec::with_capacity(1024);
    let fut = phases_entries
        .fold(payload, |mut payload, (value, phase)| {
            payload.write_u32::<BigEndian>(u32::from(phase))?;
            payload.write_all(value.as_ref())?;
            Ok::<_, Error>(payload)
        })
        .map_err(|err| err.context(ErrorKind::PhaseHeadsGeneration))
        .compat()
        .timed()
        .map(move |(stats, result)| {
            if result.is_ok() {
                scuba_logger
                    .add_future_stats(&stats)
                    .log_with_msg("Phases calculated - Success", None);
            }
            result
        })
        .compat();
    builder.set_data_future(fut);
    Ok(builder)
}

pub fn changegroup_part<CS>(
    changelogentries: CS,
    filenodeentries: Option<BoxStream<(MPath, Vec<FilenodeEntry>), Error>>,
    version: CgVersion,
) -> Result<PartEncodeBuilder>
where
    CS: Stream<Item = (HgNodeHash, HgBlobNode), Error = Error> + Send + 'static,
{
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::Changegroup)?;
    builder.add_mparam(
        "version",
        BytesNew::copy_from_slice(version.to_str().as_bytes()),
    )?;

    let changelogentries = convert_changeset_stream(changelogentries, version)
        .chain(once(Ok(Part::SectionEnd(Section::Changeset))))
        // One more SectionEnd entry is necessary because hg client excepts filelog section
        // even if it's empty. Add a fake SectionEnd part (the choice of
        // Manifest is just for convenience).
        .chain(once(Ok(Part::SectionEnd(Section::Manifest))));

    let changelogentries = if version == CgVersion::Cg3Version {
        // Changegroup V3 requires one empty chunk after manifest section
        // hence adding Part::SectionEnd below
        changelogentries
            .chain(once(Ok(Part::SectionEnd(Section::Manifest))))
            .boxify()
    } else {
        changelogentries.boxify()
    };

    let changegroup = if let Some(filenodeentries) = filenodeentries {
        changelogentries
            .chain(convert_file_stream(filenodeentries, version))
            .left_stream()
    } else {
        changelogentries.right_stream()
    };

    let changegroup = changegroup.chain(once(Ok(Part::End)));

    let cgdata = CgPacker::new(changegroup);
    builder.set_data_generated(cgdata);

    Ok(builder)
}

fn convert_changeset_stream<S>(
    changelogentries: S,
    version: CgVersion,
) -> impl Stream<Item = Part, Error = Error>
where
    S: Stream<Item = (HgNodeHash, HgBlobNode), Error = Error> + Send + 'static,
{
    changelogentries.map(move |(node, blobnode)| {
        let parents = blobnode.parents().get_nodes();
        let p1 = parents.0.unwrap_or(NULL_HASH);
        let p2 = parents.1.unwrap_or(NULL_HASH);
        let base = NULL_HASH;
        // Linknode is the same as node
        let linknode = node;
        let text = blobnode.as_blob().as_inner().clone();
        let delta = Delta::new_fulltext(text.to_vec());

        let flags = if version == CgVersion::Cg3Version {
            Some(RevFlags::REVIDX_DEFAULT_FLAGS)
        } else {
            None
        };

        let deltachunk = CgDeltaChunk {
            node,
            p1,
            p2,
            base,
            linknode,
            delta,
            flags,
        };
        Part::CgChunk(Section::Changeset, deltachunk)
    })
}

fn convert_file_stream<FS>(
    filenodeentries: FS,
    cg_version: CgVersion,
) -> impl Stream<Item = Part, Error = Error>
where
    FS: Stream<Item = (MPath, Vec<FilenodeEntry>), Error = Error> + Send + 'static,
{
    filenodeentries
        .map(move |(path, nodes)| {
            let mut items = vec![];
            for (node, hg_cs_id, blobnode, flags) in nodes {
                let parents = blobnode.parents().get_nodes();
                let p1 = parents.0.unwrap_or(NULL_HASH);
                let p2 = parents.1.unwrap_or(NULL_HASH);
                let base = NULL_HASH;
                // Linknode is the same as node
                let linknode = hg_cs_id.into_nodehash();
                let text = blobnode.as_blob().as_inner().clone();
                let delta = Delta::new_fulltext(text.to_vec());

                let deltachunk = CgDeltaChunk {
                    node: node.into_nodehash(),
                    p1,
                    p2,
                    base,
                    linknode,
                    delta,
                    flags,
                };
                if flags.is_some() && cg_version == CgVersion::Cg2Version {
                    return once(Err(Error::msg(
                        "internal error: unexpected flags in cg2 generation",
                    )))
                    .boxify();
                }
                items.push(Part::CgChunk(Section::Filelog(path.clone()), deltachunk));
            }

            items.push(Part::SectionEnd(Section::Filelog(path)));
            iter_ok(items).boxify()
        })
        .flatten()
}

pub fn replycaps_part(caps: Bytes) -> Result<PartEncodeBuilder> {
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::Replycaps)?;
    builder.set_data_fixed(Chunk::new(caps)?);

    Ok(builder)
}

pub fn common_heads_part(heads: Vec<HgChangesetId>) -> Result<PartEncodeBuilder> {
    let mut w = Vec::new();
    for h in heads {
        w.extend(h.as_bytes());
    }

    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::B2xCommonHeads)?;
    builder.set_data_fixed(Chunk::new(w)?);

    Ok(builder)
}

pub struct TreepackPartInput {
    pub node: HgNodeHash,
    pub p1: Option<HgNodeHash>,
    pub p2: Option<HgNodeHash>,
    pub content: Bytes,
    pub fullpath: Option<MPath>,
    pub linknode: HgNodeHash,
}

// Controls whether client should store trees in hg cache (which means
// they might be deleted and refetched from the server) or they should be stored
// in .hg directory (which means client should never delete them).
// Normally this should only be used for hydrated commit cloud commits, and
// with hg server deprecation it won't be necessary anymore.
#[derive(Clone, Copy)]
pub enum StoreInHgCache {
    Yes,
    No,
}

pub fn treepack_part<S>(entries: S, hg_cache_policy: StoreInHgCache) -> Result<PartEncodeBuilder>
where
    S: Stream<Item = BoxFuture<TreepackPartInput, Error>, Error = Error> + Send + 'static,
{
    treepack_part_impl(entries, PartHeaderType::B2xTreegroup2, hg_cache_policy)
}

pub fn pushrebase_treepack_part<S>(entries: S) -> Result<PartEncodeBuilder>
where
    S: Stream<Item = BoxFuture<TreepackPartInput, Error>, Error = Error> + Send + 'static,
{
    treepack_part_impl(entries, PartHeaderType::B2xRebasePack, StoreInHgCache::Yes)
}

fn treepack_part_impl<S>(
    entries: S,
    header_type: PartHeaderType,
    hg_cache_policy: StoreInHgCache,
) -> Result<PartEncodeBuilder>
where
    S: Stream<Item = BoxFuture<TreepackPartInput, Error>, Error = Error> + Send + 'static,
{
    let mut builder = PartEncodeBuilder::mandatory(header_type)?;
    builder.add_mparam("version", "1")?;
    match hg_cache_policy {
        StoreInHgCache::Yes => {
            builder.add_mparam("cache", "True")?;
        }
        StoreInHgCache::No => {
            builder.add_mparam("cache", "False")?;
        }
    };

    builder.add_mparam("category", "manifests")?;

    let mut buffer_size = tunables::tunables()
        .repo_client_gettreepack_buffer_size()
        .unwrap_or_default();
    if buffer_size <= 0 {
        buffer_size = 1000
    }
    let buffer_size: usize = buffer_size
        .try_into()
        .with_context(|| format!("invalid buffer size {}", buffer_size))?;
    let wirepack_parts = entries
        .buffered(buffer_size)
        .map(|input| {
            let path = match input.fullpath {
                Some(path) => RepoPath::DirectoryPath(path),
                None => RepoPath::RootPath,
            };

            let history_meta = wirepack::Part::HistoryMeta {
                path: path.clone(),
                entry_count: 1,
            };

            let history = wirepack::Part::History(wirepack::HistoryEntry {
                node: input.node.clone(),
                p1: input.p1.into(),
                p2: input.p2.into(),
                linknode: input.linknode,
                // No copies/renames for trees
                copy_from: None,
            });

            let data_meta = wirepack::Part::DataMeta {
                path,
                entry_count: 1,
            };

            let data = wirepack::Part::Data(wirepack::DataEntry {
                node: input.node,
                delta_base: NULL_HASH,
                delta: Delta::new_fulltext(input.content.to_vec()),
                metadata: None,
            });

            iter_ok(vec![history_meta, history, data_meta, data].into_iter())
        })
        .flatten()
        .chain(once(Ok(wirepack::Part::End)));

    let packer = WirePackPacker::new(wirepack_parts, wirepack::Kind::Tree);
    builder.set_data_generated(packer);

    Ok(builder)
}

pub enum ChangegroupApplyResult {
    Success { heads_num_diff: i64 },
    Error,
}

// Mercurial source code comments are a bit contradictory:
//
// From mercurial/changegroup.py
// Return an integer summarizing the change to this repo:
// - nothing changed or no source: 0
// - more heads than before: 1+added heads (2..n)
// - fewer heads than before: -1-removed heads (-2..-n)
// - number of heads stays the same: 1
//
// From mercurial/exchange.py
// Integer version of the changegroup push result
// - None means nothing to push
// - 0 means HTTP error
// - 1 means we pushed and remote head count is unchanged *or*
//   we have outgoing changesets but refused to push
// - other values as described by addchangegroup()
//
// We are using 0 to indicate a error, 1 + heads_num_diff if the number of heads increased,
// -1 + heads_num_diff if the number of heads decreased. Note that we may change it in the future

impl fmt::Display for ChangegroupApplyResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            ChangegroupApplyResult::Success { heads_num_diff } => {
                if heads_num_diff >= 0 {
                    write!(f, "{}", 1 + heads_num_diff)
                } else {
                    write!(f, "{}", -1 + heads_num_diff)
                }
            }
            ChangegroupApplyResult::Error => write!(f, "0"),
        }
    }
}

pub fn replychangegroup_part(
    res: ChangegroupApplyResult,
    in_reply_to: PartId,
) -> Result<PartEncodeBuilder> {
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::ReplyChangegroup)?;
    builder.add_mparam("return", format!("{}", res))?;
    builder.add_mparam("in-reply-to", format!("{}", in_reply_to))?;

    Ok(builder)
}

pub fn bookmark_pushkey_part(key: String, old: String, new: String) -> Result<PartEncodeBuilder> {
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::Pushkey)?;
    builder.add_mparam("namespace", "bookmarks")?;
    builder.add_mparam("key", key)?;
    builder.add_mparam("old", old)?;
    builder.add_mparam("new", new)?;

    Ok(builder)
}

pub fn replypushkey_part(res: bool, in_reply_to: PartId) -> Result<PartEncodeBuilder> {
    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::ReplyPushkey)?;
    if res {
        builder.add_mparam("return", "1")?;
    } else {
        builder.add_mparam("return", "0")?;
    }
    builder.add_mparam("in-reply-to", format!("{}", in_reply_to))?;

    Ok(builder)
}

pub fn obsmarkers_part<S>(
    pairs: S,
    time: DateTime,
    metadata: Vec<MetadataEntry>,
) -> Result<PartEncodeBuilder>
where
    S: 'static + Stream<Item = (HgChangesetId, Vec<HgChangesetId>), Error = Error> + Send,
{
    let stream = obsmarkers_packer_stream(pairs, time, metadata);

    let mut builder = PartEncodeBuilder::mandatory(PartHeaderType::Obsmarkers)?;
    builder.set_data_generated(stream);
    Ok(builder)
}

pub fn infinitepush_mutation_part<F>(entries: F) -> Result<PartEncodeBuilder>
where
    F: Future<Item = Vec<HgMutationEntry>, Error = Error> + Send + 'static,
{
    let mut builder = PartEncodeBuilder::advisory(PartHeaderType::B2xInfinitepushMutation)?;
    let data = entries.and_then(infinitepush_mutation_packer);
    builder.set_data_future(data);
    Ok(builder)
}

pub fn pushvars_part(push_vars: HashMap<String, BytesNew>) -> Result<PartEncodeBuilder> {
    let mut builder = PartEncodeBuilder::advisory(PartHeaderType::Pushvars)?;
    for (var, bytes) in push_vars {
        builder.add_aparam(var, bytes)?;
    }
    Ok(builder)
}


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

// Pushrebase codecs

use anyhow::Error;
use anyhow::Result;
use bytes_old::BytesMut;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use tokio_codec::Decoder;

#[derive(Debug)]
pub struct CommonHeadsUnpacker {}

impl CommonHeadsUnpacker {
    pub fn new() -> Self {
        Self {}
    }
}

impl Decoder for CommonHeadsUnpacker {
    type Item = HgChangesetId;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        if buf.len() >= 20 {
            let newcsid = buf.split_to(20).freeze();
            let nodehash = HgNodeHash::from_bytes(&newcsid)?;
            Ok(Some(HgChangesetId::new(nodehash)))
        } else {
            Ok(None)
        }
    }
}


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

//! Given a stream of wirepack entries, convert it to to other stream, using WirePackPartProcessor

use std::mem;

use anyhow::bail;
use anyhow::ensure;
use anyhow::Error;
use anyhow::Result;
use futures_old::try_ready;
use futures_old::Async;
use futures_old::Poll;
use futures_old::Stream;
use mercurial_types::RepoPath;

use super::DataEntry;
use super::HistoryEntry;
use super::Part;
use crate::errors::ErrorKind;

pub trait WirePackPartProcessor {
    type Data;

    fn history_meta(&mut self, path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>>;
    fn history(&mut self, entry: &HistoryEntry) -> Result<Option<Self::Data>>;
    fn data_meta(&mut self, path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>>;
    fn data(&mut self, data_entry: &DataEntry) -> Result<Option<Self::Data>>;
    fn end(&mut self) -> Result<Option<Self::Data>>;
}

pub struct WirePackConverter<S, P> {
    part_stream: S,
    state: State,
    processor: P,
}

impl<S, P> WirePackConverter<S, P> {
    pub fn new(part_stream: S, processor: P) -> Self {
        Self {
            part_stream,
            state: State::HistoryMeta,
            processor,
        }
    }
}

impl<S, P> Stream for WirePackConverter<S, P>
where
    S: Stream<Item = Part, Error = Error>,
    P: WirePackPartProcessor,
{
    type Item = <P as WirePackPartProcessor>::Data;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        use self::Part::*;

        if self.state == State::End {
            // The stream is over.
            return Ok(Async::Ready(None));
        }

        loop {
            match try_ready!(self.part_stream.poll()) {
                None => {
                    self.state.seen_none()?;
                    return Ok(Async::Ready(None));
                }
                Some(HistoryMeta { path, entry_count }) => {
                    if let Some(history_meta) = self.processor.history_meta(&path, entry_count)? {
                        self.state.seen_history_meta(path, entry_count)?;
                        return Ok(Async::Ready(Some(history_meta)));
                    }
                    // seen_history_meta comes afterwards because we have to transfer
                    // ownership of path
                    self.state.seen_history_meta(path, entry_count)?;
                }

                Some(History(history_entry)) => {
                    self.state.seen_history(&history_entry)?;
                    if let Some(encoded_history) = self.processor.history(&history_entry)? {
                        return Ok(Async::Ready(Some(encoded_history)));
                    }
                }

                Some(DataMeta { path, entry_count }) => {
                    if let Some(data_meta) = self.processor.data_meta(&path, entry_count)? {
                        self.state.seen_data_meta(path, entry_count)?;
                        return Ok(Async::Ready(Some(data_meta)));
                    }
                    self.state.seen_data_meta(path, entry_count)?;
                }

                Some(Data(data_entry)) => {
                    self.state.seen_data(&data_entry)?;
                    if let Some(data_entry) = self.processor.data(&data_entry)? {
                        return Ok(Async::Ready(Some(data_entry)));
                    }
                }

                Some(End) => {
                    self.state.seen_end()?;
                    if let Some(end) = self.processor.end()? {
                        return Ok(Async::Ready(Some(end)));
                    }
                }
            }
        }
    }
}

#[derive(Debug, PartialEq)]
enum State {
    HistoryMeta,
    History { path: RepoPath, entry_count: u32 },
    DataMeta { path: RepoPath },
    Data { path: RepoPath, entry_count: u32 },
    End,
    Invalid,
}

impl State {
    fn next_history_state(path: RepoPath, entry_count: u32) -> Self {
        if entry_count == 0 {
            State::DataMeta { path }
        } else {
            State::History { path, entry_count }
        }
    }

    fn next_data_state(path: RepoPath, entry_count: u32) -> Self {
        if entry_count == 0 {
            State::HistoryMeta
        } else {
            State::Data { path, entry_count }
        }
    }

    fn seen_history_meta(&mut self, path: RepoPath, entry_count: u32) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::HistoryMeta => Self::next_history_state(path, entry_count),
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: unexpected history meta entry (state: {:?})",
                    other
                )));
            }
        };
        Ok(())
    }

    fn seen_history(&mut self, entry: &HistoryEntry) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::History { path, entry_count } => {
                ensure!(
                    entry_count > 0,
                    ErrorKind::WirePackEncode(format!(
                        "invalid encode stream: saw history entry for {} after count dropped to 0",
                        entry.node
                    ))
                );
                Self::next_history_state(path, entry_count - 1)
            }
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: unexpected history entry for {} (state: {:?})",
                    entry.node, other
                )));
            }
        };
        Ok(())
    }

    fn seen_data_meta(&mut self, path: RepoPath, entry_count: u32) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::DataMeta {
                path: expected_path,
            } => {
                ensure!(
                    path == expected_path,
                    ErrorKind::WirePackEncode(format!(
                        "invalid encode stream: saw data meta for path '{}', expected path '{}'\
                         (entry_count: {})",
                        path, expected_path, entry_count
                    ))
                );
                Self::next_data_state(path, entry_count)
            }
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: saw unexpected data meta for {} (entry count: {}, \
                     state: {:?}",
                    path, entry_count, other
                ),));
            }
        };
        Ok(())
    }

    fn seen_data(&mut self, entry: &DataEntry) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::Data { path, entry_count } => {
                ensure!(
                    entry_count > 0,
                    ErrorKind::WirePackEncode(format!(
                        "invalid encode stream: saw history entry for {} after count dropped to 0",
                        entry.node
                    ))
                );
                Self::next_data_state(path, entry_count - 1)
            }
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: unexpected data entry for {} (state: {:?})",
                    entry.node, other
                )));
            }
        };
        Ok(())
    }

    fn seen_end(&mut self) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::HistoryMeta => State::End,
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: unexpected end (state: {:?})",
                    other
                )));
            }
        };
        Ok(())
    }

    fn seen_none(&mut self) -> Result<()> {
        let state = mem::replace(self, State::Invalid);
        *self = match state {
            State::End => State::End,
            other => {
                bail!(ErrorKind::WirePackEncode(format!(
                    "invalid encode stream: unexpected None (state: {:?})",
                    other
                )));
            }
        };
        Ok(())
    }
}


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

//! Wire packs. The format is currently undocumented.

use std::fmt;
use std::io::Cursor;

use anyhow::bail;
use anyhow::ensure;
use anyhow::Context;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::BufMut;
use bytes_old::BytesMut;
use mercurial_types::Delta;
use mercurial_types::HgNodeHash;
use mercurial_types::RepoPath;
use mercurial_types::NULL_HASH;
use revisionstore_types::Metadata;

use crate::delta;
use crate::errors::ErrorKind;
use crate::utils::BytesExt;

pub mod converter;
pub mod packer;
#[cfg(test)]
mod quickcheck_types;
pub mod unpacker;

/// What sort of wirepack this is.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Kind {
    /// A wire pack representing tree manifests.
    Tree,
    /// A wire pack representing file contents.
    File,
}

impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Kind::Tree => write!(f, "tree"),
            Kind::File => write!(f, "file"),
        }
    }
}

/// An atomic part returned from the wirepack.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Part {
    HistoryMeta { path: RepoPath, entry_count: u32 },
    History(HistoryEntry),
    DataMeta { path: RepoPath, entry_count: u32 },
    Data(DataEntry),
    End,
}

#[cfg(test)]
impl Part {
    pub(crate) fn unwrap_history_meta(self) -> (RepoPath, u32) {
        match self {
            Part::HistoryMeta { path, entry_count } => (path, entry_count),
            other => panic!("expected wirepack part to be HistoryMeta, was {:?}", other),
        }
    }

    pub(crate) fn unwrap_history(self) -> HistoryEntry {
        match self {
            Part::History(entry) => entry,
            other => panic!("expected wirepack part to be History, was {:?}", other),
        }
    }

    pub(crate) fn unwrap_data_meta(self) -> (RepoPath, u32) {
        match self {
            Part::DataMeta { path, entry_count } => (path, entry_count),
            other => panic!("expected wirepack part to be HistoryMeta, was {:?}", other),
        }
    }

    pub(crate) fn unwrap_data(self) -> DataEntry {
        match self {
            Part::Data(entry) => entry,
            other => panic!("expected wirepack part to be Data, was {:?}", other),
        }
    }
}

const WIREPACK_END: &[u8] = b"\0\0\0\0\0\0\0\0\0\0";

// See the history header definition in this file for the breakdown.
const HISTORY_COPY_FROM_OFFSET: usize = 20 + 20 + 20 + 20;
const HISTORY_HEADER_SIZE: usize = HISTORY_COPY_FROM_OFFSET + 2;

// See the data header definition in this file for the breakdown.
const DATA_DELTA_OFFSET: usize = 20 + 20;
const DATA_HEADER_SIZE: usize = DATA_DELTA_OFFSET + 8;

// TODO: move to mercurial_types
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct HistoryEntry {
    pub node: HgNodeHash,
    // TODO: replace with HgParents?
    pub p1: HgNodeHash,
    pub p2: HgNodeHash,
    pub linknode: HgNodeHash,
    pub copy_from: Option<RepoPath>,
}

impl HistoryEntry {
    pub(crate) fn decode(buf: &mut BytesMut, kind: Kind) -> Result<Option<Self>> {
        if buf.len() < HISTORY_HEADER_SIZE {
            return Ok(None);
        }

        // A history revision has:
        // ---
        // node: HgNodeHash (20 bytes)
        // p1: HgNodeHash (20 bytes)
        // p2: HgNodeHash (20 bytes)
        // link node: HgNodeHash (20 bytes)
        // copy from len: u16 (2 bytes) -- 0 if this revision is not a copy
        // copy from: RepoPath (<copy from len> bytes)
        // ---
        // Tree revisions are never copied, so <copy from len> is always 0.

        let copy_from_len =
            BigEndian::read_u16(&buf[HISTORY_COPY_FROM_OFFSET..HISTORY_HEADER_SIZE]) as usize;
        if buf.len() < HISTORY_HEADER_SIZE + copy_from_len {
            return Ok(None);
        }

        let node = buf.drain_node();
        let p1 = buf.drain_node();
        let p2 = buf.drain_node();
        let linknode = buf.drain_node();
        let _ = buf.drain_u16();
        let copy_from =
            if copy_from_len > 0 {
                let path = buf.drain_path(copy_from_len)?;
                match kind {
                    Kind::Tree => bail!(ErrorKind::WirePackDecode(format!(
                        "tree entry {} is marked as copied from path {}, but they cannot be copied",
                        node, path
                    ))),
                    Kind::File => Some(RepoPath::file(path).with_context(|| {
                        ErrorKind::WirePackDecode("invalid copy from path".into())
                    })?),
                }
            } else {
                None
            };
        Ok(Some(Self {
            node,
            p1,
            p2,
            linknode,
            copy_from,
        }))
    }

    // This would ideally be generic over any BufMut, but that won't be very useful until
    // https://github.com/carllerche/bytes/issues/170 is fixed.
    pub(crate) fn encode(&self, kind: Kind, buf: &mut Vec<u8>) -> Result<()> {
        self.verify(kind).with_context(|| {
            ErrorKind::WirePackEncode("attempted to encode an invalid history entry".into())
        })?;
        buf.put_slice(self.node.as_ref());
        buf.put_slice(self.p1.as_ref());
        buf.put_slice(self.p2.as_ref());
        buf.put_slice(self.linknode.as_ref());
        let path_vec = if let Some(ref path) = self.copy_from {
            path.mpath()
                .expect("verify ensures that path is always a RepoPath::FilePath")
                .to_vec()
        } else {
            vec![]
        };
        buf.put_u16_be(path_vec.len() as u16);
        buf.put_slice(&path_vec);
        Ok(())
    }

    pub fn verify(&self, kind: Kind) -> Result<()> {
        if let Some(ref path) = self.copy_from {
            match *path {
                RepoPath::RootPath => bail!(ErrorKind::InvalidWirePackEntry(format!(
                    "history entry for {} is copied from the root path, which isn't allowed",
                    self.node
                ))),
                RepoPath::DirectoryPath(ref path) => {
                    bail!(ErrorKind::InvalidWirePackEntry(format!(
                        "history entry for {} is copied from directory {}, which isn't allowed",
                        self.node, path
                    )))
                }
                RepoPath::FilePath(ref path) => {
                    ensure!(
                        kind == Kind::File,
                        ErrorKind::InvalidWirePackEntry(format!(
                            "history entry for {} is copied from file {}, but the pack is of \
                             kind {}",
                            self.node, path, kind
                        ))
                    );
                    ensure!(
                        path.len() <= (u16::max_value() as usize),
                        ErrorKind::InvalidWirePackEntry(format!(
                            "history entry for {} is copied from a path of length {} -- maximum \
                             length supported is {}",
                            self.node,
                            path.len(),
                            u16::max_value(),
                        ),)
                    );
                }
            }
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DataEntry {
    pub node: HgNodeHash,
    pub delta_base: HgNodeHash,
    pub delta: Delta,
    /// Metadata presence implies a getpack protocol version that supports it (i.e. version 2)
    pub metadata: Option<Metadata>,
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DataEntryVersion {
    V1,
    V2,
}

impl DataEntry {
    pub(crate) fn decode(buf: &mut BytesMut, version: DataEntryVersion) -> Result<Option<Self>> {
        if buf.len() < DATA_HEADER_SIZE {
            return Ok(None);
        }

        // A data revision has:
        // ---
        // node: HgNodeHash (20 bytes)
        // delta base: HgNodeHash (20 bytes) -- NULL_HASH if full text
        // delta len: u64 (8 bytes)
        // delta: Delta (<delta len> bytes)
        // On v2:
        // metadata len: u32 (4 bytes)
        // metadata: [metadata item, ...] (<metadata len> bytes)
        // metadata item:
        //  metadata key: u8 (1 byte)
        //  metadata value len: u16 (2 bytes)
        //  metadata value: Bytes (<metadata value len> bytes)
        // ---
        // There's a bit of a wart in the current format: if delta base is NULL_HASH, instead of
        // storing a delta with start = 0 and end = 0, we store the full text directly. This
        // should be fixed in a future wire protocol revision.

        // First, check that we have enough data to proceed.
        let delta_len = BigEndian::read_u64(&buf[DATA_DELTA_OFFSET..DATA_HEADER_SIZE]) as usize;
        match version {
            DataEntryVersion::V1 => {
                if buf.len() < DATA_HEADER_SIZE + delta_len {
                    return Ok(None);
                }
            }
            DataEntryVersion::V2 => {
                let meta_offset = DATA_HEADER_SIZE + delta_len;
                let meta_header_size = 4; // Metadata header is a u32.
                if buf.len() < meta_offset + meta_header_size {
                    return Ok(None);
                }

                let meta_header_slice = &buf[meta_offset..meta_offset + meta_header_size];
                let meta_size = BigEndian::read_u32(meta_header_slice) as usize;
                if buf.len() < meta_offset + meta_header_size + meta_size {
                    return Ok(None);
                }
            }
        }

        let node = buf.drain_node();
        let delta_base = buf.drain_node();
        let _ = buf.drain_u64();
        let delta = buf.split_to(delta_len);

        let delta = if delta_base == NULL_HASH {
            Delta::new_fulltext(delta.to_vec())
        } else {
            delta::decode_delta(delta)?
        };

        let metadata = match version {
            DataEntryVersion::V1 => None,
            DataEntryVersion::V2 => {
                let mut cursor = Cursor::new(buf.as_ref());
                let metadata = Metadata::read(&mut cursor)?;

                // Metadata::read doesn't consume bytes (it just reads them), but our
                // implementation here wants to consume bytes we read. We work around this here.
                let pos = cursor.position();
                let _ = buf.split_to(pos as usize);

                Some(metadata)
            }
        };

        Ok(Some(Self {
            node,
            delta_base,
            delta,
            metadata,
        }))
    }

    // This would ideally be generic over any BufMut, but that won't be very useful until
    // https://github.com/carllerche/bytes/issues/170 is fixed.
    pub(crate) fn encode(&self, buf: &mut Vec<u8>) -> Result<()> {
        self.verify()?;
        buf.put_slice(self.node.as_ref());
        buf.put_slice(self.delta_base.as_ref());
        if self.delta_base == NULL_HASH {
            // This is a fulltext -- the spec requires that instead of storing a delta with
            // start = 0 and end = 0, the fulltext be stored directly.
            let fulltext = self
                .delta
                .maybe_fulltext()
                .expect("verify will have already checked that the delta is a fulltext");
            buf.put_u64_be(fulltext.len() as u64);
            buf.put_slice(fulltext);
        } else {
            buf.put_u64_be(delta::encoded_len(&self.delta) as u64);
            delta::encode_delta(&self.delta, buf);
        }

        if let Some(ref metadata) = self.metadata {
            metadata.write(buf)?
        }

        Ok(())
    }

    pub fn verify(&self) -> Result<()> {
        // The only limitation is that the delta base being null means that the revision is a
        // fulltext.
        ensure!(
            self.delta_base != NULL_HASH || self.delta.maybe_fulltext().is_some(),
            ErrorKind::InvalidWirePackEntry(format!(
                "data entry for {} has a null base but is not a fulltext",
                self.node
            ))
        );
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use itertools::iproduct;
    use maplit::hashset;
    use mercurial_types::delta::Fragment;
    use mercurial_types_mocks::nodehash::AS_HASH;
    use mercurial_types_mocks::nodehash::BS_HASH;
    use quickcheck::quickcheck;
    use quickcheck::Arbitrary;
    use quickcheck::Gen;

    use super::*;

    #[test]
    fn test_history_verify_basic() {
        let foo_dir = RepoPath::dir("foo").unwrap();
        let bar_file = RepoPath::file("bar").unwrap();
        let root = RepoPath::root();

        let valid = hashset! {
            (Kind::Tree, None),
            (Kind::File, None),
            (Kind::File, Some(bar_file.clone())),
        };
        // Can't use arrays here because IntoIterator isn't supported for them:
        // https://github.com/rust-lang/rust/issues/25725
        let kinds = vec![Kind::Tree, Kind::File].into_iter();
        let copy_froms = vec![None, Some(foo_dir), Some(bar_file), Some(root)].into_iter();

        for pair in iproduct!(kinds, copy_froms) {
            let is_valid = valid.contains(&pair);
            let (kind, copy_from) = pair;
            let entry = make_history_entry(copy_from);
            let result = entry.verify(kind);
            if is_valid {
                result.expect("expected history entry to be valid");
            } else {
                result.expect_err("expected history entry to be invalid");
            }
        }
    }

    fn make_history_entry(copy_from: Option<RepoPath>) -> HistoryEntry {
        HistoryEntry {
            node: NULL_HASH,
            p1: NULL_HASH,
            p2: NULL_HASH,
            linknode: NULL_HASH,
            copy_from,
        }
    }

    #[test]
    fn test_history_verify_arbitrary() {
        let mut rng = Gen::new(100);
        for _n in 0..100 {
            HistoryEntry::arbitrary_kind(&mut rng, Kind::Tree)
                .verify(Kind::Tree)
                .expect("generated history entry should be valid");
            HistoryEntry::arbitrary_kind(&mut rng, Kind::File)
                .verify(Kind::File)
                .expect("generated history entry should be valid");
        }
    }

    #[test]
    fn test_history_roundtrip() {
        let mut rng = Gen::new(100);
        for _n in 0..100 {
            history_roundtrip(&mut rng, Kind::Tree);
            history_roundtrip(&mut rng, Kind::File);
        }
    }

    fn history_roundtrip(g: &mut Gen, kind: Kind) {
        let entry = HistoryEntry::arbitrary_kind(g, kind);
        let mut encoded = vec![];
        entry
            .encode(kind, &mut encoded)
            .expect("encoding this history entry should succeed");

        let mut encoded_bytes = BytesMut::from(encoded);

        // Ensure that a partial entry results in None.
        let bytes_len = encoded_bytes.len();
        let reduced_len = unsafe {
            let reduced_len = if bytes_len == 0 {
                0
            } else {
                1 + usize::arbitrary(g) % (bytes_len - 1)
            };
            encoded_bytes.set_len(reduced_len);
            reduced_len
        };
        let decoded = HistoryEntry::decode(&mut encoded_bytes, kind)
            .expect("decoding this history entry should succeed");
        assert_eq!(decoded, None);
        // Ensure that no bytes in encoded actually got read
        assert_eq!(encoded_bytes.len(), reduced_len);

        // Restore the original length for the next test.
        unsafe {
            encoded_bytes.set_len(bytes_len);
        }

        let decoded = HistoryEntry::decode(&mut encoded_bytes, kind)
            .expect("decoding this history entry should succeed");
        assert_eq!(Some(entry), decoded);
        assert_eq!(encoded_bytes.len(), 0);
    }

    #[test]
    fn test_data_verify_basic() {
        #[rustfmt::skip]
        let tests = vec![
            (NULL_HASH, vec![Fragment { start: 0, end: 0, content: vec![b'a'] }], true),
            (NULL_HASH, vec![Fragment { start: 0, end: 5, content: vec![b'b'] }], false),
            (AS_HASH, vec![Fragment { start: 0, end: 0, content: vec![b'c'] }], true),
            (AS_HASH, vec![Fragment { start: 0, end: 5, content: vec![b'd'] }], true),
        ];

        for (delta_base, frags, is_valid) in tests.into_iter() {
            let delta = Delta::new(frags).expect("test deltas should all be valid");
            let entry = DataEntry {
                node: BS_HASH,
                delta_base,
                delta,
                metadata: None,
            };
            let result = entry.verify();
            if is_valid {
                result.expect("expected data entry to be valid");
            } else {
                result.expect_err("expected data entry to be invalid");
            }
        }
    }

    quickcheck! {
        fn test_data_verify_arbitrary(entry: DataEntry) -> bool {
            entry.verify().is_ok()
        }

        fn test_data_roundtrip(entry: DataEntry) -> bool {
            let version = if entry.metadata.is_none() {
                DataEntryVersion::V1
            } else {
                DataEntryVersion::V2
            };

            let mut rng = Gen::new(100);

            let mut encoded = vec![];
            entry.encode(&mut encoded).expect("encoding this data entry should succeed");

            let mut encoded_bytes = BytesMut::from(encoded);

            // Ensure that a partial entry results in None.
            let bytes_len = encoded_bytes.len();
            let reduced_len = unsafe {
                let reduced_len = if bytes_len == 0 { 0 } else { 1 + usize::arbitrary(&mut rng) % (bytes_len - 1) };
                encoded_bytes.set_len(reduced_len);
                reduced_len
            };
            let decoded = DataEntry::decode(&mut encoded_bytes, version)
                .expect("decoding this data entry should succeed");
            assert_eq!(decoded, None);
            // Ensure that no bytes in encoded actually got read.
            assert_eq!(encoded_bytes.len(), reduced_len);

            // Restore the original length for the next test.
            unsafe {
                encoded_bytes.set_len(bytes_len);
            }

            let decoded = DataEntry::decode(&mut encoded_bytes, version)
                .expect("decoding this history entry should succeed");
            assert_eq!(Some(entry), decoded);
            assert_eq!(encoded_bytes.len(), 0);
            true
        }
    }
}


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

//! Packing wirepacks to be sent over the wire during e.g. an hg pull.
//! The format is documented at
//! <https://phab.mercurial-scm.org/diffusion/FBHGX/browse/default/remotefilelog/wirepack.py>

#![allow(deprecated)] // TODO: T29077977 convert from put_X::<BigEndian> -> put_X_be

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use byteorder::BigEndian;
use bytes_old::BufMut;
use futures_old::Poll;
use futures_old::Stream;
use mercurial_types::MPath;
use mercurial_types::RepoPath;

use super::converter::WirePackConverter;
use super::converter::WirePackPartProcessor;
use super::DataEntry;
use super::HistoryEntry;
use super::Kind;
use super::Part;
use super::WIREPACK_END;
use crate::chunk::Chunk;
use crate::errors::ErrorKind;

pub struct WirePackPacker<S> {
    stream: WirePackConverter<S, PackerProcessor>,
}

impl<S> WirePackPacker<S>
where
    S: Stream<Item = Part, Error = Error>,
{
    pub fn new(part_stream: S, kind: Kind) -> Self {
        Self {
            stream: WirePackConverter::new(part_stream, PackerProcessor { kind }),
        }
    }
}

impl<S> Stream for WirePackPacker<S>
where
    S: Stream<Item = Part, Error = Error>,
{
    type Item = Chunk;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Chunk>, Error> {
        self.stream.poll()
    }
}

struct PackerProcessor {
    kind: Kind,
}

unsafe impl Send for PackerProcessor {}
unsafe impl Sync for PackerProcessor {}

impl WirePackPartProcessor for PackerProcessor {
    type Data = Chunk;

    fn history_meta(&mut self, path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>> {
        let mut builder = ChunkBuilder::new(self.kind);
        builder.encode_filename(path)?;
        builder.encode_entry_count(entry_count);
        Ok(Some(builder.build()?))
    }

    fn history(&mut self, entry: &HistoryEntry) -> Result<Option<Self::Data>> {
        let mut builder = ChunkBuilder::new(self.kind);
        builder.encode_history(entry)?;
        Ok(Some(builder.build()?))
    }

    fn data_meta(&mut self, _path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>> {
        let mut builder = ChunkBuilder::new(self.kind);
        builder.encode_entry_count(entry_count);
        Ok(Some(builder.build()?))
    }

    fn data(&mut self, data_entry: &DataEntry) -> Result<Option<Self::Data>> {
        let mut builder = ChunkBuilder::new(self.kind);
        builder.encode_data(data_entry)?;
        Ok(Some(builder.build()?))
    }

    fn end(&mut self) -> Result<Option<Self::Data>> {
        let mut builder = ChunkBuilder::new(self.kind);
        builder.encode_end();
        Ok(Some(builder.build()?))
    }
}

#[derive(Debug)]
struct ChunkBuilder {
    kind: Kind,
    inner: Vec<u8>,
}

impl ChunkBuilder {
    pub fn new(kind: Kind) -> Self {
        Self {
            kind,
            inner: Vec::with_capacity(256),
        }
    }

    /// Encode a filename -- this should always happen before any history or data entries are
    /// encoded.
    fn encode_filename(&mut self, filename: &RepoPath) -> Result<&mut Self> {
        let mpath = match (self.kind, filename) {
            (Kind::Tree, &RepoPath::RootPath) => None,
            (Kind::File, &RepoPath::RootPath) => bail!(ErrorKind::WirePackEncode(
                "attempted to encode a zero-length filename into a file wirepack".into()
            )),
            (Kind::Tree, RepoPath::DirectoryPath(dir_path)) => Some(verify_path(dir_path)?),
            (Kind::File, RepoPath::FilePath(file_path)) => Some(verify_path(file_path)?),
            (kind, path) => bail!(ErrorKind::WirePackEncode(format!(
                "attempted to encode incompatible path into wirepack (kind: {}, path: {:?})",
                kind, path
            ))),
        };

        match mpath {
            Some(mpath) => {
                self.inner.put_u16::<BigEndian>(mpath.len() as u16);
                mpath.generate(&mut self.inner)?;
            }
            None => {
                self.inner.put_u16::<BigEndian>(0);
            }
        }
        Ok(self)
    }

    #[inline]
    fn encode_entry_count(&mut self, entry_count: u32) -> &mut Self {
        self.inner.put_u32::<BigEndian>(entry_count);
        self
    }

    #[inline]
    fn encode_history(&mut self, history_entry: &HistoryEntry) -> Result<&mut Self> {
        history_entry.encode(self.kind, &mut self.inner)?;
        Ok(self)
    }

    #[inline]
    fn encode_data(&mut self, data_entry: &DataEntry) -> Result<&mut Self> {
        data_entry.encode(&mut self.inner)?;
        Ok(self)
    }

    #[inline]
    fn encode_end(&mut self) -> &mut Self {
        self.inner.put_slice(WIREPACK_END);
        self
    }

    #[inline]
    fn build(self) -> Result<Chunk> {
        Chunk::new(self.inner)
    }
}

fn verify_path<'a>(mpath: &'a MPath) -> Result<&'a MPath> {
    let len = mpath.len();
    if len > (u16::max_value() as usize) {
        bail!(ErrorKind::WirePackEncode(format!(
            "attempted to encode a filename of length {} -- maximum length supported is {}",
            len,
            u16::max_value()
        )));
    }
    Ok(mpath)
}


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

//! QuickCheck support for wire packs.

use mercurial_types::Delta;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;
use mercurial_types::RepoPath;
use mercurial_types::NULL_HASH;
use quickcheck::empty_shrinker;
use quickcheck::Arbitrary;
use quickcheck::Gen;
use revisionstore_types::Metadata;

use super::DataEntry;
use super::HistoryEntry;
use super::Kind;

#[derive(Clone, Debug)]
pub struct WirePackPartSequence {
    pub kind: Kind,
    pub files: Vec<FileEntries>,
}

impl Arbitrary for WirePackPartSequence {
    fn arbitrary(g: &mut Gen) -> Self {
        let size = g.size();

        let kind = if bool::arbitrary(g) {
            Kind::Tree
        } else {
            Kind::File
        };

        let file_count = usize::arbitrary(g) % size;
        let files = (0..file_count)
            .map(|_| FileEntries::arbitrary_params(g, kind))
            .collect();
        Self { kind, files }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        let kind = self.kind;
        Box::new(self.files.shrink().map(move |files| Self { kind, files }))
    }
}

#[derive(Clone, Debug)]
pub struct FileEntries {
    pub filename: RepoPath,
    pub history: Vec<HistoryEntry>,
    pub data: Vec<DataEntry>,
}

impl FileEntries {
    fn arbitrary_params(g: &mut Gen, kind: Kind) -> Self {
        let size = g.size();
        let history_len = usize::arbitrary(g) % size;
        let data_len = usize::arbitrary(g) % size;

        let filename = match kind {
            Kind::Tree => {
                // 10% chance for it to be the root
                if u64::arbitrary(g) % 10 == 0 {
                    RepoPath::root()
                } else {
                    RepoPath::DirectoryPath(MPath::arbitrary(g))
                }
            }
            Kind::File => RepoPath::FilePath(MPath::arbitrary(g)),
        };
        let history = (0..history_len)
            .map(|_| HistoryEntry::arbitrary_kind(g, kind))
            .collect();
        let data = (0..data_len).map(|_| DataEntry::arbitrary(g)).collect();
        Self {
            filename,
            history,
            data,
        }
    }
}

impl Arbitrary for FileEntries {
    fn arbitrary(_g: &mut Gen) -> Self {
        // FileEntries depends on the kind of the overall wirepack, so this can't be implemented.
        unimplemented!("use WirePackPartSequence::arbitrary instead")
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        let filename = self.filename.clone();
        let self_history = self.history.clone();
        let self_data = self.data.clone();
        Box::new(
            (self_history, self_data)
                .shrink()
                .map(move |(history, data)| Self {
                    filename: filename.clone(),
                    history,
                    data,
                }),
        )
    }
}

impl HistoryEntry {
    pub fn arbitrary_kind(g: &mut Gen, kind: Kind) -> Self {
        let copy_from = match kind {
            Kind::File => {
                // 20% chance of generating copy-from info
                if *g.choose(&[0, 1, 2, 3, 4]).unwrap() == 0 {
                    Some(RepoPath::FilePath(MPath::arbitrary(g)))
                } else {
                    None
                }
            }
            Kind::Tree => None,
        };
        Self {
            node: HgNodeHash::arbitrary(g),
            p1: HgNodeHash::arbitrary(g),
            p2: HgNodeHash::arbitrary(g),
            linknode: HgNodeHash::arbitrary(g),
            copy_from,
        }
    }
}

impl Arbitrary for HistoryEntry {
    fn arbitrary(_g: &mut Gen) -> Self {
        // HistoryEntry depends on the kind of the overall wirepack, so this can't be implemented.
        unimplemented!("use WirePackPartSequence::arbitrary instead")
    }

    // Not going to get anything out of shrinking this since MPath is not shrinkable.
}

impl Arbitrary for DataEntry {
    fn arbitrary(g: &mut Gen) -> Self {
        // 20% chance of a fulltext revision
        let (delta_base, delta) = if *g.choose(&[0, 1, 2, 3, 4]).unwrap() == 0 {
            (NULL_HASH, Delta::new_fulltext(Vec::arbitrary(g)))
        } else {
            let mut delta_base = NULL_HASH;
            while delta_base == NULL_HASH {
                delta_base = HgNodeHash::arbitrary(g);
            }
            (delta_base, Delta::arbitrary(g))
        };

        // 50% chance of having metadata (i.e. being v2)
        let metadata = if bool::arbitrary(g) {
            // 50% chance of flags being present
            let flags = if bool::arbitrary(g) { Some(1) } else { None };
            // 50% chance of size being present
            let size = if bool::arbitrary(g) { Some(2) } else { None };
            Some(Metadata { flags, size })
        } else {
            None
        };

        Self {
            node: HgNodeHash::arbitrary(g),
            delta_base,
            delta,
            metadata,
        }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        // The delta is the only shrinkable here. However, we cannot shrink it if we don't have
        // base (this might generate a non-fulltext delta).
        if self.delta_base == NULL_HASH {
            empty_shrinker()
        } else {
            let node = self.node;
            let delta_base = self.delta_base;
            let metadata = self.metadata;
            Box::new(self.delta.shrink().map(move |delta| Self {
                node,
                delta_base,
                delta,
                metadata,
            }))
        }
    }
}


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

//! Unpacking wirepacks.

use std::cmp;
use std::mem;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bytes_old::BytesMut;
use mercurial_types::RepoPath;
use slog::trace;
use slog::Logger;
use tokio_codec::Decoder;

use super::DataEntry;
use super::DataEntryVersion;
use super::HistoryEntry;
use super::Kind;
use super::Part;
use super::WIREPACK_END;
use crate::errors::ErrorKind;
use crate::utils::BytesExt;

#[derive(Debug)]
pub struct WirePackUnpacker {
    state: State,
    inner: UnpackerInner,
}

impl Decoder for WirePackUnpacker {
    type Item = Part;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        match self.inner.decode_next(buf, self.state.take()) {
            Err(e) => {
                self.state = State::Invalid;
                Err(e)
            }
            Ok((ret, state)) => {
                self.state = state;
                match ret {
                    None => Ok(None),
                    Some(v) => Ok(Some(v)),
                }
            }
        }
    }

    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        match self.decode(buf)? {
            None => {
                if !buf.is_empty() {
                    let len = buf.len();
                    let bytes = &buf[..cmp::min(len, 128)];
                    let msg = format!(
                        "incomplete wirepack: {} bytes remaining in \
                         buffer. State: {:?}, First 128 bytes: {:?}",
                        len, self.state, bytes,
                    );
                    bail!(ErrorKind::WirePackDecode(msg));
                }
                if self.state != State::End {
                    let msg = format!(
                        "incomplete wirepack: expected state End, found {:?}",
                        self.state
                    );
                    bail!(ErrorKind::WirePackDecode(msg));
                }
                Ok(None)
            }
            Some(v) => Ok(Some(v)),
        }
    }
}

pub fn new(logger: Logger, kind: Kind) -> WirePackUnpacker {
    WirePackUnpacker {
        state: State::Filename,
        inner: UnpackerInner { logger, kind },
    }
}

#[derive(Debug)]
struct UnpackerInner {
    logger: Logger,
    kind: Kind,
}

impl UnpackerInner {
    fn decode_next(&mut self, buf: &mut BytesMut, state: State) -> Result<(Option<Part>, State)> {
        use self::State::*;
        let mut state = state;

        loop {
            trace!(&self.logger, "state: {:?}", state);
            match state {
                Filename => match self.decode_filename(buf)? {
                    DecodeRes::None => return Ok((None, State::Filename)),
                    DecodeRes::Some(f) => {
                        state = HistoryStart(f);
                    }
                    DecodeRes::End => return Ok((Some(Part::End), State::End)),
                },
                HistoryStart(f) => match self.decode_section_start(buf) {
                    Some(entry_count) => {
                        let next_state = self.next_history_state(&f, entry_count);
                        return Ok((
                            Some(Part::HistoryMeta {
                                path: f,
                                entry_count,
                            }),
                            next_state,
                        ));
                    }
                    None => return Ok((None, HistoryStart(f))),
                },
                History(f, entry_count) => match self.decode_history(buf)? {
                    Some(entry) => {
                        let entry_count = entry_count - 1;
                        let next_state = self.next_history_state(&f, entry_count);
                        return Ok((Some(Part::History(entry)), next_state));
                    }
                    None => return Ok((None, History(f, entry_count))),
                },
                DataStart(f) => match self.decode_section_start(buf) {
                    Some(entry_count) => {
                        let next_state = self.next_data_state(&f, entry_count);
                        return Ok((
                            Some(Part::DataMeta {
                                path: f,
                                entry_count,
                            }),
                            next_state,
                        ));
                    }
                    None => return Ok((None, DataStart(f))),
                },
                Data(f, entry_count) => match self.decode_data(buf)? {
                    Some(entry) => {
                        let entry_count = entry_count - 1;
                        let next_state = self.next_data_state(&f, entry_count);
                        return Ok((Some(Part::Data(entry)), next_state));
                    }
                    None => return Ok((None, Data(f, entry_count))),
                },
                End => return Ok((None, End)),
                Invalid => bail!(ErrorKind::WirePackDecode("byte stream corrupt".into())),
            }
        }
    }

    #[inline]
    fn next_history_state<P: Into<RepoPath>>(&mut self, filename: P, entry_count: u32) -> State {
        if entry_count == 0 {
            State::DataStart(filename.into())
        } else {
            State::History(filename.into(), entry_count)
        }
    }

    #[inline]
    fn next_data_state<P: Into<RepoPath>>(&mut self, filename: P, entry_count: u32) -> State {
        if entry_count == 0 {
            State::Filename
        } else {
            State::Data(filename.into(), entry_count)
        }
    }

    fn decode_filename(&mut self, buf: &mut BytesMut) -> Result<DecodeRes<RepoPath>> {
        // Notes:
        // - A zero-length filename indicates the root manifest for tree packs.
        //   (It is not allowed for file packs.)
        // - The end of the stream is marked with 10 null bytes (2 for the filename + 4 for
        //   history entry count + 4 for data entry count).
        // - This means that the buffer has to have at least 10 bytes in it at this stage.
        if buf.len() < 10 {
            return Ok(DecodeRes::None);
        }
        if &buf[..10] == WIREPACK_END {
            let _ = buf.split_to(10);
            return Ok(DecodeRes::End);
        }

        let filename_len = buf.peek_u16() as usize;
        if buf.len() < filename_len + 2 {
            return Ok(DecodeRes::None);
        }
        let _ = buf.split_to(2);
        let filename = if filename_len == 0 {
            match self.kind {
                Kind::Tree => RepoPath::root(),
                Kind::File => bail!(ErrorKind::WirePackDecode(
                    "file packs cannot contain zero-length filenames".into(),
                )),
            }
        } else {
            let mpath = buf.drain_path(filename_len).with_context(|| {
                let msg = format!("invalid filename of length {}", filename_len);
                ErrorKind::WirePackDecode(msg)
            })?;

            match self.kind {
                Kind::Tree => RepoPath::dir(mpath),
                Kind::File => RepoPath::file(mpath),
            }
            .with_context(|| ErrorKind::WirePackDecode("invalid filename".into()))?
        };

        trace!(&self.logger, "decoding entries for filename: {}", filename);

        Ok(DecodeRes::Some(filename))
    }

    fn decode_section_start(&mut self, buf: &mut BytesMut) -> Option<u32> {
        if buf.len() < 4 {
            None
        } else {
            Some(buf.drain_u32())
        }
    }

    #[inline]
    fn decode_history(&mut self, buf: &mut BytesMut) -> Result<Option<HistoryEntry>> {
        HistoryEntry::decode(buf, self.kind)
    }

    #[inline]
    fn decode_data(&mut self, buf: &mut BytesMut) -> Result<Option<DataEntry>> {
        DataEntry::decode(buf, DataEntryVersion::V1)
    }
}

#[derive(Debug, Eq, PartialEq)]
enum State {
    Filename,
    HistoryStart(RepoPath),
    History(RepoPath, u32),
    DataStart(RepoPath),
    Data(RepoPath, u32),
    End,
    Invalid,
}

impl State {
    pub fn take(&mut self) -> Self {
        mem::replace(self, State::Invalid)
    }
}

enum DecodeRes<T> {
    None,
    Some(T),
    End,
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use futures::compat::Future01CompatExt;
    use futures_old::Future;
    use futures_old::Stream;
    use slog::o;
    use slog::Discard;
    use tokio_codec::FramedRead;

    use super::*;

    #[tokio::test]
    async fn test_empty() {
        let logger = Logger::root(Discard, o!());

        // Absolutely nothing in here.
        let empty_1 = Cursor::new(WIREPACK_END);
        let unpacker = new(logger.clone(), Kind::Tree);
        let stream = FramedRead::new(empty_1, unpacker);
        let collect_fut = stream.collect();

        let fut = collect_fut
            .and_then(move |parts| {
                assert_eq!(parts, vec![Part::End]);

                // A file with no entries:
                // - filename b"\0\x03foo"
                // - history count: b"\0\0\0\0"
                // - data count: b"\0\0\0\0"
                // - next filename, end of stream: b"\0\0\0\0\0\0\0\0\0\0"
                let empty_2 = Cursor::new(b"\0\x03foo\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
                let unpacker = new(logger.clone(), Kind::File);
                let stream = FramedRead::new(empty_2, unpacker);
                stream.collect()
            })
            .map(|parts| {
                let foo_dir = RepoPath::file("foo").unwrap();
                assert_eq!(
                    parts,
                    vec![
                        Part::HistoryMeta {
                            path: foo_dir.clone(),
                            entry_count: 0,
                        },
                        Part::DataMeta {
                            path: foo_dir,
                            entry_count: 0,
                        },
                        Part::End,
                    ]
                );
            })
            .map_err(|err| panic!("{:#?}", err));

        fut.compat().await.unwrap();
    }
}


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

//! Encapsulation of data required to replay Mercurial bundles in the bookmark
//! update log.

use std::collections::HashMap;

use mercurial_types::HgChangesetId;
use mononoke_types::RawBundle2Id;
use mononoke_types::Timestamp;

pub struct BundleReplayData {
    pub bundle2_id: RawBundle2Id,
    pub timestamps: HashMap<HgChangesetId, Timestamp>,
}

impl BundleReplayData {
    pub fn new(bundle2_id: RawBundle2Id) -> Self {
        BundleReplayData {
            bundle2_id,
            timestamps: HashMap::new(),
        }
    }

    pub fn new_with_timestamps(
        bundle2_id: RawBundle2Id,
        timestamps: HashMap<HgChangesetId, Timestamp>,
    ) -> Self {
        BundleReplayData {
            bundle2_id,
            timestamps,
        }
    }
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "hg_mutation_entry crate
mercurial_thrift mercurial_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=hg_mutation_entry_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "hg_mutation_entry.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


/*
 * 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 metaconfig_types::RemoteDatabaseConfig;
use metaconfig_types::RemoteMetadataDatabaseConfig;
use mononoke_types::RepositoryId;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::SqlConnections;

use crate::store::SqlHgMutationStore;

const DEFAULT_MUTATION_CHAIN_LIMIT: usize = 500;

#[allow(unused)]
pub struct SqlHgMutationStoreBuilder {
    pub(crate) connections: SqlConnections,
    pub(crate) mutation_chain_limit: usize,
}

impl SqlConstruct for SqlHgMutationStoreBuilder {
    const LABEL: &'static str = "hg_mutations";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-hg-mutations.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            connections,
            mutation_chain_limit: DEFAULT_MUTATION_CHAIN_LIMIT,
        }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlHgMutationStoreBuilder {
    fn remote_database_config(
        remote: &RemoteMetadataDatabaseConfig,
    ) -> Option<&RemoteDatabaseConfig> {
        Some(&remote.mutation)
    }
}

impl SqlHgMutationStoreBuilder {
    pub fn with_repo_id(self, repo_id: RepositoryId) -> SqlHgMutationStore {
        SqlHgMutationStore::new(repo_id, self.connections, self.mutation_chain_limit)
    }

    pub fn with_mutation_limit(self, mutation_chain_limit: usize) -> Self {
        Self {
            mutation_chain_limit,
            ..self
        }
    }
}