/*
 * 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::fmt;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::Arc;

use once_cell::sync::Lazy;
use parking_lot::Mutex;
use regex::Regex;
use tracing::span::Attributes;
use tracing::span::Record;
use tracing::Callsite;
use tracing::Event;
use tracing::Id;
use tracing::Level;
use tracing::Metadata;
use tracing::Subscriber;

use super::*;

#[test]
fn test_callsite_span() {
    let callsite = create_callsite::<SpanKindType, _>((11, 0), || CallsiteInfo {
        name: "foo".to_string(),
        target: "bar".to_string(),
        level: Level::ERROR,
        file: Some("a.rs".to_string()),
        line: Some(10),
        module_path: Some("z::a".to_string()),
        field_names: vec!["x".to_string(), "y".to_string(), "z".to_string()],
    });
    assert_eq!(
        d(callsite.metadata()),
        "Metadata { name: foo, target: bar, level: Level(Error), module_path: z::a, location: a.rs:10, fields: {x, y, z}, callsite: _, kind: Kind(SPAN) }"
    );
    assert_eq!(callsite.identifier(), callsite.metadata().callsite());
    let log = capture(|| {
        let span = callsite.create_span(&[None, None, None]);
        span.record("y", &"yyy2");
        span.in_scope(|| {});
        let span = callsite.create_span(&[Some(Box::new("foo")), None, Some(Box::new(123))]);
        span.record("x", &123);
        span.in_scope(|| {});
    });
    assert_eq!(
        log,
        [
            "new_span(Attributes { metadata: Metadata { name: foo, target: bar, level: Level(Error), module_path: z::a, location: a.rs:10, fields: {x, y, z}, callsite: _, kind: Kind(SPAN) }, values: ValueSet { callsite: _ }, parent: Current } = 1",
            "record(Id(1), Record { values: ValueSet { y: yyy2, callsite: _ } })",
            "enter(Id(1))",
            "exit(Id(1))",
            "new_span(Attributes { metadata: Metadata { name: foo, target: bar, level: Level(Error), module_path: z::a, location: a.rs:10, fields: {x, y, z}, callsite: _, kind: Kind(SPAN) }, values: ValueSet { x: foo, z: 123, callsite: _ }, parent: Current } = 2",
            "record(Id(2), Record { values: ValueSet { x: 123, callsite: _ } })",
            "enter(Id(2))",
            "exit(Id(2))"
        ]
    );
}

#[test]
fn test_callsite_event() {
    let callsite = create_callsite::<EventKindType, _>((22, 0), || CallsiteInfo {
        name: "foo".to_string(),
        level: Level::ERROR,
        field_names: vec!["x".to_string(), "y".to_string(), "z".to_string()],
        ..Default::default()
    });
    assert_eq!(
        d(callsite.metadata()),
        "Metadata { name: foo, target: , level: Level(Error), fields: {x, y, z}, callsite: _, kind: Kind(EVENT) }"
    );
    assert_eq!(callsite.identifier(), callsite.metadata().callsite());
    let log = capture(|| {
        callsite.create_event(&[None, None, None]);
        callsite.create_event(&[Some(Box::new(12)), None, Some(Box::new("zz"))]);
        callsite.create_event(&[Some(Box::new("15"))]);
    });
    assert_eq!(
        log,
        [
            "event(Event { fields: ValueSet { callsite: _ }, metadata: Metadata { name: foo, target: , level: Level(Error), fields: {x, y, z}, callsite: _, kind: Kind(EVENT) }, parent: Current })",
            "event(Event { fields: ValueSet { x: 12, z: zz, callsite: _ }, metadata: Metadata { name: foo, target: , level: Level(Error), fields: {x, y, z}, callsite: _, kind: Kind(EVENT) }, parent: Current })",
            "event(Event { fields: ValueSet { x: 15, callsite: _ }, metadata: Metadata { name: foo, target: , level: Level(Error), fields: {x, y, z}, callsite: _, kind: Kind(EVENT) }, parent: Current })"
        ]
    );
}

#[test]
fn test_callsite_reuse() {
    let callsite1 = create_callsite::<EventKindType, _>((33, 1), CallsiteInfo::default);
    let callsite2 = create_callsite::<EventKindType, _>((33, 1), CallsiteInfo::default);
    assert_eq!(callsite1 as *const _, callsite2 as *const _);
}

#[test]
fn test_intern() {
    use crate::Intern;
    let s1 = "abc".intern();
    let s2 = "abc".to_string().intern();
    assert_eq!(s1.as_ptr(), s2.as_ptr());
}

/// Capture logs about tracing.
fn capture(f: impl FnOnce()) -> Vec<String> {
    // Prevent races since tests run in multiple threads.
    let _locked = THREAD_LOCK.lock();
    let sub = TestSubscriber::default();
    let out = sub.out.clone();
    tracing::subscriber::with_default(sub, f);
    let out = out.lock();
    out.clone()
}

/// Subscriber that captures calls to a string.
#[derive(Default)]
struct TestSubscriber {
    id: AtomicU64,
    out: Arc<Mutex<Vec<String>>>,
}
impl TestSubscriber {
    fn log(&self, s: String) {
        self.out.lock().push(normalize(&s));
    }
}
impl Subscriber for TestSubscriber {
    fn enabled(&self, _metadata: &Metadata) -> bool {
        true
    }
    fn new_span(&self, span: &Attributes) -> Id {
        let id = self.id.fetch_add(1, AcqRel) + 1;
        self.log(format!("new_span({:?} = {}", span, id));
        Id::from_u64(id)
    }
    fn record(&self, span: &Id, values: &Record) {
        self.log(format!("record({:?}, {:?})", span, values));
    }
    fn event(&self, event: &Event) {
        self.log(format!("event({:?})", event));
    }
    fn enter(&self, span: &Id) {
        self.log(format!("enter({:?})", span));
    }
    fn exit(&self, span: &Id) {
        self.log(format!("exit({:?})", span));
    }
    fn record_follows_from(&self, span: &Id, follows: &Id) {
        self.log(format!("record_follows_from({:?}, {:?})", span, follows));
    }
}

/// Debug format with some normalization.
fn d<T: fmt::Debug>(t: T) -> String {
    let s = format!("{:?}", t);
    normalize(&s)
}

fn normalize(s: &str) -> String {
    // Change "Identifier(...)" to "_". It has dynamic pointer.
    IDENTIFIER_RE.replace_all(&s, "_").replace('"', "")
}

static THREAD_LOCK: Lazy<Mutex<()>> = Lazy::new(Default::default);
static IDENTIFIER_RE: Lazy<Regex> = Lazy::new(|| Regex::new("Identifier\\([^)]*\\)").unwrap());


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

use std::io::Write;
use std::sync::Arc;

use once_cell::sync::OnceCell;
use sampling::SamplingConfig;
use serde::ser::SerializeMap;
use serde::ser::Serializer;
use serde_json::Serializer as JsonSerializer;
use tracing::subscriber::Interest;
use tracing::Event;
use tracing::Metadata;
use tracing::Subscriber;
use tracing_serde::fields::AsMap;
use tracing_subscriber::layer::Context;
use tracing_subscriber::layer::Filter;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::Layer;

pub struct SamplingLayer {
    config: &'static OnceCell<Option<Arc<SamplingConfig>>>,
}

impl SamplingLayer {
    pub fn new<S: Subscriber + for<'a> LookupSpan<'a>>() -> impl Layer<S> {
        Self::new_with_config(&sampling::CONFIG)
    }

    fn new_with_config<S: Subscriber + for<'a> LookupSpan<'a>>(
        config: &'static OnceCell<Option<Arc<SamplingConfig>>>,
    ) -> impl Layer<S> {
        Self { config }.with_filter(SamplingFilter { config })
    }
}

impl<S: Subscriber> Layer<S> for SamplingLayer {
    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
        let config = match self.config.get() {
            Some(Some(sc)) => sc.clone(),
            _ => return,
        };

        let category = match config.category(event.metadata().target()) {
            Some(v) => v,
            None => return,
        };

        let serialize = || -> std::io::Result<()> {
            let mut file = config.file();
            let mut serializer = JsonSerializer::new(&*file);

            let mut serializer = serializer.serialize_map(None)?;
            serializer.serialize_entry("category", category)?;
            serializer.serialize_entry("data", &event.field_map())?;
            serializer.end()?;

            file.write_all(b"\0")?;

            Ok(())
        };

        let _ = serialize();
    }
}

struct SamplingFilter {
    config: &'static OnceCell<Option<Arc<SamplingConfig>>>,
}

impl SamplingFilter {
    fn is_enabled(&self, meta: &Metadata<'_>) -> bool {
        match self.config.get() {
            Some(Some(sc)) => sc.category(meta.target()).is_some(),
            _ => false,
        }
    }
}

impl<S: Subscriber> Filter<S> for SamplingFilter {
    fn enabled(&self, meta: &Metadata<'_>, _: &Context<'_, S>) -> bool {
        self.is_enabled(meta)
    }

    fn callsite_enabled(&self, meta: &'static Metadata<'static>) -> Interest {
        if self.is_enabled(meta) {
            Interest::always()
        } else {
            Interest::never()
        }
    }
}

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

    use tempfile::tempdir;
    use tracing_subscriber::layer::SubscriberExt;
    use tracing_subscriber::Registry;

    use super::*;

    static TEST_CONFIG: OnceCell<Option<Arc<SamplingConfig>>> = OnceCell::new();

    #[test]
    fn test_sampling_layer() {
        let dir = tempdir().unwrap();
        let out_path = dir.path().join("out");

        {
            // Config not initialized yet.
            let sl = SamplingLayer::new_with_config(&TEST_CONFIG);

            let _subscriber = tracing::subscriber::set_default(Registry::default().with(sl));

            // layer not configured, shouldn't show up.
            tracing::info!(target: "banana", foo = "bar");

            let filepath = out_path.to_string_lossy();
            let config = BTreeMap::<&str, &str>::from([
                ("sampling.filepath", filepath.as_ref()),
                ("sampling.key.banana", "pear"),
                ("sampling.key.orange", "melon"),
            ]);

            TEST_CONFIG
                .set(Some(Arc::new(SamplingConfig::new(&config).unwrap())))
                .unwrap();

            // Should be picked up.
            tracing::info!(target: "banana", foo = "baz");

            // Target isn't sampled.
            tracing::info!(target: "apple", foo = "baz");

            // Should also be collected.
            tracing::info!(target: "orange", pi = 123);
        }

        assert_eq!(
            std::fs::read(&out_path).unwrap(),
            b"{\"category\":\"pear\",\"data\":{\"foo\":\"baz\"}}\0{\"category\":\"melon\",\"data\":{\"pi\":123}}\0",
        );
    }
}


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

//! Directory State.

use std::io::Write;
use std::path::Path;

use anyhow::anyhow;
use anyhow::Result;
use repolock::RepoLocker;
use types::hgid::NULL_ID;
use types::HgId;

use crate::serialization::Serializable;
use crate::store::BlockId;
use crate::treestate::TreeState;
use crate::ErrorKind;

/// A dirstate object. This maintains .hg/dirstate file
#[derive(Debug, PartialEq)]
pub struct Dirstate {
    pub p1: HgId,
    pub p2: HgId,

    pub tree_state: Option<TreeStateFields>,
}

#[derive(Debug, PartialEq)]
pub struct TreeStateFields {
    // Final component of treestate file. Normally a UUID.
    pub tree_filename: String,
    pub tree_root_id: BlockId,
    pub repack_threshold: Option<u64>,
}

pub fn flush(root: &Path, treestate: &mut TreeState, locker: &RepoLocker) -> Result<()> {
    if treestate.dirty() {
        tracing::debug!("flushing dirty treestate");
        let id = identity::must_sniff_dir(root)?;
        let dot_dir = root.join(id.dot_dir());
        let dirstate_path = dot_dir.join("dirstate");

        let _lock = locker.lock_working_copy(dot_dir.clone())?;

        let dirstate_input = util::file::read(&dirstate_path)?;
        let mut dirstate = Dirstate::deserialize(&mut dirstate_input.as_slice())?;

        // If the dirstate has changed since we last loaded it, don't flush since we might
        // overwrite data. For instance, if we start running 'hg status', it loads the dirstate and
        // treestate and starts updating the treestate.  Before status gets to this flush, if
        // another process, like 'hg checkout' writes to the dirstate/treestate, then if we let
        // this 'hg status' flush it's old data, we'd result in a dirty working copy where the
        // clean checkout data was thought to be dirty because we had old treestate data.
        //
        // In that situation, just return an error and the client can decide if that's ok or not.

        if let Some(dirstate_fields) = &dirstate.tree_state {
            if treestate.file_name()? != dirstate_fields.tree_filename
                || treestate.original_root_id() != dirstate_fields.tree_root_id
            {
                return Err(ErrorKind::TreestateOutOfDate.into());
            }
        }

        let metadata = treestate.metadata()?;
        dirstate.p1 = metadata
            .get("p1")
            .map_or(Ok(NULL_ID), |p| HgId::from_hex(p.as_bytes()))?;
        dirstate.p2 = metadata
            .get("p2")
            .map_or(Ok(NULL_ID), |p| HgId::from_hex(p.as_bytes()))?;
        let treestate_fields = dirstate.tree_state.as_mut().ok_or_else(|| {
            anyhow!(
                "Unable to flush treestate because dirstate is missing required treestate fields"
            )
        })?;

        let root_id = treestate.flush()?;
        treestate_fields.tree_filename = treestate.file_name()?;
        treestate_fields.tree_root_id = root_id;

        let mut dirstate_output: Vec<u8> = Vec::new();
        dirstate.serialize(&mut dirstate_output).unwrap();
        util::file::atomic_write(&dirstate_path, |file| file.write_all(&dirstate_output))
            .map_err(|e| anyhow!(e))
            .map(|_| ())
    } else {
        tracing::debug!("skipping treestate flush - it is not dirty");
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use types::hgid::NULL_ID;

    use super::*;
    use crate::serialization::Serializable;

    #[test]
    fn test_serialization() -> anyhow::Result<()> {
        let mut ds = Dirstate {
            p1: HgId::from_hex(b"93a7f768ac7506e31015dfa545b7f1475a76c4cf")?,
            p2: NULL_ID,
            tree_state: Some(TreeStateFields {
                tree_filename: "2c715852-5e8c-45bf-b1f2-236e25dd648b".to_string(),
                tree_root_id: BlockId(2236480),
                repack_threshold: None,
            }),
        };

        {
            let mut buf: Vec<u8> = Vec::new();
            ds.serialize(&mut buf).unwrap();
            assert_eq!(
                &buf,
                // I pulled this out of a real sparse .hg/dirstate.
                b"\x93\xa7\xf7\x68\xac\x75\x06\xe3\x10\x15\xdf\xa5\x45\xb7\xf1\x47\x5a\x76\xc4\xcf\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\ntreestate\n\0filename=2c715852-5e8c-45bf-b1f2-236e25dd648b\0rootid=2236480",
            );

            let got = Dirstate::deserialize(&mut buf.as_slice())?;
            assert_eq!(got, ds);
        }

        {
            (&mut ds.tree_state).as_mut().unwrap().repack_threshold = Some(123);

            let mut buf: Vec<u8> = Vec::new();
            ds.serialize(&mut buf).unwrap();
            assert_eq!(
                &buf,
                b"\x93\xa7\xf7\x68\xac\x75\x06\xe3\x10\x15\xdf\xa5\x45\xb7\xf1\x47\x5a\x76\xc4\xcf\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\ntreestate\n\0filename=2c715852-5e8c-45bf-b1f2-236e25dd648b\0rootid=2236480\0threshold=123",
            );

            let got = Dirstate::deserialize(&mut buf.as_slice())?;
            assert_eq!(got, ds);
        }

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

//! Errors.

use std::path::PathBuf;

use thiserror::Error;

#[derive(Debug, Error, PartialEq)]
pub enum ErrorKind {
    #[error("the provided store file is not a valid store file: {0}")]
    NotAStoreFile(PathBuf),
    #[error("tree version not supported: {0}")]
    UnsupportedTreeVersion(u32),
    #[error("store file version not supported: {0}")]
    UnsupportedVersion(u32),
    #[error("invalid store id: {0}")]
    InvalidStoreId(u64),
    #[error("store is read-only")]
    ReadOnlyStore,
    #[error("treedirstate is corrupt")]
    CorruptTree,
    #[error("callback error: {0}")]
    CallbackError(String),
    #[error("dirstate/treestate was out of date and therefore did not flush")]
    TreestateOutOfDate,
}


/*
 * 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::fs::File;
use std::io;
use std::io::BufWriter;
use std::io::Cursor;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io::Write;
use std::path::Path;

use fs2::FileExt;

pub trait FileSync {
    fn sync_all(&mut self) -> io::Result<()>;
}

pub trait FileLock {
    /// Lock exclusively (across threads and processes).
    /// Block if the lock was held by others.
    ///
    /// If the lock is held by `self`, then the function will just increase an
    /// internal lock counter and return immediately.
    ///
    /// Use `unlock` or `drop` to unlock.
    fn lock_exclusive(&mut self) -> io::Result<()>;

    /// Cancel one `lock_exclusive` invocation.
    ///
    /// If `lock_exclusive` is called `N` times, it requires `unlock` to be
    /// called `N` times to unlock.
    ///
    /// If `self` gets dropped, then the lock is automatically released.
    fn unlock(&mut self) -> io::Result<()>;

    /// Test if `lock_exclusive` is called more times than `unlock`.
    fn is_locked(&self) -> bool;
}

pub trait FileReadWrite:
    std::io::Read + std::io::Write + std::io::Seek + FileSync + FileLock + Send
{
}

pub struct FileReaderWriter {
    writer: BufWriter<File>,
    lock_file: Option<File>,
    locked: usize,
}

impl FileReaderWriter {
    pub fn new(writer: BufWriter<File>, path: &Path) -> io::Result<Self> {
        let lock_file = if cfg!(windows) {
            // On Windows, exclusive file lock prevents read. We only use
            // lock for protecting racy writes and want read to just work
            // regardless of locks. Use a separate lock file so locking
            // does not prevent read.
            let lock_path = path.with_extension("lock");
            let lock_file = std::fs::OpenOptions::new()
                .create(true)
                .read(true)
                .write(true)
                .open(lock_path)?;
            Some(lock_file)
        } else {
            None
        };
        Ok(Self {
            writer,
            lock_file,
            locked: 0,
        })
    }
}

impl Read for FileReaderWriter {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.writer.get_mut().read(buf)
    }
}

impl Write for FileReaderWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.writer.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.writer.flush()
    }
}

impl Seek for FileReaderWriter {
    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
        self.writer.seek(pos)
    }
}

impl FileSync for FileReaderWriter {
    fn sync_all(&mut self) -> Result<(), std::io::Error> {
        self.writer.get_mut().sync_all()
    }
}

impl FileLock for FileReaderWriter {
    fn lock_exclusive(&mut self) -> io::Result<()> {
        if self.locked == 0 {
            match self.lock_file.as_mut() {
                Some(file) => file.lock_exclusive()?,
                None => self.writer.get_mut().lock_exclusive()?,
            }
        }
        self.locked += 1;
        Ok(())
    }

    fn unlock(&mut self) -> io::Result<()> {
        if self.locked == 1 {
            match self.lock_file.as_mut() {
                Some(file) => file.unlock()?,
                None => self.writer.get_mut().unlock()?,
            }
        }
        if self.locked > 0 {
            self.locked -= 1;
        }
        Ok(())
    }

    fn is_locked(&self) -> bool {
        self.locked > 0
    }
}

impl FileReadWrite for FileReaderWriter {}

impl<T> FileSync for Cursor<T> {
    fn sync_all(&mut self) -> Result<(), std::io::Error> {
        Ok(())
    }
}

impl<T> FileReadWrite for Cursor<T>
where
    Cursor<T>: Write,
    T: std::convert::AsRef<[u8]> + Send,
{
}

impl<T> FileLock for Cursor<T> {
    fn lock_exclusive(&mut self) -> io::Result<()> {
        Ok(())
    }

    fn unlock(&mut self) -> io::Result<()> {
        Ok(())
    }

    fn is_locked(&self) -> bool {
        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.
 */

//! File State.

use bitflags::bitflags;

/// Information relating to a file in the dirstate.
#[derive(Debug, PartialEq, Copy, Clone)]
pub struct FileState {
    /// State of the file, as recorded by Mercurial.  Mercurial uses a single character to
    /// represent the current state of the file.  Only a single byte is used in the file, so only
    /// ASCII characters are valid here.
    pub state: u8,

    /// Mode (permissions) mask for the file.
    pub mode: u32,

    /// Size of the file.  Mercurial uses negative sizes for special values, so this must be
    /// signed.
    pub size: i32,

    /// Modification time of the file.
    pub mtime: i32,
}

impl FileState {
    pub fn new(state: u8, mode: u32, size: i32, mtime: i32) -> FileState {
        FileState {
            state,
            mode,
            size,
            mtime,
        }
    }
}

bitflags! {
    /// Bit flags for a file "state". Certain flags can be used together.
    ///
    /// Mapping to some Mercurial's concepts:
    ///
    /// |           | EXIST_P1 | EXIST_P2 | EXIST_NEXT | IGNORED |
    /// | added     | no       | no       | yes        | ?       |
    /// | merge     | yes      | yes      | yes        | ?       |
    /// | normal    | yes      | no       | yes        | ?       |
    /// | normal    | no       | yes      | yes        | ?       |
    /// | removed   | either one is yes   | no         | ?       |
    /// | untracked | no       | no       | no         | no      |
    /// | ignored   | no       | no       | no         | yes     |
    #[cfg_attr(test, derive(Default))]
    pub struct StateFlags: u16 {
        /// Exist in the first working parent.
        const EXIST_P1 = 1;

        /// Exist in a non-first working parent.
        const EXIST_P2 = 2;

        /// Will exist in the next commit.
        const EXIST_NEXT = 4;

        /// Explicitly marked as ignored.
        const IGNORED = 8;

        /// Known possibly changed. Need stat check.
        ///
        /// For non-watchman case, this is a quick way to get all mtime < 0 entries. aka. for
        /// calculating non-normal set quickly.
        ///
        /// For watchman case, this also includes untracked files and normal files with mtime >= 0,
        /// that are known changed during the last watchman check. Combined with a new watchman
        /// query since the recorded watchman clock, the caller can figure out all files that are
        /// possibly changed, and ignore files outside that list.
        const NEED_CHECK = 16;

        /// Marked as copied from another path.
        const COPIED = 32;
    }
}

impl StateFlags {
    pub fn to_bits(self) -> u16 {
        self.bits
    }
}

/// Information relating to a file in the dirstate, version 2.
/// Unlike V1, the `state` field is no longer a char defined by Mercurial,
/// but a bitflag. It also has a `copied` field.
#[derive(Debug, PartialEq, Clone)]
#[cfg_attr(test, derive(Default))]
pub struct FileStateV2 {
    /// Mode (permissions) mask for the file.
    pub mode: u32,

    /// Size of the file.  Mercurial uses negative sizes for special values, so this must be
    /// signed.
    pub size: i32,

    /// Modification time of the file.
    pub mtime: i32,

    /// State of the file.
    pub state: StateFlags,

    /// Path copied from.
    pub copied: Option<Box<[u8]>>,
}

impl FileStateV2 {
    pub fn is_executable(&self) -> bool {
        // Symlinks show as executable, but don't be fooled. "executable" and
        // "symlink" are mutually exclusive in the manifest, so it is just
        // confusing if we have dirstate entries that claim to be both.
        !self.is_symlink() && self.mode & 0o100 == 0o100
    }

    pub fn is_symlink(&self) -> bool {
        self.mode & 0o120000 == 0o120000
    }
}

#[cfg(test)]
impl rand::distributions::Distribution<FileStateV2> for rand::distributions::Standard {
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> FileStateV2 {
        let mode = rng.gen();
        let size = rng.gen();
        let mtime = rng.gen();
        let state = StateFlags::from_bits_truncate(rng.gen());
        let copied = if state.contains(StateFlags::COPIED) {
            Some(b"copied_source".to_vec().into_boxed_slice())
        } else {
            None
        };
        FileStateV2 {
            mode,
            size,
            mtime,
            state,
            copied,
        }
    }
}


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

//! Implementation of a store using file I/O.

use std::borrow::Cow;
use std::cell::RefCell;
use std::fs::OpenOptions;
use std::io::BufWriter;
use std::io::Cursor;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::bail;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;

use crate::errors::ErrorKind;
use crate::filereadwrite::FileReadWrite;
use crate::filereadwrite::FileReaderWriter;
use crate::store::BlockId;
use crate::store::ScopedLock;
use crate::store::Store;
use crate::store::StoreView;

// File storage format:
//
// Header: Magic string: 'appendstore\n'
//         Version:      BigEndian u32 (Current version: 1)
//
// Entries: Length:      BigEndian u32
//          Data:        "Length" bytes of data

const MAGIC_LEN: usize = 12;
const MAGIC: [u8; MAGIC_LEN] = *b"appendstore\n";
const VERSION: u32 = 1;
const HEADER_LEN: u64 = (MAGIC_LEN + 4) as u64;

/// Implementation of a store using file I/O to read and write blocks to a file.
pub struct FileStore {
    /// The underlying file.
    file: Arc<Mutex<Box<dyn FileReadWrite>>>,

    /// The position in the file to which new items will be written.
    position: u64,

    /// Whether the file handle is currently at the end of the file.  This is used to avoid seeking
    /// to the end each time a block is written, as seeking causes the BufWrite to flush, which
    /// hurts performance.  This is stored in a RefCell so that we can seek away from the end
    /// during reads.
    at_end: RefCell<bool>,

    /// True if the file is read-only.
    read_only: bool,

    /// Cache of data loaded from disk.  Used when iterating over the whole dirstate.
    cache: Option<Vec<u8>>,

    /// The path to underlying file.
    path: Option<PathBuf>,
}

impl FileStore {
    /// Create a new FileStore, overwriting any existing file.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<FileStore> {
        let path = path.as_ref();
        let writer = BufWriter::new(
            OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .truncate(true)
                .open(&path)?,
        );
        let mut file = FileReaderWriter::new(writer, path)?;
        file.write(&MAGIC)?;
        file.write_u32::<BigEndian>(VERSION)?;
        let file = Arc::new(Mutex::new(Box::new(file) as Box<dyn FileReadWrite>));
        Ok(FileStore {
            file,
            position: HEADER_LEN,
            at_end: RefCell::new(true),
            read_only: false,
            cache: None,
            path: Some(path.to_path_buf()),
        })
    }

    /// Create a new FileStore in memory. This is used solely for providing
    /// EdenFS a TreeState that is backed by a legacy dirstate.
    /// TODO: Remove once EdenFS has migrated to TreeState.
    pub fn in_memory() -> Result<FileStore> {
        let mut file = Cursor::new(Vec::new());
        file.write_all(&MAGIC)?;
        file.write_u32::<BigEndian>(VERSION)?;
        let file = Arc::new(Mutex::new(Box::new(file) as Box<dyn FileReadWrite>));
        Ok(FileStore {
            file,
            position: HEADER_LEN,
            at_end: RefCell::new(true),
            read_only: false,
            cache: None,
            path: None,
        })
    }

    /// Open an existing FileStore.  Attempts to open the file in read/write mode.  If write
    /// access is not permitted, falls back to opening the file in read-only mode.  When open
    /// in read-only mode, new blocks of data cannot be appended.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<FileStore> {
        let path = path.as_ref();
        let mut read_only = false;
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&path)
            .or_else(|_e| {
                read_only = true;
                OpenOptions::new().read(true).open(&path)
            })?;
        let mut file = FileReaderWriter::new(BufWriter::new(file), path)?;

        // Check the file header is as expected.
        let mut buffer = [0; MAGIC_LEN];
        file.read_exact(&mut buffer)
            .map_err(|_e| ErrorKind::NotAStoreFile(path.to_path_buf()))?;
        if buffer != MAGIC {
            bail!(ErrorKind::NotAStoreFile(path.to_path_buf()));
        }
        let version = file.read_u32::<BigEndian>()?;
        if version != VERSION {
            bail!(ErrorKind::UnsupportedVersion(version));
        }

        // Find the size of the file (and hence the position to write new blocks of data)
        // by seeking to the end.
        let position = file.seek(SeekFrom::End(0))?;
        let file = Arc::new(Mutex::new(Box::new(file) as Box<dyn FileReadWrite>));
        Ok(FileStore {
            file,
            position,
            at_end: RefCell::new(true),
            read_only,
            cache: None,
            path: Some(path.to_path_buf()),
        })
    }

    pub fn cache(&mut self) -> Result<()> {
        if self.cache.is_none() {
            let mut file = self.file.lock().unwrap();
            file.flush()?;
            file.seek(SeekFrom::Start(0))?;
            *self.at_end.get_mut() = false;
            let mut buffer = Vec::with_capacity(self.position as usize);
            unsafe {
                // This is safe as we've just allocated the buffer and are about to read into it.
                buffer.set_len(self.position as usize);
            }
            file.read_exact(buffer.as_mut_slice())?;
            file.seek(SeekFrom::Start(self.position))?;
            *self.at_end.get_mut() = true;
            self.cache = Some(buffer);
        }
        Ok(())
    }

    pub fn position(&self) -> u64 {
        self.position
    }

    pub fn path(&self) -> Option<&Path> {
        self.path.as_ref().map(|path| path.as_ref())
    }
}

impl Store for FileStore {
    fn append(&mut self, data: &[u8]) -> Result<BlockId> {
        if self.read_only {
            bail!(ErrorKind::ReadOnlyStore);
        }
        let id = BlockId(self.position);
        let mut file = self.file.lock().unwrap();
        debug_assert!(file.is_locked(), "writing to store requires locking");
        let at_end = self.at_end.get_mut();
        if !*at_end {
            file.seek(SeekFrom::Start(self.position))?;
            *at_end = true;
        }
        assert!(data.len() <= std::u32::MAX as usize, "data too long");
        file.write_u32::<BigEndian>(data.len() as u32)?;
        file.write_all(data)?;
        self.position += 4 + data.len() as u64;
        debug_assert!(self.position == file.seek(SeekFrom::End(0))?);
        Ok(id)
    }

    fn flush(&mut self) -> Result<()> {
        let mut file = self.file.lock().unwrap();
        debug_assert!(file.is_locked(), "flushing store requires locking");
        file.flush()?;
        file.sync_all()?;
        Ok(())
    }

    fn lock(&mut self) -> Result<ScopedLock> {
        let file = self.file.clone();
        {
            let mut mut_file = file.lock().unwrap();
            mut_file.lock_exclusive()?;
            // Seek to the end for appending.
            self.position = mut_file.seek(SeekFrom::End(0))?;
            *self.at_end.borrow_mut() = true;
        }
        let unlock = move || {
            let _ = file.lock().unwrap().unlock();
        };
        Ok(ScopedLock {
            unlock: Some(Box::new(unlock)),
        })
    }
}

impl StoreView for FileStore {
    fn read<'a>(&'a self, id: BlockId) -> Result<Cow<'a, [u8]>> {
        // Check the ID is in range.
        if id.0 < HEADER_LEN || id.0 > self.position - 4 {
            bail!(ErrorKind::InvalidStoreId(id.0));
        }

        if let Some(ref cache) = self.cache {
            if (id.0) < cache.len() as u64 {
                if (id.0) > cache.len() as u64 - 4 {
                    // The ID falls in the last 3 bytes of the cache.  This is invalid.
                    bail!(ErrorKind::InvalidStoreId(id.0));
                }
                let header_start = id.0 as usize;
                let data_start = header_start + 4;
                let size = BigEndian::read_u32(&cache[header_start..data_start]) as usize;
                if size as u64 > cache.len() as u64 - data_start as u64 {
                    // The stored size of this block exceeds the number of bytes left in the
                    // cache.  We must have been given an invalid ID.
                    bail!(ErrorKind::InvalidStoreId(id.0));
                }
                return Ok(Cow::from(&cache[data_start..data_start + size]));
            }
        }

        // Get mutable access to the file, and seek to the right location.
        let mut file = self.file.lock().unwrap();
        file.seek(SeekFrom::Start(id.0))?;
        *self.at_end.borrow_mut() = false;

        // Read the block of data from the file.
        let size = file.read_u32::<BigEndian>()?;
        if size as u64 > self.position - id.0 {
            // The stored size of this block exceeds the number of bytes left in the file.  We
            // must have been given an invalid ID.
            bail!(ErrorKind::InvalidStoreId(id.0));
        }
        let mut buffer: Vec<u8> = Vec::with_capacity(size as usize);
        unsafe {
            // This is safe as we've just allocated the buffer and are about to read into it.
            buffer.set_len(size as usize);
        }
        file.read_exact(&mut buffer[..])?;

        Ok(Cow::from(buffer))
    }
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::io::Write;

    use tempfile::tempdir;

    use crate::filestore::FileStore;
    use crate::store::BlockId;
    use crate::store::Store;
    use crate::store::StoreView;

    #[test]
    fn goodpath() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut s = FileStore::create(&p).expect("create store");
        let lock = s.lock().unwrap();
        let id1 = s.append("data block 1".as_bytes()).expect("write block 1");
        let id2 = s
            .append("data block two".as_bytes())
            .expect("write block 2");
        s.flush().expect("flush");
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes());
        drop((s, lock));
        let mut s = FileStore::open(&p).expect("open store");
        let lock = s.lock().unwrap();
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes());
        let id3 = s
            .append("third data block".as_bytes())
            .expect("write block 3");
        s.flush().expect("flush");
        drop((s, lock));
        let s = FileStore::open(p.clone()).expect("open store");
        assert_eq!(s.read(id3).expect("read 3"), "third data block".as_bytes());
        assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes());
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        dir.close().expect("clean up temp dir");
    }

    #[test]
    fn readonly() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut s = FileStore::create(&p).expect("create store");
        let lock = s.lock().unwrap();
        let id1 = s.append("data block 1".as_bytes()).expect("write block 1");
        s.flush().expect("flush");
        drop(lock);
        drop(s);
        let mut perms = fs::metadata(&p).unwrap().permissions();
        perms.set_readonly(true);
        fs::set_permissions(&p, perms).unwrap();
        let mut s = FileStore::open(&p).expect("open store");
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        assert_eq!(
            s.append("third data block".as_bytes())
                .unwrap_err()
                .to_string(),
            "store is read-only"
        );
    }

    #[test]
    fn empty_file() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&p)
            .unwrap();
        drop(file);
        assert!(
            FileStore::open(&p)
                .err()
                .expect("file should not be opened")
                .to_string()
                .starts_with("the provided store file is not a valid store file: ")
        );
    }

    #[test]
    fn invalid_file() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&p)
            .unwrap();
        file.write(b"not a store file").unwrap();
        drop(file);
        assert!(
            FileStore::open(&p)
                .err()
                .expect("file should not be opened")
                .to_string()
                .starts_with("the provided store file is not a valid store file")
        );
    }

    #[test]
    fn unsupported_version() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&p)
            .unwrap();
        // Version 0 is not supported.
        file.write(b"appendstore\n\x00\x00\x00\x00").unwrap();
        drop(file);
        assert_eq!(
            FileStore::open(&p)
                .err()
                .expect("file should not be opened")
                .to_string(),
            "store file version not supported: 0"
        );
    }

    #[test]
    fn cache() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut s = FileStore::create(&p).expect("create store");
        let lock = s.lock().unwrap();
        let id1 = s.append("data block 1".as_bytes()).expect("write block 1");
        let id2 = s
            .append("data block two".as_bytes())
            .expect("write block 2");
        s.flush().expect("flush");
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes());
        drop(lock);
        drop(s);

        let mut s = FileStore::open(&p).expect("open store");
        let _lock = s.lock().unwrap();
        s.cache().expect("can cache");
        assert_eq!(s.read(id1).expect("read 1"), "data block 1".as_bytes());
        assert_eq!(s.read(id2).expect("read 2"), "data block two".as_bytes());
        let id3 = s
            .append("third data block".as_bytes())
            .expect("write block 3");
        s.flush().expect("flush");
        assert_eq!(s.read(id3).expect("read 3"), "third data block".as_bytes());
        assert_eq!(
            s.read(BlockId(id3.0 - 2)).unwrap_err().to_string(),
            format!("invalid store id: {}", id3.0 - 2)
        );
    }

    #[test]
    fn invalid_store_ids() {
        let dir = tempdir().expect("create temp dir");
        let p = dir.path().join("store");
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&p)
            .unwrap();
        // Version 0 is not supported.
        file.write(b"appendstore\n\x00\x00\x00\x01\x00\x00\xff\xffdata")
            .unwrap();
        drop(file);
        let f = FileStore::open(&p).expect("file should be opened");
        // Store ID 2 is inside the header.
        assert_eq!(
            f.read(BlockId(2)).unwrap_err().to_string(),
            "invalid store id: 2"
        );
        // Store ID 16 has an invalid length.
        assert_eq!(
            f.read(BlockId(16)).unwrap_err().to_string(),
            "invalid store id: 16"
        );
        // Store ID 22 is within 4 bytes of the end of the file.
        assert_eq!(
            f.read(BlockId(22)).unwrap_err().to_string(),
            "invalid store id: 22"
        );
        // Store ID 64 is after the end of the file.
        assert_eq!(
            f.read(BlockId(64)).unwrap_err().to_string(),
            "invalid store id: 64"
        );
    }
}


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

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::str::FromStr;

use anyhow::anyhow;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use sha2::Digest;
use sha2::Sha256;
use types::hgid::WriteHgIdExt;
use types::HgId;
use util::file::atomic_write;

use crate::filestate::FileStateV2;
use crate::filestate::StateFlags;

/// Reader and writer functions to a legacy Eden dirstate. Note carefully that
/// a legacy Eden dirstate has a different binary format to a legacy dirstate
/// in a non-Eden repo.

const CURRENT_VERSION: u32 = 1;

const DIRSTATE_HEADER: u8 = 0x01;
const COPYMAP_HEADER: u8 = 0x02;
const CHECKSUM_HEADER: u8 = 0xFF;

const MERGE_NOT_APPLICABLE: i8 = 0;
const MERGE_BOTH_PARENTS: i8 = -1;
const MERGE_OTHER_PARENT: i8 = -2;

pub fn read_eden_dirstate(
    dirstate_path: &Path,
) -> Result<(BTreeMap<String, String>, HashMap<Box<[u8]>, FileStateV2>)> {
    let mut dirstate = File::open(dirstate_path)?;
    deserialize_dirstate(&mut dirstate)
}

pub fn write_eden_dirstate(
    dirstate_path: &Path,
    metadata: BTreeMap<String, String>,
    entries: HashMap<Box<[u8]>, FileStateV2>,
) -> Result<()> {
    let mut dirstate_bytes = Vec::new();
    serialize_dirstate(&mut dirstate_bytes, metadata, &entries)?;
    atomic_write(dirstate_path, |f| f.write_all(&dirstate_bytes))
        .map(|_| ())
        .map_err(|e| anyhow!(e))
}

struct Reader<T> {
    hasher: Sha256,
    inner: T,
}

impl<T: Read> Reader<T> {
    fn new(inner: T) -> Self {
        let hasher = Sha256::new();
        Self { hasher, inner }
    }

    fn verify_checksum(mut self) -> Result<bool> {
        let actual: [u8; 32] = self.hasher.finalize().into();
        let mut expected = [0; 32];
        self.inner.read_exact(&mut expected)?;
        Ok(actual == expected)
    }
}

impl<T: Read> Read for Reader<T> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let result = self.inner.read(buf)?;
        self.hasher.update(buf);
        Ok(result)
    }
}

struct Writer<T> {
    hasher: Sha256,
    inner: T,
}

impl<T: Write> Writer<T> {
    fn new(inner: T) -> Self {
        let hasher = Sha256::new();
        Self { hasher, inner }
    }

    fn write_checksum(mut self) -> std::io::Result<()> {
        let checksum = self.hasher.finalize();
        self.inner.write_all(&checksum)
    }
}

impl<T: Write> Write for Writer<T> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.hasher.update(buf);
        self.inner.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

// The serialization format of the dirstate is as follows:
// - The first 40 bytes are the hashes of the two parent pointers.
// - The next 4 bytes are the version number of the format.
// - The next section is the dirstate tuples. Each dirstate tuple is
//   represented as follows:
//   - The first byte is '\x01'.
//   - The second byte represents the status. It is the ASCII value of
//     'n', 'm', 'r', 'a', '?', as appropriate.
//   - The next four bytes are an unsigned integer representing mode_t.
//   - The seventh byte (which is signed) represents the merge state:
//     - 0 is NotApplicable
//     - -1 is BothParents
//     - -2 is OtherParent
//   - The next two bytes are an unsigned short representing the length of
//     the path, in bytes.
//   - The bytes of the path itself. Note that a path cannot contain \0.
// - The next section is the copymap. Each entry in the copymap is
//   represented as follows.
//   - The first byte is '\x02'.
//   - An unsigned short (two bytes) representing the length, followed by
//     that number of bytes, which constitutes the relative path name of the
//     *destination* of the copy.
//   - An unsigned short (two bytes) representing the length, followed by
//     that number of bytes, which constitutes the relative path name of the
//     *source* of the copy.
// - The last section is the checksum. Although the other tuples can be
//   interleaved or reordered without issue, the checksum must come last.
//   The checksum is a function of all of the bytes written up to this point
//   plus the \xFF header for the checksum section.
//   - The first byte is '\xFF' to distinguish it from the other fields.
//   - Because we use SHA-256 as the hash algorithm for the checksum, the
//     remaining 32 bytes are used for the hash.
fn serialize_dirstate(
    dirstate: impl Write,
    metadata: BTreeMap<String, String>,
    treestate: &HashMap<Box<[u8]>, FileStateV2>,
) -> Result<()> {
    let mut writer = Writer::new(dirstate);

    let p1 = metadata
        .get("p1")
        .ok_or_else(|| anyhow!("Treestate is missing necessary P1 node"))?;
    let p1 = HgId::from_str(p1)?;
    writer.write_hgid(&p1)?;

    if let Some(p2) = metadata.get("p2") {
        let p2 = HgId::from_str(p2)?;
        writer.write_hgid(&p2)?;
    } else {
        writer.write_hgid(HgId::null_id())?;
    }

    writer.write_u32::<BigEndian>(CURRENT_VERSION)?;

    for (path, state) in treestate {
        serialize_entry(&mut writer, path, state)?;
    }

    writer.write_u8(CHECKSUM_HEADER)?;
    writer.write_checksum()?;

    Ok(())
}

fn serialize_entry(mut dirstate: impl Write, path: &[u8], state: &FileStateV2) -> Result<()> {
    if state.state.contains(StateFlags::COPIED) {
        dirstate.write_u8(COPYMAP_HEADER)?;
        dirstate.write_u16::<BigEndian>(path.len().try_into()?)?;
        dirstate.write_all(path)?;

        let source_path = state.copied.as_ref().ok_or_else(|| {
            anyhow!(
                "Missing source path on copied entry ({})",
                String::from_utf8_lossy(path)
            )
        })?;
        dirstate.write_u16::<BigEndian>(source_path.len().try_into()?)?;
        dirstate.write_all(source_path)?;
    }

    dirstate.write_u8(DIRSTATE_HEADER)?;

    let (state_char, merge_state) = {
        let s =
            state.state & (StateFlags::EXIST_NEXT | StateFlags::EXIST_P1 | StateFlags::EXIST_P2);

        if s == StateFlags::EXIST_NEXT | StateFlags::EXIST_P1 | StateFlags::EXIST_P2 {
            (b'm', MERGE_BOTH_PARENTS)
        } else if s == StateFlags::EXIST_NEXT | StateFlags::EXIST_P1 {
            (b'n', MERGE_NOT_APPLICABLE)
        } else if s == StateFlags::EXIST_NEXT | StateFlags::EXIST_P2 {
            (b'n', MERGE_OTHER_PARENT)
        } else if s == StateFlags::EXIST_NEXT {
            (b'a', MERGE_BOTH_PARENTS)
        } else if s == StateFlags::EXIST_P1 | StateFlags::EXIST_P2 {
            (b'r', MERGE_BOTH_PARENTS)
        } else if s == StateFlags::EXIST_P2 {
            (b'r', MERGE_OTHER_PARENT)
        } else if s == StateFlags::EXIST_P1 {
            (b'r', MERGE_NOT_APPLICABLE)
        } else if s.is_empty() {
            (b'?', MERGE_NOT_APPLICABLE)
        } else {
            return Err(anyhow!(
                "Unhandled treestate -> dirstate conversion ({:?})",
                s
            ));
        }
    };

    dirstate.write_u8(state_char)?;
    dirstate.write_u32::<BigEndian>(0)?;
    dirstate.write_i8(merge_state)?;
    dirstate.write_u16::<BigEndian>(path.len().try_into()?)?;
    dirstate.write_all(path)?;
    Ok(())
}

fn deserialize_dirstate(
    dirstate: impl Read,
) -> Result<(BTreeMap<String, String>, HashMap<Box<[u8]>, FileStateV2>)> {
    let mut reader = Reader::new(dirstate);
    let mut p1 = [0; 20];
    reader.read_exact(&mut p1)?;
    let p1 = HgId::from(&p1);

    let mut p2_bytes = [0; 20];
    reader.read_exact(&mut p2_bytes)?;
    let p2 = HgId::from(&p2_bytes);

    let mut metadata = BTreeMap::new();
    metadata.insert("p1".to_string(), p1.to_string());
    if !p2.is_null() {
        metadata.insert("p2".to_string(), p2.to_string());
    }

    let _version = reader.read_u32::<BigEndian>()?;

    let mut entries: HashMap<Box<[u8]>, FileStateV2> = HashMap::new();
    while let Some((path, mut state)) = deserialize_entry(&mut reader)? {
        if let Some(existing_state) = entries.get_mut(&path) {
            state.state |= existing_state.state;
            if existing_state.copied.is_some() {
                state.copied = Some(existing_state.copied.take().unwrap());
            }
        }
        entries.insert(path, state);
    }

    if !reader.verify_checksum()? {
        return Err(anyhow!(
            "Dirstate contained incorrect checksum: likely a corrupt dirstate"
        ));
    }

    Ok((metadata, entries))
}

fn deserialize_entry(mut dirstate: impl Read) -> Result<Option<(Box<[u8]>, FileStateV2)>> {
    let header = dirstate.read_u8()?;
    if header == CHECKSUM_HEADER {
        // Reached checksum, no more entries in dirstate
        return Ok(None);
    }

    if header == COPYMAP_HEADER {
        let dest_size = dirstate.read_u16::<BigEndian>()?;
        let dest_path = read_path(&mut dirstate, dest_size)?;

        let source_size = dirstate.read_u16::<BigEndian>()?;
        let source_path = read_path(&mut dirstate, source_size)?;

        return Ok(Some((
            dest_path,
            FileStateV2 {
                mode: 0,
                size: 0,
                mtime: 0,
                state: StateFlags::COPIED,
                copied: Some(source_path),
            },
        )));
    }

    if header != DIRSTATE_HEADER {
        return Err(anyhow!("Unexpected header value: {}", header));
    }

    let state_char = dirstate.read_u8()?;
    let mode = dirstate.read_u32::<BigEndian>()?;
    let merge_state = dirstate.read_i8()?;
    let size = dirstate.read_u16::<BigEndian>()?;

    let state = match state_char {
        b'n' if merge_state == MERGE_OTHER_PARENT => StateFlags::EXIST_NEXT | StateFlags::EXIST_P2,
        b'n' => StateFlags::EXIST_NEXT | StateFlags::EXIST_P1,
        b'm' => StateFlags::EXIST_NEXT | StateFlags::EXIST_P1 | StateFlags::EXIST_P2,
        b'r' if merge_state == MERGE_BOTH_PARENTS => StateFlags::EXIST_P1 | StateFlags::EXIST_P2,
        b'r' if merge_state == MERGE_OTHER_PARENT => StateFlags::EXIST_P2,
        b'r' => StateFlags::EXIST_P1,
        b'a' => StateFlags::EXIST_NEXT,
        b'?' => StateFlags::empty(),
        s => return Err(anyhow!("Unexpected state: {}", s)),
    };

    let path = read_path(&mut dirstate, size)?;
    Ok(Some((
        path,
        FileStateV2 {
            mode,
            size: 0,
            mtime: 0,
            state,
            copied: None,
        },
    )))
}

fn read_path(dirstate: impl Read, size: u16) -> Result<Box<[u8]>> {
    let mut path = Vec::with_capacity(size.into());
    let mut reader = Read::take(dirstate, size.into());
    reader.read_to_end(&mut path)?;
    Ok(path.into_boxed_slice())
}

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

    use pretty_assertions::assert_eq;
    use sha2::Digest;
    use sha2::Sha256;

    use super::deserialize_dirstate;
    use super::serialize_dirstate;
    use crate::filestate::FileStateV2;
    use crate::filestate::StateFlags;
    use crate::metadata::Metadata;
    use crate::serialization::Serializable;

    // DIRSTATE corresponds to the following working copy state:
    // $ hg st -C
    // M modified
    // A added
    // A copy_dest
    //   copy_source
    // A move_after
    //   move_before
    // R move_before
    // R removed
    // ! missing
    // ? untracked
    static DIRSTATE: [u8; 215] = [
        0xc9, 0x4c, 0x85, 0xea, 0x63, 0xba, 0x78, 0x41, 0x64, 0xf8, 0x5d, 0x3c, 0x8c, 0x89, 0xb2,
        0x57, 0x06, 0xec, 0x7a, 0x7f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01,
        0x72, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x72, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x64,
        0x01, 0x61, 0x00, 0x00, 0x00, 0x00, 0xff, 0x00, 0x05, 0x61, 0x64, 0x64, 0x65, 0x64, 0x01,
        0x61, 0x00, 0x00, 0x00, 0x00, 0xff, 0x00, 0x0a, 0x6d, 0x6f, 0x76, 0x65, 0x5f, 0x61, 0x66,
        0x74, 0x65, 0x72, 0x01, 0x72, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0b, 0x6d, 0x6f, 0x76,
        0x65, 0x5f, 0x62, 0x65, 0x66, 0x6f, 0x72, 0x65, 0x01, 0x61, 0x00, 0x00, 0x00, 0x00, 0xff,
        0x00, 0x09, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x64, 0x65, 0x73, 0x74, 0x02, 0x00, 0x0a, 0x6d,
        0x6f, 0x76, 0x65, 0x5f, 0x61, 0x66, 0x74, 0x65, 0x72, 0x00, 0x0b, 0x6d, 0x6f, 0x76, 0x65,
        0x5f, 0x62, 0x65, 0x66, 0x6f, 0x72, 0x65, 0x02, 0x00, 0x09, 0x63, 0x6f, 0x70, 0x79, 0x5f,
        0x64, 0x65, 0x73, 0x74, 0x00, 0x0b, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x6f, 0x75, 0x72,
        0x63, 0x65, 0xff, 0x39, 0x39, 0x1a, 0x7d, 0x3b, 0x0b, 0xa4, 0xe7, 0xe2, 0xc3, 0xfd, 0x1b,
        0xae, 0xd3, 0x75, 0x11, 0x00, 0x14, 0x06, 0x84, 0x1b, 0xe9, 0x68, 0xd9, 0x1e, 0x65, 0xf3,
        0xf6, 0x43, 0x2a, 0xbe, 0xbf,
    ];

    fn treestate() -> (BTreeMap<String, String>, HashMap<Box<[u8]>, FileStateV2>) {
        let mut metadata_bytes = b"p1=c94c85ea63ba784164f85d3c8c89b25706ec7a7f".as_slice();
        let metadata = Metadata::deserialize(&mut metadata_bytes).unwrap();
        let entries = vec![
            (
                "added",
                FileStateV2 {
                    mode: 0,
                    size: 0,
                    mtime: 0,
                    state: StateFlags::EXIST_NEXT,
                    copied: None,
                },
            ),
            (
                "copy_dest",
                FileStateV2 {
                    mode: 0,
                    size: 0,
                    mtime: 0,
                    state: StateFlags::EXIST_NEXT | StateFlags::COPIED,
                    copied: Some(b"copy_source".to_vec().into_boxed_slice()),
                },
            ),
            (
                "move_after",
                FileStateV2 {
                    mode: 0,
                    size: 0,
                    mtime: 0,
                    state: StateFlags::EXIST_NEXT | StateFlags::COPIED,
                    copied: Some(b"move_before".to_vec().into_boxed_slice()),
                },
            ),
            (
                "move_before",
                FileStateV2 {
                    mode: 0,
                    size: 0,
                    mtime: 0,
                    state: StateFlags::EXIST_P1,
                    copied: None,
                },
            ),
            (
                "removed",
                FileStateV2 {
                    mode: 0,
                    size: 0,
                    mtime: 0,
                    state: StateFlags::EXIST_P1,
                    copied: None,
                },
            ),
        ]
        .iter()
        .map(|(key, value)| (key.as_bytes().to_vec().into_boxed_slice(), value.clone()))
        .collect::<HashMap<_, _>>();

        (metadata.0, entries)
    }

    #[test]
    fn deserialize_test() {
        let (metadata, entries) = deserialize_dirstate(&mut DIRSTATE.as_slice()).unwrap();
        let (expected_metadata, expected_entries) = treestate();
        assert_eq!(metadata, expected_metadata);

        let mut entries = entries
            .into_iter()
            .map(|(path, state)| (String::from_utf8_lossy(&path).into_owned(), state))
            .collect::<Vec<_>>();
        entries.sort_unstable_by_key(|(path, _)| path.clone());

        let mut expected_entries = expected_entries
            .into_iter()
            .map(|(path, state)| (String::from_utf8_lossy(&path).into_owned(), state))
            .collect::<Vec<_>>();
        expected_entries.sort_unstable_by_key(|(path, _)| path.clone());

        assert_eq!(entries, expected_entries);
    }

    #[test]
    fn serialize_test() {
        let (in_metadata, in_entries) = treestate();
        let mut dirstate = Vec::new();
        serialize_dirstate(&mut dirstate, in_metadata.clone(), &in_entries).unwrap();

        let (out_metadata, out_entries) = deserialize_dirstate(&mut dirstate.as_slice()).unwrap();

        assert_eq!(in_metadata, out_metadata);

        let mut in_entries = in_entries
            .into_iter()
            .map(|(path, state)| (String::from_utf8_lossy(&path).into_owned(), state))
            .collect::<Vec<_>>();
        in_entries.sort_unstable_by_key(|(path, _)| path.clone());

        let mut out_entries = out_entries
            .into_iter()
            .map(|(path, state)| (String::from_utf8_lossy(&path).into_owned(), state))
            .collect::<Vec<_>>();
        out_entries.sort_unstable_by_key(|(path, _)| path.clone());

        assert_eq!(in_entries, out_entries);
    }

    #[test]
    fn write_checksum_test() {
        let (metadata, entries) = treestate();
        let mut dirstate = Vec::new();
        serialize_dirstate(&mut dirstate, metadata, &entries).unwrap();
        assert_eq!(dirstate.len(), 215);

        let mut hasher = Sha256::new();
        hasher.update(&dirstate[0..183]);
        let expected_checksum: [u8; 32] = hasher.finalize().into();

        assert_eq!(expected_checksum, dirstate[183..215]);
    }
}


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

//! treestate - Tree-based State.
//!
//! The tree state stores a map from paths to a lightweight structure, and provides efficient
//! lookups.  In particular, for each file in the tree, it stores the mode flags, size, mtime, and
//! whether deleted or not, etc. These can be useful for source control to determine if the file
//! is tracked, or has changed, etc.

pub mod dirstate;
pub mod errors;
mod filereadwrite;
pub mod filestate;
pub mod filestore;
mod legacy_eden_dirstate;
pub mod metadata;
pub mod root;
pub mod serialization;
pub mod store;
pub mod tree;
pub mod treestate;
pub mod vecmap;
pub mod vecstack;

pub use crate::errors::*;


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

use std::collections::BTreeMap;

#[derive(Debug, PartialEq)]
pub struct Metadata(pub BTreeMap<String, String>);

#[cfg(test)]
mod test {
    use super::*;
    use crate::serialization::Serializable;

    #[test]
    fn test_serialization() -> anyhow::Result<()> {
        let roundtrip = |data: Vec<(&str, &str)>, exp: &[u8]| -> anyhow::Result<()> {
            let meta = Metadata(
                data.iter()
                    .map(|(k, v)| (k.to_string(), v.to_string()))
                    .collect(),
            );
            let mut buf: Vec<u8> = Vec::new();
            meta.serialize(&mut buf)?;
            assert_eq!(buf, exp);
            assert_eq!(Metadata::deserialize(&mut &buf[..])?, meta);

            Ok(())
        };

        roundtrip(Vec::new(), b"")?;
        roundtrip(vec![("foo", "bar")], b"foo=bar")?;
        roundtrip(vec![("abc", "123"), ("xyz", "456")], b"abc=123\0xyz=456")?;
        roundtrip(
            vec![("abc", "why=dothis"), ("ä½ ", "å¥½")],
            b"abc=why=dothis\0\xE4\xBD\xA0=\xE5\xA5\xBD",
        )?;

        Ok(())
    }

    #[test]
    fn test_serializate_empty_value() -> anyhow::Result<()> {
        let meta = Metadata(BTreeMap::from([
            ("foo".to_string(), "bar".to_string()),
            ("empty".to_string(), "".to_string()),
        ]));
        let mut buf: Vec<u8> = Vec::new();
        meta.serialize(&mut buf)?;
        assert_eq!(buf, b"foo=bar");

        Ok(())
    }

    #[test]
    fn test_serialize_errors() {
        {
            let meta = Metadata(BTreeMap::from([("foo=bar".to_string(), "baz".to_string())]));
            let mut buf: Vec<u8> = Vec::new();
            assert!(meta.serialize(&mut buf).is_err());
        }

        {
            let meta = Metadata(BTreeMap::from([(
                "foo".to_string(),
                "baz\0qux".to_string(),
            )]));
            let mut buf: Vec<u8> = Vec::new();
            assert!(meta.serialize(&mut buf).is_err());
        }

        {
            let meta = Metadata(BTreeMap::from([(
                "foo\0oops".to_string(),
                "bar".to_string(),
            )]));
            let mut buf: Vec<u8> = Vec::new();
            assert!(meta.serialize(&mut buf).is_err());
        }
    }

    #[test]
    fn test_deserialize_errors() {
        assert!(Metadata::deserialize(&mut &b"foo"[..]).is_err());
        assert!(Metadata::deserialize(&mut &b"foo=bar\0baz"[..]).is_err());
        assert!(Metadata::deserialize(&mut &b"\0"[..]).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 crate::store::BlockId;

/// `TreeStateRoot` contains block id to the root `Tree`, and other metadata.
#[derive(Default)]
pub struct TreeStateRoot {
    version: u32,
    file_count: u32,
    tree_block_id: BlockId,
    metadata: Box<[u8]>,
    dirty: bool,
}

impl TreeStateRoot {
    pub fn new(version: u32, file_count: u32, tree_block_id: BlockId, metadata: Box<[u8]>) -> Self {
        TreeStateRoot {
            version,
            file_count,
            tree_block_id,
            metadata,
            dirty: false,
        }
    }

    pub fn dirty(&self) -> bool {
        self.dirty
    }

    pub fn version(&self) -> u32 {
        self.version
    }

    pub fn set_version(&mut self, value: u32) {
        self.version = value;
        self.dirty = true;
    }

    pub fn file_count(&self) -> u32 {
        self.file_count
    }

    pub fn set_file_count(&mut self, value: u32) {
        self.file_count = value;
        self.dirty = true;
    }

    pub fn tree_block_id(&self) -> BlockId {
        self.tree_block_id
    }

    pub fn set_tree_block_id(&mut self, value: BlockId) {
        self.tree_block_id = value;
        self.dirty = true;
    }

    pub fn metadata(&self) -> &Box<[u8]> {
        &self.metadata
    }

    pub fn set_metadata(&mut self, value: Box<[u8]>) {
        self.metadata = value;
        self.dirty = 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.
 */

//! Trait for serialization and deserialization of tree data.

use std::collections::BTreeMap;
use std::hash::Hasher;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Cursor;
use std::io::Read;
use std::io::Write;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use twox_hash::XxHash;
use types::hgid::ReadHgIdExt;
use vlqencoding::VLQDecode;
use vlqencoding::VLQEncode;

use crate::dirstate::Dirstate;
use crate::dirstate::TreeStateFields;
use crate::errors::*;
use crate::filestate::FileState;
use crate::filestate::FileStateV2;
use crate::filestate::StateFlags;
use crate::metadata::Metadata;
use crate::root::TreeStateRoot;
use crate::store::BlockId;
use crate::tree::AggregatedState;
use crate::tree::Key;
use crate::tree::Node;
use crate::tree::NodeEntry;
use crate::tree::NodeEntryMap;

pub trait Serializable
where
    Self: Sized,
{
    /// Serialize the storable data to a `Write` stream.
    fn serialize(&self, w: &mut dyn Write) -> Result<()>;

    /// Deserialize a new data item from a `Read` stream.
    fn deserialize(r: &mut dyn Read) -> Result<Self>;
}

impl Serializable for FileState {
    /// Write a file entry to the store.
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_u8(self.state)?;
        w.write_vlq(self.mode)?;
        w.write_vlq(self.size)?;
        w.write_vlq(self.mtime)?;
        Ok(())
    }

    /// Read an entry from the store.
    fn deserialize(r: &mut dyn Read) -> Result<FileState> {
        let state = r.read_u8()?;
        let mode = r.read_vlq()?;
        let size = r.read_vlq()?;
        let mtime = r.read_vlq()?;
        Ok(FileState {
            state,
            mode,
            size,
            mtime,
        })
    }
}

impl Serializable for AggregatedState {
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_vlq(self.union.to_bits())?;
        w.write_vlq(self.intersection.to_bits())?;
        Ok(())
    }

    /// Read an entry from the store.
    fn deserialize(r: &mut dyn Read) -> Result<AggregatedState> {
        let state: u16 = r.read_vlq()?;
        let union = StateFlags::from_bits_truncate(state);
        let state: u16 = r.read_vlq()?;
        let intersection = StateFlags::from_bits_truncate(state);
        Ok(AggregatedState {
            union,
            intersection,
        })
    }
}

impl Serializable for Box<[u8]> {
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_vlq(self.len())?;
        w.write_all(&self)?;

        Ok(())
    }

    fn deserialize(r: &mut dyn Read) -> Result<Self> {
        let len: usize = r.read_vlq()?;
        let mut buf = vec![0; len];
        r.read_exact(&mut buf)?;
        Ok(buf.into_boxed_slice())
    }
}

impl Serializable for FileStateV2 {
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_vlq(self.state.to_bits())?;
        w.write_vlq(self.mode)?;
        w.write_vlq(self.size)?;
        w.write_vlq(self.mtime)?;

        if self.state.contains(StateFlags::COPIED) {
            if let &Some(ref copied) = &self.copied {
                copied.serialize(w)?;
            } else {
                panic!("COPIED flag set without copied path");
            }
        }
        Ok(())
    }

    fn deserialize(r: &mut dyn Read) -> Result<FileStateV2> {
        let state: u16 = r.read_vlq()?;
        let state = StateFlags::from_bits_truncate(state);
        let mode = r.read_vlq()?;
        let size = r.read_vlq()?;
        let mtime = r.read_vlq()?;
        let copied = if state.contains(StateFlags::COPIED) {
            Some(Box::<[u8]>::deserialize(r)?)
        } else {
            None
        };

        Ok(FileStateV2 {
            state,
            mode,
            size,
            mtime,
            copied,
        })
    }
}

/// Deserialize a single entry in a node's entry map.  Returns the name and the entry.
fn deserialize_node_entry<T>(r: &mut dyn Read) -> Result<(Key, NodeEntry<T>)>
where
    T: Serializable + Clone,
{
    let entry_type = r.read_u8()?;
    match entry_type {
        b'f' => {
            // File entry.
            let data = T::deserialize(r)?;
            let name_len = r.read_vlq()?;
            let mut name = Vec::with_capacity(name_len);
            unsafe {
                // Safe as we've just allocated the buffer and are about to read into it.
                name.set_len(name_len);
            }
            r.read_exact(name.as_mut_slice())?;
            Ok((name.into_boxed_slice(), NodeEntry::File(data)))
        }
        b'd' => {
            // Directory entry.
            let id = r.read_vlq()?;
            let name_len = r.read_vlq()?;
            let mut name = Vec::with_capacity(name_len);
            unsafe {
                // Safe as we've just allocated the buffer and are about to read into it.
                name.set_len(name_len);
            }
            r.read_exact(name.as_mut_slice())?;
            Ok((
                name.into_boxed_slice(),
                NodeEntry::Directory(Node::open(BlockId(id))),
            ))
        }
        _ => {
            bail!(ErrorKind::CorruptTree);
        }
    }
}

impl<T: Serializable + Clone> Serializable for NodeEntryMap<T> {
    fn deserialize(r: &mut dyn Read) -> Result<NodeEntryMap<T>> {
        let count = r.read_vlq()?;
        let mut entries = NodeEntryMap::with_capacity(count);
        for _i in 0..count {
            let (name, entry) = deserialize_node_entry(r)?;
            entries.insert_hint_end(name, entry);
        }
        Ok(entries)
    }

    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_vlq(self.len())?;
        for (name, entry) in self.iter() {
            match entry {
                &NodeEntry::File(ref file) => {
                    w.write_u8(b'f')?;
                    file.serialize(w)?;
                }
                &NodeEntry::Directory(ref node) => {
                    w.write_u8(b'd')?;
                    w.write_vlq(node.id.unwrap().0)?;
                }
            }
            w.write_vlq(name.len())?;
            w.write_all(name)?;
        }
        Ok(())
    }
}

#[inline]
fn xxhash<T: AsRef<[u8]>>(buf: T) -> u64 {
    let mut xx = XxHash::default();
    xx.write(buf.as_ref());
    xx.finish()
}

impl Serializable for TreeStateRoot {
    fn deserialize(r: &mut dyn Read) -> Result<Self> {
        let checksum = r.read_u64::<BigEndian>()?;
        let mut buf = Vec::new();
        r.read_to_end(&mut buf)?;

        if xxhash(&buf) != checksum {
            bail!(ErrorKind::CorruptTree);
        }

        let mut cur = Cursor::new(buf);
        let version = cur.read_vlq()?;
        if version != 0 {
            bail!(ErrorKind::UnsupportedTreeVersion(version));
        }

        let tree_block_id = BlockId(cur.read_vlq()?);
        let file_count = cur.read_vlq()?;
        let metadata = Box::<[u8]>::deserialize(&mut cur)?;

        Ok(TreeStateRoot::new(
            version,
            file_count,
            tree_block_id,
            metadata,
        ))
    }

    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        let mut buf = Vec::new();
        buf.write_vlq(self.version())?;
        buf.write_vlq(self.tree_block_id().0)?;
        buf.write_vlq(self.file_count())?;
        self.metadata().serialize(&mut buf)?;
        w.write_u64::<BigEndian>(xxhash(&buf))?;
        w.write_all(&buf)?;
        Ok(())
    }
}

impl Serializable for StateFlags {
    fn deserialize(r: &mut dyn Read) -> Result<Self> {
        let v = r.read_vlq()?;
        Ok(Self::from_bits_truncate(v))
    }

    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        Ok(w.write_vlq(self.to_bits())?)
    }
}

impl Serializable for Metadata {
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        for (i, (k, v)) in self.0.iter().enumerate() {
            if v.is_empty() {
                continue;
            }

            if k.contains(&['=', '\0']) || v.contains('\0') {
                return Err(anyhow!("invalid metadata: {k:?} -> {v:?}"));
            }

            write!(w, "{k}={v}")?;

            if i < self.0.len() - 1 {
                w.write_all(&[0])?;
            }
        }
        Ok(())
    }

    fn deserialize(r: &mut dyn Read) -> Result<Self> {
        let mut buf_reader = BufReader::new(r);
        let mut data = BTreeMap::<String, String>::new();
        loop {
            let mut key = Vec::<u8>::new();
            if buf_reader.read_until(b'=', &mut key)? == 0 {
                break;
            }
            if key.pop() != Some(b'=') {
                return Err(anyhow!("metadata key missing '=': {key:?}"));
            }

            let mut value = Vec::<u8>::new();
            buf_reader.read_until(b'\0', &mut value)?;
            if value.last() == Some(&0) {
                value.pop();
            }

            data.insert(String::from_utf8(key)?, String::from_utf8(value)?);
        }
        Ok(Self(data))
    }
}

const DIRSTATE_TREESTATE_HEADER: &[u8] = b"\ntreestate\n\0";

impl Serializable for Dirstate {
    fn serialize(&self, w: &mut dyn Write) -> Result<()> {
        w.write_all(self.p1.as_ref())?;
        w.write_all(self.p2.as_ref())?;

        let ts_fields = match &self.tree_state {
            Some(ts) => ts,
            None => {
                bail!("tree state fields are required for serializing dirstate")
            }
        };

        w.write_all(DIRSTATE_TREESTATE_HEADER)?;

        let mut meta = Metadata(BTreeMap::from([
            ("filename".to_string(), ts_fields.tree_filename.clone()),
            ("rootid".to_string(), ts_fields.tree_root_id.0.to_string()),
        ]));
        if let Some(threshold) = ts_fields.repack_threshold {
            meta.0
                .insert("threshold".to_string(), threshold.to_string());
        }

        meta.serialize(w)?;

        Ok(())
    }

    /// Best effort parsing of the dirstate. For non-treestate
    /// dirstates we only parse the parents.
    fn deserialize(r: &mut dyn Read) -> Result<Self> {
        let mut ds = Self {
            p1: r.read_hgid()?,
            p2: r.read_hgid()?,
            tree_state: None,
        };

        let mut header_buf = [0; DIRSTATE_TREESTATE_HEADER.len()];
        if r.read_exact(&mut header_buf).is_ok() && header_buf == DIRSTATE_TREESTATE_HEADER {
            let mut meta = Metadata::deserialize(r)?;
            ds.tree_state = Some(TreeStateFields {
                tree_filename: meta
                    .0
                    .remove("filename")
                    .ok_or_else(|| anyhow!("no treestate 'filename' in dirstate"))?,
                tree_root_id: BlockId(
                    meta.0
                        .remove("rootid")
                        .ok_or_else(|| anyhow!("no treestate 'rootid' in dirstate"))?
                        .parse()
                        .context("error parsing dirstate rootid")?,
                ),
                repack_threshold: meta
                    .0
                    .remove("threshold")
                    .map(|t| {
                        t.parse()
                            .with_context(|| format!("error parsing dirstate threshold {:?}", t))
                    })
                    .transpose()?,
            });
        }

        Ok(ds)
    }
}


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

//! Trait defining an append-only storage system.

use std::borrow::Cow;

use anyhow::bail;
use anyhow::Result;

use crate::errors::ErrorKind;

#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash)]
pub struct BlockId(pub u64);

/// Append-only storage.  Blocks of data may be stored in an instance of a Store.  Once written,
/// blocks are immutable.
pub trait Store {
    /// Append a new block of data to the store.  Returns the ID of the block.  Note that blocks
    /// may be buffered until `flush` is called.
    fn append(&mut self, data: &[u8]) -> Result<BlockId>;

    /// Flush all appended blocks to the backing store.
    fn flush(&mut self) -> Result<()>;

    /// Lock the store exclusively to block other writers.
    /// Block if the lock was taken by others.
    ///
    /// Unlock when the returned `ScopedLock` gets dropped.
    fn lock(&mut self) -> Result<ScopedLock>;
}

/// Lock for a scope.
pub struct ScopedLock {
    /// Specify how to unlock.
    pub unlock: Option<Box<dyn FnOnce()>>,
}

impl Drop for ScopedLock {
    fn drop(&mut self) {
        let unlock = self.unlock.take();
        if let Some(unlock) = unlock {
            (unlock)();
        }
    }
}

/// Read-only view of a store.
pub trait StoreView {
    /// Read a block of data from the store.  Blocks are immutiable, so the result may be a
    /// reference to the internal copy of the data in the store.
    fn read<'a>(&'a self, id: BlockId) -> Result<Cow<'a, [u8]>>;
}

/// Null implementation of a store.  This cannot be used to store new blocks of data, and returns
/// an error if any attempts to read are made.
pub struct NullStore;

impl NullStore {
    pub fn new() -> NullStore {
        NullStore
    }
}

impl Store for NullStore {
    fn append(&mut self, _: &[u8]) -> Result<BlockId> {
        panic!("append to NullStore");
    }

    fn flush(&mut self) -> Result<()> {
        Ok(())
    }

    fn lock(&mut self) -> Result<ScopedLock> {
        Ok(ScopedLock { unlock: None })
    }
}

impl StoreView for NullStore {
    fn read<'a>(&'a self, id: BlockId) -> Result<Cow<'a, [u8]>> {
        bail!(ErrorKind::InvalidStoreId(id.0))
    }
}

#[cfg(test)]
pub mod tests {
    use std::borrow::Cow;
    use std::collections::HashMap;

    use super::*;
    use crate::store::BlockId;
    use crate::store::Store;
    use crate::store::StoreView;

    /// Define a Store to be used in tests.  This doesn't store the data on disk, but rather
    /// keeps it in memory in a hash map.
    pub struct MapStore {
        next_id: BlockId,
        data: HashMap<BlockId, Vec<u8>>,
    }

    impl MapStore {
        pub fn new() -> MapStore {
            // Initial ID is set to 24 to simulate a header.
            MapStore {
                next_id: BlockId(24),
                data: HashMap::new(),
            }
        }
    }

    impl Store for MapStore {
        fn append(&mut self, data: &[u8]) -> Result<BlockId> {
            let id = self.next_id;
            self.data.insert(id, data.to_vec());
            self.next_id.0 += data.len() as u64;
            Ok(id)
        }

        fn flush(&mut self) -> Result<()> {
            Ok(())
        }

        fn lock(&mut self) -> Result<ScopedLock> {
            Ok(ScopedLock { unlock: None })
        }
    }

    impl StoreView for MapStore {
        fn read<'a>(&'a self, id: BlockId) -> Result<Cow<'a, [u8]>> {
            match self.data.get(&id) {
                Some(data) => Ok(Cow::from(data.as_slice())),
                None => bail!(ErrorKind::InvalidStoreId(id.0)),
            }
        }
    }

    #[test]
    fn basic_test() {
        let mut ms = MapStore::new();
        let key1 = ms.append("12345".as_bytes()).expect("append key1");
        let key2 = ms.append("67890".as_bytes()).expect("append key2");
        ms.flush().expect("flush");
        assert_eq!(ms.read(key2).unwrap(), "67890".as_bytes());
        assert_eq!(ms.read(key1).unwrap(), "12345".as_bytes());
        assert_eq!(
            ms.read(BlockId(999)).unwrap_err().to_string(),
            "invalid store id: 999"
        );
    }
}