/*
 * 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::ffi::OsStr;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::time::Duration;
use std::time::SystemTime;

use anyhow::Error;
use anyhow::Result;
use fs2::FileExt;
use util::lock::PathLock;
use util::path::create_shared_dir;

use crate::Entry;

pub struct FileStore {
    // Directory to write files to (normally .hg/runlog/).
    dir: PathBuf,

    // Path to runlog watchfile (normally .hg/runlog_watchfile).
    watchfile_path: PathBuf,

    // A lock file indicating we are still running.
    #[allow(dead_code)]
    lock_file: PathLock,

    // Whether the current command is boring. This determines whether
    // we touch the watchfile on updates and how aggressively we clean
    // up files.
    boring: bool,
}

const LOCK_EXT: &str = "lock";
const JSON_EXT: &str = "json";
const WATCHFILE: &str = "runlog_watchfile";
const RUNLOG_DIR: &str = "runlog";

/// FileStore is a simple runlog storage that writes JSON entries to a
/// specified directory.
impl FileStore {
    // Create a new FileStore that writes files to directory dir. dir
    // is created automatically if it doesn't exist.
    pub(crate) fn new(shared_dot_hg_dir: &Path, entry_id: &str, boring: bool) -> Result<Self> {
        let dir = shared_dot_hg_dir.join(RUNLOG_DIR);

        create_shared_dir(&dir)?;

        let lock_file = PathLock::exclusive(dir.join(entry_id).with_extension(LOCK_EXT))?;

        Ok(FileStore {
            dir,
            watchfile_path: shared_dot_hg_dir.join(WATCHFILE),
            lock_file,
            boring,
        })
    }

    pub(crate) fn save(&self, e: &Entry) -> Result<()> {
        // NB: we don't fsync so incomplete or empty JSON files are possible.
        util::file::atomic_write(&self.dir.join(&e.id).with_extension(JSON_EXT), |f| {
            serde_json::to_writer_pretty(f, e)?;
            Ok(())
        })?;

        if !self.boring && !e.command.is_empty() {
            // Contents aren't important, but it makes it easier to test.
            fs::write(&self.watchfile_path, &e.command[0])?;
        }

        Ok(())
    }

    pub(crate) fn close(&self, e: &Entry) -> Result<()> {
        // Remove inconsequential, clean-exitting runlog entries immediately.
        if self.boring && e.exit_code == Some(0) {
            let path = self.dir.join(&e.id);
            remove_file_ignore_missing(path.with_extension(LOCK_EXT))?;
            remove_file_ignore_missing(path.with_extension(JSON_EXT))?;
        }
        Ok(())
    }

    pub(crate) fn cleanup<P: AsRef<Path>>(shared_dot_hg_dir: P, threshold: Duration) -> Result<()> {
        let dir = shared_dot_hg_dir.as_ref().join(RUNLOG_DIR);

        create_shared_dir(&dir)?;

        for dir_entry in fs::read_dir(dir)? {
            let path = dir_entry?.path();

            let ext = match path.extension().and_then(OsStr::to_str) {
                Some(ext) => ext,
                _ => continue,
            };

            // Skip ".lock" files. This leaves ".json" and any stray tmp files.
            if ext == LOCK_EXT {
                continue;
            }

            // Command process is still running - don't clean up.
            if is_locked(&path)? {
                continue;
            }

            // Avoid trying to read the contents so we can clean up
            // incomplete files.
            let mtime = fs::metadata(&path)?.modified()?;
            if SystemTime::now().duration_since(mtime)? >= threshold {
                // Cleanup up ".json" or tmp file.
                remove_file_ignore_missing(&path)?;
                // Clean up lock file (we know command process isn't running anymore).
                remove_file_ignore_missing(path.with_extension(LOCK_EXT))?;
            }
        }

        Ok(())
    }

    // Iterates each entry, yielding the entry and whether the
    // associated command is still running.
    pub fn entry_iter<P: AsRef<Path>>(
        shared_dot_hg_path: P,
    ) -> Result<impl Iterator<Item = Result<(Entry, bool), Error>>> {
        let dir = shared_dot_hg_path.as_ref().join(RUNLOG_DIR);

        create_shared_dir(&dir)?;

        Ok(fs::read_dir(&dir)?.filter_map(|file| match file {
            Ok(file) => {
                // We only care about ".json" files.
                match file.path().extension().and_then(OsStr::to_str) {
                    Some(ext) if ext == JSON_EXT => {}
                    _ => return None,
                };

                match fs::File::open(file.path()) {
                    Ok(f) => Some(
                        serde_json::from_reader(&f)
                            .map_err(Error::new)
                            .and_then(|e| Ok((e, is_locked(file.path())?))),
                    ),
                    Err(err) if err.kind() == io::ErrorKind::NotFound => None,
                    Err(err) => Some(Err(Error::new(err))),
                }
            }
            Err(err) => Some(Err(Error::new(err))),
        }))
    }
}

fn remove_file_ignore_missing<P: AsRef<Path>>(path: P) -> io::Result<()> {
    fs::remove_file(&path).or_else(|err| match err.kind() {
        io::ErrorKind::NotFound => Ok(()),
        _ => Err(err),
    })
}

// Return whether path's corresponding locked file is exclusively
// locked (by running command). Return false if lock file doesn't exist.
fn is_locked<P: AsRef<Path>>(path: P) -> Result<bool> {
    match fs::File::open(path.as_ref().with_extension(LOCK_EXT)) {
        Ok(f) => Ok(f.try_lock_shared().is_err()),
        Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(false),
        Err(err) => Err(Error::new(err)),
    }
}

#[cfg(test)]
mod tests {
    #[cfg(unix)]
    use std::os::unix::prelude::MetadataExt;

    use tempfile::tempdir;

    use super::*;

    #[test]
    fn test_save() -> Result<()> {
        let td = tempdir()?;

        let fs = FileStore::new(td.path(), "some_id", false)?;
        let rl_dir = td.path().join(RUNLOG_DIR);

        // Make sure FileStore creates directory automatically.
        assert!(rl_dir.exists());

        let mut entry = Entry::new(vec!["some_command".to_string()]);

        let assert_entry = |e: &Entry| {
            let f = fs::File::open(rl_dir.join(&e.id).with_extension(JSON_EXT)).unwrap();
            let got: Entry = serde_json::from_reader(&f).unwrap();
            assert_eq!(&got, e);

            #[cfg(unix)]
            assert_eq!(0o644, 0o777 & f.metadata().unwrap().mode());
        };

        // Can create new entry.
        fs.save(&entry)?;
        assert_entry(&entry);

        // Can update existing entry.
        entry.pid = 1234;
        fs.save(&entry)?;
        assert_entry(&entry);

        Ok(())
    }

    #[test]
    fn test_cleanup() -> Result<()> {
        let td = tempdir()?;
        let e = Entry::new(vec!["foo".to_string()]);
        let entry_path = td
            .path()
            .join(RUNLOG_DIR)
            .join(&e.id)
            .with_extension(JSON_EXT);

        {
            let fs = FileStore::new(td.path(), &e.id, false)?;
            fs.save(&e)?;

            // Still locked, don't clean up.
            FileStore::cleanup(&td, Duration::ZERO)?;
            assert!(entry_path.exists());
        }

        // No longer locked since file store is closed, but haven't met threshold.
        FileStore::cleanup(&td, Duration::from_secs(3600))?;
        assert!(entry_path.exists());

        // Met threshold - delete.
        FileStore::cleanup(&td, Duration::ZERO)?;
        assert!(!entry_path.exists());

        // Don't delete the watchfile.
        assert!(td.path().join(WATCHFILE).exists());

        Ok(())
    }

    #[test]
    fn test_iter() -> Result<()> {
        let td = tempdir()?;

        let a = Entry::new(vec!["a".to_string()]);
        let a_fs = FileStore::new(td.path(), &a.id, false)?;
        a_fs.save(&a)?;

        let b = Entry::new(vec!["b".to_string()]);
        {
            let b_fs = FileStore::new(td.path(), &b.id, false)?;
            b_fs.save(&b)?;
        }

        let mut got: Vec<(Entry, bool)> = FileStore::entry_iter(td.path())?
            .map(Result::unwrap)
            .collect();
        got.sort_by(|a, b| a.0.command[0].cmp(&b.0.command[0]));

        assert_eq!(vec![(a, true), (b, false)], got);

        Ok(())
    }

    #[test]
    fn test_watchfile() -> Result<()> {
        let td = tempdir()?;

        // Should write out watchfile.
        {
            let e = Entry::new(vec!["exciting".to_string()]);
            let fs = FileStore::new(td.path(), &e.id, false)?;
            fs.save(&e)?;
        }

        let watchfile_path = td.path().join(WATCHFILE);
        assert_eq!(fs::read_to_string(&watchfile_path)?, "exciting");

        // Boring command, don't touch watchfile.
        {
            let e = Entry::new(vec!["boring".to_string()]);
            let fs = FileStore::new(td.path(), &e.id, true)?;
            fs.save(&e)?;
        }

        assert_eq!(fs::read_to_string(&watchfile_path)?, "exciting");

        // Should touch watchfile.
        {
            let e = Entry::new(vec!["amazing".to_string()]);
            let fs = FileStore::new(td.path(), &e.id, false)?;
            fs.save(&e)?;
        }

        assert_eq!(fs::read_to_string(&watchfile_path)?, "amazing");

        Ok(())
    }
}


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

mod filestore;

#[cfg(unix)]
use std::os::unix::prelude::MetadataExt;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use configmodel::Config;
use configmodel::ConfigExt;
pub use filestore::FileStore;
use parking_lot::Mutex;
use rand::distributions::Alphanumeric;
use rand::thread_rng;
use rand::Rng;
use repo::repo::Repo;
use serde::Deserialize;
use serde::Serialize;

/// Logger logs runtime information for a single hg command invocation.
pub struct Logger {
    entry: Mutex<Entry>,
    storage: Option<Mutex<FileStore>>,
}

impl Logger {
    /// Initialize a new logger and write out initial runlog entry.
    /// Respects runlog.enable config field.
    pub fn from_repo(repo: Option<&Repo>, command: Vec<String>) -> Result<Arc<Self>> {
        if let Some(repo) = repo {
            Self::new(repo.config(), repo.shared_dot_hg_path(), command)
        } else {
            Ok(Self::empty(command))
        }
    }

    fn empty(command: Vec<String>) -> Arc<Self> {
        Arc::new(Logger {
            entry: Mutex::new(Entry::new(command)),
            storage: None,
        })
    }

    fn new(config: &dyn Config, shared_path: &Path, command: Vec<String>) -> Result<Arc<Self>> {
        if !config.get_or("runlog", "enable", || false)?
            || accidentally_running_as_root(shared_path)
        {
            return Ok(Self::empty(command));
        }

        // Probabilistically clean up old entries to avoid doing the work every time.
        let cleanup_chance = config.get_or("runlog", "cleanup-chance", || 0.01)?;
        if cleanup_chance > rand::thread_rng().gen::<f64>() {
            let threshold = config.get_or("runlog", "cleanup-threshold", || 3600.0)?;
            FileStore::cleanup(shared_path, Duration::from_secs_f64(threshold))?;
        }

        let boring_commands: Vec<String> = config.get_or_default("runlog", "boring-commands")?;

        // This command is boring if it is in boring-commands, or it
        // looks like the invoker is disabling the blackbox.
        let boring = command.is_empty()
            || boring_commands.contains(&command[0])
            || (config.get_or_default::<String>("extensions", "blackbox")? == "!")
            || (config.get("blackbox", "track") == Some("".into()));

        let entry = Entry::new(command);
        let storage = Some(Mutex::new(FileStore::new(shared_path, &entry.id, boring)?));

        let logger = Self {
            entry: Mutex::new(entry),
            storage,
        };
        logger.write(&logger.entry.lock(), false)?;

        return Ok(Arc::new(logger));
    }

    pub fn close(&self, exit_code: i32) -> Result<()> {
        let mut entry = self.entry.lock();
        entry.exit_code = Some(exit_code);
        entry.end_time = Some(chrono::Utc::now());
        entry.progress = Vec::new();

        self.write(&entry, true)?;

        Ok(())
    }

    pub fn update_progress(&self, progress: Vec<Progress>) -> Result<()> {
        let mut entry = self.entry.lock();
        if entry.exit_code.is_none() && entry.update_status(progress) {
            self.write(&entry, false)?;
        }

        Ok(())
    }

    fn write(&self, e: &Entry, close: bool) -> Result<()> {
        if let Some(storage) = &self.storage {
            let storage = storage.lock();
            storage.save(e)?;

            if close {
                storage.close(e)?;
            }
        }

        Ok(())
    }
}

#[cfg(unix)]
fn accidentally_running_as_root(shared_path: &Path) -> bool {
    // Check if we are root and repo is not owned by root.

    if unsafe { libc::geteuid() } != 0 {
        return false;
    }

    match std::fs::metadata(shared_path) {
        Ok(m) => m.uid() != 0,
        // err on side of not writing files as root
        Err(_) => true,
    }
}

#[cfg(not(unix))]
fn accidentally_running_as_root(_: &Path) -> bool {
    false
}

/// Entry represents one runlog entry (i.e. a single hg command
/// execution).
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct Entry {
    pub id: String,
    pub command: Vec<String>,
    pub pid: u32,
    pub download_bytes: usize,
    pub upload_bytes: usize,
    pub start_time: chrono::DateTime<chrono::Utc>,
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    pub exit_code: Option<i32>,
    pub progress: Vec<Progress>,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct Progress {
    pub topic: String,
    pub unit: String,
    pub total: u64,
    pub position: u64,
}

impl Entry {
    fn new(command: Vec<String>) -> Self {
        Self {
            id: thread_rng()
                .sample_iter(Alphanumeric)
                .take(16)
                .map(char::from)
                .collect(),
            command,
            pid: std::process::id(),
            download_bytes: 0,
            upload_bytes: 0,
            start_time: chrono::Utc::now(),
            end_time: None,
            exit_code: None,
            progress: Vec::new(),
        }
    }

    /// Return whether anything changed in the entry
    pub fn update_status(&mut self, progress: Vec<Progress>) -> bool {
        let (download_bytes, upload_bytes, _) = hg_http::current_progress();
        macro_rules! try_to_update {
            ($original_stat:expr,$new_stat:expr) => {{
                if $original_stat == $new_stat {
                    false
                } else {
                    $original_stat = $new_stat;
                    true
                }
            }};
        }
        let progress_updated = try_to_update!(self.progress, progress);
        let downloaded_bytes_updated = try_to_update!(self.download_bytes, download_bytes);
        let upload_bytes_updated = try_to_update!(self.upload_bytes, upload_bytes);
        progress_updated || downloaded_bytes_updated || upload_bytes_updated
    }
}

impl Progress {
    pub fn new(bar: Arc<progress_model::ProgressBar>) -> Progress {
        let (position, total) = bar.position_total();
        return Progress {
            topic: bar.topic().to_string(),
            position,
            total,
            unit: bar.unit().to_string(),
        };
    }
}

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

    use tempfile::tempdir;

    use super::*;

    #[test]
    fn test_close() -> Result<()> {
        let mut cfg = BTreeMap::new();
        cfg.insert("runlog.boring-commands", "boring");
        cfg.insert("runlog.enable", "1");

        let cleaned_up_files = |cfg, name: &str, exit: i32| -> Result<bool> {
            let td = tempdir()?;
            let logger = Logger::new(cfg, td.path(), vec![name.to_string()])?;
            logger.close(exit)?;

            let got: Vec<String> = std::fs::read_dir(td.path().join("runlog"))?
                .map(|d| d.unwrap().path().to_string_lossy().to_string())
                .collect();

            Ok(got.is_empty())
        };

        // Boring commands that exit cleanly are removed immediately.
        assert!(cleaned_up_files(&cfg, "boring", 0)?);

        // Boring commands that exit uncleanly are still recorded.
        assert!(!cleaned_up_files(&cfg, "boring", 1)?);

        // Non-boring commands aren't deleted immediately.
        assert!(!cleaned_up_files(&cfg, "interesting", 0)?);

        // Infer boringness from blackbox disablement.
        let mut no_bb_cfg = cfg.clone();
        no_bb_cfg.insert("extensions.blackbox", "!");
        assert!(cleaned_up_files(&no_bb_cfg, "blackbox_disabled", 0)?);

        // Infer boringness from empty blackbox.trace.
        let mut no_bb_trace = cfg.clone();
        no_bb_trace.insert("blackbox.track", "");
        assert!(cleaned_up_files(&no_bb_trace, "blackbox_disabled", 0)?);

        Ok(())
    }

    #[test]
    fn test_empty_command() -> Result<()> {
        let td = tempdir()?;
        let mut cfg = BTreeMap::new();
        cfg.insert("runlog.enable", "1");

        // Don't crash.
        Logger::new(&cfg, td.path(), vec![])?;

        Ok(())
    }
}


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

use std::collections::HashMap;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;

use once_cell::sync::OnceCell;
use parking_lot::Mutex;
use parking_lot::MutexGuard;

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

pub fn init(config: &dyn configmodel::Config) {
    CONFIG.get_or_init(|| SamplingConfig::new(config).map(Arc::new));
}

pub fn flush() {
    if let Some(Some(sc)) = CONFIG.get() {
        let _ = sc.file().flush();
    }
}

#[derive(Debug)]
pub struct SamplingConfig {
    keys: HashMap<String, String>,
    file: Mutex<File>,
}

impl SamplingConfig {
    pub fn new(config: &dyn configmodel::Config) -> Option<Self> {
        let sample_categories: HashMap<String, String> = config
            .keys("sampling")
            .into_iter()
            .filter_map(|name| {
                if let Some(key) = name.strip_prefix("key.") {
                    if let Some(val) = config.get("sampling", &name) {
                        return Some((key.to_string(), val.to_string()));
                    }
                }
                None
            })
            .collect();
        if sample_categories.is_empty() {
            return None;
        }

        if let Some((output_file, okay_exists)) = sampling_output_file(config) {
            match OpenOptions::new()
                .create(okay_exists)
                .create_new(!okay_exists)
                .append(true)
                .open(&output_file)
            {
                Ok(file) => {
                    return Some(Self {
                        keys: sample_categories,
                        file: Mutex::new(file),
                    });
                }
                Err(err) => {
                    // This is expected for child commands that skirt the telemetry wrapper.
                    tracing::warn!(
                        ?err,
                        ?output_file,
                        "error opening sampling file (expected for child commands)"
                    );
                }
            }
        }

        None
    }

    pub fn category(&self, key: &str) -> Option<&str> {
        self.keys.get(key).map(|c| &**c)
    }

    pub fn file(&self) -> MutexGuard<File> {
        self.file.lock()
    }
}

// Returns tuple of output path and whether it's okay if the path already exists.
fn sampling_output_file(config: &dyn configmodel::Config) -> Option<(PathBuf, bool)> {
    let mut candidates: Vec<(PathBuf, bool)> = Vec::with_capacity(2);

    if let Ok(path) = std::env::var("SCM_SAMPLING_FILEPATH") {
        // Env var is not-okay-exists (i.e. only one process should respect this).
        candidates.push((path.into(), false));
    }

    if let Some(path) = config.get("sampling", "filepath") {
        // Config setting is okay to be shared across multiple commands (mainly
        // for test compat).
        candidates.push((path.to_string().into(), true));
    }

    candidates
        .into_iter()
        .find(|(path, _okay_exists)| path.parent().map_or(false, |d| d.exists()))
}


/*
 * 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::hash_map::Entry;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::hash::Hash;
use std::hash::Hasher;
use std::io;
use std::io::BufRead;
use std::io::BufReader;

use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::Future;
use once_cell::sync::Lazy;
use pathmatcher::DirectoryMatch;
use pathmatcher::Matcher as MatcherTrait;
use pathmatcher::PatternKind;
use pathmatcher::TreeMatcher;
use pathmatcher::UnionMatcher;
use regex::Regex;
use types::RepoPath;

#[derive(Default, Debug)]
pub struct Profile {
    // Where this profile came from (typically a file path).
    source: String,

    // [include], [exclude] and %include
    entries: Vec<ProfileEntry>,

    // [metadata]
    title: Option<String>,
    description: Option<String>,
    hidden: Option<String>,
    version: Option<String>,

    case_sensitive: bool,
}

/// Root represents the root sparse profile (usually .hg/sparse).
#[derive(Debug, Hash)]
pub struct Root(Profile);

#[derive(Debug, Clone, PartialEq)]
enum Pattern {
    Include(String),
    Exclude(String),
}

#[derive(Debug)]
enum ProfileEntry {
    // Pattern plus additional source for this rule (e.g. "hgrc.dynamic").
    Pattern(Pattern, Option<String>),
    Profile(String),
}

#[derive(PartialEq)]
enum SectionType {
    Include,
    Exclude,
    Metadata,
}

impl Pattern {
    fn as_str(&self) -> &str {
        match self {
            Self::Include(p) => p,
            Self::Exclude(p) => p,
        }
    }
}

impl SectionType {
    fn from_str(value: &str) -> Option<Self> {
        match value {
            "[include]" => Some(SectionType::Include),
            "[exclude]" => Some(SectionType::Exclude),
            "[metadata]" => Some(SectionType::Metadata),
            _ => None,
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error(transparent)]
    Io(#[from] io::Error),

    #[error("import cycle involving {0}")]
    ImportCycle(String),

    #[error(transparent)]
    Fetch(#[from] anyhow::Error),

    #[error("unsuppported pattern type {0}")]
    UnsupportedPattern(String),

    #[error(transparent)]
    GlobsetError(#[from] globset::Error),
}

impl Root {
    pub fn from_bytes(data: impl AsRef<[u8]>, source: String) -> Result<Self, io::Error> {
        Ok(Self(Profile::from_bytes(data, source)?))
    }

    pub async fn matcher<B: Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>(
        &self,
        mut fetch: impl FnMut(String) -> B + Send + Sync,
    ) -> Result<Matcher, Error> {
        if self.0.entries.is_empty() {
            return Ok(Matcher::always());
        }

        let mut matchers: Vec<TreeMatcher> = Vec::new();

        // List of rule origins per-matcher.
        let mut rule_origins: Vec<Vec<String>> = Vec::new();

        let mut rules: VecDeque<(Pattern, String)> = VecDeque::new();

        // Maintain the excludes-come-last ordering.
        let mut push_rule = |(pat, src)| match pat {
            Pattern::Exclude(_) => rules.push_back((pat, src)),
            Pattern::Include(_) => rules.push_front((pat, src)),
        };

        let prepare_rules =
            |rules: VecDeque<(Pattern, String)>| -> Result<(Vec<String>, Vec<String>), Error> {
                let mut matcher_rules = Vec::new();
                let mut origins = Vec::new();

                for (pat, src) in rules {
                    match sparse_pat_to_matcher_rule(&pat) {
                        Err(err) => {
                            tracing::error!(%err, ?pat, %src, "ignoring unsupported sparse pattern");
                        }
                        Ok(rules) => {
                            for expanded_rule in rules {
                                matcher_rules.push(expanded_rule);
                                origins.push(src.clone());
                            }
                        }
                    }
                }

                Ok((matcher_rules, origins))
            };

        let mut only_v1 = true;
        for entry in self.0.entries.iter() {
            match entry {
                ProfileEntry::Pattern(p, src) => push_rule((
                    p.clone(),
                    join_source(self.0.source.clone(), src.as_deref()),
                )),
                ProfileEntry::Profile(child_path) => {
                    let child = match fetch(child_path.clone()).await? {
                        Some(data) => Profile::from_bytes(data, child_path.clone())?,
                        None => continue,
                    };

                    let child_rules: VecDeque<(Pattern, String)> = child
                        .rules(&mut fetch)
                        .await?
                        .into_iter()
                        .map(|(p, s)| (p, format!("{} -> {}", self.0.source, s)))
                        .collect();

                    if child.is_v2() {
                        only_v1 = false;

                        let (matcher_rules, origins) = prepare_rules(child_rules)?;
                        matchers.push(TreeMatcher::from_rules(
                            matcher_rules.iter(),
                            self.0.case_sensitive,
                        )?);
                        rule_origins.push(origins);
                    } else {
                        for rule in child_rules {
                            push_rule(rule);
                        }
                    }
                }
            }
        }

        // If all user specified rules are exclude rules, add an
        // implicit "**" to provide the default include of everything.
        if only_v1 && (rules.is_empty() || matches!(&rules[0].0, Pattern::Exclude(_))) {
            rules.push_front((Pattern::Include("**".to_string()), "(builtin)".to_string()))
        }

        // This is for files such as .hgignore and .hgsparse-base, unrelated to the .hg directory.
        rules.push_front((
            Pattern::Include("glob:.hg*".to_string()),
            "(builtin)".to_string(),
        ));

        let (matcher_rules, origins) = prepare_rules(rules)?;
        matchers.push(TreeMatcher::from_rules(
            matcher_rules.iter(),
            self.0.case_sensitive,
        )?);
        rule_origins.push(origins);

        Ok(Matcher::new(matchers, rule_origins))
    }
}

impl Profile {
    fn from_bytes(data: impl AsRef<[u8]>, source: String) -> Result<Self, io::Error> {
        let mut prof = Profile {
            // Case insensitive matcher support breaks w/ too many rules, so
            // leave it disabled for now. This may need to be fixed if sparse
            // profiles are supposed to support case insensitivity.
            case_sensitive: true,
            ..Default::default()
        };
        let mut current_metadata_val: Option<&mut String> = None;
        let mut section_type = SectionType::Include;
        let mut dynamic_source: Option<String> = None;

        for (mut line_num, line) in BufReader::new(data.as_ref()).lines().enumerate() {
            line_num += 1;

            let line = line?;
            let trimmed = line.trim();

            // Ingore comments and emtpy lines.
            let mut chars = trimmed.chars();
            match chars.next() {
                None => continue,
                Some('#' | ';') => {
                    let comment = chars.as_str().trim();
                    if let Some((l, r)) = comment.split_once(&['=', ':']) {
                        match (l.trim(), r.trim()) {
                            // Allow a magic comment to specify additional
                            // source information for particular rules. This way
                            // it is backwards compatible with the python code
                            // if a config like this ever gets written out.
                            ("source", "") => dynamic_source = None,
                            ("source", src) => dynamic_source = Some(src.to_string()),
                            _ => {}
                        }
                    }
                    continue;
                }
                _ => {}
            }

            if let Some(p) = trimmed.strip_prefix("%include ") {
                let p = p.trim();
                if p.ends_with('/') {
                    tracing::warn!(%line, %source, line_num, "ignoring sparse %include ending with /");
                    continue;
                }

                prof.entries
                    .push(ProfileEntry::Profile(p.trim().to_string()));
            } else if let Some(section_start) = SectionType::from_str(trimmed) {
                section_type = section_start;
                current_metadata_val = None;
            } else if section_type == SectionType::Metadata {
                if line.starts_with(&[' ', '\t']) {
                    // Continuation of multiline value.
                    if let Some(ref mut val) = current_metadata_val {
                        val.push('\n');
                        val.push_str(trimmed);
                    } else {
                        tracing::warn!(%line, %source, line_num, "orphan metadata line");
                    }
                } else {
                    current_metadata_val = None;
                    if let Some((key, val)) = trimmed.split_once(&['=', ':']) {
                        let prof_val = match key.trim() {
                            "description" => &mut prof.description,
                            "title" => &mut prof.title,
                            "hidden" => &mut prof.hidden,
                            "version" => &mut prof.version,
                            _ => {
                                tracing::warn!(%line, %source, line_num, "ignoring uninteresting metadata key");
                                continue;
                            }
                        };

                        current_metadata_val = Some(prof_val.insert(val.trim().to_string()));
                    }
                }
            } else {
                if trimmed.starts_with('/') {
                    tracing::warn!(%line, %source, line_num, "ignoring sparse rule starting with /");
                    continue;
                }

                if section_type == SectionType::Include {
                    prof.entries.push(ProfileEntry::Pattern(
                        Pattern::Include(trimmed.to_string()),
                        dynamic_source.clone(),
                    ));
                } else {
                    prof.entries.push(ProfileEntry::Pattern(
                        Pattern::Exclude(trimmed.to_string()),
                        dynamic_source.clone(),
                    ));
                }
            }
        }

        prof.source = source;

        Ok(prof)
    }

    fn is_v2(&self) -> bool {
        if let Some(version) = &self.version {
            version == "2"
        } else {
            false
        }
    }

    // Recursively flatten this profile into a DFS ordered list of rules.
    // %import statements are resolved by fetching the imported profile's
    // contents using the fetch callback. Returns a vec of each Pattern paired
    // with a String describing its provenance.
    async fn rules<B: Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>(
        &self,
        mut fetch: impl FnMut(String) -> B + Send + Sync,
    ) -> Result<Vec<(Pattern, String)>, Error> {
        fn rules_inner<'a, B: Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>(
            prof: &'a Profile,
            fetch: &'a mut (dyn FnMut(String) -> B + Send + Sync),
            rules: &'a mut Vec<(Pattern, String)>,
            source: Option<&'a str>,
            // path => (contents, in_progress)
            seen: &'a mut HashMap<String, (Vec<u8>, bool)>,
        ) -> BoxFuture<'a, Result<(), Error>> {
            async move {
                let source = match source {
                    Some(history) => format!("{} -> {}", history, prof.source),
                    None => prof.source.clone(),
                };

                for entry in prof.entries.iter() {
                    match entry {
                        ProfileEntry::Pattern(p, psrc) => {
                            rules.push((p.clone(), join_source(source.clone(), psrc.as_deref())))
                        }
                        ProfileEntry::Profile(child_path) => {
                            let entry = seen.entry(child_path.clone());
                            let data = match entry {
                                Entry::Occupied(e) => match e.into_mut() {
                                    (_, true) => {
                                        return Err(Error::ImportCycle(child_path.clone()));
                                    }
                                    (data, false) => data,
                                },
                                Entry::Vacant(e) => {
                                    if let Some(data) = fetch(child_path.clone()).await? {
                                        &e.insert((data, true)).0
                                    } else {
                                        continue;
                                    }
                                }
                            };

                            let mut child = Profile::from_bytes(&data, child_path.clone())?;
                            rules_inner(&mut child, fetch, rules, Some(&source), seen).await?;

                            if let Some((_, in_progress)) = seen.get_mut(child_path) {
                                *in_progress = false;
                            }
                        }
                    }
                }

                Ok(())
            }
            .boxed()
        }

        let mut rules = Vec::new();
        rules_inner(self, &mut fetch, &mut rules, None, &mut HashMap::new()).await?;
        Ok(rules)
    }
}

impl Hash for Profile {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        for entry in self.entries.iter() {
            if let ProfileEntry::Pattern(pat, _) = entry {
                match pat {
                    Pattern::Include(_) => "include",
                    Pattern::Exclude(_) => "exclude",
                }
                .hash(state);
                pat.as_str().hash(state);
            }
        }
    }
}

fn join_source(main_source: String, opt_source: Option<&str>) -> String {
    match opt_source {
        None => main_source,
        Some(opt) => format!("{} ({})", main_source, opt),
    }
}

pub struct Matcher {
    always: bool,
    matchers: Vec<TreeMatcher>,
    // List of rule origins per-matcher.
    rule_origins: Vec<Vec<String>>,
}

impl Matcher {
    pub fn matches(&self, path: &RepoPath) -> anyhow::Result<bool> {
        if self.always {
            Ok(true)
        } else {
            let result = UnionMatcher::matches_file(self.matchers.iter(), path);
            tracing::trace!(%path, ?result, "matches");
            result
        }
    }

    pub fn explain(&self, path: &RepoPath) -> anyhow::Result<(bool, String)> {
        if self.always {
            return Ok((true, "implicit match due to empty profile".to_string()));
        }

        for (i, m) in self.matchers.iter().enumerate() {
            if let Some(idx) = m.matching_rule_indexes(path.as_str()).last() {
                let rule_origin = self
                    .rule_origins
                    .get(i)
                    .and_then(|o| o.get(*idx))
                    .map_or("(unknown)".to_string(), |o| o.clone());
                return Ok((m.matches(path.as_str()), rule_origin));
            }
        }

        Ok((false, "no rules matched".to_string()))
    }
}

impl MatcherTrait for Matcher {
    fn matches_directory(&self, path: &RepoPath) -> anyhow::Result<DirectoryMatch> {
        if self.always {
            Ok(DirectoryMatch::Everything)
        } else {
            let result = UnionMatcher::matches_directory(self.matchers.iter(), path);
            tracing::trace!(%path, ?result, "matches_directory");
            result
        }
    }

    fn matches_file(&self, path: &RepoPath) -> anyhow::Result<bool> {
        self.matches(path)
    }
}

impl Matcher {
    fn new(matchers: Vec<TreeMatcher>, rule_origins: Vec<Vec<String>>) -> Self {
        Self {
            always: false,
            matchers,
            rule_origins,
        }
    }
    fn always() -> Self {
        Self {
            always: true,
            rule_origins: Vec::new(),
            matchers: Vec::new(),
        }
    }
}

// Convert a sparse profile pattern into what the tree matcher
// expects. We only support "glob" and "path" pattern types.
fn sparse_pat_to_matcher_rule(pat: &Pattern) -> Result<Vec<String>, Error> {
    let (pat_type, pat_text) = pathmatcher::split_pattern(pat.as_str(), PatternKind::Glob);
    match pat_type {
        PatternKind::Glob | PatternKind::Path => {} // empty
        PatternKind::RE => match convert_regex_to_glob(pat_text) {
            Some(globs) => {
                return Ok(globs);
            }
            None => return Err(Error::UnsupportedPattern(pat_type.name().to_string())),
        },
        _ => {
            return Err(Error::UnsupportedPattern(pat_type.name().to_string()));
        }
    };

    let pats = match pat_type {
        PatternKind::Glob => pathmatcher::expand_curly_brackets(pat_text)
            .iter()
            .map(|s| pathmatcher::normalize_glob(s))
            .collect(),
        PatternKind::Path => vec![pathmatcher::normalize_glob(
            pathmatcher::plain_to_glob(pat_text).as_str(),
        )],
        _ => unreachable!(),
    };

    Ok(pats
        .into_iter()
        // Adjust glob to ensure sparse rules match everything below them.
        .map(make_recursive)
        .map(|p| match pat {
            Pattern::Exclude(_) => format!("!{}", p),
            Pattern::Include(_) => p,
        })
        .collect())
}

fn make_recursive(p: impl Into<String>) -> String {
    let p = p.into();
    if p.is_empty() || p.ends_with('/') {
        p + "**"
    } else {
        p + "/**"
    }
}

// Match patterns like "foo/(?!bar/)" which mean "include foo/ except foo/bar/".
static EXCLUDE_DIR_RE: Lazy<Regex> =
    Lazy::new(|| Regex::new(r"^\^?([\w._/]+)/\(\?!([\w._/]+)\)$").unwrap());

// Match patterns like "foo/(?:.*/)?bar(?:/|$)".
static ANY_DIR_RE: Lazy<Regex> =
    Lazy::new(|| Regex::new(r"^\^?([\w._/]+)/\(\?:\.\*/\)\?([\w._]+)(\(\?:/\|\$\))?$").unwrap());

// Match patterns like "foo/\..*/ to match dotfiles.
static DOT_FILES_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^\^?([\w._/]+)/\\.\.\*$").unwrap());

// Attempt to convert given regex pattern to glob(s). Only certain
// cases are handled to give best effort support for checking out
// older commits that still use regex patterns.
fn convert_regex_to_glob(pat: &str) -> Option<Vec<String>> {
    if let Some(caps) = EXCLUDE_DIR_RE.captures(pat) {
        let prefix = caps.get(1).unwrap().as_str();
        let excluded = caps.get(2).unwrap().as_str();
        return Some(vec![
            make_recursive(prefix),
            make_recursive(format!("!{}/{}", prefix, excluded)),
        ]);
    }

    if let Some(caps) = ANY_DIR_RE.captures(pat) {
        let prefix = caps.get(1).unwrap().as_str();
        let name = caps.get(2).unwrap().as_str();

        // Turn trailing (?:/|$) into trailing "/**".
        let end = match caps.get(3) {
            Some(_) => "/**",
            None => "",
        };
        return Some(vec![format!("{}/**/{}{}", prefix, name, end)]);
    }

    if let Some(caps) = DOT_FILES_RE.captures(pat) {
        let prefix = caps.get(1).unwrap().as_str();
        return Some(vec![format!("{}/.*/**", prefix)]);
    }

    None
}

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

    use super::*;

    // Returns a profile's (includes, excludes, profiles).
    fn split_prof(prof: &Profile) -> (Vec<&str>, Vec<&str>, Vec<&str>) {
        let (mut inc, mut exc, mut profs) = (vec![], vec![], vec![]);
        for entry in &prof.entries {
            match entry {
                ProfileEntry::Pattern(Pattern::Include(p), _) => inc.push(p.as_ref()),
                ProfileEntry::Pattern(Pattern::Exclude(p), _) => exc.push(p.as_ref()),
                ProfileEntry::Profile(p) => profs.push(p.as_ref()),
            }
        }
        (inc, exc, profs)
    }

    #[test]
    fn test_parsing() {
        let got = Profile::from_bytes(
            b"
; hello
  # there

a
[metadata]
boring = banana
title  =   foo
[include]
glob:b/**/z
/skip/me
%include  other.sparse
 [exclude]
c
/skip/me

[metadata]
	skip me
description:howdy
 doody
version : 123
hidden=your eyes
	only

",
            "test".to_string(),
        )
        .unwrap();

        assert_eq!(got.source, "test");

        let (inc, exc, profs) = split_prof(&got);
        assert_eq!(inc, vec!["a", "glob:b/**/z"]);
        assert_eq!(exc, vec!["c"]);
        assert_eq!(profs, vec!["other.sparse"]);

        assert_eq!(got.title.unwrap(), "foo");
        assert_eq!(got.description.unwrap(), "howdy\ndoody");
        assert_eq!(got.hidden.unwrap(), "your eyes\nonly");
        assert_eq!(got.version.unwrap(), "123");
    }

    #[tokio::test]
    async fn test_rules() -> anyhow::Result<()> {
        let base = b"
%include child

[include]
a

[metadata]
title = base
";

        let child = b"
%include grand_child

[include]
b

[metadata]
title = child
";

        let grand_child = b"
[include]
c

[metadata]
title = grand_child
";

        let base_prof = Profile::from_bytes(base, "test".to_string()).unwrap();

        let rules = base_prof
            .rules(|path| async move {
                match path.as_ref() {
                    "child" => Ok(Some(child.to_vec())),
                    "grand_child" => Ok(Some(grand_child.to_vec())),
                    _ => Err(anyhow!("not found")),
                }
            })
            .await?;

        assert_eq!(
            rules,
            vec![
                (
                    Pattern::Include("c".to_string()),
                    "test -> child -> grand_child".to_string()
                ),
                (
                    Pattern::Include("b".to_string()),
                    "test -> child".to_string()
                ),
                (Pattern::Include("a".to_string()), "test".to_string())
            ]
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_recursive_imports() {
        let a = b"%include b";
        let b = b"%include a";

        let a_prof = Profile::from_bytes(a, "test".to_string()).unwrap();

        let res = a_prof
            .rules(|path| async move {
                match path.as_ref() {
                    "a" => Ok(Some(a.to_vec())),
                    "b" => Ok(Some(b.to_vec())),
                    _ => Err(anyhow!("not found")),
                }
            })
            .await;

        assert_eq!(format!("{}", res.unwrap_err()), "import cycle involving b");
    }

    #[tokio::test]
    async fn test_resolve_imports_caching() {
        let a = b"
%include b
%include b
";

        let a_prof = Profile::from_bytes(a, "test".to_string()).unwrap();

        let mut fetch_count = 0;

        // Make sure we cache results from the callback.
        let res = a_prof
            .rules(|_path| {
                fetch_count += 1;
                assert_eq!(fetch_count, 1);
                async { Ok(Some(vec![])) }
            })
            .await;

        assert!(res.is_ok());
    }

    #[test]
    fn test_sparse_pat_to_matcher_rule() {
        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include("a_valid:path/bar".to_string())).unwrap(),
            vec!["a_valid:path/bar/**"]
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include("path:/foo/bar".to_string())).unwrap(),
            vec!["/foo/bar/**"]
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include("path:/foo//bar".to_string())).unwrap(),
            vec!["/foo/bar/**"]
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include("/foo/*/bar{1,{2,3}}/".to_string()))
                .unwrap(),
            vec!["/foo/*/bar1/**", "/foo/*/bar2/**", "/foo/*/bar3/**"],
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include("path:/foo/*/bar{1,{2,3}}/".to_string()))
                .unwrap(),
            vec!["/foo/\\*/bar\\{1,\\{2,3\\}\\}/**"],
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Exclude("glob:**".to_string())).unwrap(),
            vec!["!**/**"],
        );

        assert!(sparse_pat_to_matcher_rule(&Pattern::Include("re:.*".to_string())).is_err());

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include(r"re:foo/\..*".to_string())).unwrap(),
            vec![r"foo/.*/**"],
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include(r"re:foo/(?!bar/)".to_string())).unwrap(),
            vec![r"foo/**", "!foo/bar/**"],
        );

        assert_eq!(
            sparse_pat_to_matcher_rule(&Pattern::Include(
                r"re:foo/(?:.*/)?baz.txt(?:/|$)".to_string()
            ))
            .unwrap(),
            vec!["foo/**/baz.txt/**"],
        );

        // Don't unescape asterisks accidentally.
        assert!(sparse_pat_to_matcher_rule(&Pattern::Include(r"re:\*".to_string())).is_err());

        // Giver up on regex exclude patterns.
        assert!(sparse_pat_to_matcher_rule(&Pattern::Exclude(r"re:foo".to_string())).is_err());
    }

    #[tokio::test]
    async fn test_matcher_implicit_include() -> anyhow::Result<()> {
        let config = b"
[exclude]
path:exc
";

        let prof = Root::from_bytes(config, "test".to_string()).unwrap();

        let matcher = prof.matcher(|_| async { Ok(Some(vec![])) }).await?;

        // Show we got an implicit rule that includes everything.
        assert!(matcher.matches("a/b".try_into()?)?);

        // Sanity that exclude works.
        assert!(!matcher.matches("exc/foo".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_matcher_v1() -> anyhow::Result<()> {
        let base = b"
%include child

[exclude]
path:a/exc

[include]
path:a
";

        let child = b"
[exclude]
path:b/exc

[include]
path:b
";

        let prof = Root::from_bytes(base, "test".to_string())?;
        let matcher = prof.matcher(|_| async { Ok(Some(child.to_vec())) }).await?;

        // Exclude rule "wins" for v1 despite order in confing.
        assert!(!matcher.matches("a/exc".try_into()?)?);
        assert!(!matcher.matches("b/exc".try_into()?)?);
        assert!(matcher.matches("a/inc".try_into()?)?);
        assert!(matcher.matches("b/inc".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_matcher_v2() -> anyhow::Result<()> {
        let base = b"
%include child_1
%include child_2

[exclude]
path:a/exc
path:c

[include]
path:a
";

        let child_1 = b"
[include]
path:c

[metadata]
version = 2
";

        let child_2 = b"
[exclude]
path:b/exc
path:c

[include]
path:b

[metadata]
version = 2
";

        let prof = Root::from_bytes(base, "test".to_string())?;
        let matcher = prof
            .matcher(|path| async move {
                match path.as_ref() {
                    "child_1" => Ok(Some(child_1.to_vec())),
                    "child_2" => Ok(Some(child_2.to_vec())),
                    _ => unreachable!(),
                }
            })
            .await?;

        // Rules directly in root profile still get excludes-go-last ordering.
        assert!(!matcher.matches("a/exc".try_into()?)?);
        assert!(matcher.matches("a/inc".try_into()?)?);

        // Order for v2 child profile is maintained - include rule wins.
        assert!(matcher.matches("b/exc".try_into()?)?);
        assert!(matcher.matches("b/inc".try_into()?)?);

        // "c" is included due to unioning of v2 profiles.
        assert!(matcher.matches("c".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_matcher_missing_include() -> anyhow::Result<()> {
        let config = b"
%include banana
foo
";

        let prof = Root::from_bytes(config, "test".to_string()).unwrap();

        let matcher = prof.matcher(|_| async { Ok(None) }).await?;

        // We ignore missing includes so that things don't completely
        // break if someone accidentally deletes an in-use sparse
        // profile.
        assert!(matcher.matches("foo".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_matcher_unsupported_patterns() -> anyhow::Result<()> {
        let config = b"
re:.*
listfile0:/tmp/oops
foo
";

        let prof = Root::from_bytes(config, "test".to_string()).unwrap();

        // Can still get a matcher, skipping unsupported patterns.
        let matcher = prof.matcher(|_| async { Ok(None) }).await?;

        assert!(matcher.matches("foo".try_into()?)?);
        assert!(!matcher.matches("bar".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_regex_patterns() -> anyhow::Result<()> {
        let config = b"
[metadata]
version = 2

[include]
re:foo/\\..*
re:^bar/(?!bad/)
re:^bar/bad/(?:.*/)?IMPORTANT.ext(?:/|$)
";

        let prof = Root::from_bytes(b"%include foo", "test".to_string()).unwrap();
        let matcher = prof
            .matcher(|_| async { Ok(Some(config.to_vec())) })
            .await?;

        assert!(matcher.matches("foo/.blah".try_into()?)?);
        assert!(!matcher.matches("foo/not-dot".try_into()?)?);

        assert!(matcher.matches("bar/ok".try_into()?)?);
        assert!(!matcher.matches("bar/bad/nono".try_into()?)?);
        assert!(matcher.matches("bar/bad/well/jk/IMPORTANT.ext".try_into()?)?);

        Ok(())
    }

    #[tokio::test]
    async fn test_explain_empty() {
        let prof = Root::from_bytes(b"", "test".to_string()).unwrap();
        let matcher = prof
            .matcher(|_| async move { Ok(Some(vec![])) })
            .await
            .unwrap();

        assert_eq!(
            matcher.explain("a/b".try_into().unwrap()).unwrap(),
            (true, "implicit match due to empty profile".to_string())
        );
    }

    #[tokio::test]
    async fn test_explain_no_match() {
        let prof = Root::from_bytes(b"a", "test".to_string()).unwrap();
        let matcher = prof
            .matcher(|_| async move { Ok(Some(vec![])) })
            .await
            .unwrap();

        assert_eq!(
            matcher.explain("b".try_into().unwrap()).unwrap(),
            (false, "no rules matched".to_string())
        );
    }

    #[tokio::test]
    async fn test_explain_chain() {
        let base = b"%include child_1";
        let child_1 = b"%include child_2";
        let child_2 = b"
[include]
glob:{a,b,c}

[exclude]
path:d
";

        let prof = Root::from_bytes(base, "base".to_string()).unwrap();
        let matcher = prof
            .matcher(|path| async move {
                match path.as_ref() {
                    "child_1" => Ok(Some(child_1.to_vec())),
                    "child_2" => Ok(Some(child_2.to_vec())),
                    _ => unreachable!(),
                }
            })
            .await
            .unwrap();

        assert_eq!(
            matcher.explain("b".try_into().unwrap()).unwrap(),
            (true, "base -> child_1 -> child_2".to_string())
        );

        assert_eq!(
            matcher.explain("d".try_into().unwrap()).unwrap(),
            (false, "base -> child_1 -> child_2".to_string())
        );
    }

    #[tokio::test]
    async fn test_dynamic_rule_source() {
        let config = b"
one

# source = banana
two
three

# source =
four
";

        let prof = Root::from_bytes(config, "base".to_string()).unwrap();

        let matcher = prof.matcher(|_| async { Ok(Some(vec![])) }).await.unwrap();

        assert_eq!(
            matcher.explain("one".try_into().unwrap()).unwrap(),
            (true, "base".to_string())
        );

        assert_eq!(
            matcher.explain("two".try_into().unwrap()).unwrap(),
            (true, "base (banana)".to_string())
        );

        assert_eq!(
            matcher.explain("three".try_into().unwrap()).unwrap(),
            (true, "base (banana)".to_string())
        );

        assert_eq!(
            matcher.explain("four".try_into().unwrap()).unwrap(),
            (true, "base".to_string())
        );
    }
}


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

use std::process::Command;
use std::time::SystemTime;

use spawn_ext::CommandExt;

fn main() {
    let exe_path = std::env::current_exe().unwrap();
    if std::env::args().len() == 1 {
        let pid = unsafe { libc::getpid() };
        println!("parent pid: {}", pid);
        let mut cmd = Command::new(exe_path);
        let clock = SystemTime::now();
        cmd.arg("child").avoid_inherit_handles().new_session();
        println!("avoid_inherit_handles took: {:?}", clock.elapsed().unwrap());
        let child_pid = cmd.arg("child").spawn_detached().unwrap().id();
        println!("child pid: {}", child_pid);
        println!("spawn took: {:?}", clock.elapsed().unwrap());
        println!("Both processes are sleeping.");
        println!();
        if cfg!(windows) {
            println!("On Windows, use Process Hacker to check handles.");
            println!("Inheritable handles are highlighted in cyan.");
        } else {
            println!(
                "On Linux, use 'ls -l /proc/{{{},{}}}/fd/' to check fds.",
                pid, child_pid
            );
        }
        println!();
        println!("The child should not have more file handles than the parent.");
    }
    std::thread::sleep(std::time::Duration::from_secs(300))
}


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

//! This crate extends the `process::Command` interface in stdlib.
//! - `avoid_inherit_handles` is similar to `close_fds=True` in Python.
//! - `new_session` uses `CREATE_NEW_PROCESS_GROUP` on Windows, and `setsid` on
//!   Unix.
//! - `spawn_detached` is a quicker way to spawn and forget.

use std::io;
use std::process::Child;
use std::process::Command;
use std::process::Stdio;

/// Extensions to `std::process::Command`.
pub trait CommandExt {
    /// Attempt to avoid inheriting file handles.
    /// Call this before setting up redirections!
    fn avoid_inherit_handles(&mut self) -> &mut Self;

    /// Use a new session for the new process.
    /// Call this after `avoid_inherit_handles`!
    fn new_session(&mut self) -> &mut Self;

    /// Spawn a process with stdio redirected to null and forget about it.
    /// Return the process id.
    fn spawn_detached(&mut self) -> io::Result<Child>;
}

impl CommandExt for Command {
    fn avoid_inherit_handles(&mut self) -> &mut Self {
        #[cfg(unix)]
        unix::avoid_inherit_handles(self);

        #[cfg(windows)]
        windows::avoid_inherit_handles(self);

        self
    }

    fn new_session(&mut self) -> &mut Self {
        #[cfg(unix)]
        unix::new_session(self);

        #[cfg(windows)]
        windows::new_session(self);

        self
    }

    fn spawn_detached(&mut self) -> io::Result<Child> {
        self.avoid_inherit_handles()
            .new_session()
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
    }
}

#[cfg(windows)]
mod windows {
    use std::os::windows::process::CommandExt;

    use winapi::um::handleapi::SetHandleInformation;
    use winapi::um::winbase::CREATE_NEW_PROCESS_GROUP;
    use winapi::um::winbase::CREATE_NO_WINDOW;
    use winapi::um::winbase::HANDLE_FLAG_INHERIT;

    use super::*;

    // A larger value like 8192 adds visible overhead (ex. >5ms).
    // At first we had 2048, however it wasn't enough in some cases
    // https://fburl.com/px16lb62. We bumped it to 4096.
    const MAX_HANDLE: usize = 4096;

    pub fn avoid_inherit_handles(command: &mut Command) {
        // Attempt to mark handles as "not inheritable".
        //
        // Practically only a few handles are accidentally inheritable.
        // Use Process Hacker to examine handles [1]. Inheritable handles
        // are highlighted in cyan background. Example:
        //
        //  File, \Device\ConDrv, 0x58
        //  File, \Device\ConDrv, 0x5c
        //  File, \Device\ConDrv, 0x60
        //  File, \Device\Afd, 0x348
        //
        // [1]: https://github.com/processhacker/processhacker/
        for handle in 1..=MAX_HANDLE {
            // According to https://devblogs.microsoft.com/oldnewthing/20050121-00/?p=36633
            // kernel handles are always a multiple of 4
            let handle = unsafe { std::mem::transmute(handle * 4) };
            unsafe { SetHandleInformation(handle, HANDLE_FLAG_INHERIT, 0) };
        }
        // A cleaner way might be setting bInheritHandles to FALSE at
        // CreateProcessW time. However the Rust stdlib does not expose an
        // interface to set bInheritHandles, and bInheritHandles=FALSE
        // could break file redirections (with possible solutions [2] [3]).
        //
        // [2]: https://github.com/python/cpython/commit/b2a6083eb0384f38839d
        // [3]: https://devblogs.microsoft.com/oldnewthing/20111216-00/?p=8873

        // CREATE_NO_WINDOW forbids allocating stdio handles.
        command.creation_flags(CREATE_NO_WINDOW);
    }

    pub fn new_session(command: &mut Command) {
        command.creation_flags(CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP);
    }
}

#[cfg(unix)]
mod unix {
    use std::os::unix::process::CommandExt;

    use super::*;

    // Linux by default has max fd limited to 1024.
    // 2048 is practically more than enough with about 2ms overhead.
    const MAXFD: i32 = 2048;

    pub fn avoid_inherit_handles(command: &mut Command) {
        // There are some constraints for this function.
        // See std::os::unix::process::CommandExt::pre_exec.
        // Namely, do not allocate.
        unsafe { command.pre_exec(pre_exec_close_fds) };
    }

    pub fn new_session(command: &mut Command) {
        unsafe { command.pre_exec(pre_exec_setsid) };
    }

    fn pre_exec_close_fds() -> io::Result<()> {
        // Set FD_CLOEXEC on files.
        // Note: using `close` might break error reporting if exec fails.
        for fd in 3..=MAXFD {
            unsafe { libc::fcntl(fd, libc::F_SETFD, libc::FD_CLOEXEC) };
        }
        Ok(())
    }

    fn pre_exec_setsid() -> io::Result<()> {
        // Create a new session.
        unsafe { libc::setsid() };
        Ok(())
    }
}

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

    // It's hard to test the real effects. Here we just check command still runs.
    // Use `cargo run --example spawn` to manually check the close_fds behavior.
    #[test]
    fn smoke_test_command_still_runs() {
        let dir = tempfile::tempdir().unwrap();

        let args = if cfg!(unix) {
            vec!["/bin/sh", "-c", "echo foo > a"]
        } else {
            vec!["cmd.exe", "/c", "echo foo > a"]
        };
        let mut command = if cfg!(unix) {
            Command::new(&args[0])
        } else {
            Command::new(&args[0])
        };
        let mut child = command
            .args(&args[1..])
            .current_dir(&dir.path())
            .spawn_detached()
            .unwrap();
        child.wait().unwrap();

        assert_eq!(&std::fs::read(dir.path().join("a")).unwrap()[..3], b"foo")
    }
}


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

use std::collections::HashMap;
use std::fmt;
use std::path::Path;
use std::path::PathBuf;

use types::RepoPath;
use types::RepoPathBuf;

#[derive(Default, Clone, PartialEq, Eq)]
pub struct Status {
    all: HashMap<RepoPathBuf, FileStatus>,

    // Non-utf8 or otherwise invalid paths. Up to the caller how to handle these, if at all.
    invalid_path: Vec<Vec<u8>>,

    // Invalid/unsupported file types.
    invalid_type: Vec<RepoPathBuf>,
}

pub struct StatusBuilder(Status);

impl StatusBuilder {
    pub fn new() -> Self {
        Self(Status::default())
    }

    pub fn build(self) -> Status {
        self.0
    }

    pub fn modified(mut self, modified: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, modified, FileStatus::Modified);
        self
    }

    pub fn added(mut self, added: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, added, FileStatus::Added);
        self
    }

    pub fn removed(mut self, removed: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, removed, FileStatus::Removed);
        self
    }

    pub fn deleted(mut self, deleted: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, deleted, FileStatus::Deleted);
        self
    }

    pub fn unknown(mut self, unknown: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, unknown, FileStatus::Unknown);
        self
    }

    pub fn ignored(mut self, ignored: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, ignored, FileStatus::Ignored);
        self
    }

    pub fn clean(mut self, clean: Vec<RepoPathBuf>) -> Self {
        Self::index(&mut self.0.all, clean, FileStatus::Clean);
        self
    }

    pub fn invalid_path(mut self, invalid: Vec<Vec<u8>>) -> Self {
        self.0.invalid_path = invalid;
        self
    }

    pub fn invalid_type(mut self, invalid: Vec<RepoPathBuf>) -> Self {
        self.0.invalid_type = invalid;
        self
    }

    pub fn iter(&self) -> impl Iterator<Item = (&RepoPath, FileStatus)> {
        self.0.iter()
    }

    // This fn has to take 'deconstructed' self, because you can't borrow &mut self and &self.xxx at the same time
    fn index(
        all: &mut HashMap<RepoPathBuf, FileStatus>,
        files: Vec<RepoPathBuf>,
        status: FileStatus,
    ) {
        for file in files {
            all.insert(file, status);
        }
    }
}

impl Status {
    // modified() and other functions intentionally return Iterator<> and not &Vec
    // Those functions can be used if someone needs a list of files in certain category
    // If someone need to check what is the status of a file, they should use status(file) because it handles case sensitivity properly
    pub fn modified(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Modified)
    }

    pub fn added(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Added)
    }

    pub fn removed(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Removed)
    }

    pub fn deleted(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Deleted)
    }

    pub fn unknown(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Unknown)
    }

    pub fn ignored(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Ignored)
    }

    pub fn clean(&self) -> impl Iterator<Item = &RepoPathBuf> {
        self.filter_status(FileStatus::Clean)
    }

    pub fn invalid_path(&self) -> &[Vec<u8>] {
        &self.invalid_path
    }

    pub fn invalid_type(&self) -> &[RepoPathBuf] {
        &self.invalid_type
    }

    pub fn status(&self, file: &RepoPath) -> Option<FileStatus> {
        self.all.get(file).copied()
    }

    fn filter_status(&self, status: FileStatus) -> impl Iterator<Item = &RepoPathBuf> {
        self.all
            .iter()
            .filter_map(move |(f, s)| if *s == status { Some(f) } else { None })
    }

    pub fn iter(&self) -> impl Iterator<Item = (&RepoPath, FileStatus)> {
        self.all.iter().map(|(f, s)| (f.as_repo_path(), *s))
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FileStatus {
    /// The file has been modified.
    Modified,
    /// The file has been added to the tree.
    Added,
    /// The file has been removed to the tree.
    Removed,
    /// The file is in the tree, but it isn't in the working copy (it's "missing").
    Deleted,
    /// The file isn't in the tree, but it exists in the working copy and it isn't ignored.
    Unknown,
    /// The file isn't in the tree and it exists in the working copy, but it is ignored.
    Ignored,
    /// The file has not been modified.
    Clean,
}

impl FileStatus {
    pub fn py_letter(&self) -> &'static str {
        match self {
            FileStatus::Modified => "M",
            FileStatus::Added => "A",
            FileStatus::Removed => "R",
            FileStatus::Deleted => "!",
            FileStatus::Unknown => "?",
            FileStatus::Ignored => "I",
            FileStatus::Clean => "C",
        }
    }
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (file, status) in self.all.iter() {
            write!(f, "{} {}", status.py_letter(), file)?;
        }
        Ok(())
    }
}

pub fn needs_morestatus_extension(hg_dir: &Path, parent_count: usize) -> bool {
    if parent_count > 1 {
        return true;
    }

    for path in [
        PathBuf::from("bisect.state"),
        PathBuf::from("graftstate"),
        PathBuf::from("histedit-state"),
        PathBuf::from("merge/state2"),
        PathBuf::from("rebasestate"),
        PathBuf::from("unshelverebasestate"),
        PathBuf::from("updatestate"),
    ] {
        if hg_dir.join(path).is_file() {
            return true;
        }
    }

    false
}


/*
 * 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::str;
use std::str::FromStr;

use anyhow::bail;
use anyhow::Result;
use minibytes::Bytes;
use types::HgId;
use types::Key;
use types::RepoPath;

/// Mercurial may embed the copy-from information into the blob itself, in which case, the `Delta`
/// would look like:
///
///   \1
///   copy: path
///   copyrev: sha1
///   \1
///   blob
///
/// If the blob starts with \1\n too, it's escaped by adding \1\n\1\n at the beginning.
pub fn strip_metadata(data: &Bytes) -> Result<(Bytes, Option<Key>)> {
    let (blob, copy_from) = separate_metadata(data)?;
    if copy_from.len() > 0 {
        let slice = copy_from.as_ref();
        let slice = &slice[2..copy_from.len() - 2];
        let mut path = None;
        let mut hgid = None;

        for line in slice.split(|c| c == &b'\n') {
            if line.is_empty() {
                continue;
            }
            if line.starts_with(b"copy: ") {
                path = Some(RepoPath::from_str(str::from_utf8(&line[6..])?)?.to_owned());
            } else if line.starts_with(b"copyrev: ") {
                hgid = Some(HgId::from_str(str::from_utf8(&line[9..])?)?);
            } else {
                bail!("Unknown metadata in data: {:?}", line);
            }
        }

        let key = match (path, hgid) {
            (None, Some(_)) => bail!("missing 'copyrev' metadata"),
            (Some(_), None) => bail!("missing 'copy' metadata"),

            (None, None) => None,
            (Some(path), Some(hgid)) => Some(Key::new(path, hgid)),
        };

        Ok((blob, key))
    } else {
        Ok((blob, None))
    }
}

pub fn separate_metadata(data: &Bytes) -> Result<(Bytes, Bytes)> {
    let slice = data.as_ref();
    if !slice.starts_with(b"\x01\n") {
        return Ok((data.clone(), Bytes::new()));
    }
    let slice = &slice[2..];
    if let Some(pos) = slice.windows(2).position(|needle| needle == b"\x01\n") {
        let split_pos = 2 + pos + 2;
        Ok((data.slice(split_pos..), data.slice(..split_pos)))
    } else {
        Ok((data.clone(), Bytes::new()))
    }
}

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

    use super::*;

    #[test]
    fn test_strip_separate_metadata() -> Result<()> {
        let key = key("foo/bar/baz", "1234");
        let data = Bytes::copy_from_slice(
            format!(
                "\x01\ncopy: {}\ncopyrev: {}\n\x01\nthis is a blob",
                key.path, key.hgid
            )
            .as_bytes(),
        );
        let (split_data, path) = strip_metadata(&data)?;
        assert_eq!(split_data, Bytes::from(&b"this is a blob"[..]));
        assert_eq!(path, Some(key.clone()));

        let (blob, copy_from) = separate_metadata(&data)?;
        assert_eq!(blob, Bytes::from(&b"this is a blob"[..]));
        assert_eq!(
            copy_from,
            Bytes::copy_from_slice(
                format!("\x01\ncopy: {}\ncopyrev: {}\n\x01\n", key.path, key.hgid).as_bytes(),
            )
        );

        let data = Bytes::from(&b"\x01\n\x01\nthis is a blob"[..]);
        let (split_data, path) = strip_metadata(&data)?;
        assert_eq!(split_data, Bytes::from(&b"this is a blob"[..]));
        assert_eq!(path, None);

        let (blob, copy_from) = separate_metadata(&data)?;
        assert_eq!(blob, Bytes::from(&b"this is a blob"[..]));
        assert_eq!(copy_from, &b"\x01\n\x01\n"[..]);

        let data = Bytes::from(&b"\x01\nthis is a blob"[..]);
        let (split_data, path) = strip_metadata(&data)?;
        assert_eq!(split_data, data);
        assert_eq!(path, None);

        let (blob, copy_from) = separate_metadata(&data)?;
        assert_eq!(blob, data);
        assert_eq!(copy_from, Bytes::new());

        Ok(())
    }
}


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

//! Utilities interacting with hg store.

mod filestore_util;

pub use filestore_util::separate_metadata;
pub use filestore_util::strip_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.
 */

//! # storemodel
//!
//! Provides trait definitions for storage needs.
//! Useful to de-couple from heavyweight storage implementations.
//!
//! Traits defined by this crate are ideally tech-debt free and
//! consider both hg and git use-cases. This mainly means:
//! - APIs do expose hg details like filelog "copy from" or LFS pointer.
//! - History-related APIs should avoid linkrev or linknode, which do not exist
//!   in git.
//!
//! For flexibility, different features might be defined in different traits.
//! Traits can be combined later. For example, reading file content, metadata,
//! and history should probably be 3 different traits.

use async_trait::async_trait;
pub use futures;
use futures::stream::BoxStream;
pub use minibytes;
pub use types;
use types::HgId;
use types::Key;
use types::RepoPath;

#[async_trait]
#[auto_impl::auto_impl(Arc)]
pub trait ReadFileContents {
    type Error;

    /// Read the content of specified files.
    ///
    /// The returned content should be just the file contents. This means:
    /// - The returned content does not contain the "copy from" header.
    /// - The returned content does not contain raw LFS content. LFS pointer
    ///   is resolved transparently.
    /// - If the file content is redacted, it's an error instead of a placeholder
    ///   of dummy data.
    async fn read_file_contents(
        &self,
        keys: Vec<Key>,
    ) -> BoxStream<Result<(minibytes::Bytes, Key), Self::Error>>;

    /// Read rename metadata of sepcified files.
    ///
    /// The result is a vector of (key, Option<rename_from_key>) pairs for success case.
    async fn read_rename_metadata(
        &self,
        keys: Vec<Key>,
    ) -> BoxStream<Result<(Key, Option<Key>), Self::Error>>;
}

pub trait RefreshableReadFileContents: ReadFileContents {
    fn refresh(&self) -> Result<(), Self::Error>;
}

#[async_trait]
pub trait ReadRootTreeIds {
    /// Read root tree nodes of given commits.
    /// Return `(commit_id, tree_id)` list.
    async fn read_root_tree_ids(&self, commits: Vec<HgId>) -> anyhow::Result<Vec<(HgId, HgId)>>;
}

/// The `TreeStore` is an abstraction layer for the tree manifest that decouples how or where the
/// data is stored. This allows more easy iteration on serialization format. It also simplifies
/// writing storage migration.
#[auto_impl::auto_impl(Arc)]
pub trait TreeStore {
    /// Read the contents of a directory.
    ///
    /// The result is opaque bytes data, encoded using the format specified by `format()`.
    /// To parse the bytes consider `manifest_tree::TreeEntry`.
    fn get(&self, path: &RepoPath, hgid: HgId) -> anyhow::Result<minibytes::Bytes>;

    fn insert(&self, path: &RepoPath, hgid: HgId, data: minibytes::Bytes) -> anyhow::Result<()>;

    /// Indicate to the store that we will be attempting to access the given
    /// tree nodes soon. Some stores (especially ones that may perform network
    /// I/O) may use this information to prepare for these accesses (e.g., by
    /// by prefetching the nodes in bulk). For some stores this operation does
    /// not make sense, so the default implementation is a no-op.
    fn prefetch(&self, _keys: Vec<Key>) -> anyhow::Result<()> {
        Ok(())
    }

    /// Decides whether the store uses git or hg format.
    fn format(&self) -> TreeFormat {
        TreeFormat::Hg
    }
}

pub trait RefreshableTreeStore: TreeStore {
    fn refresh(&self) -> anyhow::Result<()>;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub enum TreeFormat {
    // NAME '\0' HEX_SHA1 MODE '\n'
    // MODE: 't' (tree), 'l' (symlink), 'x' (executable)
    Hg,

    // MODE ' ' NAME '\0' BIN_SHA1
    // MODE: '40000' (tree), '100644' (regular), '100755' (executable),
    //       '120000' (symlink), '160000' (gitlink)
    Git,
}


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

use std::collections::HashMap;
use std::collections::VecDeque;
use std::fmt::Debug;
use std::hash::Hash;
use std::pin::Pin;

use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::task::Context;
use futures::task::Poll;
use futures::Stream;

/// Resolve a stream of `I`s (inputs) into a stream of `O`s (outputs).
///
/// The resolution uses a "resolver" that can resolve `I`s in both locally and
/// remotely. `I` is attempted to be resolved locally, if that fails, a batch
/// of `I`s will be sent to remote to resolve.
pub struct HybridStream<I, O, E>(BoxStream<'static, Result<(I, O), E>>);

/// Internal state of the stream.
struct HybridStreamState<I, O, E> {
    /// Remaining inputs to resolve.
    input: BoxStream<'static, Result<I, E>>,

    /// Buffered inputs to resolve (in order).
    /// The 1st item should be unknown locally.
    buffer: VecDeque<ResolveState<I, O>>,

    /// Buffer size before sending a request.
    buffer_size: usize,

    /// Pending remote request. The stream populates `response`.
    request: Option<BoxStream<'static, Result<(I, O), E>>>,

    /// Result from consumed `request` stream.
    response: HashMap<I, O>,

    /// Defines how to resolve I to O.
    resolver: Box<dyn HybridResolver<I, O, E> + Send + Sync + 'static>,

    /// Retry attempted.
    attempt: usize,

    /// Maximum retry count.
    retry_limit: usize,
}

/// Defines how to resolve input to output using local data and remote data.
/// The output stream preserves the order of the input stream.
#[async_trait]
pub trait HybridResolver<I, O, E> {
    /// Attempt to resolve I using local data.
    fn resolve_local(&mut self, input: &I) -> Result<Option<O>, E>;

    /// Resolve I using remote data in batch. The output stream can be out-of-order.
    async fn resolve_remote(&self, input: &[I])
    -> Result<BoxStream<'static, Result<(I, O), E>>, E>;

    /// Raise an error if the server did not response to input multiple times.
    fn retry_error(&self, attempt: usize, input: &[I]) -> E;
}

#[derive(Debug)]
enum ResolveState<I, O> {
    Resolved(I, O),
    NotResolved(I),
}

impl<I, O, E> HybridStream<I, O, E>
where
    I: Eq + Hash + Clone + Send + Sync + Debug + 'static,
    O: Send + Sync + Debug + 'static,
    E: 'static,
{
    /// Create a new `HybridStream`.
    pub fn new(
        stream: BoxStream<'static, Result<I, E>>,
        resolver: impl HybridResolver<I, O, E> + Send + Sync + 'static,
        buffer_size: usize,
        retry_limit: usize,
    ) -> Self {
        let state = HybridStreamState {
            input: stream,
            buffer: Default::default(),
            response: Default::default(),
            buffer_size: buffer_size.max(1),
            request: Default::default(),
            resolver: Box::new(resolver),
            attempt: 0,
            retry_limit,
        };
        let stream = futures::stream::unfold(state, |mut state| async {
            let item = state.next_item().await;
            item.transpose().map(|i| (i, state))
        });
        Self(Box::pin(stream.fuse()))
    }
}

impl<I, O, E> Stream for HybridStream<I, O, E>
where
    I: Eq + Hash + Clone + Send + Sync + Debug + 'static,
    O: Send + Sync + Debug + 'static,
    E: 'static,
{
    type Item = Result<(I, O), E>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        Stream::poll_next(self.0.as_mut(), cx)
    }
}

impl<I, O, E> HybridStreamState<I, O, E>
where
    I: Eq + Hash + Clone + Debug,
    O: Debug,
{
    /// A future to produce one `next` item.
    async fn next_item(&mut self) -> Result<Option<(I, O)>, E> {
        loop {
            let item = self.buffer.pop_front();
            match item {
                None => {
                    let count = self.fill_buffer().await?;
                    if count == 0 {
                        // Empty buffer. Input stream has ended.
                        break Ok(None);
                    }
                }
                Some(ResolveState::Resolved(i, o)) => break Ok(Some((i, o))),
                Some(ResolveState::NotResolved(i)) => {
                    if let Some(o) = self.response.remove(&i) {
                        break Ok(Some((i, o)));
                    } else {
                        self.buffer.push_front(ResolveState::NotResolved(i));
                        self.poll_remote().await?;
                    }
                }
            }
        }
    }

    /// Prepare a buffer of inputs. Part of them resolved locally. The remaining
    /// are to be resolved remotely.
    ///
    /// Returns number of items pushed to the buffer.
    ///
    /// Consumes items from `self.input`. Updates `self.buffer`.
    async fn fill_buffer(&mut self) -> Result<usize, E> {
        let mut count = 0;
        while self.buffer.len() < self.buffer_size {
            let next_input = self.input.next().await.transpose()?;
            match next_input {
                Some(input) => {
                    // Attempt to resolve it locally.
                    let state = match self.resolver.resolve_local(&input)? {
                        Some(output) => ResolveState::Resolved(input, output),
                        None => ResolveState::NotResolved(input),
                    };
                    self.buffer.push_back(state);
                    // Reset attempt counter.
                    self.attempt = 0;
                    count += 1;
                }
                // Reached the end.
                None => break,
            }
        }
        Ok(count)
    }

    /// Make progress related to the remote request.
    ///
    /// If there is no pending request, send a new one if necessary.
    /// If there is an existing request, read from it (`self.request`),
    /// and updates `self.response`.
    async fn poll_remote(&mut self) -> Result<(), E> {
        // Send a batch request if any input is unresolved and there is no
        // pending request.
        match self.request {
            None => {
                let batch: Vec<I> = self.remote_input();
                if !batch.is_empty() {
                    if self.attempt > self.retry_limit {
                        return Err(self.resolver.retry_error(self.attempt, &batch));
                    }
                    self.attempt += 1;
                    let request = self.resolver.resolve_remote(&batch).await?;
                    self.request = Some(request);
                }
            }
            Some(ref mut stream) => match stream.next().await {
                None => self.request = None,
                Some(Ok((i, o))) => {
                    self.response.insert(i, o);
                }
                Some(Err(e)) => {
                    self.request = None;
                    return Err(e);
                }
            },
        }
        Ok(())
    }

    /// Input for a remote request.
    fn remote_input(&self) -> Vec<I> {
        self.buffer
            .iter()
            .filter_map(|i| match i {
                ResolveState::NotResolved(i) => {
                    if self.response.contains_key(&i) {
                        // The item was fetched previously.
                        // This can happen when retry requests were sent.
                        None
                    } else {
                        Some(i)
                    }
                }
                ResolveState::Resolved(_, _) => None,
            })
            .cloned()
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::Mutex;

    use futures::stream;
    use futures::stream::StreamExt;
    use tokio::time::sleep;
    use tokio::time::Duration;

    use super::*;

    type I = usize;
    type O = String;
    type E = std::io::Error;

    #[derive(Default)]
    struct Resolver {
        cached: Arc<Mutex<HashMap<usize, String>>>,
    }

    #[async_trait]
    impl HybridResolver<I, O, E> for Resolver {
        fn resolve_local(&mut self, input: &I) -> Result<Option<O>, E> {
            Ok(self.cached.lock().unwrap().get(input).cloned())
        }

        async fn resolve_remote(
            &self,
            input: &[I],
        ) -> Result<BoxStream<'static, Result<(I, O), E>>, E> {
            let cached = self.cached.clone();
            // Exercise ".await" in this function.
            sleep(Duration::from_millis(1)).await;
            // Return nothing for 404.
            let output_iter = input
                .to_vec()
                .into_iter()
                .filter(|&i| i != 404)
                .map(move |i| {
                    let o = i.to_string();
                    cached.lock().unwrap().insert(i, o.clone());
                    // Return an error for 500.
                    if i == 500 {
                        Err(error("cannot resolve 500"))
                    } else {
                        Ok((i, o))
                    }
                });
            Ok(Box::pin(stream::iter(output_iter)))
        }

        fn retry_error(&self, attempt: usize, input: &[I]) -> E {
            error(format!(
                "give up after {} attempts for input {:?}",
                attempt, input
            ))
        }
    }

    fn error(msg: impl ToString) -> std::io::Error {
        std::io::Error::new(std::io::ErrorKind::Other, msg.to_string())
    }

    fn u(v: Option<Result<(I, O), E>>) -> (I, O) {
        v.unwrap().unwrap()
    }

    #[tokio::test]
    async fn test_hybrid_stream() {
        for buffer_size in [0, 1, 2, 10] {
            let input = stream::iter(vec![0, 1, 3, 5, 10].into_iter().map(Ok));
            let resolver = Resolver::default();
            resolver.cached.lock().unwrap().insert(1, "one".to_string());
            let mut stream = HybridStream::new(Box::pin(input), resolver, buffer_size, 0);
            assert_eq!(u(stream.next().await), (0, "0".to_string()));
            assert_eq!(u(stream.next().await), (1, "one".to_string()));
            assert_eq!(u(stream.next().await), (3, "3".to_string()));
            assert_eq!(u(stream.next().await), (5, "5".to_string()));
            assert_eq!(u(stream.next().await), (10, "10".to_string()));
            assert!(stream.next().await.is_none());
            assert!(stream.next().await.is_none());
        }
    }

    #[tokio::test]
    async fn test_hybrid_stream_retry() {
        for retry_limit in [0, 3] {
            let input = stream::iter(vec![0, 404, 1, 2].into_iter().map(Ok));
            let resolver = Resolver::default();
            let mut stream = HybridStream::new(Box::pin(input), resolver, 3, retry_limit);
            assert_eq!(u(stream.next().await), (0, "0".to_string()));
            assert_eq!(
                stream.next().await.unwrap().unwrap_err().to_string(),
                format!("give up after {} attempts for input [404]", retry_limit + 1)
            );
        }
    }

    #[tokio::test]
    async fn test_hybrid_stream_error() {
        let input = stream::iter(vec![0, 500, 1, 2].into_iter().map(Ok));
        let resolver = Resolver::default();
        let mut stream = HybridStream::new(Box::pin(input), resolver, 3, 1);
        assert_eq!(u(stream.next().await), (0, "0".to_string()));
        assert_eq!(
            stream.next().await.unwrap().unwrap_err().to_string(),
            "cannot resolve 500",
        );
    }
}


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

//! `streams` provides some generic streams that can be useful in other places.
//! - `HybridStream` provides a way to use local data (with a single point get
//!   API) and remote data (with an unordered batch get API) to resolve a stream
//!   of input into a stream of output.
//! - `SelectDrop` provides a version of `futures::stream::Select` which drops
//!   each of the combined streams after it terminates. This is useful for
//!   preventing deadlocks when one stream is waiting on another to be dropped
//!   to complete.

mod hybrid;
mod select_drop;

pub use hybrid::HybridResolver;
pub use hybrid::HybridStream;
pub use select_drop::select_drop;
pub use select_drop::SelectDrop;


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

//! Provides a version of `futures::stream::select` which drops each stream when it terminates.
//! This code is derived from `futures::stream::select`, with the streams wrapped in `Option`.

use std::pin::Pin;

use futures::stream::Fuse;
use futures::stream::FusedStream;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;

/// Stream for the [`select_drop()`] function.
#[pin_project]
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct SelectDrop<St1, St2> {
    #[pin]
    stream1: Option<Fuse<St1>>,
    #[pin]
    stream2: Option<Fuse<St2>>,
    poll_stream2_first: bool,
}

/// This function will attempt to pull items from both streams. Each
/// stream will be polled in a round-robin fashion, and whenever a stream is
/// ready to yield an item that item is yielded.
///
/// After one of the two input stream completes, the remaining one will be
/// polled exclusively. The returned stream completes when both input
/// streams have completed.
///
/// Note that this function consumes both streams and returns a wrapped
/// version of them.
///
/// This function is identical to `futures::stream::select`, except that
/// it drops each stream once it terminates.
pub fn select_drop<St1, St2>(stream1: St1, stream2: St2) -> SelectDrop<St1, St2>
where
    St1: Stream,
    St2: Stream<Item = St1::Item>,
{
    SelectDrop {
        stream1: Some(stream1.fuse()),
        stream2: Some(stream2.fuse()),
        poll_stream2_first: false,
    }
}

impl<St1, St2> SelectDrop<St1, St2> {
    /// Acquires a reference to the underlying streams that this combinator is
    /// pulling from.
    pub fn get_ref(&self) -> (Option<&St1>, Option<&St2>) {
        (
            self.stream1.as_ref().map(|s| s.get_ref()),
            self.stream2.as_ref().map(|s| s.get_ref()),
        )
    }

    /// Acquires a mutable reference to the underlying streams that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_mut(&mut self) -> (Option<&mut St1>, Option<&mut St2>) {
        (
            self.stream1.as_mut().map(|s| s.get_mut()),
            self.stream2.as_mut().map(|s| s.get_mut()),
        )
    }

    /// Acquires a pinned mutable reference to the underlying streams that this
    /// combinator is pulling from.
    ///
    /// Note that care must be taken to avoid tampering with the state of the
    /// stream which may otherwise confuse this combinator.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> (Option<Pin<&mut St1>>, Option<Pin<&mut St2>>) {
        let this = self.project();
        (
            this.stream1.as_pin_mut().map(|s| s.get_pin_mut()),
            this.stream2.as_pin_mut().map(|s| s.get_pin_mut()),
        )
    }

    /// Consumes this combinator, returning the underlying streams.
    ///
    /// Note that this may discard intermediate state of this combinator, so
    /// care should be taken to avoid losing resources when this is called.
    pub fn into_inner(self) -> (Option<St1>, Option<St2>) {
        (
            self.stream1.map(|s| s.into_inner()),
            self.stream2.map(|s| s.into_inner()),
        )
    }
}

impl<St1, St2> FusedStream for SelectDrop<St1, St2>
where
    St1: Stream,
    St2: Stream<Item = St1::Item>,
{
    fn is_terminated(&self) -> bool {
        self.stream1.as_ref().map_or(true, |s| s.is_terminated())
            && self.stream2.as_ref().map_or(true, |s| s.is_terminated())
    }
}

impl<St1, St2> Stream for SelectDrop<St1, St2>
where
    St1: Stream,
    St2: Stream<Item = St1::Item>,
{
    type Item = St1::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St1::Item>> {
        let this = self.project();
        if !*this.poll_stream2_first {
            poll_inner(this.poll_stream2_first, this.stream1, this.stream2, cx)
        } else {
            poll_inner(this.poll_stream2_first, this.stream2, this.stream1, cx)
        }
    }
}

fn poll_inner<St1, St2>(
    poll_stream2_first: &mut bool,
    mut a: Pin<&mut Option<St1>>,
    mut b: Pin<&mut Option<St2>>,
    cx: &mut Context<'_>,
) -> Poll<Option<St1::Item>>
where
    St1: Stream,
    St2: Stream<Item = St1::Item>,
{
    // `a.as_mut().as_pin_mut()` to go from `&mut Pin<&mut Option<St1>>`
    // to `Option<Pin<&mut St1>>` without moving out of `a`, so we can later
    // do `a.set(None)` to drop the completed stream.
    let a_done = match a.as_mut().as_pin_mut().map(|a| a.poll_next(cx)) {
        Some(Poll::Ready(Some(item))) => {
            *poll_stream2_first = !*poll_stream2_first;
            return Poll::Ready(Some(item));
        }
        Some(Poll::Pending) => false,
        Some(Poll::Ready(None)) | None => true,
    };

    if a_done {
        a.set(None);
    }

    let b_done = match b.as_mut().as_pin_mut().map(|b| b.poll_next(cx)) {
        Some(Poll::Ready(Some(item))) => {
            return Poll::Ready(Some(item));
        }
        Some(Poll::Pending) => false,
        Some(Poll::Ready(None)) | None => true,
    };

    if b_done {
        b.set(None);
    }

    if a_done && b_done {
        Poll::Ready(None)
    } else {
        Poll::Pending
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use futures::channel::mpsc::channel;
    use futures::stream;
    use futures::SinkExt;
    use futures::StreamExt;
    use tokio::time::timeout;
    use tokio::time::Duration;

    use super::*;

    #[tokio::test]
    async fn test_no_deadlock() -> Result<()> {
        let (sender, receiver) = channel(10);
        let producer = stream::iter(vec![1u8, 2, 3, 4, 5, 6, 7]).filter_map(move |num| {
            let mut sender = sender.clone();
            async move {
                if num == 4 {
                    match sender.send(num).await {
                        Ok(()) => None,
                        Err(_) => Some(0),
                    }
                } else {
                    Some(num)
                }
            }
        });
        let combined = select_drop(producer, receiver);
        let mut collected = timeout(Duration::from_secs(1), combined.collect::<Vec<u8>>())
            .await
            .expect(
                "stream timed out, select_drop didn't prevent deadlock or something went wrong",
            );
        collected.sort_unstable();
        assert_eq!(collected, vec![1u8, 2, 3, 4, 5, 6, 7]);
        Ok(())
    }
}