<commit_before>#![allow(dead_code)]

extern crate tempdir;

use std::env;
use std::fs::{self, File};
use std::io::{Read, Write, Result};
#[cfg(unix)]
use std::os::unix::fs::symlink as symlink_file;
#[cfg(windows)]
use std::os::windows::fs::symlink_file;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::str::from_utf8;
use std::ffi::OsStr;
use self::tempdir::TempDir;
use std::rc::Rc;

#[cfg(windows)]
static PROGNAME: &'static str = "target\\debug\\uutils.exe";
#[cfg(windows)]
static FIXTURES_DIR: &'static str = "tests\\fixtures";
#[cfg(not(windows))]
static PROGNAME: &'static str = "target/debug/uutils";
#[cfg(not(windows))]
static FIXTURES_DIR: &'static str = "tests/fixtures";
static ALREADY_RUN: &'static str = " you have already run this UCommand, if you want to run \
                                    another command in the same test, use TestSet::new instead of \
                                    testing();";

#[macro_export]
macro_rules! assert_empty_stderr(
    ($cond:expr) => (
        if $cond.stderr.len() > 0 {
            panic!(format!("stderr: {}", $cond.stderr))
        }
    );
);

#[macro_export]
macro_rules! assert_no_error(
    ($cond:expr) => (
        assert!($cond.success);
        if $cond.stderr.len() > 0 {
            panic!(format!("stderr: {}", $cond.stderr))
        }
    );
);

pub struct CmdResult {
    pub success: bool,
    pub stdout: String,
    pub stderr: String,
}

pub fn log_info<T: AsRef<str>, U: AsRef<str>>(msg: T, par: U) {
    println!("{}: {}", msg.as_ref(), par.as_ref());
}

pub fn repeat_str(s: &str, n: u32) -> String {
    let mut repeated = String::new();
    for _ in 0..n {
        repeated.push_str(s);
    }
    repeated
}

pub fn recursive_copy(src: &Path, dest: &Path) -> Result<()> {
    if try!(fs::metadata(src)).is_dir() {
        for entry in try!(fs::read_dir(src)) {
            let entry = try!(entry);
            let mut new_dest = PathBuf::from(dest);
            new_dest.push(entry.file_name());
            if try!(fs::metadata(entry.path())).is_dir() {
                try!(fs::create_dir(&new_dest));
                try!(recursive_copy(&entry.path(), &new_dest));
            } else {
                try!(fs::copy(&entry.path(), new_dest));
            }
        }
    }
    Ok(())
}

pub struct AtPath {
    pub subdir: PathBuf,
}
impl AtPath {
    pub fn new(subdir: &Path) -> AtPath {
        AtPath { subdir: PathBuf::from(subdir) }
    }
    pub fn as_string(&self) -> String {
        self.subdir.to_str().unwrap().to_owned()
    }
    pub fn plus(&self, name: &str) -> PathBuf {
        let mut pathbuf = self.subdir.clone();
        pathbuf.push(name);
        pathbuf
    }
    pub fn plus_as_string(&self, name: &str) -> String {
        String::from(self.plus(name).to_str().unwrap())
    }
    fn minus(&self, name: &str) -> PathBuf {
        // relative_from is currently unstable
        let prefixed = PathBuf::from(name);
        if prefixed.starts_with(&self.subdir) {
            let mut unprefixed = PathBuf::new();
            for component in prefixed.components()
                                     .skip(self.subdir.components().count()) {
                unprefixed.push(component.as_ref().to_str().unwrap());
            }
            unprefixed
        } else {
            prefixed
        }
    }
    pub fn minus_as_string(&self, name: &str) -> String {
        String::from(self.minus(name).to_str().unwrap())
    }
    pub fn open(&self, name: &str) -> File {
        log_info("open", self.plus_as_string(name));
        File::open(self.plus(name)).unwrap()
    }
    pub fn read(&self, name: &str) -> String {
        let mut f = self.open(name);
        let mut contents = String::new();
        let _ = f.read_to_string(&mut contents);
        contents
    }
    pub fn write(&self, name: &str, contents: &str) {
        let mut f = self.open(name);
        let _ = f.write(contents.as_bytes());
    }
    pub fn mkdir(&self, dir: &str) {
        log_info("mkdir", self.plus_as_string(dir));
        fs::create_dir(&self.plus(dir)).unwrap();
    }
    pub fn mkdir_all(&self, dir: &str) {
        log_info("mkdir_all", self.plus_as_string(dir));
        fs::create_dir_all(self.plus(dir)).unwrap();
    }
    pub fn make_file(&self, name: &str) -> File {
        match File::create(&self.plus(name)) {
            Ok(f) => f,
            Err(e) => panic!("{}", e),
        }
    }
    pub fn touch(&self, file: &str) {
        log_info("touch", self.plus_as_string(file));
        File::create(&self.plus(file)).unwrap();
    }
    pub fn symlink(&self, src: &str, dst: &str) {
        log_info("symlink",
                 &format!("{},{}", self.plus_as_string(src), self.plus_as_string(dst)));
        symlink_file(&self.plus(src), &self.plus(dst)).unwrap();
    }
    pub fn is_symlink(&self, path: &str) -> bool {
        log_info("is_symlink", self.plus_as_string(path));
        match fs::symlink_metadata(&self.plus(path)) {
            Ok(m) => m.file_type().is_symlink(),
            Err(_) => false,
        }
    }

    pub fn resolve_link(&self, path: &str) -> String {
        log_info("resolve_link", self.plus_as_string(path));
        match fs::read_link(&self.plus(path)) {
            Ok(p) => {
                self.minus_as_string(p.to_str().unwrap())
            }
            Err(_) => "".to_string(),
        }
    }

    pub fn metadata(&self, path: &str) -> fs::Metadata {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m,
            Err(e) => panic!("{}", e),
        }
    }

    pub fn file_exists(&self, path: &str) -> bool {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m.is_file(),
            Err(_) => false,
        }
    }

    pub fn dir_exists(&self, path: &str) -> bool {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m.is_dir(),
            Err(_) => false,
        }
    }

    pub fn cleanup(&self, path: &'static str) {
        let p = &self.plus(path);
        match fs::metadata(p) {
            Ok(m) => if m.is_file() {
                fs::remove_file(&p).unwrap();
            } else {
                fs::remove_dir(&p).unwrap();
            },
            Err(_) => {}
        }
    }
    pub fn root_dir(&self) -> String {
        log_info("current_directory", "");
        self.subdir.to_str().unwrap().to_owned()
    }
}

pub struct TestSet {
    bin_path: PathBuf,
    util_name: String,
    pub fixtures: AtPath,
    tmpd: Rc<TempDir>,
}
impl TestSet {
    pub fn new(util_name: &str) -> TestSet {
        let tmpd = Rc::new(TempDir::new("uutils").unwrap());
        let ts = TestSet {
            bin_path: {
                let mut bin_path_builder = env::current_dir().unwrap();
                bin_path_builder.push(PathBuf::from(PROGNAME));
                bin_path_builder
            },
            util_name: String::from(util_name),
            fixtures: AtPath::new(&tmpd.as_ref().path()),
            tmpd: tmpd,
        };
        let mut fixture_path_builder = env::current_dir().unwrap();
        fixture_path_builder.push(PathBuf::from(FIXTURES_DIR));
        fixture_path_builder.push(PathBuf::from(util_name));
        match fs::metadata(&fixture_path_builder) {
            Ok(m) => if m.is_dir() {
                recursive_copy(&fixture_path_builder, &ts.fixtures.subdir).unwrap();
            },
            Err(_) => {}
        }
        ts
    }
    pub fn util_cmd(&self) -> UCommand {
        let mut cmd = self.cmd(&self.bin_path);
        cmd.arg(&self.util_name);
        cmd
    }
    pub fn cmd<S: AsRef<OsStr>>(&self, bin: S) -> UCommand {
        UCommand::new_from_tmp(bin, self.tmpd.clone(), true)
    }
    // different names are used rather than an argument
    // because the need to keep the environment is exceedingly rare.
    pub fn util_cmd_keepenv(&self) -> UCommand {
        let mut cmd = self.cmd_keepenv(&self.bin_path);
        cmd.arg(&self.util_name);
        cmd
    }
    pub fn cmd_keepenv<S: AsRef<OsStr>>(&self, bin: S) -> UCommand {
        UCommand::new_from_tmp(bin, self.tmpd.clone(), false)
    }
}

pub struct UCommand {
    pub raw: Command,
    comm_string: String,
    tmpd: Option<Rc<TempDir>>,
    has_run: bool,
}
impl UCommand {
    pub fn new<T: AsRef<OsStr>, U: AsRef<OsStr>>(arg: T, curdir: U, env_clear: bool) -> UCommand {
        UCommand {
            tmpd: None,
            has_run: false,
            raw: {
                let mut cmd = Command::new(arg.as_ref());
                cmd.current_dir(curdir.as_ref());
                if env_clear {
                    cmd.env_clear();
                }
                cmd
            },
            comm_string: String::from(arg.as_ref().to_str().unwrap()),
        }
    }
    pub fn new_from_tmp<T: AsRef<OsStr>>(arg: T, tmpd: Rc<TempDir>, env_clear: bool) -> UCommand {
        let tmpd_path_buf = String::from(&(*tmpd.as_ref().path().to_str().unwrap()));
        let mut ucmd: UCommand = UCommand::new(arg.as_ref(), tmpd_path_buf, env_clear);
        ucmd.tmpd = Some(tmpd);
        ucmd
    }
    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> Box<&mut UCommand> {
        if self.has_run {
            panic!(ALREADY_RUN);
        }
        self.comm_string.push_str(" ");
        self.comm_string.push_str(arg.as_ref().to_str().unwrap());
        self.raw.arg(arg.as_ref());
        Box::new(self)
    }

    pub fn args<S: AsRef<OsStr>>(&mut self, args: &[S]) -> Box<&mut UCommand> {
        if self.has_run {
            panic!(ALREADY_RUN);
        }
        for s in args {
            self.comm_string.push_str(" ");
            self.comm_string.push_str(s.as_ref().to_str().unwrap());
        }

        self.raw.args(args.as_ref());
        Box::new(self)
    }
    pub fn run(&mut self) -> CmdResult {
        self.has_run = true;
        log_info("run", &self.comm_string);
        let prog = self.raw.output().unwrap();
        CmdResult {
            success: prog.status.success(),
            stdout: from_utf8(&prog.stdout).unwrap().to_string(),
            stderr: from_utf8(&prog.stderr).unwrap().to_string(),
        }
    }
    pub fn run_piped_stdin<T: AsRef<[u8]>>(&mut self, input: T) -> CmdResult {
        self.has_run = true;
        log_info("run_piped_stdin", &self.comm_string);
        let mut result = self.raw
                             .stdin(Stdio::piped())
                             .stdout(Stdio::piped())
                             .stderr(Stdio::piped())
                             .spawn()
                             .unwrap();

        result.stdin
              .take()
              .unwrap_or_else(|| panic!("Could not take child process stdin"))
              .write_all(input.as_ref())
              .unwrap_or_else(|e| panic!("{}", e));

        let prog = result.wait_with_output().unwrap();
        CmdResult {
            success: prog.status.success(),
            stdout: from_utf8(&prog.stdout).unwrap().to_string(),
            stderr: from_utf8(&prog.stderr).unwrap().to_string(),
        }
    }
}

// returns a testSet and a ucommand initialized to the utility binary
// operating in the fixtures directory with a cleared environment
pub fn testset_and_ucommand(utilname: &str) -> (TestSet, UCommand) {
    let ts = TestSet::new(utilname);
    let ucmd = ts.util_cmd();
    (ts, ucmd)
}
pub fn testing(utilname: &str) -> (AtPath, UCommand) {
    let ts = TestSet::new(utilname);
    let ucmd = ts.util_cmd();
    (ts.fixtures, ucmd)
}
<commit_msg>tests: add UCommand::env()<commit_after>#![allow(dead_code)]

extern crate tempdir;

use std::env;
use std::fs::{self, File};
use std::io::{Read, Write, Result};
#[cfg(unix)]
use std::os::unix::fs::symlink as symlink_file;
#[cfg(windows)]
use std::os::windows::fs::symlink_file;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::str::from_utf8;
use std::ffi::OsStr;
use self::tempdir::TempDir;
use std::rc::Rc;

#[cfg(windows)]
static PROGNAME: &'static str = "target\\debug\\uutils.exe";
#[cfg(windows)]
static FIXTURES_DIR: &'static str = "tests\\fixtures";
#[cfg(not(windows))]
static PROGNAME: &'static str = "target/debug/uutils";
#[cfg(not(windows))]
static FIXTURES_DIR: &'static str = "tests/fixtures";
static ALREADY_RUN: &'static str = " you have already run this UCommand, if you want to run \
                                    another command in the same test, use TestSet::new instead of \
                                    testing();";

#[macro_export]
macro_rules! assert_empty_stderr(
    ($cond:expr) => (
        if $cond.stderr.len() > 0 {
            panic!(format!("stderr: {}", $cond.stderr))
        }
    );
);

#[macro_export]
macro_rules! assert_no_error(
    ($cond:expr) => (
        assert!($cond.success);
        if $cond.stderr.len() > 0 {
            panic!(format!("stderr: {}", $cond.stderr))
        }
    );
);

pub struct CmdResult {
    pub success: bool,
    pub stdout: String,
    pub stderr: String,
}

pub fn log_info<T: AsRef<str>, U: AsRef<str>>(msg: T, par: U) {
    println!("{}: {}", msg.as_ref(), par.as_ref());
}

pub fn repeat_str(s: &str, n: u32) -> String {
    let mut repeated = String::new();
    for _ in 0..n {
        repeated.push_str(s);
    }
    repeated
}

pub fn recursive_copy(src: &Path, dest: &Path) -> Result<()> {
    if try!(fs::metadata(src)).is_dir() {
        for entry in try!(fs::read_dir(src)) {
            let entry = try!(entry);
            let mut new_dest = PathBuf::from(dest);
            new_dest.push(entry.file_name());
            if try!(fs::metadata(entry.path())).is_dir() {
                try!(fs::create_dir(&new_dest));
                try!(recursive_copy(&entry.path(), &new_dest));
            } else {
                try!(fs::copy(&entry.path(), new_dest));
            }
        }
    }
    Ok(())
}

pub struct AtPath {
    pub subdir: PathBuf,
}
impl AtPath {
    pub fn new(subdir: &Path) -> AtPath {
        AtPath { subdir: PathBuf::from(subdir) }
    }
    pub fn as_string(&self) -> String {
        self.subdir.to_str().unwrap().to_owned()
    }
    pub fn plus(&self, name: &str) -> PathBuf {
        let mut pathbuf = self.subdir.clone();
        pathbuf.push(name);
        pathbuf
    }
    pub fn plus_as_string(&self, name: &str) -> String {
        String::from(self.plus(name).to_str().unwrap())
    }
    fn minus(&self, name: &str) -> PathBuf {
        // relative_from is currently unstable
        let prefixed = PathBuf::from(name);
        if prefixed.starts_with(&self.subdir) {
            let mut unprefixed = PathBuf::new();
            for component in prefixed.components()
                                     .skip(self.subdir.components().count()) {
                unprefixed.push(component.as_ref().to_str().unwrap());
            }
            unprefixed
        } else {
            prefixed
        }
    }
    pub fn minus_as_string(&self, name: &str) -> String {
        String::from(self.minus(name).to_str().unwrap())
    }
    pub fn open(&self, name: &str) -> File {
        log_info("open", self.plus_as_string(name));
        File::open(self.plus(name)).unwrap()
    }
    pub fn read(&self, name: &str) -> String {
        let mut f = self.open(name);
        let mut contents = String::new();
        let _ = f.read_to_string(&mut contents);
        contents
    }
    pub fn write(&self, name: &str, contents: &str) {
        let mut f = self.open(name);
        let _ = f.write(contents.as_bytes());
    }
    pub fn mkdir(&self, dir: &str) {
        log_info("mkdir", self.plus_as_string(dir));
        fs::create_dir(&self.plus(dir)).unwrap();
    }
    pub fn mkdir_all(&self, dir: &str) {
        log_info("mkdir_all", self.plus_as_string(dir));
        fs::create_dir_all(self.plus(dir)).unwrap();
    }
    pub fn make_file(&self, name: &str) -> File {
        match File::create(&self.plus(name)) {
            Ok(f) => f,
            Err(e) => panic!("{}", e),
        }
    }
    pub fn touch(&self, file: &str) {
        log_info("touch", self.plus_as_string(file));
        File::create(&self.plus(file)).unwrap();
    }
    pub fn symlink(&self, src: &str, dst: &str) {
        log_info("symlink",
                 &format!("{},{}", self.plus_as_string(src), self.plus_as_string(dst)));
        symlink_file(&self.plus(src), &self.plus(dst)).unwrap();
    }
    pub fn is_symlink(&self, path: &str) -> bool {
        log_info("is_symlink", self.plus_as_string(path));
        match fs::symlink_metadata(&self.plus(path)) {
            Ok(m) => m.file_type().is_symlink(),
            Err(_) => false,
        }
    }

    pub fn resolve_link(&self, path: &str) -> String {
        log_info("resolve_link", self.plus_as_string(path));
        match fs::read_link(&self.plus(path)) {
            Ok(p) => {
                self.minus_as_string(p.to_str().unwrap())
            }
            Err(_) => "".to_string(),
        }
    }

    pub fn metadata(&self, path: &str) -> fs::Metadata {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m,
            Err(e) => panic!("{}", e),
        }
    }

    pub fn file_exists(&self, path: &str) -> bool {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m.is_file(),
            Err(_) => false,
        }
    }

    pub fn dir_exists(&self, path: &str) -> bool {
        match fs::metadata(&self.plus(path)) {
            Ok(m) => m.is_dir(),
            Err(_) => false,
        }
    }

    pub fn cleanup(&self, path: &'static str) {
        let p = &self.plus(path);
        match fs::metadata(p) {
            Ok(m) => if m.is_file() {
                fs::remove_file(&p).unwrap();
            } else {
                fs::remove_dir(&p).unwrap();
            },
            Err(_) => {}
        }
    }
    pub fn root_dir(&self) -> String {
        log_info("current_directory", "");
        self.subdir.to_str().unwrap().to_owned()
    }
}

pub struct TestSet {
    bin_path: PathBuf,
    util_name: String,
    pub fixtures: AtPath,
    tmpd: Rc<TempDir>,
}
impl TestSet {
    pub fn new(util_name: &str) -> TestSet {
        let tmpd = Rc::new(TempDir::new("uutils").unwrap());
        let ts = TestSet {
            bin_path: {
                let mut bin_path_builder = env::current_dir().unwrap();
                bin_path_builder.push(PathBuf::from(PROGNAME));
                bin_path_builder
            },
            util_name: String::from(util_name),
            fixtures: AtPath::new(&tmpd.as_ref().path()),
            tmpd: tmpd,
        };
        let mut fixture_path_builder = env::current_dir().unwrap();
        fixture_path_builder.push(PathBuf::from(FIXTURES_DIR));
        fixture_path_builder.push(PathBuf::from(util_name));
        match fs::metadata(&fixture_path_builder) {
            Ok(m) => if m.is_dir() {
                recursive_copy(&fixture_path_builder, &ts.fixtures.subdir).unwrap();
            },
            Err(_) => {}
        }
        ts
    }
    pub fn util_cmd(&self) -> UCommand {
        let mut cmd = self.cmd(&self.bin_path);
        cmd.arg(&self.util_name);
        cmd
    }
    pub fn cmd<S: AsRef<OsStr>>(&self, bin: S) -> UCommand {
        UCommand::new_from_tmp(bin, self.tmpd.clone(), true)
    }
    // different names are used rather than an argument
    // because the need to keep the environment is exceedingly rare.
    pub fn util_cmd_keepenv(&self) -> UCommand {
        let mut cmd = self.cmd_keepenv(&self.bin_path);
        cmd.arg(&self.util_name);
        cmd
    }
    pub fn cmd_keepenv<S: AsRef<OsStr>>(&self, bin: S) -> UCommand {
        UCommand::new_from_tmp(bin, self.tmpd.clone(), false)
    }
}

pub struct UCommand {
    pub raw: Command,
    comm_string: String,
    tmpd: Option<Rc<TempDir>>,
    has_run: bool,
}
impl UCommand {
    pub fn new<T: AsRef<OsStr>, U: AsRef<OsStr>>(arg: T, curdir: U, env_clear: bool) -> UCommand {
        UCommand {
            tmpd: None,
            has_run: false,
            raw: {
                let mut cmd = Command::new(arg.as_ref());
                cmd.current_dir(curdir.as_ref());
                if env_clear {
                    cmd.env_clear();
                }
                cmd
            },
            comm_string: String::from(arg.as_ref().to_str().unwrap()),
        }
    }
    pub fn new_from_tmp<T: AsRef<OsStr>>(arg: T, tmpd: Rc<TempDir>, env_clear: bool) -> UCommand {
        let tmpd_path_buf = String::from(&(*tmpd.as_ref().path().to_str().unwrap()));
        let mut ucmd: UCommand = UCommand::new(arg.as_ref(), tmpd_path_buf, env_clear);
        ucmd.tmpd = Some(tmpd);
        ucmd
    }
    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> Box<&mut UCommand> {
        if self.has_run {
            panic!(ALREADY_RUN);
        }
        self.comm_string.push_str(" ");
        self.comm_string.push_str(arg.as_ref().to_str().unwrap());
        self.raw.arg(arg.as_ref());
        Box::new(self)
    }

    pub fn args<S: AsRef<OsStr>>(&mut self, args: &[S]) -> Box<&mut UCommand> {
        if self.has_run {
            panic!(ALREADY_RUN);
        }
        for s in args {
            self.comm_string.push_str(" ");
            self.comm_string.push_str(s.as_ref().to_str().unwrap());
        }

        self.raw.args(args.as_ref());
        Box::new(self)
    }

    pub fn env<K, V>(&mut self, key: K, val: V) -> Box<&mut UCommand> where K: AsRef<OsStr>, V: AsRef<OsStr> {
        if self.has_run {
            panic!(ALREADY_RUN);
        }
        self.raw.env(key, val);
        Box::new(self)
    }

    pub fn run(&mut self) -> CmdResult {
        self.has_run = true;
        log_info("run", &self.comm_string);
        let prog = self.raw.output().unwrap();
        CmdResult {
            success: prog.status.success(),
            stdout: from_utf8(&prog.stdout).unwrap().to_string(),
            stderr: from_utf8(&prog.stderr).unwrap().to_string(),
        }
    }
    pub fn run_piped_stdin<T: AsRef<[u8]>>(&mut self, input: T) -> CmdResult {
        self.has_run = true;
        log_info("run_piped_stdin", &self.comm_string);
        let mut result = self.raw
                             .stdin(Stdio::piped())
                             .stdout(Stdio::piped())
                             .stderr(Stdio::piped())
                             .spawn()
                             .unwrap();

        result.stdin
              .take()
              .unwrap_or_else(|| panic!("Could not take child process stdin"))
              .write_all(input.as_ref())
              .unwrap_or_else(|e| panic!("{}", e));

        let prog = result.wait_with_output().unwrap();
        CmdResult {
            success: prog.status.success(),
            stdout: from_utf8(&prog.stdout).unwrap().to_string(),
            stderr: from_utf8(&prog.stderr).unwrap().to_string(),
        }
    }
}

// returns a testSet and a ucommand initialized to the utility binary
// operating in the fixtures directory with a cleared environment
pub fn testset_and_ucommand(utilname: &str) -> (TestSet, UCommand) {
    let ts = TestSet::new(utilname);
    let ucmd = ts.util_cmd();
    (ts, ucmd)
}
pub fn testing(utilname: &str) -> (AtPath, UCommand) {
    let ts = TestSet::new(utilname);
    let ucmd = ts.util_cmd();
    (ts.fixtures, ucmd)
}
<|endoftext|>


<commit_before>// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use std::fmt;

/// Function implementations' parameter data types.
///
/// It is similar to the `EvalType` in TiDB, but doesn't provide type `Timestamp`, which is
/// handled by the same type as `DateTime` here, instead of a new type. Also, `String` is
/// called `Bytes` here to be less confusing.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum EvalType {
    Int,
    Real,
    Decimal,
    Bytes,
    DateTime,
    Duration,
    Json,
    Enum,
    Set,
}

impl EvalType {
    /// Converts `EvalType` into one of the compatible `FieldTypeTp`s.
    ///
    /// This function should be only useful in test scenarios that only cares about `EvalType` but
    /// accepts a `FieldTypeTp`.
    pub fn into_certain_field_type_tp_for_test(self) -> crate::FieldTypeTp {
        match self {
            EvalType::Int => crate::FieldTypeTp::LongLong,
            EvalType::Real => crate::FieldTypeTp::Double,
            EvalType::Decimal => crate::FieldTypeTp::NewDecimal,
            EvalType::Bytes => crate::FieldTypeTp::String,
            EvalType::DateTime => crate::FieldTypeTp::DateTime,
            EvalType::Duration => crate::FieldTypeTp::Duration,
            EvalType::Json => crate::FieldTypeTp::JSON,
            EvalType::Enum => crate::FieldTypeTp::Enum,
            EvalType::Set => crate::FieldTypeTp::Set,
        }
    }
}

impl fmt::Display for EvalType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}

impl std::convert::TryFrom<crate::FieldTypeTp> for EvalType {
    type Error = crate::DataTypeError;

    // Succeeds for all field types supported as eval types, fails for unsupported types.
    fn try_from(tp: crate::FieldTypeTp) -> Result<Self, crate::DataTypeError> {
        let eval_type = match tp {
            crate::FieldTypeTp::Tiny
            | crate::FieldTypeTp::Short
            | crate::FieldTypeTp::Int24
            | crate::FieldTypeTp::Long
            | crate::FieldTypeTp::LongLong
            | crate::FieldTypeTp::Year => EvalType::Int,
            crate::FieldTypeTp::Float | crate::FieldTypeTp::Double => EvalType::Real,
            crate::FieldTypeTp::NewDecimal => EvalType::Decimal,
            crate::FieldTypeTp::Timestamp
            | crate::FieldTypeTp::Date
            | crate::FieldTypeTp::DateTime => EvalType::DateTime,
            crate::FieldTypeTp::Duration => EvalType::Duration,
            crate::FieldTypeTp::JSON => EvalType::Json,
            crate::FieldTypeTp::VarChar
            | crate::FieldTypeTp::TinyBlob
            | crate::FieldTypeTp::MediumBlob
            | crate::FieldTypeTp::LongBlob
            | crate::FieldTypeTp::Blob
            | crate::FieldTypeTp::VarString
            | crate::FieldTypeTp::String => EvalType::Bytes,
            _ => {
                // In TiDB, Bit's eval type is Int, but it is not yet supported in TiKV.
                // TODO: we need to handle FieldTypeTp::{Enum, Set} after we implement encode and decode.
                return Err(crate::DataTypeError::UnsupportedType {
                    name: tp.to_string(),
                });
            }
        };
        Ok(eval_type)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FieldTypeAccessor;
    use crate::FieldTypeTp::*;
    use std::convert::TryFrom;

    #[test]
    fn test_fieldtype_to_evaltype() {
        let cases = vec![
            (Unspecified, None),
            (Tiny, Some(EvalType::Int)),
            (Short, Some(EvalType::Int)),
            (Long, Some(EvalType::Int)),
            (Float, Some(EvalType::Real)),
            (Double, Some(EvalType::Real)),
            (Null, None),
            (Timestamp, Some(EvalType::DateTime)),
            (LongLong, Some(EvalType::Int)),
            (Int24, Some(EvalType::Int)),
            (Date, Some(EvalType::DateTime)),
            (Duration, Some(EvalType::Duration)),
            (DateTime, Some(EvalType::DateTime)),
            (Year, Some(EvalType::Int)),
            (NewDate, None),
            (VarChar, Some(EvalType::Bytes)),
            (Bit, None),
            (JSON, Some(EvalType::Json)),
            (NewDecimal, Some(EvalType::Decimal)),
            (Enum, None),
            (Set, None),
            (TinyBlob, Some(EvalType::Bytes)),
            (MediumBlob, Some(EvalType::Bytes)),
            (LongBlob, Some(EvalType::Bytes)),
            (Blob, Some(EvalType::Bytes)),
            (VarString, Some(EvalType::Bytes)),
            (String, Some(EvalType::Bytes)),
            (Geometry, None),
        ];

        for (tp, etype) in cases {
            let mut ft = tipb::FieldType::default();
            ft.set_tp(tp as i32);

            let ftt = EvalType::try_from(ft.as_accessor().tp());

            if let Some(etype) = etype {
                assert_eq!(ftt.unwrap(), etype);
            } else {
                assert!(ftt.is_err());
            }
        }
    }
}
<commit_msg>copr: use EvalType::Bytes as the eval type for Null type. (#9632)<commit_after>// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use std::fmt;

/// Function implementations' parameter data types.
///
/// It is similar to the `EvalType` in TiDB, but doesn't provide type `Timestamp`, which is
/// handled by the same type as `DateTime` here, instead of a new type. Also, `String` is
/// called `Bytes` here to be less confusing.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum EvalType {
    Int,
    Real,
    Decimal,
    Bytes,
    DateTime,
    Duration,
    Json,
    Enum,
    Set,
}

impl EvalType {
    /// Converts `EvalType` into one of the compatible `FieldTypeTp`s.
    ///
    /// This function should be only useful in test scenarios that only cares about `EvalType` but
    /// accepts a `FieldTypeTp`.
    pub fn into_certain_field_type_tp_for_test(self) -> crate::FieldTypeTp {
        match self {
            EvalType::Int => crate::FieldTypeTp::LongLong,
            EvalType::Real => crate::FieldTypeTp::Double,
            EvalType::Decimal => crate::FieldTypeTp::NewDecimal,
            EvalType::Bytes => crate::FieldTypeTp::String,
            EvalType::DateTime => crate::FieldTypeTp::DateTime,
            EvalType::Duration => crate::FieldTypeTp::Duration,
            EvalType::Json => crate::FieldTypeTp::JSON,
            EvalType::Enum => crate::FieldTypeTp::Enum,
            EvalType::Set => crate::FieldTypeTp::Set,
        }
    }
}

impl fmt::Display for EvalType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}

impl std::convert::TryFrom<crate::FieldTypeTp> for EvalType {
    type Error = crate::DataTypeError;

    // Succeeds for all field types supported as eval types, fails for unsupported types.
    fn try_from(tp: crate::FieldTypeTp) -> Result<Self, crate::DataTypeError> {
        let eval_type = match tp {
            crate::FieldTypeTp::Tiny
            | crate::FieldTypeTp::Short
            | crate::FieldTypeTp::Int24
            | crate::FieldTypeTp::Long
            | crate::FieldTypeTp::LongLong
            | crate::FieldTypeTp::Year => EvalType::Int,
            crate::FieldTypeTp::Float | crate::FieldTypeTp::Double => EvalType::Real,
            crate::FieldTypeTp::NewDecimal => EvalType::Decimal,
            crate::FieldTypeTp::Timestamp
            | crate::FieldTypeTp::Date
            | crate::FieldTypeTp::DateTime => EvalType::DateTime,
            crate::FieldTypeTp::Duration => EvalType::Duration,
            crate::FieldTypeTp::JSON => EvalType::Json,
            crate::FieldTypeTp::VarChar
            | crate::FieldTypeTp::TinyBlob
            | crate::FieldTypeTp::MediumBlob
            | crate::FieldTypeTp::LongBlob
            | crate::FieldTypeTp::Blob
            | crate::FieldTypeTp::VarString
            | crate::FieldTypeTp::String
            | crate::FieldTypeTp::Null => EvalType::Bytes,
            _ => {
                // In TiDB, Bit's eval type is Int, but it is not yet supported in TiKV.
                // TODO: we need to handle FieldTypeTp::{Enum, Set} after we implement encode and decode.
                return Err(crate::DataTypeError::UnsupportedType {
                    name: tp.to_string(),
                });
            }
        };
        Ok(eval_type)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FieldTypeAccessor;
    use crate::FieldTypeTp::*;
    use std::convert::TryFrom;

    #[test]
    fn test_fieldtype_to_evaltype() {
        let cases = vec![
            (Unspecified, None),
            (Tiny, Some(EvalType::Int)),
            (Short, Some(EvalType::Int)),
            (Long, Some(EvalType::Int)),
            (Float, Some(EvalType::Real)),
            (Double, Some(EvalType::Real)),
            (Null, Some(EvalType::Bytes)),
            (Timestamp, Some(EvalType::DateTime)),
            (LongLong, Some(EvalType::Int)),
            (Int24, Some(EvalType::Int)),
            (Date, Some(EvalType::DateTime)),
            (Duration, Some(EvalType::Duration)),
            (DateTime, Some(EvalType::DateTime)),
            (Year, Some(EvalType::Int)),
            (NewDate, None),
            (VarChar, Some(EvalType::Bytes)),
            (Bit, None),
            (JSON, Some(EvalType::Json)),
            (NewDecimal, Some(EvalType::Decimal)),
            (Enum, None),
            (Set, None),
            (TinyBlob, Some(EvalType::Bytes)),
            (MediumBlob, Some(EvalType::Bytes)),
            (LongBlob, Some(EvalType::Bytes)),
            (Blob, Some(EvalType::Bytes)),
            (VarString, Some(EvalType::Bytes)),
            (String, Some(EvalType::Bytes)),
            (Geometry, None),
        ];

        for (tp, etype) in cases {
            let mut ft = tipb::FieldType::default();
            ft.set_tp(tp as i32);

            let ftt = EvalType::try_from(ft.as_accessor().tp());

            if let Some(etype) = etype {
                assert_eq!(ftt.unwrap(), etype);
            } else {
                assert!(ftt.is_err());
            }
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>Init Nodes with their max capacity for fewer reallocs<commit_after><|endoftext|>


<commit_before>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity, github_updater, pubsubhubbub};
use db::{connect_db, update_search_index};

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    let dbopts = opts();

    // check paths once
    dbopts.check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(dbopts.prefix.join("cratesfyi.pid"))
            .expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every minute
    thread::Builder::new().name("crates.io reader".to_string()).spawn(move || {
        // space this out to prevent it from clashing against the queue-builder thread on launch
        thread::sleep(Duration::from_secs(30));
        loop {
            let opts = opts();
            let mut doc_builder = DocBuilder::new(opts);

            debug!("Checking new crates");
            match doc_builder.get_new_crates() {
                Ok(n) => debug!("{} crates added to queue", n),
                Err(e) => error!("Failed to get new crates: {}", e),
            }

            thread::sleep(Duration::from_secs(60));
        }
    }).unwrap();

    // build new crates every minute
    thread::Builder::new().name("build queue reader".to_string()).spawn(move || {
        let opts = opts();
        let mut doc_builder = DocBuilder::new(opts);

        /// Represents the current state of the builder thread.
        enum BuilderState {
            /// The builder thread has just started, and hasn't built any crates yet.
            Fresh,
            /// The builder has just seen an empty build queue.
            EmptyQueue,
            /// The builder has just seen the lock file.
            Locked,
            /// The builder has just finished building a crate. The enclosed count is the number of
            /// crates built since the caches have been refreshed.
            QueueInProgress(usize),
        }

        let mut status = BuilderState::Fresh;

        loop {
            if !status.is_in_progress() {
                thread::sleep(Duration::from_secs(60));
            }

            // check lock file
            if doc_builder.is_locked() {
                warn!("Lock file exits, skipping building new crates");
                status = BuilderState::Locked;
                continue;
            }

            if status.count() >= 10 {
                // periodically, we need to flush our caches and ping the hubs
                debug!("10 builds in a row; flushing caches");
                status = BuilderState::QueueInProgress(0);

                match pubsubhubbub::ping_hubs() {
                    Err(e) => error!("Failed to ping hub: {}", e),
                    Ok(n) => debug!("Succesfully pinged {} hubs", n)
                }

                if let Err(e) = doc_builder.load_cache() {
                    error!("Failed to load cache: {}", e);
                }

                if let Err(e) = doc_builder.save_cache() {
                    error!("Failed to save cache: {}", e);
                }

                if let Err(e) = update_sources() {
                    error!("Failed to update sources: {}", e);
                    continue;
                }
            }

            // Only build crates if there are any to build
            debug!("Checking build queue");
            match doc_builder.get_queue_count() {
                Err(e) => {
                    error!("Failed to read the number of crates in the queue: {}", e);
                    continue;
                }
                Ok(0) => {
                    if status.count() > 0 {
                        // ping the hubs before continuing
                        match pubsubhubbub::ping_hubs() {
                            Err(e) => error!("Failed to ping hub: {}", e),
                            Ok(n) => debug!("Succesfully pinged {} hubs", n)
                        }

                        if let Err(e) = doc_builder.save_cache() {
                            error!("Failed to save cache: {}", e);
                        }
                    }
                    debug!("Queue is empty, going back to sleep");
                    status = BuilderState::EmptyQueue;
                    continue;
                }
                Ok(queue_count) => {
                    info!("Starting build with {} crates in queue (currently on a {} crate streak)",
                          queue_count, status.count());
                }
            }

            // if we're starting a new batch, reload our caches and sources
            if !status.is_in_progress() {
                if let Err(e) = doc_builder.load_cache() {
                    error!("Failed to load cache: {}", e);
                    continue;
                }

                if let Err(e) = update_sources() {
                    error!("Failed to update sources: {}", e);
                    continue;
                }
            }

            // Run build_packages_queue under `catch_unwind` to catch panics
            // This only panicked twice in the last 6 months but its just a better
            // idea to do this.
            let res = catch_unwind(AssertUnwindSafe(|| {
                match doc_builder.build_next_queue_package() {
                    Err(e) => error!("Failed to build crate from queue: {}", e),
                    Ok(crate_built) => if crate_built {
                        status.increment();
                    }

                }
            }));

            if let Err(e) = res {
                error!("GRAVE ERROR Building new crates panicked: {:?}", e);
            }
        }

        impl BuilderState {
            fn count(&self) -> usize {
                match *self {
                    BuilderState::QueueInProgress(n) => n,
                    _ => 0,
                }
            }

            fn is_in_progress(&self) -> bool {
                match *self {
                    BuilderState::QueueInProgress(_) => true,
                    _ => false,
                }
            }

            fn increment(&mut self) {
                *self = BuilderState::QueueInProgress(self.count() + 1);
            }
        }
    }).unwrap();

    // update release activity everyday at 23:55
    thread::Builder::new().name("release activity updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 23 && now.tm_min == 55 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    }).unwrap();

    // update search index every 3 hours
    thread::Builder::new().name("search index updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60 * 60 * 3));
            let conn = connect_db().expect("Failed to connect database");
            if let Err(e) = update_search_index(&conn) {
                error!("Failed to update search index: {}", e);
            }
        }
    }).unwrap();

    // update github stats every 6 hours
    thread::Builder::new().name("github stat updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60 * 60 * 6));
            if let Err(e) = github_updater() {
                error!("Failed to update github fields: {}", e);
            }
        }
    }).unwrap();

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
        .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<commit_msg>don't check new crates if lock file exists<commit_after>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity, github_updater, pubsubhubbub};
use db::{connect_db, update_search_index};

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    let dbopts = opts();

    // check paths once
    dbopts.check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(dbopts.prefix.join("cratesfyi.pid"))
            .expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every minute
    thread::Builder::new().name("crates.io reader".to_string()).spawn(move || {
        // space this out to prevent it from clashing against the queue-builder thread on launch
        thread::sleep(Duration::from_secs(30));
        loop {
            let opts = opts();
            let mut doc_builder = DocBuilder::new(opts);

            if doc_builder.is_locked() {
                debug!("Lock file exists, skipping checking new crates");
            } else {
                debug!("Checking new crates");
                match doc_builder.get_new_crates() {
                    Ok(n) => debug!("{} crates added to queue", n),
                    Err(e) => error!("Failed to get new crates: {}", e),
                }
            }

            thread::sleep(Duration::from_secs(60));
        }
    }).unwrap();

    // build new crates every minute
    thread::Builder::new().name("build queue reader".to_string()).spawn(move || {
        let opts = opts();
        let mut doc_builder = DocBuilder::new(opts);

        /// Represents the current state of the builder thread.
        enum BuilderState {
            /// The builder thread has just started, and hasn't built any crates yet.
            Fresh,
            /// The builder has just seen an empty build queue.
            EmptyQueue,
            /// The builder has just seen the lock file.
            Locked,
            /// The builder has just finished building a crate. The enclosed count is the number of
            /// crates built since the caches have been refreshed.
            QueueInProgress(usize),
        }

        let mut status = BuilderState::Fresh;

        loop {
            if !status.is_in_progress() {
                thread::sleep(Duration::from_secs(60));
            }

            // check lock file
            if doc_builder.is_locked() {
                warn!("Lock file exits, skipping building new crates");
                status = BuilderState::Locked;
                continue;
            }

            if status.count() >= 10 {
                // periodically, we need to flush our caches and ping the hubs
                debug!("10 builds in a row; flushing caches");
                status = BuilderState::QueueInProgress(0);

                match pubsubhubbub::ping_hubs() {
                    Err(e) => error!("Failed to ping hub: {}", e),
                    Ok(n) => debug!("Succesfully pinged {} hubs", n)
                }

                if let Err(e) = doc_builder.load_cache() {
                    error!("Failed to load cache: {}", e);
                }

                if let Err(e) = doc_builder.save_cache() {
                    error!("Failed to save cache: {}", e);
                }

                if let Err(e) = update_sources() {
                    error!("Failed to update sources: {}", e);
                    continue;
                }
            }

            // Only build crates if there are any to build
            debug!("Checking build queue");
            match doc_builder.get_queue_count() {
                Err(e) => {
                    error!("Failed to read the number of crates in the queue: {}", e);
                    continue;
                }
                Ok(0) => {
                    if status.count() > 0 {
                        // ping the hubs before continuing
                        match pubsubhubbub::ping_hubs() {
                            Err(e) => error!("Failed to ping hub: {}", e),
                            Ok(n) => debug!("Succesfully pinged {} hubs", n)
                        }

                        if let Err(e) = doc_builder.save_cache() {
                            error!("Failed to save cache: {}", e);
                        }
                    }
                    debug!("Queue is empty, going back to sleep");
                    status = BuilderState::EmptyQueue;
                    continue;
                }
                Ok(queue_count) => {
                    info!("Starting build with {} crates in queue (currently on a {} crate streak)",
                          queue_count, status.count());
                }
            }

            // if we're starting a new batch, reload our caches and sources
            if !status.is_in_progress() {
                if let Err(e) = doc_builder.load_cache() {
                    error!("Failed to load cache: {}", e);
                    continue;
                }

                if let Err(e) = update_sources() {
                    error!("Failed to update sources: {}", e);
                    continue;
                }
            }

            // Run build_packages_queue under `catch_unwind` to catch panics
            // This only panicked twice in the last 6 months but its just a better
            // idea to do this.
            let res = catch_unwind(AssertUnwindSafe(|| {
                match doc_builder.build_next_queue_package() {
                    Err(e) => error!("Failed to build crate from queue: {}", e),
                    Ok(crate_built) => if crate_built {
                        status.increment();
                    }

                }
            }));

            if let Err(e) = res {
                error!("GRAVE ERROR Building new crates panicked: {:?}", e);
            }
        }

        impl BuilderState {
            fn count(&self) -> usize {
                match *self {
                    BuilderState::QueueInProgress(n) => n,
                    _ => 0,
                }
            }

            fn is_in_progress(&self) -> bool {
                match *self {
                    BuilderState::QueueInProgress(_) => true,
                    _ => false,
                }
            }

            fn increment(&mut self) {
                *self = BuilderState::QueueInProgress(self.count() + 1);
            }
        }
    }).unwrap();

    // update release activity everyday at 23:55
    thread::Builder::new().name("release activity updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 23 && now.tm_min == 55 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    }).unwrap();

    // update search index every 3 hours
    thread::Builder::new().name("search index updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60 * 60 * 3));
            let conn = connect_db().expect("Failed to connect database");
            if let Err(e) = update_search_index(&conn) {
                error!("Failed to update search index: {}", e);
            }
        }
    }).unwrap();

    // update github stats every 6 hours
    thread::Builder::new().name("github stat updater".to_string()).spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60 * 60 * 6));
            if let Err(e) = github_updater() {
                error!("Failed to update github fields: {}", e);
            }
        }
    }).unwrap();

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
        .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<|endoftext|>


<commit_before>use associations::{HasTable, Identifiable};
use dsl::Find;
use query_dsl::FindDsl;
use query_source::Table;

#[doc(hidden)]
#[derive(Debug)]
pub struct UpdateTarget<Table, WhereClause> {
    pub table: Table,
    pub where_clause: WhereClause,
}

/// A type which can be passed to `update`.
///
/// Apps will never need to implement this type directly. There are three kinds
/// of things which implement this trait. Tables, queries which have only had
/// `filter` called on them, and types which implement `Identifiable`.
///
/// When a table is passed to `update`, every row in the table will be updated.
/// You can scope this down by calling `filter` before passing it, which will
/// result in `UPDATE your_table SET ... WHERE args_to_filter`. Passing a type
/// which implements `Identifiable` is the same as passing
/// `SomeStruct::table().find(some_struct)`.
pub trait IntoUpdateTarget: HasTable {
    type WhereClause;

    fn into_update_target(self) -> UpdateTarget<Self::Table, Self::WhereClause>;
}

impl<T, Tab, V> IntoUpdateTarget for T
where
    T: Identifiable<Table = Tab>,
    Tab: Table + FindDsl<T::Id>,
    Find<Tab, T::Id>: IntoUpdateTarget<Table = Tab, WhereClause = V>,
{
    type WhereClause = V;

    fn into_update_target(self) -> UpdateTarget<Self::Table, Self::WhereClause> {
        T::table().find(self.id()).into_update_target()
    }
}
<commit_msg>Document all the items of `IntoUpdateTarget`<commit_after>use associations::{HasTable, Identifiable};
use dsl::Find;
use query_dsl::FindDsl;
use query_source::Table;

#[doc(hidden)]
#[derive(Debug)]
pub struct UpdateTarget<Table, WhereClause> {
    pub table: Table,
    pub where_clause: WhereClause,
}

/// A type which can be passed to `update`.
///
/// Apps will never need to implement this type directly. There are three kinds
/// which implement this trait. Tables, queries which have only had `filter`
/// called on them, and types which implement `Identifiable`.
///
/// When a table is passed to `update`, every row in the table will be updated.
/// You can scope this down by calling `filter` before passing it, which will
/// result in `UPDATE your_table SET ... WHERE args_to_filter`. Passing a type
/// which implements `Identifiable` is the same as passing
/// `SomeStruct::table().find(some_struct)`.
pub trait IntoUpdateTarget: HasTable {
    /// What is the `WHERE` clause of this target?
    type WhereClause;

    /// Decomposes `self` into the table and where clause.
    fn into_update_target(self) -> UpdateTarget<Self::Table, Self::WhereClause>;
}

impl<T, Tab, V> IntoUpdateTarget for T
where
    T: Identifiable<Table = Tab>,
    Tab: Table + FindDsl<T::Id>,
    Find<Tab, T::Id>: IntoUpdateTarget<Table = Tab, WhereClause = V>,
{
    type WhereClause = V;

    fn into_update_target(self) -> UpdateTarget<Self::Table, Self::WhereClause> {
        T::table().find(self.id()).into_update_target()
    }
}
<|endoftext|>


<commit_before>use crate::prelude::*;

use fixedbitset::FixedBitSet;
use std::collections::HashSet;
use std::marker::PhantomData;

use crate::data::DataMap;
use crate::visit::{Data, NodeCompactIndexable, NodeCount};
use crate::visit::{
    GraphBase, GraphProp, IntoEdgeReferences, IntoEdges, IntoEdgesDirected, IntoNeighbors,
    IntoNeighborsDirected, IntoNodeIdentifiers, IntoNodeReferences, NodeIndexable, NodeRef,
    VisitMap, Visitable,
};

/// A graph filter for nodes.
pub trait FilterNode<N> {
    /// Return true to have the node be part of the graph
    fn include_node(&self, node: N) -> bool;
}

impl<F, N> FilterNode<N> for F
where
    F: Fn(N) -> bool,
{
    fn include_node(&self, n: N) -> bool {
        (*self)(n)
    }
}

/// This filter includes the nodes that are contained in the set.
impl<N> FilterNode<N> for FixedBitSet
where
    FixedBitSet: VisitMap<N>,
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

/// This filter includes the nodes that are contained in the set.
impl<N, S> FilterNode<N> for HashSet<N, S>
where
    HashSet<N, S>: VisitMap<N>,
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

/// A node-filtering graph adaptor.
#[derive(Copy, Clone, Debug)]
pub struct NodeFiltered<G, F>(pub G, pub F);

impl<F, G> NodeFiltered<G, F>
where
    G: GraphBase,
    F: Fn(G::NodeId) -> bool,
{
    /// Create an `NodeFiltered` adaptor from the closure `filter`.
    pub fn from_fn(graph: G, filter: F) -> Self {
        NodeFiltered(graph, filter)
    }
}

impl<G, F> GraphBase for NodeFiltered<G, F>
where
    G: GraphBase,
{
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<'a, G, F> IntoNeighbors for &'a NodeFiltered<G, F>
where
    G: IntoNeighbors,
    F: FilterNode<G::NodeId>,
{
    type Neighbors = NodeFilteredNeighbors<'a, G::Neighbors, F>;
    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
        NodeFilteredNeighbors {
            include_source: self.1.include_node(n),
            iter: self.0.neighbors(n),
            f: &self.1,
        }
    }
}

/// A filtered neighbors iterator.
pub struct NodeFilteredNeighbors<'a, I, F: 'a> {
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, I, F> Iterator for NodeFilteredNeighbors<'a, I, F>
where
    I: Iterator,
    I::Item: Copy,
    F: FilterNode<I::Item>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        if !self.include_source {
            None
        } else {
            self.iter.find(move |&target| f.include_node(target))
        }
    }
}

impl<'a, G, F> IntoNeighborsDirected for &'a NodeFiltered<G, F>
where
    G: IntoNeighborsDirected,
    F: FilterNode<G::NodeId>,
{
    type NeighborsDirected = NodeFilteredNeighbors<'a, G::NeighborsDirected, F>;
    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
        NodeFilteredNeighbors {
            include_source: self.1.include_node(n),
            iter: self.0.neighbors_directed(n, dir),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNodeIdentifiers for &'a NodeFiltered<G, F>
where
    G: IntoNodeIdentifiers,
    F: FilterNode<G::NodeId>,
{
    type NodeIdentifiers = NodeFilteredNeighbors<'a, G::NodeIdentifiers, F>;
    fn node_identifiers(self) -> Self::NodeIdentifiers {
        NodeFilteredNeighbors {
            include_source: true,
            iter: self.0.node_identifiers(),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNodeReferences for &'a NodeFiltered<G, F>
where
    G: IntoNodeReferences,
    F: FilterNode<G::NodeId>,
{
    type NodeRef = G::NodeRef;
    type NodeReferences = NodeFilteredNodes<'a, G::NodeReferences, F>;
    fn node_references(self) -> Self::NodeReferences {
        NodeFilteredNodes {
            include_source: true,
            iter: self.0.node_references(),
            f: &self.1,
        }
    }
}

/// A filtered node references iterator.
pub struct NodeFilteredNodes<'a, I, F: 'a> {
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, I, F> Iterator for NodeFilteredNodes<'a, I, F>
where
    I: Iterator,
    I::Item: Copy + NodeRef,
    F: FilterNode<<I::Item as NodeRef>::NodeId>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        if !self.include_source {
            None
        } else {
            self.iter.find(move |&target| f.include_node(target.id()))
        }
    }
}

impl<'a, G, F> IntoEdgeReferences for &'a NodeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: FilterNode<G::NodeId>,
{
    type EdgeRef = G::EdgeRef;
    type EdgeReferences = NodeFilteredEdgeReferences<'a, G, G::EdgeReferences, F>;
    fn edge_references(self) -> Self::EdgeReferences {
        NodeFilteredEdgeReferences {
            graph: PhantomData,
            iter: self.0.edge_references(),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct NodeFilteredEdgeReferences<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for NodeFilteredEdgeReferences<'a, G, I, F>
where
    F: FilterNode<G::NodeId>,
    G: IntoEdgeReferences,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        self.iter
            .find(move |&edge| f.include_node(edge.source()) && f.include_node(edge.target()))
    }
}

impl<'a, G, F> IntoEdges for &'a NodeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterNode<G::NodeId>,
{
    type Edges = NodeFilteredEdges<'a, G, G::Edges, F>;
    fn edges(self, a: G::NodeId) -> Self::Edges {
        NodeFilteredEdges {
            graph: PhantomData,
            include_source: self.1.include_node(a),
            iter: self.0.edges(a),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct NodeFilteredEdges<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for NodeFilteredEdges<'a, G, I, F>
where
    F: FilterNode<G::NodeId>,
    G: IntoEdges,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        if !self.include_source {
            None
        } else {
            let f = self.f;
            self.iter.find(move |&edge| f.include_node(edge.target()))
        }
    }
}

impl<G, F> DataMap for NodeFiltered<G, F>
where
    G: DataMap,
    F: FilterNode<G::NodeId>,
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        if self.1.include_node(id) {
            self.0.node_weight(id)
        } else {
            None
        }
    }

    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.0.edge_weight(id)
    }
}

macro_rules! access0 {
    ($e:expr) => {
        $e.0
    };
}

Data! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
NodeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
GraphProp! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
Visitable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}

/// A graph filter for edges
pub trait FilterEdge<Edge> {
    /// Return true to have the edge be part of the graph
    fn include_edge(&self, edge: Edge) -> bool;
}

impl<F, N> FilterEdge<N> for F
where
    F: Fn(N) -> bool,
{
    fn include_edge(&self, n: N) -> bool {
        (*self)(n)
    }
}

/// An edge-filtering graph adaptor.
///
/// The adaptor may filter out edges. The filter implements the trait
/// `FilterEdge`. Closures of type `Fn(G::EdgeRef) -> bool` already
/// implement this trait.
///
/// The filter may use edge source, target, id, and weight to select whether to
/// include the edge or not.
#[derive(Copy, Clone, Debug)]
pub struct EdgeFiltered<G, F>(pub G, pub F);

impl<F, G> EdgeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: Fn(G::EdgeRef) -> bool,
{
    /// Create an `EdgeFiltered` adaptor from the closure `filter`.
    pub fn from_fn(graph: G, filter: F) -> Self {
        EdgeFiltered(graph, filter)
    }
}

impl<G, F> GraphBase for EdgeFiltered<G, F>
where
    G: GraphBase,
{
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<'a, G, F> IntoNeighbors for &'a EdgeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterEdge<G::EdgeRef>,
{
    type Neighbors = EdgeFilteredNeighbors<'a, G, F>;
    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
        EdgeFilteredNeighbors {
            iter: self.0.edges(n),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNeighborsDirected for &'a EdgeFiltered<G, F>
where
    G: IntoEdgesDirected,
    F: FilterEdge<G::EdgeRef>,
{
    type NeighborsDirected = EdgeFilteredNeighborsDirected<'a, G, F>;
    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
        EdgeFilteredNeighborsDirected {
            iter: self.0.edges_directed(n, dir),
            f: &self.1,
            from: n,
        }
    }
}

/// A filtered neighbors iterator.
pub struct EdgeFilteredNeighbors<'a, G, F: 'a>
where
    G: IntoEdges,
{
    iter: G::Edges,
    f: &'a F,
}

impl<'a, G, F> Iterator for EdgeFilteredNeighbors<'a, G, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdges,
{
    type Item = G::NodeId;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        (&mut self.iter)
            .filter_map(move |edge| {
                if f.include_edge(edge) {
                    Some(edge.target())
                } else {
                    None
                }
            })
            .next()
    }
}

impl<'a, G, F> IntoEdgeReferences for &'a EdgeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: FilterEdge<G::EdgeRef>,
{
    type EdgeRef = G::EdgeRef;
    type EdgeReferences = EdgeFilteredEdges<'a, G, G::EdgeReferences, F>;
    fn edge_references(self) -> Self::EdgeReferences {
        EdgeFilteredEdges {
            graph: PhantomData,
            iter: self.0.edge_references(),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoEdges for &'a EdgeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterEdge<G::EdgeRef>,
{
    type Edges = EdgeFilteredEdges<'a, G, G::Edges, F>;
    fn edges(self, n: G::NodeId) -> Self::Edges {
        EdgeFilteredEdges {
            graph: PhantomData,
            iter: self.0.edges(n),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct EdgeFilteredEdges<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for EdgeFilteredEdges<'a, G, I, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdgeReferences,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        self.iter.find(move |&edge| f.include_edge(edge))
    }
}

/// A filtered neighbors-directed iterator.
pub struct EdgeFilteredNeighborsDirected<'a, G, F: 'a>
where
    G: IntoEdgesDirected,
{
    iter: G::EdgesDirected,
    f: &'a F,
    from: G::NodeId,
}

impl<'a, G, F> Iterator for EdgeFilteredNeighborsDirected<'a, G, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdgesDirected,
{
    type Item = G::NodeId;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        let from = self.from;
        (&mut self.iter)
            .filter_map(move |edge| {
                if f.include_edge(edge) {
                    if edge.source() != from {
                        Some(edge.source())
                    } else {
                        Some(edge.target()) // includes case where from == source == target
                    }
                } else {
                    None
                }
            })
            .next()
    }
}

Data! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
GraphProp! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
IntoNodeIdentifiers! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
IntoNodeReferences! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
NodeCompactIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
NodeCount! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
NodeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
Visitable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
<commit_msg>[visit] add FilterNode impls for &FixedBitSet and &HashSet<commit_after>use crate::prelude::*;

use fixedbitset::FixedBitSet;
use std::collections::HashSet;
use std::marker::PhantomData;

use crate::data::DataMap;
use crate::visit::{Data, NodeCompactIndexable, NodeCount};
use crate::visit::{
    GraphBase, GraphProp, IntoEdgeReferences, IntoEdges, IntoEdgesDirected, IntoNeighbors,
    IntoNeighborsDirected, IntoNodeIdentifiers, IntoNodeReferences, NodeIndexable, NodeRef,
    VisitMap, Visitable,
};

/// A graph filter for nodes.
pub trait FilterNode<N> {
    /// Return true to have the node be part of the graph
    fn include_node(&self, node: N) -> bool;
}

impl<F, N> FilterNode<N> for F
where
    F: Fn(N) -> bool,
{
    fn include_node(&self, n: N) -> bool {
        (*self)(n)
    }
}

/// This filter includes the nodes that are contained in the set.
impl<N> FilterNode<N> for FixedBitSet
where
    FixedBitSet: VisitMap<N>,
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

/// This filter includes the nodes that are contained in the set.
impl<N, S> FilterNode<N> for HashSet<N, S>
where
    HashSet<N, S>: VisitMap<N>,
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

// Can't express these as a generic impl over all references since that would conflict with the
// impl for Fn.
impl<N> FilterNode<N> for &FixedBitSet
    where FixedBitSet: VisitMap<N>
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

impl<N, S> FilterNode<N> for &HashSet<N, S>
    where HashSet<N, S>: VisitMap<N>,
{
    fn include_node(&self, n: N) -> bool {
        self.is_visited(&n)
    }
}

/// A node-filtering graph adaptor.
#[derive(Copy, Clone, Debug)]
pub struct NodeFiltered<G, F>(pub G, pub F);

impl<F, G> NodeFiltered<G, F>
where
    G: GraphBase,
    F: Fn(G::NodeId) -> bool,
{
    /// Create an `NodeFiltered` adaptor from the closure `filter`.
    pub fn from_fn(graph: G, filter: F) -> Self {
        NodeFiltered(graph, filter)
    }
}

impl<G, F> GraphBase for NodeFiltered<G, F>
where
    G: GraphBase,
{
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<'a, G, F> IntoNeighbors for &'a NodeFiltered<G, F>
where
    G: IntoNeighbors,
    F: FilterNode<G::NodeId>,
{
    type Neighbors = NodeFilteredNeighbors<'a, G::Neighbors, F>;
    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
        NodeFilteredNeighbors {
            include_source: self.1.include_node(n),
            iter: self.0.neighbors(n),
            f: &self.1,
        }
    }
}

/// A filtered neighbors iterator.
pub struct NodeFilteredNeighbors<'a, I, F: 'a> {
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, I, F> Iterator for NodeFilteredNeighbors<'a, I, F>
where
    I: Iterator,
    I::Item: Copy,
    F: FilterNode<I::Item>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        if !self.include_source {
            None
        } else {
            self.iter.find(move |&target| f.include_node(target))
        }
    }
}

impl<'a, G, F> IntoNeighborsDirected for &'a NodeFiltered<G, F>
where
    G: IntoNeighborsDirected,
    F: FilterNode<G::NodeId>,
{
    type NeighborsDirected = NodeFilteredNeighbors<'a, G::NeighborsDirected, F>;
    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
        NodeFilteredNeighbors {
            include_source: self.1.include_node(n),
            iter: self.0.neighbors_directed(n, dir),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNodeIdentifiers for &'a NodeFiltered<G, F>
where
    G: IntoNodeIdentifiers,
    F: FilterNode<G::NodeId>,
{
    type NodeIdentifiers = NodeFilteredNeighbors<'a, G::NodeIdentifiers, F>;
    fn node_identifiers(self) -> Self::NodeIdentifiers {
        NodeFilteredNeighbors {
            include_source: true,
            iter: self.0.node_identifiers(),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNodeReferences for &'a NodeFiltered<G, F>
where
    G: IntoNodeReferences,
    F: FilterNode<G::NodeId>,
{
    type NodeRef = G::NodeRef;
    type NodeReferences = NodeFilteredNodes<'a, G::NodeReferences, F>;
    fn node_references(self) -> Self::NodeReferences {
        NodeFilteredNodes {
            include_source: true,
            iter: self.0.node_references(),
            f: &self.1,
        }
    }
}

/// A filtered node references iterator.
pub struct NodeFilteredNodes<'a, I, F: 'a> {
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, I, F> Iterator for NodeFilteredNodes<'a, I, F>
where
    I: Iterator,
    I::Item: Copy + NodeRef,
    F: FilterNode<<I::Item as NodeRef>::NodeId>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        if !self.include_source {
            None
        } else {
            self.iter.find(move |&target| f.include_node(target.id()))
        }
    }
}

impl<'a, G, F> IntoEdgeReferences for &'a NodeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: FilterNode<G::NodeId>,
{
    type EdgeRef = G::EdgeRef;
    type EdgeReferences = NodeFilteredEdgeReferences<'a, G, G::EdgeReferences, F>;
    fn edge_references(self) -> Self::EdgeReferences {
        NodeFilteredEdgeReferences {
            graph: PhantomData,
            iter: self.0.edge_references(),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct NodeFilteredEdgeReferences<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for NodeFilteredEdgeReferences<'a, G, I, F>
where
    F: FilterNode<G::NodeId>,
    G: IntoEdgeReferences,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        self.iter
            .find(move |&edge| f.include_node(edge.source()) && f.include_node(edge.target()))
    }
}

impl<'a, G, F> IntoEdges for &'a NodeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterNode<G::NodeId>,
{
    type Edges = NodeFilteredEdges<'a, G, G::Edges, F>;
    fn edges(self, a: G::NodeId) -> Self::Edges {
        NodeFilteredEdges {
            graph: PhantomData,
            include_source: self.1.include_node(a),
            iter: self.0.edges(a),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct NodeFilteredEdges<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    include_source: bool,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for NodeFilteredEdges<'a, G, I, F>
where
    F: FilterNode<G::NodeId>,
    G: IntoEdges,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        if !self.include_source {
            None
        } else {
            let f = self.f;
            self.iter.find(move |&edge| f.include_node(edge.target()))
        }
    }
}

impl<G, F> DataMap for NodeFiltered<G, F>
where
    G: DataMap,
    F: FilterNode<G::NodeId>,
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        if self.1.include_node(id) {
            self.0.node_weight(id)
        } else {
            None
        }
    }

    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.0.edge_weight(id)
    }
}

macro_rules! access0 {
    ($e:expr) => {
        $e.0
    };
}

Data! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
NodeIndexable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
GraphProp! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}
Visitable! {delegate_impl [[G, F], G, NodeFiltered<G, F>, access0]}

/// A graph filter for edges
pub trait FilterEdge<Edge> {
    /// Return true to have the edge be part of the graph
    fn include_edge(&self, edge: Edge) -> bool;
}

impl<F, N> FilterEdge<N> for F
where
    F: Fn(N) -> bool,
{
    fn include_edge(&self, n: N) -> bool {
        (*self)(n)
    }
}

/// An edge-filtering graph adaptor.
///
/// The adaptor may filter out edges. The filter implements the trait
/// `FilterEdge`. Closures of type `Fn(G::EdgeRef) -> bool` already
/// implement this trait.
///
/// The filter may use edge source, target, id, and weight to select whether to
/// include the edge or not.
#[derive(Copy, Clone, Debug)]
pub struct EdgeFiltered<G, F>(pub G, pub F);

impl<F, G> EdgeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: Fn(G::EdgeRef) -> bool,
{
    /// Create an `EdgeFiltered` adaptor from the closure `filter`.
    pub fn from_fn(graph: G, filter: F) -> Self {
        EdgeFiltered(graph, filter)
    }
}

impl<G, F> GraphBase for EdgeFiltered<G, F>
where
    G: GraphBase,
{
    type NodeId = G::NodeId;
    type EdgeId = G::EdgeId;
}

impl<'a, G, F> IntoNeighbors for &'a EdgeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterEdge<G::EdgeRef>,
{
    type Neighbors = EdgeFilteredNeighbors<'a, G, F>;
    fn neighbors(self, n: G::NodeId) -> Self::Neighbors {
        EdgeFilteredNeighbors {
            iter: self.0.edges(n),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoNeighborsDirected for &'a EdgeFiltered<G, F>
where
    G: IntoEdgesDirected,
    F: FilterEdge<G::EdgeRef>,
{
    type NeighborsDirected = EdgeFilteredNeighborsDirected<'a, G, F>;
    fn neighbors_directed(self, n: G::NodeId, dir: Direction) -> Self::NeighborsDirected {
        EdgeFilteredNeighborsDirected {
            iter: self.0.edges_directed(n, dir),
            f: &self.1,
            from: n,
        }
    }
}

/// A filtered neighbors iterator.
pub struct EdgeFilteredNeighbors<'a, G, F: 'a>
where
    G: IntoEdges,
{
    iter: G::Edges,
    f: &'a F,
}

impl<'a, G, F> Iterator for EdgeFilteredNeighbors<'a, G, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdges,
{
    type Item = G::NodeId;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        (&mut self.iter)
            .filter_map(move |edge| {
                if f.include_edge(edge) {
                    Some(edge.target())
                } else {
                    None
                }
            })
            .next()
    }
}

impl<'a, G, F> IntoEdgeReferences for &'a EdgeFiltered<G, F>
where
    G: IntoEdgeReferences,
    F: FilterEdge<G::EdgeRef>,
{
    type EdgeRef = G::EdgeRef;
    type EdgeReferences = EdgeFilteredEdges<'a, G, G::EdgeReferences, F>;
    fn edge_references(self) -> Self::EdgeReferences {
        EdgeFilteredEdges {
            graph: PhantomData,
            iter: self.0.edge_references(),
            f: &self.1,
        }
    }
}

impl<'a, G, F> IntoEdges for &'a EdgeFiltered<G, F>
where
    G: IntoEdges,
    F: FilterEdge<G::EdgeRef>,
{
    type Edges = EdgeFilteredEdges<'a, G, G::Edges, F>;
    fn edges(self, n: G::NodeId) -> Self::Edges {
        EdgeFilteredEdges {
            graph: PhantomData,
            iter: self.0.edges(n),
            f: &self.1,
        }
    }
}

/// A filtered edges iterator.
pub struct EdgeFilteredEdges<'a, G, I, F: 'a> {
    graph: PhantomData<G>,
    iter: I,
    f: &'a F,
}

impl<'a, G, I, F> Iterator for EdgeFilteredEdges<'a, G, I, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdgeReferences,
    I: Iterator<Item = G::EdgeRef>,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        self.iter.find(move |&edge| f.include_edge(edge))
    }
}

/// A filtered neighbors-directed iterator.
pub struct EdgeFilteredNeighborsDirected<'a, G, F: 'a>
where
    G: IntoEdgesDirected,
{
    iter: G::EdgesDirected,
    f: &'a F,
    from: G::NodeId,
}

impl<'a, G, F> Iterator for EdgeFilteredNeighborsDirected<'a, G, F>
where
    F: FilterEdge<G::EdgeRef>,
    G: IntoEdgesDirected,
{
    type Item = G::NodeId;
    fn next(&mut self) -> Option<Self::Item> {
        let f = self.f;
        let from = self.from;
        (&mut self.iter)
            .filter_map(move |edge| {
                if f.include_edge(edge) {
                    if edge.source() != from {
                        Some(edge.source())
                    } else {
                        Some(edge.target()) // includes case where from == source == target
                    }
                } else {
                    None
                }
            })
            .next()
    }
}

Data! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
GraphProp! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
IntoNodeIdentifiers! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
IntoNodeReferences! {delegate_impl [['a, G, F], G, &'a EdgeFiltered<G, F>, access0]}
NodeCompactIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
NodeCount! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
NodeIndexable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
Visitable! {delegate_impl [[G, F], G, EdgeFiltered<G, F>, access0]}
<|endoftext|>


<commit_before>use crate::common::initiate_validator_exit;
use types::{BeaconStateError as Error, *};

/// Slash the validator with index ``index``.
///
/// Spec v0.6.1
pub fn slash_validator<T: EthSpec>(
    state: &mut BeaconState<T>,
    slashed_index: usize,
    opt_whistleblower_index: Option<usize>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    if slashed_index >= state.validator_registry.len() || slashed_index >= state.balances.len() {
        return Err(BeaconStateError::UnknownValidator);
    }

    let current_epoch = state.current_epoch();

    initiate_validator_exit(state, slashed_index, spec)?;

    state.validator_registry[slashed_index].slashed = true;
    state.validator_registry[slashed_index].withdrawable_epoch =
        current_epoch + Epoch::from(T::latest_slashed_exit_length());
    let slashed_balance = state.get_effective_balance(slashed_index, spec)?;

    state.set_slashed_balance(
        current_epoch,
        state.get_slashed_balance(current_epoch)? + slashed_balance,
    )?;

    let proposer_index =
        state.get_beacon_proposer_index(state.slot, RelativeEpoch::Current, spec)?;
    let whistleblower_index = opt_whistleblower_index.unwrap_or(proposer_index);
    let whistleblowing_reward = slashed_balance / spec.whistleblowing_reward_quotient;
    let proposer_reward = whistleblowing_reward / spec.proposer_reward_quotient;

    safe_add_assign!(state.balances[proposer_index], proposer_reward);
    safe_add_assign!(
        state.balances[whistleblower_index],
        whistleblowing_reward - proposer_reward
    );
    safe_sub_assign!(state.balances[slashed_index], whistleblowing_reward);

    Ok(())
}
<commit_msg>Add saturating sub to slash_validator<commit_after>use crate::common::initiate_validator_exit;
use types::{BeaconStateError as Error, *};

/// Slash the validator with index ``index``.
///
/// Spec v0.6.1
pub fn slash_validator<T: EthSpec>(
    state: &mut BeaconState<T>,
    slashed_index: usize,
    opt_whistleblower_index: Option<usize>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    if slashed_index >= state.validator_registry.len() || slashed_index >= state.balances.len() {
        return Err(BeaconStateError::UnknownValidator);
    }

    let current_epoch = state.current_epoch();

    initiate_validator_exit(state, slashed_index, spec)?;

    state.validator_registry[slashed_index].slashed = true;
    state.validator_registry[slashed_index].withdrawable_epoch =
        current_epoch + Epoch::from(T::latest_slashed_exit_length());
    let slashed_balance = state.get_effective_balance(slashed_index, spec)?;

    state.set_slashed_balance(
        current_epoch,
        state.get_slashed_balance(current_epoch)? + slashed_balance,
    )?;

    let proposer_index =
        state.get_beacon_proposer_index(state.slot, RelativeEpoch::Current, spec)?;
    let whistleblower_index = opt_whistleblower_index.unwrap_or(proposer_index);
    let whistleblowing_reward = slashed_balance / spec.whistleblowing_reward_quotient;
    let proposer_reward = whistleblowing_reward / spec.proposer_reward_quotient;

    safe_add_assign!(state.balances[proposer_index], proposer_reward);
    safe_add_assign!(
        state.balances[whistleblower_index],
        whistleblowing_reward.saturating_sub(proposer_reward)
    );
    safe_sub_assign!(state.balances[slashed_index], whistleblowing_reward);

    Ok(())
}
<|endoftext|>