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

use std::collections::HashSet;

use anyhow::Result;
use regex_automata::dense::Builder;
use regex_automata::DenseDFA;
use regex_automata::DFA;
use regex_syntax::ast::parse;
use regex_syntax::ast::AssertionKind;
use regex_syntax::ast::Ast;
use types::RepoPath;

use crate::DirectoryMatch;
use crate::Matcher;

/// Pattern matcher constructed by an regular expression.
///
/// The regular expression syntax is same as regex crate with below limitations
/// due to the underlying RE lib (regex_automata):
///     1. Anchors \A, \z. RegexMatcher checks for a match only at the beginning
///        of the string by default, so '^' is not required.
///     2. Word boundary assertions such as \b and \B.
///     3. Lacks a few features like look around and backreferences (e.g. '?!').
///
/// The [RegexMatcher::match_prefix] method can be used to rule out
/// unnecessary directory visit early.
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub struct RegexMatcher {
    // Transformed regular expression pattern after replacing '$' (End-Of-Line) with '\0'.
    //
    // The underlying regex engine doesn't support '$' (EOL), we workaround this limitation
    // by replacing '$' with '\0' (which cannot occur in a path) in the pattern. Then have
    // `matches` API add a '\0' for testing when needed.
    pattern: String,

    // A table-based deterministic finite automaton (DFA) constructed by regular
    // expression.
    //
    // The RegexBuilder generally constructs two DFAs, where one is responsible for
    // finding the end of a match and the other is responsible for finding the
    // start of a match. Since we only need to detect whether something matched,
    // we use a dense::Builder to construct a single DFA here, which is cheaper
    // than building two DFAs.
    dfa: DenseDFA<Vec<usize>, usize>,
}

impl RegexMatcher {
    pub fn new(pattern: &str, case_sensitive: bool) -> Result<Self> {
        let pattern = handle_sol_eol(pattern)?;

        // The RE library doesn't support ^, we use this `Builder::anchored` to
        // make the search anchored at the beginning of the input. By default,
        // the regex will act as if the pattern started with a .*?, which enables
        // a match to appear anywhere.
        let dfa = Builder::new()
            .anchored(true)
            .case_insensitive(!case_sensitive)
            .build(&pattern)?;

        Ok(RegexMatcher { pattern, dfa })
    }

    /// Return `Some(bool)` if the end state of the DFA is 'match' or 'dead' state.
    /// Return `None` otherwise.
    pub fn match_prefix(&self, dir: &str) -> Option<bool> {
        // empty dir is for testing the "root" dir
        if dir.is_empty() {
            return None;
        }
        let bytes = dir.as_bytes();
        let mut state = self.dfa.start_state();

        for b in bytes {
            state = self.dfa.next_state(state, *b);
            if self.dfa.is_match_or_dead_state(state) {
                break;
            }
        }
        // Adding trailing '/' to handle cases like: if the pattern is `aa/bb`, then
        // it should return `Some(false)` for input "a"
        if !self.dfa.is_match_or_dead_state(state) && bytes.last() != Some(&b'/') {
            state = self.dfa.next_state(state, b'/');
        }

        if self.dfa.is_dead_state(state) {
            return Some(false);
        } else if self.dfa.is_match_state(state) {
            return Some(true);
        }
        None
    }

    /// Return if `path` matches with the matcher.
    pub fn matches(&self, path: &str) -> bool {
        let bytes = path.as_bytes();
        let mut state = self.dfa.start_state();

        // This handles two special cases:
        // 1. empty regex pattern (""), which describes the empty language. The empty language is
        //    a sub-language of every other language. So it will true without checking the input.
        // 2. it is possible to write a regex that is the opposite of the empty set, i.e., one that
        //    will not match anything. You could write it like so: [a&&b] -- intersection of a and b,
        //    which is empty. Currently though, such patterns won't compile in Rust Regex parser.
        if self.dfa.is_match_or_dead_state(state) {
            return self.dfa.is_match_state(state);
        }

        for &b in bytes {
            // We are using `next_state_unchecked` method here for speed, this follows
            // the implementation of DFA.is_match_at API:
            // https://github.com/BurntSushi/regex-automata/blob/0.1.10/src/dfa.rs#L220
            state = unsafe { self.dfa.next_state_unchecked(state, b) };
            if self.dfa.is_match_or_dead_state(state) {
                return self.dfa.is_match_state(state);
            }
        }

        // At this point, it means we are not in 'match' or 'dead' state, then
        // we add '\0' to check if the pattern is expecting EOL. Check the
        // comment of [RegexMatcher.pattern] for more details.
        state = unsafe { self.dfa.next_state_unchecked(state, b'\0') };
        self.dfa.is_match_state(state)
    }
}

/// Handle sol (start-of-line) and eol (end-of-line) since regex-automata doesn't support '^' and '$'.
///   1. sol ('^'), we just remove it, RegexMatcher will only match at the beginning of the string.
///   2. eol ('$'), we replace it with '\0'
fn handle_sol_eol(pattern: &str) -> Result<String> {
    fn traverse_ast(ast: &Ast, pattern: &mut String, sol_indices: &mut HashSet<usize>) {
        match ast {
            Ast::Group(group) => traverse_ast(&group.ast, pattern, sol_indices),
            Ast::Assertion(assertion) => {
                if assertion.kind == AssertionKind::EndLine {
                    let start = assertion.span.start.offset;
                    let end = assertion.span.end.offset;
                    assert_eq!(start + 1, end, "$ (end of line) should be 1 char");
                    pattern.replace_range(start..end, "\0");
                } else if assertion.kind == AssertionKind::StartLine {
                    let start = assertion.span.start.offset;
                    let end = assertion.span.end.offset;
                    assert_eq!(start + 1, end, "^ (start of line) should be 1 char");
                    sol_indices.insert(start);
                }
            }
            Ast::Repetition(repeat) => {
                traverse_ast(&repeat.ast, pattern, sol_indices);
            }
            Ast::Alternation(alternation) => {
                for t in &alternation.asts {
                    traverse_ast(t, pattern, sol_indices);
                }
            }
            Ast::Concat(concat) => {
                for t in &concat.asts {
                    traverse_ast(t, pattern, sol_indices);
                }
            }
            // Ast::Empty, Ast::Flags, Ast::Literal, Ast::Dot, Ast::Class
            _ => {}
        }
    }

    let ast = parse::Parser::new().parse(pattern)?;
    let mut new_pattern = pattern.to_string();
    let mut sol_indices = HashSet::new();

    traverse_ast(&ast, &mut new_pattern, &mut sol_indices);

    if sol_indices.is_empty() {
        Ok(new_pattern)
    } else {
        let bytes = new_pattern
            .as_bytes()
            .iter()
            .enumerate()
            .filter(|(i, _)| !sol_indices.contains(i))
            .map(|(_, e)| *e)
            .collect::<Vec<u8>>();
        Ok(String::from_utf8(bytes)?)
    }
}

impl Matcher for RegexMatcher {
    fn matches_directory(&self, path: &RepoPath) -> Result<DirectoryMatch> {
        let dm = match self.match_prefix(path.as_str()) {
            Some(true) => DirectoryMatch::Everything,
            Some(false) => DirectoryMatch::Nothing,
            None => DirectoryMatch::ShouldTraverse,
        };
        Ok(dm)
    }

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

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

    #[test]
    fn test_re_empty_pattern() {
        let m = RegexMatcher::new("", true).unwrap();
        assert!(m.matches(""));
        assert!(m.matches("a"));
        assert!(m.matches("abc"));
    }

    #[test]
    fn test_re_literal_path() {
        let m = RegexMatcher::new(r"(?:a/.*|b/.*)", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), Some(true));
        assert_eq!(m.match_prefix("a/"), Some(true));
        assert_eq!(m.match_prefix("b"), Some(true));
        assert_eq!(m.match_prefix("b/"), Some(true));
        assert_eq!(m.match_prefix("c"), Some(false));

        assert!(m.matches("a/c"));
        assert!(m.matches("a/b/c"));
        assert!(m.matches("b/c"));
        assert!(!m.matches("c/a"));
    }

    #[test]
    fn test_re_dir_boundary() {
        let m = RegexMatcher::new(r"aa/t\d+", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), Some(false));
        assert_eq!(m.match_prefix("aa"), None);
        assert_eq!(m.match_prefix("aa/t123"), Some(true));
        assert_eq!(m.match_prefix("aa/t123/b"), Some(true));
    }

    #[test]
    fn test_re_simple_pattern() {
        let m = RegexMatcher::new(r"a/t\d+/", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), None);
        assert_eq!(m.match_prefix("a/t123"), Some(true));
        assert_eq!(m.match_prefix("a/t123/"), Some(true));
        assert_eq!(m.match_prefix("a/test"), Some(false));

        assert!(!m.matches("b"));
        assert!(!m.matches("a/t"));
        assert!(!m.matches("a/tt"));
        assert!(!m.matches("a/c"));
        assert!(m.matches("a/t1/"));
        assert!(m.matches("a/t123/"));
    }

    #[test]
    fn test_re_without_eol() {
        let m = RegexMatcher::new(r"a/t\d+.py", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), None);
        assert_eq!(m.match_prefix("b"), Some(false));

        assert!(!m.matches("a/tt.py"));
        assert!(m.matches("a/t1.py"));
        assert!(m.matches("a/t1.pyc"));
    }

    #[test]
    fn test_re_with_eol() {
        let m = RegexMatcher::new(r"(:?a/t\d+.py$|a\d*.txt)", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), None);
        assert_eq!(m.match_prefix("b"), Some(false));

        assert!(m.matches("a/t123.py"));
        assert!(m.matches("a.txt"));
        assert!(m.matches("a1.txt"));
        assert!(m.matches("a1.txt1"));
        assert!(!m.matches("a/tt.py"));
        assert!(!m.matches("a/t123.pyc"));
    }

    #[test]
    fn test_re_with_sol() {
        let m = RegexMatcher::new(r"(:?^a/t.py$|^a\^.txt)", true).unwrap();
        assert_eq!(m.match_prefix(""), None);
        assert_eq!(m.match_prefix("a"), None);
        assert_eq!(m.match_prefix("b"), Some(false));

        assert!(m.matches("a/t.py"));
        assert!(!m.matches("ba/t.py"));
        assert!(!m.matches("a/t.pyc"));
        assert!(m.matches(r"a^.txt"));
        assert!(m.matches(r"a^.txt1"));
    }

    #[test]
    fn test_re_handle_sol_eol() {
        assert!(handle_sol_eol(r"*").is_err());
        assert_eq!(handle_sol_eol(r"a.py").unwrap(), "a.py");
        assert_eq!(handle_sol_eol(r"a.py\$").unwrap(), r"a.py\$");
        assert_eq!(handle_sol_eol(r"a.py$").unwrap(), "a.py\0");
        assert_eq!(handle_sol_eol(r"a.py$|a.txt").unwrap(), "a.py\0|a.txt");
        assert_eq!(handle_sol_eol(r"a.py$|a.txt$").unwrap(), "a.py\0|a.txt\0");
        assert_eq!(
            handle_sol_eol(r"(a.py$|a.txt$)|b.py").unwrap(),
            "(a.py\0|a.txt\0)|b.py"
        );
        assert_eq!(
            handle_sol_eol(r"(a$[b$]c$|d\$)e$").unwrap(),
            "(a\0[b$]c\0|d\\$)e\0"
        );

        assert_eq!(
            handle_sol_eol(r"(?:^a[^xyz]\^b$|^abc)").unwrap(),
            "(?:a[^xyz]\\^b\0|abc)"
        );
        // this is the edge case of current implementation of '^' support, which
        // should be rare since the regular expression is actually "wrong".
        assert_eq!(handle_sol_eol(r"^ab^c").unwrap(), "abc");
    }

    #[test]
    fn test_case_insensitive() {
        let case_sensitive = [true, false];
        for sensitive in case_sensitive {
            let m = RegexMatcher::new(r"(?:a/.*|b/.*)", sensitive).unwrap();
            assert_eq!(m.match_prefix(""), None);
            assert_eq!(m.match_prefix("A"), Some(!sensitive));
            assert_eq!(m.match_prefix("A/"), Some(!sensitive));
            assert_eq!(m.match_prefix("B"), Some(!sensitive));
            assert_eq!(m.match_prefix("B/"), Some(!sensitive));

            assert_eq!(m.matches("A/c"), !sensitive);
            assert_eq!(m.matches("A/b/c"), !sensitive);
            assert_eq!(m.matches("B/c"), !sensitive);
        }
    }
}


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

//! Tree-aware pattern matcher
//!
//! [TreeMatcher] is the main structure.

use std::path::Path;

use anyhow::Result;
use bitflags::bitflags;
use globset::Glob;
use globset::GlobBuilder;
use globset::GlobSet;
use globset::GlobSetBuilder;
use types::RepoPath;

use crate::DirectoryMatch;
use crate::Matcher;

bitflags! {
    struct RuleFlags: u8 {
        // A negative rule.
        const NEGATIVE = 1;

        // Auto-generated rule because the user specified a subpath.
        const PARENT = 2;

        // Mark a rule as "recursive" (ex. ending with "/**").
        const RECURSIVE = 4;
    }
}

#[derive(Clone, Debug)]
struct RuleInfo {
    flags: RuleFlags,
    orig_idx: usize,
}

/// Pattern matcher constructed by an ordered list of positive and negative
/// glob patterns. Negative patterns are prefixed with `!`.
///
/// The syntax is quite similar to gitignore, with some difference to avoid
/// inefficient uses. See [`TreeMatcher::from_rules`] for details about the
/// patterns.
///
/// The [TreeMatcher::match_recursive] method can be used to rule out
/// unnecessary directory visit early.
#[derive(Clone, Debug)]
pub struct TreeMatcher {
    // The [GlobSet] takes care of many algorithm stuff.  It can match a path
    // against multiple patterns and return the pattern indexes.
    glob_set: GlobSet,

    // Flags (ex. negative rule or is it a parent directory) for additional
    // information matching the pattern indexes.
    rule_info: Vec<RuleInfo>,
}

impl TreeMatcher {
    /// Create [TreeMatcher] using an ordered list of patterns.
    ///
    /// The patterns are glob patterns supported by the `globset` crate.
    /// Like gitignore, negative patterns are supported. They are prefixed
    /// with `!`. Special characters can be escaped by prefixing `\\`.
    ///
    /// Patterns are ordered. A later pattern always overrides an earlier
    /// pattern. Invalid patterns are ignored.
    ///
    /// Unlike gitignore, all patterns are treated as using absolute paths.
    /// That is, `*.c` is treated the same as `/*.c` and does not match `a/b.c`.
    /// Similarly, `!*.c` will be treated as `!/*.c`, in gitignore's sense.
    /// Use `**/*.c` to match files recursively. Note the `**` in the middle
    /// of a pattern effectively disable fast paths provided by `match_recursive`.
    ///
    /// Patterns do not match recursively.
    ///
    /// For example, both `/a/b` and `/a*/b*` do NOT match `/a/b/c/d`. Append
    /// `/**` to make rules recursive. The matcher works best if all rules end
    /// with `**`.
    pub fn from_rules(
        rules: impl Iterator<Item = impl AsRef<str>>,
        case_sensitive: bool,
    ) -> Result<Self, globset::Error> {
        let mut builder = GlobSetBuilder::new();
        let mut rule_info = Vec::new();

        for (idx, rule) in rules.enumerate() {
            let rule = rule.as_ref();
            let (negative, rule) = if rule.starts_with("!") {
                (true, &rule[1..])
            } else {
                (false, rule)
            };

            // Strip a leading "/". More friendly to gitignore users.
            let rule = if rule.starts_with("/") {
                &rule[1..]
            } else {
                rule
            };

            // "{", "}" do not have special meaning in gitignore, while
            // globset treats them differently.
            //
            // For now, workaround it by escaping. In the future, this can
            // possibly be done by tweaking a GlobBuilder option in
            // build_globs().
            //
            // See https://github.com/BurntSushi/ripgrep/issues/1183.
            let rule = escape_curly_brackets(&rule);

            // Add flags to the rule_id
            let mut flag = if negative {
                RuleFlags::NEGATIVE
            } else {
                RuleFlags::empty()
            };

            // Insert "parent" rules so match_recursive won't return "None"
            // incorrectly.
            let mut sep_index = 0;
            let rule_bytes = rule.as_ref();
            while let Some(index) = next_path_separator(rule_bytes, sep_index) {
                if index > 0 && index < rule_bytes.len() - 1 {
                    let parent_rule = &rule[..index];
                    for glob in build_globs(parent_rule, case_sensitive)? {
                        builder.add(glob);
                        rule_info.push(RuleInfo {
                            flags: flag | RuleFlags::PARENT,
                            orig_idx: idx,
                        });
                    }
                }
                sep_index = index + 1;
            }
            // Mark the rule as recursive so fast paths (i.e. claim everything
            // matches or nothing matches) can be used.
            if rule.ends_with("/**") || rule.ends_with("**/*") || rule == "**" {
                flag |= RuleFlags::RECURSIVE;
            }
            // Insert the rule.
            // NOTE: This crate depends on the fact that "a/**" matches "a", although
            // the documentation of globset might say otherwise.
            for glob in build_globs(&rule, case_sensitive)? {
                builder.add(glob);
                rule_info.push(RuleInfo {
                    flags: flag,
                    orig_idx: idx,
                });
            }
        }

        let glob_set = builder.build()?;
        let matcher = Self {
            glob_set,
            rule_info,
        };
        Ok(matcher)
    }

    /// Create [TreeMatcher] that matches nothing.
    pub fn never() -> Self {
        let rules: [&str; 0] = [];
        TreeMatcher::from_rules(rules.iter(), true).unwrap()
    }

    /// Create [TreeMatcher] that matches everything.
    pub fn always() -> Self {
        let rules: [&str; 1] = ["**"];
        TreeMatcher::from_rules(rules.iter(), true).unwrap()
    }

    /// Return `Some(bool)` if for all path inside the given `dir`,
    /// `matches(path)` will return `bool`.
    ///
    /// Return `None` if there is no fast path.
    ///
    /// `/` should be used as the path separator, regardless of system.
    pub fn match_recursive(&self, dir: impl AsRef<Path>) -> Option<bool> {
        let dir = dir.as_ref();
        // A subpath may match - cannot return Some(false)
        let mut subpath_may_match = false;
        // A subpath may mismatch - cannot return Some(true)
        let mut subpath_may_mismatch = false;
        for id in self.glob_set.matches(dir).into_iter().rev() {
            let flag = self.rule_info[id].flags;
            if flag.contains(RuleFlags::PARENT) {
                // An auto-generated parent rule matches.
                if flag.contains(RuleFlags::NEGATIVE) {
                    subpath_may_mismatch = true;
                } else {
                    subpath_may_match = true;
                }
            } else {
                // If it is not RECURSIVE, then fast paths (i.e. claim everything
                // matches, or nothing matches) cannot be used.
                if !flag.contains(RuleFlags::RECURSIVE) {
                    subpath_may_match = true;
                    subpath_may_mismatch = true;
                }
                // A non-parent rule matches.
                if flag.contains(RuleFlags::NEGATIVE) {
                    if subpath_may_match {
                        return None;
                    } else {
                        return Some(false);
                    }
                } else {
                    if subpath_may_mismatch {
                        return None;
                    } else {
                        return Some(true);
                    }
                }
            }
        }

        if subpath_may_match {
            None
        } else if !self.rule_info.is_empty() && dir.to_str() == Some("") {
            // Special case: empty dir
            None
        } else {
            Some(false)
        }
    }

    /// Return if `path` matches with the matcher.
    ///
    /// `/` should be used as the path separator, regardless of system.
    pub fn matches(&self, path: impl AsRef<Path>) -> bool {
        for id in self.glob_set.matches(path).into_iter().rev() {
            let flag = self.rule_info[id].flags;
            if flag.contains(RuleFlags::PARENT) {
                // For full path matches, parent rules do not count.
                continue;
            } else if flag.contains(RuleFlags::NEGATIVE) {
                // Not matched.
                return false;
            } else {
                // Matched.
                return true;
            }
        }
        // No rule matches
        false
    }

    /// Similar to matches, but return rule indexes matching the given path.
    /// Includes both positive and negative rules.
    pub fn matching_rule_indexes(&self, path: impl AsRef<Path>) -> Vec<usize> {
        let mut idxs: Vec<usize> = self
            .glob_set
            .matches(path)
            .into_iter()
            .filter_map(|idx: usize| {
                let info = &self.rule_info[idx];
                if info.flags.contains(RuleFlags::PARENT) {
                    None
                } else {
                    Some(info.orig_idx)
                }
            })
            .collect();
        idxs.dedup();
        idxs
    }
}

impl Matcher for TreeMatcher {
    fn matches_directory(&self, path: &RepoPath) -> Result<DirectoryMatch> {
        let dm = match self.match_recursive(path.as_str()) {
            Some(true) => DirectoryMatch::Everything,
            Some(false) => DirectoryMatch::Nothing,
            None => DirectoryMatch::ShouldTraverse,
        };
        Ok(dm)
    }

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

fn build_globs(pat: &str, case_sensitive: bool) -> Result<Vec<Glob>, globset::Error> {
    // Fast path (maybe).
    if pat.ends_with("/**") {
        let prefix = &pat[..pat.len() - 3];
        if !prefix.contains("?") && !prefix.contains("*") {
            // Rewrite "foo/**" (literal_separator=true) to
            // "foo" (literal_separator=false) and
            // "foo/*" (literal_separator=false) so
            // MatchStrategy::Literal and MatchStrategy::Prefix
            // instead of MatchStrategy::Regex) might be used.
            // See globset::glob::Glob::{literal,prefix}.
            let rules = [prefix, &pat[..pat.len() - 1]];
            return rules
                .iter()
                .map(|r| {
                    GlobBuilder::new(r)
                        .backslash_escape(true)
                        .case_insensitive(!case_sensitive)
                        .build()
                })
                .collect();
        } else {
            // Still append a rule without "/**".
            // We want pre-#1756 (https://github.com/BurntSushi/ripgrep/pull/1756) behavior.
            return Ok(vec![
                GlobBuilder::new(prefix)
                    .backslash_escape(true)
                    .literal_separator(true)
                    .case_insensitive(!case_sensitive)
                    .build()?,
                GlobBuilder::new(pat)
                    .backslash_escape(true)
                    .literal_separator(true)
                    .case_insensitive(!case_sensitive)
                    .build()?,
            ]);
        }
    }

    // General path.
    let glob = GlobBuilder::new(pat)
        .literal_separator(true) // `*` or `?` should not match `/`
        .backslash_escape(true)
        .case_insensitive(!case_sensitive)
        .build()?;
    Ok(vec![glob])
}

/// Find the next path separator in a pattern. Respect escaping rules.
/// Return the index (>= `start`), or None if there are no remaining path separator.
fn next_path_separator(pat: &[u8], start: usize) -> Option<usize> {
    let mut in_box_brackets = false;
    let mut escaped = false;

    for (i, ch) in pat.iter().skip(start).enumerate() {
        if escaped {
            match ch {
                _ => escaped = false,
            }
        } else if in_box_brackets {
            match ch {
                b']' => in_box_brackets = false,
                _ => {}
            }
        } else {
            match ch {
                b'\\' => escaped = true,
                b'[' => in_box_brackets = true,
                b'/' => return Some(i + start),
                _ => {}
            }
        }
    }

    None
}

/// Escape `{` and `}` so they no longer have special meanings to `globset`.
fn escape_curly_brackets(pat: &str) -> String {
    if pat.contains('{') || pat.contains('}') {
        let mut result = String::with_capacity(pat.len() * 2);
        for ch in pat.chars() {
            match ch {
                '{' => result.push_str("\\{"),
                '}' => result.push_str("\\}"),
                ch => result.push(ch),
            }
        }
        result
    } else {
        // No escaping is needed
        pat.to_string()
    }
}

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

    #[test]
    fn test_never_matcher() {
        let m = TreeMatcher::never();
        assert_eq!(m.match_recursive(""), Some(false));
        assert_eq!(m.match_recursive("a"), Some(false));
        assert_eq!(m.match_recursive("a/b"), Some(false));
        assert_eq!(m.matches(""), false);
        assert_eq!(m.matches("a/b"), false);
    }

    #[test]
    fn test_always_matcher() {
        let m = TreeMatcher::always();
        assert_eq!(m.match_recursive(""), Some(true));
        assert_eq!(m.match_recursive("a"), Some(true));
        assert_eq!(m.match_recursive("a/b"), Some(true));
        assert_eq!(m.matches(""), true);
        assert_eq!(m.matches("a/b"), true);
    }

    #[test]
    fn test_literal_paths() {
        let m = TreeMatcher::from_rules(["/a/**", "b/c/d/**", "\\e/\\f/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive(""), None);
        assert_eq!(m.match_recursive("a"), Some(true));
        assert_eq!(m.match_recursive("a/b"), Some(true));
        assert_eq!(m.match_recursive("b"), None);
        assert_eq!(m.match_recursive("b/c"), None);
        assert_eq!(m.match_recursive("b/c/d"), Some(true));
        assert_eq!(m.match_recursive("b/c/d/e"), Some(true));
        assert_eq!(m.match_recursive("e/f/g"), Some(true));
        assert_eq!(m.match_recursive("c"), Some(false));
        assert_eq!(m.match_recursive("c/a"), Some(false));
        assert_eq!(m.matches(""), false);
        assert_eq!(m.matches("a/b"), true);
        assert_eq!(m.matches("b/x"), false);
        assert_eq!(m.matches("b/c/d/e"), true);
        assert_eq!(m.matches("e"), false);
        assert_eq!(m.matches("e/f1"), false);
        assert_eq!(m.matches("e/f/g"), true);
    }

    #[test]
    fn test_simple_glob() {
        let m = TreeMatcher::from_rules(["a/*[cd][ef]/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a"), None);
        assert_eq!(m.match_recursive("b"), Some(false));
        assert_eq!(m.match_recursive("a/x"), Some(false));
        assert_eq!(m.match_recursive("a/xde"), Some(true));
        assert_eq!(m.match_recursive("a/xde/x"), Some(true));
        assert_eq!(m.matches("a/12df"), true);
        assert_eq!(m.matches("a/12df/12df"), true);
    }

    #[test]
    fn test_complex_glob() {
        let m = TreeMatcher::from_rules(["a/v/**/*.c/**", "a/**/w/*.c/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a"), None);
        assert_eq!(m.match_recursive("b"), Some(false));
        assert_eq!(m.match_recursive("a/v/.c"), Some(true));
        assert_eq!(m.match_recursive("a/v/.c/z"), Some(true));
        assert_eq!(m.match_recursive("a/z"), None);
        assert_eq!(m.matches("v/.c"), false);
        assert_eq!(m.matches("a/v/.c"), true);
        assert_eq!(m.matches("a/w/.c"), true);
        assert_eq!(m.matches("a/v/c/v/c/v/c/v/c/v.c"), true);
        assert_eq!(m.matches("a/c/c/c/c/w/w.c"), true);
        assert_eq!(m.matches("a/w/v/w.c"), false);

        // "{" has no special meaning
        let m = TreeMatcher::from_rules(["a/{b,c/d}/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a/{b,c/d}"), Some(true));
        assert_eq!(m.match_recursive("a/{b,c"), None);
        assert_eq!(m.match_recursive("a/{b,d"), Some(false));
    }

    #[test]
    fn test_mixed_literal_and_simple_glob() {
        let m = TreeMatcher::from_rules(["b/c/d/**", "b/*c/**", "b/1c/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive(""), None);
        assert_eq!(m.match_recursive("b/c/d/e"), Some(true));
        assert_eq!(m.match_recursive("b/1c"), Some(true));
        assert_eq!(m.match_recursive("b/xc/yc"), Some(true));
        assert_eq!(m.match_recursive("b/xc"), Some(true));
        assert_eq!(m.match_recursive("b/d"), Some(false));
        assert_eq!(m.matches("b/c/d/e/f"), true);
        assert_eq!(m.matches("b/fc"), true);
        assert_eq!(m.matches("b/ce"), false);
        assert_eq!(m.matches("b/c/e"), true);
    }

    #[test]
    fn test_mixed_literal_and_complex_glob() {
        let m = TreeMatcher::from_rules(["b/c/d/**", "b/**/c/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("b/c/d/e"), Some(true));
        assert_eq!(m.match_recursive("b/d"), None);
        assert_eq!(m.match_recursive("b/c"), Some(true));
        assert_eq!(m.match_recursive("b/x/c/y"), Some(true));
        assert_eq!(m.matches("b/c/d/e/f"), true);
        assert_eq!(m.matches("b/c/d"), true);
        assert_eq!(m.matches("b/c"), true);
        assert_eq!(m.matches("b"), false);
        assert_eq!(m.matches("b/x/y/c/x/y"), true);
    }

    #[test]
    fn test_empty_negative() {
        let m = TreeMatcher::from_rules(["!a/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive(""), None); // better answer is Some(false)
        assert_eq!(m.match_recursive("a"), Some(false));
        assert_eq!(m.match_recursive("a/b"), Some(false));
        assert_eq!(m.matches(""), false);
        assert_eq!(m.matches("a/b"), false);
    }

    #[test]
    fn test_literal_negative() {
        let m = TreeMatcher::from_rules(["a/**", "!a/b/**", "a/b/c/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a"), None);
        assert_eq!(m.match_recursive("a/c"), Some(true));
        assert_eq!(m.match_recursive("a/b"), None);
        assert_eq!(m.match_recursive("a/b/d"), Some(false));
        assert_eq!(m.match_recursive("a/b/c"), Some(true));
        assert_eq!(m.matches("a"), true);
        assert_eq!(m.matches("a/b"), false);
        assert_eq!(m.matches("a/b/c/d"), true);
        assert_eq!(m.matches("a/b/d"), false);
        assert_eq!(m.matches("a/c"), true);
        assert_eq!(m.matches("z"), false);
    }

    #[test]
    fn test_negative_override() {
        let m = TreeMatcher::from_rules(["a/**", "!a/**", "!b/**", "b/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a/b"), Some(false));
        assert_eq!(m.match_recursive("b/c"), Some(true));
        assert_eq!(m.matches("a"), false);
        assert_eq!(m.matches("b"), true);
    }

    #[test]
    fn test_mixed_negative_literal_simple_glob() {
        let m = TreeMatcher::from_rules(["a*/**", "!a1/**", "a1/a/**", "!a1/a*c/**"].iter(), true)
            .unwrap();
        assert_eq!(m.match_recursive("b"), Some(false));
        assert_eq!(m.match_recursive("a1/a"), Some(true));
        assert_eq!(m.matches("a"), true);
        assert_eq!(m.matches("a1"), false);
        assert_eq!(m.matches("a1/a"), true);
        assert_eq!(m.matches("a1/b"), false);
        assert_eq!(m.matches("a1/a1c"), false);
        assert_eq!(m.matches("a2"), true);
        assert_eq!(m.matches("b"), false);
    }

    #[test]
    fn test_fast_paths() {
        // Some interesting fast paths
        let m = TreeMatcher::from_rules(["a/**/b/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a/b"), Some(true));
        assert_eq!(m.match_recursive("a/1/2/3/b"), Some(true));

        let m = TreeMatcher::from_rules(["a/**/b/**", "!a/**/b/*/**"].iter(), true).unwrap();
        assert_eq!(m.match_recursive("a/b/1"), Some(false));
        assert_eq!(m.match_recursive("a/1/2/3/b/2"), Some(false));
    }

    #[test]
    fn test_non_recursive_patterns() {
        let m = TreeMatcher::from_rules(["a/*"].iter(), true).unwrap();
        assert!(m.matches("a/a"));
        assert!(!m.matches("b/a"));
        assert!(!m.matches("a"));
        assert_eq!(m.match_recursive("a"), None);
        assert_eq!(m.match_recursive("a/b"), None);
        assert_eq!(m.match_recursive("a/b/c"), Some(false));
        assert_eq!(m.match_recursive("b"), Some(false));

        let m = TreeMatcher::from_rules(["*a"].iter(), true).unwrap();
        assert!(m.matches("aa"));
        assert!(!m.matches("aa/b"));
        assert!(!m.matches("b"));
        assert_eq!(m.match_recursive("aa"), None);
        assert_eq!(m.match_recursive("a/a"), Some(false));
        assert_eq!(m.match_recursive("b"), Some(false));

        let m = TreeMatcher::from_rules(["b*/**/*a"].iter(), true).unwrap();
        assert!(m.matches("b1/aa"));
        assert!(!m.matches("c/aa"));
        assert!(!m.matches("b1/aa/11"));
        assert!(!m.matches("b/a/b"));
        assert!(m.matches("b/a/b/a"));
        assert_eq!(m.match_recursive("aa"), Some(false));
        assert_eq!(m.match_recursive("b/a/b"), None);
        assert_eq!(m.match_recursive("b/a/b/a"), None);
    }

    #[test]
    fn test_next_path_separator() {
        assert_eq!(next_path_separator(b"/a/b", 0), Some(0));
        assert_eq!(next_path_separator(b"/a/b", 1), Some(2));
        assert_eq!(next_path_separator(b"/a/b", 2), Some(2));
        assert_eq!(next_path_separator(b"/a/b", 3), None);
        assert_eq!(next_path_separator(b"[/]a\\/b", 0), None);
        assert_eq!(next_path_separator(b"\\[/]a", 0), Some(2));
    }

    #[test]
    fn test_t62872044() {
        let patterns = ["fbandroid/libraries/fbjni/**", "xplat/hermes/**"];
        let m1 = TreeMatcher::from_rules(patterns.iter(), true).unwrap();
        let m2 = TreeMatcher::from_rules(patterns.iter().rev(), true).unwrap();
        // aho-corasick 0.7.x where x < 9 can fail this test.
        // Update aho-corasick to 0.7.9 can solve it.
        // See https://github.com/BurntSushi/aho-corasick/issues/53.
        assert!(m1.matches("fbandroid/libraries/fbjni/x"));
        assert!(m2.matches("fbandroid/libraries/fbjni/x"));
    }

    #[test]
    fn test_matching_rule_indexes() {
        let pats = ["foo", "bar/baz", "qux/**", "z/**/z"];
        let m = TreeMatcher::from_rules(pats.iter(), true).unwrap();
        assert_eq!(m.matching_rule_indexes("banana"), vec![]);
        assert_eq!(m.matching_rule_indexes("foo"), vec![0]);
        assert_eq!(m.matching_rule_indexes("bar/baz"), vec![1]);
        assert_eq!(m.matching_rule_indexes("qux/some/thing"), vec![2]);
        assert_eq!(m.matching_rule_indexes("z/1/z"), vec![3]);
        assert_eq!(m.matching_rule_indexes("z/1"), vec![]);
    }

    #[test]
    fn test_case_insensitive() {
        let case_sensitive = [true, false];
        for sensitive in case_sensitive {
            let pats = ["foo", "bar/baz", "qux/**", "z/**/z"];
            let m = TreeMatcher::from_rules(pats.iter(), sensitive).unwrap();
            assert_eq!(m.matches("FOO"), !sensitive);
            assert_eq!(m.matches("BAR/baz"), !sensitive);
            assert_eq!(m.matches("bar/BAZ"), !sensitive);
            assert_eq!(m.matches("QUX/some/thing"), !sensitive);
            assert_eq!(m.matches("qux/SOME/thing"), true);
            assert_eq!(m.matches("qux/some/THING"), true);
            assert_eq!(m.matches("z/1/Z"), !sensitive);
            assert_eq!(m.matches("Z/1"), 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.
 */

//! Utility functions

/// Expand csh style brace expressions (`{` `}`) used in a glob pattern.
/// Return multiple glob patterns. If the brackets do not match, return
/// an empty vector.
///
/// TODO: Fix error handling so it returns an error when the input is
/// illegal.
///
/// Examples:
///
/// ```
/// use pathmatcher::expand_curly_brackets;
///
/// assert_eq!(expand_curly_brackets("foo"), vec!["foo"]);
/// assert_eq!(expand_curly_brackets("foo{a,b,}"), vec!["fooa", "foob", "foo"]);
/// assert_eq!(expand_curly_brackets("a{b,c{d,e}f}g"), vec!["abg", "acdfg", "acefg"]);
/// assert_eq!(expand_curly_brackets("{a,b}{}{c,d}{{e}}"), vec!["ace", "ade", "bce", "bde"]);
/// assert_eq!(expand_curly_brackets("\\{a\\}"), vec!["\\{a\\}"]);
/// assert_eq!(expand_curly_brackets("[{a}]"), vec!["[{a}]"]);
/// assert!(expand_curly_brackets("a}").is_empty());
/// assert!(expand_curly_brackets("{a").is_empty());
/// ```
pub fn expand_curly_brackets(pat: &str) -> Vec<String> {
    // A DAG of string segments. Vec indexes are used as identities.
    #[derive(Default, Debug)]
    struct StrNode(String, Vec<usize>);
    let mut dag = vec![StrNode::default()];

    // Convert the pattern to a DAG. For example, "a{b,c{d,e}f}g" is
    // converted to:
    //   dag[0] = ("a", [1, 2])
    //   dag[1] = ("b", [6])
    //   dag[2] = ("c", [3, 4])
    //   dag[3] = ("d", [5])
    //   dag[4] = ("e", [5])
    //   dag[5] = ("f", [6])
    //   dag[6] = ("g", [])

    let mut in_box_brackets = false;
    let mut escaped = false;

    // "Current" StrNode id used before "{"
    let mut bracket_stack: Vec<usize> = Vec::new();

    for ch in pat.chars() {
        let mut need_write = true;
        if escaped {
            match ch {
                _ => escaped = false,
            }
        } else if in_box_brackets {
            match ch {
                ']' => in_box_brackets = false,
                _ => {}
            }
        } else {
            match ch {
                '\\' => escaped = true,
                '[' => in_box_brackets = true,
                '{' => {
                    let next_id = dag.len();
                    let current_id = next_id - 1;
                    dag.push(StrNode::default());
                    bracket_stack.push(current_id);
                    dag[current_id].1.push(next_id);
                    need_write = false;
                }
                '}' => {
                    if bracket_stack.is_empty() {
                        // ill-formed pattern - '}' without '{'
                        return Vec::new();
                    }
                    // "Merge" all "heads" in "{ ... }" into one node
                    let next_id = dag.len();
                    dag.push(StrNode::default());
                    let last_id = bracket_stack.pop().unwrap();
                    for id in last_id + 1..next_id {
                        let is_head = dag[id].1.is_empty();
                        if is_head {
                            dag[id].1.push(next_id);
                        }
                    }
                    need_write = false;
                }
                ',' if !bracket_stack.is_empty() => {
                    // Start another "head"
                    let next_id = dag.len();
                    let last_id: usize = *bracket_stack.last().unwrap();
                    dag[last_id].1.push(next_id);
                    dag.push(StrNode::default());
                    need_write = false;
                }
                _ => {}
            }
        }

        // Write to the "current" node. It's always the last one.
        if need_write {
            dag.last_mut().unwrap().0.push(ch);
        }
    }

    if !bracket_stack.is_empty() {
        // '{' and '}' mismatched
        return Vec::new();
    }

    // Traverse the DAG to get all expanded strings
    let mut result = Vec::new();
    fn visit(dag: &Vec<StrNode>, result: &mut Vec<String>, prefix: String, id: usize) {
        let prefix = prefix + &dag[id].0;
        if id == dag.len() - 1 {
            assert!(dag[id].1.is_empty());
            result.push(prefix);
        } else {
            for child_id in dag[id].1.iter().cloned() {
                visit(dag, result, prefix.clone(), child_id);
            }
        }
    }
    visit(&dag, &mut result, String::new(), 0);
    result
}

/// Normalize a less strict glob pattern to a strict glob pattern.
///
/// In a strict glob pattern, `**` can only be a single directory component.
pub fn normalize_glob(pat: &str) -> String {
    let mut result = String::with_capacity(pat.len());
    let chars: Vec<_> = pat.chars().collect();
    for (i, &ch) in chars.iter().enumerate() {
        if ch == '*'
            && chars.get(i + 1).cloned() == Some('*')
            && !result.ends_with('\\')
            && !result.is_empty()
            && !result.ends_with('/')
        {
            // Change 'a**' to 'a*/**'
            result += "*/";
        }
        if ch == '/' && i > 0 && chars[i - 1] == '/' {
            continue;
        }
        result.push(ch);
        if ch == '*'
            && i > 0
            && chars[i - 1] == '*'
            && chars.get(i + 1) != None
            && chars.get(i + 1).cloned() != Some('/')
        {
            // Change '**a' to '**/*a'
            result += "/*";
        }
    }
    result
}

/// Escape special characters in a plain pattern so it can be used
/// as a glob pattern.
pub fn plain_to_glob(plain: &str) -> String {
    let mut result = String::with_capacity(plain.len());
    if plain.starts_with("!") {
        result.push('\\');
    }
    for ch in plain.chars() {
        match ch {
            '\\' | '*' | '{' | '}' | '[' | ']' => result.push('\\'),
            _ => {}
        }
        result.push(ch);
    }
    result
}

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

    #[test]
    fn test_normalize_glob() {
        assert_eq!(normalize_glob("**"), "**");
        assert_eq!(normalize_glob("a//b"), "a/b");
        assert_eq!(normalize_glob("a/**"), "a/**");
        assert_eq!(normalize_glob("a/b**"), "a/b*/**");
        assert_eq!(normalize_glob("a/b\\**"), "a/b\\**");
        assert_eq!(normalize_glob("a/**/c"), "a/**/c");
        assert_eq!(normalize_glob("a/**c"), "a/**/*c");
    }

    #[test]
    fn test_plain_to_glob() {
        assert_eq!(plain_to_glob("a[b{c*d\\e}]"), "a\\[b\\{c\\*d\\\\e\\}\\]");
        assert_eq!(plain_to_glob(""), "");
        assert_eq!(plain_to_glob("!a!"), "\\!a!");
    }
}


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

//! # `pprint`
//!
//! Pretty-print a serde value.
//!
//! - Support non-utf8 bytes.
//! - Support non-string keys in maps.
//! - Auto convert 20 or 32-byte hashes to hex format.
//! - Output is valid Python syntax.

use serde::Serialize;
pub use serde_cbor::Value;

/// Pretty-print a serializable `value` to a string.
pub fn pformat<T: Serialize>(value: &T) -> serde_cbor::Result<String> {
    let value = serde_cbor::value::to_value(value)?;
    Ok(pformat_value(&value))
}

/// Pretty-print `value` to a string.
pub fn pformat_value(value: &Value) -> String {
    let mut out = String::new();
    format_value(value, 0, &mut out);
    out
}

/// Print byte sequence with escapes.
fn format_bytes(value: &[u8], out: &mut String) {
    out.push_str("b\"");
    for &b in value {
        match b {
            0 => out.push_str("\\0"),
            b'"' => out.push_str("\\\""),
            b'\\' => out.push_str("\\\\"),
            b'\n' => out.push_str("\\n"),
            b'\r' => out.push_str("\\r"),
            b'\t' => out.push_str("\\t"),
            _ => {
                if b >= b' ' && b < 0x7f {
                    out.push(b as char)
                } else {
                    out.push_str("\\x");
                    out.push_str(&to_hex(&[b]));
                }
            }
        }
    }
    out.push('"');
}

/// Main entry. Write `value` to `out`.
/// For added new lines, prefix them with `indent` spaces.
fn format_value(value: &Value, indent: usize, out: &mut String) {
    use Value::Array;
    use Value::Bool;
    use Value::Bytes;
    use Value::Float;
    use Value::Integer;
    use Value::Map;
    use Value::Null;
    use Value::Text;
    match value {
        Null => out.push_str("None"),
        Bool(v) => out.push_str(if *v { "True" } else { "False" }),
        Integer(v) => out.push_str(&format!("{}", v)),
        Float(v) => out.push_str(&format!("{}", v)),
        Bytes(v) => {
            if [20, 32].contains(&v.len()) {
                out.push_str(&format!("bin({:?})", to_hex(v)));
            } else {
                format_bytes(&v, out);
            }
        }
        Text(v) => out.push_str(&format!("{:?}", v)),
        Array(a) => {
            out.push_str("[");
            for (i, v) in a.iter().enumerate() {
                if i > 0 {
                    out.push('\n');
                    out.push_str(&" ".repeat(indent + 1));
                }
                format_value(v, indent + 1, out);
                if i + 1 < a.len() {
                    out.push(',');
                }
            }
            out.push(']');
        }
        Map(m) => {
            out.push('{');
            for (i, (k, v)) in m.iter().enumerate() {
                if i > 0 {
                    out.push('\n');
                    out.push_str(&" ".repeat(indent + 1));
                }
                let kfmt = pformat_value(k);
                out.push_str(&kfmt);
                out.push_str(": ");
                format_value(v, indent + kfmt.len() + 3, out);
                if i + 1 < m.len() {
                    out.push(',');
                }
            }
            out.push('}')
        }
        _ => {}
    }
}

fn to_hex(slice: &[u8]) -> String {
    const HEX_CHARS: &[u8] = b"0123456789abcdef";
    let mut v = Vec::with_capacity(slice.len() * 2);
    for &byte in slice {
        v.push(HEX_CHARS[(byte >> 4) as usize]);
        v.push(HEX_CHARS[(byte & 0xf) as usize]);
    }
    unsafe { String::from_utf8_unchecked(v) }
}

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

    fn p(v: &impl Serialize) -> String {
        pformat(v).unwrap()
    }

    #[test]
    fn test_pformat() {
        assert_eq!(p(&[1, 2]), "[1,\n 2]");
        assert_eq!(p(&["a\n", "b"]), "[\"a\\n\",\n \"b\"]");
        assert_eq!(
            p(&[
                Value::Bytes(b"\0\x01".to_vec()),
                Value::Bytes(b"3\n4\t5".to_vec())
            ]),
            "[b\"\\0\\x01\",\n b\"3\\n4\\t5\"]"
        );
        assert_eq!(p(&[true, false]), "[True,\n False]");
        assert_eq!(p(&[Some(33), None]), "[33,\n None]");

        use std::collections::HashMap;
        let mut m = HashMap::new();
        m.insert(10, vec![2, 3]);
        m.insert(5, vec![4, 5]);
        assert_eq!(p(&m), "{5: [4,\n     5],\n 10: [2,\n      3]}");
    }

    #[test]
    fn test_hex_quote() {
        let v = Value::Bytes(b"12345678901234567890".to_vec());
        assert_eq!(
            p(&[&v]),
            "[bin(\"3132333435363738393031323334353637383930\")]"
        );
    }
}


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

fn main() {
    #[cfg(target_os = "macos")]
    {
        cc::Build::new().file("darwin.c").compile("darwin");
    }
}


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

fn main() {
    let mut ppid = procinfo::parent_pid(0);
    while ppid != 0 {
        let name = procinfo::exe_name(ppid);
        println!("Parent PID: {:8}  Name: {}", ppid, name);
        ppid = procinfo::parent_pid(ppid);
    }
}


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

fn main() {
    let bytes = std::env::args().map(|a| a.parse().unwrap_or(0)).sum();
    println!("Allocating {} bytes", bytes);
    let _vec = vec![3u8; bytes];
    {
        println!("Allocating another {} bytes", bytes);
        let _vec = vec![7u8; bytes];
        println!("Releasing {} bytes", bytes);
    }
    println!(
        "Max RSS: {} bytes (expected: around {} bytes)",
        procinfo::max_rss_bytes(),
        bytes * 2
    );
}


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

#[cfg(target_os = "macos")]
extern "C" {
    fn darwin_ppid(pid: u32) -> u32;
    fn darwin_exepath(pid: u32) -> *const libc::c_char;
}

#[cfg(windows)]
mod windows {
    use std::ffi::OsString;
    use std::mem::size_of;
    use std::mem::zeroed;
    use std::os::windows::ffi::OsStringExt;

    use winapi::shared::minwindef::DWORD;
    use winapi::um::handleapi::CloseHandle;
    use winapi::um::handleapi::INVALID_HANDLE_VALUE;
    use winapi::um::tlhelp32::CreateToolhelp32Snapshot;
    use winapi::um::tlhelp32::Process32FirstW;
    use winapi::um::tlhelp32::Process32NextW;
    use winapi::um::tlhelp32::PROCESSENTRY32W;
    use winapi::um::tlhelp32::TH32CS_SNAPPROCESS;
    use winapi::um::winnt::HANDLE;

    pub(crate) struct Snapshot {
        handle: HANDLE,
    }

    impl Drop for Snapshot {
        fn drop(&mut self) {
            unsafe { CloseHandle(self.handle) };
        }
    }

    impl Snapshot {
        fn new_pe32() -> PROCESSENTRY32W {
            let mut pe: PROCESSENTRY32W = unsafe { zeroed() };
            pe.dwSize = size_of::<PROCESSENTRY32W>() as DWORD;
            pe
        }

        pub(crate) fn new() -> Result<Snapshot, ()> {
            let snapshot_handle: HANDLE =
                unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };

            if snapshot_handle == INVALID_HANDLE_VALUE {
                return Err(());
            }

            Ok(Snapshot {
                handle: snapshot_handle,
            })
        }

        fn find<F>(&self, condition: F) -> Result<PROCESSENTRY32W, ()>
        where
            F: Fn(&PROCESSENTRY32W) -> bool,
        {
            let mut pe32 = Snapshot::new_pe32();
            if unsafe { Process32FirstW(self.handle, &mut pe32) } == 0 {
                return Err(());
            }

            loop {
                if condition(&pe32) {
                    return Ok(pe32);
                }

                if unsafe { Process32NextW(self.handle, &mut pe32) } == 0 {
                    break;
                }
            }

            Err(())
        }

        fn find_by_pid(&self, pid: DWORD) -> Result<PROCESSENTRY32W, ()> {
            self.find(|pe32| pe32.th32ProcessID == pid)
        }

        pub(crate) fn get_parent_process_id(&self, process_id: DWORD) -> Result<DWORD, ()> {
            self.find_by_pid(process_id)
                .map(|pe32| pe32.th32ParentProcessID)
        }

        pub(crate) fn get_process_executable_name(&self, process_id: DWORD) -> Result<String, ()> {
            self.find_by_pid(process_id)
                .map(|pe32| pe32.szExeFile)
                .map(|cs| {
                    cs.iter()
                        .take_while(|&&i| i != 0)
                        .map(|&i| i as u16)
                        .collect::<Vec<u16>>()
                })
                .map(|ref v| OsString::from_wide(v).into_string().unwrap_or("".into()))
        }
    }
}

/// Get the max RSS usage of the current process in bytes.
/// Return 0 on unsupported platform.
pub fn max_rss_bytes() -> u64 {
    #[cfg(unix)]
    {
        let usage = unsafe {
            let mut usage: libc::rusage = std::mem::zeroed();
            libc::getrusage(libc::RUSAGE_SELF, &mut usage as *mut _);
            usage
        };
        // POSIX didn't specify unit of ru_maxrss. Linux uses KB while BSD and
        // OSX use bytes (despite their manpages might say differently).
        let scale = if cfg!(target_os = "linux") {
            1024
        } else {
            // Assume BSD-ish
            1
        };
        return usage.ru_maxrss as u64 * scale;
    }

    #[cfg(windows)]
    {
        use winapi::shared::minwindef::DWORD;
        use winapi::um::processthreadsapi::GetCurrentProcess;
        use winapi::um::psapi::K32GetProcessMemoryInfo;
        use winapi::um::psapi::PROCESS_MEMORY_COUNTERS;
        let mut pmc: PROCESS_MEMORY_COUNTERS = unsafe { std::mem::zeroed() };
        pmc.cb = std::mem::size_of_val(&pmc) as DWORD;
        return match unsafe { K32GetProcessMemoryInfo(GetCurrentProcess(), &mut pmc, pmc.cb) } {
            0 => 0,
            _ => pmc.PeakWorkingSetSize as u64,
        };
    }

    #[allow(unreachable_code)]
    0
}

/// Get the parent pid. Return 0 on error or unsupported platform.
/// If pid is 0, return the parent pid of the current process.
pub fn parent_pid(pid: u32) -> u32 {
    if pid == 0 {
        #[cfg(unix)]
        return unsafe { libc::getppid() as u32 };

        #[cfg(not(unix))]
        return parent_pid(unsafe { libc::getpid() as u32 });
    }

    #[cfg(target_os = "macos")]
    unsafe {
        return crate::darwin_ppid(pid);
    }

    #[cfg(target_os = "linux")]
    {
        if let Ok(content) = std::fs::read_to_string(format!("/proc/{}/status", pid)) {
            let prefix = "PPid:";
            for line in content.lines() {
                if line.starts_with(prefix) {
                    if let Ok(ppid) = line[prefix.len()..].trim().parse() {
                        return ppid;
                    }
                }
            }
        }
        return 0;
    }

    #[cfg(windows)]
    {
        if let Ok(snapshot) = windows::Snapshot::new() {
            if let Ok(ppid) = snapshot.get_parent_process_id(pid) {
                return ppid;
            }
        }
    }

    #[allow(unreachable_code)]
    0
}

/// Get the executable name of the specified pid. Return an empty string on error.
pub fn exe_name(pid: u32) -> String {
    #[cfg(target_os = "macos")]
    {
        use std::ffi::CStr;
        let slice = unsafe { CStr::from_ptr(crate::darwin_exepath(pid)) };
        return slice.to_string_lossy().into_owned();
    }

    #[cfg(target_os = "linux")]
    {
        if let Ok(path) = std::fs::read_link(format!("/proc/{}/exe", pid)) {
            return path.into_os_string().to_string_lossy().into_owned();
        }
    }

    #[cfg(windows)]
    {
        if let Ok(snapshot) = windows::Snapshot::new() {
            if let Ok(name) = snapshot.get_process_executable_name(pid) {
                return name;
            }
        }
    }

    #[allow(unreachable_code)]
    String::new()
}


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

use std::borrow::Cow;
use std::fmt;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;

/// Statistics for cache data. For example, "files: hit 1000, miss 2".
pub struct CacheStats {
    /// Topic (ex. "Files").
    topic: Cow<'static, str>,

    /// Total cache hit.
    hit: AtomicUsize,

    /// Total read.
    miss: AtomicUsize,
}

impl CacheStats {
    /// Create a [`CacheStats`].
    pub fn new(topic: impl Into<Cow<'static, str>>) -> Arc<Self> {
        let topic = topic.into();
        let stats = Self {
            topic,
            hit: Default::default(),
            miss: Default::default(),
        };
        Arc::new(stats)
    }

    /// The topic of the [`CacheStats`].
    pub fn topic(&self) -> &str {
        &self.topic
    }

    /// Increase the cache miss count.
    pub fn increase_miss(&self, count: usize) {
        self.miss.fetch_add(count, Relaxed);
    }

    /// Increase the cache hit count.
    pub fn increase_hit(&self, count: usize) {
        self.hit.fetch_add(count, Relaxed);
    }

    /// Get the cache miss count.
    pub fn miss(&self) -> usize {
        self.miss.load(Relaxed)
    }

    /// Get the cache hit count.
    pub fn hit(&self) -> usize {
        self.hit.load(Relaxed)
    }

    /// Get total read count.
    pub fn total(&self) -> usize {
        self.miss() + self.hit()
    }
}

impl fmt::Debug for CacheStats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "[{} hit {} miss {}]",
            &self.topic,
            self.hit(),
            self.miss()
        )
    }
}


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

use std::fmt;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::Relaxed;
use std::time::Duration;
use std::time::Instant;

use once_cell::sync::Lazy;

static INSTANT_START: Lazy<Instant> = Lazy::new(Instant::now);

#[derive(Clone, Debug)]
pub struct IoSample {
    /// Total input bytes.
    input_bytes: u64,

    /// Total output bytes.
    output_bytes: u64,

    /// Count (ex. requests).
    count: u64,

    /// Timestamp the sample is taken.
    at: Instant,
}

/// A sample of total read / write bytes.
#[derive(Default)]
pub(crate) struct MutableIoSample {
    /// Total input bytes.
    input_bytes: AtomicU64,

    /// Total output bytes.
    output_bytes: AtomicU64,

    /// Count (ex. requests).
    count: AtomicU64,

    /// The sample is taken at this timestamp (milliseconds after UNIX_EPOCH).
    epoch: AtomicU64,
}

impl fmt::Debug for MutableIoSample {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{:?}|{:?}|{:?}",
            self.input_bytes, self.output_bytes, self.count
        )
    }
}

impl IoSample {
    pub fn from_io_bytes(input_bytes: u64, output_bytes: u64) -> Self {
        Self::from_io_bytes_count(input_bytes, output_bytes, 0)
    }

    pub fn from_io_bytes_count(input_bytes: u64, output_bytes: u64, count: u64) -> Self {
        Self {
            input_bytes,
            output_bytes,
            count,
            at: Instant::now(),
        }
    }

    /// Set `at` for testing.
    pub(crate) fn test_at(mut self, millis: u64) -> Self {
        self.at = *INSTANT_START + Duration::from_millis(millis);
        self
    }
}

impl MutableIoSample {
    /// Update the sample to the specified data.
    pub(crate) fn set(&self, sample: IoSample) {
        self.input_bytes.store(sample.input_bytes, Relaxed);
        self.output_bytes.store(sample.output_bytes, Relaxed);
        self.count.store(sample.count, Relaxed);
        let millis = if sample.at <= *INSTANT_START {
            0
        } else {
            sample.at.duration_since(*INSTANT_START).as_millis()
        };
        self.epoch.store(millis as u64, Relaxed);
    }

    /// Sum of input and output bytes.
    pub(crate) fn total_bytes(&self) -> u64 {
        self.input_bytes.load(Relaxed) + self.output_bytes.load(Relaxed)
    }

    pub(crate) fn input_bytes(&self) -> u64 {
        self.input_bytes.load(Relaxed)
    }

    pub(crate) fn output_bytes(&self) -> u64 {
        self.output_bytes.load(Relaxed)
    }

    /// Number of (ex. requests).
    pub(crate) fn count(&self) -> u64 {
        self.count.load(Relaxed)
    }

    /// Milliseconds since `other`. Normalize to at least 1.
    fn millis_since(&self, other: &Self) -> u64 {
        self.epoch
            .load(Relaxed)
            .saturating_sub(other.epoch.load(Relaxed))
            .max(1)
    }

    /// Speed calculated from the previous state.
    /// Return (input_speed, output_speed), bytes per second.
    pub(crate) fn bytes_per_second(&self, prev: &Self) -> (u64, u64) {
        let input_delta = self
            .input_bytes
            .load(Relaxed)
            .saturating_sub(prev.input_bytes.load(Relaxed));
        let output_delta = self
            .output_bytes
            .load(Relaxed)
            .saturating_sub(prev.output_bytes.load(Relaxed));
        let millis_delta = self.millis_since(prev);
        (
            input_delta * 1000 / millis_delta,
            output_delta * 1000 / millis_delta,
        )
    }
}

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

    impl IoSample {
        pub(crate) fn to_mutable(self) -> MutableIoSample {
            let mutable_sample = MutableIoSample::default();
            mutable_sample.set(self);
            mutable_sample
        }
    }

    #[test]
    fn test_total_bytes() {
        let a = IoSample::from_io_bytes(10, 20).test_at(0).to_mutable();

        assert_eq!(a.total_bytes(), 30);
    }

    #[test]
    fn test_bytes_per_second() {
        let a = IoSample::from_io_bytes(10, 20).test_at(0).to_mutable();
        let b = IoSample::from_io_bytes(100, 200).test_at(2002).to_mutable();

        assert_eq!(b.millis_since(&a), 2002);
        assert_eq!(a.millis_since(&b), 1);

        assert_eq!(b.bytes_per_second(&a), (44, 89));
        assert_eq!(a.bytes_per_second(&b), (0, 0)); // no panic (saturating_sub)
        assert_eq!(a.bytes_per_second(&a), (0, 0)); // no panic
    }
}


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

//! Data model related to progress reporting.
//!
//! - Pure data. Minimal state just enough for rendering.
//! - Separate from rendering.
//! - Lock-free (nice to have).

mod cache_stats;
mod io_sample;
mod progress_bar;
mod registry;
mod time_series;

pub use cache_stats::CacheStats;
pub use io_sample::IoSample;
pub use progress_bar::AggregatingProgressBar;
pub use progress_bar::ProgressBar;
pub use registry::Registry;
pub use time_series::IoTimeSeries;
pub use time_series::TimeSeriesMode;


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

use std::borrow::Cow;
use std::fmt;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Release;
use std::sync::Arc;
use std::sync::Weak;
use std::time::Duration;
use std::time::Instant;

use arc_swap::ArcSwapOption;
use parking_lot::Mutex;

use crate::Registry;

/// A progress bar. It has multiple `Metric`s and a `Metric`.
///
/// ```plain,ignore
/// topic [ message ] [ pos / total unit1 ], [ pos / total unit2 ], ...
/// ```
pub struct ProgressBar {
    topic: Cow<'static, str>,
    message: ArcSwapOption<String>,
    pos: AtomicU64,
    total: AtomicU64,
    unit: Cow<'static, str>,
    created_at: Instant,
}

impl ProgressBar {
    /// Create a new progress bar of the given topic (ex. "writing files").
    pub fn new(
        topic: impl Into<Cow<'static, str>>,
        total: u64,
        unit: impl Into<Cow<'static, str>>,
    ) -> Arc<Self> {
        let bar = Self {
            topic: topic.into(),
            unit: unit.into(),
            total: AtomicU64::new(total),
            pos: Default::default(),
            message: Default::default(),
            created_at: Instant::now(),
        };
        Arc::new(bar)
    }

    /// Create a new progress bar and register with default registry.
    pub fn register_new(
        topic: impl Into<Cow<'static, str>>,
        total: u64,
        unit: impl Into<Cow<'static, str>>,
    ) -> Arc<Self> {
        let bar = Self::new(topic, total, unit);
        Registry::main().register_progress_bar(&bar);
        bar
    }

    /// Get the progress bar topic.
    pub fn topic(&self) -> &str {
        &self.topic
    }

    /// Get the progress message.
    pub fn message(&self) -> Option<Arc<String>> {
        self.message.load_full()
    }

    /// Set the progress message.
    pub fn set_message(&self, message: String) {
        self.message.store(Some(Arc::new(message)));
    }

    /// Obtain the position and total.
    pub fn position_total(&self) -> (u64, u64) {
        let pos = self.pos.load(Acquire);
        let total = self.total.load(Acquire);
        (pos, total)
    }

    /// Set total.
    pub fn set_total(&self, total: u64) {
        self.total.store(total, Release);
    }

    /// Set position.
    pub fn set_position(&self, pos: u64) {
        self.pos.store(pos, Release);
    }

    /// Increase position.
    pub fn increase_position(&self, inc: u64) {
        self.pos.fetch_add(inc, AcqRel);
    }

    /// Increase total.
    pub fn increase_total(&self, inc: u64) {
        self.total.fetch_add(inc, AcqRel);
    }

    /// Obtain unit.
    pub fn unit(&self) -> &str {
        &self.unit
    }

    /// Time since the creation of the progress bar.
    pub fn elapsed(&self) -> Duration {
        self.created_at.elapsed()
    }
}

impl fmt::Debug for ProgressBar {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (pos, total) = self.position_total();
        write!(f, "[{} {}/{} {}", self.topic(), pos, total, self.unit())?;
        if let Some(message) = self.message() {
            write!(f, " {}", message)?;
        }
        Ok(())
    }
}

pub struct AggregatingProgressBar {
    bar: Mutex<Weak<ProgressBar>>,
    topic: Cow<'static, str>,
    unit: Cow<'static, str>,
}

/// AggregatingProgressBar allows sharing a progress bar across
/// concurrent uses when otherwise inconvenient. For example, it lets
/// you display a single progress bar via a low level client object
/// when that client is used by multiple high level threads.
impl AggregatingProgressBar {
    pub fn new(
        topic: impl Into<Cow<'static, str>>,
        unit: impl Into<Cow<'static, str>>,
    ) -> Arc<Self> {
        Arc::new(AggregatingProgressBar {
            bar: Mutex::new(Weak::new()),
            topic: topic.into(),
            unit: unit.into(),
        })
    }

    /// If progress bar exists, increase its total, otherwise create a
    /// new progress bar. You should avoid calling set_position or
    /// set_total on the returned ProgressBar.
    pub fn create_or_extend(&self, additional_total: u64) -> Arc<ProgressBar> {
        let mut bar = self.bar.lock();

        match bar.upgrade() {
            Some(bar) => {
                bar.increase_total(additional_total);
                bar
            }
            None => {
                let new_bar = ProgressBar::register_new(
                    self.topic.clone(),
                    additional_total,
                    self.unit.clone(),
                );
                *bar = Arc::downgrade(&new_bar);
                new_bar
            }
        }
    }
}

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

    #[test]
    fn test_aggregating_bar() {
        let agg = AggregatingProgressBar::new("eat", "apples");

        {
            let bar1 = agg.create_or_extend(10);
            bar1.increase_position(5);
            assert_eq!((5, 10), agg.create_or_extend(0).position_total());

            {
                let bar2 = agg.create_or_extend(5);
                bar2.increase_position(5);
                assert_eq!((10, 15), agg.create_or_extend(0).position_total());
            }

            assert_eq!((10, 15), agg.create_or_extend(0).position_total());
        }

        Registry::main().remove_orphan_progress_bar();

        assert_eq!((0, 0), agg.create_or_extend(0).position_total());
    }
}


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

//! Progress rendering configuration.

use std::borrow::Cow;
use std::time::Duration;

use unicode_width::UnicodeWidthChar;
use unicode_width::UnicodeWidthStr;

pub struct RenderingConfig {
    /// Delay before showing a newly created bar.
    pub delay: Duration,

    /// Maximum number of bars to show.
    pub max_bar_count: usize,

    /// Terminal width.
    pub term_width: usize,

    /// Use CJK width (some characters are treated as 2-char width, instead of 1-char width).
    /// Practically, some CJK fonts would work better with this set to true.
    pub cjk_width: bool,
}

impl Default for RenderingConfig {
    fn default() -> Self {
        Self {
            delay: Duration::from_secs(3),
            max_bar_count: 8,
            term_width: 80,
            cjk_width: false,
        }
    }
}

#[cfg(test)]
impl RenderingConfig {
    pub fn for_testing() -> Self {
        Self {
            delay: Duration::from_secs(0),
            max_bar_count: 5,
            term_width: 60,
            cjk_width: false,
        }
    }
}

impl RenderingConfig {
    /// Truncate a single line.
    pub(crate) fn truncate_line<'a>(&self, line: &'a str) -> Cow<'a, str> {
        self.truncate_by_width(line, self.term_width, "â€¦")
    }

    /// Truncate `text` to fit in the given width.
    /// `suffix` is appended if `text` is truncated.
    pub(crate) fn truncate_by_width<'a>(
        &self,
        text: &'a str,
        width: usize,
        suffix: &str,
    ) -> Cow<'a, str> {
        if self.width_str(text) >= width {
            let mut current_width = 0;
            let suffix_width = self.width_str(suffix);
            for (i, ch) in text.char_indices() {
                let next_width = current_width + self.width_char(ch);
                if next_width + suffix_width >= width {
                    // Cannot take this char.
                    return format!("{}{}", &text[..i], suffix).into();
                }
                current_width = next_width;
            }
        }
        return Cow::Borrowed(text);
    }

    pub(crate) fn max_topic_len(&self) -> usize {
        match self.term_width {
            w if w >= 100 => 20,
            w if w >= 80 => 16,
            _ => 12,
        }
    }

    fn width_str(&self, text: &str) -> usize {
        if self.cjk_width {
            text.width_cjk()
        } else {
            text.width()
        }
    }

    fn width_char(&self, ch: char) -> usize {
        if self.cjk_width {
            ch.width_cjk()
        } else {
            ch.width()
        }
        .unwrap_or_default()
    }
}


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

//! Progress rendering.

mod config;
pub mod simple;
mod unit;

pub use config::RenderingConfig;

#[cfg(test)]
mod tests;


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

pub fn unit_phrase(unit: &str, pos: u64, total: u64) -> String {
    match unit {
        "%" => {
            let total = total.max(1);
            format!("{}%", pos.min(total) * 100 / total)
        }
        "bytes" | "B" => {
            if total == 0 {
                human_bytes(pos as _)
            } else {
                format!("{}/{}", human_bytes(pos as _), human_bytes(total as _))
            }
        }
        _ => {
            if total == 0 {
                if pos == 0 {
                    String::new()
                } else {
                    format!("{} {}", pos, unit)
                }
            } else {
                format!("{}/{} {}", pos, total, unit)
            }
        }
    }
}

pub fn human_bytes(bytes: u64) -> String {
    if bytes < 5000 {
        format!("{}B", bytes)
    } else if bytes < 5_000_000 {
        format!("{}KB", bytes / 1000)
    } else if bytes < 5_000_000_000 {
        format!("{}MB", bytes / 1000000)
    } else {
        format!("{}GB", bytes / 1000000000)
    }
}


/*
 * 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 minibench::bench;
use minibench::elapsed;
use radixbuf::key::FixedKey;
use radixbuf::key::KeyId;
use radixbuf::radix::radix_insert;
use radixbuf::radix::radix_lookup;
use rand::ChaChaRng;
use rand::RngCore;

const N: usize = 204800;

/// Generate random buffer
fn gen_buf(size: usize) -> Vec<u8> {
    let mut buf = vec![0u8; size];
    ChaChaRng::new_unseeded().fill_bytes(buf.as_mut());
    buf
}

fn batch_insert_radix_buf(key_buf: &Vec<u8>, count: usize) -> Vec<u32> {
    let mut radix_buf = vec![0u32; 16];
    for i in 0..count {
        let key_id: KeyId = ((i * 20) as u32).into();
        radix_insert(&mut radix_buf, 0, key_id, FixedKey::read, key_buf).expect("insert");
    }
    radix_buf
}

fn main() {
    bench("index insertion", || {
        let key_buf = gen_buf(20 * N);
        elapsed(|| {
            batch_insert_radix_buf(&key_buf, N);
        })
    });

    bench("index lookup", || {
        let key_buf = gen_buf(20 * N);
        let radix_buf = batch_insert_radix_buf(&key_buf, N);
        elapsed(move || {
            for i in 0..N {
                let key_id = (i as u32 * 20).into();
                let key = FixedKey::read(&key_buf, key_id).unwrap();
                radix_lookup(&radix_buf, 0, &key, FixedKey::read, &key_buf).expect("lookup");
            }
        })
    });
}


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

//! Base16 iterator
//!
//! The main radix tree uses base16 to support hex string prefix lookups and
//! make the space usage more efficient.

#[derive(Debug, Copy, Clone)]
pub struct Base16Iter<'a, T>(&'a T, usize, usize);

impl<'a, T: AsRef<[u8]>> Base16Iter<'a, T> {
    /// Convert base256 binary sequence to a base16 iterator.
    pub fn from_bin(binary: &'a T) -> Self {
        let len = binary.as_ref().len() * 2;
        Base16Iter(binary, 0, len)
    }
}

/// Base16 iterator for [u8]
impl<'a, T: AsRef<[u8]>> Iterator for Base16Iter<'a, T> {
    type Item = u8;

    #[inline]
    fn next(&mut self) -> Option<u8> {
        if self.2 <= self.1 {
            None
        } else {
            let i = self.1;
            self.1 = i + 1;
            let v = self.0.as_ref()[i / 2];
            if i & 1 == 0 { v >> 4 } else { v & 0xf }.into()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }

    #[inline]
    fn count(self) -> usize {
        self.len()
    }
}

impl<'a, T: AsRef<[u8]>> DoubleEndedIterator for Base16Iter<'a, T> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.2 <= self.1 {
            None
        } else {
            let i = self.2 - 1;
            self.2 = i;
            let v = self.0.as_ref()[i / 2];
            if i & 1 == 0 { v >> 4 } else { v & 0xf }.into()
        }
    }
}

impl<'a, T: AsRef<[u8]>> ExactSizeIterator for Base16Iter<'a, T> {
    #[inline]
    fn len(&self) -> usize {
        self.2 - self.1
    }
}

impl<'a, T: AsRef<[u8]>> Base16Iter<'a, T> {
    #[inline]
    pub fn skip(self, n: usize) -> Self {
        Base16Iter(self.0, self.1 + n, self.2)
    }

    #[inline]
    pub fn take(self, n: usize) -> Self {
        // TODO: Use (self.1 + n).min(self.2) once ord_max_min is available (Rust 1.22)
        let mut end = self.1 + n;
        if self.2 < end {
            end = self.2
        }
        Base16Iter(self.0, self.1, end)
    }
}

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

    use super::*;

    quickcheck! {
        fn check_skip_rev(src: Vec<u8>) -> bool {
            let iter = Base16Iter::from_bin(&src);
            let full: Vec<u8> = iter.clone().collect();
            let rev: Vec<u8> = iter.clone().rev().collect();
            (0..full.len()).all(|i| {
                let v: Vec<u8> = iter.clone().skip(i).collect();
                let r: Vec<u8> = iter.clone().skip(i).rev().rev().rev().collect();
                v[..] == full[i..] && r[..] == rev[..(rev.len() - i)]
            })
        }
    }

    // The below patterns (skip, zip, rev; skip, take, rev) are used in radix.rs.
    // Make sure they work at iterator level without needing an extra container.

    #[test]
    fn test_zip_skip_rev() {
        let x = [0x12, 0x34, 0x56, 0x21u8];
        let y = [0x78, 0x90, 0xab, 0xcdu8];
        let i = Base16Iter::from_bin(&x)
            .skip(2)
            .zip(Base16Iter::from_bin(&y).skip(3))
            .rev(); // .rev() works directly
        let v: Vec<(u8, u8)> = i.collect();
        assert_eq!(v.capacity(), v.len());
        assert_eq!(v, vec![(2, 0xd), (6, 0xc), (5, 0xb), (4, 0xa), (3, 0)]);
    }

    #[test]
    fn test_skip_take_rev() {
        let x = [0x12, 0x34, 0x56u8];
        let i = Base16Iter::from_bin(&x).skip(3).take(3).rev(); // .rev() works directly
        let v: Vec<u8> = i.collect();
        assert_eq!(v, vec![6, 5, 4]);
    }
}


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

use crate::key::KeyId;

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ErrorKind {
    #[error("offset {0} is out of range")]
    OffsetOverflow(u64),
    #[error("ambiguous prefix")]
    AmbiguousPrefix,
    #[error("{0:?} cannot be a prefix of {1:?}")]
    PrefixConflict(KeyId, KeyId),
    #[error("{0:?} cannot be resolved")]
    InvalidKeyId(KeyId),
    #[error("{0} is not a base16 value")]
    InvalidBase16(u8),
}


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

//! Types and helper functions to convert a `KeyId` to a Key (`[u8]`). The radix index is
//! responsible for the other direction - `[u8]` to `KeyId`. This mapping is to remove the
//! need of storing full keys in the radix index.
//!
//! Typically, there is an append-only file which is the source of truth of some information.
//! The `KeyId` could be offsets in that file.
//!
//! The radix index APIs only care about how to read a key, given a `KeyId`. It does not
//! care about how to write a key. This is more flexible. For example, the source of truth
//! file may have keys stored already before the index gets built. The `append` methods provided
//! below are only for convenience.
//!
//! There could be cases where the `KeyId`s are not offsets, or the keys are not stored as-is.
//! For example, `KeyId` could be mercurial revision numbers instead of offsets for non-inlined
//! revlog .i files. The radix index has a limitation (to make the format slightly more compact)
//! that one key cannot be the prefix of another different key. This could be guaranteed by using
//! the fixed length key, or appending `'\0'` to strings that cannot have `'\0'`. But in case
//! of variant-length binary key, some escaping is needed to avoid the prefix conflict. The
//! escaping logic might be implemented as key read/write functions.

use std::io::Write;

use vlqencoding::VLQDecodeAt;
use vlqencoding::VLQEncode;

use crate::errors::ErrorKind;
use crate::traits::Resize;
use crate::Result;

/// Integer that maps to a key (`[u8]`).
#[derive(Debug, Eq, PartialEq, PartialOrd, Ord, Clone, Copy, Default)]
pub struct KeyId(u32);

macro_rules! impl_convert {
    ($T: ty) => {
        impl From<$T> for KeyId {
            #[allow(unused_comparisons)]
            #[inline]
            fn from(v: $T) -> Self {
                if v > 0xffff_ffff || v < 0 {
                    panic!("KeyId out of range")
                }
                KeyId(v as u32)
            }
        }

        impl Into<$T> for KeyId {
            #[inline]
            fn into(self) -> $T {
                self.0 as $T
            }
        }
    };
}

impl_convert!(u32);
impl_convert!(u64);

// Make sure `usize` contains at least 4 bytes (`KeyId` size).
const _SIZE_TEST: usize = 0xffff_ffff;
impl_convert!(usize);

/// Keys with fixed 20 bytes length.
pub struct FixedKey;

/// Keys with variant-length. Serialized as a VLQ-encoded length, followed by the actual key.
pub struct VariantKey;

impl FixedKey {
    #[inline]
    pub fn read<'a, K: AsRef<[u8]>>(key_buf: &'a K, key_id: KeyId) -> Result<&'a [u8]> {
        let key_buf = key_buf.as_ref();
        let start_pos: usize = key_id.into();
        // LANG: Consider making 20 a type parameter once supported.
        let end_pos = start_pos + 20;
        if key_buf.len() < end_pos {
            return Err(ErrorKind::InvalidKeyId(key_id));
        }
        Ok(&key_buf[start_pos..end_pos])
    }

    #[inline]
    pub fn append<B: AsMut<[u8]> + Resize<u8>, K: AsRef<[u8]>>(key_buf: &mut B, key: &K) -> KeyId {
        let key = key.as_ref();
        assert_eq!(key.len(), 20);

        let pos = key_buf.as_mut().len();
        key_buf.resize(pos + 20, 0);
        let key_buf = key_buf.as_mut();
        key_buf[pos..pos + 20].copy_from_slice(key);
        pos.into()
    }
}

impl VariantKey {
    #[inline]
    pub fn read<'a, K: AsRef<[u8]>>(key_buf: &'a K, key_id: KeyId) -> Result<&'a [u8]> {
        let key_buf = key_buf.as_ref();
        let (key_len, vlq_size): (usize, _) = key_buf
            .read_vlq_at(key_id.0 as usize)
            .map_err(|_| ErrorKind::InvalidKeyId(key_id))?;
        let start_pos = key_id.0 as usize + vlq_size;
        let end_pos = start_pos + key_len;
        if key_buf.len() < end_pos {
            return Err(ErrorKind::InvalidKeyId(key_id));
        }
        Ok(&key_buf[start_pos as usize..end_pos as usize])
    }

    #[inline]
    pub fn append<B: AsMut<[u8]> + Resize<u8>, K: AsRef<[u8]>>(key_buf: &mut B, key: &K) -> KeyId {
        let key = key.as_ref();
        // PERF: the Vec allocation may be avoided with a more complex implementation.
        // Most of the time, key length could be encoded in 2-byte VLQ. Pre-allocate 2 bytes.
        let mut buf = Vec::<u8>::with_capacity(key.len() + 2);
        buf.write_vlq(key.len()).expect("write len");
        buf.write_all(key).expect("write key");

        let pos = key_buf.as_mut().len();
        key_buf.resize(pos + buf.len(), 0);
        let key_buf = key_buf.as_mut();
        key_buf[pos..pos + buf.len()].copy_from_slice(&buf[..]);
        (pos as u64).into()
    }
}

#[cfg(test)]
mod tests {
    use rand::distributions::Alphanumeric;
    use rand::thread_rng;
    use rand::Rng;
    use rand::RngCore;

    use super::*;

    #[test]
    fn test_variant_key_round_trip() {
        let mut buf = Vec::<u8>::new();
        let keys: Vec<Vec<u8>> = (0..1000usize)
            .map(|i| {
                let rng = thread_rng();
                rng.sample_iter(&Alphanumeric).take(i % 40).collect()
            })
            .collect();
        let key_ids: Vec<KeyId> = keys
            .iter()
            .map(|key| VariantKey::append(&mut buf, key))
            .collect();
        let keys_retrieved: Vec<Vec<u8>> = key_ids
            .iter()
            .map(|&id| VariantKey::read(&buf, id).unwrap().to_vec())
            .collect();
        assert_eq!(keys, keys_retrieved);
    }

    #[test]
    fn test_fixed_key_round_trip() {
        let mut rng = thread_rng();
        let mut buf = Vec::<u8>::new();
        let keys: Vec<[u8; 20]> = (0..1000usize)
            .map(|_| {
                let mut bytes = [0u8; 20];
                rng.fill_bytes(&mut bytes);
                bytes
            })
            .collect();
        let key_ids: Vec<KeyId> = keys
            .iter()
            .map(|key| FixedKey::append(&mut buf, key))
            .collect();
        assert_eq!(buf.len(), 1000 * 20);
        let keys_retrieved: Vec<[u8; 20]> = key_ids
            .iter()
            .map(|&id| {
                let mut bytes = [0u8; 20];
                bytes.copy_from_slice(FixedKey::read(&buf, id).unwrap());
                bytes
            })
            .collect();
        assert_eq!(keys, keys_retrieved);
    }
}


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

pub mod base16;
pub mod errors;
pub mod key;
pub mod radix;
pub mod traits;

pub use errors::ErrorKind as Error;
pub type Result<T> = std::result::Result<T, Error>;