<commit_before>// Std
use std::collections::HashMap;
use std::ffi::OsString;
use std::mem;
use std::ops::Deref;

// Internal
use crate::builder::{Arg, ArgPredicate, Command};
use crate::parser::AnyValue;
use crate::parser::{ArgMatches, MatchedArg, SubCommand, ValueSource};
use crate::util::Id;
use crate::INTERNAL_ERROR_MSG;

#[derive(Debug, Default)]
pub(crate) struct ArgMatcher(ArgMatches);

impl ArgMatcher {
    pub(crate) fn new(_cmd: &Command) -> Self {
        ArgMatcher(ArgMatches {
            #[cfg(debug_assertions)]
            valid_args: {
                let args = _cmd.get_arguments().map(|a| a.id.clone());
                let groups = _cmd.get_groups().map(|g| g.id.clone());
                args.chain(groups).collect()
            },
            #[cfg(debug_assertions)]
            valid_subcommands: _cmd.get_subcommands().map(|sc| sc.get_id()).collect(),
            // HACK: Allow an external subcommand's ArgMatches be a stand-in for any ArgMatches
            // since users can't detect it and avoid the asserts.
            //
            // See clap-rs/clap#3263
            #[cfg(debug_assertions)]
            #[cfg(not(feature = "unstable-v4"))]
            disable_asserts: _cmd.is_allow_external_subcommands_set(),
            #[cfg(debug_assertions)]
            #[cfg(feature = "unstable-v4")]
            disable_asserts: false,
            ..Default::default()
        })
    }

    pub(crate) fn into_inner(self) -> ArgMatches {
        self.0
    }

    pub(crate) fn propagate_globals(&mut self, global_arg_vec: &[Id]) {
        debug!(
            "ArgMatcher::get_global_values: global_arg_vec={:?}",
            global_arg_vec
        );
        let mut vals_map = HashMap::new();
        self.fill_in_global_values(global_arg_vec, &mut vals_map);
    }

    fn fill_in_global_values(
        &mut self,
        global_arg_vec: &[Id],
        vals_map: &mut HashMap<Id, MatchedArg>,
    ) {
        for global_arg in global_arg_vec {
            if let Some(ma) = self.get(global_arg) {
                // We have to check if the parent's global arg wasn't used but still exists
                // such as from a default value.
                //
                // For example, `myprog subcommand --global-arg=value` where --global-arg defines
                // a default value of `other` myprog would have an existing MatchedArg for
                // --global-arg where the value is `other`, however the occurs will be 0.
                let to_update = if let Some(parent_ma) = vals_map.get(global_arg) {
                    if parent_ma.get_occurrences() > 0 && ma.get_occurrences() == 0 {
                        parent_ma
                    } else {
                        ma
                    }
                } else {
                    ma
                }
                .clone();
                vals_map.insert(global_arg.clone(), to_update);
            }
        }
        if let Some(ref mut sc) = self.0.subcommand {
            let mut am = ArgMatcher(mem::take(&mut sc.matches));
            am.fill_in_global_values(global_arg_vec, vals_map);
            mem::swap(&mut am.0, &mut sc.matches);
        }

        for (name, matched_arg) in vals_map.iter_mut() {
            self.0.args.insert(name.clone(), matched_arg.clone());
        }
    }

    pub(crate) fn get(&self, arg: &Id) -> Option<&MatchedArg> {
        self.0.args.get(arg)
    }

    pub(crate) fn get_mut(&mut self, arg: &Id) -> Option<&mut MatchedArg> {
        self.0.args.get_mut(arg)
    }

    pub(crate) fn remove(&mut self, arg: &Id) {
        self.0.args.swap_remove(arg);
    }

    pub(crate) fn contains(&self, arg: &Id) -> bool {
        self.0.args.contains_key(arg)
    }

    pub(crate) fn arg_names(&self) -> indexmap::map::Keys<Id, MatchedArg> {
        self.0.args.keys()
    }

    pub(crate) fn entry(&mut self, arg: &Id) -> indexmap::map::Entry<Id, MatchedArg> {
        self.0.args.entry(arg.clone())
    }

    pub(crate) fn subcommand(&mut self, sc: SubCommand) {
        self.0.subcommand = Some(Box::new(sc));
    }

    pub(crate) fn subcommand_name(&self) -> Option<&str> {
        self.0.subcommand_name()
    }

    pub(crate) fn iter(&self) -> indexmap::map::Iter<Id, MatchedArg> {
        self.0.args.iter()
    }

    pub(crate) fn check_explicit<'a>(&self, arg: &Id, predicate: ArgPredicate<'a>) -> bool {
        self.get(arg).map_or(false, |a| a.check_explicit(predicate))
    }

    pub(crate) fn start_custom_arg(&mut self, arg: &Arg, source: ValueSource) {
        let id = &arg.id;
        debug!(
            "ArgMatcher::start_custom_arg: id={:?}, source={:?}",
            id, source
        );
        let ma = self.entry(id).or_insert(MatchedArg::new_arg(arg));
        debug_assert_eq!(ma.type_id(), Some(arg.get_value_parser().type_id()));
        ma.set_source(source);
        ma.new_val_group();
    }

    pub(crate) fn start_custom_group(&mut self, id: &Id, source: ValueSource) {
        debug!(
            "ArgMatcher::start_custom_arg: id={:?}, source={:?}",
            id, source
        );
        let ma = self.entry(id).or_insert(MatchedArg::new_group());
        debug_assert_eq!(ma.type_id(), None);
        ma.set_source(source);
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_arg(&mut self, arg: &Arg) {
        let id = &arg.id;
        debug!("ArgMatcher::start_occurrence_of_arg: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new_arg(arg));
        debug_assert_eq!(ma.type_id(), Some(arg.get_value_parser().type_id()));
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_group(&mut self, id: &Id) {
        debug!("ArgMatcher::start_occurrence_of_group: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new_group());
        debug_assert_eq!(ma.type_id(), None);
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_external(&mut self, cmd: &crate::Command) {
        let id = &Id::empty_hash();
        debug!("ArgMatcher::start_occurrence_of_external: id={:?}", id,);
        let ma = self.entry(id).or_insert(MatchedArg::new_external(cmd));
        debug_assert_eq!(
            ma.type_id(),
            Some(
                cmd.get_external_subcommand_value_parser()
                    .expect(INTERNAL_ERROR_MSG)
                    .type_id()
            )
        );
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn add_val_to(&mut self, arg: &Id, val: AnyValue, raw_val: OsString) {
        let ma = self.get_mut(arg).expect(INTERNAL_ERROR_MSG);
        ma.append_val(val, raw_val);
    }

    pub(crate) fn add_index_to(&mut self, arg: &Id, idx: usize) {
        let ma = self.get_mut(arg).expect(INTERNAL_ERROR_MSG);
        ma.push_index(idx);
    }

    pub(crate) fn needs_more_vals(&self, o: &Arg) -> bool {
        debug!("ArgMatcher::needs_more_vals: o={}", o.name);
        if let Some(ma) = self.get(&o.id) {
            let current_num = ma.num_vals();
            if let Some(num) = o.num_vals {
                debug!("ArgMatcher::needs_more_vals: num_vals...{}", num);
                return if o.is_multiple_occurrences_set() {
                    (current_num % num) != 0
                } else {
                    num != current_num
                };
            } else if let Some(num) = o.max_vals {
                debug!("ArgMatcher::needs_more_vals: max_vals...{}", num);
                return current_num < num;
            } else if o.min_vals.is_some() {
                debug!("ArgMatcher::needs_more_vals: min_vals...true");
                return true;
            }
            return o.is_multiple_values_set();
        }
        true
    }
}

impl Deref for ArgMatcher {
    type Target = ArgMatches;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<commit_msg>refactor(parser): Allow more match-state<commit_after>// Std
use std::collections::HashMap;
use std::ffi::OsString;
use std::mem;
use std::ops::Deref;

// Internal
use crate::builder::{Arg, ArgPredicate, Command};
use crate::parser::AnyValue;
use crate::parser::{ArgMatches, MatchedArg, SubCommand, ValueSource};
use crate::util::Id;
use crate::INTERNAL_ERROR_MSG;

#[derive(Debug, Default)]
pub(crate) struct ArgMatcher {
    matches: ArgMatches,
}

impl ArgMatcher {
    pub(crate) fn new(_cmd: &Command) -> Self {
        ArgMatcher {
            matches: ArgMatches {
                #[cfg(debug_assertions)]
                valid_args: {
                    let args = _cmd.get_arguments().map(|a| a.id.clone());
                    let groups = _cmd.get_groups().map(|g| g.id.clone());
                    args.chain(groups).collect()
                },
                #[cfg(debug_assertions)]
                valid_subcommands: _cmd.get_subcommands().map(|sc| sc.get_id()).collect(),
                // HACK: Allow an external subcommand's ArgMatches be a stand-in for any ArgMatches
                // since users can't detect it and avoid the asserts.
                //
                // See clap-rs/clap#3263
                #[cfg(debug_assertions)]
                #[cfg(not(feature = "unstable-v4"))]
                disable_asserts: _cmd.is_allow_external_subcommands_set(),
                #[cfg(debug_assertions)]
                #[cfg(feature = "unstable-v4")]
                disable_asserts: false,
                ..Default::default()
            },
        }
    }

    pub(crate) fn into_inner(self) -> ArgMatches {
        self.matches
    }

    pub(crate) fn propagate_globals(&mut self, global_arg_vec: &[Id]) {
        debug!(
            "ArgMatcher::get_global_values: global_arg_vec={:?}",
            global_arg_vec
        );
        let mut vals_map = HashMap::new();
        self.fill_in_global_values(global_arg_vec, &mut vals_map);
    }

    fn fill_in_global_values(
        &mut self,
        global_arg_vec: &[Id],
        vals_map: &mut HashMap<Id, MatchedArg>,
    ) {
        for global_arg in global_arg_vec {
            if let Some(ma) = self.get(global_arg) {
                // We have to check if the parent's global arg wasn't used but still exists
                // such as from a default value.
                //
                // For example, `myprog subcommand --global-arg=value` where --global-arg defines
                // a default value of `other` myprog would have an existing MatchedArg for
                // --global-arg where the value is `other`, however the occurs will be 0.
                let to_update = if let Some(parent_ma) = vals_map.get(global_arg) {
                    if parent_ma.get_occurrences() > 0 && ma.get_occurrences() == 0 {
                        parent_ma
                    } else {
                        ma
                    }
                } else {
                    ma
                }
                .clone();
                vals_map.insert(global_arg.clone(), to_update);
            }
        }
        if let Some(ref mut sc) = self.matches.subcommand {
            let mut am = ArgMatcher {
                matches: mem::take(&mut sc.matches),
            };
            am.fill_in_global_values(global_arg_vec, vals_map);
            mem::swap(&mut am.matches, &mut sc.matches);
        }

        for (name, matched_arg) in vals_map.iter_mut() {
            self.matches.args.insert(name.clone(), matched_arg.clone());
        }
    }

    pub(crate) fn get(&self, arg: &Id) -> Option<&MatchedArg> {
        self.matches.args.get(arg)
    }

    pub(crate) fn get_mut(&mut self, arg: &Id) -> Option<&mut MatchedArg> {
        self.matches.args.get_mut(arg)
    }

    pub(crate) fn remove(&mut self, arg: &Id) {
        self.matches.args.swap_remove(arg);
    }

    pub(crate) fn contains(&self, arg: &Id) -> bool {
        self.matches.args.contains_key(arg)
    }

    pub(crate) fn arg_names(&self) -> indexmap::map::Keys<Id, MatchedArg> {
        self.matches.args.keys()
    }

    pub(crate) fn entry(&mut self, arg: &Id) -> indexmap::map::Entry<Id, MatchedArg> {
        self.matches.args.entry(arg.clone())
    }

    pub(crate) fn subcommand(&mut self, sc: SubCommand) {
        self.matches.subcommand = Some(Box::new(sc));
    }

    pub(crate) fn subcommand_name(&self) -> Option<&str> {
        self.matches.subcommand_name()
    }

    pub(crate) fn iter(&self) -> indexmap::map::Iter<Id, MatchedArg> {
        self.matches.args.iter()
    }

    pub(crate) fn check_explicit<'a>(&self, arg: &Id, predicate: ArgPredicate<'a>) -> bool {
        self.get(arg).map_or(false, |a| a.check_explicit(predicate))
    }

    pub(crate) fn start_custom_arg(&mut self, arg: &Arg, source: ValueSource) {
        let id = &arg.id;
        debug!(
            "ArgMatcher::start_custom_arg: id={:?}, source={:?}",
            id, source
        );
        let ma = self.entry(id).or_insert(MatchedArg::new_arg(arg));
        debug_assert_eq!(ma.type_id(), Some(arg.get_value_parser().type_id()));
        ma.set_source(source);
        ma.new_val_group();
    }

    pub(crate) fn start_custom_group(&mut self, id: &Id, source: ValueSource) {
        debug!(
            "ArgMatcher::start_custom_arg: id={:?}, source={:?}",
            id, source
        );
        let ma = self.entry(id).or_insert(MatchedArg::new_group());
        debug_assert_eq!(ma.type_id(), None);
        ma.set_source(source);
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_arg(&mut self, arg: &Arg) {
        let id = &arg.id;
        debug!("ArgMatcher::start_occurrence_of_arg: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new_arg(arg));
        debug_assert_eq!(ma.type_id(), Some(arg.get_value_parser().type_id()));
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_group(&mut self, id: &Id) {
        debug!("ArgMatcher::start_occurrence_of_group: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new_group());
        debug_assert_eq!(ma.type_id(), None);
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn start_occurrence_of_external(&mut self, cmd: &crate::Command) {
        let id = &Id::empty_hash();
        debug!("ArgMatcher::start_occurrence_of_external: id={:?}", id,);
        let ma = self.entry(id).or_insert(MatchedArg::new_external(cmd));
        debug_assert_eq!(
            ma.type_id(),
            Some(
                cmd.get_external_subcommand_value_parser()
                    .expect(INTERNAL_ERROR_MSG)
                    .type_id()
            )
        );
        ma.set_source(ValueSource::CommandLine);
        ma.inc_occurrences();
        ma.new_val_group();
    }

    pub(crate) fn add_val_to(&mut self, arg: &Id, val: AnyValue, raw_val: OsString) {
        let ma = self.get_mut(arg).expect(INTERNAL_ERROR_MSG);
        ma.append_val(val, raw_val);
    }

    pub(crate) fn add_index_to(&mut self, arg: &Id, idx: usize) {
        let ma = self.get_mut(arg).expect(INTERNAL_ERROR_MSG);
        ma.push_index(idx);
    }

    pub(crate) fn needs_more_vals(&self, o: &Arg) -> bool {
        debug!("ArgMatcher::needs_more_vals: o={}", o.name);
        if let Some(ma) = self.get(&o.id) {
            let current_num = ma.num_vals();
            if let Some(num) = o.num_vals {
                debug!("ArgMatcher::needs_more_vals: num_vals...{}", num);
                return if o.is_multiple_occurrences_set() {
                    (current_num % num) != 0
                } else {
                    num != current_num
                };
            } else if let Some(num) = o.max_vals {
                debug!("ArgMatcher::needs_more_vals: max_vals...{}", num);
                return current_num < num;
            } else if o.min_vals.is_some() {
                debug!("ArgMatcher::needs_more_vals: min_vals...true");
                return true;
            }
            return o.is_multiple_values_set();
        }
        true
    }
}

impl Deref for ArgMatcher {
    type Target = ArgMatches;

    fn deref(&self) -> &Self::Target {
        &self.matches
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#![deny(warnings)]
#![deny(rust_2018_idioms)]
#![deny(clippy::all)]

use std::collections::hash_map::HashMap;
use std::fmt;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

#[derive(Default)]
pub struct BufferedFileSystem {
    files: HashMap<PathBuf, Option<String>>,
}

///
/// A filesystem wrapper that buffers file reads and writes until `commit()` is called:
///
/// - Write: it only writes to the cache.
/// - Delete: it only writes to the cache and mark the entry as deleted.
/// - Read/Exists: it first checks the cache, if there is no entry in the cache, delegates to the standard library.
/// - Commit: it clears the cache, write all contents to the disk, and remove all files that are marked as deleted.
///

impl BufferedFileSystem {
    pub fn new() -> Self {
        BufferedFileSystem {
            files: Default::default(),
        }
    }

    pub fn read_to_string<P: AsRef<Path>>(&self, path: P) -> io::Result<String> {
        if let Some(content) = self.files.get(path.as_ref()) {
            match content {
                None => Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Trying to read a deleted file.",
                )),
                Some(content) => Ok(content.clone()),
            }
        } else {
            fs::read_to_string(path)
        }
    }

    pub fn write<P: AsRef<Path>>(&mut self, path: P, content: String) {
        self.files
            .insert(path.as_ref().to_path_buf(), Some(content));
    }

    pub fn remove_file<P: AsRef<Path>>(&mut self, path: P) {
        self.files.insert(path.as_ref().to_path_buf(), None);
    }

    pub fn has_changes(&self) -> bool {
        !self.files.is_empty()
    }

    pub fn exists<P: AsRef<Path>>(&self, path: P) -> bool {
        let path = path.as_ref();
        match self.files.get(path) {
            None => path.exists(),
            Some(data) => data.is_some(),
        }
    }

    // TODO: Optimize, make concurrent writes
    pub fn commit(&mut self) -> io::Result<(Vec<PathBuf>, Vec<PathBuf>)> {
        let mut removed = vec![];
        let mut added = vec![];
        for (path, data) in self.files.drain() {
            match data {
                None => {
                    fs::remove_file(&path)?;
                    removed.push(path);
                }
                Some(data) => {
                    fs::write(&path, data)?;
                    added.push(path);
                }
            }
        }
        Ok((added, removed))
    }
}

impl fmt::Debug for BufferedFileSystem {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut added = vec![];
        let mut updated = vec![];
        let mut removed = vec![];
        for (path, data) in self.files.iter() {
            match data {
                None => {
                    removed.push(path);
                }
                Some(_) => {
                    if path.exists() {
                        updated.push(path);
                    } else {
                        added.push(path);
                    }
                }
            }
        }
        write!(
            f,
            "Added: {:#?}\nRemoved: {:#?}\nUpdated: {:#?}",
            added, removed, updated
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::env::temp_dir;
    const CONTENTS: &str = "TEST🤯";

    #[test]
    fn test_write() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_write");
        let mut fs = BufferedFileSystem::new();

        assert!(!fs.has_changes());
        assert!(!fs.exists(&path));
        fs.write(&path, CONTENTS.to_owned());
        assert!(fs.has_changes());
        assert!(fs.exists(&path));
        dbg!(&fs);

        let result = fs.read_to_string(path)?;
        assert_eq!(result, CONTENTS);
        Ok(())
    }

    #[test]
    fn test_commit() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_commit");
        if path.exists() {
            std::fs::remove_file(&path)?;
        }
        let mut fs = BufferedFileSystem::new();
        fs.write(&path, CONTENTS.to_owned());
        let (added, removed) = fs.commit()?;
        assert_eq!(added.len(), 1);
        assert!(removed.is_empty());

        // Read from the disk
        let result = std::fs::read_to_string(path)?;
        assert_eq!(result, CONTENTS);
        Ok(())
    }

    #[test]
    fn test_remove() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_remove");
        let mut fs = BufferedFileSystem::new();
        fs.write(&path, CONTENTS.to_owned());
        assert!(fs.has_changes());
        assert!(fs.exists(&path));
        fs.remove_file(&path);
        assert!(fs.has_changes());
        assert!(!fs.exists(&path));
        Ok(())
    }

    #[test]
    fn test_commit_remove() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_commit_remove");
        std::fs::write(&path, CONTENTS)?;

        let mut fs = BufferedFileSystem::new();
        assert!(fs.exists(&path));
        fs.remove_file(&path);
        assert!(fs.has_changes());
        assert!(!fs.exists(&path));

        let (added, removed) = fs.commit()?;
        assert_eq!(removed.len(), 1);
        assert!(added.is_empty());

        // Should be deleted from the disk
        assert!(!path.exists());
        Ok(())
    }
}
<commit_msg>remove leftover dbg!<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#![deny(warnings)]
#![deny(rust_2018_idioms)]
#![deny(clippy::all)]

use std::collections::hash_map::HashMap;
use std::fmt;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

#[derive(Default)]
pub struct BufferedFileSystem {
    files: HashMap<PathBuf, Option<String>>,
}

///
/// A filesystem wrapper that buffers file reads and writes until `commit()` is called:
///
/// - Write: it only writes to the cache.
/// - Delete: it only writes to the cache and mark the entry as deleted.
/// - Read/Exists: it first checks the cache, if there is no entry in the cache, delegates to the standard library.
/// - Commit: it clears the cache, write all contents to the disk, and remove all files that are marked as deleted.
///

impl BufferedFileSystem {
    pub fn new() -> Self {
        BufferedFileSystem {
            files: Default::default(),
        }
    }

    pub fn read_to_string<P: AsRef<Path>>(&self, path: P) -> io::Result<String> {
        if let Some(content) = self.files.get(path.as_ref()) {
            match content {
                None => Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Trying to read a deleted file.",
                )),
                Some(content) => Ok(content.clone()),
            }
        } else {
            fs::read_to_string(path)
        }
    }

    pub fn write<P: AsRef<Path>>(&mut self, path: P, content: String) {
        self.files
            .insert(path.as_ref().to_path_buf(), Some(content));
    }

    pub fn remove_file<P: AsRef<Path>>(&mut self, path: P) {
        self.files.insert(path.as_ref().to_path_buf(), None);
    }

    pub fn has_changes(&self) -> bool {
        !self.files.is_empty()
    }

    pub fn exists<P: AsRef<Path>>(&self, path: P) -> bool {
        let path = path.as_ref();
        match self.files.get(path) {
            None => path.exists(),
            Some(data) => data.is_some(),
        }
    }

    // TODO: Optimize, make concurrent writes
    pub fn commit(&mut self) -> io::Result<(Vec<PathBuf>, Vec<PathBuf>)> {
        let mut removed = vec![];
        let mut added = vec![];
        for (path, data) in self.files.drain() {
            match data {
                None => {
                    fs::remove_file(&path)?;
                    removed.push(path);
                }
                Some(data) => {
                    fs::write(&path, data)?;
                    added.push(path);
                }
            }
        }
        Ok((added, removed))
    }
}

impl fmt::Debug for BufferedFileSystem {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut added = vec![];
        let mut updated = vec![];
        let mut removed = vec![];
        for (path, data) in self.files.iter() {
            match data {
                None => {
                    removed.push(path);
                }
                Some(_) => {
                    if path.exists() {
                        updated.push(path);
                    } else {
                        added.push(path);
                    }
                }
            }
        }
        write!(
            f,
            "Added: {:#?}\nRemoved: {:#?}\nUpdated: {:#?}",
            added, removed, updated
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::env::temp_dir;
    const CONTENTS: &str = "TEST🤯";

    #[test]
    fn test_write() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_write");
        let mut fs = BufferedFileSystem::new();

        assert!(!fs.has_changes());
        assert!(!fs.exists(&path));
        fs.write(&path, CONTENTS.to_owned());
        assert!(fs.has_changes());
        assert!(fs.exists(&path));

        let result = fs.read_to_string(path)?;
        assert_eq!(result, CONTENTS);
        Ok(())
    }

    #[test]
    fn test_commit() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_commit");
        if path.exists() {
            std::fs::remove_file(&path)?;
        }
        let mut fs = BufferedFileSystem::new();
        fs.write(&path, CONTENTS.to_owned());
        let (added, removed) = fs.commit()?;
        assert_eq!(added.len(), 1);
        assert!(removed.is_empty());

        // Read from the disk
        let result = std::fs::read_to_string(path)?;
        assert_eq!(result, CONTENTS);
        Ok(())
    }

    #[test]
    fn test_remove() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_remove");
        let mut fs = BufferedFileSystem::new();
        fs.write(&path, CONTENTS.to_owned());
        assert!(fs.has_changes());
        assert!(fs.exists(&path));
        fs.remove_file(&path);
        assert!(fs.has_changes());
        assert!(!fs.exists(&path));
        Ok(())
    }

    #[test]
    fn test_commit_remove() -> io::Result<()> {
        let mut path = temp_dir();
        path.push("buffered_fs_test_commit_remove");
        std::fs::write(&path, CONTENTS)?;

        let mut fs = BufferedFileSystem::new();
        assert!(fs.exists(&path));
        fs.remove_file(&path);
        assert!(fs.has_changes());
        assert!(!fs.exists(&path));

        let (added, removed) = fs.commit()?;
        assert_eq!(removed.len(), 1);
        assert!(added.is_empty());

        // Should be deleted from the disk
        assert!(!path.exists());
        Ok(())
    }
}
<|endoftext|>


<commit_before>use hash_map::HashMap;
use std::borrow::Borrow;
use std::hash::{BuildHasher, Hash};

use FailedAllocationError;

#[cfg(target_pointer_width = "32")]
const CANARY: usize = 0x42cafe99;
#[cfg(target_pointer_width = "64")]
const CANARY: usize = 0x42cafe9942cafe99;

#[derive(Clone, Debug)]
enum JournalEntry {
    Insert(usize),
    GOIW(usize),
    Remove(usize),
    DidClear(usize),
}

#[derive(Clone, Debug)]
pub struct DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    map: HashMap<K, (usize, V), S>,
    journal: Vec<JournalEntry>,
    readonly: bool,
}

impl<K: Hash + Eq, V, S: BuildHasher> DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    #[inline(always)]
    pub fn inner(&self) -> &HashMap<K, (usize, V), S> {
        &self.map
    }

    #[inline(never)]
    pub fn begin_mutation(&mut self) {
        self.map.verify();
        assert!(self.readonly);
        self.readonly = false;
        self.verify();
    }

    #[inline(never)]
    pub fn end_mutation(&mut self) {
        self.map.verify();
        assert!(!self.readonly);
        self.readonly = true;
        self.verify();
    }

    fn verify(&self) {
        let mut position = 0;
        let mut bad_canary: Option<(usize, *const usize)> = None;
        for (_,v) in self.map.iter() {
            let canary_ref = &v.0;
            if *canary_ref == CANARY {
                position += 1;
                continue;
            }
            bad_canary = Some((*canary_ref, canary_ref));
        }
        if let Some(c) = bad_canary {
            self.report_corruption(c.0, c.1, position);
        }
    }

    #[inline(always)]
    pub fn with_hasher(hash_builder: S) -> Self {
        Self {
            map: HashMap::<K, (usize, V), S>::with_hasher(hash_builder),
            journal: Vec::new(),
            readonly: true,
        }
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.map.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    #[inline(always)]
    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        self.map.contains_key(k)
    }

    #[inline(always)]
    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        self.map.get(k).map(|v| &v.1)
    }

    #[inline(always)]
    pub fn try_get_or_insert_with<F: FnOnce() -> V>(
        &mut self,
        key: K,
        default: F
    ) -> Result<&mut V, FailedAllocationError> {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::GOIW(self.map.make_hash(&key).inspect()));
        let entry = self.map.try_entry(key)?;
        Ok(&mut entry.or_insert_with(|| (CANARY, default())).1)
    }

    #[inline(always)]
    pub fn try_insert(&mut self, k: K, v: V) -> Result<Option<V>, FailedAllocationError> {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::Insert(self.map.make_hash(&k).inspect()));
        let old = self.map.try_insert(k, (CANARY, v))?;
        Ok(old.map(|x| x.1))
    }

    #[inline(always)]
    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::Remove(self.map.make_hash(k).inspect()));
        self.map.remove(k).map(|x| x.1)
    }

    #[inline(always)]
    pub fn clear(&mut self) where K: 'static, V: 'static  {
        // We handle scoped mutations for the caller here, since callsites that
        // invoke clear() don't benefit from the coalescing we do around insertion.
        self.begin_mutation();
        self.journal.clear();
        self.journal.push(JournalEntry::DidClear(self.map.raw_capacity()));
        self.map.clear();
        self.end_mutation();
    }

    #[inline(never)]
    fn report_corruption(
        &self,
        canary: usize,
        canary_addr: *const usize,
        position: usize
    ) {
        use ::std::ffi::CString;
        let key = b"HashMapJournal\0";
        let value = CString::new(format!("{:?}", self.journal)).unwrap();
        unsafe {
            Gecko_AnnotateCrashReport(
                key.as_ptr() as *const ::std::os::raw::c_char,
                value.as_ptr(),
            );
        }
        panic!(
            "HashMap Corruption (sz={}, cap={}, pairsz={}, cnry={:#x}, pos={}, base_addr={:?}, cnry_addr={:?}, jrnl_len={})",
            self.map.len(),
            self.map.raw_capacity(),
            ::std::mem::size_of::<(K, (usize, V))>(),
            canary,
            position,
            self.map.raw_buffer(),
            canary_addr,
            self.journal.len(),
        );
    }
}

impl<K, V, S> PartialEq for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          V: PartialEq,
          S: BuildHasher
{
    fn eq(&self, other: &Self) -> bool {
        self.map.eq(&other.map)
    }
}

impl<K, V, S> Eq for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          V: Eq,
          S: BuildHasher
{
}

impl<K, V, S> Default for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher + Default
{
    fn default() -> Self {
        Self {
            map: HashMap::default(),
            journal: Vec::new(),
            readonly: true,
        }
    }
}

impl<K: Hash + Eq, V, S: BuildHasher> Drop for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    fn drop(&mut self) {
        self.map.verify();
        debug_assert!(self.readonly, "Dropped while mutating");
        self.verify();
    }
}

extern "C" {
    pub fn Gecko_AnnotateCrashReport(key_str: *const ::std::os::raw::c_char,
                                     value_str: *const ::std::os::raw::c_char);
}
<commit_msg>Auto merge of #18861 - upsuper:diagnose-fix, r=Manishearth<commit_after>use hash_map::HashMap;
use std::borrow::Borrow;
use std::hash::{BuildHasher, Hash};
use std::ptr;

use FailedAllocationError;

#[cfg(target_pointer_width = "32")]
const CANARY: usize = 0x42cafe99;
#[cfg(target_pointer_width = "64")]
const CANARY: usize = 0x42cafe9942cafe99;

#[cfg(target_pointer_width = "32")]
const POISON: usize = 0xdeadbeef;
#[cfg(target_pointer_width = "64")]
const POISON: usize = 0xdeadbeefdeadbeef;

#[derive(Clone, Debug)]
enum JournalEntry {
    Insert(usize),
    GOIW(usize),
    Remove(usize),
    DidClear(usize),
}

#[derive(Clone, Debug)]
pub struct DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    map: HashMap<K, (usize, V), S>,
    journal: Vec<JournalEntry>,
    readonly: bool,
}

impl<K: Hash + Eq, V, S: BuildHasher> DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    #[inline(always)]
    pub fn inner(&self) -> &HashMap<K, (usize, V), S> {
        &self.map
    }

    #[inline(never)]
    pub fn begin_mutation(&mut self) {
        self.map.verify();
        assert!(self.readonly);
        self.readonly = false;
        self.verify();
    }

    #[inline(never)]
    pub fn end_mutation(&mut self) {
        self.map.verify();
        assert!(!self.readonly);
        self.readonly = true;
        self.verify();
    }

    fn verify(&self) {
        let mut position = 0;
        let mut count = 0;
        let mut bad_canary = None;
        for (_,v) in self.map.iter() {
            let canary_ref = &v.0;
            position += 1;
            if *canary_ref == CANARY {
                continue;
            }
            count += 1;
            bad_canary = Some((*canary_ref, canary_ref, position));
        }
        if let Some(c) = bad_canary {
            self.report_corruption(c.0, c.1, c.2, count);
        }
    }

    #[inline(always)]
    pub fn with_hasher(hash_builder: S) -> Self {
        Self {
            map: HashMap::<K, (usize, V), S>::with_hasher(hash_builder),
            journal: Vec::new(),
            readonly: true,
        }
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.map.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    #[inline(always)]
    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        self.map.contains_key(k)
    }

    #[inline(always)]
    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        self.map.get(k).map(|v| &v.1)
    }

    #[inline(always)]
    pub fn try_get_or_insert_with<F: FnOnce() -> V>(
        &mut self,
        key: K,
        default: F
    ) -> Result<&mut V, FailedAllocationError> {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::GOIW(self.map.make_hash(&key).inspect()));
        let entry = self.map.try_entry(key)?;
        Ok(&mut entry.or_insert_with(|| (CANARY, default())).1)
    }

    #[inline(always)]
    pub fn try_insert(&mut self, k: K, v: V) -> Result<Option<V>, FailedAllocationError> {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::Insert(self.map.make_hash(&k).inspect()));
        let old = self.map.try_insert(k, (CANARY, v))?;
        Ok(old.map(|x| x.1))
    }

    #[inline(always)]
    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>,
              Q: Hash + Eq
    {
        assert!(!self.readonly);
        self.journal.push(JournalEntry::Remove(self.map.make_hash(k).inspect()));
        if let Some(v) = self.map.get_mut(k) {
            unsafe { ptr::write_volatile(&mut v.0, POISON); }
        }
        self.map.remove(k).map(|x| x.1)
    }

    #[inline(always)]
    pub fn clear(&mut self) where K: 'static, V: 'static  {
        // We handle scoped mutations for the caller here, since callsites that
        // invoke clear() don't benefit from the coalescing we do around insertion.
        self.begin_mutation();
        self.journal.clear();
        self.journal.push(JournalEntry::DidClear(self.map.raw_capacity()));
        self.map.clear();
        self.end_mutation();
    }

    #[inline(never)]
    fn report_corruption(
        &self,
        canary: usize,
        canary_addr: *const usize,
        position: usize,
        count: usize,
    ) {
        use ::std::ffi::CString;
        let key = b"HashMapJournal\0";
        let value = CString::new(format!("{:?}", self.journal)).unwrap();
        unsafe {
            Gecko_AnnotateCrashReport(
                key.as_ptr() as *const ::std::os::raw::c_char,
                value.as_ptr(),
            );
        }
        panic!(
            concat!("HashMap Corruption (sz={}, cap={}, pairsz={}, cnry={:#x}, count={}, ",
                    "last_pos={}, base_addr={:?}, cnry_addr={:?}, jrnl_len={})"),
            self.map.len(),
            self.map.raw_capacity(),
            ::std::mem::size_of::<(K, (usize, V))>(),
            canary,
            count,
            position,
            self.map.raw_buffer(),
            canary_addr,
            self.journal.len(),
        );
    }
}

impl<K, V, S> PartialEq for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          V: PartialEq,
          S: BuildHasher
{
    fn eq(&self, other: &Self) -> bool {
        self.map.eq(&other.map)
    }
}

impl<K, V, S> Eq for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          V: Eq,
          S: BuildHasher
{
}

impl<K, V, S> Default for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher + Default
{
    fn default() -> Self {
        Self {
            map: HashMap::default(),
            journal: Vec::new(),
            readonly: true,
        }
    }
}

impl<K: Hash + Eq, V, S: BuildHasher> Drop for DiagnosticHashMap<K, V, S>
    where K: Eq + Hash,
          S: BuildHasher
{
    fn drop(&mut self) {
        self.map.verify();
        debug_assert!(self.readonly, "Dropped while mutating");
        self.verify();
    }
}

extern "C" {
    pub fn Gecko_AnnotateCrashReport(key_str: *const ::std::os::raw::c_char,
                                     value_str: *const ::std::os::raw::c_char);
}
<|endoftext|>


<commit_before>// # Custom collector example
//
// This example shows how you can implement your own
// collector. As an example, we will compute a collector
// that computes the standard deviation of a given fast field.
//
// Of course, you can have a look at the tantivy's built-in collectors
// such as the `CountCollector` for more examples.

// ---
// Importing tantivy...
use std::marker::PhantomData;

use crate::collector::{Collector, SegmentCollector};
use crate::fastfield::FastFieldReader;
use crate::schema::Field;
use crate::{Score, SegmentReader, TantivyError};

/// The `TopDocs` collector keeps track of the top `K` documents
/// sorted by their score.
///
/// The implementation is based on a `BinaryHeap`.
/// The theorical complexity for collecting the top `K` out of `n` documents
/// is `O(n log K)`.
///
/// This collector guarantees a stable sorting in case of a tie on the
/// document score. As such, it is suitable to implement pagination.
///
/// ```rust
/// use tantivy::collector::{TopDocs, FilterCollector};
/// use tantivy::query::QueryParser;
/// use tantivy::schema::{Schema, TEXT, INDEXED, FAST};
/// use tantivy::{doc, DocAddress, Index};
///
/// let mut schema_builder = Schema::builder();
/// let title = schema_builder.add_text_field("title", TEXT);
/// let price = schema_builder.add_u64_field("price", INDEXED | FAST);
/// let schema = schema_builder.build();
/// let index = Index::create_in_ram(schema);
///
/// let mut index_writer = index.writer_with_num_threads(1, 10_000_000).unwrap();
/// index_writer.add_document(doc!(title => "The Name of the Wind", price => 30_200u64));
/// index_writer.add_document(doc!(title => "The Diary of Muadib", price => 29_240u64));
/// index_writer.add_document(doc!(title => "A Dairy Cow", price => 21_240u64));
/// index_writer.add_document(doc!(title => "The Diary of a Young Girl", price => 20_120u64));
/// assert!(index_writer.commit().is_ok());
///
/// let reader = index.reader().unwrap();
/// let searcher = reader.searcher();
///
/// let query_parser = QueryParser::for_index(&index, vec![title]);
/// let query = query_parser.parse_query("diary").unwrap();
/// let no_filter_collector = FilterCollector::new(price, &|value| true, TopDocs::with_limit(2));
/// let top_docs = searcher.search(&query, &no_filter_collector).unwrap();
///
/// assert_eq!(top_docs.len(), 2);
/// assert_eq!(top_docs[0].1, DocAddress(0, 1));
/// assert_eq!(top_docs[1].1, DocAddress(0, 3));
///
/// let filter_all_collector = FilterCollector::new(price, &|value| false, TopDocs::with_limit(2));
/// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap();
///
/// assert_eq!(filtered_top_docs.len(), 0);
/// ```
pub struct FilterCollector<TCollector, TSegmentCollector> {
    field: Field,
    collector: TCollector,
    predicate: &'static (dyn Fn(u64) -> bool + Send + Sync),
    phantom: PhantomData<TSegmentCollector>
}

impl<TCollector, TSegmentCollector> FilterCollector<TCollector, TSegmentCollector>
where
    TCollector: Collector<Child = TSegmentCollector> + Send + Sync,
    TSegmentCollector: 'static + SegmentCollector + Send + Sync {
    pub fn new(field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector) -> FilterCollector<TCollector, TSegmentCollector> {
        FilterCollector { field, predicate, collector, phantom: PhantomData }
    }
}

impl<TCollector, TSegmentCollector> Collector for FilterCollector<TCollector, TSegmentCollector>
where
    TSegmentCollector: 'static + SegmentCollector + Send + Sync,
    TCollector: Collector<Child = TSegmentCollector> + Send + Sync {
    // That's the type of our result.
    // Our standard deviation will be a float.
    type Fruit = TCollector::Fruit;

    type Child = FilterSegmentCollector<TSegmentCollector>;

    fn for_segment(
        &self,
        segment_local_id: u32,
        segment_reader: &SegmentReader,
    ) -> crate::Result<FilterSegmentCollector<TSegmentCollector>> {
        let fast_field_reader = segment_reader
            .fast_fields()
            .u64(self.field)
            .ok_or_else(|| {
                let field_name = segment_reader.schema().get_field_name(self.field);
                TantivyError::SchemaError(format!(
                    "Field {:?} is not a u64 fast field.",
                    field_name
                ))
            })?;
        let child_segment_collector = self.collector.for_segment(segment_local_id, segment_reader);
        match child_segment_collector {
            Ok(segment_collector) => {
                Ok(FilterSegmentCollector::<TSegmentCollector> {
                    fast_field_reader,
                    segment_collector: segment_collector,
                    predicate: self.predicate
                })
            },
            Err(_) => {
                Err(TantivyError::SystemError("Could not open segment: ".to_owned()))
            },
        }

    }

    fn requires_scoring(&self) -> bool {
        self.collector.requires_scoring()
    }

    fn merge_fruits(&self, segment_fruits: Vec<<TCollector::Child as SegmentCollector>::Fruit>) -> crate::Result<TCollector::Fruit> {
        self.collector.merge_fruits(segment_fruits)
    }
}

pub struct FilterSegmentCollector<TSegmentCollector> {
    fast_field_reader: FastFieldReader<u64>,
    segment_collector: TSegmentCollector,
    predicate: &'static (dyn Fn(u64) -> bool + Send + Sync)
}

impl<TSegmentCollector> SegmentCollector for FilterSegmentCollector<TSegmentCollector>
where
    TSegmentCollector: 'static + SegmentCollector + Send + Sync {
    type Fruit = TSegmentCollector::Fruit;

    fn collect(&mut self, doc: u32, score: Score) {
        let value = self.fast_field_reader.get(doc);
        if (self.predicate)(value) {
            self.segment_collector.collect(doc, score)
        }
    }

    fn harvest(self) -> <TSegmentCollector as SegmentCollector>::Fruit {
        self.segment_collector.harvest()
    }
}
<commit_msg>removed unnecessary static liftimes<commit_after>// # Custom collector example
//
// This example shows how you can implement your own
// collector. As an example, we will compute a collector
// that computes the standard deviation of a given fast field.
//
// Of course, you can have a look at the tantivy's built-in collectors
// such as the `CountCollector` for more examples.

// ---
// Importing tantivy...
use std::marker::PhantomData;

use crate::collector::{Collector, SegmentCollector};
use crate::fastfield::FastFieldReader;
use crate::schema::Field;
use crate::{Score, SegmentReader, TantivyError};

/// The `TopDocs` collector keeps track of the top `K` documents
/// sorted by their score.
///
/// The implementation is based on a `BinaryHeap`.
/// The theorical complexity for collecting the top `K` out of `n` documents
/// is `O(n log K)`.
///
/// This collector guarantees a stable sorting in case of a tie on the
/// document score. As such, it is suitable to implement pagination.
///
/// ```rust
/// use tantivy::collector::{TopDocs, FilterCollector};
/// use tantivy::query::QueryParser;
/// use tantivy::schema::{Schema, TEXT, INDEXED, FAST};
/// use tantivy::{doc, DocAddress, Index};
///
/// let mut schema_builder = Schema::builder();
/// let title = schema_builder.add_text_field("title", TEXT);
/// let price = schema_builder.add_u64_field("price", INDEXED | FAST);
/// let schema = schema_builder.build();
/// let index = Index::create_in_ram(schema);
///
/// let mut index_writer = index.writer_with_num_threads(1, 10_000_000).unwrap();
/// index_writer.add_document(doc!(title => "The Name of the Wind", price => 30_200u64));
/// index_writer.add_document(doc!(title => "The Diary of Muadib", price => 29_240u64));
/// index_writer.add_document(doc!(title => "A Dairy Cow", price => 21_240u64));
/// index_writer.add_document(doc!(title => "The Diary of a Young Girl", price => 20_120u64));
/// assert!(index_writer.commit().is_ok());
///
/// let reader = index.reader().unwrap();
/// let searcher = reader.searcher();
///
/// let query_parser = QueryParser::for_index(&index, vec![title]);
/// let query = query_parser.parse_query("diary").unwrap();
/// let no_filter_collector = FilterCollector::new(price, &|value| true, TopDocs::with_limit(2));
/// let top_docs = searcher.search(&query, &no_filter_collector).unwrap();
///
/// assert_eq!(top_docs.len(), 2);
/// assert_eq!(top_docs[0].1, DocAddress(0, 1));
/// assert_eq!(top_docs[1].1, DocAddress(0, 3));
///
/// let filter_all_collector = FilterCollector::new(price, &|value| false, TopDocs::with_limit(2));
/// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap();
///
/// assert_eq!(filtered_top_docs.len(), 0);
/// ```
pub struct FilterCollector<TCollector, TSegmentCollector> {
    field: Field,
    collector: TCollector,
    predicate: &'static (dyn Fn(u64) -> bool + Send + Sync),
    phantom: PhantomData<TSegmentCollector>
}

impl<TCollector, TSegmentCollector> FilterCollector<TCollector, TSegmentCollector>
where
    TCollector: Collector<Child = TSegmentCollector> + Send + Sync,
    TSegmentCollector: SegmentCollector + Send + Sync {
    pub fn new(field: Field, predicate: &'static (dyn Fn(u64) -> bool + Send + Sync), collector: TCollector) -> FilterCollector<TCollector, TSegmentCollector> {
        FilterCollector { field, predicate, collector, phantom: PhantomData }
    }
}

impl<TCollector, TSegmentCollector> Collector for FilterCollector<TCollector, TSegmentCollector>
where
    TSegmentCollector: SegmentCollector + Send + Sync,
    TCollector: Collector<Child = TSegmentCollector> + Send + Sync {
    // That's the type of our result.
    // Our standard deviation will be a float.
    type Fruit = TCollector::Fruit;

    type Child = FilterSegmentCollector<TSegmentCollector>;

    fn for_segment(
        &self,
        segment_local_id: u32,
        segment_reader: &SegmentReader,
    ) -> crate::Result<FilterSegmentCollector<TSegmentCollector>> {
        let fast_field_reader = segment_reader
            .fast_fields()
            .u64(self.field)
            .ok_or_else(|| {
                let field_name = segment_reader.schema().get_field_name(self.field);
                TantivyError::SchemaError(format!(
                    "Field {:?} is not a u64 fast field.",
                    field_name
                ))
            })?;
        let child_segment_collector = self.collector.for_segment(segment_local_id, segment_reader);
        match child_segment_collector {
            Ok(segment_collector) => {
                Ok(FilterSegmentCollector::<TSegmentCollector> {
                    fast_field_reader,
                    segment_collector: segment_collector,
                    predicate: self.predicate
                })
            },
            Err(_) => {
                Err(TantivyError::SystemError("Could not open segment: ".to_owned()))
            },
        }

    }

    fn requires_scoring(&self) -> bool {
        self.collector.requires_scoring()
    }

    fn merge_fruits(&self, segment_fruits: Vec<<TCollector::Child as SegmentCollector>::Fruit>) -> crate::Result<TCollector::Fruit> {
        self.collector.merge_fruits(segment_fruits)
    }
}

pub struct FilterSegmentCollector<TSegmentCollector> {
    fast_field_reader: FastFieldReader<u64>,
    segment_collector: TSegmentCollector,
    predicate: &'static (dyn Fn(u64) -> bool + Send + Sync)
}

impl<TSegmentCollector> SegmentCollector for FilterSegmentCollector<TSegmentCollector>
where
    TSegmentCollector: SegmentCollector + Send + Sync {
    type Fruit = TSegmentCollector::Fruit;

    fn collect(&mut self, doc: u32, score: Score) {
        let value = self.fast_field_reader.get(doc);
        if (self.predicate)(value) {
            self.segment_collector.collect(doc, score)
        }
    }

    fn harvest(self) -> <TSegmentCollector as SegmentCollector>::Fruit {
        self.segment_collector.harvest()
    }
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    sync::Arc,
    time::Duration,
};

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote: upstream::Upstream> {
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote: upstream::Upstream> DnsRelay<Remote> {
    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;
        // Start querying name servers
        debug!(
            "attempting lookup of {:?} {} with ns {:?} and {:?}",
            query.query_type(), query.name(), local, remote
        );

        let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
        let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                debug!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                debug!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        debug!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        debug!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        debug!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    debug!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                        RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, tx) = socket.split();
    let tx = Arc::new(tokio::sync::Mutex::new(tx));

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let relay = Arc::new(DnsRelay {
        context: context.clone(),
        remote_upstream: upstream::ProxyTcpUpstream {
            context: context.clone(),
            svr_cfg: move || balancer.pick_server().server_config().clone(),
            ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
        }
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let relay = relay.clone();
        let tx = tx.clone();

        tokio::spawn(async move {
            let message = relay.resolve(request).await;
            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = tx.lock().await.send_to(&res_buffer, &src).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Split UDP logic into run_udp<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote: upstream::Upstream> {
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote: upstream::Upstream> DnsRelay<Remote> {
    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;
        // Start querying name servers
        debug!(
            "attempting lookup of {:?} {} with ns {:?} and {:?}",
            query.query_type(), query.name(), local, remote
        );

        let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
        let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                debug!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                debug!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        debug!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        debug!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        debug!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    debug!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                        RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

async fn run_udp<Remote: upstream::Upstream + Send + Sync + 'static>(
    relay: Arc<DnsRelay<Remote>>,
    bind_addr: SocketAddr
) -> io::Result<()> {
    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, tx) = socket.split();
    let tx = Arc::new(tokio::sync::Mutex::new(tx));

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let relay = relay.clone();
        let tx = tx.clone();

        tokio::spawn(async move {
            let message = relay.resolve(request).await;
            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = tx.lock().await.send_to(&res_buffer, &src).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };
    let bind_addr = local_addr.bind_addr(&context).await?;

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let relay = Arc::new(DnsRelay {
        context: context.clone(),
        remote_upstream: upstream::ProxyTcpUpstream {
            context: context.clone(),
            svr_cfg: move || balancer.pick_server().server_config().clone(),
            ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
        }
    });

    tokio::spawn(run_udp(relay, bind_addr.clone())).await?
}
<|endoftext|>


<commit_before>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::{MutableAntichain, Antichain};
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::count_map::CountMap;
use progress::timestamp::RootTimestamp;
use progress::nested::product::Product;

use timely_communication::Allocate;
use {Data, Push};
use dataflow::channels::Content;
use dataflow::channels::pushers::{Tee, Counter as PushCounter};
use dataflow::channels::pushers::buffer::{Buffer as PushBuffer, AutoflushSession};

use dataflow::operators::Capability;
use dataflow::operators::capability::mint as mint_capability;

use dataflow::{Stream, Scope};
use dataflow::scopes::{Child, Root};

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait UnorderedInput<G: Scope> {
    /// Create a new capability-based `Stream` and `Handle` through which to supply input. This
    /// input supports multiple open epochs (timestamps) at the same time.
    ///
    /// The `new_unordered_input` method returns `((Handle, Capability), Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, `Handle` and `Capability` are later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Capability` returned is for the default value of the timestamp type in use. The
    /// capability can be dropped to inform the system that the input has advanced beyond the
    /// capability's timestamp. To retain the ability to send, a new capability at a later timestamp
    /// should be obtained first, via the `delayed` function for `Capability`.
    ///
    /// To communicate the end-of-input drop all available capabilities.
    fn new_unordered_input<D:Data>(&mut self) -> ((UnorderedHandle<G, D>, Capability<G::Timestamp>), Stream<G, D>);
}

impl<G: Scope> UnorderedInput<G> for G {
    fn new_unordered_input<D:Data>(&mut self) -> ((UnorderedHandle<G, D>, Capability<G::Timestamp>), Stream<G, D>) {

        let (output, registrar) = Tee::<G::Timestamp, D>::new();
        let internal = Rc::new(RefCell::new(CountMap::new()));
        let produced = Rc::new(RefCell::new(CountMap::new()));
        let cap = mint_capability(Default::default(), internal.clone());
        let helper = UnorderedHandle::new(PushCounter::new(output, produced.clone()));
        let peers = self.peers();

        let index = self.add_operator(UnorderedOperator {
            internal: internal.clone(),
            produced: produced.clone(),
            peers:    peers,
        });

        return ((helper, cap), Stream::new(Source { index: index, port: 0 }, registrar, self.clone()));
    }
}

struct UnorderedOperator<T:Timestamp> {
    internal:   Rc<RefCell<CountMap<T>>>,
    produced:   Rc<RefCell<CountMap<T>>>,
    peers:     usize,
}

impl<T:Timestamp> Operate<T> for UnorderedOperator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<T as Timestamp>::Summary>>>,
                                           Vec<CountMap<T>>) {
        let mut internal = CountMap::new();
        // augment the counts for each reserved capability.
        for &(ref time, count) in self.internal.borrow().iter() {
            internal.update(time, count * (self.peers as i64 - 1));
        }

        // drain the changes to empty out, and complete the counts for internal.
        self.internal.borrow_mut().drain_into(&mut internal);
        (Vec::new(), vec![internal])
    }

    fn pull_internal_progress(&mut self,_consumed: &mut [CountMap<T>],
                                         internal: &mut [CountMap<T>],
                                         produced: &mut [CountMap<T>]) -> bool
    {
        self.produced.borrow_mut().drain_into(&mut produced[0]);
        self.internal.borrow_mut().drain_into(&mut internal[0]);
        return false;
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct UnorderedHandle<G: Scope, D: Data> {
    buffer: PushBuffer<G::Timestamp, D, PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>>,
}

impl<G: Scope, D: Data> UnorderedHandle<G, D> {
    fn new(pusher: PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>) -> UnorderedHandle<G, D> {
        UnorderedHandle {
            buffer: PushBuffer::new(pusher),
        }
    }

    pub fn session<'b>(&'b mut self, cap: Capability<G::Timestamp>) -> AutoflushSession<'b, G::Timestamp, D, PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>> {
        self.buffer.autoflush_session(cap)
    }
}

impl<G: Scope, D: Data> Drop for UnorderedHandle<G, D> {
    fn drop(&mut self) {
        // TODO: explode if not all capabilities were given up?
    }
}
<commit_msg>added example / test<commit_after>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::Antichain;
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::count_map::CountMap;

use timely_communication::Allocate;
use Data;
use dataflow::channels::pushers::{Tee, Counter as PushCounter};
use dataflow::channels::pushers::buffer::{Buffer as PushBuffer, AutoflushSession};

use dataflow::operators::Capability;
use dataflow::operators::capability::mint as mint_capability;

use dataflow::{Stream, Scope};

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait UnorderedInput<G: Scope> {
    /// Create a new capability-based `Stream` and `Handle` through which to supply input. This
    /// input supports multiple open epochs (timestamps) at the same time.
    ///
    /// The `new_unordered_input` method returns `((Handle, Capability), Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, `Handle` and `Capability` are later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Capability` returned is for the default value of the timestamp type in use. The
    /// capability can be dropped to inform the system that the input has advanced beyond the
    /// capability's timestamp. To retain the ability to send, a new capability at a later timestamp
    /// should be obtained first, via the `delayed` function for `Capability`.
    ///
    /// To communicate the end-of-input drop all available capabilities.
    ///
    /// #Examples
    ///
    /// ```
    /// use std::sync::{Arc, Mutex};
    ///
    /// use timely::*;
    /// use timely::dataflow::operators::*;
    /// use timely::dataflow::operators::capture::Extract;
    /// use timely::dataflow::{Stream, Scope};
    /// use timely::progress::timestamp::RootTimestamp;
    ///
    /// // get send and recv endpoints, wrap send to share
    /// let (send, recv) = ::std::sync::mpsc::channel();
    /// let send = Arc::new(Mutex::new(send));
    ///
    /// timely::execute(Configuration::Thread, move |root| {
    ///
    ///     // this is only to validate the output.
    ///     let send = send.lock().unwrap().clone();
    ///
    ///     // create and capture the unordered input.
    ///     let (mut input, mut cap) = root.scoped(|scope| {
    ///         let (input, stream) = scope.new_unordered_input();
    ///         stream.capture_into(send);
    ///         input
    ///     });
    ///
    ///     // feed values 0..10 at times 0..10.
    ///     for round in 0..10 {
    ///         input.session(cap.clone()).give(round);
    ///         cap = cap.delayed(&RootTimestamp::new(round + 1));
    ///         root.step();
    ///     }
    /// }).unwrap();
    ///
    /// let extract = recv.extract();
    /// for i in 0..10 {
    ///     assert_eq!(extract[i], (RootTimestamp::new(i), vec![i]));
    /// }
    /// ```
    fn new_unordered_input<D:Data>(&mut self) -> ((UnorderedHandle<G, D>, Capability<G::Timestamp>), Stream<G, D>);
}

impl<G: Scope> UnorderedInput<G> for G {
    fn new_unordered_input<D:Data>(&mut self) -> ((UnorderedHandle<G, D>, Capability<G::Timestamp>), Stream<G, D>) {

        let (output, registrar) = Tee::<G::Timestamp, D>::new();
        let internal = Rc::new(RefCell::new(CountMap::new()));
        let produced = Rc::new(RefCell::new(CountMap::new()));
        let cap = mint_capability(Default::default(), internal.clone());
        let helper = UnorderedHandle::new(PushCounter::new(output, produced.clone()));
        let peers = self.peers();

        let index = self.add_operator(UnorderedOperator {
            internal: internal.clone(),
            produced: produced.clone(),
            peers:    peers,
        });

        return ((helper, cap), Stream::new(Source { index: index, port: 0 }, registrar, self.clone()));
    }
}

struct UnorderedOperator<T:Timestamp> {
    internal:   Rc<RefCell<CountMap<T>>>,
    produced:   Rc<RefCell<CountMap<T>>>,
    peers:     usize,
}

impl<T:Timestamp> Operate<T> for UnorderedOperator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<T as Timestamp>::Summary>>>,
                                           Vec<CountMap<T>>) {
        let mut internal = CountMap::new();
        // augment the counts for each reserved capability.
        for &(ref time, count) in self.internal.borrow().iter() {
            internal.update(time, count * (self.peers as i64 - 1));
        }

        // drain the changes to empty out, and complete the counts for internal.
        self.internal.borrow_mut().drain_into(&mut internal);
        (Vec::new(), vec![internal])
    }

    fn pull_internal_progress(&mut self,_consumed: &mut [CountMap<T>],
                                         internal: &mut [CountMap<T>],
                                         produced: &mut [CountMap<T>]) -> bool
    {
        self.produced.borrow_mut().drain_into(&mut produced[0]);
        self.internal.borrow_mut().drain_into(&mut internal[0]);
        return false;
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct UnorderedHandle<G: Scope, D: Data> {
    buffer: PushBuffer<G::Timestamp, D, PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>>,
}

impl<G: Scope, D: Data> UnorderedHandle<G, D> {
    fn new(pusher: PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>) -> UnorderedHandle<G, D> {
        UnorderedHandle {
            buffer: PushBuffer::new(pusher),
        }
    }

    pub fn session<'b>(&'b mut self, cap: Capability<G::Timestamp>) -> AutoflushSession<'b, G::Timestamp, D, PushCounter<G::Timestamp, D, Tee<G::Timestamp, D>>> {
        self.buffer.autoflush_session(cap)
    }
}

impl<G: Scope, D: Data> Drop for UnorderedHandle<G, D> {
    fn drop(&mut self) {
        // TODO: explode if not all capabilities were given up?
    }
}
<|endoftext|>


<commit_before>use byteorder::{
    LittleEndian,
    WriteBytesExt
};

use std::io::{
    Cursor
};

use slog::{
    Discard,
    Logger
};

// TODO also generate bgzf blocks with extra gzip subfields
fn bgzf_block(deflated_payload: Vec<u8>, inflated_payload_size: u32, inflated_payload_crc32: u32, deflated_payload_size_delta: i32) -> Vec<u8> {
    let mut data = vec![
        0x1f, 0x8b,             // gzip identifier
        0x08,                   // method (deflate)
        0x04,                   // flags (FEXTRA)
        0x00, 0x00, 0x00, 0x00, // modification time
        0x00,                   // extra flags
        0xff,                   // operating system (unknown)
        0x06, 0x00,             // extra field length (6 bytes)
        0x42, 0x43,             // bgzf identifier
        0x02, 0x00,             // extra subfield length (2 bytes)
    ];

    data.write_u16::<LittleEndian>((deflated_payload.len() as i32 + 25i32 + deflated_payload_size_delta) as u16).unwrap(); // bgzf block size, minus 1

    data.extend(deflated_payload);
    data.write_u32::<LittleEndian>(inflated_payload_crc32).unwrap();
    data.write_u32::<LittleEndian>(inflated_payload_size).unwrap();

    data
}

pub fn empty_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0x03, 0x00 // deflated empty string
    ], 0, 0, 0)
}

pub fn regular_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 5, 907060870, 0)
}

pub fn bad_inflated_payload_crc32_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0x25, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 5, 907060870, 0)
}

pub fn bad_inflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, 0)
}

pub fn too_small_deflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, -5i32)
}

pub fn too_large_deflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, -5i32)
}

pub fn null_logger() -> Logger {
    Logger::root(Discard, o!())
}

pub fn empty_file() -> Cursor<Vec<u8>> {
    Cursor::new(vec![])
}

pub fn empty_bam() -> Cursor<Vec<u8>> {
    Cursor::new(empty_bgzf_block())
}

pub fn single_block_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn two_blocks_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_empty_inside_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn single_block_missing_empty_bam() -> Cursor<Vec<u8>> {
    Cursor::new(regular_bgzf_block())
}

pub fn single_block_missing_gzip_identifier_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    data[0] = 42;
    Cursor::new(data)
}

pub fn single_block_missing_bgzf_identifier_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    data[12] = 21;
    Cursor::new(data)
}

pub fn two_blocks_missing_empty_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bad_inflated_payload_crc32_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(bad_inflated_payload_crc32_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bad_inflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(bad_inflated_payload_size_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_too_small_deflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(too_small_deflated_payload_size_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_too_large_deflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(too_large_deflated_payload_size_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}
<commit_msg>Use the advertised number of blocks in tests<commit_after>use byteorder::{
    LittleEndian,
    WriteBytesExt
};

use std::io::{
    Cursor
};

use slog::{
    Discard,
    Logger
};

// TODO also generate bgzf blocks with extra gzip subfields
fn bgzf_block(deflated_payload: Vec<u8>, inflated_payload_size: u32, inflated_payload_crc32: u32, deflated_payload_size_delta: i32) -> Vec<u8> {
    let mut data = vec![
        0x1f, 0x8b,             // gzip identifier
        0x08,                   // method (deflate)
        0x04,                   // flags (FEXTRA)
        0x00, 0x00, 0x00, 0x00, // modification time
        0x00,                   // extra flags
        0xff,                   // operating system (unknown)
        0x06, 0x00,             // extra field length (6 bytes)
        0x42, 0x43,             // bgzf identifier
        0x02, 0x00,             // extra subfield length (2 bytes)
    ];

    data.write_u16::<LittleEndian>((deflated_payload.len() as i32 + 25i32 + deflated_payload_size_delta) as u16).unwrap(); // bgzf block size, minus 1

    data.extend(deflated_payload);
    data.write_u32::<LittleEndian>(inflated_payload_crc32).unwrap();
    data.write_u32::<LittleEndian>(inflated_payload_size).unwrap();

    data
}

pub fn empty_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0x03, 0x00 // deflated empty string
    ], 0, 0, 0)
}

pub fn regular_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 5, 907060870, 0)
}

pub fn bad_inflated_payload_crc32_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0x25, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 5, 907060870, 0)
}

pub fn bad_inflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, 0)
}

pub fn too_small_deflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, -5i32)
}

pub fn too_large_deflated_payload_size_bgzf_block() -> Vec<u8> {
    bgzf_block(vec![
        0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00 // deflated "hello"
    ], 25, 907060870, -5i32)
}

pub fn null_logger() -> Logger {
    Logger::root(Discard, o!())
}

pub fn empty_file() -> Cursor<Vec<u8>> {
    Cursor::new(vec![])
}

pub fn empty_bam() -> Cursor<Vec<u8>> {
    Cursor::new(empty_bgzf_block())
}

pub fn single_block_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn two_blocks_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_empty_inside_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn single_block_missing_empty_bam() -> Cursor<Vec<u8>> {
    Cursor::new(regular_bgzf_block())
}

pub fn single_block_missing_gzip_identifier_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    data[0] = 42;
    Cursor::new(data)
}

pub fn single_block_missing_bgzf_identifier_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(empty_bgzf_block());
    data[12] = 21;
    Cursor::new(data)
}

pub fn two_blocks_missing_empty_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(regular_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bad_inflated_payload_crc32_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(bad_inflated_payload_crc32_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_bad_inflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = regular_bgzf_block();
    data.extend(bad_inflated_payload_size_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_too_small_deflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = too_small_deflated_payload_size_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}

pub fn three_blocks_too_large_deflated_payload_size_bam() -> Cursor<Vec<u8>> {
    let mut data = too_large_deflated_payload_size_bgzf_block();
    data.extend(regular_bgzf_block());
    data.extend(regular_bgzf_block());
    data.extend(empty_bgzf_block());
    Cursor::new(data)
}
<|endoftext|>


<commit_before>use sea_canal::{Analyze, Analyzer};
use sea_canal::pattern::{CustomPatternElem, Pattern};
use sea_canal::pattern::PatternElem::*;

fn pow4(i: i32, j: i32) -> bool {
    i * i * i * i == j
}

fn root4(i: i32, j: i32) -> bool {
    j * j * j * j == i
}

#[test]
fn find_any_pattern() {
    let pow4_pattern = CustomPatternElem::new(pow4, "^4");
    let root4_pattern = CustomPatternElem::new(root4, "root 4");

    let slice = &[1, 2, 16, 2, 3, 81, 3];
    let analyzer = Analyzer::with_custom_patterns(slice, vec![pow4_pattern.clone(), root4_pattern.clone()]);

    assert_eq!(None, analyzer.find_any_pattern(1));
    assert_eq!(Some(pat![Plus(1), Custom(pow4_pattern), Custom(root4_pattern)]), analyzer.find_any_pattern(4));
}
<commit_msg>remove unnecessary assertion<commit_after>use sea_canal::{Analyze, Analyzer};
use sea_canal::pattern::{CustomPatternElem, Pattern};
use sea_canal::pattern::PatternElem::*;

fn pow4(i: i32, j: i32) -> bool {
    i * i * i * i == j
}

fn root4(i: i32, j: i32) -> bool {
    j * j * j * j == i
}

#[test]
fn find_any_pattern() {
    let pow4_pattern = CustomPatternElem::new(pow4, "^4");
    let root4_pattern = CustomPatternElem::new(root4, "root 4");

    let slice = &[1, 2, 16, 2, 3, 81, 3];
    let analyzer = Analyzer::with_custom_patterns(slice, vec![pow4_pattern.clone(), root4_pattern.clone()]);

    assert_eq!(Some(pat![Plus(1), Custom(pow4_pattern), Custom(root4_pattern)]), analyzer.find_any_pattern(4));
}
<|endoftext|>