<commit_before>// Feature: Format String Completion
//
// `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//
// The following postfix snippets are available:
//
// - `format` -> `format!(...)`
// - `panic` -> `panic!(...)`
// - `println` -> `println!(...)`
// - `log`:
//   + `logd` -> `log::debug!(...)`
//   + `logt` -> `log::trace!(...)`
//   + `logi` -> `log::info!(...)`
//   + `logw` -> `log::warn!(...)`
//   + `loge` -> `log::error!(...)`
//
// image::https://user-images.githubusercontent.com/48062697/113020656-b560f500-917a-11eb-87de-02991f61beb8.gif[]

use ide_db::helpers::SnippetCap;
use syntax::ast::{self, AstToken};

use crate::{completions::postfix::postfix_snippet, context::CompletionContext, Completions};

/// Mapping ("postfix completion item" => "macro to use")
static KINDS: &[(&str, &str)] = &[
    ("format", "format!"),
    ("panic", "panic!"),
    ("println", "println!"),
    ("eprintln", "eprintln!"),
    ("logd", "log::debug!"),
    ("logt", "log::trace!"),
    ("logi", "log::info!"),
    ("logw", "log::warn!"),
    ("loge", "log::error!"),
];

pub(crate) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &ast::String,
) {
    let input = match string_literal_contents(receiver_text) {
        // It's not a string literal, do not parse input.
        Some(input) => input,
        None => return,
    };

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for (label, macro_name) in KINDS {
            let snippet = parser.into_suggestion(macro_name);

            postfix_snippet(ctx, cap, &dot_receiver, label, macro_name, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn string_literal_contents(item: &ast::String) -> Option<String> {
    let item = item.text();
    if item.len() >= 2 && item.starts_with('\"') && item.ends_with('\"') {
        return Some(item[1..item.len() - 1].to_owned());
    }

    None
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub(crate) struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub(crate) fn new(input: String) -> Self {
        Self {
            input: input,
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub(crate) fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        let mut chars = self.input.chars().peekable();
        while let Some(chr) = chars.next() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') if chars.peek().copied() == Some(':') => {
                    // path seperator
                    current_expr.push_str("::");
                    chars.next();
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub(crate) fn into_suggestion(&self, macro_name: &str) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let expressions_as_string = self.extracted_expressions.join(", ");
        format!(r#"{}("{}", {})"#, macro_name, self.output, expressions_as_string)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use expect_test::{expect, Expect};

    fn check(input: &str, expect: &Expect) {
        let mut parser = FormatStrParser::new((*input).to_owned());
        let outcome_repr = if parser.parse().is_ok() {
            // Parsing should be OK, expected repr is "string; expr_1, expr_2".
            if parser.extracted_expressions.is_empty() {
                parser.output
            } else {
                format!("{}; {}", parser.output, parser.extracted_expressions.join(", "))
            }
        } else {
            // Parsing should fail, expected repr is "-".
            "-".to_owned()
        };

        expect.assert_eq(&outcome_repr);
    }

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", expect![["no expressions"]]),
            ("{expr} is {2 + 2}", expect![["{} is {}; expr, 2 + 2"]]),
            ("{expr:?}", expect![["{:?}; expr"]]),
            ("{malformed", expect![["-"]]),
            ("malformed}", expect![["-"]]),
            ("{{correct", expect![["{{correct"]]),
            ("correct}}", expect![["correct}}"]]),
            ("{correct}}}", expect![["{}}}; correct"]]),
            ("{correct}}}}}", expect![["{}}}}}; correct"]]),
            ("{incorrect}}", expect![["-"]]),
            ("placeholders {} {}", expect![["placeholders {} {}; $1, $2"]]),
            ("mixed {} {2 + 2} {}", expect![["mixed {} {} {}; $1, 2 + 2, $2"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                expect![["{}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{expr:?} is {2.32f64:.5}", expect![["{:?} is {:.5}; expr, 2.32f64"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                expect![["{:?}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{     2 + 2        }", expect![["{}; 2 + 2"]]),
            ("{strsim::jaro_winkle(a)}", expect![["{}; strsim::jaro_winkle(a)"]]),
            ("{foo::bar::baz()}", expect![["{}; foo::bar::baz()"]]),
            ("{foo::bar():?}", expect![["{:?}; foo::bar()"]]),
        ];

        for (input, output) in test_vector {
            check(input, output)
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            ("println!", "{}", r#"println!("{}", $1)"#),
            ("eprintln!", "{}", r#"eprintln!("{}", $1)"#),
            (
                "log::info!",
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $1, expr, $2, 2 + 2)"#,
            ),
            ("format!", "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new((*input).to_owned());
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<commit_msg>Fix incorrect list nesting in format_like feature description<commit_after>// Feature: Format String Completion
//
// `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//
// The following postfix snippets are available:
//
// * `format` -> `format!(...)`
// * `panic` -> `panic!(...)`
// * `println` -> `println!(...)`
// * `log`:
// ** `logd` -> `log::debug!(...)`
// ** `logt` -> `log::trace!(...)`
// ** `logi` -> `log::info!(...)`
// ** `logw` -> `log::warn!(...)`
// ** `loge` -> `log::error!(...)`
//
// image::https://user-images.githubusercontent.com/48062697/113020656-b560f500-917a-11eb-87de-02991f61beb8.gif[]

use ide_db::helpers::SnippetCap;
use syntax::ast::{self, AstToken};

use crate::{completions::postfix::postfix_snippet, context::CompletionContext, Completions};

/// Mapping ("postfix completion item" => "macro to use")
static KINDS: &[(&str, &str)] = &[
    ("format", "format!"),
    ("panic", "panic!"),
    ("println", "println!"),
    ("eprintln", "eprintln!"),
    ("logd", "log::debug!"),
    ("logt", "log::trace!"),
    ("logi", "log::info!"),
    ("logw", "log::warn!"),
    ("loge", "log::error!"),
];

pub(crate) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &ast::String,
) {
    let input = match string_literal_contents(receiver_text) {
        // It's not a string literal, do not parse input.
        Some(input) => input,
        None => return,
    };

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for (label, macro_name) in KINDS {
            let snippet = parser.into_suggestion(macro_name);

            postfix_snippet(ctx, cap, &dot_receiver, label, macro_name, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn string_literal_contents(item: &ast::String) -> Option<String> {
    let item = item.text();
    if item.len() >= 2 && item.starts_with('\"') && item.ends_with('\"') {
        return Some(item[1..item.len() - 1].to_owned());
    }

    None
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub(crate) struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub(crate) fn new(input: String) -> Self {
        Self {
            input: input,
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub(crate) fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        let mut chars = self.input.chars().peekable();
        while let Some(chr) = chars.next() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') if chars.peek().copied() == Some(':') => {
                    // path seperator
                    current_expr.push_str("::");
                    chars.next();
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub(crate) fn into_suggestion(&self, macro_name: &str) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let expressions_as_string = self.extracted_expressions.join(", ");
        format!(r#"{}("{}", {})"#, macro_name, self.output, expressions_as_string)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use expect_test::{expect, Expect};

    fn check(input: &str, expect: &Expect) {
        let mut parser = FormatStrParser::new((*input).to_owned());
        let outcome_repr = if parser.parse().is_ok() {
            // Parsing should be OK, expected repr is "string; expr_1, expr_2".
            if parser.extracted_expressions.is_empty() {
                parser.output
            } else {
                format!("{}; {}", parser.output, parser.extracted_expressions.join(", "))
            }
        } else {
            // Parsing should fail, expected repr is "-".
            "-".to_owned()
        };

        expect.assert_eq(&outcome_repr);
    }

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", expect![["no expressions"]]),
            ("{expr} is {2 + 2}", expect![["{} is {}; expr, 2 + 2"]]),
            ("{expr:?}", expect![["{:?}; expr"]]),
            ("{malformed", expect![["-"]]),
            ("malformed}", expect![["-"]]),
            ("{{correct", expect![["{{correct"]]),
            ("correct}}", expect![["correct}}"]]),
            ("{correct}}}", expect![["{}}}; correct"]]),
            ("{correct}}}}}", expect![["{}}}}}; correct"]]),
            ("{incorrect}}", expect![["-"]]),
            ("placeholders {} {}", expect![["placeholders {} {}; $1, $2"]]),
            ("mixed {} {2 + 2} {}", expect![["mixed {} {} {}; $1, 2 + 2, $2"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                expect![["{}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{expr:?} is {2.32f64:.5}", expect![["{:?} is {:.5}; expr, 2.32f64"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                expect![["{:?}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{     2 + 2        }", expect![["{}; 2 + 2"]]),
            ("{strsim::jaro_winkle(a)}", expect![["{}; strsim::jaro_winkle(a)"]]),
            ("{foo::bar::baz()}", expect![["{}; foo::bar::baz()"]]),
            ("{foo::bar():?}", expect![["{:?}; foo::bar()"]]),
        ];

        for (input, output) in test_vector {
            check(input, output)
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            ("println!", "{}", r#"println!("{}", $1)"#),
            ("eprintln!", "{}", r#"eprintln!("{}", $1)"#),
            (
                "log::info!",
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $1, expr, $2, 2 + 2)"#,
            ),
            ("format!", "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new((*input).to_owned());
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate regex;

#[macro_use]
extern crate lazy_static;

use std::io::prelude::*;
use std::fs::File;
use regex::Regex;
use std::collections::HashMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Position {
    x: usize,
    y: usize,
}

#[derive(Debug)]
struct NodeInfo {
    size: usize,
    used: usize,
    avail: usize,
    use_percent: usize,
}

#[derive(Debug)]
struct NodeListing {
    pos: Position,
    info: NodeInfo,
}

impl std::str::FromStr for NodeListing {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        lazy_static! {
            static ref NODE_RE: Regex =
                Regex::new(
                    r"/dev/grid/node-x(\d+)-y(\d+)\s+(\d+)T\s+(\d+)T\s+(\d+)T\s+(\d+)%")
                .expect("bad regex NODE_RE");
        }
        if let Some(cap) = NODE_RE.captures(s) {
            Ok(NodeListing {
                pos: Position {
                    x: cap[1].parse().unwrap(),
                    y: cap[2].parse().unwrap(),
                },
                info: NodeInfo {
                    size: cap[3].parse().unwrap(),
                    used: cap[4].parse().unwrap(),
                    avail: cap[5].parse().unwrap(),
                    use_percent: cap[6].parse().unwrap(),
                }
            })
        } else {
            Err(format!("Could not parse NodeListing from {}", s))
        }
    }
}

fn valid_pairs(nodes: &HashMap<Position, NodeInfo>) -> Vec<(Position, Position)> {
    let mut pairs = Vec::new();
    for (&pos_a, ref info_a) in nodes.iter() {
        for (&pos_b, ref info_b) in nodes.iter() {
            if info_a.used != 0 && pos_a != pos_b && info_a.used <= info_b.avail {
                pairs.push((pos_a, pos_b));
            }
        }
    }
    pairs
}

fn parse_args() -> std::io::Result<String> {
    let matches = clap::App::new("Day 22")
        .author("Devon Hollowood")
        .arg(clap::Arg::with_name("nodes")
            .index(1)
            .help("file containing output of nodes `df`. Reads from stdin otherwise"))
        .get_matches();
    let source = matches.value_of_os("nodes");
    let mut contents = String::new();
    match source {
        Some(filename) => {
            let mut input = File::open(filename)?;
            input.read_to_string(&mut contents)?;
        }
        None => {
            let mut input = std::io::stdin();
            input.read_to_string(&mut contents)?;
        }
    }
    Ok(contents)
}

fn main() {
    let input = parse_args().unwrap_or_else(|err| panic!("Error reading args: {}", err));
    let nodes: HashMap<Position, NodeInfo> = input.lines()
        .map(|line| line.trim().parse::<NodeListing>())
        .filter_map(|r| r.ok())
        .map(|listing| (listing.pos, listing.info))
        .collect();
    println!("number of valid node pairs: {}", valid_pairs(&nodes).len());
}
<commit_msg>Day 22 part 2 attempt 1<commit_after>extern crate clap;
extern crate regex;

#[macro_use]
extern crate lazy_static;

use std::io::prelude::*;
use std::fs::File;
use regex::Regex;
use std::collections::{HashMap, VecDeque};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Position {
    x: usize,
    y: usize,
}

impl std::fmt::Display for Position {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(f, "node-y{}-x{}", self.y, self.x)
    }
}

impl Position {
    fn neighbors(&self) -> Vec<Position> {
        let &Position { x, y } = self;
        let mut neighbors = vec![Position { x: x + 1, y: y }, Position { x: x, y: y + 1 }];
        if x != 0 {
            neighbors.push(Position { x: x - 1, y: y });
        }
        if y != 0 {
            neighbors.push(Position { x: x, y: y - 1 });
        }
        neighbors
    }
}

#[derive(Debug, Clone)]
struct NodeInfo {
    size: usize,
    used: usize,
    avail: usize,
}

#[derive(Debug)]
struct NodeListing {
    pos: Position,
    info: NodeInfo,
}

impl std::str::FromStr for NodeListing {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        lazy_static! {
            static ref NODE_RE: Regex =
                Regex::new(
                    r"/dev/grid/node-x(\d+)-y(\d+)\s+(\d+)T\s+(\d+)T\s+(\d+)T\s+\d+%")
                .expect("bad regex NODE_RE");
        }
        if let Some(cap) = NODE_RE.captures(s) {
            Ok(NodeListing {
                pos: Position {
                    x: cap[1].parse().unwrap(),
                    y: cap[2].parse().unwrap(),
                },
                info: NodeInfo {
                    size: cap[3].parse().unwrap(),
                    used: cap[4].parse().unwrap(),
                    avail: cap[5].parse().unwrap(),
                },
            })
        } else {
            Err(format!("Could not parse NodeListing from {}", s))
        }
    }
}

fn valid_pairs(nodes: &HashMap<Position, NodeInfo>) -> Vec<(Position, Position)> {
    let mut pairs = Vec::new();
    for (&pos_a, ref info_a) in nodes.iter() {
        for (&pos_b, ref info_b) in nodes.iter() {
            if info_a.used != 0 && pos_a != pos_b && info_a.used <= info_b.avail {
                pairs.push((pos_a, pos_b));
            }
        }
    }
    pairs
}

fn possible_moves(nodes: &HashMap<Position, NodeInfo>) -> Vec<(Position, Position)> {
    let mut possibilities = Vec::new();
    for (&pos_a, ref info_a) in nodes.iter() {
        if info_a.used == 0 {
            // no point in moving no data
            continue;
        }
        for pos_b in pos_a.neighbors() {
            if let Some(info_b) = nodes.get(&pos_b) {
                if info_a.used <= info_b.avail {
                    possibilities.push((pos_a, pos_b));
                }
            }
        }
    }
    possibilities
}

fn move_data(nodes: &mut HashMap<Position, NodeInfo>, from: Position, to: Position) {
    let from_info = nodes.get(&from)
        .unwrap_or_else(|| panic!("Invalid move: {} does not exist", from))
        .clone();
    if let Some(ref mut to_info) = nodes.get_mut(&to) {
        if from_info.used > to_info.avail {
            panic!("Invalid move: {} (avail: {}) cannot hold {} (used: {})",
                   to,
                   to_info.avail,
                   from,
                   from_info.avail);
        }
        to_info.used += from_info.used;
        to_info.avail -= from_info.used;
    } else {
        panic!("Invalid move: {} does not exist", to);
    }
    let mut from_info_mut = nodes.get_mut(&from).unwrap();
    from_info_mut.avail = from_info.size;
    from_info_mut.used = 0;
}

fn solve(nodes: &HashMap<Position, NodeInfo>) -> Option<Vec<(Position, Position)>> {
    let mut state_queue = VecDeque::new();
    let initial_target_loc = nodes.keys()
        .filter(|pos| pos.y == 0)
        .max_by_key(|pos| pos.x)
        .expect("Invalid target state")
        .clone();
    state_queue.push_back((nodes.clone(), initial_target_loc, Vec::new()));
    while let Some((state, target_loc, moves)) = state_queue.pop_front() {
        // print!("considering target = {} moves = [", target_loc);
        // for &(from, to) in &moves {
        //    print!("{} -> {}, ", from, to);
        // }
        // println!("] ({} moves from here.)", possible_moves(&state).len());
        for (from, to) in possible_moves(&state) {
            let mut new_state = state.clone();
            let mut new_moves = moves.clone();
            new_moves.push((from, to));
            move_data(&mut new_state, from, to);
            let new_target_loc = if from == target_loc { to } else { target_loc };
            if new_target_loc.x == 0 && new_target_loc.y == 0 {
                return Some(new_moves);
            }
            state_queue.push_back((new_state, new_target_loc, new_moves));
        }
    }
    None
}

fn parse_args() -> std::io::Result<String> {
    let matches = clap::App::new("Day 22")
        .author("Devon Hollowood")
        .arg(clap::Arg::with_name("nodes")
            .index(1)
            .help("file containing output of nodes `df`. Reads from stdin otherwise"))
        .get_matches();
    let source = matches.value_of_os("nodes");
    let mut contents = String::new();
    match source {
        Some(filename) => {
            let mut input = File::open(filename)?;
            input.read_to_string(&mut contents)?;
        }
        None => {
            let mut input = std::io::stdin();
            input.read_to_string(&mut contents)?;
        }
    }
    Ok(contents)
}

fn main() {
    let input = parse_args().unwrap_or_else(|err| panic!("Error reading args: {}", err));
    let nodes: HashMap<Position, NodeInfo> = input.lines()
        .map(|line| line.trim().parse::<NodeListing>())
        .filter_map(|r| r.ok())
        .map(|listing| (listing.pos, listing.info))
        .collect();
    println!("number of valid node pairs: {}", valid_pairs(&nodes).len());
    if let Some(solution) = solve(&nodes) {
        println!("Solution:");
        for mov in &solution {
            println!("{} -> {}", mov.0, mov.1);
        }
        println!("Total moves taken: {}", solution.len())
    }
}

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

    #[test]
    fn example() {
        let nodes = ["/dev/grid/node-x0-y0   10T    8T     2T   80%",
                     "/dev/grid/node-x0-y1   11T    6T     5T   54%",
                     "/dev/grid/node-x0-y2   32T   28T     4T   87%",
                     "/dev/grid/node-x1-y0    9T    7T     2T   77%",
                     "/dev/grid/node-x1-y1    8T    0T     8T    0%",
                     "/dev/grid/node-x1-y2   11T    7T     4T   63%",
                     "/dev/grid/node-x2-y0   10T    6T     4T   60%",
                     "/dev/grid/node-x2-y1    9T    8T     1T   88%",
                     "/dev/grid/node-x2-y2    9T    6T     3T   66%"]
            .into_iter()
            .map(|s| s.parse::<NodeListing>().unwrap())
            .map(|listing| (listing.pos, listing.info))
            .collect();
        assert_eq!(solve(&nodes).unwrap().len(), 7);
    }
}
<|endoftext|>


<commit_before>extern mod std;
use str;

// TODO http://dev.w3.org/csswg/css3-syntax/#the-input-byte-stream
// fn decode(input: &[u8], protocol_encoding: &str, link_encoding: &str,
//           document_encoding: &str) -> ~str
// Use empty strings for “no such encoding information”?

// http://dev.w3.org/csswg/css3-syntax/#preprocessing-the-input-stream
fn preprocess(input: &str) -> ~str {
    // TODO: Is this faster if done in one pass?
    str::replace(str::replace(str::replace(input,
    "\r\n", "\n"),
    "\r", "\n"),
    "\x00", "\uFFFD")
}

#[test]
fn test_preprocess() {
    assert preprocess("") == ~"";
    assert preprocess("Lorem\r\n\t\x00ipusm\ndoror\uFFFD\r")
        == ~"Lorem\n\t\uFFFDipusm\ndoror\uFFFD\n";
}
<commit_msg>Add token types.<commit_after>extern mod std;
use str;

// TODO http://dev.w3.org/csswg/css3-syntax/#the-input-byte-stream
// fn decode(input: &[u8], protocol_encoding: &str, link_encoding: &str,
//           document_encoding: &str) -> ~str
// Use empty strings for “no such encoding information”?

// http://dev.w3.org/csswg/css3-syntax/#preprocessing-the-input-stream
fn preprocess(input: &str) -> ~str {
    // TODO: Is this faster if done in one pass?
    str::replace(str::replace(str::replace(input,
    "\r\n", "\n"),
    "\r", "\n"),
    "\x00", "\uFFFD")
}

#[test]
fn test_preprocess() {
    assert preprocess("") == ~"";
    assert preprocess("Lorem\r\n\t\x00ipusm\ndoror\uFFFD\r")
        == ~"Lorem\n\t\uFFFDipusm\ndoror\uFFFD\n";
}

// http://dev.w3.org/csswg/css3-syntax/#tokenization
//
// The output of the tokenization step is a series of zero or more
// of the following tokens:
// ident, function, at-keyword, hash, string, bad-string, url, bad-url,
// delim, number, percentage, dimension, unicode-range, whitespace, comment,
// cdo, cdc, colon, semicolon, [, ], (, ), {, }.
//
// ident, function, at-keyword, hash, string, and url tokens
// have a value composed of zero or more characters.
// Delim tokens have a value composed of a single character.
// Number, percentage, and dimension tokens have a representation
// composed of 1 or more character, a numeric value,
// and a type flag set to either "integer" or "number".
// The type flag defaults to "integer" if not otherwise set.
// Dimension tokens additionally have a unit
// composed of one or more characters.
// Unicode-range tokens have a range of characters.

enum NumericValue {
    Integer(int),
    // The spec calls this "number".
    // Use "float" instead to reduce term overloading with "number token".
    Float(float),
}
// TODO: add a NumberValue.as_float() method.

enum Token {
    Ident(~str),
    Function(~str),
    AtKeyword(~str),
    Hash(~str),
    String(~str),
    BadString,
    URL(~str),
    BadURL,
    Delim(char),
    Number(NumericValue, ~str),  // value, representation
    Percentage(NumericValue, ~str),  // value, representation
    Dimension(NumericValue, ~str, ~str),  // value, representation, unit
    UnicodeRange(char, char),  // start, end. XXX what about empty ranges?
    WhiteSpace,
    Comment,
    CDO,  // <!--
    CDC,  // -->
    Colon,  // ,
    Semicolon,  // ;
    OpenBraket, // [
    OpenParen, // (
    OpenBrace, // {
    CloseBraket, // ]
    CloseParen, // (
    CloseBrace, // }
}
<|endoftext|>


<commit_before>#![deny(rust_2018_idioms)]

use std::env;
use std::process::Command;
use std::path::{Path, PathBuf};
use std::fs;

struct Test {
    repo: &'static str,
    name: &'static str,
    sha: &'static str,
    lock: Option<&'static str>,
    packages: &'static [&'static str],
}

const TEST_REPOS: &'static [Test] = &[
    Test {
        name: "iron",
        repo: "https://github.com/iron/iron",
        sha: "21c7dae29c3c214c08533c2a55ac649b418f2fe3",
        lock: Some(include_str!("lockfiles/iron-Cargo.lock")),
        packages: &[],
    },
    Test {
        name: "ripgrep",
        repo: "https://github.com/BurntSushi/ripgrep",
        sha: "ad9befbc1d3b5c695e7f6b6734ee1b8e683edd41",
        lock: None,
        packages: &[],
    },
    Test {
        name: "tokei",
        repo: "https://github.com/XAMPPRocky/tokei",
        sha: "5e11c4852fe4aa086b0e4fe5885822fbe57ba928",
        lock: None,
        packages: &[],
    },
    Test {
        name: "treeify",
        repo: "https://github.com/dzamlo/treeify",
        sha: "999001b223152441198f117a68fb81f57bc086dd",
        lock: None,
        packages: &[],
    },
    Test {
        name: "xsv",
        repo: "https://github.com/BurntSushi/xsv",
        sha: "66956b6bfd62d6ac767a6b6499c982eae20a2c9f",
        lock: None,
        packages: &[],
    },
    Test {
        name: "servo",
        repo: "https://github.com/servo/servo",
        sha: "987e376ca7a4245dbc3e0c06e963278ee1ac92d1",
        lock: None,
        // Only test Stylo a.k.a. Quantum CSS, the parts of Servo going into Firefox.
        // This takes much less time to build than all of Servo and supports stable Rust.
        packages: &["selectors"],
    },
    Test {
        name: "webrender",
        repo: "https://github.com/servo/webrender",
        sha: "57250b2b8fa63934f80e5376a29f7dcb3f759ad6",
        lock: None,
        packages: &[],
    },
];

fn main() {
    let args = env::args().collect::<Vec<_>>();
    let ref cargo = args[1];
    let out_dir = Path::new(&args[2]);
    let ref cargo = Path::new(cargo);

    for test in TEST_REPOS.iter().rev() {
        test_repo(cargo, out_dir, test);
    }
}

fn test_repo(cargo: &Path, out_dir: &Path, test: &Test) {
    println!("testing {}", test.repo);
    let dir = clone_repo(test, out_dir);
    if let Some(lockfile) = test.lock {
        fs::write(&dir.join("Cargo.lock"), lockfile).unwrap();
    }
    if !run_cargo_test(cargo, &dir, test.packages) {
        panic!("tests failed for {}", test.repo);
    }
}

fn clone_repo(test: &Test, out_dir: &Path) -> PathBuf {
    let out_dir = out_dir.join(test.name);

    if !out_dir.join(".git").is_dir() {
        let status = Command::new("git")
                         .arg("init")
                         .arg(&out_dir)
                         .status()
                         .expect("");
        assert!(status.success());
    }

    // Try progressively deeper fetch depths to find the commit
    let mut found = false;
    for depth in &[0, 1, 10, 100, 1000, 100000] {
        if *depth > 0 {
            let status = Command::new("git")
                             .arg("fetch")
                             .arg(test.repo)
                             .arg("master")
                             .arg(&format!("--depth={}", depth))
                             .current_dir(&out_dir)
                             .status()
                             .expect("");
            assert!(status.success());
        }

        let status = Command::new("git")
                         .arg("reset")
                         .arg(test.sha)
                         .arg("--hard")
                         .current_dir(&out_dir)
                         .status()
                         .expect("");

        if status.success() {
            found = true;
            break;
        }
    }

    if !found {
        panic!("unable to find commit {}", test.sha)
    }
    let status = Command::new("git")
                     .arg("clean")
                     .arg("-fdx")
                     .current_dir(&out_dir)
                     .status()
                     .unwrap();
    assert!(status.success());

    out_dir
}

fn run_cargo_test(cargo_path: &Path, crate_path: &Path, packages: &[&str]) -> bool {
    let mut command = Command::new(cargo_path);
    command.arg("test");
    for name in packages {
        command.arg("-p").arg(name);
    }
    let status = command
        // Disable rust-lang/cargo's cross-compile tests
        .env("CFG_DISABLE_CROSS_TESTS", "1")
        // Relax #![deny(warnings)] in some crates
        .env("RUSTFLAGS", "--cap-lints warn")
        .current_dir(crate_path)
        .status()
        .expect("");

    status.success()
}
<commit_msg>Update webrender we're testing<commit_after>#![deny(rust_2018_idioms)]

use std::env;
use std::process::Command;
use std::path::{Path, PathBuf};
use std::fs;

struct Test {
    repo: &'static str,
    name: &'static str,
    sha: &'static str,
    lock: Option<&'static str>,
    packages: &'static [&'static str],
}

const TEST_REPOS: &'static [Test] = &[
    Test {
        name: "iron",
        repo: "https://github.com/iron/iron",
        sha: "21c7dae29c3c214c08533c2a55ac649b418f2fe3",
        lock: Some(include_str!("lockfiles/iron-Cargo.lock")),
        packages: &[],
    },
    Test {
        name: "ripgrep",
        repo: "https://github.com/BurntSushi/ripgrep",
        sha: "ad9befbc1d3b5c695e7f6b6734ee1b8e683edd41",
        lock: None,
        packages: &[],
    },
    Test {
        name: "tokei",
        repo: "https://github.com/XAMPPRocky/tokei",
        sha: "5e11c4852fe4aa086b0e4fe5885822fbe57ba928",
        lock: None,
        packages: &[],
    },
    Test {
        name: "treeify",
        repo: "https://github.com/dzamlo/treeify",
        sha: "999001b223152441198f117a68fb81f57bc086dd",
        lock: None,
        packages: &[],
    },
    Test {
        name: "xsv",
        repo: "https://github.com/BurntSushi/xsv",
        sha: "66956b6bfd62d6ac767a6b6499c982eae20a2c9f",
        lock: None,
        packages: &[],
    },
    Test {
        name: "servo",
        repo: "https://github.com/servo/servo",
        sha: "987e376ca7a4245dbc3e0c06e963278ee1ac92d1",
        lock: None,
        // Only test Stylo a.k.a. Quantum CSS, the parts of Servo going into Firefox.
        // This takes much less time to build than all of Servo and supports stable Rust.
        packages: &["selectors"],
    },
    Test {
        name: "webrender",
        repo: "https://github.com/servo/webrender",
        sha: "cdadd068f4c7218bd983d856981d561e605270ab",
        lock: None,
        packages: &[],
    },
];

fn main() {
    let args = env::args().collect::<Vec<_>>();
    let ref cargo = args[1];
    let out_dir = Path::new(&args[2]);
    let ref cargo = Path::new(cargo);

    for test in TEST_REPOS.iter().rev() {
        test_repo(cargo, out_dir, test);
    }
}

fn test_repo(cargo: &Path, out_dir: &Path, test: &Test) {
    println!("testing {}", test.repo);
    let dir = clone_repo(test, out_dir);
    if let Some(lockfile) = test.lock {
        fs::write(&dir.join("Cargo.lock"), lockfile).unwrap();
    }
    if !run_cargo_test(cargo, &dir, test.packages) {
        panic!("tests failed for {}", test.repo);
    }
}

fn clone_repo(test: &Test, out_dir: &Path) -> PathBuf {
    let out_dir = out_dir.join(test.name);

    if !out_dir.join(".git").is_dir() {
        let status = Command::new("git")
                         .arg("init")
                         .arg(&out_dir)
                         .status()
                         .expect("");
        assert!(status.success());
    }

    // Try progressively deeper fetch depths to find the commit
    let mut found = false;
    for depth in &[0, 1, 10, 100, 1000, 100000] {
        if *depth > 0 {
            let status = Command::new("git")
                             .arg("fetch")
                             .arg(test.repo)
                             .arg("master")
                             .arg(&format!("--depth={}", depth))
                             .current_dir(&out_dir)
                             .status()
                             .expect("");
            assert!(status.success());
        }

        let status = Command::new("git")
                         .arg("reset")
                         .arg(test.sha)
                         .arg("--hard")
                         .current_dir(&out_dir)
                         .status()
                         .expect("");

        if status.success() {
            found = true;
            break;
        }
    }

    if !found {
        panic!("unable to find commit {}", test.sha)
    }
    let status = Command::new("git")
                     .arg("clean")
                     .arg("-fdx")
                     .current_dir(&out_dir)
                     .status()
                     .unwrap();
    assert!(status.success());

    out_dir
}

fn run_cargo_test(cargo_path: &Path, crate_path: &Path, packages: &[&str]) -> bool {
    let mut command = Command::new(cargo_path);
    command.arg("test");
    for name in packages {
        command.arg("-p").arg(name);
    }
    let status = command
        // Disable rust-lang/cargo's cross-compile tests
        .env("CFG_DISABLE_CROSS_TESTS", "1")
        // Relax #![deny(warnings)] in some crates
        .env("RUSTFLAGS", "--cap-lints warn")
        .current_dir(crate_path)
        .status()
        .expect("");

    status.success()
}
<|endoftext|>


<commit_before>use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

struct Test {
    repo: &'static str,
    name: &'static str,
    sha: &'static str,
    lock: Option<&'static str>,
    packages: &'static [&'static str],
}

const TEST_REPOS: &[Test] = &[
    Test {
        name: "iron",
        repo: "https://github.com/iron/iron",
        sha: "cf056ea5e8052c1feea6141e40ab0306715a2c33",
        lock: None,
        packages: &[],
    },
    Test {
        name: "ripgrep",
        repo: "https://github.com/BurntSushi/ripgrep",
        sha: "ad9befbc1d3b5c695e7f6b6734ee1b8e683edd41",
        lock: None,
        packages: &[],
    },
    Test {
        name: "tokei",
        repo: "https://github.com/XAMPPRocky/tokei",
        sha: "a950ff128d5a435a8083b1c7577c0431f98360ca",
        lock: None,
        packages: &[],
    },
    Test {
        name: "xsv",
        repo: "https://github.com/BurntSushi/xsv",
        sha: "66956b6bfd62d6ac767a6b6499c982eae20a2c9f",
        lock: None,
        packages: &[],
    },
    Test {
        name: "servo",
        repo: "https://github.com/servo/servo",
        sha: "caac107ae8145ef2fd20365e2b8fadaf09c2eb3b",
        lock: None,
        // Only test Stylo a.k.a. Quantum CSS, the parts of Servo going into Firefox.
        // This takes much less time to build than all of Servo and supports stable Rust.
        packages: &["selectors"],
    },
];

fn main() {
    let args = env::args().collect::<Vec<_>>();
    let cargo = &args[1];
    let out_dir = Path::new(&args[2]);
    let cargo = &Path::new(cargo);

    for test in TEST_REPOS.iter().rev() {
        test_repo(cargo, out_dir, test);
    }
}

fn test_repo(cargo: &Path, out_dir: &Path, test: &Test) {
    println!("testing {}", test.repo);
    let dir = clone_repo(test, out_dir);
    if let Some(lockfile) = test.lock {
        fs::write(&dir.join("Cargo.lock"), lockfile).unwrap();
    }
    if !run_cargo_test(cargo, &dir, test.packages) {
        panic!("tests failed for {}", test.repo);
    }
}

fn clone_repo(test: &Test, out_dir: &Path) -> PathBuf {
    let out_dir = out_dir.join(test.name);

    if !out_dir.join(".git").is_dir() {
        let status = Command::new("git").arg("init").arg(&out_dir).status().unwrap();
        assert!(status.success());
    }

    // Try progressively deeper fetch depths to find the commit
    let mut found = false;
    for depth in &[0, 1, 10, 100, 1000, 100000] {
        if *depth > 0 {
            let status = Command::new("git")
                .arg("fetch")
                .arg(test.repo)
                .arg("master")
                .arg(&format!("--depth={}", depth))
                .current_dir(&out_dir)
                .status()
                .unwrap();
            assert!(status.success());
        }

        let status = Command::new("git")
            .arg("reset")
            .arg(test.sha)
            .arg("--hard")
            .current_dir(&out_dir)
            .status()
            .unwrap();

        if status.success() {
            found = true;
            break;
        }
    }

    if !found {
        panic!("unable to find commit {}", test.sha)
    }
    let status =
        Command::new("git").arg("clean").arg("-fdx").current_dir(&out_dir).status().unwrap();
    assert!(status.success());

    out_dir
}

fn run_cargo_test(cargo_path: &Path, crate_path: &Path, packages: &[&str]) -> bool {
    let mut command = Command::new(cargo_path);
    command.arg("test");
    for name in packages {
        command.arg("-p").arg(name);
    }
    let status = command
        // Disable rust-lang/cargo's cross-compile tests
        .env("CFG_DISABLE_CROSS_TESTS", "1")
        // Relax #![deny(warnings)] in some crates
        .env("RUSTFLAGS", "--cap-lints warn")
        .current_dir(crate_path)
        .status()
        .unwrap();

    status.success()
}
<commit_msg>update tokei and ripgrep in cargotest<commit_after>use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

struct Test {
    repo: &'static str,
    name: &'static str,
    sha: &'static str,
    lock: Option<&'static str>,
    packages: &'static [&'static str],
}

const TEST_REPOS: &[Test] = &[
    Test {
        name: "iron",
        repo: "https://github.com/iron/iron",
        sha: "cf056ea5e8052c1feea6141e40ab0306715a2c33",
        lock: None,
        packages: &[],
    },
    Test {
        name: "ripgrep",
        repo: "https://github.com/BurntSushi/ripgrep",
        sha: "3de31f752729525d85a3d1575ac1978733b3f7e7",
        lock: None,
        packages: &[],
    },
    Test {
        name: "tokei",
        repo: "https://github.com/XAMPPRocky/tokei",
        sha: "fdf3f8cb279a7aeac0696c87e5d8b0cd946e4f9e",
        lock: None,
        packages: &[],
    },
    Test {
        name: "xsv",
        repo: "https://github.com/BurntSushi/xsv",
        sha: "66956b6bfd62d6ac767a6b6499c982eae20a2c9f",
        lock: None,
        packages: &[],
    },
    Test {
        name: "servo",
        repo: "https://github.com/servo/servo",
        sha: "caac107ae8145ef2fd20365e2b8fadaf09c2eb3b",
        lock: None,
        // Only test Stylo a.k.a. Quantum CSS, the parts of Servo going into Firefox.
        // This takes much less time to build than all of Servo and supports stable Rust.
        packages: &["selectors"],
    },
];

fn main() {
    let args = env::args().collect::<Vec<_>>();
    let cargo = &args[1];
    let out_dir = Path::new(&args[2]);
    let cargo = &Path::new(cargo);

    for test in TEST_REPOS.iter().rev() {
        test_repo(cargo, out_dir, test);
    }
}

fn test_repo(cargo: &Path, out_dir: &Path, test: &Test) {
    println!("testing {}", test.repo);
    let dir = clone_repo(test, out_dir);
    if let Some(lockfile) = test.lock {
        fs::write(&dir.join("Cargo.lock"), lockfile).unwrap();
    }
    if !run_cargo_test(cargo, &dir, test.packages) {
        panic!("tests failed for {}", test.repo);
    }
}

fn clone_repo(test: &Test, out_dir: &Path) -> PathBuf {
    let out_dir = out_dir.join(test.name);

    if !out_dir.join(".git").is_dir() {
        let status = Command::new("git").arg("init").arg(&out_dir).status().unwrap();
        assert!(status.success());
    }

    // Try progressively deeper fetch depths to find the commit
    let mut found = false;
    for depth in &[0, 1, 10, 100, 1000, 100000] {
        if *depth > 0 {
            let status = Command::new("git")
                .arg("fetch")
                .arg(test.repo)
                .arg("master")
                .arg(&format!("--depth={}", depth))
                .current_dir(&out_dir)
                .status()
                .unwrap();
            assert!(status.success());
        }

        let status = Command::new("git")
            .arg("reset")
            .arg(test.sha)
            .arg("--hard")
            .current_dir(&out_dir)
            .status()
            .unwrap();

        if status.success() {
            found = true;
            break;
        }
    }

    if !found {
        panic!("unable to find commit {}", test.sha)
    }
    let status =
        Command::new("git").arg("clean").arg("-fdx").current_dir(&out_dir).status().unwrap();
    assert!(status.success());

    out_dir
}

fn run_cargo_test(cargo_path: &Path, crate_path: &Path, packages: &[&str]) -> bool {
    let mut command = Command::new(cargo_path);
    command.arg("test");
    for name in packages {
        command.arg("-p").arg(name);
    }
    let status = command
        // Disable rust-lang/cargo's cross-compile tests
        .env("CFG_DISABLE_CROSS_TESTS", "1")
        // Relax #![deny(warnings)] in some crates
        .env("RUSTFLAGS", "--cap-lints warn")
        .current_dir(crate_path)
        .status()
        .unwrap();

    status.success()
}
<|endoftext|>


<commit_before>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint, RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float, FromPrimitive};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::iter::{FromIterator};
use std::clone::Clone;
use std::ops::{Shr, BitAnd};
use std::bool;

static pow_offsets : &'static[uint] = &[0u, 4u, 6u, 10u, 12u, 16u];

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Rabin-Miller primality test, probabilistic
fn try_composite(a : &BigUint, d : &BigUint, n : &BigUint, s : &BigUint) -> bool {

    let one : BigUint = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if big(i) == *s {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn rm_is_prime(n : &BigUint) -> bool {

    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_biguint_range(&two, n);
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn fermat_is_prime(n : &BigUint) -> bool {
    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    if n == &two {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&two, &(n-one), n) == one {
        true
    } else {
        false
    }
}

fn is_valid_pow(prime : &BigUint) -> bool {
    for offset in [big(0u), big(4u), big(6u), big(10u), big(12u), big(16u)].iter() {
        if !fermat_is_prime(&(*prime + *offset)) {
            return false;
        }
    }
    true
}

// Simple sieve
fn simple_sieve(primes : &mut Vec<bool>, max_sieve : uint, verbose : bool) -> Vec<uint> {
    #[inline]
    fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

    if primes.len() <= max_sieve {
        *primes = Vec::from_elem(max_sieve+1, true);
        *primes.get_mut(0) = false;
        *primes.get_mut(1) = false;
        for prime in range(2, int_sqrt(max_sieve) + 1) {
            if *primes.get(prime) {
                for multiple in iter::range_step(prime * prime, max_sieve + 1, prime) {
                    *primes.get_mut(multiple) = false
                }
                if verbose { println!("Sieved {}", prime) }
            }
        }
    }
    if verbose { println!("Finished sieving") }
    range(2, max_sieve+1).filter(|&n| *primes.get(n)).collect()
}

fn gen_prime(base_val : &BigUint, max_val : &BigUint, max_sieve : uint, verbose : bool) -> Vec<BigUint> {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }

    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);

    let mut primes : Vec<bool> = Vec::new();

    for &i in simple_sieve(&mut primes, max_sieve, verbose).iter() {
        if i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(i), big(primorial));
        primorial = primorial * i;
    }
    offsets.retain(|o| o >= base_val && is_valid_pow(o));
    offsets
}

#[cfg(test)]
mod test_primes {
    use super::{big, gen_prime, simple_sieve};

    #[test]
    fn gen_prime_is_correct() {
        assert!(gen_prime(&big(0), &big(100000000), 29, false).len() == 81u);
    }

    #[test]
    fn simple_sieve_is_correct() {
        assert!(simple_sieve(&mut Vec::new(), 5000000, false).len() == 348513);
    }
}

fn main() {
    println!("{}", gen_prime(&big(0), &big(100000000), 29, true).len());
}<commit_msg>Add tasks to prime cluster verification<commit_after>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint, RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float, FromPrimitive};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::iter::{FromIterator};
use std::clone::Clone;
use std::ops::{Shr, BitAnd};
use std::bool;

static pow_offsets : &'static[uint] = &[0u, 4u, 6u, 10u, 12u, 16u];

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Rabin-Miller primality test, probabilistic
fn try_composite(a : &BigUint, d : &BigUint, n : &BigUint, s : &BigUint) -> bool {

    let one : BigUint = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if big(i) == *s {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn rm_is_prime(n : &BigUint) -> bool {

    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_biguint_range(&two, n);
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn fermat_is_prime(n : &BigUint) -> bool {
    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    if n == &two {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&two, &(n-one), n) == one {
        true
    } else {
        false
    }
}

fn is_valid_pow(prime : &BigUint) -> bool {
    for offset in [big(0u), big(4u), big(6u), big(10u), big(12u), big(16u)].iter() {
        if !fermat_is_prime(&(*prime + *offset)) {
            return false;
        }
    }
    true
}

// Simple sieve
fn simple_sieve(primes : &mut Vec<bool>, max_sieve : uint, verbose : bool) -> Vec<uint> {
    #[inline]
    fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

    if primes.len() <= max_sieve {
        *primes = Vec::from_elem(max_sieve+1, true);
        *primes.get_mut(0) = false;
        *primes.get_mut(1) = false;
        for prime in range(2, int_sqrt(max_sieve) + 1) {
            if *primes.get(prime) {
                for multiple in iter::range_step(prime * prime, max_sieve + 1, prime) {
                    *primes.get_mut(multiple) = false
                }
                if verbose { println!("Sieved {}", prime) }
            }
        }
    }
    if verbose { println!("Finished sieving") }
    range(2, max_sieve+1).filter(|&n| *primes.get(n)).collect()
}

fn gen_prime(base_val : &BigUint, max_val : &BigUint, max_sieve : uint, verbose : bool) -> Vec<BigUint> {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }

    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);

    let mut primes : Vec<bool> = Vec::new();

    for &i in simple_sieve(&mut primes, max_sieve, verbose).iter() {
        if i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(i), big(primorial));
        primorial = primorial * i;
    }
    offsets.retain(|o| o >= base_val && is_valid_pow(o));
    offsets
}

fn gen_prime_par(base_val : &BigUint, max_val : &BigUint, max_sieve : uint, verbose : bool) -> Vec<BigUint> {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }

    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);

    let mut primes : Vec<bool> = Vec::new();

    for &i in simple_sieve(&mut primes, max_sieve, verbose).iter() {
        if i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(i), big(primorial));
        primorial = primorial * i;
    }

    offsets.retain(|o| o >= base_val);

    let rxs = offsets.iter().map( |o| {
        let (tx, rx) = channel();
        let b = o.clone();
        spawn(proc() {
            tx.send(is_valid_pow(&b));
        });
        rx
    });

    // Wait on each port, accumulating the results
    let mut results : Vec<BigUint> = Vec::new();
    for o in rxs.zip(offsets.iter()).filter_map(|(rx, o)| if rx.recv() {Some(o)} else {None}) {
        results.push(o.clone());
    }
    results
}

#[cfg(test)]
mod test_primes {
    use super::{big, gen_prime, simple_sieve};

    #[test]
    fn gen_prime_is_correct() {
        assert!(gen_prime(&big(0), &big(100000000), 29, false).len() == 81u);
    }

    #[test]
    fn simple_sieve_is_correct() {
        assert!(simple_sieve(&mut Vec::new(), 5000000, false).len() == 348513);
    }
}

fn main() {
    println!("{}", gen_prime(&big(0), &big(1000000000), 29, true).len());
}<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::collections::BTreeMap;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

pub struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    pub fn one() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    pub fn prime(p: u64) -> Factors {
        debug_assert!(miller_rabin::is_prime(p));
        let mut f = Factors::one();
        f.push(p);
        f
    }

    pub fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;
    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        // TODO: Optimise with 32 and 64b versions
        _factor::<A>(n, f)
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery>(n, factors)
    } else {
        _factor::<Montgomery>(n, factors)
    }
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }
}
<commit_msg>factor::Factors: Generalise invariant check (from prime() to add())<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::collections::BTreeMap;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

pub struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    pub fn one() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    pub fn prime(p: u64) -> Factors {
        let mut f = Factors::one();
        f.push(p);
        f
    }

    pub fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(miller_rabin::is_prime(prime));
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;
    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        // TODO: Optimise with 32 and 64b versions
        _factor::<A>(n, f)
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery>(n, factors)
    } else {
        _factor::<Montgomery>(n, factors)
    }
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::collections::BTreeMap;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    pub fn one() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    pub fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(miller_rabin::is_prime(prime));
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    }
}

#[cfg(test)]
mod tests {
    use super::{factor, Factors};
    use quickcheck::quickcheck;

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> bool {
            factor(f.product()) == f
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
        use rand::Rng;
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorisation.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = gen.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}
<commit_msg>factor: Refactor and optimise the `Factors` datastructure (#1572)<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::cell::RefCell;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

type Exponent = u8;

#[derive(Clone, Debug)]
struct Decomposition(Vec<(u64, Exponent)>);

impl Decomposition {
    fn one() -> Decomposition {
        Decomposition(Vec::new())
    }

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);

        if let Some((_, e)) = self.0.iter_mut().find(|(f, _)| *f == factor) {
            *e += exp;
        } else {
            self.0.push((factor, exp))
        }
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }

    fn get(&self, p: u64) -> Option<&(u64, u8)> {
        self.0.iter().find(|(q, _)| *q == p)
    }
}

impl PartialEq for Decomposition {
    fn eq(&self, other: &Decomposition) -> bool {
        for p in &self.0 {
            if other.get(p.0) != Some(p) {
                return false;
            }
        }

        for p in &other.0 {
            if self.get(p.0) != Some(p) {
                return false;
            }
        }

        true
    }
}
impl Eq for Decomposition {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors(RefCell<Decomposition>);

impl Factors {
    pub fn one() -> Factors {
        Factors(RefCell::new(Decomposition::one()))
    }

    pub fn add(&mut self, prime: u64, exp: Exponent) {
        debug_assert!(miller_rabin::is_prime(prime));
        self.0.borrow_mut().add(prime, exp)
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let v = &mut (self.0).borrow_mut().0;
        v.sort_unstable();

        for (p, exp) in v.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as Exponent);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    }
}

#[cfg(test)]
mod tests {
    use super::{factor, Factors};
    use quickcheck::quickcheck;

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> bool {
            assert_eq!(factor(f.product()), f);
            true
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
        use rand::Rng;
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorisation.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = gen.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use smallvec::SmallVec;
use std::cell::RefCell;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

type Exponent = u8;

#[derive(Clone, Debug)]
struct Decomposition(SmallVec<[(u64, Exponent); NUM_FACTORS_INLINE]>);

// spell-checker:ignore (names) Erdős–Kac * Erdős Kac
// The number of factors to inline directly into a `Decomposition` object.
// As a consequence of the Erdős–Kac theorem, the average number of prime factors
// of integers < 10²⁵ ≃ 2⁸³ is 4, so we can use a slightly higher value.
const NUM_FACTORS_INLINE: usize = 5;

impl Decomposition {
    fn one() -> Decomposition {
        Decomposition(SmallVec::new())
    }

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);

        if let Some((_, e)) = self.0.iter_mut().find(|(f, _)| *f == factor) {
            *e += exp;
        } else {
            self.0.push((factor, exp))
        }
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }

    fn get(&self, p: u64) -> Option<&(u64, u8)> {
        self.0.iter().find(|(q, _)| *q == p)
    }
}

impl PartialEq for Decomposition {
    fn eq(&self, other: &Decomposition) -> bool {
        for p in &self.0 {
            if other.get(p.0) != Some(p) {
                return false;
            }
        }

        for p in &other.0 {
            if self.get(p.0) != Some(p) {
                return false;
            }
        }

        true
    }
}
impl Eq for Decomposition {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors(RefCell<Decomposition>);

impl Factors {
    pub fn one() -> Factors {
        Factors(RefCell::new(Decomposition::one()))
    }

    pub fn add(&mut self, prime: u64, exp: Exponent) {
        debug_assert!(miller_rabin::is_prime(prime));
        self.0.borrow_mut().add(prime, exp)
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let v = &mut (self.0).borrow_mut().0;
        v.sort_unstable();

        for (p, exp) in v.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            #[cfg(feature = "coz")]
            coz::progress!("factor found");
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    #[cfg(feature = "coz")]
    coz::begin!("factorization");
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let n_zeros = n.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }

    if n == 1 {
        #[cfg(feature = "coz")]
        coz::end!("factorization");
        return factors;
    }

    table::factor(&mut n, &mut factors);

    #[allow(clippy::let_and_return)]
    let r = if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    };

    #[cfg(feature = "coz")]
    coz::end!("factorization");

    r
}

#[cfg(test)]
mod tests {
    use super::{factor, Decomposition, Exponent, Factors};
    use quickcheck::quickcheck;
    use smallvec::smallvec;
    use std::cell::RefCell;

    #[test]
    fn factor_2044854919485649() {
        let f = Factors(RefCell::new(Decomposition(smallvec![
            (503, 1),
            (2423, 1),
            (40961, 2)
        ])));
        assert_eq!(factor(f.product()), f);
    }

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's code path handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> () {
            assert_eq!(factor(f.product()), f);
        }

        fn exponentiate_factors(f: Factors, e: Exponent) -> () {
            if e == 0 { return; }
            if let Some(fe) = f.product().checked_pow(e.into()) {
                assert_eq!(factor(fe), f ^ e);
            }
        }
    }
}

#[cfg(test)]
use rand::{
    distributions::{Distribution, Standard},
    Rng,
};
#[cfg(test)]
impl Distribution<Factors> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Factors {
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // spell-checker:ignore (names) Adam Kalai * Kalai's
        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorization.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = rng.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Self {
        g.gen()
    }
}

#[cfg(test)]
impl std::ops::BitXor<Exponent> for Factors {
    type Output = Self;

    fn bitxor(self, rhs: Exponent) -> Factors {
        debug_assert_ne!(rhs, 0);
        let mut r = Factors::one();
        for (p, e) in self.0.borrow().0.iter() {
            r.add(*p, rhs * e);
        }

        debug_assert_eq!(r.product(), self.product().pow(rhs.into()));
        r
    }
}
<commit_msg>tests: update factor Distribution<Factors> sample to take range<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use smallvec::SmallVec;
use std::cell::RefCell;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

type Exponent = u8;

#[derive(Clone, Debug)]
struct Decomposition(SmallVec<[(u64, Exponent); NUM_FACTORS_INLINE]>);

// spell-checker:ignore (names) Erdős–Kac * Erdős Kac
// The number of factors to inline directly into a `Decomposition` object.
// As a consequence of the Erdős–Kac theorem, the average number of prime factors
// of integers < 10²⁵ ≃ 2⁸³ is 4, so we can use a slightly higher value.
const NUM_FACTORS_INLINE: usize = 5;

impl Decomposition {
    fn one() -> Decomposition {
        Decomposition(SmallVec::new())
    }

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);

        if let Some((_, e)) = self.0.iter_mut().find(|(f, _)| *f == factor) {
            *e += exp;
        } else {
            self.0.push((factor, exp))
        }
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }

    fn get(&self, p: u64) -> Option<&(u64, u8)> {
        self.0.iter().find(|(q, _)| *q == p)
    }
}

impl PartialEq for Decomposition {
    fn eq(&self, other: &Decomposition) -> bool {
        for p in &self.0 {
            if other.get(p.0) != Some(p) {
                return false;
            }
        }

        for p in &other.0 {
            if self.get(p.0) != Some(p) {
                return false;
            }
        }

        true
    }
}
impl Eq for Decomposition {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors(RefCell<Decomposition>);

impl Factors {
    pub fn one() -> Factors {
        Factors(RefCell::new(Decomposition::one()))
    }

    pub fn add(&mut self, prime: u64, exp: Exponent) {
        debug_assert!(miller_rabin::is_prime(prime));
        self.0.borrow_mut().add(prime, exp)
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let v = &mut (self.0).borrow_mut().0;
        v.sort_unstable();

        for (p, exp) in v.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            #[cfg(feature = "coz")]
            coz::progress!("factor found");
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    #[cfg(feature = "coz")]
    coz::begin!("factorization");
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let n_zeros = n.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }

    if n == 1 {
        #[cfg(feature = "coz")]
        coz::end!("factorization");
        return factors;
    }

    table::factor(&mut n, &mut factors);

    #[allow(clippy::let_and_return)]
    let r = if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    };

    #[cfg(feature = "coz")]
    coz::end!("factorization");

    r
}

#[cfg(test)]
mod tests {
    use super::{factor, Decomposition, Exponent, Factors};
    use quickcheck::quickcheck;
    use smallvec::smallvec;
    use std::cell::RefCell;

    #[test]
    fn factor_2044854919485649() {
        let f = Factors(RefCell::new(Decomposition(smallvec![
            (503, 1),
            (2423, 1),
            (40961, 2)
        ])));
        assert_eq!(factor(f.product()), f);
    }

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's code path handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> () {
            assert_eq!(factor(f.product()), f);
        }

        fn exponentiate_factors(f: Factors, e: Exponent) -> () {
            if e == 0 { return; }
            if let Some(fe) = f.product().checked_pow(e.into()) {
                assert_eq!(factor(fe), f ^ e);
            }
        }
    }
}

#[cfg(test)]
use rand::{
    distributions::{Distribution, Standard},
    Rng,
};
#[cfg(test)]
impl Distribution<Factors> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Factors {
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // spell-checker:ignore (names) Adam Kalai * Kalai's
        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorization.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = rng.gen_range(1..n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Self {
        g.gen()
    }
}

#[cfg(test)]
impl std::ops::BitXor<Exponent> for Factors {
    type Output = Self;

    fn bitxor(self, rhs: Exponent) -> Factors {
        debug_assert_ne!(rhs, 0);
        let mut r = Factors::one();
        for (p, e) in self.0.borrow().0.iter() {
            r.add(*p, rhs * e);
        }

        debug_assert_eq!(r.product(), self.product().pow(rhs.into()));
        r
    }
}
<|endoftext|>