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

use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use ::types::Key;
use ::types::Node;
use ::types::RepoPath;
use ::types::RepoPathBuf;
use anyhow::anyhow;
use anyhow::bail;
use anyhow::Result;
use crossbeam::channel::unbounded;
use minibytes::Bytes;
use tracing::field;

pub mod types;

use crate::datastore::HgIdDataStore;
use crate::datastore::RemoteDataStore;
use crate::indexedlogdatastore::Entry;
use crate::indexedlogdatastore::IndexedLogHgIdDataStore;
use crate::memcache::MEMCACHE_DELAY;
use crate::scmstore::fetch::CommonFetchState;
use crate::scmstore::fetch::FetchErrors;
use crate::scmstore::fetch::FetchMode;
use crate::scmstore::fetch::FetchResults;
use crate::scmstore::fetch::KeyFetchError;
use crate::scmstore::file::FileStore;
use crate::scmstore::tree::types::LazyTree;
use crate::scmstore::tree::types::StoreTree;
use crate::scmstore::tree::types::TreeAttributes;
use crate::util;
use crate::ContentDataStore;
use crate::ContentMetadata;
use crate::ContentStore;
use crate::Delta;
use crate::EdenApiTreeStore;
use crate::HgIdMutableDeltaStore;
use crate::LegacyStore;
use crate::LocalStore;
use crate::MemcacheStore;
use crate::Metadata;
use crate::RepackLocation;
use crate::StoreKey;
use crate::StoreResult;

#[derive(Clone)]
pub struct TreeStore {
    /// The "local" indexedlog store. Stores content that is created locally.
    pub indexedlog_local: Option<Arc<IndexedLogHgIdDataStore>>,

    /// The "cache" indexedlog store (previously called "shared"). Stores content downloaded from
    /// a remote store.
    pub indexedlog_cache: Option<Arc<IndexedLogHgIdDataStore>>,

    /// If cache_to_local_cache is true, data found by falling back to a remote store
    /// will the written to indexedlog_cache.
    pub cache_to_local_cache: bool,

    /// If provided, memcache will be checked before other remote stores
    pub memcache: Option<Arc<MemcacheStore>>,

    /// If cache_to_memcache is true, data found by falling back to another remote store
    // will be written to memcache.
    pub cache_to_memcache: bool,

    /// An EdenApi Client, EdenApiTreeStore provides the tree-specific subset of EdenApi functionality
    /// used by TreeStore.
    pub edenapi: Option<Arc<EdenApiTreeStore>>,

    /// Hook into the legacy storage architecture, if we fall back to this and succeed, we
    /// should alert / log something, as this should never happen if TreeStore is implemented
    /// correctly.
    pub contentstore: Option<Arc<ContentStore>>,

    /// A FileStore, which can be used for fetching and caching file aux data for a tree.
    pub filestore: Option<Arc<FileStore>>,

    pub creation_time: Instant,

    pub flush_on_drop: bool,
}

impl Drop for TreeStore {
    fn drop(&mut self) {
        if self.flush_on_drop {
            let _ = self.flush();
        }
    }
}

impl TreeStore {
    pub fn fetch_batch(
        &self,
        reqs: impl Iterator<Item = Key>,
        fetch_mode: FetchMode,
    ) -> FetchResults<StoreTree> {
        let (found_tx, found_rx) = unbounded();
        let found_tx2 = found_tx.clone();
        let mut common: CommonFetchState<StoreTree> =
            CommonFetchState::new(reqs, TreeAttributes::CONTENT, found_tx);

        let keys_len = common.pending_len();

        let indexedlog_cache = self.indexedlog_cache.clone();
        let indexedlog_local = self.indexedlog_local.clone();
        let memcache = self.memcache.clone();
        let edenapi = self.edenapi.clone();

        let contentstore = self.contentstore.clone();
        let creation_time = self.creation_time;
        let cache_to_memcache = self.cache_to_memcache;
        let cache_to_local_cache = self.cache_to_local_cache;
        let (aux_local, aux_cache) = if let Some(ref filestore) = self.filestore {
            (filestore.aux_local.clone(), filestore.aux_cache.clone())
        } else {
            (None, None)
        };
        let process_func = move || -> Result<()> {
            if let Some(ref indexedlog_cache) = indexedlog_cache {
                let pending: Vec<_> = common
                    .pending(TreeAttributes::CONTENT, false)
                    .map(|(key, _attrs)| key.clone())
                    .collect();
                for key in pending.into_iter() {
                    if let Some(entry) = indexedlog_cache.get_entry(key)? {
                        tracing::trace!("{:?} found in cache", &entry.key());
                        common.found(entry.key().clone(), LazyTree::IndexedLog(entry).into());
                    }
                }
            }

            if let Some(ref indexedlog_local) = indexedlog_local {
                let pending: Vec<_> = common
                    .pending(TreeAttributes::CONTENT, false)
                    .map(|(key, _attrs)| key.clone())
                    .collect();
                for key in pending.into_iter() {
                    if let Some(entry) = indexedlog_local.get_entry(key)? {
                        tracing::trace!("{:?} found in local", &entry.key());
                        common.found(entry.key().clone(), LazyTree::IndexedLog(entry).into());
                    }
                }
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if use_memcache(creation_time) {
                    if let Some(ref memcache) = memcache {
                        let pending: Vec<_> = common
                            .pending(TreeAttributes::CONTENT, false)
                            .map(|(key, _attrs)| key.clone())
                            .collect();

                        if !pending.is_empty() {
                            for entry in memcache.get_data_iter(&pending)? {
                                let entry = entry?;
                                let key = entry.key.clone();
                                let entry = LazyTree::Memcache(entry);
                                if indexedlog_cache.is_some() && cache_to_local_cache {
                                    if let Some(entry) =
                                        entry.indexedlog_cache_entry(key.clone())?
                                    {
                                        indexedlog_cache.as_ref().unwrap().put_entry(entry)?;
                                    }
                                }
                                tracing::trace!("{:?} found in memcache", &key);
                                common.found(key, entry.into());
                            }
                        }
                    }
                }
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if let Some(ref edenapi) = edenapi {
                    let pending: Vec<_> = common
                        .pending(TreeAttributes::CONTENT, false)
                        .map(|(key, _attrs)| key.clone())
                        .collect();
                    if !pending.is_empty() {
                        let span = tracing::info_span!(
                            "fetch_edenapi",
                            downloaded = field::Empty,
                            uploaded = field::Empty,
                            requests = field::Empty,
                            time = field::Empty,
                            latency = field::Empty,
                            download_speed = field::Empty,
                        );
                        let _enter = span.enter();
                        tracing::debug!(
                            "attempt to fetch {} keys from edenapi ({:?})",
                            pending.len(),
                            edenapi.url()
                        );
                        let attributes = if aux_local.is_some() {
                            Some(edenapi_types::TreeAttributes {
                                child_metadata: true,
                                ..edenapi_types::TreeAttributes::default()
                            })
                        } else {
                            None
                        };
                        let response = edenapi
                            .trees_blocking(pending, attributes)
                            .map_err(|e| e.tag_network())?;
                        for entry in response.entries {
                            let entry = entry?;
                            let key = entry.key.clone();
                            let entry = LazyTree::EdenApi(entry);

                            if let Some(ref aux_local) = aux_local {
                                let aux_data = entry.aux_data();
                                for (hgid, aux) in aux_data.into_iter() {
                                    if let Some(ref aux_cache) = aux_cache {
                                        aux_cache.put(hgid, &aux.into())?;
                                    } else {
                                        aux_local.put(hgid, &aux.into())?;
                                    }
                                }
                            }
                            if indexedlog_cache.is_some() && cache_to_local_cache {
                                if let Some(entry) = entry.indexedlog_cache_entry(key.clone())? {
                                    indexedlog_cache.as_ref().unwrap().put_entry(entry)?;
                                }
                            }
                            if memcache.is_some()
                                && cache_to_memcache
                                && use_memcache(creation_time)
                            {
                                if let Some(entry) = entry.indexedlog_cache_entry(key.clone())? {
                                    memcache.as_ref().unwrap().add_mcdata(entry.try_into()?);
                                }
                            }
                            common.found(key, entry.into());
                        }
                        util::record_edenapi_stats(&span, &response.stats);
                    }
                } else {
                    tracing::debug!("no EdenApi associated with TreeStore");
                }
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if let Some(ref contentstore) = contentstore {
                    let pending: Vec<_> = common
                        .pending(TreeAttributes::CONTENT, false)
                        .map(|(key, _attrs)| StoreKey::HgId(key.clone()))
                        .collect();
                    if !pending.is_empty() {
                        tracing::debug!(
                            "attempt to fetch {} keys from contentstore",
                            pending.len()
                        );
                        contentstore.prefetch(&pending)?;

                        let pending = pending.into_iter().map(|key| match key {
                            StoreKey::HgId(key) => key,
                            // Safe because we constructed pending with only StoreKey::HgId above
                            // we're just re-using the already allocated paths in the keys
                            _ => unreachable!("unexpected non-HgId StoreKey"),
                        });

                        for key in pending {
                            let store_key = StoreKey::HgId(key.clone());
                            let blob = match contentstore.get(store_key.clone())? {
                                StoreResult::Found(v) => Some(v),
                                StoreResult::NotFound(_k) => None,
                            };
                            let meta = match contentstore.get_meta(store_key)? {
                                StoreResult::Found(v) => Some(v),
                                StoreResult::NotFound(_k) => None,
                            };

                            if let (Some(blob), Some(meta)) = (blob, meta) {
                                // We don't write to local indexedlog or memcache for contentstore fallbacks because
                                // contentstore handles that internally.
                                tracing::trace!("{:?} found in contentstore", &key);
                                common.found(key, LazyTree::ContentStore(blob.into(), meta).into());
                            }
                        }
                    }
                }
            }

            // TODO(meyer): Report incomplete / not found, handle errors better instead of just always failing the batch, etc
            common.results(FetchErrors::new());
            Ok(())
        };
        let process_func_errors = move || {
            if let Err(err) = process_func() {
                let _ = found_tx2.send(Err(KeyFetchError::Other(err)));
            }
        };

        // Only kick off a thread if there's a substantial amount of work.
        if keys_len > 1000 {
            std::thread::spawn(process_func_errors);
        } else {
            process_func_errors();
        }

        FetchResults::new(Box::new(found_rx.into_iter()))
    }

    fn write_batch(&self, entries: impl Iterator<Item = (Key, Bytes, Metadata)>) -> Result<()> {
        if let Some(ref indexedlog_local) = self.indexedlog_local {
            for (key, bytes, meta) in entries {
                indexedlog_local.put_entry(Entry::new(key, bytes, meta))?;
            }
        }
        Ok(())
    }

    pub fn empty() -> Self {
        TreeStore {
            indexedlog_local: None,

            indexedlog_cache: None,
            cache_to_local_cache: true,

            memcache: None,
            cache_to_memcache: true,

            edenapi: None,

            contentstore: None,

            filestore: None,
            creation_time: Instant::now(),
            flush_on_drop: true,
        }
    }

    #[allow(unused_must_use)]
    pub fn flush(&self) -> Result<()> {
        let mut result = Ok(());
        let mut handle_error = |error| {
            tracing::error!(%error);
            result = Err(error);
        };

        if let Some(ref indexedlog_local) = self.indexedlog_local {
            indexedlog_local.flush_log().map_err(&mut handle_error);
        }

        if let Some(ref indexedlog_cache) = self.indexedlog_cache {
            indexedlog_cache.flush_log().map_err(&mut handle_error);
        }

        result
    }

    pub fn refresh(&self) -> Result<()> {
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.refresh()?;
        }
        self.flush()
    }

    pub fn with_content_store(&self, cs: Arc<ContentStore>) -> Self {
        let mut clone = self.clone();
        clone.contentstore = Some(cs);
        clone
    }
}

fn use_memcache(creation_time: Instant) -> bool {
    // Only use memcache if the process has been around a while. It takes 2s to setup, which
    // hurts responiveness for short commands.
    creation_time.elapsed() > MEMCACHE_DELAY
}

impl LegacyStore for TreeStore {
    /// Returns only the local cache / shared stores, in place of the local-only stores, such that writes will go directly to the local cache.
    /// For compatibility with ContentStore::get_shared_mutable
    fn get_shared_mutable(&self) -> Arc<dyn HgIdMutableDeltaStore> {
        // this is infallible in ContentStore so panic if there are no shared/cache stores.
        assert!(
            self.indexedlog_cache.is_some(),
            "cannot get shared_mutable, no shared / local cache stores available"
        );
        Arc::new(TreeStore {
            indexedlog_local: self.indexedlog_cache.clone(),
            indexedlog_cache: None,
            cache_to_local_cache: false,

            memcache: None,
            cache_to_memcache: false,

            edenapi: None,
            contentstore: None,

            filestore: None,
            creation_time: Instant::now(),
            flush_on_drop: true,
        })
    }

    fn get_logged_fetches(&self) -> HashSet<RepoPathBuf> {
        unimplemented!(
            "get_logged_fetches is not implemented for trees, it should only ever be falled for files"
        );
    }

    fn get_file_content(&self, _key: &Key) -> Result<Option<Bytes>> {
        unimplemented!(
            "get_file_content is not implemented for trees, it should only ever be falled for files"
        );
    }

    // If ContentStore is available, these call into ContentStore. Otherwise, implement these
    // methods on top of scmstore (though they should still eventaully be removed).
    fn add_pending(
        &self,
        key: &Key,
        data: Bytes,
        meta: Metadata,
        location: RepackLocation,
    ) -> Result<()> {
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.add_pending(key, data, meta, location)
        } else {
            let delta = Delta {
                data,
                base: None,
                key: key.clone(),
            };

            match location {
                RepackLocation::Local => self.add(&delta, &meta),
                RepackLocation::Shared => self.get_shared_mutable().add(&delta, &meta),
            }
        }
    }

    fn commit_pending(&self, location: RepackLocation) -> Result<Option<Vec<PathBuf>>> {
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.commit_pending(location)
        } else {
            self.flush()?;
            Ok(None)
        }
    }
}

impl HgIdDataStore for TreeStore {
    // Fetch the raw content of a single TreeManifest blob
    fn get(&self, key: StoreKey) -> Result<StoreResult<Vec<u8>>> {
        Ok(
            match self
                .fetch_batch(std::iter::once(key.clone()).filter_map(StoreKey::maybe_into_key), FetchMode::AllowRemote)
                .single()?
            {
                Some(entry) => StoreResult::Found(entry.content.expect("content attribute not found despite being requested and returned as complete").hg_content()?.into_vec()),
                None => StoreResult::NotFound(key),
            },
        )
    }

    fn get_meta(&self, key: StoreKey) -> Result<StoreResult<Metadata>> {
        Ok(
            match self
                .fetch_batch(
                    std::iter::once(key.clone()).filter_map(StoreKey::maybe_into_key),
                    FetchMode::AllowRemote,
                )
                .single()?
            {
                // This is currently in a bit of an awkward state, as revisionstore metadata is no longer used for trees
                // (it should always be default), but the get_meta function should return StoreResult::Found
                // only when the content is available. Thus, we request the tree content, but ignore it and just
                // return default metadata when it's found, and otherwise report StoreResult::NotFound.
                // TODO(meyer): Replace this with an presence check once support for separate fetch and return attrs
                // is added.
                Some(_e) => StoreResult::Found(Metadata::default()),
                None => StoreResult::NotFound(key),
            },
        )
    }

    fn refresh(&self) -> Result<()> {
        self.refresh()
    }
}

impl RemoteDataStore for TreeStore {
    fn prefetch(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        Ok(self
            .fetch_batch(
                keys.iter().cloned().filter_map(StoreKey::maybe_into_key),
                FetchMode::AllowRemote,
            )
            .missing()?
            .into_iter()
            .map(StoreKey::HgId)
            .collect())
    }

    fn upload(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        Ok(keys.to_vec())
    }
}

impl LocalStore for TreeStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        let mut missing: Vec<_> = keys.to_vec();

        missing = if let Some(ref indexedlog_cache) = self.indexedlog_cache {
            missing
                .into_iter()
                .filter(
                    |sk| match sk.maybe_as_key().map(|k| indexedlog_cache.get_raw_entry(k)) {
                        Some(Ok(Some(_))) => false,
                        None | Some(Err(_)) | Some(Ok(None)) => true,
                    },
                )
                .collect()
        } else {
            missing
        };

        missing = if let Some(ref indexedlog_local) = self.indexedlog_local {
            missing
                .into_iter()
                .filter(
                    |sk| match sk.maybe_as_key().map(|k| indexedlog_local.get_raw_entry(k)) {
                        Some(Ok(Some(_))) => false,
                        None | Some(Err(_)) | Some(Ok(None)) => true,
                    },
                )
                .collect()
        } else {
            missing
        };

        Ok(missing)
    }
}

impl HgIdMutableDeltaStore for TreeStore {
    fn add(&self, delta: &Delta, metadata: &Metadata) -> Result<()> {
        if let Delta {
            data,
            base: None,
            key,
        } = delta.clone()
        {
            self.write_batch(std::iter::once((key, data, metadata.clone())))
        } else {
            bail!("Deltas with non-None base are not supported")
        }
    }

    fn flush(&self) -> Result<Option<Vec<PathBuf>>> {
        if let Some(ref indexedlog_local) = self.indexedlog_local {
            indexedlog_local.flush_log()?;
        }
        if let Some(ref indexedlog_cache) = self.indexedlog_cache {
            indexedlog_cache.flush_log()?;
        }
        Ok(None)
    }
}

// TODO(meyer): Content addressing not supported at all for trees. I could look for HgIds present here and fetch with
// that if available, but I feel like there's probably something wrong if this is called for trees. Do we need to implement
// this at all for trees?
impl ContentDataStore for TreeStore {
    fn blob(&self, key: StoreKey) -> Result<StoreResult<Bytes>> {
        Ok(StoreResult::NotFound(key))
    }

    fn metadata(&self, key: StoreKey) -> Result<StoreResult<ContentMetadata>> {
        Ok(StoreResult::NotFound(key))
    }
}

impl storemodel::TreeStore for TreeStore {
    fn get(&self, path: &RepoPath, node: Node) -> Result<minibytes::Bytes> {
        if node.is_null() {
            return Ok(Default::default());
        }

        let key = Key::new(path.to_owned(), node);
        match self
            .fetch_batch(std::iter::once(key.clone()), FetchMode::AllowRemote)
            .single()?
        {
            Some(entry) => Ok(entry.content.expect("no tree content").hg_content()?),
            None => Err(anyhow!("key {:?} not found in manifest", key)),
        }
    }

    fn prefetch(&self, keys: Vec<Key>) -> Result<()> {
        self.fetch_batch(keys.into_iter(), FetchMode::AllowRemote)
            .consume();
        Ok(())
    }

    fn insert(&self, _path: &RepoPath, _node: Node, _data: Bytes) -> Result<()> {
        unimplemented!("not needed yet");
    }
}

impl storemodel::RefreshableTreeStore for TreeStore {
    fn refresh(&self) -> Result<()> {
        self.refresh()
    }
}


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

use std::ops::BitAnd;
use std::ops::BitOr;
use std::ops::Not;
use std::ops::Sub;

use crate::scmstore::attrs::StoreAttrs;

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct TreeAttributes {
    pub content: bool,
}

impl StoreAttrs for TreeAttributes {
    const NONE: Self = TreeAttributes { content: false };

    /// Returns all the attributes which are present or can be computed from present attributes.
    fn with_computable(&self) -> TreeAttributes {
        *self
    }
}

impl TreeAttributes {
    pub const CONTENT: Self = TreeAttributes { content: true };
}

impl Not for TreeAttributes {
    type Output = Self;

    fn not(self) -> Self::Output {
        TreeAttributes {
            content: !self.content,
        }
    }
}

impl BitAnd for TreeAttributes {
    type Output = Self;

    fn bitand(self, rhs: Self) -> Self::Output {
        TreeAttributes {
            content: self.content & rhs.content,
        }
    }
}

impl BitOr for TreeAttributes {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        TreeAttributes {
            content: self.content | rhs.content,
        }
    }
}

/// The subtraction operator is implemented here to mean "set difference" aka relative complement.
impl Sub for TreeAttributes {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        self & !rhs
    }
}


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

use std::collections::HashMap;

use anyhow::Result;
use edenapi_types::TreeChildEntry;
use edenapi_types::TreeEntry;
use manifest_tree::TreeEntry as ManifestTreeEntry;
use minibytes::Bytes;
use storemodel::TreeFormat;
use types::HgId;
use types::Key;

use crate::indexedlogdatastore::Entry;
use crate::memcache::McData;
use crate::scmstore::file::FileAuxData;
use crate::Metadata;

/// A minimal tree enum that simply wraps the possible underlying tree types,
/// with no processing.
#[derive(Debug)]
pub(crate) enum LazyTree {
    /// A response from calling into the legacy storage API
    ContentStore(Bytes, Metadata),

    /// An entry from a local IndexedLog. The contained Key's path might not match the requested Key's path.
    IndexedLog(Entry),

    /// An EdenApi TreeEntry.
    EdenApi(TreeEntry),

    /// A memcache entry, convertable to Entry. In this case the Key's path should match the requested Key's path.
    Memcache(McData),
}

impl LazyTree {
    #[allow(dead_code)]
    fn hgid(&self) -> Option<HgId> {
        use LazyTree::*;
        match self {
            ContentStore(_, _) => None,
            IndexedLog(ref entry) => Some(entry.key().hgid),
            EdenApi(ref entry) => Some(entry.key().hgid),
            Memcache(ref entry) => Some(entry.key.hgid),
        }
    }

    /// The tree content, as would be encoded in the Mercurial blob
    pub(crate) fn hg_content(&mut self) -> Result<Bytes> {
        use LazyTree::*;
        Ok(match self {
            IndexedLog(ref mut entry) => entry.content()?,
            ContentStore(ref blob, _) => blob.clone(),
            EdenApi(ref entry) => entry.data()?.into(),
            Memcache(ref entry) => entry.data.clone(),
        })
    }

    /// Convert the LazyTree to an indexedlog Entry, if it should ever be written to IndexedLog cache
    pub(crate) fn indexedlog_cache_entry(&self, key: Key) -> Result<Option<Entry>> {
        use LazyTree::*;
        Ok(match self {
            IndexedLog(ref entry) => Some(entry.clone().with_key(key)),
            EdenApi(ref entry) => Some(Entry::new(key, entry.data()?.into(), Metadata::default())),
            // TODO(meyer): We shouldn't ever need to replace the key with Memcache, can probably just clone this.
            Memcache(ref entry) => Some({
                let entry: Entry = entry.clone().into();
                entry.with_key(key)
            }),
            // ContentStore handles caching internally
            ContentStore(_, _) => None,
        })
    }

    pub fn manifest_tree_entry(&mut self) -> Result<ManifestTreeEntry> {
        // TODO(meyer): Make manifest-tree crate use minibytes::Bytes
        // Currently revisionstore is only for hg format.
        let format = TreeFormat::Hg;
        Ok(ManifestTreeEntry(
            self.hg_content()?.into_vec().into(),
            format,
        ))
    }

    pub fn aux_data(&self) -> HashMap<HgId, FileAuxData> {
        use LazyTree::*;
        match self {
            EdenApi(entry) => entry
                .children
                .as_ref()
                .map_or_else(HashMap::new, |childrens| {
                    childrens
                        .iter()
                        .filter_map(|entry| {
                            let child_entry = match entry {
                                Err(err) => {
                                    tracing::warn!("Error fetching child entry: {:?}", err);
                                    return None;
                                }
                                Ok(file_entry) => file_entry,
                            };
                            // TODO: Also return directory aux data.
                            if let TreeChildEntry::File(file_entry) = child_entry {
                                file_entry.file_metadata.map(|file_metadata| {
                                    (
                                        file_entry.key.hgid,
                                        FileAuxData {
                                            total_size: file_metadata.size.unwrap(),
                                            content_id: file_metadata.content_id.unwrap(),
                                            content_sha1: file_metadata.content_sha1.unwrap(),
                                            content_sha256: file_metadata
                                                .content_sha256
                                                .unwrap()
                                                .into_byte_array()
                                                .into(),
                                            content_seeded_blake3: file_metadata
                                                .content_seeded_blake3,
                                        },
                                    )
                                })
                            } else {
                                None
                            }
                        })
                        .collect::<HashMap<_, _>>()
                }),
            _ => HashMap::new(),
        }
    }
}


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

mod attrs;
mod lazy_tree;
mod store_tree;

pub use self::attrs::TreeAttributes;
pub(crate) use self::lazy_tree::LazyTree;
pub use self::store_tree::StoreTree;


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

use std::collections::HashMap;
use std::ops::BitOr;

use anyhow::anyhow;
use anyhow::Result;
use manifest_tree::TreeEntry as ManifestTreeEntry;
use types::HgId;

use crate::scmstore::file::FileAuxData;
use crate::scmstore::tree::types::LazyTree;
use crate::scmstore::tree::types::TreeAttributes;
use crate::scmstore::value::StoreValue;

#[derive(Debug, Default)]
pub struct StoreTree {
    pub(crate) content: Option<LazyTree>,
}

impl StoreTree {
    pub fn manifest_tree_entry(&mut self) -> Result<ManifestTreeEntry> {
        self.content
            .as_mut()
            .ok_or_else(|| anyhow!("no content available"))?
            .manifest_tree_entry()
    }

    pub fn aux_data(&self) -> Result<HashMap<HgId, FileAuxData>> {
        Ok(self
            .content
            .as_ref()
            .ok_or_else(|| anyhow!("no content available"))?
            .aux_data())
    }
}

impl StoreValue for StoreTree {
    type Attrs = TreeAttributes;

    /// Returns which attributes are present in this StoreTree
    fn attrs(&self) -> TreeAttributes {
        TreeAttributes {
            content: self.content.is_some(),
        }
    }

    /// Return a StoreTree with only the specified subset of attributes
    fn mask(self, attrs: TreeAttributes) -> Self {
        StoreTree {
            content: if attrs.content { self.content } else { None },
        }
    }
}

impl BitOr for StoreTree {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        StoreTree {
            content: self.content.or(rhs.content),
        }
    }
}

impl From<LazyTree> for StoreTree {
    fn from(v: LazyTree) -> Self {
        StoreTree { content: Some(v) }
    }
}


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

use std::io::Cursor;
use std::io::Write;

use anyhow::format_err;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck_arbitrary_derive::Arbitrary;
use serde_derive::Deserialize;
use serde_derive::Serialize;

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
pub struct Metadata {
    pub size: Option<u64>,
    pub flags: Option<u64>,
}

impl Metadata {
    pub const LFS_FLAG: u64 = 0x2000;

    /// Returns true if the blob retrieved from `DataStore::get` is an LFS pointer.
    pub fn is_lfs(&self) -> bool {
        match self.flags {
            None => false,
            Some(flag) => (flag & Metadata::LFS_FLAG) == Metadata::LFS_FLAG,
        }
    }

    pub fn write<T: Write>(&self, writer: &mut T) -> Result<()> {
        let mut buf = vec![];
        if let Some(flags) = self.flags {
            if flags != 0 {
                Metadata::write_meta(b'f', flags, &mut buf)?;
            }
        }
        if let Some(size) = self.size {
            Metadata::write_meta(b's', size, &mut buf)?;
        }

        writer.write_u32::<BigEndian>(buf.len() as u32)?;
        writer.write_all(buf.as_ref())?;
        Ok(())
    }

    fn write_meta<T: Write>(flag: u8, value: u64, writer: &mut T) -> Result<()> {
        writer.write_u8(flag as u8)?;
        writer.write_u16::<BigEndian>(u64_to_bin_len(value))?;
        u64_to_bin(value, writer)?;
        Ok(())
    }

    pub fn read(cur: &mut Cursor<&[u8]>) -> Result<Metadata> {
        let metadata_len = cur.read_u32::<BigEndian>()? as u64;
        let mut size: Option<u64> = None;
        let mut flags: Option<u64> = None;
        let start_offset = cur.position();
        while cur.position() < start_offset + metadata_len {
            let key = cur.read_u8()?;
            let value_len = cur.read_u16::<BigEndian>()? as usize;
            match key {
                b'f' => {
                    let buf = cur.get_ref();
                    flags = Some(bin_to_u64(
                        &buf[cur.position() as usize..cur.position() as usize + value_len],
                    ));
                }
                b's' => {
                    let buf = cur.get_ref();
                    size = Some(bin_to_u64(
                        &buf[cur.position() as usize..cur.position() as usize + value_len],
                    ));
                }
                _ => return Err(format_err!("invalid metadata format '{:?}'", key)),
            }

            let cur_pos = cur.position();
            cur.set_position(cur_pos + value_len as u64);
        }

        Ok(Metadata { flags, size })
    }
}

/// Precompute the size of a u64 when it is serialized
fn u64_to_bin_len(value: u64) -> u16 {
    let mut value = value;
    let mut count = 0;
    while value > 0 {
        count += 1;
        value >>= 8;
    }
    count
}

/// Converts an integer into a buffer using a special format used in the datapack format.
fn u64_to_bin<T: Write>(value: u64, writer: &mut T) -> Result<()> {
    let mut value = value;
    let mut buf = [0; 8];
    let len = u64_to_bin_len(value) as usize;
    let mut pos = len;
    while value > 0 {
        pos -= 1;
        buf[pos] = value as u8;
        value >>= 8;
    }
    assert!(value == 0 && pos == 0);

    writer.write_all(&buf[0..len])?;
    Ok(())
}

/// Converts a buffer to an integer using a special format used in the datapack format.
fn bin_to_u64(buf: &[u8]) -> u64 {
    let mut n: u64 = 0;
    for byte in buf.iter() {
        n <<= 8;
        n |= *byte as u64;
    }
    n
}

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

    use super::*;
    quickcheck! {
        fn test_roundtrip_bin_to_u64(value: u64) -> bool {
            let mut buf: Vec<u8> = vec![];
            u64_to_bin(value, &mut buf).unwrap();
            if buf.len() != u64_to_bin_len(value) as usize {
                return false;
            }
            let new_value = bin_to_u64(&buf);
            value == new_value
        }

        fn test_roundtrip_metadata(size: Option<u64>, flags: Option<u64>) -> bool {
            let meta = Metadata { size, flags };
            let mut buf: Vec<u8> = vec![];
            meta.write(&mut buf).expect("write");
            let read_meta = Metadata::read(&mut Cursor::new(&buf)).expect("read");

            meta.size == read_meta.size && (meta.flags == read_meta.flags || meta.flags.map_or(false, |v| v == 0))
        }
    }
}


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

//! revisionstore_types - Data types used by the revisionstore crate

mod datastore;

pub use crate::datastore::Metadata;


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

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

use dag::errors::NotFoundError;
use dag::Id;
use dag::Vertex;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum RevlogIndexError {
    #[error(transparent)]
    CommitNotFound(#[from] CommitNotFound),

    #[error(transparent)]
    RevNotFound(#[from] RevNotFound),

    #[error("ambiguous prefix")]
    AmbiguousPrefix,

    // Collapse different kinds of corruption into one variant.
    // This helps keeping the enum sane.
    #[error(transparent)]
    Corruption(Box<CorruptionError>),

    #[error("unsupported: {0}")]
    Unsupported(String),
}

#[derive(Debug, Error)]
pub enum CorruptionError {
    #[error("{0}")]
    Generic(String),

    #[error(transparent)]
    Io(#[from] io::Error),

    #[error(transparent)]
    RadixTree(#[from] radixbuf::Error),

    #[error(transparent)]
    Lz4(#[from] lz4_pyframe::Error),

    #[error(transparent)]
    IndexedLog(#[from] indexedlog::Error),
}

impl From<radixbuf::Error> for RevlogIndexError {
    fn from(err: radixbuf::Error) -> Self {
        match err {
            radixbuf::Error::AmbiguousPrefix => Self::AmbiguousPrefix,
            _ => Self::Corruption(Box::new(err.into())),
        }
    }
}

impl From<lz4_pyframe::Error> for RevlogIndexError {
    fn from(err: lz4_pyframe::Error) -> Self {
        Self::Corruption(Box::new(err.into()))
    }
}

impl From<indexedlog::Error> for RevlogIndexError {
    fn from(err: indexedlog::Error) -> Self {
        Self::Corruption(Box::new(err.into()))
    }
}

// Currently, consider io::Error as a corruption error in this crate.
impl From<io::Error> for RevlogIndexError {
    fn from(err: io::Error) -> Self {
        Self::Corruption(Box::new(err.into()))
    }
}

pub fn corruption<T>(s: impl ToString) -> crate::Result<T> {
    Err(RevlogIndexError::Corruption(Box::new(
        CorruptionError::Generic(s.to_string()),
    )))
}

pub fn unsupported<T>(s: impl ToString) -> crate::Result<T> {
    Err(RevlogIndexError::Unsupported(s.to_string()))
}

impl From<CorruptionError> for RevlogIndexError {
    fn from(err: CorruptionError) -> RevlogIndexError {
        RevlogIndexError::Corruption(Box::new(err))
    }
}

#[derive(Debug)]
pub struct CommitNotFound(pub Vertex);

#[derive(Debug)]
pub struct RevNotFound(pub u32);

impl fmt::Display for CommitNotFound {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "00changelog.i@{:.8?}: not found", &self.0)
    }
}

impl std::error::Error for CommitNotFound {}

impl fmt::Display for RevNotFound {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "00changelog.i: rev {:.8?} not found", &self.0)
    }
}

impl std::error::Error for RevNotFound {}

impl From<RevlogIndexError> for dag::Error {
    fn from(err: RevlogIndexError) -> dag::Error {
        use dag::errors::BackendError;
        use RevlogIndexError as R;

        match err {
            R::CommitNotFound(CommitNotFound(vertex)) => vertex.not_found_error(),
            R::RevNotFound(RevNotFound(id)) => Id(id as _).not_found_error(),
            R::AmbiguousPrefix => {
                dag::Error::Bug("AmbiguousPrefix should not be translated".into())
            }
            R::Corruption(err) => match *err {
                CorruptionError::Generic(message) => BackendError::Generic(message),
                CorruptionError::Io(e) => BackendError::Io(e),
                CorruptionError::IndexedLog(e) => BackendError::IndexedLog(e),
                CorruptionError::RadixTree(e) => BackendError::Other(e.into()),
                CorruptionError::Lz4(e) => BackendError::Other(e.into()),
            }
            .into(),
            R::Unsupported(message) => dag::Error::Programming(message),
        }
    }
}


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

#![allow(dead_code)]

pub mod errors;
pub mod nodemap;
mod revlogindex;

pub use crate::errors::RevlogIndexError as Error;
pub use crate::nodemap::NodeRevMap;
pub use crate::revlogindex::RevlogEntry;
pub use crate::revlogindex::RevlogIndex;
pub type Result<T> = std::result::Result<T, Error>;


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

use radixbuf::key::KeyId;
use radixbuf::radix::radix_insert;
use radixbuf::radix::radix_lookup;
use radixbuf::radix::radix_lookup_unchecked;
use radixbuf::radix::radix_prefix_lookup;
use radixbuf::radix::RADIX_NCHILDREN;

use crate::errors::corruption;
use crate::Result;
use crate::RevlogEntry;

/// An index for node to rev lookups.
///
/// The index depends entirely on an append-only changelog.i source
/// of truth. It does not support in-memory overrides, which could be
/// implemented at a higher level.
///
/// ```text
///
///     changelogi
///   +------------+
///   | ... | node | < rev 0  \
///   +------------+           |
///   | ... | node | < rev 1   |> included in the main (on-disk) index
///   +------------+           |
///   | .......... | ...      /
///   +------------+
///   | ... | node | < next_index_rev         \
///   +------------+                           |
///   | ... | node | < next_index_rev + 1      |  will be built on-demand
///   +------------+                           |> in the side (in-memory)
///   | .......... | ...                       |  index
///   +------------+                           |
///   | ... | node | < next_changelog_rev - 1 /
///   +------------+
///                  < next_changelog_rev
/// ```
///
/// The main index is an immutable, periodically-rebuilt, on-disk radix buffer
/// with an extra metadata about what's the next revision unknown to the index.
/// The side index covers remaining revisions in changelogi, built on-demand and
/// is in-memory only. The side index is usually much smaller than the main index
/// so it can be built quickly.
///
/// ```text
///         main index               side index
///   +---------------------+  +----------------------+
///   | next_index_rev: u32 |  | (small radix buffer) |
///   +---------------------+  +----------------------+
///   |                     |      (in-memory only)
///   |(large radix buffer) |
///   |                     |
///   +---------------------+
///    (backed by filesystem)
/// ```
///
/// Having the side index allows us to make the main index immutable for most
/// of the time even if the source of truth has changed. It's possible to update
/// the main index in-place. But that requires extra efforts to deal with possible
/// filesystem issues like locking, or unexpected poweroff.
#[derive(Clone)]
pub struct NodeRevMap<C, I> {
    pub(crate) changelogi: C,
    main_index: I,             // Immutable main index
    side_index: Arc<Vec<u32>>, // Mutable side index
}

// Offsets in the main radix and key buffers
const RADIX_NEXT_REV_OFFSET: usize = 0;
const RADIX_HEADER_LEN: usize = RADIX_NEXT_REV_OFFSET + 1;

// Offsets of root nodes in radix buffers
const MAIN_RADIX_OFFSET: u32 = 1;
const SIDE_RADIX_OFFSET: u32 = 0;

impl<C: AsRef<[RevlogEntry]>, I: AsRef<[u32]>> NodeRevMap<C, I> {
    /// Initialize NodeMap from a non-inlined version of changelog.i and an incomplete index.
    pub fn new(changelogi: C, main_index: I) -> Result<Self> {
        // Sanity check if the index is corrupted or not.

        // The index must contain at least 17 elements. index[0] tracks the last rev the index has.
        // index[1..17] is the root radix node.
        if main_index.as_ref().len() < RADIX_HEADER_LEN + RADIX_NCHILDREN {
            return corruption("revlog radix index corrupted (main index too small)");
        }

        // Check if the index is behind and build incrementally
        let next_rev = u32::from_be(main_index.as_ref()[RADIX_NEXT_REV_OFFSET]);
        let end_rev = changelog_end_rev(&changelogi);

        if next_rev > end_rev {
            // next_rev cannot be larger than what changelogi has.
            return corruption("revlog radix index corrupted (next_rev > end_rev)");
        } else if next_rev > 0 {
            // Sanity check: if the last node stored in the index does not match the changelogi,
            // the index is broken and needs rebuilt. That could happen if strip happens.
            let rev: KeyId = (next_rev - 1).into();
            let node = rev_to_node(&changelogi, rev)?;
            if let Ok(Some(id)) = radix_lookup_unchecked(&main_index, MAIN_RADIX_OFFSET, &node) {
                if id != rev {
                    return corruption("revlog radix index corrupted (revlog out-of-sync)");
                }
            } else {
                return corruption("revlog radix index corrupted (revlog out-of-sync)");
            }
        }

        // Build side_index for the revisions not in the main index
        let mut side_index = vec![0u32; RADIX_NCHILDREN];
        build(
            &changelogi,
            &mut side_index,
            SIDE_RADIX_OFFSET,
            next_rev,
            end_rev,
        )?;
        let side_index = Arc::new(side_index);

        Ok(NodeRevMap {
            changelogi,
            main_index,
            side_index,
        })
    }

    /// Convert hex prefix to node.
    pub fn hex_prefix_to_node<T: AsRef<[u8]>>(&self, hex_prefix: T) -> Result<Option<&[u8]>> {
        let bin_prefix = match hex_to_bin_base16(hex_prefix) {
            Some(v) => v,
            None => return Ok(None),
        };
        let iter = bin_prefix.iter().cloned();
        let cl = &self.changelogi;
        let main_res = radix_prefix_lookup(
            &self.main_index,
            MAIN_RADIX_OFFSET,
            iter.clone(),
            rev_to_node,
            cl,
        )?;
        let side_res = radix_prefix_lookup(
            self.side_index.as_ref(),
            SIDE_RADIX_OFFSET,
            iter,
            rev_to_node,
            cl,
        )?;
        match (main_res, side_res) {
            (Some(_), Some(_)) => Err(crate::Error::AmbiguousPrefix),
            (Some(rev), None) | (None, Some(rev)) => Ok(Some(rev_to_node(&self.changelogi, rev)?)),
            _ => Ok(None),
        }
    }

    /// Convert node to rev.
    pub fn node_to_rev<T: AsRef<[u8]>>(&self, node: T) -> Result<Option<u32>> {
        let cl = &self.changelogi;
        if let Some(rev) = radix_lookup(&self.main_index, 1, &node, rev_to_node, cl)? {
            Ok(Some(rev.into()))
        } else if let Some(rev) = radix_lookup(self.side_index.as_ref(), 0, &node, rev_to_node, cl)?
        {
            Ok(Some(rev.into()))
        } else {
            Ok(None)
        }
    }

    /// How many revisions the side index has.
    pub fn lag(&self) -> u32 {
        let next_rev = u32::from_be(self.main_index.as_ref()[0]);
        let end_rev = changelog_end_rev(&self.changelogi);
        end_rev - next_rev
    }

    /// Incrementally build the main index based on the existing one.
    /// Note: this will memcpy the immutable main index so the new buffer
    /// could be written and resized.
    pub fn build_incrementally(&self) -> Result<Vec<u32>> {
        // Copy the main index since we need to modify it.
        let mut index = self.main_index.as_ref().to_vec();
        let end_rev = changelog_end_rev(&self.changelogi);
        let next_rev = u32::from_be(index[0]);
        build(
            &self.changelogi,
            &mut index,
            MAIN_RADIX_OFFSET,
            next_rev,
            end_rev,
        )?;
        index[0] = end_rev.to_be();
        Ok(index)
    }
}

/// Return the minimal revision number the changelog.i does not have.
fn changelog_end_rev<T: AsRef<[RevlogEntry]>>(changelogi: &T) -> u32 {
    let changelogi = changelogi.as_ref();
    let rev = changelogi.len();
    if rev as u64 > u32::MAX as u64 {
        panic!("rev exceeds 32 bit integers")
    }
    rev as u32
}

/// Read the given range of revisions (from `start_rev` (inclusive) to
/// `end_rev` (exclusive)) from changelogi. Insert them to the radix
/// index.
fn build<T>(
    changelogi: &T,
    index: &mut Vec<u32>,
    radix_offset: u32,
    start_rev: u32,
    end_rev: u32,
) -> Result<()>
where
    T: AsRef<[RevlogEntry]>,
{
    // Reserve the approximate size needed for the index - 28 bytes for each revision.
    // See D1291 for a table of number of revisions and index sizes.
    index.reserve(7 * (end_rev - start_rev) as usize);
    for i in start_rev..end_rev {
        let _ = radix_insert(index, radix_offset, i.into(), rev_to_node, changelogi)?;
    }
    Ok(())
}

/// Helper method similar to `radixbuf::key::FixedKey::read`, but takes a revision number instead.
fn rev_to_node<K: AsRef<[RevlogEntry]>>(changelogi: &K, rev: KeyId) -> radixbuf::Result<&[u8]> {
    let buf = changelogi.as_ref();
    let rev_usize: usize = rev.into();
    let entry = &buf[rev_usize];
    Ok(&entry.node[..])
}

/// Convert hex base16 sequence to binary base16 sequence.
fn hex_to_bin_base16<T: AsRef<[u8]>>(base16: T) -> Option<Vec<u8>> {
    let base16 = base16.as_ref();
    let len = base16.len();
    let mut result = vec![0u8; len];
    for (i, &ch) in base16.iter().enumerate() {
        result[i] = match ch {
            b'a'..=b'f' => ch - b'a' + 10,
            b'A'..=b'F' => ch - b'A' + 10,
            b'0'..=b'9' => ch - b'0',
            _ => return None,
        }
    }
    Some(result)
}

/// Return an empty index that can be used as "main_index" when passed to `new`.
pub fn empty_index_buffer() -> Vec<u8> {
    return vec![0u8; 4 * (RADIX_HEADER_LEN + RADIX_NCHILDREN)];
}