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

use std::cell::Cell;
use std::collections::Bound;
use std::io::Cursor;
use std::io::Read;
use std::io::Write;

use anyhow::Result;
use types::RepoPath;

use crate::filestate::FileState;
use crate::filestate::FileStateV2;
use crate::filestate::StateFlags;
use crate::serialization::Serializable;
use crate::store::BlockId;
use crate::store::Store;
use crate::store::StoreView;
use crate::vecmap::VecMap;
use crate::vecstack::VecStack;

/// A node entry is an entry in a directory, either a file or another directory.
#[derive(Debug)]
pub(crate) enum NodeEntry<T> {
    Directory(Node<T>),
    File(T),
}

/// Filenames are buffers of bytes.  They're not stored in Strings as they may not be UTF-8.
pub type Key = Box<[u8]>;
pub type KeyRef<'a> = &'a [u8];

/// Result of a "visitor" function.  Specify whether a file is changed or not. Used to mark
/// parent directory as "dirty" recursively.
#[derive(Debug, Eq, PartialEq, Clone, Copy)]
pub enum VisitorResult {
    NotChanged,
    Changed,
}

/// Store the node entries in an ordered map from name to node entry.
pub(crate) type NodeEntryMap<T> = VecMap<Key, NodeEntry<T>>;

/// The aggregated state. Useful for fast decision about whether to visit a directory recursively
/// or not.
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct AggregatedState {
    pub union: StateFlags,
    pub intersection: StateFlags,
}

impl AggregatedState {
    pub fn merge(&self, rhs: AggregatedState) -> AggregatedState {
        AggregatedState {
            union: self.union | rhs.union,
            intersection: self.intersection & rhs.intersection,
        }
    }

    /// Adjust `intersection` so it does not have bits outside `union`.
    pub fn normalized(&self) -> AggregatedState {
        AggregatedState {
            union: self.union,
            intersection: self.intersection & self.union,
        }
    }
}

impl From<StateFlags> for AggregatedState {
    fn from(state: StateFlags) -> Self {
        AggregatedState {
            union: state,
            intersection: state,
        }
    }
}

impl Default for AggregatedState {
    fn default() -> Self {
        AggregatedState {
            union: StateFlags::empty(),
            intersection: StateFlags::all(),
        }
    }
}

/// The contents of a directory.
#[derive(Debug)]
pub struct Node<T> {
    /// The ID of the directory in the store.  If None, this directory has not yet been
    /// written to the back-end store in its current state.
    pub(crate) id: Option<BlockId>,

    /// The set of files and directories in this directory, indexed by their name.  If None,
    /// then the ID must not be None, and the entries are yet to be loaded from the back-end
    /// store.
    pub(crate) entries: Option<NodeEntryMap<T>>,

    /// Aggregated state flags. This is useful for quickly test whether there is a file matching
    /// given state or not in this tree (recursively). `None` means it is not calculated yet.
    aggregated_state: Cell<Option<AggregatedState>>,

    /// Optional cache about name filtering result. See `FilteredKeyCache` and `get_filtered_key`
    /// for details.
    filtered_keys: Option<FilteredKeyCache>,
}

/// A map from keys that have been filtered through a case-folding filter function to the
/// original key.  This is used for case-folded look-ups.  Filtered values are cached here.
///
/// The map is associated with an integer, the identity of the filter function used. So the
/// map could be invalidated correctly if filter function changes.
///
/// If a filtered key maps to multiple keys. All of them are stored, sorted by alphabet order.
#[derive(Debug)]
struct FilteredKeyCache {
    filter_id: u64,
    map: VecMap<Key, Vec<Key>>,
}

/// The root of the tree.  The count of files in the tree is maintained for fast size
/// determination.
pub struct Tree<T> {
    root: Node<T>,
    file_count: u32,
}

/// Utility enum for recursing through trees.
enum PathRecurse<'name, 'node, T: 'node> {
    Directory(KeyRef<'name>, KeyRef<'name>, &'node mut Node<T>),
    ExactDirectory(KeyRef<'name>, &'node mut Node<T>),
    MissingDirectory(KeyRef<'name>, KeyRef<'name>),
    File(KeyRef<'name>, &'node mut T),
    MissingFile(KeyRef<'name>),
    ConflictingFile(KeyRef<'name>, KeyRef<'name>, &'node mut T),
}

/// Splits a key into the first path element and the remaining path elements (if any).  Doesn't
/// split the key if it just contains an exact file or directory name.
fn split_key<'a>(key: KeyRef<'a>) -> (KeyRef<'a>, Option<KeyRef<'a>>) {
    if key.len() == 0 {
        return (key, None);
    }
    // Skip the last character.  Even if it's a '/' we don't want to split on it.
    for (index, value) in key.iter().take(key.len() - 1).enumerate() {
        if *value == b'/' {
            return (&key[..index + 1], Some(&key[index + 1..]));
        }
    }
    (key, None)
}

/// Splits a key into the first path element and the remaining path elements (if any).  Splits
/// the key even if it contains an exact file or directory name.
fn split_key_exact<'a>(key: KeyRef<'a>) -> (KeyRef<'a>, Option<KeyRef<'a>>) {
    for (index, value) in key.iter().enumerate() {
        if *value == b'/' {
            return (&key[..index + 1], Some(&key[index + 1..]));
        }
    }
    (key, None)
}

/// Compatiblity layer - difference between `FileState` and `FileStateV2`
pub trait CompatExt<T> {
    /// Load extra fields. Extends `load`.
    fn load_ext(&self, data: &mut dyn Read) -> Result<()>;

    /// Write extra fields. Extends `write_entries`.
    fn write_ext(&self, writer: &mut dyn Write) -> Result<()>;

    /// Calculate `aggregated_state` if it's not calculated yet.
    fn calculate_aggregated_state(&self) -> AggregatedState;

    /// Calculate `aggregated_state` if it's not calculated yet, recursively.
    fn calculate_aggregated_state_recursive(
        &mut self,
        _store: &dyn StoreView,
    ) -> Result<AggregatedState>;
}

impl CompatExt<FileState> for Node<FileState> {
    fn load_ext(&self, _: &mut dyn Read) -> Result<()> {
        Ok(())
    }

    fn write_ext(&self, _: &mut dyn Write) -> Result<()> {
        Ok(())
    }

    fn calculate_aggregated_state(&self) -> AggregatedState {
        AggregatedState::default().normalized()
    }

    fn calculate_aggregated_state_recursive(
        &mut self,
        _store: &dyn StoreView,
    ) -> Result<AggregatedState> {
        Ok(AggregatedState::default().normalized())
    }
}

impl CompatExt<FileStateV2> for Node<FileStateV2> {
    fn write_ext(&self, writer: &mut dyn Write) -> Result<()> {
        let state = self.calculate_aggregated_state();
        state.serialize(writer)?;
        Ok(())
    }

    fn calculate_aggregated_state(&self) -> AggregatedState {
        match self.aggregated_state.get() {
            None => {
                let state = self
                    .entries
                    .as_ref()
                    .expect("entries should exist")
                    .iter()
                    .fold(AggregatedState::default(), |acc, (_, x)| match x {
                        &NodeEntry::Directory(ref x) => {
                            acc.merge(x.aggregated_state.get().expect("should be ready now"))
                        }
                        &NodeEntry::File(ref x) => acc.merge(x.state.into()),
                    });
                self.aggregated_state.set(Some(state));
                state
            }
            Some(state) => state,
        }
    }

    fn calculate_aggregated_state_recursive(
        &mut self,
        store: &dyn StoreView,
    ) -> Result<AggregatedState> {
        self.load_aggregated_state(store)?;
        if self.aggregated_state.get().is_none() {
            for (_name, entry) in self.load_entries(store)?.iter_mut() {
                if let &mut NodeEntry::Directory(ref mut node) = entry {
                    node.calculate_aggregated_state_recursive(store)?;
                }
            }
        }
        Ok(self.calculate_aggregated_state())
    }

    fn load_ext(&self, data: &mut dyn Read) -> Result<()> {
        self.aggregated_state
            .set(Some(AggregatedState::deserialize(data)?));
        Ok(())
    }
}

impl<T: Serializable + Clone> Node<T> {
    /// Create a new empty Node.  This has no ID as it is not yet written to the store.
    fn new() -> Node<T> {
        Node {
            id: None,
            entries: Some(NodeEntryMap::new()),
            filtered_keys: None,
            aggregated_state: Cell::new(None),
        }
    }

    /// Create a new Node for an existing entry in the store.  The entries are not loaded until
    /// the load method is called.
    pub(crate) fn open(id: BlockId) -> Node<T> {
        Node {
            id: Some(id),
            entries: None,
            filtered_keys: None,
            aggregated_state: Cell::new(None),
        }
    }

    /// Return the aggregated file state that is the "bitwise-or" of all subentries,
    /// or `None` if it's not calculated.
    pub fn get_aggregated_state(&self) -> Option<AggregatedState> {
        self.aggregated_state.get()
    }

    /// Return if the node is changed in memory and is not written to disk yet.
    pub fn is_changed(&self) -> bool {
        self.id.is_none()
    }
}

impl<T: Serializable + Clone> Node<T>
where
    Self: CompatExt<T>,
{
    /// Attempt to load a node from a store.
    fn load(&mut self, store: &dyn StoreView) -> Result<()> {
        if self.entries.is_some() {
            // Already loaded.
            return Ok(());
        }
        let id = self.id.expect("Node must have a valid ID to be loaded");
        let data = store.read(id)?;
        let mut cur = Cursor::new(data);
        self.load_ext(&mut cur)?;
        self.entries = Some(NodeEntryMap::<T>::deserialize(&mut cur)?);
        Ok(())
    }

    /// Load only the aggregated_state without entries.
    fn load_aggregated_state(&mut self, store: &dyn StoreView) -> Result<()> {
        if self.entries.is_some() {
            // No need to load aggregated_state, since it was loaded before.
            return Ok(());
        } else if self.aggregated_state.get().is_some() {
            // Already loaded.
            return Ok(());
        }
        let id = self
            .id
            .expect("Node must have a valid ID to load aggregated_state");
        let data = store.read(id)?;
        let mut cur = Cursor::new(data);
        self.load_ext(&mut cur)?;
        Ok(())
    }

    /// Get access to the node entries, ensuring they are loaded first.
    #[inline]
    fn load_entries(&mut self, store: &dyn StoreView) -> Result<&mut NodeEntryMap<T>> {
        self.load(store)?;
        let entries = self
            .entries
            .as_mut()
            .expect("Entries should have been populated by loading");
        Ok(entries)
    }

    /// Writes all entries for this node to the store.  Any child directory entries must have
    /// had IDs assigned to them.
    fn write_entries(&mut self, store: &mut dyn Store) -> Result<()> {
        let mut data = Vec::new();
        self.write_ext(&mut data)?;
        {
            let entries = self
                .entries
                .as_ref()
                .expect("Node should have entries populated before writing out.");
            entries.serialize(&mut data)?;
        }
        self.id = Some(store.append(&data)?);
        Ok(())
    }

    /// Perform a full write of the node and its children to the store.  Old entries are
    /// loaded from the old_store before being written back to the new store.
    fn write_full(&mut self, store: &mut dyn Store, old_store: &dyn StoreView) -> Result<()> {
        // Write out all the child nodes.
        for (_name, entry) in self.load_entries(old_store)?.iter_mut() {
            if let &mut NodeEntry::Directory(ref mut node) = entry {
                node.write_full(store, old_store)?;
            }
        }
        // Write out this node.
        self.write_entries(store)
    }

    /// Perform a delta write of the node and its children to the store.  Entries that are
    /// already in the store will not be written again.
    fn write_delta<S: Store + StoreView>(&mut self, store: &mut S) -> Result<()> {
        if self.id.is_none() {
            // This node has been modified, write out a new copy of any children who have
            // also changed.  The entries list must already have been populated when the node
            // was modified, so no need to load it here.
            {
                let entries = self
                    .entries
                    .as_mut()
                    .expect("Node should have entries populated if it was modified.");
                for (_name, entry) in entries.iter_mut() {
                    if let &mut NodeEntry::Directory(ref mut node) = entry {
                        node.write_delta(store)?;
                    }
                }
            }

            // This is needed. Sometimes subentries have `id` set but not aggregated_state.
            // That happens with `Node::open`.
            {
                self.calculate_aggregated_state_recursive(store)?;
            }

            // Write out this node.
            self.write_entries(store)
        } else {
            // This node and its descendents have not been modified.
            Ok(())
        }
    }

    /// Visit all of the files in under this node, by calling the visitor function on each one.
    ///
    /// `visit_dir` will be called to test if a directory is worth visiting or not.
    /// `visit_file` will be called to test if a file is worth visiting or not.
    ///
    /// The visitor can change the file, in which case it must return `VisitorResult::Changed` so
    /// parent nodes can be marked as "changed" correctly.
    ///
    /// Return a `VisitorResult` indicating whether this node is changed or not so nodes can be
    /// marked "changed" recursively.
    fn visit<'a, F, VD, VF>(
        &'a mut self,
        store: &dyn StoreView,
        path: &mut VecStack<'a, [u8]>,
        visitor: &mut F,
        visit_dir: &VD,
        visit_file: &VF,
    ) -> Result<VisitorResult>
    where
        F: FnMut(&Vec<KeyRef>, &mut T) -> Result<VisitorResult>,
        VD: Fn(&Vec<KeyRef>, &Node<T>) -> bool,
        VF: Fn(&Vec<KeyRef>, &T) -> bool,
    {
        // visit_dir wants aggregated_state to be populated to do quick filtering.
        self.load_aggregated_state(store)?;
        if !visit_dir(path.as_ref(), self) {
            return Ok(VisitorResult::NotChanged);
        }

        let mut result = VisitorResult::NotChanged;

        let entries: &mut NodeEntryMap<T> = {
            self.load_entries(store)?;
            self.entries.as_mut().unwrap()
        };

        for (name, entry) in entries.iter_mut() {
            let mut path = path.push(name);
            let sub_result = match entry {
                &mut NodeEntry::Directory(ref mut node) => {
                    node.visit(store, &mut path, visitor, visit_dir, visit_file)?
                }
                &mut NodeEntry::File(ref mut file) => {
                    if visit_file(path.as_ref(), file) {
                        visitor(path.as_ref(), file)?
                    } else {
                        VisitorResult::NotChanged
                    }
                }
            };
            if sub_result == VisitorResult::Changed {
                result = VisitorResult::Changed;
            }
        }

        if result == VisitorResult::Changed {
            self.id = None;
            self.aggregated_state.set(None);
        }
        Ok(result)
    }

    /// Get the first file in the subtree under this node.  If the subtree is not empty, returns a
    /// pair containing the path to the file as a reversed vector of key references for each path
    /// element, and a reference to the file.
    fn get_first<'node>(
        &'node mut self,
        store: &dyn StoreView,
    ) -> Result<Option<(Vec<KeyRef<'node>>, &'node T)>> {
        for (name, entry) in self.load_entries(store)?.iter_mut() {
            match entry {
                &mut NodeEntry::Directory(ref mut node) => {
                    if let Some((mut next_name, next_file)) = node.get_first(store)? {
                        next_name.push(name);
                        return Ok(Some((next_name, next_file)));
                    }
                }
                &mut NodeEntry::File(ref file) => {
                    return Ok(Some((vec![name], file)));
                }
            }
        }
        Ok(None)
    }

    /// Get the next file after a particular file in the tree.  Returns a pair containing the path
    /// to the file as a reversed vector of key references for each path element, and a reference
    /// to the file, or None if there are no more files.
    fn get_next<'node>(
        &'node mut self,
        store: &dyn StoreView,
        name: KeyRef,
    ) -> Result<Option<(Vec<KeyRef<'node>>, &'node T)>> {
        // Find the entry within this list, and what the remainder of the path is.
        let (elem, mut path) = split_key(name);

        // Get the next entry after the current one.  We need to look inside directories as we go.
        // The subpath we obtained from split_key is only relevant if we are looking inside the
        // directory the path refers to.
        for (entry_name, entry) in self
            .load_entries(store)?
            .range_mut((Bound::Included(elem), Bound::Unbounded))
        {
            match entry {
                &mut NodeEntry::Directory(ref mut node) => {
                    // The entry is a directory, check inside it.
                    if elem != &entry_name[..] {
                        // This directory is not the one we were initially looking inside.  We
                        // have moved on past that directory, so the rest of the path is no
                        // longer relevant.
                        path = None
                    }
                    let next = if let Some(path) = path {
                        // Find the next file after the given subpath.
                        node.get_next(store, path)?
                    } else {
                        // Find the first file in this subtree.
                        node.get_first(store)?
                    };
                    if let Some((mut next_name, next_file)) = next {
                        next_name.push(entry_name);
                        return Ok(Some((next_name, next_file)));
                    }
                }
                &mut NodeEntry::File(ref file) => {
                    // This entry is a file.  Skip over it if it is the original file.
                    if elem != &entry_name[..] {
                        return Ok(Some((vec![entry_name], file)));
                    }
                }
            }
        }
        Ok(None)
    }

    /// Utility function for recursing through subdirectories.  Returns the appropriate
    /// PathRecurse variant for the current position in the file tree given by name.
    fn path_recurse<'name, 'node>(
        &'node mut self,
        store: &dyn StoreView,
        name: KeyRef<'name>,
    ) -> Result<PathRecurse<'name, 'node, T>> {
        let (elem, path) = split_key(name);
        let res = if let Some(path) = path {
            // The name is for a subdirectory.
            match self.load_entries(store)?.get_mut(elem) {
                Some(&mut NodeEntry::Directory(ref mut node)) => {
                    PathRecurse::Directory(elem, path, node)
                }
                Some(&mut NodeEntry::File(ref mut file)) => {
                    PathRecurse::ConflictingFile(elem, path, file)
                }
                None => PathRecurse::MissingDirectory(elem, path),
            }
        } else {
            // The name is for a file or directory in this directory.
            match self.load_entries(store)?.get_mut(elem) {
                Some(&mut NodeEntry::Directory(ref mut node)) => {
                    PathRecurse::ExactDirectory(elem, node)
                }
                Some(&mut NodeEntry::File(ref mut file)) => PathRecurse::File(elem, file),
                None => PathRecurse::MissingFile(elem),
            }
        };
        Ok(res)
    }

    /// Get a file's state.
    fn get<'node>(
        &'node mut self,
        store: &dyn StoreView,
        name: KeyRef,
    ) -> Result<Option<&'node T>> {
        match self.path_recurse(store, name)? {
            PathRecurse::Directory(_dir, path, node) => node.get(store, path),
            PathRecurse::ExactDirectory(_dir, _node) => Ok(None),
            PathRecurse::MissingDirectory(_dir, _path) => Ok(None),
            PathRecurse::File(_name, file) => Ok(Some(file)),
            PathRecurse::MissingFile(_name) => Ok(None),
            PathRecurse::ConflictingFile(_name, _path, _file) => Ok(None),
        }
    }

    /// Returns true if the given path is a directory.
    fn has_dir(&mut self, store: &dyn StoreView, name: KeyRef) -> Result<bool> {
        // This directory exists, without checking entries.
        if name == b"/" {
            return Ok(true);
        }
        match self.path_recurse(store, name)? {
            PathRecurse::Directory(_dir, path, node) => node.has_dir(store, path),
            PathRecurse::ExactDirectory(_dir, _node) => Ok(true),
            PathRecurse::MissingDirectory(_dir, _path) => Ok(false),
            PathRecurse::File(_name, _file) => Ok(false),
            PathRecurse::MissingFile(_name) => Ok(false),
            PathRecurse::ConflictingFile(_name, _path, _file) => Ok(false),
        }
    }

    /// Returns `Some(AggregatedState)` if the given path is a directory, or `None`.
    fn get_dir(&mut self, store: &dyn StoreView, name: KeyRef) -> Result<Option<AggregatedState>> {
        if name == b"/" {
            return Ok(Some(self.calculate_aggregated_state_recursive(store)?));
        }

        match self.path_recurse(store, name)? {
            PathRecurse::Directory(_dir, path, node) => node.get_dir(store, path),
            PathRecurse::ExactDirectory(_dir, node) => node.get_dir(store, b"/"),
            PathRecurse::MissingDirectory(_dir, _path) => Ok(None),
            PathRecurse::File(_name, _file) => Ok(None),
            PathRecurse::MissingFile(_name) => Ok(None),
            PathRecurse::ConflictingFile(_name, _path, _file) => Ok(None),
        }
    }

    /// Add a file to the node.  The name may contain a path, in which case sufficient
    /// subdirectories are updated to add or update the file.
    fn add(&mut self, store: &dyn StoreView, name: KeyRef, info: &T) -> Result<bool> {
        // Construct a RepoPath so we match the core path validation logic.
        let _ = RepoPath::from_utf8(name)?;
        let (new_entry, file_added) = match self.path_recurse(store, name)? {
            PathRecurse::Directory(_dir, path, node) => {
                // The file is in a subdirectory.  Add it to the subdirectory.
                let file_added = node.add(store, path, info)?;
                (None, file_added)
            }
            PathRecurse::ExactDirectory(_dir, _node) => {
                panic!("Adding file which matches the name of a directory.");
            }
            PathRecurse::MissingDirectory(dir, path) => {
                // The file is in a new subdirectory.  Create the directory and add the file.
                let mut node = Node::new();
                let file_added = node.add(store, path, info)?;
                (
                    Some((dir.to_vec().into_boxed_slice(), NodeEntry::Directory(node))),
                    file_added,
                )
            }
            PathRecurse::File(_name, file) => {
                // The file is in this directory.  Update it.
                file.clone_from(info);
                (None, false)
            }
            PathRecurse::MissingFile(ref name) => {
                // The file should be in this directory.  Add it.
                if name.is_empty() || name[name.len() - 1] == b'/' {
                    panic!("Adding file with tailing slash");
                }
                (
                    Some((
                        name.to_vec().into_boxed_slice(),
                        NodeEntry::File(info.clone()),
                    )),
                    true,
                )
            }
            PathRecurse::ConflictingFile(_name, _path, _file) => {
                panic!("Adding file with path prefix that matches the name of a file.")
            }
        };
        if let Some((new_key, new_entry)) = new_entry {
            self.load_entries(store)?.insert(new_key, new_entry);
            self.filtered_keys = None;
        }
        // Reset aggregated_state so it needs recalculation.
        self.aggregated_state.set(None);
        self.id = None;
        Ok(file_added)
    }

    /// Remove a file from the node.  The name may contain a path, in which case sufficient
    /// subdirectories are updated to remove the file.
    ///
    /// Returns a pair of booleans (file_removed, now_empty) indicating whether the file
    /// was removed, and whether the diectory is now empty.
    fn remove(&mut self, store: &dyn StoreView, name: KeyRef) -> Result<(bool, bool)> {
        let (file_removed, remove_entry) = match self.path_recurse(store, name)? {
            PathRecurse::Directory(dir, path, node) => {
                let (file_removed, now_empty) = node.remove(store, path)?;
                (file_removed, if now_empty { Some(dir) } else { None })
            }
            PathRecurse::ExactDirectory(_dir, _node) => (false, None),
            PathRecurse::MissingDirectory(_dir, _path) => (false, None),
            PathRecurse::File(name, _file) => (true, Some(name)),
            PathRecurse::MissingFile(_name) => (false, None),
            PathRecurse::ConflictingFile(_name, _path, _file) => (false, None),
        };
        if let Some(entry) = remove_entry {
            self.load_entries(store)?.remove(entry);
            self.filtered_keys = None;
            self.id = None;
        }
        if file_removed {
            self.aggregated_state.set(None);
            self.id = None;
        }
        Ok((file_removed, self.load_entries(store)?.is_empty()))
    }

    /// Performs a key lookup using filtered keys.
    ///
    /// Applies the filter function to each key in the node, then returns the real key that
    /// matches the name provided.  The name may contain a path, in which case the subdirectories
    /// of this node are also queried.
    ///
    /// Returns a list of reversed vector of key references for each path element.
    ///
    /// `filter_id` should be different for logically different `filter` functions. It is used for
    /// cache invalidation.
    fn get_filtered_key<'a, F>(
        &'a mut self,
        store: &dyn StoreView,
        name: KeyRef,
        filter: &mut F,
        filter_id: u64,
    ) -> Result<Vec<Vec<Key>>>
    where
        F: FnMut(KeyRef) -> Result<Key>,
    {
        let (elem, path) = split_key(name);
        let elem = trim_separator(elem);
        if self.filtered_keys.is_none()
            || self.filtered_keys.as_ref().unwrap().filter_id != filter_id
        {
            let new_map = {
                let entries = self.load_entries(store)?;
                let mut new_map: VecMap<Key, Vec<Key>> = VecMap::with_capacity(entries.len());
                for (k, _v) in entries.iter() {
                    let filtered = filter(trim_separator(k))?;
                    let inserted = match new_map.get_mut(&filtered) {
                        Some(keys) => {
                            keys.push(k.to_vec().into_boxed_slice());
                            true
                        }
                        None => false,
                    };
                    if !inserted {
                        new_map.insert(filtered, vec![k.to_vec().into_boxed_slice()]);
                    }
                }
                new_map
            };
            self.filtered_keys = Some(FilteredKeyCache {
                filter_id,
                map: new_map,
            });
        }
        if let Some(path) = path {
            let mut result = Vec::new();
            if let Some(mapped_elems) = self.filtered_keys.as_ref().unwrap().map.get(elem) {
                let mut entries = self.entries.as_mut().unwrap();
                let entries = &mut entries;
                for mapped_elem in mapped_elems {
                    if let Some(&mut NodeEntry::Directory(ref mut node)) =
                        entries.get_mut(mapped_elem)
                    {
                        for mut mapped_path in
                            node.get_filtered_key(store, path, filter, filter_id)?
                        {
                            mapped_path.push(mapped_elem.clone());
                            result.push(mapped_path);
                        }
                    }
                }
            }
            Ok(result)
        } else {
            Ok(self
                .filtered_keys
                .as_ref()
                .unwrap()
                .map
                .get(elem)
                .cloned()
                .unwrap_or_else(|| Vec::new())
                .iter()
                .map(|e| vec![trim_separator(e).to_vec().into_boxed_slice()])
                .collect())
        }
    }

    /// Checks if a path is suitable for completion, in that it contains a file that matches
    /// the acceptable conditions.
    fn path_complete_check<FA>(&mut self, store: &dyn StoreView, acceptable: &FA) -> Result<bool>
    where
        FA: Fn(&T) -> bool,
    {
        for (_name, entry) in self.load_entries(store)?.iter_mut() {
            match entry {
                &mut NodeEntry::Directory(ref mut node) => {
                    if node.path_complete_check(store, acceptable)? {
                        return Ok(true);
                    }
                }
                &mut NodeEntry::File(ref mut file) => {
                    if acceptable(file) {
                        return Ok(true);
                    }
                }
            }
        }
        Ok(false)
    }

    /// Attempt to complete a path prefix.
    ///
    /// If full_paths is true, calls the visitor for every file that matches the prefix where
    /// the file's state returns true when passed to acceptable.
    ///
    /// If full_paths is false, the first matching directory that matches the prefix is used as
    /// long as there is at least one file under the directory that is acceptable.
    fn path_complete<'a, FA, FV>(
        &'a mut self,
        store: &dyn StoreView,
        path: &mut VecStack<'a, [u8]>,
        prefix: KeyRef<'a>,
        full_paths: bool,
        acceptable: &FA,
        visitor: &mut FV,
    ) -> Result<()>
    where
        FA: Fn(&T) -> bool,
        FV: FnMut(&Vec<KeyRef>) -> Result<()>,
    {
        let (elem, subpath) = split_key_exact(prefix);
        if let Some(subpath) = subpath {
            // Prefix part is for a directory, so look for that directory.
            let entry = self.load_entries(store)?.get_mut(elem);
            if let Some(&mut NodeEntry::Directory(ref mut node)) = entry {
                let mut path = path.push(elem);
                node.path_complete(store, &mut path, subpath, full_paths, acceptable, visitor)?;
            }
        } else {
            // Prefix part is for a entry in this directory.  Iterate across all matching entries.
            for (entry_name, entry) in self
                .load_entries(store)?
                .range_mut((Bound::Included(elem), Bound::Unbounded))
            {
                if entry_name.len() < elem.len() || &entry_name[..elem.len()] != elem {
                    // This entry is no longer a prefix.
                    break;
                }
                match entry {
                    &mut NodeEntry::Directory(ref mut node) => {
                        if full_paths {
                            let mut path = path.push(entry_name);
                            // The entry is a directory, and the caller has asked for full paths.
                            // Visit every entry inside the directory.
                            let mut visit_adapter = |filepath: &Vec<KeyRef>, state: &mut T| {
                                if acceptable(state) {
                                    visitor(filepath)?;
                                }
                                Ok(VisitorResult::NotChanged)
                            };
                            node.visit(
                                store,
                                &mut path,
                                &mut visit_adapter,
                                &|_, _| true,
                                &|_, _| true,
                            )?;
                        } else {
                            // The entry is a directory, and the caller has asked for matching
                            // directories.  Check there is an acceptable file under the
                            // directory.
                            if node.path_complete_check(store, acceptable)? {
                                let path = path.push(entry_name);
                                visitor(path.as_ref())?;
                            }
                        }
                    }
                    &mut NodeEntry::File(ref mut file) => {
                        // This entry is a file.
                        if acceptable(file) {
                            let path = path.push(entry_name);
                            visitor(path.as_ref())?;
                        }
                    }
                }
            }
        }
        Ok(())
    }
}

impl<T: Serializable + Clone> Tree<T>
where
    Node<T>: CompatExt<T>,
{
    /// Create a new empty tree.
    pub fn new() -> Tree<T> {
        Tree {
            root: Node::new(),
            file_count: 0,
        }
    }

    /// Create a tree that references an existing root node.
    pub fn open(root_id: BlockId, file_count: u32) -> Tree<T> {
        Tree {
            root: Node::open(root_id),
            file_count,
        }
    }

    /// Clear all entries in the tree.
    pub fn clear(&mut self) {
        self.root = Node::new();
        self.file_count = 0;
    }

    pub fn root_id(&self) -> Option<BlockId> {
        self.root.id
    }

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

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

    pub fn write_full(
        &mut self,
        store: &mut dyn Store,
        old_store: &dyn StoreView,
    ) -> Result<BlockId> {
        self.root.write_full(store, old_store)?;
        Ok(self.root.id.unwrap())
    }

    pub fn write_delta<S: Store + StoreView>(&mut self, store: &mut S) -> Result<BlockId> {
        self.root.write_delta(store)?;
        Ok(self.root.id.unwrap())
    }

    pub fn get<'a>(&'a mut self, store: &dyn StoreView, name: KeyRef) -> Result<Option<&'a T>> {
        Ok(self.root.get(store, name)?)
    }

    pub fn visit_advanced<F, VD, VF>(
        &mut self,
        store: &dyn StoreView,
        visitor: &mut F,
        visit_dir: &VD,
        visit_file: &VF,
    ) -> Result<()>
    where
        F: FnMut(&Vec<KeyRef>, &mut T) -> Result<VisitorResult>,
        VD: Fn(&Vec<KeyRef>, &Node<T>) -> bool,
        VF: Fn(&Vec<KeyRef>, &T) -> bool,
    {
        let mut path = Vec::new();
        let mut path = VecStack::new(&mut path);
        self.root
            .visit(store, &mut path, visitor, visit_dir, visit_file)?;
        Ok(())
    }

    pub fn visit<F>(&mut self, store: &dyn StoreView, visitor: &mut F) -> Result<()>
    where
        F: FnMut(&Vec<KeyRef>, &mut T) -> Result<VisitorResult>,
    {
        self.visit_advanced(store, visitor, &|_, _| true, &|_, _| true)
    }

    pub fn visit_changed<F>(&mut self, store: &dyn StoreView, visitor: &mut F) -> Result<()>
    where
        F: FnMut(&Vec<KeyRef>, &mut T) -> Result<VisitorResult>,
    {
        self.visit_advanced(
            store,
            visitor,
            &|_, dir: &Node<T>| dir.is_changed(),
            &|_, _| true,
        )
    }

    pub fn get_first<'a>(&'a mut self, store: &dyn StoreView) -> Result<Option<(Key, &'a T)>> {
        Ok(self.root.get_first(store)?.map(|(mut path, file)| {
            path.reverse();
            (path.concat().into_boxed_slice(), file)
        }))
    }

    pub fn get_next<'a>(
        &'a mut self,
        store: &dyn StoreView,
        name: KeyRef,
    ) -> Result<Option<(Key, &'a T)>> {
        Ok(self.root.get_next(store, name)?.map(|(mut path, file)| {
            path.reverse();
            (path.concat().into_boxed_slice(), file)
        }))
    }

    pub fn has_dir(&mut self, store: &dyn StoreView, name: KeyRef) -> Result<bool> {
        Ok(self.root.has_dir(store, name)?)
    }

    pub fn get_dir(
        &mut self,
        store: &dyn StoreView,
        name: KeyRef,
    ) -> Result<Option<AggregatedState>> {
        Ok(self.root.get_dir(store, name)?)
    }

    pub fn add(&mut self, store: &dyn StoreView, name: KeyRef, file: &T) -> Result<()> {
        if self.root.add(store, name, file)? {
            self.file_count += 1;
        }
        Ok(())
    }

    pub fn remove(&mut self, store: &dyn StoreView, name: KeyRef) -> Result<bool> {
        let removed = self.root.remove(store, name)?.0;
        if removed {
            assert!(self.file_count > 0);
            self.file_count -= 1;
        }
        Ok(removed)
    }

    pub fn get_filtered_key<F>(
        &mut self,
        store: &dyn StoreView,
        name: KeyRef,
        filter: &mut F,
        filter_id: u64,
    ) -> Result<Vec<Key>>
    where
        F: FnMut(KeyRef) -> Result<Key>,
    {
        Ok(self
            .root
            .get_filtered_key(store, name, filter, filter_id)?
            .iter_mut()
            .map(|path| {
                path.reverse();
                path.concat().into_boxed_slice()
            })
            .collect())
    }

    pub fn path_complete<FA, FV>(
        &mut self,
        store: &dyn StoreView,
        prefix: KeyRef,
        full_paths: bool,
        acceptable: &FA,
        visitor: &mut FV,
    ) -> Result<()>
    where
        FA: Fn(&T) -> bool,
        FV: FnMut(&Vec<KeyRef>) -> Result<()>,
    {
        let mut path = Vec::new();
        let mut path = VecStack::new(&mut path);
        self.root
            .path_complete(store, &mut path, prefix, full_paths, acceptable, visitor)
    }
}

fn trim_separator(path: &[u8]) -> &[u8] {
    // Strip trailing slashes
    if path.last() == Some(&b'/') {
        &path[..path.len() - 1]
    } else {
        path
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::store::tests::MapStore;
    use crate::store::NullStore;

    // Test files in order.  Note lexicographic ordering of file9 and file10.
    static TEST_FILES: [(&[u8], u32, i32, i32); 16] = [
        (b"dirA/subdira/file1", 0o644, 1, 10001),
        (b"dirA/subdira/file2", 0o644, 2, 10002),
        (b"dirA/subdirb/file3", 0o644, 3, 10003),
        (b"dirB/subdira/file4", 0o644, 4, 10004),
        (b"dirB/subdira/subsubdirx/file5", 0o644, 5, 10005),
        (b"dirB/subdira/subsubdiry/file6", 0o644, 6, 10006),
        (b"dirB/subdira/subsubdirz/file7", 0o755, 7, 10007),
        (b"dirB/subdira/subsubdirz/file8", 0o755, 8, 10008),
        (b"dirB/subdirb/file10", 0o644, 10, 10010),
        (b"dirB/subdirb/file9", 0o644, 9, 10009),
        (b"dirC/file11", 0o644, 11, 10011),
        (b"dirC/file12", 0o644, 12, 10012),
        (b"dirC/file13", 0o644, 13, 10013),
        (b"dirC/file14", 0o644, 14, 10014),
        (b"dirC/file15", 0o644, 15, 10015),
        (b"file16", 0o644, 16, 10016),
    ];

    fn populate(t: &mut Tree<FileState>, s: &MapStore) {
        for &(name, mode, size, mtime) in TEST_FILES.iter() {
            t.add(s, name, &FileState::new(b'n', mode, size, mtime))
                .expect("can add file");
        }
    }

    #[test]
    fn count_get_and_remove() {
        let ms = MapStore::new();
        let mut t = Tree::new();
        assert_eq!(t.file_count(), 0);
        assert_eq!(
            t.get(&ms, b"dirB/subdira/subsubdirz/file7")
                .expect("can get"),
            None
        );
        populate(&mut t, &ms);
        assert_eq!(t.file_count(), 16);
        assert_eq!(
            t.get(&ms, b"dirB/subdira/subsubdirz/file7")
                .expect("can get"),
            Some(&FileState::new(b'n', 0o755, 7, 10007))
        );
        t.remove(&ms, b"dirB/subdirb/file9").expect("can remove");
        assert_eq!(t.file_count(), 15);
        t.remove(&ms, b"dirB/subdirb/file10").expect("can remove");
        assert_eq!(t.file_count(), 14);
        assert_eq!(
            t.get(&ms, b"dirB/subdira/subsubdirz/file7")
                .expect("can get"),
            Some(&FileState::new(b'n', 0o755, 7, 10007))
        );
        assert_eq!(t.get(&ms, b"dirB/subdirb/file9").expect("can get"), None);
    }

    #[test]
    fn iterate() {
        let ms = MapStore::new();
        let mut t = Tree::new();
        assert_eq!(t.get_first(&ms).expect("can get first"), None);
        populate(&mut t, &ms);
        let mut expect_iter = TEST_FILES.iter();
        let expected = expect_iter.next().unwrap();
        let mut filename = expected.0.to_vec();
        assert_eq!(
            t.get_first(&ms).expect("can get first"),
            Some((
                filename.clone().into_boxed_slice(),
                &FileState::new(b'n', expected.1, expected.2, expected.3)
            ))
        );
        while let Some(expected) = expect_iter.next() {
            let actual = t.get_next(&ms, &filename).expect("can get next");
            filename = expected.0.to_vec();
            assert_eq!(
                actual,
                Some((
                    filename.clone().into_boxed_slice(),
                    &FileState::new(b'n', expected.1, expected.2, expected.3)
                ))
            );
        }
        assert_eq!(t.get_next(&ms, &filename).expect("can get next"), None);
    }

    #[test]
    fn has_dir() {
        let ms = MapStore::new();
        let mut t = Tree::new();
        assert_eq!(
            t.has_dir(&ms, b"anything/").expect("can check has_dir"),
            false
        );
        populate(&mut t, &ms);
        assert_eq!(
            t.has_dir(&ms, b"something else/")
                .expect("can check has_dir"),
            false
        );
        assert_eq!(t.has_dir(&ms, b"dirB/").expect("can check has_dir"), true);
        assert_eq!(
            t.has_dir(&ms, b"dirB/subdira/").expect("can check has_dir"),
            true
        );
        assert_eq!(
            t.has_dir(&ms, b"dirB/subdira/subsubdirz/")
                .expect("can check has_dir"),
            true
        );
        assert_eq!(
            t.has_dir(&ms, b"dirB/subdira/subsubdirz/file7")
                .expect("can check has_dir"),
            false
        );
        assert_eq!(
            t.has_dir(&ms, b"dirB/subdira/subsubdirz/file7/")
                .expect("can check has_dir"),
            false
        );
    }

    #[test]
    fn write_empty() {
        let ns = NullStore::new();
        let mut ms = MapStore::new();
        let mut t = Tree::<FileState>::new();
        t.write_full(&mut ms, &ns).expect("can write full");
        t.write_delta(&mut ms).expect("can write delta");
        let mut ms2 = MapStore::new();
        t.write_full(&mut ms2, &ms).expect("can write full");
        let t_root = t.root_id().unwrap();
        let t_count = t.file_count();
        let mut t2 = Tree::<FileState>::open(t_root, t_count);
        assert_eq!(t2.get_first(&ms2).expect("can get first"), None);
    }

    #[test]
    fn write() {
        let ns = NullStore::new();
        let mut ms = MapStore::new();
        let mut t = Tree::new();
        populate(&mut t, &ms);
        t.write_full(&mut ms, &ns).expect("can write full");
        t.write_delta(&mut ms).expect("can write delta");
        let mut ms2 = MapStore::new();
        t.write_full(&mut ms2, &ms).expect("can write full");
        let t_root = t.root_id().unwrap();
        let t_count = t.file_count();
        let mut t2 = Tree::open(t_root, t_count);
        assert_eq!(
            t2.get(&ms2, b"dirB/subdira/subsubdirz/file7")
                .expect("can get"),
            Some(&FileState::new(b'n', 0o755, 7, 10007))
        );
    }

    #[test]
    fn visit() {
        let mut ms = MapStore::new();
        let mut t = Tree::new();
        populate(&mut t, &ms);
        let mut files = Vec::new();
        {
            let mut v = |path: &Vec<KeyRef>, _fs: &mut FileState| {
                files.push(path.concat());
                Ok(VisitorResult::NotChanged)
            };
            t.visit(&mut ms, &mut v).expect("can visit");
        }
        assert_eq!(
            files,
            TEST_FILES
                .iter()
                .map(|t| t.0.to_vec())
                .collect::<Vec<Vec<u8>>>()
        );
    }

    #[test]
    fn visit_changed() {
        let ns = NullStore::new();
        let mut ms = MapStore::new();
        let mut t = Tree::new();
        populate(&mut t, &ms);
        t.write_full(&mut ms, &ns).expect("can write full");

        // Touch file5.  This file, and any file in an ancestor directory (file4, file5 and file16)
        // will be in directories marked as changed.
        t.add(
            &ms,
            b"dirB/subdira/subsubdirx/file5",
            &FileState::new(b'm', 0o644, 200, 2000),
        )
        .expect("can add");

        let mut files = Vec::new();
        {
            let mut v = |path: &Vec<KeyRef>, _fs: &mut FileState| {
                files.push(path.concat());
                Ok(VisitorResult::NotChanged)
            };
            t.visit_changed(&mut ms, &mut v).expect("can visit_changed");
        }
        assert_eq!(
            files,
            vec![
                b"dirB/subdira/file4".to_vec(),
                b"dirB/subdira/subsubdirx/file5".to_vec(),
                b"file16".to_vec(),
            ]
        );
    }

    #[test]
    fn filtered_keys() {
        let ms = MapStore::new();
        let mut t = Tree::new();
        populate(&mut t, &ms);

        // Define a mapping function that upper-cases 'A' characters:
        fn map_upper_a(k: KeyRef) -> Result<Key> {
            Ok(k.iter()
                .map(|c| if *c == b'a' { b'A' } else { *c })
                .collect::<Vec<u8>>()
                .into_boxed_slice())
        }

        // Another map function that does nothing.
        fn map_noop(k: KeyRef) -> Result<Key> {
            Ok(Vec::from(k).into_boxed_slice())
        }

        // Look-up with normalized name should give non-normalized version.
        assert_eq!(
            t.get_filtered_key(&ms, b"dirA/subdirA/file1", &mut map_upper_a, 0)
                .expect("should succeed"),
            vec![b"dirA/subdira/file1".to_vec().into_boxed_slice()]
        );

        // Look-up with non-normalized name should match nothing.
        assert_eq!(
            t.get_filtered_key(&ms, b"dirA/subdira/file1", &mut map_upper_a, 0)
                .expect("should succeed"),
            vec![]
        );

        // Change filter function should invalid existing cache.
        assert_eq!(
            t.get_filtered_key(&ms, b"dirA/subdirA/file1", &mut map_noop, 1)
                .unwrap(),
            vec![]
        );
    }

    #[test]
    fn filtered_keys_surjective() {
        let mut t = Tree::new();
        let ms = MapStore::new();
        t.add(&ms, b"a/a/A", &FileState::new(b'a', 0, 0, 0))
            .unwrap();
        t.add(&ms, b"A/a/A", &FileState::new(b'a', 0, 0, 0))
            .unwrap();
        t.add(&ms, b"A/A/a", &FileState::new(b'a', 0, 0, 0))
            .unwrap();

        fn map_upper_a(key: KeyRef) -> Result<Key> {
            Ok(key
                .iter()
                .map(|c| if *c == b'a' { b'A' } else { *c })
                .collect::<Vec<u8>>()
                .into_boxed_slice())
        }

        assert_eq!(
            t.get_filtered_key(&ms, b"A/A", &mut map_upper_a, 0)
                .unwrap(),
            vec![
                b"A/A".to_vec().into_boxed_slice(),
                b"A/a".to_vec().into_boxed_slice(),
                b"a/a".to_vec().into_boxed_slice(),
            ]
        );

        assert_eq!(
            t.get_filtered_key(&ms, b"A/A/A", &mut map_upper_a, 0)
                .unwrap(),
            vec![
                b"A/A/a".to_vec().into_boxed_slice(),
                b"A/a/A".to_vec().into_boxed_slice(),
                b"a/a/A".to_vec().into_boxed_slice(),
            ]
        );
    }
}


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

use std::borrow::Cow;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::fmt;
use std::io::Cursor;
use std::iter::Iterator;
use std::ops::Deref;
use std::path::Path;
use std::path::PathBuf;
use std::str::FromStr;

use anyhow::anyhow;
use anyhow::Result;
use types::HgId;
use util::path::create_dir;

use crate::filestate::FileStateV2;
use crate::filestate::StateFlags;
use crate::filestore::FileStore;
use crate::legacy_eden_dirstate::read_eden_dirstate;
use crate::legacy_eden_dirstate::write_eden_dirstate;
use crate::metadata::Metadata;
use crate::root::TreeStateRoot;
use crate::serialization::Serializable;
use crate::store::BlockId;
use crate::store::Store;
use crate::store::StoreView;
use crate::tree::AggregatedState;
use crate::tree::Key;
use crate::tree::KeyRef;
use crate::tree::Node;
use crate::tree::Tree;
use crate::tree::VisitorResult;

const FILTER_LOWERCASE: u64 = 0x1;
/// `TreeState` uses a single tree to track an extended state of `TreeDirstate`.
/// See the comment about `FileStateV2` for the difference.
/// In short, `TreeState` combines dirstate and fsmonitor state.
pub struct TreeState {
    store: FileStore,
    tree: Tree<FileStateV2>,
    root: TreeStateRoot,
    original_root_id: BlockId,
    // eden_dirstate_path is only used in the case the case that the treestate is
    // wrapping a legacy eden dirstate which is necessary for EdenFS compatility.
    // TODO: Remove once EdenFS has migrated to treestate.
    eden_dirstate_path: Option<PathBuf>,
    case_sensitive: bool,
}

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

impl TreeState {
    /// Read `TreeState` from a file, or create an empty new `TreeState` if `root_id` is None.
    pub fn open<P: AsRef<Path>>(path: P, root_id: BlockId, case_sensitive: bool) -> Result<Self> {
        let path = path.as_ref();
        tracing::trace!(target: "treestate::open", "creating filestore at {path:?}");
        let store = FileStore::open(path)?;
        let root = {
            tracing::trace!(target: "treestate::open", "reading root data");
            let mut root_buf = Cursor::new(store.read(root_id)?);
            tracing::trace!(target: "treestate::open", "deserializing root data");
            TreeStateRoot::deserialize(&mut root_buf)?
        };
        tracing::trace!(target: "treestate::open", "constructing tree");
        let tree = Tree::open(root.tree_block_id(), root.file_count());
        Ok(TreeState {
            store,
            tree,
            root,
            original_root_id: root_id,
            eden_dirstate_path: None,
            case_sensitive,
        })
    }

    pub fn new(directory: &Path, case_sensitive: bool) -> Result<(Self, BlockId)> {
        tracing::trace!(target: "treestate::create", "creating directory {directory:?}");
        create_dir(directory)?;
        let name = format!("{:x}", uuid::Uuid::new_v4());
        let path = directory.join(&name);
        tracing::trace!(target: "treestate::create", "creating filestore {path:?}");
        let mut store = FileStore::create(&path)?;
        let _lock = store.lock()?;
        let root = TreeStateRoot::default();
        let tree = Tree::new();
        let mut treestate = TreeState {
            store,
            tree,
            root,
            original_root_id: BlockId(0),
            eden_dirstate_path: None,
            case_sensitive,
        };
        tracing::trace!(target: "treestate::create", "flushing treestate");
        let root_id = treestate.flush()?;

        tracing::trace!(target: "treestate::create", "treestate created");
        Ok((treestate, root_id))
    }

    /// Provides the ability to populate a treestate from a legacy eden dirstate.
    /// Clean up once EdenFS has been migrated from legacy dirstate to
    /// treestate.
    /// N.B: A legacy eden dirstate has a different binary format to a legacy
    /// dirstate.
    pub fn from_eden_dirstate<P: AsRef<Path>>(
        eden_dirstate_path: P,
        case_sensitive: bool,
    ) -> Result<Self> {
        let store = FileStore::in_memory()?;
        let root = TreeStateRoot::default();
        let tree = Tree::new();

        let (metadata, entries) = read_eden_dirstate(eden_dirstate_path.as_ref())?;

        let path = eden_dirstate_path.as_ref().to_path_buf();
        let mut treestate = TreeState {
            store,
            tree,
            root,
            original_root_id: BlockId(0),
            eden_dirstate_path: Some(path),
            case_sensitive,
        };

        treestate.set_metadata(metadata)?;

        for (key, state) in entries {
            treestate.insert(key, &state)?;
        }

        Ok(treestate)
    }

    pub fn path(&self) -> Option<&Path> {
        self.store.path()
    }

    pub fn file_name(&self) -> Result<String> {
        Ok(self
            .path()
            .ok_or_else(|| anyhow!("missing store path for TreeState"))?
            .file_name()
            .ok_or_else(|| anyhow!("missing file name for TreeState"))?
            .to_string_lossy()
            .to_string())
    }

    /// The root_id from when the treestate was loaded or last saved. Gets updated upon flush.
    pub fn original_root_id(&self) -> BlockId {
        self.original_root_id
    }

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

    /// Flush dirty entries. Return new `root_id` that can be passed to `open`.
    pub fn flush(&mut self) -> Result<BlockId> {
        let _lock = self.store.lock()?;
        let tree_block_id = { self.tree.write_delta(&mut self.store)? };
        self.write_root(tree_block_id)
    }

    /// Save as a new file.
    pub fn write_new<P: AsRef<Path>>(&mut self, directory: P) -> Result<BlockId> {
        let name = format!("{:x}", uuid::Uuid::new_v4());
        let path = directory.as_ref().join(&name);
        let mut new_store = FileStore::create(path)?;
        let _lock = new_store.lock()?;
        let tree_block_id = self.tree.write_full(&mut new_store, &self.store)?;
        self.store = new_store;
        let root_id = self.write_root(tree_block_id)?;
        Ok(root_id)
    }

    fn write_root(&mut self, tree_block_id: BlockId) -> Result<BlockId> {
        self.root.set_tree_block_id(tree_block_id);
        self.root.set_file_count(self.len() as u32);

        let mut root_buf = Vec::new();
        self.root.serialize(&mut root_buf)?;
        let result = self.store.append(&root_buf)?;
        self.store.flush()?;

        // TODO: Clean up once we migrate EdenFS to TreeState and no longer
        // need to write to legacy eden dirstate format.
        if let Some(eden_dirstate_path) = self.eden_dirstate_path.clone() {
            let metadata = self.metadata()?;
            let entries = self.flatten_tree()?;
            write_eden_dirstate(&eden_dirstate_path, metadata, entries)?;
        }

        self.original_root_id = result;
        Ok(result)
    }

    fn flatten_tree(&mut self) -> Result<HashMap<Box<[u8]>, FileStateV2>> {
        let mut results = HashMap::with_capacity(self.len());
        self.visit(
            &mut |path_components, state| {
                results.insert(path_components.concat().into_boxed_slice(), state.clone());
                Ok(VisitorResult::NotChanged)
            },
            &|_, _| true, // Visit all directories
            &|_, _| true, // Visit all files
        )?;

        Ok(results)
    }

    /// Create or replace the existing entry.
    pub fn insert<K: AsRef<[u8]>>(&mut self, path: K, state: &FileStateV2) -> Result<()> {
        self.tree.add(&self.store, path.as_ref(), state)
    }

    pub fn remove<K: AsRef<[u8]>>(&mut self, path: K) -> Result<bool> {
        self.tree.remove(&self.store, path.as_ref())
    }

    pub fn get<K: AsRef<[u8]>>(&mut self, path: K) -> Result<Option<&FileStateV2>> {
        self.tree.get(&self.store, path.as_ref())
    }

    pub fn get_keys_ignorecase<K: AsRef<[u8]>>(&mut self, path: K) -> Result<Vec<Key>> {
        fn map_lowercase(k: KeyRef) -> Result<Key> {
            let s = std::str::from_utf8(k);
            Ok(if let Ok(s) = s {
                s.to_lowercase().into_bytes().into_boxed_slice()
            } else {
                k.to_vec().into_boxed_slice()
            })
        }
        self.get_filtered_key(
            &map_lowercase(path.as_ref())?,
            &mut map_lowercase,
            FILTER_LOWERCASE,
        )
    }

    pub fn normalize_path_and_get<'a>(
        &mut self,
        path: &'a [u8],
    ) -> Result<(Cow<'a, [u8]>, Option<FileStateV2>)> {
        if self.case_sensitive {
            return Ok((Cow::Borrowed(path), self.get(path)?.cloned()));
        }

        // Optimistic check for normal case where path matches exactly.
        // Case insensitive filesystems are normally still case preserving.
        if let Some(state) = self.get(path)? {
            if state.state.intersects(StateFlags::EXIST_NEXT) {
                return Ok((Cow::Borrowed(path), Some(state.clone())));
            }
        }

        let mut best = None;
        for key in self.get_keys_ignorecase(path)? {
            let state = match self.get(&key)? {
                None => continue,
                Some(state) => state.clone(),
            };

            // If there are multiple matches, prefer the format for the version that still exists.
            if best.is_none() || state.state.intersects(StateFlags::EXIST_NEXT) {
                best = Some((Cow::Owned(key.into()), state));
            }
        }

        match best {
            Some((path, state)) => Ok((path, Some(state))),
            None => Ok((Cow::Borrowed(path), None)),
        }
    }

    pub fn normalize_path<'a>(&mut self, path: &'a [u8]) -> Result<Cow<'a, [u8]>> {
        Ok(self.normalize_path_and_get(path)?.0)
    }

    pub fn normalized_get<K: AsRef<[u8]>>(&mut self, path: K) -> Result<Option<FileStateV2>> {
        Ok(self.normalize_path_and_get(path.as_ref())?.1)
    }

    /// Get the aggregated state of a directory. This is useful, for example, to tell if a
    /// directory only contains removed files.
    pub fn get_dir<K: AsRef<[u8]>>(&mut self, path: K) -> Result<Option<AggregatedState>> {
        self.tree.get_dir(&self.store, path.as_ref())
    }

    pub fn len(&self) -> usize {
        self.tree.file_count() as usize
    }

    pub fn set_metadata_bytes<T: AsRef<[u8]>>(&mut self, metadata: T) {
        self.root
            .set_metadata(Vec::from(metadata.as_ref()).into_boxed_slice());
    }

    pub fn set_metadata(&mut self, metadata: BTreeMap<String, String>) -> Result<()> {
        let mut buf = Vec::new();
        Metadata(metadata).serialize(&mut buf)?;
        self.set_metadata_bytes(buf);
        Ok(())
    }

    pub fn metadata_bytes(&self) -> &[u8] {
        self.root.metadata().deref()
    }

    pub fn metadata(&self) -> Result<BTreeMap<String, String>> {
        let metadata = Metadata::deserialize(&mut self.metadata_bytes())?;
        Ok(metadata.0)
    }

    pub fn update_metadata(&mut self, new: &[(String, Option<String>)]) -> Result<()> {
        let mut metadata = self.metadata()?;

        for (key, value) in new.iter() {
            match value {
                Some(value) => metadata.insert(key.to_string(), value.to_string()),
                None => metadata.remove(key),
            };
        }

        let mut buf = Vec::new();
        Metadata(metadata).serialize(&mut buf)?;
        self.root.set_metadata(buf.into_boxed_slice());
        Ok(())
    }

    pub fn parents<'a>(&'a self) -> impl Iterator<Item = Result<HgId>> + 'a {
        (1..).map_while(|i| match self.metadata() {
            Err(err) => Some(Err(err)),
            Ok(metadata) => metadata
                .get(&format!("p{}", i))
                .map(|parent_hash| HgId::from_str(parent_hash).map_err(|e| e.into())),
        })
    }

    pub fn set_parents(&mut self, parents: &mut dyn Iterator<Item = &HgId>) -> Result<()> {
        let mut values: Vec<(String, Option<String>)> = Vec::with_capacity(2);
        for (i, p) in parents.enumerate() {
            // i+1 because parents are 1-indexed
            values.push((format!("p{}", i + 1), Some(p.to_string())));
        }
        // Set p1 or p2 to None to remove it from the metadata if necessary.
        if values.len() == 0 {
            values.push(("p1".to_string(), None));
        }
        if values.len() == 1 {
            values.push(("p2".to_string(), None));
        }
        self.update_metadata(&values)
    }

    pub fn has_dir<P: AsRef<[u8]>>(&mut self, path: P) -> Result<bool> {
        self.tree.has_dir(&self.store, path.as_ref())
    }

    pub fn visit<F, VD, VF>(
        &mut self,
        visitor: &mut F,
        visit_dir: &VD,
        visit_file: &VF,
    ) -> Result<()>
    where
        F: FnMut(&Vec<&[u8]>, &mut FileStateV2) -> Result<VisitorResult>,
        VD: Fn(&Vec<KeyRef>, &Node<FileStateV2>) -> bool,
        VF: Fn(&Vec<KeyRef>, &FileStateV2) -> bool,
    {
        self.tree
            .visit_advanced(&self.store, visitor, visit_dir, visit_file)
    }

    pub fn get_filtered_key<F>(
        &mut self,
        name: KeyRef,
        filter: &mut F,
        filter_id: u64,
    ) -> Result<Vec<Key>>
    where
        F: FnMut(KeyRef) -> Result<Key>,
    {
        self.tree
            .get_filtered_key(&self.store, name, filter, filter_id)
    }

    pub fn path_complete<FA, FV>(
        &mut self,
        prefix: KeyRef,
        full_paths: bool,
        acceptable: &FA,
        visitor: &mut FV,
    ) -> Result<()>
    where
        FA: Fn(&FileStateV2) -> bool,
        FV: FnMut(&Vec<KeyRef>) -> Result<()>,
    {
        self.tree
            .path_complete(&self.store, prefix, full_paths, acceptable, visitor)
    }

    // Distrust changed files with a mtime of `fsnow`. Rewrite their mtime to -1.
    // See mercurial/pure/parsers.py:pack_dirstate in core Mercurial for motivation.
    // Basically, this is required for the following case:
    //
    // $ hg update rev; write foo; hg commit -m update-foo
    //
    //   Time (second) | 0   | 1           |
    //   hg update       ...----|
    //   write foo               |--|
    //   hg commit                   |---...
    //
    // If "write foo" changes a file without changing its mtime and size, the file
    // change won't be detected. Therefore if mtime is `fsnow`, reset it to a different
    // value and mark it as NEED_CHECK, at the end of update to workaround the issue.
    // Here, hg assumes nobody else is touching the working copy when it holds wlock
    // (ex. during second 0).
    //
    // This is used before "flush" or "saveas".
    //
    // Note: In TreeState's case, NEED_CHECK might mean "perform a quick mtime check",
    // or "perform a content check" depending on the caller. Be careful when removing
    // "mtime = -1" statement.
    pub fn invalidate_mtime(&mut self, now: i32) -> Result<()> {
        self.visit(
            &mut |_, state| {
                if state.mtime >= now {
                    state.mtime = -1;
                    state.state |= StateFlags::NEED_CHECK;
                    Ok(VisitorResult::Changed)
                } else {
                    Ok(VisitorResult::NotChanged)
                }
            },
            &|_, dir| {
                if !dir.is_changed() {
                    false
                } else {
                    match dir.get_aggregated_state() {
                        Some(x) => x
                            .union
                            .intersects(StateFlags::EXIST_P1 | StateFlags::EXIST_P2),
                        None => true,
                    }
                }
            },
            &|_, file| {
                file.state
                    .intersects(StateFlags::EXIST_P1 | StateFlags::EXIST_P2)
            },
        )
    }
}

#[cfg(test)]
mod tests {
    use rand::Rng;
    use rand::SeedableRng;
    use rand_chacha::ChaChaRng;
    use tempfile::tempdir;

    use super::*;
    use crate::filestate::StateFlags;

    #[test]
    fn test_new() {
        let dir = tempdir().expect("tempdir");
        let state = TreeState::new(dir.path(), true).expect("open").0;
        assert!(state.metadata_bytes().is_empty());
        assert_eq!(state.len(), 0);
    }

    #[test]
    fn test_empty_flush() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.path(), true).expect("open").0;
        let block_id = state.flush().expect("flush");
        let state = TreeState::open(
            dir.path().join(state.file_name().unwrap()),
            block_id.into(),
            true,
        )
        .expect("open");
        assert!(state.metadata_bytes().is_empty());
        assert_eq!(state.len(), 0);
    }

    #[test]
    fn test_empty_write_as() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.path(), true).expect("open").0;
        let block_id = state.write_new(dir.path()).expect("write_as");
        let state = TreeState::open(
            dir.path().join(state.file_name().unwrap()),
            block_id.into(),
            true,
        )
        .expect("open");
        assert!(state.metadata_bytes().is_empty());
        assert_eq!(state.len(), 0);
    }

    #[test]
    fn test_set_metadata() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.path(), true).expect("open").0;
        state.set_metadata_bytes(b"foobar");
        let orig_name = state.file_name().unwrap();
        let block_id1 = state.flush().expect("flush");
        let block_id2 = state.write_new(dir.path()).expect("write_as");
        let new_name = state.file_name().unwrap();
        let state =
            TreeState::open(dir.path().join(orig_name), block_id1.into(), true).expect("open");
        assert_eq!(state.metadata_bytes()[..], b"foobar"[..]);
        let state =
            TreeState::open(dir.path().join(new_name), block_id2.into(), true).expect("open");
        assert_eq!(state.metadata_bytes()[..], b"foobar"[..]);
    }

    #[test]
    fn test_update_metadata() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.path(), true).expect("open").0;
        state
            .update_metadata(&[
                ("key1".to_string(), Some("value1".to_string())),
                ("key2".to_string(), Some("value2".to_string())),
            ])
            .unwrap();
        assert_eq!(
            state.metadata().unwrap().get("key1").cloned(),
            Some("value1".to_string())
        );
        assert_eq!(
            state.metadata().unwrap().get("key2").cloned(),
            Some("value2".to_string())
        );

        state
            .update_metadata(&[("key1".to_string(), Some("value1.b".to_string()))])
            .unwrap();
        assert_eq!(
            state.metadata().unwrap().get("key1").cloned(),
            Some("value1.b".to_string())
        );
        assert_eq!(
            state.metadata().unwrap().get("key2").cloned(),
            Some("value2".to_string())
        );

        state
            .update_metadata(&[
                ("key1".to_string(), Some("value1.c".to_string())),
                ("key2".to_string(), None),
            ])
            .unwrap();
        assert_eq!(
            state.metadata().unwrap().get("key1").cloned(),
            Some("value1.c".to_string())
        );
        assert_eq!(state.metadata().unwrap().get("key2"), None);
    }

    // Some random paths extracted from fb-ext, plus some manually added entries, shuffled.
    const SAMPLE_PATHS: [&[u8]; 21] = [
        b".fbarcanist",
        b"phabricator/phabricator_graphql_client_urllib.pyc",
        b"ext3rd/__init__.py",
        b"ext3rd/.git/objects/14/8f179e7e702ddedb54c53f2726e7f81b14a33f",
        b"rust/radixbuf/.git/objects/pack/pack-c0bc37a255e59f5563de9a76013303d8df46a659.idx",
        b".hg/shelved/default-106.patch",
        b"rust/radixbuf/.git/objects/20/94e0274ba1ef2ec30de884e3ca4d7093838064",
        b"rust/radixbuf/.git/hooks/prepare-commit-msg.sample",
        b"rust/radixbuf/.git/objects/b3/9acb828290b77704cc44e748d6e7d4a528d6ae",
        b"scripts/lint.py",
        b".fbarcanist/unit/MercurialTestEngine.php",
        b".hg/shelved/default-37.patch",
        b"rust/radixbuf/.git/objects/01/d8e75b3bae0819c4095ae96ebdc889e9e5d806",
        b"ext3rd/fastannotate/error.py",
        b"rust/radixbuf/.git/objects/pack/pack-c0bc37a255e59f5563de9a76013303d8df46a659.pack",
        b"distutils_rust/__init__.py",
        b".editorconfig",
        b"rust/radixbuf/.git/objects/01/89a583d7e9aff802cdfed3ff3cc3a473253281",
        b"ext3rd/fastannotate/commands.py",
        b"distutils_rust/__init__.pyc",
        b"rust/radixbuf/.git/objects/b3/9b2824f47b66462e92ffa4f978bc95f5fdad2e",
    ];

    fn new_treestate<P: AsRef<Path>>(directory: P) -> TreeState {
        let mut state = TreeState::new(directory.as_ref(), true).expect("open").0;
        let mut rng = ChaChaRng::from_seed([0; 32]);
        for path in &SAMPLE_PATHS {
            let file = rng.gen();
            state.insert(path, &file).expect("insert");
        }
        state
    }

    #[test]
    fn test_insert() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        let mut rng = ChaChaRng::from_seed([0; 32]);
        for path in &SAMPLE_PATHS {
            let file: FileStateV2 = rng.gen();
            assert_eq!(state.get(path).unwrap().unwrap(), &file);
        }
        assert_eq!(state.len(), SAMPLE_PATHS.len());
    }

    #[test]
    fn test_remove() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        for path in &SAMPLE_PATHS {
            assert!(state.remove(path).unwrap())
        }
        for path in &SAMPLE_PATHS {
            assert!(!state.remove(path).unwrap())
        }
        assert_eq!(state.len(), 0);
    }

    #[test]
    fn test_non_empty_flush() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        let block_id = state.flush().expect("flush");
        let mut state = TreeState::open(
            dir.path().join(state.file_name().unwrap()),
            block_id.into(),
            true,
        )
        .expect("open");
        let mut rng = ChaChaRng::from_seed([0; 32]);
        for path in &SAMPLE_PATHS {
            let file: FileStateV2 = rng.gen();
            assert_eq!(state.get(path).unwrap().unwrap(), &file);
        }
        assert_eq!(state.len(), SAMPLE_PATHS.len());
    }

    #[test]
    fn test_non_empty_write_as() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        let block_id = state.write_new(dir.path()).expect("write_as");
        let mut state = TreeState::open(
            dir.path().join(state.file_name().unwrap()),
            block_id.into(),
            true,
        )
        .expect("open");
        let mut rng = ChaChaRng::from_seed([0; 32]);
        for path in &SAMPLE_PATHS {
            let file: FileStateV2 = rng.gen();
            assert_eq!(state.get(path).unwrap().unwrap(), &file);
        }
        assert_eq!(state.len(), SAMPLE_PATHS.len());
    }

    #[test]
    fn test_has_dir() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        assert!(state.has_dir(b"/").unwrap());
        assert!(state.has_dir(b"ext3rd/").unwrap());
        assert!(!state.has_dir(b"ext4th/").unwrap());
        assert!(state.has_dir(b"rust/radixbuf/.git/objects/").unwrap());
        assert!(!state.has_dir(b"rust/radixbuf/.git2/objects/").unwrap());
    }

    #[test]
    fn test_get_keys_ignorecase() {
        let dir = tempdir().expect("tempdir");
        let mut state = new_treestate(dir.path());
        let expected = vec![b"ext3rd/__init__.py".to_vec().into_boxed_slice()];
        assert_eq!(
            state.get_keys_ignorecase(b"ext3rd/__init__.py").unwrap(),
            expected
        );
        assert_eq!(
            state.get_keys_ignorecase(b"exT3rd/__init__.py").unwrap(),
            expected
        );
        assert_eq!(
            state.get_keys_ignorecase(b"ext3rd/__Init__.py").unwrap(),
            expected
        );
    }

    #[test]
    fn test_normalize_casesensitive() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.as_ref(), true).expect("open").0;

        let mut rng = ChaChaRng::from_seed([0; 32]);
        let mut file = rng.gen();
        state.insert(b"dir/file", &file).unwrap();
        assert_eq!(
            state.normalize_path(b"dir/file").unwrap().as_ref(),
            b"dir/file"
        );
        assert_eq!(
            state.normalize_path(b"dir/FILE").unwrap().as_ref(),
            b"dir/FILE"
        );
        assert_eq!(
            state.normalize_path(b"DIR/file").unwrap().as_ref(),
            b"DIR/file"
        );

        file.state = StateFlags::EXIST_NEXT;
        state.insert(b"dir/RENAME", &file).unwrap();
        file.state = StateFlags::EXIST_P1;
        state.insert(b"dir/rename", &file).unwrap();
        assert_eq!(
            state.normalize_path(b"dir/rename").unwrap().as_ref(),
            b"dir/rename"
        );
        assert_eq!(
            state.normalize_path(b"dir/RENAME").unwrap().as_ref(),
            b"dir/RENAME"
        );
    }

    #[test]
    fn test_normalize_incasesensitive() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.as_ref(), false).expect("open").0;

        let mut rng = ChaChaRng::from_seed([0; 32]);
        let mut file = rng.gen();
        state.insert(b"dir/file", &file).unwrap();
        assert_eq!(
            std::str::from_utf8(state.normalize_path(b"dir/file").unwrap().as_ref()).unwrap(),
            "dir/file"
        );
        assert_eq!(
            std::str::from_utf8(state.normalize_path(b"dir/FILE").unwrap().as_ref()).unwrap(),
            "dir/file"
        );
        assert_eq!(
            std::str::from_utf8(state.normalize_path(b"DIR/file").unwrap().as_ref()).unwrap(),
            "dir/file"
        );

        // Choose paths that will exist over paths that used to exist.
        file.state = StateFlags::EXIST_NEXT;
        state.insert(b"dir/RENAME", &file).unwrap();
        file.state = StateFlags::EXIST_P1;
        state.insert(b"dir/rename", &file).unwrap();
        assert_eq!(
            std::str::from_utf8(state.normalize_path(b"dir/rename").unwrap().as_ref()).unwrap(),
            "dir/RENAME"
        );
        assert_eq!(
            std::str::from_utf8(state.normalize_path(b"dir/RENAME").unwrap().as_ref()).unwrap(),
            "dir/RENAME"
        );
    }

    #[test]
    fn test_parents() {
        let dir = tempdir().expect("tempdir");
        let mut state = TreeState::new(dir.path(), true).expect("open").0;
        let orig_name = state.file_name().unwrap();
        let mut rng = ChaChaRng::from_seed([0; 32]);

        let p1 = HgId::random(&mut rng);
        let p2 = HgId::random(&mut rng);
        let p3 = HgId::random(&mut rng);

        state.set_parents(&mut [p1].iter()).unwrap();
        assert_eq!(
            state.parents().collect::<Result<Vec<_>>>().unwrap(),
            [p1].to_vec()
        );

        state.set_parents(&mut [p1, p2].iter()).unwrap();
        assert_eq!(
            state.parents().collect::<Result<Vec<_>>>().unwrap(),
            [p1, p2].to_vec()
        );

        state.set_parents(&mut [p1, p3].iter()).unwrap();
        assert_eq!(
            state.parents().collect::<Result<Vec<_>>>().unwrap(),
            [p1, p3].to_vec()
        );

        let block_id = state.flush().expect("flush");

        let state =
            TreeState::open(dir.path().join(orig_name), block_id.into(), true).expect("open");
        assert_eq!(
            state.parents().collect::<Result<Vec<_>>>().unwrap(),
            [p1, p3].to_vec()
        );
    }

    #[test]
    fn test_concurrent_writes() {
        check_concurrent_writes(&[b"a"], &[b"b"]);
        check_concurrent_writes(&[b"a/1"], &[b"a/2"]);

        let paths1 = SAMPLE_PATHS.into_iter().take(10).collect::<Vec<_>>();
        let paths2 = SAMPLE_PATHS.into_iter().rev().take(1).collect::<Vec<_>>();
        check_concurrent_writes(&paths1, &paths2);

        let paths2 = SAMPLE_PATHS.into_iter().rev().take(10).collect::<Vec<_>>();
        check_concurrent_writes(&paths1, &paths2);
    }

    // Test appending paths1, and paths2 "concurrently".
    // Paths should not overlap.
    fn check_concurrent_writes(paths1: &[&'static [u8]], paths2: &[&'static [u8]]) {
        let dir = tempdir().unwrap();
        let dir_path = dir.path();

        // Prepare initial state.
        let (path, root_id) = {
            let mut state = new_treestate(dir_path);
            let root_id = state.flush().unwrap();
            (dir_path.join(state.file_name().unwrap()), root_id)
        };

        // Start M threads for concurrent writes.
        // Each thread writes N updates for paths, and returns the root_id and state to check for
        // the paths.
        const M: usize = 10;
        const N: usize = 20;
        let threads = (0..M)
            .map(|i| {
                let path = path.to_path_buf();
                let paths = if i & 1 == 0 { paths1 } else { paths2 };
                let paths = paths.to_vec();
                std::thread::spawn(move || {
                    let mut state = TreeState::open(&path, root_id, true).unwrap();
                    let mut file_state = FileStateV2 {
                        size: (i * N) as i32,
                        ..Default::default()
                    };
                    let mut result = Vec::new();
                    for _ in 0..N {
                        // Change `file_state` so `result` contains different `file_state`s to
                        // check. This makes the test more interesting.
                        file_state.mtime += 1;
                        for p in &paths {
                            state.insert(p, &file_state).unwrap();
                        }
                        let root_id = state.flush().unwrap();
                        result.push((root_id, file_state.clone()));
                    }
                    (paths, result)
                })
            })
            .collect::<Vec<_>>();

        let to_check: Vec<(Vec<&[u8]>, Vec<(BlockId, FileStateV2)>)> =
            threads.into_iter().map(|t| t.join().unwrap()).collect();

        // Check that things can be read properly (aka. no "invalid store id: ..." errors),
        // and are written properly (file_state1 and file_state2).
        for (paths, root_id_file_states) in to_check {
            for (root_id, file_state) in root_id_file_states {
                let mut state = TreeState::open(&path, root_id, true).unwrap();
                for p in &paths {
                    let got = state.get(p).unwrap().unwrap();
                    assert_eq!(got, &file_state);
                }
            }
        }
    }
}


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

//! Ordered map implementation using a sorted vector

use std::borrow::Borrow;
use std::collections::Bound;
use std::collections::Bound::*;
use std::mem;
use std::slice::Iter as VecIter;
use std::slice::IterMut as VecIterMut;

#[derive(Debug)]
pub struct VecMap<K, V> {
    vec: Vec<(K, V)>,
}

pub struct Iter<'a, K: 'a, V: 'a>(VecIter<'a, (K, V)>);

pub struct IterMut<'a, K: 'a, V: 'a>(VecIterMut<'a, (K, V)>);

impl<K, V> VecMap<K, V>
where
    K: Ord,
{
    /// Creates a new, empty VecMap.
    pub fn new() -> VecMap<K, V> {
        VecMap { vec: Vec::new() }
    }

    /// Creates a new, empty VecMap, with capacity for `capacity` entries.
    pub fn with_capacity(capacity: usize) -> VecMap<K, V> {
        VecMap {
            vec: Vec::with_capacity(capacity),
        }
    }

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

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

    /// Utility function to binary search for an index using the key.
    fn find_index<Q: ?Sized>(&self, q: &Q) -> Result<usize, usize>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        self.vec.binary_search_by(|e| e.0.borrow().cmp(q))
    }

    /// Inserts a key-value pair into the map.  If the key already had a value present in the
    /// map, that value is replaced and returned.  Otherwise, `None` is returned.
    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        let mut v = v;
        match self.find_index(&k) {
            Ok(index) => {
                mem::swap(&mut self.vec[index].1, &mut v);
                Some(v)
            }
            Err(index) => {
                self.vec.insert(index, (k, v));
                None
            }
        }
    }

    /// Inserts a key-value pair into the map.  Fast-path for when the key is not already
    /// present and is at the end of the map.
    pub fn insert_hint_end(&mut self, k: K, v: V) -> Option<V> {
        let len = self.vec.len();
        if len == 0 || self.vec[len - 1].0 < k {
            self.vec.push((k, v));
            None
        } else {
            self.insert(k, v)
        }
    }

    /// Removes a key-value pair from the map, returning the value if the key was previously in
    /// the map.
    pub fn remove<Q: ?Sized>(&mut self, q: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        match self.find_index(q) {
            Ok(index) => {
                let (_k, v) = self.vec.remove(index);
                Some(v)
            }
            Err(_index) => None,
        }
    }

    /// Returns a reference to the value corresponding to the key.
    pub fn get<'a, Q: ?Sized>(&'a self, q: &Q) -> Option<&'a V>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        match self.find_index(q) {
            Ok(index) => Some(&self.vec[index].1),
            Err(_index) => None,
        }
    }

    /// Returns a mutable reference to the value corresponding to the key.
    pub fn get_mut<'a, Q: ?Sized>(&'a mut self, q: &Q) -> Option<&'a mut V>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        match self.find_index(q) {
            Ok(index) => Some(&mut self.vec[index].1),
            Err(_index) => None,
        }
    }

    // Returns an iterator over the pairs of entries in the map.
    pub fn iter(&self) -> Iter<K, V> {
        Iter(self.vec.iter())
    }

    /// Returns a mutable iterator over the pairs of entries in the map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        IterMut(self.vec.iter_mut())
    }

    /// Utility function for implementing `range` and `range_mut`.  Convert a range boundary for
    /// the start of a range into a slice index suitable for use in a range expression.
    fn range_index_start<Q: ?Sized>(&self, b: Bound<&Q>) -> usize
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        match b {
            Unbounded => 0,
            Included(q) => match self.find_index(q) {
                Ok(index) => index,
                Err(index) => index,
            },
            Excluded(q) => match self.find_index(q) {
                Ok(index) => index + 1,
                Err(index) => index,
            },
        }
    }

    /// Utility function for implementing `range` and `range_mut`.  Convert a range boundary for
    /// the end of a range into a slice index suitable for use in a range expression.
    fn range_index_end<Q: ?Sized>(&self, b: Bound<&Q>) -> usize
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        match b {
            Unbounded => self.vec.len(),
            Included(q) => match self.find_index(q) {
                Ok(index) => index + 1,
                Err(index) => index,
            },
            Excluded(q) => match self.find_index(q) {
                Ok(index) => index,
                Err(index) => index,
            },
        }
    }

    /// Returns an iterator over the given range of keys.
    pub fn range<Q>(&self, range: (Bound<&Q>, Bound<&Q>)) -> Iter<K, V>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        let start = self.range_index_start(range.0);
        let end = self.range_index_end(range.1);
        Iter(self.vec[start..end].iter())
    }

    /// Returns a mutuable iterator over the given range of keys.
    pub fn range_mut<Q>(&mut self, range: (Bound<&Q>, Bound<&Q>)) -> IterMut<K, V>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        let start = self.range_index_start(range.0);
        let end = self.range_index_end(range.1);
        IterMut(self.vec[start..end].iter_mut())
    }
}

// Wrap `Iter` and `IterMut` for `VecMap` types.  These implementations adapt the `next` methods,
// converting their yielded types from `Option<&(K, V)>` to `Option<(&K, &V)>`, and from
// `Option<&mut (K, V)>` to `Option<(&K, &mut V)>`.  This allows `VecMap` iterators to be used
// in the same way as other map iterators to iterate over key-value pairs, and prevents callers
// from using the mutable iterator to mutate keys.

impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|&(ref k, ref v)| (k, v))
    }
}

impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
    type Item = (&'a K, &'a mut V);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|&mut (ref k, ref mut v)| (k, v))
    }
}

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

    use quickcheck::quickcheck;

    use crate::vecmap::Iter;
    use crate::vecmap::VecMap;

    #[test]
    fn insert_get_remove() {
        let mut vm = VecMap::new();
        assert_eq!(vm.insert("test1", "value1".to_string()), None);
        assert_eq!(vm.insert("test2", "value2".to_string()), None);
        assert_eq!(vm.insert_hint_end("test4", "value4".to_string()), None);
        assert_eq!(vm.insert_hint_end("test3", "value3".to_string()), None);
        assert_eq!(
            vm.insert("test1", "value1b".to_string()),
            Some("value1".to_string())
        );
        assert_eq!(vm.get(&"test1"), Some(&"value1b".to_string()));
        if let Some(v) = vm.get_mut(&"test1") {
            *v = "value1c".to_string();
        }
        assert_eq!(vm.get(&"test1"), Some(&"value1c".to_string()));
        assert_eq!(vm.remove("test2"), Some("value2".to_string()));
        assert_eq!(vm.remove("test2"), None);
        assert_eq!(vm.get(&"test2"), None);
        assert_eq!(vm.get_mut(&"never"), None);
    }

    #[test]
    fn iter() {
        let mut vm = VecMap::with_capacity(4);
        assert!(vm.is_empty());
        vm.insert(2, "value2");
        vm.insert(1, "value1");
        vm.insert(4, "value4");
        vm.insert(3, "value3");
        assert!(!vm.is_empty());
        assert_eq!(vm.len(), 4);
        {
            let mut im = vm.iter_mut();
            im.next();
            let e2 = im.next().unwrap();
            *e2.1 = "value2 - modified";
        }
        let mut i = vm.iter();
        assert_eq!(i.next(), Some((&1, &"value1")));
        assert_eq!(i.next(), Some((&2, &"value2 - modified")));
        assert_eq!(i.next(), Some((&3, &"value3")));
        assert_eq!(i.next(), Some((&4, &"value4")));
        assert_eq!(i.next(), None);
    }

    #[test]
    fn range() {
        let mut vm: VecMap<i32, i32> = VecMap::new();
        for n in 0..20 {
            vm.insert(n * 2, n * 4);
        }

        fn check_iter(mut x: Iter<i32, i32>, start: i32, end: i32) {
            let mut i = start;
            while i < end {
                assert_eq!(x.next(), Some((&i, &(i * 2))));
                i += 2;
            }
            assert_eq!(x.next(), None);
        }

        check_iter(vm.range((Unbounded, Unbounded)), 0, 39);
        check_iter(vm.range((Unbounded, Included(&2))), 0, 3);
        check_iter(vm.range((Unbounded, Excluded(&2))), 0, 1);
        check_iter(vm.range((Unbounded, Excluded(&7))), 0, 7);
        check_iter(vm.range((Unbounded, Included(&13))), 0, 13);
        check_iter(vm.range((Included(&4), Included(&13))), 4, 13);
        check_iter(vm.range((Included(&5), Included(&14))), 6, 15);
        check_iter(vm.range((Excluded(&5), Included(&20))), 6, 21);
        check_iter(vm.range((Excluded(&6), Included(&60))), 8, 39);
        check_iter(vm.range((Excluded(&-30), Unbounded)), 0, 39);
        check_iter(vm.range((Included(&-1), Unbounded)), 0, 39);

        assert_eq!(vm.get(&16), Some(&32));
        {
            let mut im = vm.range_mut((Included(&16), Excluded(&18)));
            *im.next().unwrap().1 *= 2;
            assert_eq!(im.next(), None);
        }
        assert_eq!(vm.get(&16), Some(&64));
    }

    fn vecmap_from_btreemap<K: Ord + Clone, V: Clone>(b: &BTreeMap<K, V>) -> VecMap<K, V> {
        let mut vm = VecMap::new();
        for (k, v) in b.iter() {
            vm.insert(k.clone(), v.clone());
        }
        vm
    }

    quickcheck! {
        fn like_btreemap_is_empty (b: BTreeMap<u32, u32>) -> bool {
            let vm = vecmap_from_btreemap(&b);
            vm.is_empty() == b.is_empty()
        }

        fn like_btreemap_len (b: BTreeMap<u32, u32>) -> bool {
            let vm = vecmap_from_btreemap(&b);
            vm.len() == b.len()
        }

        fn like_btreemap_iter (b: BTreeMap<u32, u32>) -> bool {
            let vm = vecmap_from_btreemap(&b);
            itertools::equal(vm.iter(), b.iter())
        }

        fn like_btreemap_range (b: BTreeMap<u32, u32>, key1: u32, key2: u32) -> bool {
            // range requires start key is not after end key.
            let (start, end) = if key1 <= key2 { (key1, key2) } else { (key2, key1) };
            let vm = vecmap_from_btreemap(&b);
            let range = (Included(&start), Excluded(&end));
            itertools::equal(vm.range(range), b.range(range))
        }
    }
}


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

/// Manage a `Vec` of references. so elements with different lifetimes can be
/// pushed manually, and poped automatically. Elements pushed later must have
/// a more narrow lifetime than elements pushed earlier.
///
/// Practically, instead of `vec![&'a T, &'a T, &'a, T, &'a T, ...]`,
/// have something like `vec![&'a T, &'b T, &'c, T, &'d T, ...]`.
/// where `'a` covers `'b`, `'b` covers `'c`....
/// The first function call gets `vec![&'a T]`, the second (nested) one
/// gets `vec![&'b T, &'b T], and so on.
///
/// An example is:
///
/// ```
/// # use treestate::vecstack::VecStack;
/// let mut vec = vec![];
/// let mut stack = VecStack::new(&mut vec);
/// {
///     let str1 = String::from("1");
///     let mut stack = stack.push(&str1);
///     assert_eq!(stack.as_ref()[..], vec!["1"][..]);
///     {
///         // This could also be a recursive function call
///         let str2 = String::from("2");
///         let mut stack = stack.push(&str2);
///         assert_eq!(stack.as_ref()[..], vec!["1", "2"][..]);
///         {
///             let str3 = String::from("3");
///             let stack = stack.push(&str3);
///             assert_eq!(stack.as_ref()[..], vec!["1", "2", "3"][..]);
///         }
///         assert_eq!(stack.as_ref()[..], vec!["1", "2"][..]);
///     }
///     assert_eq!(stack.as_ref()[..], vec!["1"][..]);
/// }
/// assert!(stack.as_ref().is_empty());
/// ```
pub struct VecStack<'a, T: 'a>
where
    T: ?Sized,
{
    // Multiple &mut is undefined behavior.
    // Use a pointer to disable optimization.
    inner: *mut Vec<&'a T>,

    // `inner.len()` required to be able to do `push`
    prepush_len: usize,

    // `inner.len()` after `drop`
    postdrop_len: usize,
}

impl<'a, T: ?Sized> Drop for VecStack<'a, T> {
    fn drop(&mut self) {
        unsafe { self.inner.as_mut().unwrap() }.truncate(self.postdrop_len);
    }
}

impl<'outer, T: ?Sized> VecStack<'outer, T> {
    /// Construct a `VecStack` from a `Vec`. Once constructed, the `Vec`
    /// becomes fully managed by this `VecStack` until this `VecStack` is
    /// dropped. Use `as_ref` for reading, `push` for writing.
    pub fn new(vec: &'outer mut Vec<&'outer T>) -> VecStack<'outer, T> {
        let len = vec.len();
        VecStack {
            inner: vec,
            prepush_len: len,
            postdrop_len: len,
        }
    }

    /// Push a reference with a narrowed lifetime. Return a new `VecStack`.
    ///
    /// Panic if `push` is called when a previously pushed element is
    /// not popped (ex. the result of a previous `push` is not dropped).
    ///
    /// ```
    /// # use std::sync::Mutex;
    /// # use treestate::vecstack::VecStack;
    /// let mut vec = vec![];
    /// let mut stack0 = Mutex::new(VecStack::new(&mut vec));
    /// let str1 = String::from("1");
    /// let mut stack1 = stack0.lock().unwrap().push(&str1);
    /// let str2 = String::from("2");
    /// let result = std::panic::catch_unwind(|| {
    ///     let mut stack2 = stack0.lock().unwrap().push(&str2); // panic since `stack1` is alive
    /// });
    /// assert!(result.is_err());
    /// ```
    pub fn push<'inner>(&mut self, elem: &'inner T) -> VecStack<'inner, T>
    where
        'outer: 'inner,
    {
        // This casts `vec` from `Vec<&'outer T>` to `Vec<&'inner T>`. This is safe because:
        // - `'inner` is a narrow lifetime. Guaranteed by the where clause. So reading `'outer`
        //   using `'inner` lifetime is fine.
        // - As long as the new `VecStack` is alive, `'inner` is valid. Guaranteed by the lifetime
        //   checker.
        // - Once `VecStack` is dead, `'inner` reference will be removed from `vec`. Guaranteed by
        //   `Drop`.
        // - As long as the new `VecStack` is alive, reading (`as_ref`) via `'outer` is forbidden.
        //   Guaranteed by the assert check in `as_ref`.
        // - As long as the new `VecStack` is alive, writing (`push`) via `'outer` is forbidden.
        //   Guaranteed by the assert check here.
        assert_eq!(
            unsafe { self.inner.as_ref().unwrap() }.len(),
            self.prepush_len,
            "cannot push if vec is changed"
        );
        let vec: &'inner mut Vec<&'inner T> =
            unsafe { mem::transmute(self.inner as *mut Vec<&'outer T>) };
        let postdrop_len = vec.len();
        vec.push(elem);
        let prepush_len = vec.len();
        VecStack {
            inner: vec,
            prepush_len,
            postdrop_len,
        }
    }
}

impl<'a, T: ?Sized> AsRef<Vec<&'a T>> for VecStack<'a, T> {
    /// Get a read-only reference of the actual `Vec`.
    ///
    /// When a reference returned by `as_ref` is alive, it's not allowed
    /// to push again. This is checked statically.
    ///
    /// ```compile_fail
    /// # use treestate::vecstack::VecStack;
    /// let mut vec = vec![];
    /// let mut stack = VecStack::new(&mut vec);
    /// let vecref = stack.as_ref();
    /// let str1 = String::from("1");
    /// let _ = stack.push(&str1); // cannot push when `vecref` is alive.
    /// let __ = vecref;
    /// ```
    ///
    /// Panic if `as_ref` is called when a previously pushed element is
    /// not poped (ex. the result of a previous `push` is not dropped).
    ///
    /// ```
    /// # use treestate::vecstack::VecStack;
    /// let mut vec = vec![];
    /// let mut stack0 = VecStack::new(&mut vec);
    /// let str1 = String::from("1");
    /// let mut stack1 = stack0.push(&str1);
    /// let result = std::panic::catch_unwind(|| {
    ///     let ref0 = stack0.as_ref(); // panic since `stack1` is alive
    /// });
    /// assert!(result.is_err());
    /// ```
    fn as_ref(&self) -> &Vec<&'a T> {
        let result = unsafe { self.inner.as_ref().unwrap() };
        assert_eq!(
            result.len(),
            self.prepush_len,
            "cannot get reference if vec is changed"
        );
        result
    }
}


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

use anyhow::Error;
use thiserror::Error;

#[derive(Debug, Error)]
#[error("Key Error: {0:?}")]
pub struct KeyError(#[source] Error);

impl KeyError {
    pub fn new(err: Error) -> Self {
        KeyError(err)
    }
}

/// NeworkError is a wrapper/tagging error meant for libraries to use
/// to mark errors that may imply a network problem.
#[derive(Debug, Error)]
#[error("Network Error: {0:?}")]
pub struct NetworkError(#[source] pub Error);

impl NetworkError {
    pub fn wrap(err: impl Into<Error>) -> Error {
        Self(err.into()).into()
    }
}

pub fn is_network_error(err: &Error) -> bool {
    has_in_chain::<NetworkError>(err)
}

fn has_in_chain<E>(err: &Error) -> bool
where
    E: std::error::Error + 'static,
{
    err.chain().any(|e| e.is::<E>())
}

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

    use super::*;

    #[test]
    fn test_network_error() {
        let network: &dyn std::error::Error =
            &NetworkError(io::Error::from(io::ErrorKind::Other).into());
        assert!(network.is::<NetworkError>());
        assert!(network.source().unwrap().is::<io::Error>());

        let network: Error = NetworkError(io::Error::from(io::ErrorKind::Other).into()).into();
        assert!(network.is::<NetworkError>());
        assert!(network.source().unwrap().is::<io::Error>());
        assert!(is_network_error(&network));

        assert_eq!(format!("{}", network), "Network Error: other error");

        let with_context = network.context("hello");
        assert!(is_network_error(&with_context));

        let wrapped: Error = KeyError(with_context).into();
        assert!(is_network_error(&wrapped));
    }
}