<commit_before>use std::collections::{HashSet, VecDeque};
use std::env;
use std::ffi::OsStr;
use std::fs;
use std::hash::Hash;
use std::io::{BufReader, BufRead, Error, ErrorKind};
use std::os::unix;
use std::path::{Component, Path, PathBuf};
use std::process::Command;

struct NonRepeatingQueue<T> {
    queue: VecDeque<T>,
    seen: HashSet<T>,
}

impl<T> NonRepeatingQueue<T> {
    fn new() -> NonRepeatingQueue<T> {
        NonRepeatingQueue {
            queue: VecDeque::new(),
            seen: HashSet::new(),
        }
    }
}

impl<T: Clone + Eq + Hash> NonRepeatingQueue<T> {
    fn push_back(&mut self, value: T) -> bool {
        if self.seen.contains(&value) {
            false
        } else {
            self.seen.insert(value.clone());
            self.queue.push_back(value);
            true
        }
    }

    fn pop_front(&mut self) -> Option<T> {
        self.queue.pop_front()
    }
}

fn patch_elf<S: AsRef<OsStr>, P: AsRef<OsStr>>(mode: S, path: P) -> Result<String, Error> {
    let output = Command::new("patchelf")
        .arg(&mode)
        .arg(&path)
        .output()?;
    if output.status.success() {
        Ok(String::from_utf8(output.stdout).expect("Failed to parse output"))
    } else {
        Err(Error::new(ErrorKind::Other, format!("failed: patchelf {:?} {:?}", OsStr::new(&mode), OsStr::new(&path))))
    }
}

fn copy_file<P: AsRef<Path> + AsRef<OsStr>, S: AsRef<Path> + AsRef<OsStr>>(
    source: P,
    target: S,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    fs::copy(&source, &target)?;

    if !Command::new("ldd").arg(&source).output()?.status.success() {
        // Not dynamically linked - no need to recurse
        return Ok(());
    }

    let rpath_string = patch_elf("--print-rpath", &source)?;
    let needed_string = patch_elf("--print-needed", &source)?;
    // Shared libraries don't have an interpreter
    if let Ok(interpreter_string) = patch_elf("--print-interpreter", &source) {
        queue.push_back(Box::from(Path::new(&interpreter_string.trim())));
    }

    let rpath = rpath_string.trim().split(":").map(|p| Box::<Path>::from(Path::new(p))).collect::<Vec<_>>();

    for line in needed_string.lines() {
        let mut found = false;
        for path in &rpath {
            let lib = path.join(line);
            if lib.exists() {
                // No need to recurse. The queue will bring it back round.
                queue.push_back(Box::from(lib.as_path()));
                found = true;
                break;
            }
        }
        if !found {
            // glibc makes it tricky to make this an error because
            // none of the files have a useful rpath.
            println!("Warning: Couldn't satisfy dependency {} for {:?}", line, OsStr::new(&source));
        }
    }

    // Make file writable to strip it
    let mut permissions = fs::metadata(&target)?.permissions();
    permissions.set_readonly(false);
    fs::set_permissions(&target, permissions)?;

    // Strip further than normal
    if !Command::new("strip").arg("--strip-all").arg(OsStr::new(&target)).output()?.status.success() {
        println!("{:?} was not successfully stripped.", OsStr::new(&target));
    }

    Ok(())
}

fn queue_dir<P: AsRef<Path>>(
    source: P,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    for entry in fs::read_dir(source)? {
        let entry = entry?;
        // No need to recurse. The queue will bring us back round here on its own.
        queue.push_back(Box::from(entry.path().as_path()));
    }

    Ok(())
}

fn handle_path(
    root: &Path,
    p: &Path,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    let mut source = PathBuf::new();
    let mut target = Path::new(root).to_path_buf();
    let mut iter = p.components().peekable();
    while let Some(comp) = iter.next() {
        match comp {
            Component::Prefix(_) => panic!("This tool is not meant for Windows"),
            Component::RootDir => {
                target.clear();
                target.push(root);
                source.clear();
                source.push("/");
            }
            Component::CurDir => {}
            Component::ParentDir => {
                // Don't over-pop the target if the path has too many ParentDirs
                if source.pop() {
                    target.pop();
                }
            }
            Component::Normal(name) => {
                target.push(name);
                source.push(name);
                let typ = fs::symlink_metadata(&source)?.file_type();
                if typ.is_file() && !target.exists() {
                    copy_file(&source, &target, queue)?;
                } else if typ.is_symlink() {
                    let link_target = fs::read_link(&source)?;

                    // Create the link, then push its target to the queue
                    if !target.exists() {
                        unix::fs::symlink(&link_target, &target)?;
                    }
                    source.pop();
                    source.push(link_target);
                    while let Some(c) = iter.next() {
                        source.push(c);
                    }
                    let link_target_path = source.as_path();
                    if link_target_path.exists() {
                        queue.push_back(Box::from(link_target_path));
                    }
                    break;
                } else if typ.is_dir() {
                    if !target.exists() {
                        fs::create_dir(&target)?;
                    }

                    // Only recursively copy if the directory is the target object
                    if iter.peek().is_none() {
                        queue_dir(&source, queue)?;
                    }
                }
            }
        }
    }

    Ok(())
}

fn main() -> Result<(), Error> {
    let args: Vec<String> = env::args().collect();
    let input = fs::File::open(&args[1])?;
    let output = &args[2];
    let out_path = Path::new(output);

    let mut queue = NonRepeatingQueue::<Box<Path>>::new();

    let mut lines = BufReader::new(input).lines();
    while let Some(obj) = lines.next() {
        // Lines should always come in pairs
        let obj = obj?;
        let sym = lines.next().unwrap()?;

        let obj_path = Path::new(&obj);
        queue.push_back(Box::from(obj_path));
        if !sym.is_empty() {
            println!("{} -> {}", &sym, &obj);
            // We don't care about preserving symlink structure here
            // nearly as much as for the actual objects.
            let link_string = format!("{}/{}", output, sym);
            let link_path = Path::new(&link_string);
            let mut link_parent = link_path.to_path_buf();
            link_parent.pop();
            fs::create_dir_all(link_parent)?;
            unix::fs::symlink(obj_path, link_path)?;
        }
    }
    while let Some(obj) = queue.pop_front() {
        handle_path(out_path, &*obj, &mut queue)?;
    }

    Ok(())
}
<commit_msg>make-initrd-ng: rustfmt<commit_after>use std::collections::{HashSet, VecDeque};
use std::env;
use std::ffi::OsStr;
use std::fs;
use std::hash::Hash;
use std::io::{BufRead, BufReader, Error, ErrorKind};
use std::os::unix;
use std::path::{Component, Path, PathBuf};
use std::process::Command;

struct NonRepeatingQueue<T> {
    queue: VecDeque<T>,
    seen: HashSet<T>,
}

impl<T> NonRepeatingQueue<T> {
    fn new() -> NonRepeatingQueue<T> {
        NonRepeatingQueue {
            queue: VecDeque::new(),
            seen: HashSet::new(),
        }
    }
}

impl<T: Clone + Eq + Hash> NonRepeatingQueue<T> {
    fn push_back(&mut self, value: T) -> bool {
        if self.seen.contains(&value) {
            false
        } else {
            self.seen.insert(value.clone());
            self.queue.push_back(value);
            true
        }
    }

    fn pop_front(&mut self) -> Option<T> {
        self.queue.pop_front()
    }
}

fn patch_elf<S: AsRef<OsStr>, P: AsRef<OsStr>>(mode: S, path: P) -> Result<String, Error> {
    let output = Command::new("patchelf").arg(&mode).arg(&path).output()?;
    if output.status.success() {
        Ok(String::from_utf8(output.stdout).expect("Failed to parse output"))
    } else {
        Err(Error::new(
            ErrorKind::Other,
            format!(
                "failed: patchelf {:?} {:?}",
                OsStr::new(&mode),
                OsStr::new(&path)
            ),
        ))
    }
}

fn copy_file<P: AsRef<Path> + AsRef<OsStr>, S: AsRef<Path> + AsRef<OsStr>>(
    source: P,
    target: S,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    fs::copy(&source, &target)?;

    if !Command::new("ldd").arg(&source).output()?.status.success() {
        // Not dynamically linked - no need to recurse
        return Ok(());
    }

    let rpath_string = patch_elf("--print-rpath", &source)?;
    let needed_string = patch_elf("--print-needed", &source)?;
    // Shared libraries don't have an interpreter
    if let Ok(interpreter_string) = patch_elf("--print-interpreter", &source) {
        queue.push_back(Box::from(Path::new(&interpreter_string.trim())));
    }

    let rpath = rpath_string
        .trim()
        .split(":")
        .map(|p| Box::<Path>::from(Path::new(p)))
        .collect::<Vec<_>>();

    for line in needed_string.lines() {
        let mut found = false;
        for path in &rpath {
            let lib = path.join(line);
            if lib.exists() {
                // No need to recurse. The queue will bring it back round.
                queue.push_back(Box::from(lib.as_path()));
                found = true;
                break;
            }
        }
        if !found {
            // glibc makes it tricky to make this an error because
            // none of the files have a useful rpath.
            println!(
                "Warning: Couldn't satisfy dependency {} for {:?}",
                line,
                OsStr::new(&source)
            );
        }
    }

    // Make file writable to strip it
    let mut permissions = fs::metadata(&target)?.permissions();
    permissions.set_readonly(false);
    fs::set_permissions(&target, permissions)?;

    // Strip further than normal
    if !Command::new("strip")
        .arg("--strip-all")
        .arg(OsStr::new(&target))
        .output()?
        .status
        .success()
    {
        println!("{:?} was not successfully stripped.", OsStr::new(&target));
    }

    Ok(())
}

fn queue_dir<P: AsRef<Path>>(
    source: P,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    for entry in fs::read_dir(source)? {
        let entry = entry?;
        // No need to recurse. The queue will bring us back round here on its own.
        queue.push_back(Box::from(entry.path().as_path()));
    }

    Ok(())
}

fn handle_path(
    root: &Path,
    p: &Path,
    queue: &mut NonRepeatingQueue<Box<Path>>,
) -> Result<(), Error> {
    let mut source = PathBuf::new();
    let mut target = Path::new(root).to_path_buf();
    let mut iter = p.components().peekable();
    while let Some(comp) = iter.next() {
        match comp {
            Component::Prefix(_) => panic!("This tool is not meant for Windows"),
            Component::RootDir => {
                target.clear();
                target.push(root);
                source.clear();
                source.push("/");
            }
            Component::CurDir => {}
            Component::ParentDir => {
                // Don't over-pop the target if the path has too many ParentDirs
                if source.pop() {
                    target.pop();
                }
            }
            Component::Normal(name) => {
                target.push(name);
                source.push(name);
                let typ = fs::symlink_metadata(&source)?.file_type();
                if typ.is_file() && !target.exists() {
                    copy_file(&source, &target, queue)?;
                } else if typ.is_symlink() {
                    let link_target = fs::read_link(&source)?;

                    // Create the link, then push its target to the queue
                    if !target.exists() {
                        unix::fs::symlink(&link_target, &target)?;
                    }
                    source.pop();
                    source.push(link_target);
                    while let Some(c) = iter.next() {
                        source.push(c);
                    }
                    let link_target_path = source.as_path();
                    if link_target_path.exists() {
                        queue.push_back(Box::from(link_target_path));
                    }
                    break;
                } else if typ.is_dir() {
                    if !target.exists() {
                        fs::create_dir(&target)?;
                    }

                    // Only recursively copy if the directory is the target object
                    if iter.peek().is_none() {
                        queue_dir(&source, queue)?;
                    }
                }
            }
        }
    }

    Ok(())
}

fn main() -> Result<(), Error> {
    let args: Vec<String> = env::args().collect();
    let input = fs::File::open(&args[1])?;
    let output = &args[2];
    let out_path = Path::new(output);

    let mut queue = NonRepeatingQueue::<Box<Path>>::new();

    let mut lines = BufReader::new(input).lines();
    while let Some(obj) = lines.next() {
        // Lines should always come in pairs
        let obj = obj?;
        let sym = lines.next().unwrap()?;

        let obj_path = Path::new(&obj);
        queue.push_back(Box::from(obj_path));
        if !sym.is_empty() {
            println!("{} -> {}", &sym, &obj);
            // We don't care about preserving symlink structure here
            // nearly as much as for the actual objects.
            let link_string = format!("{}/{}", output, sym);
            let link_path = Path::new(&link_string);
            let mut link_parent = link_path.to_path_buf();
            link_parent.pop();
            fs::create_dir_all(link_parent)?;
            unix::fs::symlink(obj_path, link_path)?;
        }
    }
    while let Some(obj) = queue.pop_front() {
        handle_path(out_path, &*obj, &mut queue)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use crate::grammar::ast::{
    BinaryExpression, BinaryOp, BooleanLit, CharLit, Expression, Identifier, NumLit, Parens,
    StringLit, Underscore,
};
use crate::grammar::model::Fragment;

use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete::char;
use nom::character::complete::space0;
use nom::combinator::map;
use nom::combinator::opt;
use nom::error::context;
use nom::multi::fold_many0;
use nom::sequence::delimited;
use nom::sequence::pair;
use nom::IResult;

impl<'s> Expression<'s> {
    /// Get the expression's fragment
    pub fn frag(&self) -> Fragment<'s> {
        match self {
            Expression::NumLit(inner) => inner.frag,
            Expression::CharLit(inner) => inner.frag,
            Expression::StringLit(inner) => inner.frag,
            Expression::BooleanLit(inner) => inner.frag,
            Expression::Identifier(inner) => inner.frag,
            Expression::Underscore(inner) => inner.frag,
            Expression::Parens(inner) => inner.frag,
            Expression::BinaryExpression(inner) => inner.frag,
        }
    }

    fn new_bin_expr(
        frag: Fragment<'s>,
        left: Expression<'s>,
        op: BinaryOp,
        right: Expression<'s>,
    ) -> Expression<'s> {
        Expression::BinaryExpression(BinaryExpression {
            frag,
            left: Box::new(left),
            op,
            right: Box::new(right),
        })
    }

    fn parse_parens(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression> {
        map(
            delimited(
                space0,
                delimited(char('('), Expression::parse, char(')')),
                space0,
            ),
            |inner| {
                Expression::Parens(Parens {
                    frag: input,
                    inner: Box::new(inner),
                })
            },
        )(input)
    }

    fn parse_atom(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression<'s>> {
        macro_rules! generalize_parser {
            ($parse_fcn:expr, $wrap_fcn:expr) => {
                |input| $parse_fcn(input).map(|(input, inner)| (input, $wrap_fcn(inner)))
            };
        }

        delimited(
            space0,
            alt((
                generalize_parser!(BooleanLit::parse, Expression::BooleanLit),
                generalize_parser!(CharLit::parse, Expression::CharLit),
                generalize_parser!(Identifier::parse, Expression::Identifier),
                generalize_parser!(NumLit::parse, Expression::NumLit),
                generalize_parser!(StringLit::parse, Expression::StringLit),
                generalize_parser!(Underscore::parse, Expression::Underscore),
            )),
            space0,
        )(input)
    }

    fn parse_factor(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression> {
        alt((Self::parse_parens, Self::parse_atom))(input)
    }

    fn parse_mul_prec(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_factor(input)?;

        fold_many0(
            pair(
                alt((
                    map(char('*'), |_| BinaryOp::Mul),
                    map(char('/'), |_| BinaryOp::Div),
                    map(char('%'), |_| BinaryOp::Mod),
                )),
                Self::parse_factor,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_add_prec(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_mul_prec(input)?;

        fold_many0(
            pair(
                alt((
                    map(char('+'), |_| BinaryOp::Add),
                    map(char('-'), |_| BinaryOp::Sub),
                )),
                Self::parse_mul_prec,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_and(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_add_prec(input)?;

        fold_many0(
            pair(map(char('&'), |_| BinaryOp::And), Self::parse_add_prec),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_xor(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_and(input)?;

        fold_many0(
            pair(map(char('^'), |_| BinaryOp::Xor), Self::parse_and),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_or(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_xor(input)?;

        fold_many0(
            pair(map(char('|'), |_| BinaryOp::Or), Self::parse_xor),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_cmp(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_or(input)?;

        fold_many0(
            pair(
                alt((
                    map(tag("<="), |_| BinaryOp::Le),
                    map(tag(">="), |_| BinaryOp::Ge),
                    map(tag("=="), |_| BinaryOp::Eq),
                    map(tag("!="), |_| BinaryOp::NotEq),
                    map(tag("<"), |_| BinaryOp::Lt),
                    map(tag(">"), |_| BinaryOp::Gt),
                )),
                Self::parse_or,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_and_and(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_cmp(input)?;

        fold_many0(
            pair(map(tag("&&"), |_| BinaryOp::AndAnd), Self::parse_cmp),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_or_or(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_and_and(input)?;

        fold_many0(
            pair(map(tag("||"), |_| BinaryOp::OrOr), Self::parse_and_and),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_dot_dot(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(
            pair(
                Self::parse_or_or,
                opt(pair(
                    map(tag(".."), |_| BinaryOp::DotDot),
                    Self::parse_or_or,
                )),
            ),
            |(left, opt)| match opt {
                Some((op, right)) => Self::new_bin_expr(input, left, op, right),
                None => left,
            },
        )(input)
    }

    fn parse_eq(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(
            pair(
                Self::parse_dot_dot,
                opt(pair(map(char('='), |_| BinaryOp::Eq), Self::parse_dot_dot)),
            ),
            |(left, opt)| match opt {
                Some((op, right)) => Self::new_bin_expr(input, left, op, right),
                None => left,
            },
        )(input)
    }

    /// Parse a binary expression
    pub fn parse(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        context("expected expression", Self::parse_eq)(input)
    }
}
<commit_msg>Update comment<commit_after>use crate::grammar::ast::{
    BinaryExpression, BinaryOp, BooleanLit, CharLit, Expression, Identifier, NumLit, Parens,
    StringLit, Underscore,
};
use crate::grammar::model::Fragment;

use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete::char;
use nom::character::complete::space0;
use nom::combinator::map;
use nom::combinator::opt;
use nom::error::context;
use nom::multi::fold_many0;
use nom::sequence::delimited;
use nom::sequence::pair;
use nom::IResult;

impl<'s> Expression<'s> {
    /// Get the expression's fragment
    pub fn frag(&self) -> Fragment<'s> {
        match self {
            Expression::NumLit(inner) => inner.frag,
            Expression::CharLit(inner) => inner.frag,
            Expression::StringLit(inner) => inner.frag,
            Expression::BooleanLit(inner) => inner.frag,
            Expression::Identifier(inner) => inner.frag,
            Expression::Underscore(inner) => inner.frag,
            Expression::Parens(inner) => inner.frag,
            Expression::BinaryExpression(inner) => inner.frag,
        }
    }

    fn new_bin_expr(
        frag: Fragment<'s>,
        left: Expression<'s>,
        op: BinaryOp,
        right: Expression<'s>,
    ) -> Expression<'s> {
        Expression::BinaryExpression(BinaryExpression {
            frag,
            left: Box::new(left),
            op,
            right: Box::new(right),
        })
    }

    fn parse_parens(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression> {
        map(
            delimited(
                space0,
                delimited(char('('), Expression::parse, char(')')),
                space0,
            ),
            |inner| {
                Expression::Parens(Parens {
                    frag: input,
                    inner: Box::new(inner),
                })
            },
        )(input)
    }

    fn parse_atom(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression<'s>> {
        macro_rules! generalize_parser {
            ($parse_fcn:expr, $wrap_fcn:expr) => {
                |input| $parse_fcn(input).map(|(input, inner)| (input, $wrap_fcn(inner)))
            };
        }

        delimited(
            space0,
            alt((
                generalize_parser!(BooleanLit::parse, Expression::BooleanLit),
                generalize_parser!(CharLit::parse, Expression::CharLit),
                generalize_parser!(Identifier::parse, Expression::Identifier),
                generalize_parser!(NumLit::parse, Expression::NumLit),
                generalize_parser!(StringLit::parse, Expression::StringLit),
                generalize_parser!(Underscore::parse, Expression::Underscore),
            )),
            space0,
        )(input)
    }

    fn parse_factor(input: Fragment<'s>) -> IResult<Fragment<'s>, Expression> {
        alt((Self::parse_parens, Self::parse_atom))(input)
    }

    fn parse_mul_prec(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_factor(input)?;

        fold_many0(
            pair(
                alt((
                    map(char('*'), |_| BinaryOp::Mul),
                    map(char('/'), |_| BinaryOp::Div),
                    map(char('%'), |_| BinaryOp::Mod),
                )),
                Self::parse_factor,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_add_prec(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_mul_prec(input)?;

        fold_many0(
            pair(
                alt((
                    map(char('+'), |_| BinaryOp::Add),
                    map(char('-'), |_| BinaryOp::Sub),
                )),
                Self::parse_mul_prec,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_and(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_add_prec(input)?;

        fold_many0(
            pair(map(char('&'), |_| BinaryOp::And), Self::parse_add_prec),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_xor(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_and(input)?;

        fold_many0(
            pair(map(char('^'), |_| BinaryOp::Xor), Self::parse_and),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_or(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_xor(input)?;

        fold_many0(
            pair(map(char('|'), |_| BinaryOp::Or), Self::parse_xor),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_cmp(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_or(input)?;

        fold_many0(
            pair(
                alt((
                    map(tag("<="), |_| BinaryOp::Le),
                    map(tag(">="), |_| BinaryOp::Ge),
                    map(tag("=="), |_| BinaryOp::Eq),
                    map(tag("!="), |_| BinaryOp::NotEq),
                    map(tag("<"), |_| BinaryOp::Lt),
                    map(tag(">"), |_| BinaryOp::Gt),
                )),
                Self::parse_or,
            ),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_and_and(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_cmp(input)?;

        fold_many0(
            pair(map(tag("&&"), |_| BinaryOp::AndAnd), Self::parse_cmp),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_or_or(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        let (input, left) = Self::parse_and_and(input)?;

        fold_many0(
            pair(map(tag("||"), |_| BinaryOp::OrOr), Self::parse_and_and),
            left,
            |left, (op, right)| Self::new_bin_expr(left.frag(), left, op, right),
        )(input)
    }

    fn parse_dot_dot(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(
            pair(
                Self::parse_or_or,
                opt(pair(
                    map(tag(".."), |_| BinaryOp::DotDot),
                    Self::parse_or_or,
                )),
            ),
            |(left, opt)| match opt {
                Some((op, right)) => Self::new_bin_expr(input, left, op, right),
                None => left,
            },
        )(input)
    }

    fn parse_eq(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        map(
            pair(
                Self::parse_dot_dot,
                opt(pair(map(char('='), |_| BinaryOp::Eq), Self::parse_dot_dot)),
            ),
            |(left, opt)| match opt {
                Some((op, right)) => Self::new_bin_expr(input, left, op, right),
                None => left,
            },
        )(input)
    }

    /// Parse an expression
    pub fn parse(input: Fragment<'s>) -> IResult<Fragment<'s>, Self> {
        context("expected expression", Self::parse_eq)(input)
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Pareto distribution.

use Rng;
use distributions::{Distribution, OpenClosed01};

/// Samples floating-point numbers according to the Pareto distribution
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand::distributions::Pareto;
///
/// let val: f64 = SmallRng::from_entropy().sample(Pareto::new(1., 2.));
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Pareto {
    scale: f64,
    shape: f64,
}

impl Pareto {
    /// Construct a new Pareto distribution with given `scale` and `shape`.
    ///
    /// In the literature, `scale` is commonly written as xm or k and
    /// `shape` is often written as α.
    ///
    /// # Panics
    ///
    /// `scale` and `shape` have to be non-zero and positive.
    pub fn new(scale: f64, shape: f64) -> Pareto {
        assert!((scale > 0.) & (shape > 0.));
        Pareto { scale, shape }
    }
}

impl Distribution<f64> for Pareto {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        let u: f64 = rng.sample(OpenClosed01);
        self.scale * u.powf(-1.0 / self.shape)
    }
}

#[cfg(test)]
mod tests {
    use distributions::Distribution;
    use super::Pareto;

    #[test]
    #[should_panic]
    fn invalid() {
        Pareto::new(0., 0.);
    }

    #[test]
    fn sample() {
        let d = Pareto::new(1.0, 2.0);
        let mut rng = ::test::rng(1);
        for _ in 0..1000 {
            d.sample(&mut rng);
        }
    }
}
<commit_msg>Make `Pareto::sample` faster by moving code into the constructor<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Pareto distribution.

use Rng;
use distributions::{Distribution, OpenClosed01};

/// Samples floating-point numbers according to the Pareto distribution
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand::distributions::Pareto;
///
/// let val: f64 = SmallRng::from_entropy().sample(Pareto::new(1., 2.));
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Pareto {
    scale: f64,
    inv_neg_shape: f64,
}

impl Pareto {
    /// Construct a new Pareto distribution with given `scale` and `shape`.
    ///
    /// In the literature, `scale` is commonly written as xm or k and
    /// `shape` is often written as α.
    ///
    /// # Panics
    ///
    /// `scale` and `shape` have to be non-zero and positive.
    pub fn new(scale: f64, shape: f64) -> Pareto {
        assert!((scale > 0.) & (shape > 0.));
        Pareto { scale, inv_neg_shape: -1.0 / shape }
    }
}

impl Distribution<f64> for Pareto {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        let u: f64 = rng.sample(OpenClosed01);
        self.scale * u.powf(self.inv_neg_shape)
    }
}

#[cfg(test)]
mod tests {
    use distributions::Distribution;
    use super::Pareto;

    #[test]
    #[should_panic]
    fn invalid() {
        Pareto::new(0., 0.);
    }

    #[test]
    fn sample() {
        let d = Pareto::new(1.0, 2.0);
        let mut rng = ::test::rng(1);
        for _ in 0..1000 {
            d.sample(&mut rng);
        }
    }
}
<|endoftext|>


<commit_before>// BatWarn
// Low Battery Warning Utility
// Polls battery status and pops up warning when low.

#![allow(unstable)]
extern crate collections;
extern crate regex;
use std::io::{IoResult, IoError};
use std::io::timer::sleep;
use std::io::fs::File;
use std::io::process::{Command, Process, ProcessOutput};
use std::time::duration::Duration;
use std::num::from_str_radix;
use collections::string::String;
use regex::Regex;

static PERCENT_DANGER:     i32 = 20;
static PERCENT_CRITICAL:   i32 = 8;

#[derive(Debug)]
struct BatteryState {
    discharging: bool,
    percent: i32,
}

fn main() {
    let poll_delay: Duration = Duration::minutes(5);

    let mut nagproc: Option<Process> = None;

    loop {
        // Kill existing warnings
        nagproc = nagproc.and_then(|mut nagproc| {
            match nagproc.signal_kill() {
                Err(err) => {
                    println!("ERROR: {}", err);
                    Some(nagproc)
                },
                Ok(()) => None
            }
        });

        // Check battery status
        match acpi_battery_state() {
            Err(err) => {
                println!("ERROR: {}", err)
            },
            Ok(batstat) => {
                println!("{:?}", batstat);

                if !batstat.discharging {
                    // Do nothing
                } else if batstat.percent <= PERCENT_CRITICAL {
                    // Battery critically low.
                    let msg = format!("WARNING: Battery critically low!");
                    match show_warning(msg, true) {
                        Err(err) =>
                            println!("ERROR: {}", err),
                        Ok(child) =>
                            nagproc = Some(child),
                    };
                } else if batstat.percent <= PERCENT_DANGER {
                    // Battery low.
                    let msg = format!("WARNING: Battery low!");
                    match show_warning(msg, false) {
                        Err(err) =>
                            println!("ERROR: {}", err),
                        Ok(child) =>
                            nagproc = Some(child),
                    };
                } else {
                    // Battery discharging normally.
                }

            },
        }

        sleep(poll_delay);
    }
}

fn show_warning(msg: String, critical: bool) -> Result<Process, IoError> {
    let warning_level = match critical {
        false => "warning",
        true => "error",
    };
    Command::new("i3-nagbar")
        .arg("-t")
        .arg(warning_level)
        .arg(format!("-m {}", msg))
        .spawn()
}

// Parse output from `acpi --battery`
fn acpi_battery_state() -> Result<BatteryState, String> {
    // ACPI output format: "Battery #{number}: #{state}, #{percent}%..."
    match acpi_battery_string() {
        Err(err) => Err(err),
        Ok(s) => {
            let re = Regex::new(r"Battery \d+: (\w+), (\d+)%.*").unwrap();
            match re.captures(s.as_slice()) {
                None => Err(String::from_str("malformed acpi output")),
                Some(captures) => {
                    let state = captures.at(1).unwrap();
                    let percent_str = captures.at(2).unwrap();
                    let percent = from_str_radix::<i32>(percent_str, 10).unwrap();
                    Ok(BatteryState {
                        discharging: state == "Discharging",
                        percent: percent,})
                }
            }
        },
    }
}

// Get output from `acpi --battery`
fn acpi_battery_string() -> Result<String, String> {
    let mut cmd = Command::new("acpi");
    cmd.arg("--battery");
    match cmd.output() {
        Err(err) => Err(String::from_str(err.desc)),
        Ok(ProcessOutput { status: exit, output: stdout, error: _ }) => {
            match exit.success() {
                false => Err(String::from_str("acpi returned with non-zero exit status")),
                true => {
                    let stdout = String::from_utf8(stdout).unwrap();
                    Ok(stdout)
                },
            }
        },
    }
}

// Read battery state from /proc.
#[allow(dead_code)]
fn read_battery_state_string() -> IoResult<String> {
    let path = Path::new("/proc/acpi/battery/BAT0/state");
    match File::open(&path) {
        Err(err) => Err(err),
        Ok(mut file) => file.read_to_string(),
    }
}

<commit_msg>add comment<commit_after>// BatWarn
// Low Battery Warning Utility
// Polls battery status and pops up warning when low.

#![allow(unstable)]
extern crate collections;
extern crate regex;
use std::io::{IoResult, IoError};
use std::io::timer::sleep;
use std::io::fs::File;
use std::io::process::{Command, Process, ProcessOutput};
use std::time::duration::Duration;
use std::num::from_str_radix;
use collections::string::String;
use regex::Regex;

static PERCENT_DANGER:     i32 = 20;
static PERCENT_CRITICAL:   i32 = 8;

#[derive(Debug)]
struct BatteryState {
    discharging: bool,
    percent: i32,
}

fn main() {
    let poll_delay: Duration = Duration::minutes(5);

    let mut nagproc: Option<Process> = None;

    loop {
        // Kill existing warnings
        nagproc = nagproc.and_then(|mut nagproc| {
            match nagproc.signal_kill() {
                Err(err) => {
                    // This can occur if the process was closed by the user.
                    println!("ERROR: {}", err);
                    Some(nagproc)
                },
                Ok(()) => None
            }
        });

        // Check battery status
        match acpi_battery_state() {
            Err(err) => {
                println!("ERROR: {}", err)
            },
            Ok(batstat) => {
                println!("{:?}", batstat);

                if !batstat.discharging {
                    // Do nothing
                } else if batstat.percent <= PERCENT_CRITICAL {
                    // Battery critically low.
                    let msg = format!("WARNING: Battery critically low!");
                    match show_warning(msg, true) {
                        Err(err) =>
                            println!("ERROR: {}", err),
                        Ok(child) =>
                            nagproc = Some(child),
                    };
                } else if batstat.percent <= PERCENT_DANGER {
                    // Battery low.
                    let msg = format!("WARNING: Battery low!");
                    match show_warning(msg, false) {
                        Err(err) =>
                            println!("ERROR: {}", err),
                        Ok(child) =>
                            nagproc = Some(child),
                    };
                } else {
                    // Battery discharging normally.
                }

            },
        }

        sleep(poll_delay);
    }
}

fn show_warning(msg: String, critical: bool) -> Result<Process, IoError> {
    let warning_level = match critical {
        false => "warning",
        true => "error",
    };
    Command::new("i3-nagbar")
        .arg("-t")
        .arg(warning_level)
        .arg(format!("-m {}", msg))
        .spawn()
}

// Parse output from `acpi --battery`
fn acpi_battery_state() -> Result<BatteryState, String> {
    // ACPI output format: "Battery #{number}: #{state}, #{percent}%..."
    match acpi_battery_string() {
        Err(err) => Err(err),
        Ok(s) => {
            let re = Regex::new(r"Battery \d+: (\w+), (\d+)%.*").unwrap();
            match re.captures(s.as_slice()) {
                None => Err(String::from_str("malformed acpi output")),
                Some(captures) => {
                    let state = captures.at(1).unwrap();
                    let percent_str = captures.at(2).unwrap();
                    let percent = from_str_radix::<i32>(percent_str, 10).unwrap();
                    Ok(BatteryState {
                        discharging: state == "Discharging",
                        percent: percent,})
                }
            }
        },
    }
}

// Get output from `acpi --battery`
fn acpi_battery_string() -> Result<String, String> {
    let mut cmd = Command::new("acpi");
    cmd.arg("--battery");
    match cmd.output() {
        Err(err) => Err(String::from_str(err.desc)),
        Ok(ProcessOutput { status: exit, output: stdout, error: _ }) => {
            match exit.success() {
                false => Err(String::from_str("acpi returned with non-zero exit status")),
                true => {
                    let stdout = String::from_utf8(stdout).unwrap();
                    Ok(stdout)
                },
            }
        },
    }
}

// Read battery state from /proc.
#[allow(dead_code)]
fn read_battery_state_string() -> IoResult<String> {
    let path = Path::new("/proc/acpi/battery/BAT0/state");
    match File::open(&path) {
        Err(err) => Err(err),
        Ok(mut file) => file.read_to_string(),
    }
}

<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Calculate alignments with a generalized variant of the Smith Waterman algorithm.
//! Complexity: O(n * m) for strings of length m and n.
//!
//! # Example
//!
//! ```
//! use bio::alignment::pairwise::Aligner;
//! use bio::alignment::AlignmentOperation::{Match, Subst, Ins, Del};
//! let x = b"ACCGTGGAT";
//! let y = b"AAAAACCGTTGAT";
//! let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
//! let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
//! let alignment = aligner.semiglobal(x, y);
//! assert_eq!(alignment.ystart, 4);
//! assert_eq!(alignment.xstart, 0);
//! assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
//! ```

use std::i32;
use std::iter::repeat;
use std::cmp::max;

use alignment::{Alignment, AlignmentOperation};
use data_structures::bitenc::BitEnc;

/// Current internal state of alignment.
struct AlignmentState {
    m: usize,
    n: usize,
    best: i32,
    best_i: usize,
    best_j: usize,
    i: usize,
    j: usize,
    score: i32,
    col: usize
}

macro_rules! align {
    (
        $aligner:ident, $x:ident, $y:ident, $state:ident,
        $init:block, $inner:block, $outer:block, $ret:block
    ) => (
        {
            let mut $state = AlignmentState {
                m: $x.len(), n: $y.len(),
                best: 0, best_i: 0, best_j: 0,
                i: 1, j: 1,
                score: 0,
                col: 0
            };

            while $state.i <= $state.n {
                $state.col = $state.i % 2;
                let prev = 1 - $state.col;

                // init code
                $init

                // read next y symbol
                let b = $y[$state.i - 1];
                $state.j = 1;

                while $state.j <= $state.m {
                    // read next x symbol
                    let a = $x[$state.j - 1];

                    // score for deletion
                    let d_score = max(
                        $aligner.S[prev][$state.j] + $aligner.gap_open,
                        $aligner.D[prev][$state.j] + $aligner.gap_extend
                    );
                    // score for insertion
                    let i_score = max(
                        $aligner.S[$state.col][$state.j-1] + $aligner.gap_open,
                        $aligner.I[$state.col][$state.j-1] + $aligner.gap_extend
                    );
                    // score for substitution
                    $state.score = $aligner.S[prev][$state.j-1] + ($aligner.score)(a, b);

                    if d_score > $state.score {
                        $state.score = d_score;
                        $aligner.traceback.del($state.i, $state.j);
                    }
                    else if i_score > $state.score {
                        $state.score = i_score;
                        $aligner.traceback.ins($state.i, $state.j);
                    }
                    else {
                        $aligner.traceback.subst($state.i, $state.j);
                    }

                    // inner code
                    $inner

                    $aligner.S[$state.col][$state.j] = $state.score;
                    $aligner.D[$state.col][$state.j] = d_score;
                    $aligner.I[$state.col][$state.j] = i_score;

                    $state.j += 1;
                }

                // outer code
                $outer

                $state.i += 1;
            }

            // return code
            $ret
        }
    );
}

/// A generalized Smith-Waterman aligner.
#[allow(non_snake_case)]
pub struct Aligner<F> where F: Fn(u8, u8) -> i32 {
    S: [Vec<i32>; 2],
    I: [Vec<i32>; 2],
    D: [Vec<i32>; 2],
    traceback: Traceback,
    gap_open: i32,
    gap_extend: i32,
    score: F,
}

impl<F> Aligner<F> where F: Fn(u8, u8) -> i32 {
    /// Create new aligner instance. The size hints help to
    /// avoid unnecessary memory allocations.
    ///
    /// # Arguments
    ///
    /// * `m` - the expected size of x
    /// * `n` - the expected size of y
    /// * `gap_open` - the score for opening a gap (should be negative)
    /// * `gap_extend` - the score for extending a gap (should be negative)
    /// * `score` - function that returns the score for substitutions
    ///
    pub fn with_capacity(m: usize, n: usize, gap_open: i32, gap_extend: i32, score: F) -> Self {
        let get_vec = || Vec::with_capacity(m + 1);
        Aligner {
            S: [get_vec(), get_vec()],
            I: [get_vec(), get_vec()],
            D: [get_vec(), get_vec()],
            traceback: Traceback::with_capacity(m, n),
            gap_open: gap_open,
            gap_extend: gap_extend,
            score: score
        }
    }

    fn init(&mut self, m: usize, global: bool) {
        // set minimum score to -inf, and allow to add gap_extend
        // without overflow
        let min_score = i32::MIN - self.gap_extend;
        for k in 0..2 {
            self.S[k].clear();
            self.I[k].clear();
            self.D[k].clear();
            self.I[k].extend(repeat(min_score).take(m + 1));
            self.D[k].extend(repeat(min_score).take(m + 1));
            if global {
                let ref mut s = self.S[k];
                let mut score = self.gap_open;
                for _ in 0..m+1 {
                    s.push(score);
                    score += self.gap_extend;
                }
            }
            else {
                self.S[k].extend(repeat(0).take(m + 1))
            }
        }
    }

    /// Calculate global alignment.
    pub fn global(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, true);
        self.traceback.init(m, n, true);

        align!(
            self, x, y, state,
            {
                self.S[state.col][0] = self.gap_open + (state.i as i32 - 1) * self.gap_extend;
                self.traceback.del(state.i, 0);
            }, {}, {},
            {
                self.traceback.alignment(state.n, state.m, x, y, state.score)
            }
        )
    }

    /// Calculate semiglobal alignment.
    pub fn semiglobal(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, false);
        self.traceback.init(m, n, false);

        align!(
            self, x, y, state,
            { self.S[state.col][0] = 0; },
            {},
            {
                if state.score > state.best {
                    state.best = state.score;
                    state.best_i = state.i;
                    state.best_j = state.m;
                }
            },
            { self.traceback.alignment(state.best_i, state.best_j, x, y, state.best) }
        )
    }

    /// Calculate local alignment.
    pub fn local(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, false);
        self.traceback.init(m, n, false);

        align!(
            self, x, y, state,
            { self.S[state.col][0] = 0; },
            {
                if state.score < 0 {
                    self.traceback.start(state.i, state.j);
                    state.score = 0;
                }
                else if state.score > state.best {
                    state.best = state.score;
                    state.best_i = state.i;
                    state.best_j = state.j;
                }
            },
            {},
            { self.traceback.alignment(state.best_i, state.best_j, x, y, state.best) }
        )
    }
}

/// Internal traceback.
struct Traceback {
    matrix: Vec<BitEnc>
}

const TBSTART: u8 = 0b00;
const TBSUBST: u8 = 0b01;
const TBINS: u8   = 0b10;
const TBDEL: u8   = 0b11;

impl Traceback {

    fn with_capacity(m: usize, n: usize) -> Self {
        let mut matrix = Vec::with_capacity(n+1);
        for _ in 0..n+1 {
            matrix.push(BitEnc::with_capacity(2, m + 1));
        }
        Traceback {
            matrix: matrix
        }
    }

    fn init(&mut self, m: usize, n: usize, global: bool) {
        if global {
            // set the first cell to start, the rest to deletions
            for i in 0..n+1 {
                self.matrix[i].clear();
                self.matrix[i].push_values(m + 1, TBDEL);
            }
            self.matrix[0].set(0, TBSTART);
        }
        else {
            for i in 0..n+1 {
                self.matrix[i].clear();
                self.matrix[i].push_values(m + 1, TBSTART);
            }
        }
    }

    fn start(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBSTART);
    }

    fn subst(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBSUBST);
    }

    fn del(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBDEL);
    }

    fn ins(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBINS);
    }

    fn get(&self, i: usize, j: usize) -> u8 {
        self.matrix[i].get(j).unwrap()
    }

    fn alignment(&self, mut i: usize, mut j: usize, x: &[u8], y: &[u8], score: i32) -> Alignment {
        let mut ops = Vec::with_capacity(x.len());

        loop {
            let (ii, jj, op) = match self.get(i, j) {
                TBSUBST => {
                    let op = if y[i-1] == x[j-1] {
                        AlignmentOperation::Match
                    }
                    else {
                        AlignmentOperation::Subst
                    };
                    (i - 1, j - 1, op)
                }
                TBDEL => {
                    (i - 1, j, AlignmentOperation::Del)
                }
                TBINS => {
                    (i, j - 1, AlignmentOperation::Ins)
                }
                _ => {
                    break;
                }
            };

            ops.push(op);
            i = ii;
            j = jj;
        }

        ops.reverse();
        Alignment { ystart: i, xstart: j, xlen: x.len(), operations: ops, score: score}
    }
}

#[cfg(test)]
mod tests {
    use super::Aligner;
    use alignment::AlignmentOperation::{Match, Subst, Del};

    #[test]
    fn test_semiglobal() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.semiglobal(x, y);
        println!("{:?}", alignment);
        assert_eq!(alignment.ystart, 4);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }

    #[test]
    fn test_local() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.local(x, y);
        assert_eq!(alignment.ystart, 4);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }

    #[test]
    fn test_global() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.global(x, y);
        assert_eq!(alignment.ystart, 0);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Del, Del, Del, Del, Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }
}
<commit_msg>Fixed for issue #11<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Calculate alignments with a generalized variant of the Smith Waterman algorithm.
//! Complexity: O(n * m) for strings of length m and n.
//!
//! # Example
//!
//! ```
//! use bio::alignment::pairwise::Aligner;
//! use bio::alignment::AlignmentOperation::{Match, Subst, Ins, Del};
//! let x = b"ACCGTGGAT";
//! let y = b"AAAAACCGTTGAT";
//! let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
//! let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
//! let alignment = aligner.semiglobal(x, y);
//! assert_eq!(alignment.ystart, 4);
//! assert_eq!(alignment.xstart, 0);
//! assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
//! ```

use std::i32;
use std::iter::repeat;
use std::cmp::max;

use alignment::{Alignment, AlignmentOperation};
use data_structures::bitenc::BitEnc;

/// Current internal state of alignment.
struct AlignmentState {
    m: usize,
    n: usize,
    best: i32,
    best_i: usize,
    best_j: usize,
    i: usize,
    j: usize,
    score: i32,
    col: usize
}

macro_rules! align {
    (
        $aligner:ident, $x:ident, $y:ident, $state:ident,
        $init:block, $inner:block, $outer:block, $ret:block
    ) => (
        {
            let mut $state = AlignmentState {
                m: $x.len(), n: $y.len(),
                best: 0, best_i: 0, best_j: 0,
                i: 1, j: 1,
                score: 0,
                col: 0
            };

            while $state.i <= $state.n {
                $state.col = $state.i % 2;
                let prev = 1 - $state.col;

                // init code
                $init

                // read next y symbol
                let b = $y[$state.i - 1];
                $state.j = 1;

                while $state.j <= $state.m {
                    // read next x symbol
                    let a = $x[$state.j - 1];

                    // score for deletion
                    let d_score = max(
                        $aligner.S[prev][$state.j] + $aligner.gap_open,
                        $aligner.D[prev][$state.j] + $aligner.gap_extend
                    );
                    // score for insertion
                    let i_score = max(
                        $aligner.S[$state.col][$state.j-1] + $aligner.gap_open,
                        $aligner.I[$state.col][$state.j-1] + $aligner.gap_extend
                    );
                    // score for substitution
                    $state.score = $aligner.S[prev][$state.j-1] + ($aligner.score)(a, b);

                    if d_score > $state.score {
                        $state.score = d_score;
                        $aligner.traceback.del($state.i, $state.j);
                        println!("del {} {}", $state.i, $state.j);
                    }
                    else if i_score > $state.score {
                        $state.score = i_score;
                        $aligner.traceback.ins($state.i, $state.j);
                        println!("ins {} {}", $state.i, $state.j);
                    }
                    else {
                        $aligner.traceback.subst($state.i, $state.j);
                        println!("subst {} {}", $state.i, $state.j);
                    }

                    // inner code
                    $inner

                    $aligner.S[$state.col][$state.j] = $state.score;
                    $aligner.D[$state.col][$state.j] = d_score;
                    $aligner.I[$state.col][$state.j] = i_score;

                    $state.j += 1;
                }

                // outer code
                $outer

                $state.i += 1;
            }

            // return code
            $ret
        }
    );
}

/// A generalized Smith-Waterman aligner.
#[allow(non_snake_case)]
pub struct Aligner<F> where F: Fn(u8, u8) -> i32 {
    S: [Vec<i32>; 2],
    I: [Vec<i32>; 2],
    D: [Vec<i32>; 2],
    traceback: Traceback,
    gap_open: i32,
    gap_extend: i32,
    score: F,
}

impl<F> Aligner<F> where F: Fn(u8, u8) -> i32 {
    /// Create new aligner instance. The size hints help to
    /// avoid unnecessary memory allocations.
    ///
    /// # Arguments
    ///
    /// * `m` - the expected size of x
    /// * `n` - the expected size of y
    /// * `gap_open` - the score for opening a gap (should be negative)
    /// * `gap_extend` - the score for extending a gap (should be negative)
    /// * `score` - function that returns the score for substitutions
    ///
    pub fn with_capacity(m: usize, n: usize, gap_open: i32, gap_extend: i32, score: F) -> Self {
        let get_vec = || Vec::with_capacity(m + 1);
        Aligner {
            S: [get_vec(), get_vec()],
            I: [get_vec(), get_vec()],
            D: [get_vec(), get_vec()],
            traceback: Traceback::with_capacity(m, n),
            gap_open: gap_open,
            gap_extend: gap_extend,
            score: score
        }
    }

    fn init(&mut self, m: usize, global: bool) {
        // set minimum score to -inf, and allow to add gap_extend
        // without overflow
        let min_score = i32::MIN - self.gap_extend;
        for k in 0..2 {
            self.S[k].clear();
            self.I[k].clear();
            self.D[k].clear();
            self.I[k].extend(repeat(min_score).take(m + 1));
            self.D[k].extend(repeat(min_score).take(m + 1));
            if global {
                let ref mut s = self.S[k];
                let mut score = self.gap_open;
                for _ in 0..m+1 {
                    s.push(score);
                    score += self.gap_extend;
                }
            }
            else {
                self.S[k].extend(repeat(0).take(m + 1))
            }
        }
    }

    /// Calculate global alignment.
    pub fn global(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, true);
        self.traceback.init(m, n, true);

        align!(
            self, x, y, state,
            {
                self.S[state.col][0] = self.gap_open + (state.i as i32 - 1) * self.gap_extend;
                self.traceback.del(state.i, 0);
            }, {}, {},
            {
                self.traceback.alignment(state.n, state.m, x, y, state.score)
            }
        )
    }

    /// Calculate semiglobal alignment.
    pub fn semiglobal(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, false);
        self.traceback.init(m, n, false);

        align!(
            self, x, y, state,
            { self.S[state.col][0] = 0; },
            {},
            {
                if state.score > state.best {
                    state.best = state.score;
                    state.best_i = state.i;
                    state.best_j = state.m;
                }
            },
            { self.traceback.alignment(state.best_i, state.best_j, x, y, state.best) }
        )
    }

    /// Calculate local alignment.
    pub fn local(&mut self, x: &[u8], y: &[u8]) -> Alignment {
        let (m, n) = (x.len(), y.len());
        self.init(m, false);
        self.traceback.init(m, n, false);

        align!(
            self, x, y, state,
            { self.S[state.col][0] = 0; },
            {
                if state.score < 0 {
                    self.traceback.start(state.i, state.j);
                    state.score = 0;
                }
                else if state.score > state.best {
                    state.best = state.score;
                    state.best_i = state.i;
                    state.best_j = state.j;
                }
            },
            {},
            { self.traceback.alignment(state.best_i, state.best_j, x, y, state.best) }
        )
    }
}

/// Internal traceback.
struct Traceback {
    matrix: Vec<BitEnc>
}

const TBSTART: u8 = 0b00;
const TBSUBST: u8 = 0b01;
const TBINS: u8   = 0b10;
const TBDEL: u8   = 0b11;

impl Traceback {

    fn with_capacity(m: usize, n: usize) -> Self {
        let mut matrix = Vec::with_capacity(n+1);
        for _ in 0..n+1 {
            matrix.push(BitEnc::with_capacity(2, m + 1));
        }
        Traceback {
            matrix: matrix
        }
    }

    fn init(&mut self, m: usize, n: usize, global: bool) {
        if global {
            // set the first cell to start, the rest to insertions
            for i in 0..n+1 {
                self.matrix[i].clear();
                self.matrix[i].push_values(m + 1, TBINS);
            }
            self.matrix[0].set(0, TBSTART);
        }
        else {
            for i in 0..n+1 {
                self.matrix[i].clear();
                self.matrix[i].push_values(m + 1, TBSTART);
            }
        }
    }

    fn start(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBSTART);
    }

    fn subst(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBSUBST);
    }

    fn del(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBDEL);
    }

    fn ins(&mut self, i: usize, j: usize) {
        self.matrix[i].set(j, TBINS);
    }

    fn get(&self, i: usize, j: usize) -> u8 {
        self.matrix[i].get(j).unwrap()
    }

    fn alignment(&self, mut i: usize, mut j: usize, x: &[u8], y: &[u8], score: i32) -> Alignment {
        let mut ops = Vec::with_capacity(x.len());

        loop {
            let (ii, jj, op) = match self.get(i, j) {
                TBSUBST => {
                    let op = if y[i-1] == x[j-1] {
                        println!("tbmatch {} {}", i, j);
                        AlignmentOperation::Match
                    }
                    else {
                        println!("tbsubst {} {}", i, j);
                        AlignmentOperation::Subst
                    };
                    (i - 1, j - 1, op)
                }
                TBDEL => {
                    println!("tbdel {} {}", i, j);
                    (i - 1, j, AlignmentOperation::Del)
                }
                TBINS => {
                    println!("tbins {} {}", i, j);
                    (i, j - 1, AlignmentOperation::Ins)
                }
                _ => {
                    break;
                }
            };

            ops.push(op);
            i = ii;
            j = jj;
        }

        ops.reverse();
        Alignment { ystart: i, xstart: j, xlen: x.len(), operations: ops, score: score}
    }
}

#[cfg(test)]
mod tests {
    use super::Aligner;
    use alignment::AlignmentOperation::{Match, Subst, Del, Ins};

    #[test]
    fn test_semiglobal() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.semiglobal(x, y);
        println!("{:?}", alignment);
        assert_eq!(alignment.ystart, 4);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }

    #[test]
    fn test_local() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.local(x, y);
        assert_eq!(alignment.ystart, 4);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }

    #[test]
    fn test_global() {
        let x = b"ACCGTGGAT";
        let y = b"AAAAACCGTTGAT";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.global(x, y);
        assert_eq!(alignment.ystart, 0);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Del, Del, Del, Del, Match, Match, Match, Match, Match, Subst, Match, Match, Match]);
    }

    #[test]
    fn test_issue_11() {
        let y = b"TACC";//GTGGAC";
        let x = b"AAAAACC";//GTTGACGCAA";
        let score = |a: u8, b: u8| if a == b {1i32} else {-1i32};
        let mut aligner = Aligner::with_capacity(x.len(), y.len(), -5, -1, score);
        let alignment = aligner.global(x, y);
        assert_eq!(alignment.ystart, 0);
        assert_eq!(alignment.xstart, 0);
        assert_eq!(alignment.operations, [Ins, Ins, Ins, Subst, Match, Match, Match]);
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::collections::HashSet;
use std::vec::Vec;

use analysis::needed_upcast::needed_upcast;
use analysis::out_parameters;
use analysis::return_value;
use analysis::rust_type::*;
use analysis::upcasts::Upcasts;
use env::Env;
use library::{self, Nullable};
use nameutil;
use traits::*;
use version::Version;

//TODO: change use Parameter to reference?
pub struct Info {
    pub name: String,
    pub glib_name: String,
    pub kind: library::FunctionKind,
    pub comented: bool,
    pub class_name: Result,
    pub parameters: Vec<library::Parameter>,
    pub ret: return_value::Info,
    pub upcasts: Upcasts,
    pub outs: out_parameters::Info,
    pub version: Option<Version>,
}

pub fn analyze(env: &Env, klass: &library::Class, class_tid: library::TypeId,
    non_nullable_overrides: &[String], used_types: &mut HashSet<String>) -> Vec<Info> {
    let mut funcs = Vec::new();

    for func in &klass.functions {
        let info = analyze_function(env, func, class_tid, non_nullable_overrides, used_types);
        funcs.push(info);
    }

    funcs
}

fn analyze_function(env: &Env, func: &library::Function, class_tid: library::TypeId,
    non_nullable_overrides: &[String], all_used_types: &mut HashSet<String>) -> Info {
    let mut commented = false;
    let mut upcasts: Upcasts = Default::default();
    let mut used_types: Vec<String> = Vec::with_capacity(4);

    let ret = return_value::analyze(env, func, class_tid, non_nullable_overrides, &mut used_types);
    commented |= ret.commented;

    let mut parameters = func.parameters.clone();

    for (pos, par) in parameters.iter_mut().enumerate() {
        assert!(!par.instance_parameter || pos == 0,
            "Wrong instance parameter in {}", func.c_identifier.as_ref().unwrap());
        if let Ok(s) = used_rust_type(env, par.typ) {
            used_types.push(s);
        }
        if !par.instance_parameter {
            if let Cow::Owned(mangled) = nameutil::mangle_keywords(&*par.name) {
                par.name = mangled;
            }
        }
        if !par.instance_parameter && needed_upcast(&env.library, par.typ) {
            let type_name = rust_type(env, par.typ);
            if !upcasts.add_parameter(&par.name, type_name.as_str()) {
                panic!("Too many parameters upcasts for {}", func.c_identifier.as_ref().unwrap())
            }
        }
        if parameter_rust_type(env, par.typ, par.direction, Nullable(false)).is_err() {
            commented = true;
        }
    }

    if !commented {
        for s in used_types {
            all_used_types.insert(s);
        }
    }

    let outs = out_parameters::analyze(func);

    Info {
        name: nameutil::mangle_keywords(&*func.name).into_owned(),
        glib_name: func.c_identifier.as_ref().unwrap().clone(),
        kind: func.kind,
        comented: commented,
        class_name: rust_type(env, class_tid),
        parameters: parameters,
        ret: ret,
        upcasts: upcasts,
        outs: outs,
        version: func.version,
    }
}
<commit_msg>When collecting used_types only take the crate name for external ones<commit_after>use std::borrow::Cow;
use std::collections::HashSet;
use std::vec::Vec;

use analysis::needed_upcast::needed_upcast;
use analysis::out_parameters;
use analysis::return_value;
use analysis::rust_type::*;
use analysis::upcasts::Upcasts;
use env::Env;
use library::{self, Nullable};
use nameutil;
use traits::*;
use version::Version;

//TODO: change use Parameter to reference?
pub struct Info {
    pub name: String,
    pub glib_name: String,
    pub kind: library::FunctionKind,
    pub comented: bool,
    pub class_name: Result,
    pub parameters: Vec<library::Parameter>,
    pub ret: return_value::Info,
    pub upcasts: Upcasts,
    pub outs: out_parameters::Info,
    pub version: Option<Version>,
}

pub fn analyze(env: &Env, klass: &library::Class, class_tid: library::TypeId,
    non_nullable_overrides: &[String], used_types: &mut HashSet<String>) -> Vec<Info> {
    let mut funcs = Vec::new();

    for func in &klass.functions {
        let info = analyze_function(env, func, class_tid, non_nullable_overrides, used_types);
        funcs.push(info);
    }

    funcs
}

fn analyze_function(env: &Env, func: &library::Function, class_tid: library::TypeId,
    non_nullable_overrides: &[String], all_used_types: &mut HashSet<String>) -> Info {
    let mut commented = false;
    let mut upcasts: Upcasts = Default::default();
    let mut used_types: Vec<String> = Vec::with_capacity(4);

    let ret = return_value::analyze(env, func, class_tid, non_nullable_overrides, &mut used_types);
    commented |= ret.commented;

    let mut parameters = func.parameters.clone();

    for (pos, par) in parameters.iter_mut().enumerate() {
        assert!(!par.instance_parameter || pos == 0,
            "Wrong instance parameter in {}", func.c_identifier.as_ref().unwrap());
        if let Ok(s) = used_rust_type(env, par.typ) {
            used_types.push(s);
        }
        if !par.instance_parameter {
            if let Cow::Owned(mangled) = nameutil::mangle_keywords(&*par.name) {
                par.name = mangled;
            }
        }
        if !par.instance_parameter && needed_upcast(&env.library, par.typ) {
            let type_name = rust_type(env, par.typ);
            if !upcasts.add_parameter(&par.name, type_name.as_str()) {
                panic!("Too many parameters upcasts for {}", func.c_identifier.as_ref().unwrap())
            }
        }
        if parameter_rust_type(env, par.typ, par.direction, Nullable(false)).is_err() {
            commented = true;
        }
    }

    if !commented {
        for s in used_types {
            if let Some(i) = s.find("::") {
                all_used_types.insert(s[..i].into());
            }
            else {
                all_used_types.insert(s);
            }
        }
    }

    let outs = out_parameters::analyze(func);

    Info {
        name: nameutil::mangle_keywords(&*func.name).into_owned(),
        glib_name: func.c_identifier.as_ref().unwrap().clone(),
        kind: func.kind,
        comented: commented,
        class_name: rust_type(env, class_tid),
        parameters: parameters,
        ret: ret,
        upcasts: upcasts,
        outs: outs,
        version: func.version,
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    fs::File,
    io::{self, Error, ErrorKind},
};

#[path = "win/punch_hole.rs"]
mod punch_hole;

#[path = "win/write_zeros.rs"]
mod write_zeros;

use super::write_zeroes::punch_hole::execute_punch_hole;

/// A trait for deallocating space in a file.
pub trait PunchHole {
    /// Replace a range of bytes with a hole.
    fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>;
}

impl PunchHole for File {
    fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()> {
        execute_punch_hole(self, offset, length)
    }
}

/// A trait for writing zeroes to an arbitrary position in a file.
pub trait WriteZeroesAt {
    /// Write up to `length` bytes of zeroes starting at `offset`, returning how many bytes were
    /// written.
    fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>;

    /// Write zeroes starting at `offset` until `length` bytes have been written.
    ///
    /// This method will continuously call `write_zeroes_at` until the requested
    /// `length` is satisfied or an error is encountered.
    fn write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()> {
        while length > 0 {
            match self.write_zeroes_at(offset, length) {
                Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
                Ok(bytes_written) => {
                    length = length
                        .checked_sub(bytes_written)
                        .ok_or_else(|| Error::from(ErrorKind::Other))?;
                    offset = offset
                        .checked_add(bytes_written as u64)
                        .ok_or_else(|| Error::from(ErrorKind::Other))?;
                }
                Err(e) => {
                    if e.kind() != ErrorKind::Interrupted {
                        return Err(e);
                    }
                }
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{
        fs::OpenOptions,
        io::{Read, Seek, SeekFrom, Write},
    };
    use tempfile::TempDir;

    #[cfg_attr(all(target_os = "windows", target_env = "gnu"), ignore)]
    #[test]
    fn simple_test() {
        let tempdir = TempDir::new().unwrap();
        let mut path = tempdir.path().to_owned();
        path.push("file");
        let mut f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(&path)
            .unwrap();
        f.set_len(16384).unwrap();

        // Write buffer of non-zero bytes to offset 1234
        let orig_data = [0x55u8; 5678];
        f.seek(SeekFrom::Start(1234)).unwrap();
        f.write_all(&orig_data).unwrap();

        // Read back the data plus some overlap on each side
        let mut readback = [0u8; 16384];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // Bytes before the write should still be 0
        for read in &readback[0..1234] {
            assert_eq!(*read, 0);
        }
        // Bytes that were just written should be 0x55
        for read in &readback[1234..(1234 + 5678)] {
            assert_eq!(*read, 0x55);
        }
        // Bytes after the written area should still be 0
        for read in &readback[(1234 + 5678)..] {
            assert_eq!(*read, 0);
        }

        // Overwrite some of the data with zeroes
        f.write_zeroes_all_at(2345, 4321)
            .expect("write_zeroes failed");

        // Read back the data and verify that it is now zero
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // Bytes before the write should still be 0
        for read in &readback[0..1234] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist before the write_zeroes region
        for read in &readback[1234..2345] {
            assert_eq!(*read, 0x55);
        }
        // The write_zeroes region should now be zero
        for read in &readback[2345..(2345 + 4321)] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist after the write_zeroes region
        for read in &readback[(2345 + 4321)..(1234 + 5678)] {
            assert_eq!(*read, 0x55);
        }
        // The rest of the file should still be 0
        for read in &readback[(1234 + 5678)..] {
            assert_eq!(*read, 0);
        }
    }

    #[test]
    #[cfg_attr(all(target_os = "windows", target_env = "gnu"), ignore)]
    fn large_write_zeroes() {
        let tempdir = TempDir::new().unwrap();
        let mut path = tempdir.path().to_owned();
        path.push("file");
        let mut f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(&path)
            .unwrap();
        f.set_len(16384).unwrap();

        // Write buffer of non-zero bytes
        let orig_data = [0x55u8; 0x20000];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.write_all(&orig_data).unwrap();

        // Overwrite some of the data with zeroes
        f.write_zeroes_all_at(0, 0x10001)
            .expect("write_zeroes failed");

        // Read back the data and verify that it is now zero
        let mut readback = [0u8; 0x20000];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // The write_zeroes region should now be zero
        for read in &readback[0..0x10001] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist after the write_zeroes region
        for read in &readback[0x10001..0x20000] {
            assert_eq!(*read, 0x55);
        }
    }
}
<commit_msg>base: enable WriteZeroes tests on wine<commit_after>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    fs::File,
    io::{self, Error, ErrorKind},
};

#[path = "win/punch_hole.rs"]
mod punch_hole;

#[path = "win/write_zeros.rs"]
mod write_zeros;

use super::write_zeroes::punch_hole::execute_punch_hole;

/// A trait for deallocating space in a file.
pub trait PunchHole {
    /// Replace a range of bytes with a hole.
    fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()>;
}

impl PunchHole for File {
    fn punch_hole(&mut self, offset: u64, length: u64) -> io::Result<()> {
        execute_punch_hole(self, offset, length)
    }
}

/// A trait for writing zeroes to an arbitrary position in a file.
pub trait WriteZeroesAt {
    /// Write up to `length` bytes of zeroes starting at `offset`, returning how many bytes were
    /// written.
    fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>;

    /// Write zeroes starting at `offset` until `length` bytes have been written.
    ///
    /// This method will continuously call `write_zeroes_at` until the requested
    /// `length` is satisfied or an error is encountered.
    fn write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()> {
        while length > 0 {
            match self.write_zeroes_at(offset, length) {
                Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
                Ok(bytes_written) => {
                    length = length
                        .checked_sub(bytes_written)
                        .ok_or_else(|| Error::from(ErrorKind::Other))?;
                    offset = offset
                        .checked_add(bytes_written as u64)
                        .ok_or_else(|| Error::from(ErrorKind::Other))?;
                }
                Err(e) => {
                    if e.kind() != ErrorKind::Interrupted {
                        return Err(e);
                    }
                }
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{
        fs::OpenOptions,
        io::{Read, Seek, SeekFrom, Write},
    };
    use tempfile::TempDir;

    #[test]
    fn simple_test() {
        let tempdir = TempDir::new().unwrap();
        let mut path = tempdir.path().to_owned();
        path.push("file");
        let mut f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(&path)
            .unwrap();

        // Extend and fill the file with zero bytes.
        // This is not using set_len() because that fails on wine.
        let init_data = [0x00u8; 16384];
        f.write_all(&init_data).unwrap();

        // Write buffer of non-zero bytes to offset 1234
        let orig_data = [0x55u8; 5678];
        f.seek(SeekFrom::Start(1234)).unwrap();
        f.write_all(&orig_data).unwrap();

        // Read back the data plus some overlap on each side
        let mut readback = [0u8; 16384];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // Bytes before the write should still be 0
        for read in &readback[0..1234] {
            assert_eq!(*read, 0);
        }
        // Bytes that were just written should be 0x55
        for read in &readback[1234..(1234 + 5678)] {
            assert_eq!(*read, 0x55);
        }
        // Bytes after the written area should still be 0
        for read in &readback[(1234 + 5678)..] {
            assert_eq!(*read, 0);
        }

        // Overwrite some of the data with zeroes
        f.write_zeroes_all_at(2345, 4321)
            .expect("write_zeroes failed");

        // Read back the data and verify that it is now zero
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // Bytes before the write should still be 0
        for read in &readback[0..1234] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist before the write_zeroes region
        for read in &readback[1234..2345] {
            assert_eq!(*read, 0x55);
        }
        // The write_zeroes region should now be zero
        for read in &readback[2345..(2345 + 4321)] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist after the write_zeroes region
        for read in &readback[(2345 + 4321)..(1234 + 5678)] {
            assert_eq!(*read, 0x55);
        }
        // The rest of the file should still be 0
        for read in &readback[(1234 + 5678)..] {
            assert_eq!(*read, 0);
        }
    }

    #[test]
    fn large_write_zeroes() {
        let tempdir = TempDir::new().unwrap();
        let mut path = tempdir.path().to_owned();
        path.push("file");
        let mut f = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .open(&path)
            .unwrap();

        // Extend and fill the file with zero bytes.
        // This is not using set_len() because that fails on wine.
        let init_data = [0x00u8; 16384];
        f.write_all(&init_data).unwrap();

        // Write buffer of non-zero bytes
        let orig_data = [0x55u8; 0x20000];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.write_all(&orig_data).unwrap();

        // Overwrite some of the data with zeroes
        f.write_zeroes_all_at(0, 0x10001)
            .expect("write_zeroes failed");

        // Read back the data and verify that it is now zero
        let mut readback = [0u8; 0x20000];
        f.seek(SeekFrom::Start(0)).unwrap();
        f.read_exact(&mut readback).unwrap();
        // The write_zeroes region should now be zero
        for read in &readback[0..0x10001] {
            assert_eq!(*read, 0);
        }
        // Original data should still exist after the write_zeroes region
        for read in &readback[0x10001..0x20000] {
            assert_eq!(*read, 0x55);
        }
    }
}
<|endoftext|>