<commit_before>#[derive(Debug, PartialEq)]
pub struct NumericLiteral(pub f64);

#[derive(Debug, PartialEq)]
pub struct StringLiteral<'a>(pub &'a str);

pub mod parsing {

    use nom::{IResult, digit, hex_digit, oct_digit};
    use parser::helpers::parsing::*;
    use parser::js::terminals::*;

    named!(pub numeric_literal<&str, NumericLiteral>, conv!(NumericLiteral(do_parse!(
        take_while_s!(is_whitespace) >>
        numeric_literal: alt!(
            decimal_literal
            |
            hex_integer_literal
            |
            octal_integer_literal
        ) >>
        (numeric_literal)
    ))));

    named!(decimal_literal<&str, f64>, do_parse!(
        dec_lit: recognize!(alt!(
            do_parse!(
                decimal_integer_literal >>
                do_parse!(
                    opt!(complete!(do_parse!(
                        tag_s!(".") >>
                        opt!(complete!(digit)) >>
                        ()
                    ))) >>
                    opt!(complete!(exponential_part)) >>
                    ()
                ) >>
                ()
            )
            |
            do_parse!(
                tag_s!(".") >>
                complete!(digit) >>
                opt!(complete!(exponential_part)) >>
                ()
            )
        )) >>
        (dec_lit.parse::<f64>().unwrap())
    ));

    named!(decimal_integer_literal<&str, &str>, recognize!(alt!(
        do_parse!(
            tag_s!("0") >>
            not!(digit) >>
            not!(tag_s!("x")) >>
            not!(tag_s!("X")) >>
            ()
        )
        |
        do_parse!(
            not!(tag_s!("0")) >>
            digit >>
            ()
        )
    )));

    named!(exponential_part<&str, &str>, recognize!(do_parse!(
        alt!(tag_s!("e") | tag_s!("E")) >>
        opt!(alt!(tag_s!("+") | tag_s!("-"))) >>
        digit >>
        ()
    )));

    named!(hex_integer_literal<&str, f64>, do_parse!(
        alt!(tag_s!("0x") | tag_s!("0X")) >>
        hex: hex_digit >>
        (i64::from_str_radix(hex, 16).unwrap() as f64)
    ));

    named!(octal_integer_literal<&str, f64>, do_parse!(
        tag_s!("0") >>
        octal: oct_digit >>
        (i64::from_str_radix(octal, 8).unwrap() as f64)
    ));

    named!(pub string_literal<&str, StringLiteral>, do_parse!(
        delimiter: alt!(
            keyword!("\"")
            |
            keyword!("'")
        ) >>
        string: take_until_s!(delimiter) >> // TODO: escape chars
        keyword!(delimiter) >>
        (StringLiteral(string))
    ));

    named!(pub identifier<&str, &str>, do_parse!(
        take_while_s!(is_whitespace) >>
        identifier: recognize!(do_parse!(
            xid_start >>
            many0!(xid_continue) >>
            ()
        )) >>
        (identifier)
    ));

    named!(pub js_identifier<&str, &str>, alt!(
        keyword!("property")
        |
        keyword!("signal")
        |
        keyword!("readonly")
        |
        keyword!("on")
        |
        keyword!("get")
        |
        keyword!("set")
        |
        identifier
    ));

    named!(pub reserved_identifier<&str, &str>, alt!(
        keyword!("break")
        |
        keyword!("case")
        |
        keyword!("catch")
        |
        keyword!("continue")
        |
        keyword!("default")
        |
        keyword!("delete")
        |
        keyword!("do")
        |
        keyword!("else")
        |
        keyword!("false")
        |
        keyword!("finally")
        |
        keyword!("for")
        |
        keyword!("function")
        |
        keyword!("if")
        |
        keyword!("in")
        |
        keyword!("instanceof")
        |
        keyword!("new")
        |
        keyword!("null")
        |
        keyword!("return")
        |
        keyword!("switch")
        |
        keyword!("this")
        |
        keyword!("throw")
        |
        keyword!("true")
        |
        keyword!("try")
        |
        keyword!("typeof")
        |
        keyword!("var")
        |
        keyword!("void")
        |
        keyword!("while")
        |
        keyword!("const")
        |
        keyword!("debugger")
        |
        keyword!("with")
    ));

    named!(pub property_identifier<&str, &str>, alt!(
        reserved_identifier
        |
        js_identifier
    ));

    named!(pub assignment_operator<&str, &str>, alt!(
        keyword!("=")
        |
        keyword!("*=")
        |
        keyword!("/=")
        |
        keyword!("%=")
        |
        keyword!("+=")
        |
        keyword!("-=")
        |
        keyword!("<<=")
        |
        keyword!(">>>=")
        |
        keyword!(">>=")
        |
        keyword!("&=")
        |
        keyword!("^=")
        |
        keyword!("|=")
    ));

    named!(pub automatic_semicolon<&str, &str>, alt!(
        keyword!(";")
        |
        peek!(line_terminator)
        |
        peek!(keyword!("}"))
        |
        eof
    ));

    fn eof(i: &str) -> IResult<&str, &str> {
        eof!(i,)
    }

    named!(pub line_terminator<&str, &str>, alt!(
        tag_s!("\r\n")
        |
        tag_s!("\n")
        |
        tag_s!("\u{2028}")
        |
        tag_s!("\u{2029}")
    ));

    #[cfg(test)]
    mod tests {
        use nom::{ErrorKind, Needed};
        use super::*;

        #[test]
        fn numeric_literal() {
            assert_eq!(super::numeric_literal(""), IResult::Incomplete(Needed::Size(1)));

            // Decimal
            assert_eq!(super::numeric_literal("42"), IResult::Done("", NumericLiteral(42_f64)));
            assert_eq!(super::numeric_literal(" 42"), IResult::Done("", NumericLiteral(42_f64)));
            assert_eq!(super::numeric_literal("0"), IResult::Done("", NumericLiteral(0_f64)));
            assert_eq!(super::numeric_literal("42.42e-4"), IResult::Done("", NumericLiteral(42.42e-4_f64)));
            assert_eq!(super::numeric_literal(".42"), IResult::Done("", NumericLiteral(0.42_f64)));
            assert_eq!(super::numeric_literal("42."), IResult::Done("", NumericLiteral(42_f64)));

            assert_eq!(super::numeric_literal("+42"), IResult::Error(ErrorKind::Alt)); // Identity is unary expression
            assert_eq!(super::numeric_literal("-42"), IResult::Error(ErrorKind::Alt)); // Negation is unary expression
            assert_eq!(super::numeric_literal("."), IResult::Error(ErrorKind::Alt));

            // Octal
            assert_eq!(super::numeric_literal("042"), IResult::Done("", NumericLiteral(34_f64)));

            // Hex
            assert_eq!(super::numeric_literal("0x42"), IResult::Done("", NumericLiteral(66_f64)));
            assert_eq!(super::numeric_literal("0X42"), IResult::Done("", NumericLiteral(66_f64)));
        }

        #[test]
        fn string_literal() {
            assert_eq!(super::string_literal("\"hi\""), IResult::Done("", StringLiteral("hi")));
            assert_eq!(super::string_literal(" 'hi'"), IResult::Done("", StringLiteral("hi")));

            assert_eq!(super::string_literal(""), IResult::Incomplete(Needed::Size(1)));

            assert_eq!(super::string_literal("\"hi"), IResult::Error(ErrorKind::TakeUntil)); // FIXME: https://github.com/Geal/nom/issues/397
            assert_eq!(super::string_literal("'hi"), IResult::Error(ErrorKind::TakeUntil));
        }

    }
}<commit_msg>Add unit test for identifier<commit_after>#[derive(Debug, PartialEq)]
pub struct NumericLiteral(pub f64);

#[derive(Debug, PartialEq)]
pub struct StringLiteral<'a>(pub &'a str);

pub mod parsing {

    use nom::{IResult, digit, hex_digit, oct_digit};
    use parser::helpers::parsing::*;
    use parser::js::terminals::*;

    named!(pub numeric_literal<&str, NumericLiteral>, conv!(NumericLiteral(do_parse!(
        take_while_s!(is_whitespace) >>
        numeric_literal: alt!(
            decimal_literal
            |
            hex_integer_literal
            |
            octal_integer_literal
        ) >>
        (numeric_literal)
    ))));

    named!(decimal_literal<&str, f64>, do_parse!(
        dec_lit: recognize!(alt!(
            do_parse!(
                decimal_integer_literal >>
                do_parse!(
                    opt!(complete!(do_parse!(
                        tag_s!(".") >>
                        opt!(complete!(digit)) >>
                        ()
                    ))) >>
                    opt!(complete!(exponential_part)) >>
                    ()
                ) >>
                ()
            )
            |
            do_parse!(
                tag_s!(".") >>
                complete!(digit) >>
                opt!(complete!(exponential_part)) >>
                ()
            )
        )) >>
        (dec_lit.parse::<f64>().unwrap())
    ));

    named!(decimal_integer_literal<&str, &str>, recognize!(alt!(
        do_parse!(
            tag_s!("0") >>
            not!(digit) >>
            not!(tag_s!("x")) >>
            not!(tag_s!("X")) >>
            ()
        )
        |
        do_parse!(
            not!(tag_s!("0")) >>
            digit >>
            ()
        )
    )));

    named!(exponential_part<&str, &str>, recognize!(do_parse!(
        alt!(tag_s!("e") | tag_s!("E")) >>
        opt!(alt!(tag_s!("+") | tag_s!("-"))) >>
        digit >>
        ()
    )));

    named!(hex_integer_literal<&str, f64>, do_parse!(
        alt!(tag_s!("0x") | tag_s!("0X")) >>
        hex: hex_digit >>
        (i64::from_str_radix(hex, 16).unwrap() as f64)
    ));

    named!(octal_integer_literal<&str, f64>, do_parse!(
        tag_s!("0") >>
        octal: oct_digit >>
        (i64::from_str_radix(octal, 8).unwrap() as f64)
    ));

    named!(pub string_literal<&str, StringLiteral>, do_parse!(
        delimiter: alt!(
            keyword!("\"")
            |
            keyword!("'")
        ) >>
        string: take_until_s!(delimiter) >> // TODO: escape chars
        keyword!(delimiter) >>
        (StringLiteral(string))
    ));

    named!(pub identifier<&str, &str>, do_parse!(
        take_while_s!(is_whitespace) >>
        identifier: recognize!(do_parse!(
            xid_start >>
            many0!(xid_continue) >>
            ()
        )) >>
        (identifier)
    ));

    named!(pub js_identifier<&str, &str>, alt!(
        keyword!("property")
        |
        keyword!("signal")
        |
        keyword!("readonly")
        |
        keyword!("on")
        |
        keyword!("get")
        |
        keyword!("set")
        |
        identifier
    ));

    named!(pub reserved_identifier<&str, &str>, alt!(
        keyword!("break")
        |
        keyword!("case")
        |
        keyword!("catch")
        |
        keyword!("continue")
        |
        keyword!("default")
        |
        keyword!("delete")
        |
        keyword!("do")
        |
        keyword!("else")
        |
        keyword!("false")
        |
        keyword!("finally")
        |
        keyword!("for")
        |
        keyword!("function")
        |
        keyword!("if")
        |
        keyword!("in")
        |
        keyword!("instanceof")
        |
        keyword!("new")
        |
        keyword!("null")
        |
        keyword!("return")
        |
        keyword!("switch")
        |
        keyword!("this")
        |
        keyword!("throw")
        |
        keyword!("true")
        |
        keyword!("try")
        |
        keyword!("typeof")
        |
        keyword!("var")
        |
        keyword!("void")
        |
        keyword!("while")
        |
        keyword!("const")
        |
        keyword!("debugger")
        |
        keyword!("with")
    ));

    named!(pub property_identifier<&str, &str>, alt!(
        reserved_identifier
        |
        js_identifier
    ));

    named!(pub assignment_operator<&str, &str>, alt!(
        keyword!("=")
        |
        keyword!("*=")
        |
        keyword!("/=")
        |
        keyword!("%=")
        |
        keyword!("+=")
        |
        keyword!("-=")
        |
        keyword!("<<=")
        |
        keyword!(">>>=")
        |
        keyword!(">>=")
        |
        keyword!("&=")
        |
        keyword!("^=")
        |
        keyword!("|=")
    ));

    named!(pub automatic_semicolon<&str, &str>, alt!(
        keyword!(";")
        |
        peek!(line_terminator)
        |
        peek!(keyword!("}"))
        |
        eof
    ));

    fn eof(i: &str) -> IResult<&str, &str> {
        eof!(i,)
    }

    named!(pub line_terminator<&str, &str>, alt!(
        tag_s!("\r\n")
        |
        tag_s!("\n")
        |
        tag_s!("\u{2028}")
        |
        tag_s!("\u{2029}")
    ));

    #[cfg(test)]
    mod tests {
        use nom::{ErrorKind, Needed};
        use super::*;

        #[test]
        fn numeric_literal() {
            assert_eq!(super::numeric_literal(""), IResult::Incomplete(Needed::Size(1)));

            // Decimal
            assert_eq!(super::numeric_literal("42"), IResult::Done("", NumericLiteral(42_f64)));
            assert_eq!(super::numeric_literal(" 42"), IResult::Done("", NumericLiteral(42_f64)));
            assert_eq!(super::numeric_literal("0"), IResult::Done("", NumericLiteral(0_f64)));
            assert_eq!(super::numeric_literal("42.42e-4"), IResult::Done("", NumericLiteral(42.42e-4_f64)));
            assert_eq!(super::numeric_literal(".42"), IResult::Done("", NumericLiteral(0.42_f64)));
            assert_eq!(super::numeric_literal("42."), IResult::Done("", NumericLiteral(42_f64)));

            assert_eq!(super::numeric_literal("+42"), IResult::Error(ErrorKind::Alt)); // Identity is unary expression
            assert_eq!(super::numeric_literal("-42"), IResult::Error(ErrorKind::Alt)); // Negation is unary expression
            assert_eq!(super::numeric_literal("."), IResult::Error(ErrorKind::Alt));

            // Octal
            assert_eq!(super::numeric_literal("042"), IResult::Done("", NumericLiteral(34_f64)));

            // Hex
            assert_eq!(super::numeric_literal("0x42"), IResult::Done("", NumericLiteral(66_f64)));
            assert_eq!(super::numeric_literal("0X42"), IResult::Done("", NumericLiteral(66_f64)));
        }

        #[test]
        fn string_literal() {
            assert_eq!(super::string_literal("\"hi\""), IResult::Done("", StringLiteral("hi")));
            assert_eq!(super::string_literal(" 'hi'"), IResult::Done("", StringLiteral("hi")));

            assert_eq!(super::string_literal(""), IResult::Incomplete(Needed::Size(1)));

            assert_eq!(super::string_literal("\"hi"), IResult::Error(ErrorKind::TakeUntil)); // FIXME: https://github.com/Geal/nom/issues/397
            assert_eq!(super::string_literal("'hi"), IResult::Error(ErrorKind::TakeUntil));
        }

        #[test]
        fn identifier() {
            assert_eq!(super::identifier("foo"), IResult::Done("", "foo"));
            assert_eq!(super::identifier(" foo"), IResult::Done("", "foo"));
            assert_eq!(super::identifier("foo bar"), IResult::Done(" bar", "foo"));

            assert_eq!(super::identifier(""), IResult::Incomplete(Needed::Size(1)));

            assert_eq!(super::identifier("42foo"), IResult::Error(ErrorKind::Verify));
        }
    }
}<|endoftext|>


<commit_before>use crate::parser::parse;

fn assert_err_msg(input: &str, needles: &[&str]) {
    let res = parse(input);
    assert!(res.is_err());
    let err = res.unwrap_err();
    let err_msg = err.to_string();
    println!("{}", err_msg);
    println!("Looking for:");
    for needle in needles {
        println!("{}", needle);
        assert!(err_msg.contains(needle));
    }
}

#[test]
fn invalid_number() {
    assert_err_msg(
        "{{ 1.2.2 }}",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn invalid_op() {
    assert_err_msg("{{ 1.2 >+ 3 }}", &["1:9", "expected a comparison value"]);
}

#[test]
fn wrong_start_block() {
    assert_err_msg(
        "{{ if true %}",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn wrong_end_block() {
    assert_err_msg(
        "{{ hey %}",
        &[
            "1:9",
            "expected an integer, a float, a string, `true` or `false`, an identifier (must start with a-z), a dotted identifier (identifiers separated by `.`), a square bracketed identifier (identifiers separated by `.` or `[]`s), a concatenation of strings, or an expression"
        ],
    );
}

#[test]
fn unterminated_variable_block() {
    assert_err_msg(
        "{{ hey",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_string() {
    assert_err_msg(r#"{{ "hey }}"#, &["1:4", "expected any expressions"]);
}

#[test]
fn unterminated_if_tag() {
    assert_err_msg(
        r#"{% if true %}sd"#,
        &[
            "1:16",
            r#"expected tag, an include tag (`{% include "..." %}`), a comment tag (`{#...#}`), a variable tag (`{{ ... }}`), or some text"#
        ],
    );
}

#[test]
fn unterminated_filter_section() {
    assert_err_msg(
        r#"{% filter uppercase %}sd"#,
        &[
            "1:25",
            r#"expected tag, an include tag (`{% include "..." %}`), a comment tag (`{#...#}`), a variable tag (`{{ ... }}`), or some text"#
        ],
    );
}

#[test]
fn invalid_filter_section_missing_name() {
    assert_err_msg(
        r#"{% filter %}sd{% endfilter %}"#,
        &["1:11", "expected an identifier (must start with a-z) or a function call"],
    );
}

#[test]
fn invalid_macro_content() {
    assert_err_msg(
        r#"
{% macro input(label, type) %}
    {% macro nested() %}
    {% endmacro nested %}
{% endmacro input %}
    "#,
        &["3:5", "unexpected tag; expected `{% endmacro %}` or the macro content"],
    );
}

#[test]
fn invalid_macro_default_arg_value() {
    assert_err_msg(
        r#"
{% macro input(label=something) %}
{% endmacro input %}
    "#,
        &["2:22", "expected an integer, a float, a string, or `true` or `false`"],
    );
}

#[test]
fn invalid_elif() {
    assert_err_msg(
        r#"
{% if true %}
{% else %}
{% elif false %}
{% endif %}
    "#,
        &["4:1", "unexpected tag; expected an endif tag (`{% endif %}`) or some content"],
    );
}

#[test]
fn invalid_else() {
    assert_err_msg(
        r#"
{% if true %}
{% else %}
{% else %}
{% endif %}
    "#,
        &["4:1", "unexpected tag; expected an endif tag (`{% endif %}`) or some content"],
    );
}

#[test]
fn invalid_extends_position() {
    assert_err_msg(
        r#"
hello
{% extends "hey.html" %}
    "#,
        &["3:1", "unexpected tag; expected end of input or some content"],
    );
}

#[test]
fn invalid_operator() {
    assert_err_msg(
        "{{ hey =! }}",
        &[
            "1:8",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn missing_expression_with_not() {
    assert_err_msg("{% if not %}", &["1:11", "expected an expression"]);
}

#[test]
fn missing_expression_in_if() {
    assert_err_msg("{% if %}", &["1:7", "expected any expression"]);
}

#[test]
fn missing_container_name_in_forloop() {
    assert_err_msg("{% for i in %}", &["1:13", "expected an expression with an optional filter"]);
}

#[test]
fn missing_variable_name_in_set() {
    assert_err_msg("{% set = 1 %}", &["1:8", "expected an identifier (must start with a-z)"]);
}

#[test]
fn missing_value_in_set() {
    assert_err_msg("{% set a =  %}", &["1:13", "expected any expressions"]);
}

#[test]
fn unterminated_fn_call() {
    assert_err_msg(
        "{{ a | slice( }}",
        &["1:15", "expected a keyword argument: `key=value` where `value` can be any expression"],
    );
}

#[test]
fn invalid_fn_call_missing_value() {
    assert_err_msg("{{ a | slice(start=) }}", &["1:20", "expected any expressions"]);
}

#[test]
fn unterminated_macro_call() {
    assert_err_msg(
        "{{ my::macro( }}",
        &["1:15", "expected a keyword argument: `key=value` where `value` can be any expression"],
    );
}

#[test]
fn invalid_macro_call() {
    assert_err_msg(
        "{{ my:macro() }}",
        &[
            "1:6",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_include() {
    assert_err_msg("{% include %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_include_no_string() {
    assert_err_msg("{% include 1 %}", &["1:12", "expected a string"]);
}

#[test]
fn unterminated_extends() {
    assert_err_msg("{% extends %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_extends_no_string() {
    assert_err_msg("{% extends 1 %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_import_macros_missing_filename() {
    assert_err_msg("{% import as macros %}", &["1:11", "expected a string"]);
}

#[test]
fn invalid_import_macros_missing_namespace() {
    assert_err_msg(
        r#"{% import "hello" as %}"#,
        &["1:22", "expected an identifier (must start with a-z)"],
    );
}

#[test]
fn invalid_block_missing_name() {
    assert_err_msg(r#"{% block %}"#, &["1:10", "expected an identifier (must start with a-z)"]);
}

#[test]
fn unterminated_test() {
    assert_err_msg(
        r#"{% if a is odd( %}"#,
        &["1:17", "expected a list of test arguments (any expressions)"],
    );
}

#[test]
fn invalid_test_argument() {
    assert_err_msg(
        r#"{% if a is odd(key=1) %}"#,
        &[
            "1:19",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_raw_tag() {
    assert_err_msg(r#"{% raw %}sd"#, &["1:12", "expected `{% endraw %}`"]);
}

#[test]
fn invalid_break_outside_loop() {
    assert_err_msg(r#"{% break %}"#, &["1:1", "{% break %}", "expected a template"]);
}

#[test]
fn invalid_continue_outside_loop() {
    assert_err_msg(r#"{% continue %}"#, &["1:1", "{% continue %}", "expected a template"]);
}
<commit_msg>Slightly worse error messages of parsers with new pest version<commit_after>use crate::parser::parse;

fn assert_err_msg(input: &str, needles: &[&str]) {
    let res = parse(input);
    assert!(res.is_err());
    let err = res.unwrap_err();
    let err_msg = err.to_string();
    println!("{}", err_msg);
    println!("Looking for:");
    for needle in needles {
        println!("{}", needle);
        assert!(err_msg.contains(needle));
    }
}

#[test]
fn invalid_number() {
    assert_err_msg(
        "{{ 1.2.2 }}",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn invalid_op() {
    assert_err_msg("{{ 1.2 >+ 3 }}", &["1:9", "expected an expression"]);
}

#[test]
fn wrong_start_block() {
    assert_err_msg(
        "{{ if true %}",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn wrong_end_block() {
    assert_err_msg(
        "{{ hey %}",
        &[
            "1:9",
            "expected an integer, a float, a string, `true` or `false`, an identifier (must start with a-z), a dotted identifier (identifiers separated by `.`), a square bracketed identifier (identifiers separated by `.` or `[]`s), a concatenation of strings, or an expression"
        ],
    );
}

#[test]
fn unterminated_variable_block() {
    assert_err_msg(
        "{{ hey",
        &[
            "1:7",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_string() {
    assert_err_msg(r#"{{ "hey }}"#, &["1:4", "expected a value that can be negated"]);
}

#[test]
fn unterminated_if_tag() {
    assert_err_msg(
        r#"{% if true %}sd"#,
        &[
            "1:16",
            r#"expected tag or some content"#
        ],
    );
}

#[test]
fn unterminated_filter_section() {
    assert_err_msg(
        r#"{% filter uppercase %}sd"#,
        &[
            "1:25",
            r#"expected tag or the filter section content"#
        ],
    );
}

#[test]
fn invalid_filter_section_missing_name() {
    assert_err_msg(
        r#"{% filter %}sd{% endfilter %}"#,
        &["1:11", "expected an identifier (must start with a-z)"],
    );
}

#[test]
fn invalid_macro_content() {
    assert_err_msg(
        r#"
{% macro input(label, type) %}
    {% macro nested() %}
    {% endmacro nested %}
{% endmacro input %}
    "#,
        &["3:5", "unexpected tag; expected `{% endmacro %}` or the macro content"],
    );
}

#[test]
fn invalid_macro_default_arg_value() {
    assert_err_msg(
        r#"
{% macro input(label=something) %}
{% endmacro input %}
    "#,
        &["2:22", "expected an integer, a float, a string, or `true` or `false`"],
    );
}

#[test]
fn invalid_elif() {
    assert_err_msg(
        r#"
{% if true %}
{% else %}
{% elif false %}
{% endif %}
    "#,
        &["4:1", "unexpected tag; expected an endif tag (`{% endif %}`) or some content"],
    );
}

#[test]
fn invalid_else() {
    assert_err_msg(
        r#"
{% if true %}
{% else %}
{% else %}
{% endif %}
    "#,
        &["4:1", "unexpected tag; expected an endif tag (`{% endif %}`) or some content"],
    );
}

#[test]
fn invalid_extends_position() {
    assert_err_msg(
        r#"
hello
{% extends "hey.html" %}
    "#,
        &["3:1", "unexpected tag; expected end of input or some content"],
    );
}

#[test]
fn invalid_operator() {
    assert_err_msg(
        "{{ hey =! }}",
        &[
            "1:8",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn missing_expression_with_not() {
    assert_err_msg("{% if not %}", &["1:11", "expected an expression"]);
}

#[test]
fn missing_expression_in_if() {
    assert_err_msg("{% if %}", &["1:7", "expected a value that can be negated"]);
}

#[test]
fn missing_container_name_in_forloop() {
    assert_err_msg("{% for i in %}", &["1:13", "expected an expression or an array of values"]);
}

#[test]
fn missing_variable_name_in_set() {
    assert_err_msg("{% set = 1 %}", &["1:8", "expected an identifier (must start with a-z)"]);
}

#[test]
fn missing_value_in_set() {
    assert_err_msg("{% set a =  %}", &["1:13", "expected a value that can be negated or an array of values"]);
}

#[test]
fn unterminated_fn_call() {
    assert_err_msg(
        "{{ a | slice( }}",
        &["1:15", "expected an identifier (must start with a-z)"],
    );
}

#[test]
fn invalid_fn_call_missing_value() {
    assert_err_msg("{{ a | slice(start=) }}", &["1:20", "expected a value that can be negated or an array of values"]);
}

#[test]
fn unterminated_macro_call() {
    assert_err_msg(
        "{{ my::macro( }}",
        &["1:15", "expected an identifier (must start with a-z)"],
    );
}

#[test]
fn invalid_macro_call() {
    assert_err_msg(
        "{{ my:macro() }}",
        &[
            "1:6",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_include() {
    assert_err_msg("{% include %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_include_no_string() {
    assert_err_msg("{% include 1 %}", &["1:12", "expected a string"]);
}

#[test]
fn unterminated_extends() {
    assert_err_msg("{% extends %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_extends_no_string() {
    assert_err_msg("{% extends 1 %}", &["1:12", "expected a string"]);
}

#[test]
fn invalid_import_macros_missing_filename() {
    assert_err_msg("{% import as macros %}", &["1:11", "expected a string"]);
}

#[test]
fn invalid_import_macros_missing_namespace() {
    assert_err_msg(
        r#"{% import "hello" as %}"#,
        &["1:22", "expected an identifier (must start with a-z)"],
    );
}

#[test]
fn invalid_block_missing_name() {
    assert_err_msg(r#"{% block %}"#, &["1:10", "expected an identifier (must start with a-z)"]);
}

#[test]
fn unterminated_test() {
    assert_err_msg(
        r#"{% if a is odd( %}"#,
        &["1:17", "expected a value that can be negated"],
    );
}

#[test]
fn invalid_test_argument() {
    assert_err_msg(
        r#"{% if a is odd(key=1) %}"#,
        &[
            "1:19",
            "expected `or`, `and`, `<=`, `>=`, `<`, `>`, `==`, `!=`, `+`, `-`, `*`, `/`, `%`, or a filter"
        ],
    );
}

#[test]
fn unterminated_raw_tag() {
    assert_err_msg(r#"{% raw %}sd"#, &["1:12", "expected tag"]);
}

#[test]
fn invalid_break_outside_loop() {
    assert_err_msg(r#"{% break %}"#, &["1:1", "{% break %}", "expected a template"]);
}

#[test]
fn invalid_continue_outside_loop() {
    assert_err_msg(r#"{% continue %}"#, &["1:1", "{% continue %}", "expected a template"]);
}
<|endoftext|>


<commit_before>//! tracksegment handles parsing of GPX-spec track segments.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;

use parser::waypoint;
use parser::Context;

use TrackSegment;

enum TrackSegmentEvent {
    StartTrkSeg,
    StartTrkPt,
    EndTrkSeg,
    Ignore,
}

/// consume consumes a GPX track segment from the `reader` until it ends.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<TrackSegment> {
    let mut segment: TrackSegment = Default::default();

    loop {
        // Peep into the reader and see what type of event is next. Based on
        // that information, we'll either forward the event to a downstream
        // module or take the information for ourselves.
        let event: Result<TrackSegmentEvent> = {
            if let Some(next) = context.reader.peek() {
                match next {
                    &Ok(XmlEvent::StartElement { ref name, .. }) => {
                        match name.local_name.as_ref() {
                            "trkseg" => Ok(TrackSegmentEvent::StartTrkSeg),
                            "trkpt" => Ok(TrackSegmentEvent::StartTrkPt),
                            child => Err(Error::from(ErrorKind::InvalidChildElement(
                                String::from(child),
                                "tracksegment",
                            ))),
                        }
                    }

                    &Ok(XmlEvent::EndElement { .. }) => Ok(TrackSegmentEvent::EndTrkSeg),

                    _ => Ok(TrackSegmentEvent::Ignore),
                }
            } else {
                break;
            }
        };

        match event.chain_err(|| Error::from("error while parsing track segment event"))? {
            TrackSegmentEvent::StartTrkSeg => {
                context.reader.next();
            }

            TrackSegmentEvent::StartTrkPt => {
                segment.points.push(waypoint::consume(context, "trkpt")?);
            }

            TrackSegmentEvent::EndTrkSeg => {
                context.reader.next();

                return Ok(segment);
            }

            TrackSegmentEvent::Ignore => {
                context.reader.next();
            }
        }
    }

    return Err("no end tag for track segment".into());
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;
    use geo::length::Length;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_full_trkseg() {
        let segment = consume!(
            "
            <trkseg>
                <trkpt lon=\"-77.0365\" lat=\"38.8977\">
                    <name>The White House</name>
                </trkpt>
                <trkpt lon=\"-71.063611\" lat=\"42.358056\">
                    <name>Boston, Massachusetts</name>
                </trkpt>
                <trkpt lon=\"-69.7832\" lat=\"44.31055\">
                    <name>Augusta, Maine</name>
                </trkpt>
            </trkseg>",
            GpxVersion::Gpx11
        );

        assert!(segment.is_ok());
        let segment = segment.unwrap();

        assert_eq!(segment.points.len(), 3);

        let linestring = segment.linestring();
        assert_approx_eq!(linestring.length(), 9.2377437);
    }

    #[test]
    fn consume_empty() {
        let segment = consume!("<trkseg></trkseg>", GpxVersion::Gpx11);

        assert!(segment.is_ok());
        let segment = segment.unwrap();

        assert_eq!(segment.points.len(), 0);
    }
}
<commit_msg>Rewrite tracksegment::consume cleaner and more strict - Do not accept arbitrary amounts of opening trkseg tags, we   require exactly one - Do not accept any closing tag as tag that closes the tracksegment,   require it to be trkseg - Use new custom errors InvalidClosingTag and MissingClosingTag - Avoid indirection of TrackSegmentEvents, handle XmlEvents more directly<commit_after>//! tracksegment handles parsing of GPX-spec track segments.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;

use parser::waypoint;
use parser::Context;
use parser::verify_starting_tag;

use TrackSegment;

/// consume consumes a GPX track segment from the `reader` until it ends.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<TrackSegment> {
    let mut segment: TrackSegment = Default::default();
    verify_starting_tag(context, "trkseg")?;

    loop {
        let next_event = {
            if let Some(next) = context.reader.peek() {
                next.clone()
            } else {
                break;
            }
        };

        match next_event.chain_err(|| Error::from("error while parsing tracksegment event"))? {
            XmlEvent::StartElement { ref name, .. } => match name.local_name.as_ref() {
                "trkpt" => segment.points.push(waypoint::consume(context, "trkpt")?),
                child => {
                    bail!(ErrorKind::InvalidChildElement(
                        String::from(child),
                        "tracksegment"
                    ));
                }
            },
            XmlEvent::EndElement { ref name } => {
                ensure!(
                    name.local_name == "trkseg",
                    ErrorKind::InvalidClosingTag(name.local_name.clone(), "trksegment")
                );
                context.reader.next(); //consume the end tag
                return Ok(segment);
            }
            _ => {
                context.reader.next(); //consume and ignore this event
            }
        }
    }

    bail!(ErrorKind::MissingClosingTag("tracksegment"));
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;
    use geo::length::Length;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_full_trkseg() {
        let segment = consume!(
            "
            <trkseg>
                <trkpt lon=\"-77.0365\" lat=\"38.8977\">
                    <name>The White House</name>
                </trkpt>
                <trkpt lon=\"-71.063611\" lat=\"42.358056\">
                    <name>Boston, Massachusetts</name>
                </trkpt>
                <trkpt lon=\"-69.7832\" lat=\"44.31055\">
                    <name>Augusta, Maine</name>
                </trkpt>
            </trkseg>",
            GpxVersion::Gpx11
        );

        assert!(segment.is_ok());
        let segment = segment.unwrap();

        assert_eq!(segment.points.len(), 3);

        let linestring = segment.linestring();
        assert_approx_eq!(linestring.length(), 9.2377437);
    }

    #[test]
    fn consume_empty() {
        let segment = consume!("<trkseg></trkseg>", GpxVersion::Gpx11);

        assert!(segment.is_ok());
        let segment = segment.unwrap();

        assert_eq!(segment.points.len(), 0);
    }
}
<|endoftext|>


<commit_before>/*
 * Region resolution. This pass runs before typechecking and resolves region
 * names to the appropriate block.
 */

import driver::session::session;
import middle::ty;
import syntax::{ast, visit};
import util::common::new_def_hash;

import std::list;
import std::list::list;
import std::map;
import std::map::hashmap;

/* Represents the type of the most immediate parent node. */
enum parent {
    pa_item(ast::node_id),
    pa_block(ast::node_id),
    pa_nested_fn(ast::node_id),
    pa_crate
}

/* Records the binding site of a region name. */
type binding = {
    name: str,
    id: ast::def_id
};

type region_map = {
    /* Mapping from a block/function expression to its parent. */
    parents: hashmap<ast::node_id,ast::node_id>,
    /* Mapping from a region type in the AST to its resolved region. */
    ast_type_to_region: hashmap<ast::node_id,ty::region>,
    /* Mapping from a local variable to its containing block. */
    local_blocks: hashmap<ast::node_id,ast::node_id>,
    /* Mapping from a region name to its function. */
    region_name_to_fn: hashmap<ast::def_id,ast::node_id>,
};

type ctxt = {
    sess: session,
    def_map: resolve::def_map,
    region_map: @region_map,
    mut bindings: @list<binding>,

    /*
     * A list of local IDs that will be parented to the next block we
     * traverse. This is used when resolving `alt` statements. Since we see
     * the pattern before the associated block, upon seeing a pattern we must
     * parent all the bindings in that pattern to the next block we see.
     */
    mut queued_locals: [ast::node_id],

    parent: parent,

    /* True if we're within the pattern part of an alt, false otherwise. */
    in_alt: bool
};

fn region_to_scope(region_map: @region_map, region: ty::region)
        -> ast::node_id {
    ret alt region {
        ty::re_caller(def_id) { def_id.node }
        ty::re_named(def_id) { region_map.region_name_to_fn.get(def_id) }
        ty::re_block(node_id) { node_id }
    };
}

// Returns true if `subscope` is equal to or is lexically nested inside
// `superscope` and false otherwise.
fn scope_contains(region_map: @region_map, superscope: ast::node_id,
                  subscope: ast::node_id) -> bool {
    let subscope = subscope;
    while superscope != subscope {
        alt region_map.parents.find(subscope) {
            none { ret false; }
            some(scope) { subscope = scope; }
        }
    }
    ret true;
}

fn resolve_ty(ty: @ast::ty, cx: ctxt, visitor: visit::vt<ctxt>) {
    alt ty.node {
        ast::ty_rptr({id: region_id, node: node}, _) {
            let region;
            alt node {
                ast::re_inferred {
                    // We infer to the caller region if we're at item scope
                    // and to the block region if we're at block scope.
                    //
                    // TODO: What do we do if we're in an alt?

                    alt cx.parent {
                        pa_item(item_id) | pa_nested_fn(item_id) {
                            let def_id = {crate: ast::local_crate,
                                          node: item_id};
                            region = ty::re_caller(def_id);
                        }
                        pa_block(block_id) {
                            region = ty::re_block(block_id);
                        }
                        pa_crate {
                            cx.sess.span_bug(ty.span, "inferred region at " +
                                             "crate level?!");
                        }
                    }
                }

                ast::re_named(ident) {
                    // If at item scope, introduce or reuse a binding. If at
                    // block scope, require that the binding be introduced.
                    let bindings = cx.bindings;
                    alt list::find(*bindings, {|b| ident == b.name}) {
                        some(binding) { region = ty::re_named(binding.id); }
                        none {
                            let def_id = {crate: ast::local_crate,
                                          node: region_id};
                            let binding = {name: ident, id: def_id};
                            cx.bindings = @list::cons(binding, cx.bindings);
                            region = ty::re_named(def_id);

                            alt cx.parent {
                                pa_item(fn_id) | pa_nested_fn(fn_id) {
                                    let rf = cx.region_map.region_name_to_fn;
                                    rf.insert(def_id, fn_id);
                                }
                                pa_block(_) {
                                    cx.sess.span_err(ty.span,
                                                     "unknown region `" +
                                                     ident + "`");
                                }
                                pa_crate {
                                    cx.sess.span_bug(ty.span, "named " +
                                                     "region at crate " +
                                                     "level?!");
                                }
                            }
                        }
                    }
                }

                ast::re_self {
                    // For blocks, "self" means "the current block".
                    //
                    // TODO: What do we do in an alt?
                    //
                    // FIXME: Doesn't work in type items.

                    alt cx.parent {
                        pa_item(item_id) | pa_nested_fn(item_id) {
                            let def_id = {crate: ast::local_crate,
                                          node: item_id};
                            region = ty::re_caller(def_id);
                        }
                        pa_block(block_id) {
                            region = ty::re_block(block_id);
                        }
                        pa_crate {
                            cx.sess.span_bug(ty.span,
                                             "region type outside item?!");
                        }
                    }
                }

            }

            cx.region_map.ast_type_to_region.insert(region_id, region);
        }
        _ { /* nothing to do */ }
    }

    visit::visit_ty(ty, cx, visitor);
}

fn record_parent(cx: ctxt, child_id: ast::node_id) {
    alt cx.parent {
        pa_item(parent_id) | pa_block(parent_id) | pa_nested_fn(parent_id) {
            cx.region_map.parents.insert(child_id, parent_id);
        }
        pa_crate { /* no-op */ }
    }
}

fn resolve_block(blk: ast::blk, cx: ctxt, visitor: visit::vt<ctxt>) {
    // Record the parent of this block.
    record_parent(cx, blk.node.id);

    // Resolve queued locals to this block.
    for local_id in cx.queued_locals {
        cx.region_map.local_blocks.insert(local_id, blk.node.id);
    }

    // Descend.
    let new_cx: ctxt = {parent: pa_block(blk.node.id),
                        mut queued_locals: [],
                        in_alt: false with cx};
    visit::visit_block(blk, new_cx, visitor);
}

fn resolve_arm(arm: ast::arm, cx: ctxt, visitor: visit::vt<ctxt>) {
    let new_cx: ctxt = {mut queued_locals: [], in_alt: true with cx};
    visit::visit_arm(arm, new_cx, visitor);
}

fn resolve_pat(pat: @ast::pat, cx: ctxt, visitor: visit::vt<ctxt>) {
    alt pat.node {
        ast::pat_ident(path, _) {
            let defn_opt = cx.def_map.find(pat.id);
            alt defn_opt {
                some(ast::def_variant(_,_)) {
                    /* Nothing to do; this names a variant. */
                }
                _ {
                    /*
                     * This names a local. Enqueue it or bind it to the
                     * containing block, depending on whether we're in an alt
                     * or not.
                     */
                    if cx.in_alt {
                        vec::push(cx.queued_locals, pat.id);
                    } else {
                        alt cx.parent {
                            pa_block(block_id) {
                                let local_blocks = cx.region_map.local_blocks;
                                local_blocks.insert(pat.id, block_id);
                            }
                            _ {
                                cx.sess.span_bug(pat.span,
                                                 "unexpected parent");
                            }
                        }
                    }
                }
            }
        }
        _ { /* no-op */ }
    }

    visit::visit_pat(pat, cx, visitor);
}

fn resolve_expr(expr: @ast::expr, cx: ctxt, visitor: visit::vt<ctxt>) {
    alt expr.node {
        ast::expr_fn(_, _, _, _) | ast::expr_fn_block(_, _) {
            record_parent(cx, expr.id);
            let new_cx = {parent: pa_nested_fn(expr.id),
                          in_alt: false with cx};
            visit::visit_expr(expr, new_cx, visitor);
        }
        _ { visit::visit_expr(expr, cx, visitor); }
    }
}

fn resolve_item(item: @ast::item, cx: ctxt, visitor: visit::vt<ctxt>) {
    // Items create a new outer block scope as far as we're concerned.
    let new_cx: ctxt = {bindings: @list::nil,
                        parent: pa_item(item.id),
                        in_alt: false
                        with cx};
    visit::visit_item(item, new_cx, visitor);
}

fn resolve_crate(sess: session, def_map: resolve::def_map, crate: @ast::crate)
        -> @region_map {
    let cx: ctxt = {sess: sess,
                    def_map: def_map,
                    region_map: @{parents: map::new_int_hash(),
                                  ast_type_to_region: map::new_int_hash(),
                                  local_blocks: map::new_int_hash(),
                                  region_name_to_fn: new_def_hash()},
                    mut bindings: @list::nil,
                    mut queued_locals: [],
                    parent: pa_crate,
                    in_alt: false};
    let visitor = visit::mk_vt(@{
        visit_block: resolve_block,
        visit_item: resolve_item,
        visit_ty: resolve_ty,
        visit_arm: resolve_arm,
        visit_pat: resolve_pat,
        visit_expr: resolve_expr
        with *visit::default_visitor()
    });
    visit::visit_crate(*crate, cx, visitor);
    ret cx.region_map;
}

<commit_msg>rustc: Record what `&` would resolve to for every AST type<commit_after>/*
 * Region resolution. This pass runs before typechecking and resolves region
 * names to the appropriate block.
 */

import driver::session::session;
import middle::ty;
import syntax::{ast, visit};
import util::common::new_def_hash;

import std::list;
import std::list::list;
import std::map;
import std::map::hashmap;

/* Represents the type of the most immediate parent node. */
enum parent {
    pa_item(ast::node_id),
    pa_block(ast::node_id),
    pa_nested_fn(ast::node_id),
    pa_crate
}

/* Records the binding site of a region name. */
type binding = {
    name: str,
    id: ast::def_id
};

type region_map = {
    /* Mapping from a block/function expression to its parent. */
    parents: hashmap<ast::node_id,ast::node_id>,
    /* Mapping from a region type in the AST to its resolved region. */
    ast_type_to_region: hashmap<ast::node_id,ty::region>,
    /* Mapping from a local variable to its containing block. */
    local_blocks: hashmap<ast::node_id,ast::node_id>,
    /* Mapping from a region name to its function. */
    region_name_to_fn: hashmap<ast::def_id,ast::node_id>,
    /* Mapping from an AST type node to the region that `&` resolves to in it. */
    ast_type_to_inferred_region: hashmap<ast::node_id,ty::region>
};

type ctxt = {
    sess: session,
    def_map: resolve::def_map,
    region_map: @region_map,
    mut bindings: @list<binding>,

    /*
     * A list of local IDs that will be parented to the next block we
     * traverse. This is used when resolving `alt` statements. Since we see
     * the pattern before the associated block, upon seeing a pattern we must
     * parent all the bindings in that pattern to the next block we see.
     */
    mut queued_locals: [ast::node_id],

    parent: parent,

    /* True if we're within the pattern part of an alt, false otherwise. */
    in_alt: bool
};

fn region_to_scope(region_map: @region_map, region: ty::region)
        -> ast::node_id {
    ret alt region {
        ty::re_caller(def_id) { def_id.node }
        ty::re_named(def_id) { region_map.region_name_to_fn.get(def_id) }
        ty::re_block(node_id) { node_id }
    };
}

// Returns true if `subscope` is equal to or is lexically nested inside
// `superscope` and false otherwise.
fn scope_contains(region_map: @region_map, superscope: ast::node_id,
                  subscope: ast::node_id) -> bool {
    let subscope = subscope;
    while superscope != subscope {
        alt region_map.parents.find(subscope) {
            none { ret false; }
            some(scope) { subscope = scope; }
        }
    }
    ret true;
}

fn get_inferred_region(cx: ctxt, sp: syntax::codemap::span) -> ty::region {
    // We infer to the caller region if we're at item scope
    // and to the block region if we're at block scope.
    //
    // TODO: What do we do if we're in an alt?

    ret alt cx.parent {
        pa_item(item_id) | pa_nested_fn(item_id) {
            ty::re_caller({crate: ast::local_crate, node: item_id})
        }
        pa_block(block_id) { ty::re_block(block_id) }
        pa_crate { cx.sess.span_bug(sp, "inferred region at crate level?!"); }
    }
}

fn resolve_ty(ty: @ast::ty, cx: ctxt, visitor: visit::vt<ctxt>) {
    let inferred_region = get_inferred_region(cx, ty.span);
    cx.region_map.ast_type_to_inferred_region.insert(ty.id, inferred_region);

    alt ty.node {
        ast::ty_rptr({id: region_id, node: node}, _) {
            let region;
            alt node {
                ast::re_inferred { region = inferred_region; }
                ast::re_named(ident) {
                    // If at item scope, introduce or reuse a binding. If at
                    // block scope, require that the binding be introduced.
                    let bindings = cx.bindings;
                    alt list::find(*bindings, {|b| ident == b.name}) {
                        some(binding) { region = ty::re_named(binding.id); }
                        none {
                            let def_id = {crate: ast::local_crate,
                                          node: region_id};
                            let binding = {name: ident, id: def_id};
                            cx.bindings = @list::cons(binding, cx.bindings);
                            region = ty::re_named(def_id);

                            alt cx.parent {
                                pa_item(fn_id) | pa_nested_fn(fn_id) {
                                    let rf = cx.region_map.region_name_to_fn;
                                    rf.insert(def_id, fn_id);
                                }
                                pa_block(_) {
                                    cx.sess.span_err(ty.span,
                                                     "unknown region `" +
                                                     ident + "`");
                                }
                                pa_crate {
                                    cx.sess.span_bug(ty.span, "named " +
                                                     "region at crate " +
                                                     "level?!");
                                }
                            }
                        }
                    }
                }

                ast::re_self {
                    // For blocks, "self" means "the current block".
                    //
                    // TODO: What do we do in an alt?
                    //
                    // FIXME: Doesn't work in type items.

                    alt cx.parent {
                        pa_item(item_id) | pa_nested_fn(item_id) {
                            let def_id = {crate: ast::local_crate,
                                          node: item_id};
                            region = ty::re_caller(def_id);
                        }
                        pa_block(block_id) {
                            region = ty::re_block(block_id);
                        }
                        pa_crate {
                            cx.sess.span_bug(ty.span,
                                             "region type outside item?!");
                        }
                    }
                }

            }

            cx.region_map.ast_type_to_region.insert(region_id, region);
        }
        _ { /* nothing to do */ }
    }

    visit::visit_ty(ty, cx, visitor);
}

fn record_parent(cx: ctxt, child_id: ast::node_id) {
    alt cx.parent {
        pa_item(parent_id) | pa_block(parent_id) | pa_nested_fn(parent_id) {
            cx.region_map.parents.insert(child_id, parent_id);
        }
        pa_crate { /* no-op */ }
    }
}

fn resolve_block(blk: ast::blk, cx: ctxt, visitor: visit::vt<ctxt>) {
    // Record the parent of this block.
    record_parent(cx, blk.node.id);

    // Resolve queued locals to this block.
    for local_id in cx.queued_locals {
        cx.region_map.local_blocks.insert(local_id, blk.node.id);
    }

    // Descend.
    let new_cx: ctxt = {parent: pa_block(blk.node.id),
                        mut queued_locals: [],
                        in_alt: false with cx};
    visit::visit_block(blk, new_cx, visitor);
}

fn resolve_arm(arm: ast::arm, cx: ctxt, visitor: visit::vt<ctxt>) {
    let new_cx: ctxt = {mut queued_locals: [], in_alt: true with cx};
    visit::visit_arm(arm, new_cx, visitor);
}

fn resolve_pat(pat: @ast::pat, cx: ctxt, visitor: visit::vt<ctxt>) {
    alt pat.node {
        ast::pat_ident(path, _) {
            let defn_opt = cx.def_map.find(pat.id);
            alt defn_opt {
                some(ast::def_variant(_,_)) {
                    /* Nothing to do; this names a variant. */
                }
                _ {
                    /*
                     * This names a local. Enqueue it or bind it to the
                     * containing block, depending on whether we're in an alt
                     * or not.
                     */
                    if cx.in_alt {
                        vec::push(cx.queued_locals, pat.id);
                    } else {
                        alt cx.parent {
                            pa_block(block_id) {
                                let local_blocks = cx.region_map.local_blocks;
                                local_blocks.insert(pat.id, block_id);
                            }
                            _ {
                                cx.sess.span_bug(pat.span,
                                                 "unexpected parent");
                            }
                        }
                    }
                }
            }
        }
        _ { /* no-op */ }
    }

    visit::visit_pat(pat, cx, visitor);
}

fn resolve_expr(expr: @ast::expr, cx: ctxt, visitor: visit::vt<ctxt>) {
    alt expr.node {
        ast::expr_fn(_, _, _, _) | ast::expr_fn_block(_, _) {
            record_parent(cx, expr.id);
            let new_cx = {parent: pa_nested_fn(expr.id),
                          in_alt: false with cx};
            visit::visit_expr(expr, new_cx, visitor);
        }
        _ { visit::visit_expr(expr, cx, visitor); }
    }
}

fn resolve_item(item: @ast::item, cx: ctxt, visitor: visit::vt<ctxt>) {
    // Items create a new outer block scope as far as we're concerned.
    let new_cx: ctxt = {bindings: @list::nil,
                        parent: pa_item(item.id),
                        in_alt: false
                        with cx};
    visit::visit_item(item, new_cx, visitor);
}

fn resolve_crate(sess: session, def_map: resolve::def_map, crate: @ast::crate)
        -> @region_map {
    let cx: ctxt = {sess: sess,
                    def_map: def_map,
                    region_map: @{parents: map::new_int_hash(),
                                  ast_type_to_region: map::new_int_hash(),
                                  local_blocks: map::new_int_hash(),
                                  region_name_to_fn: new_def_hash(),
                                  ast_type_to_inferred_region:
                                    map::new_int_hash()},
                    mut bindings: @list::nil,
                    mut queued_locals: [],
                    parent: pa_crate,
                    in_alt: false};
    let visitor = visit::mk_vt(@{
        visit_block: resolve_block,
        visit_item: resolve_item,
        visit_ty: resolve_ty,
        visit_arm: resolve_arm,
        visit_pat: resolve_pat,
        visit_expr: resolve_expr
        with *visit::default_visitor()
    });
    visit::visit_crate(*crate, cx, visitor);
    ret cx.region_map;
}

<|endoftext|>


<commit_before>//! Provides a queue of block indices that the sampler positions can be initialized
//! from for the worker threads. The queue itself is not changed after creation
//! we simply work through it with an atomic counter to track the index of the next
//! block to work on

use std::vec::Vec;
use std::sync::atomic::{AtomicUint, AcqRel};
use sampler::morton;

/// The queue of blocks to be worked on shared immutably between worker threads.
pub struct BlockQueue {
    /// The block indices of blocks to work on for the image
    blocks: Vec<(u32, u32)>,
    /// Get the dimensions of an individual block
    dimensions: (u32, u32),
    /// Index of the next block to be worked on
    next: AtomicUint,
}

impl BlockQueue {
    /// Create a block queue for the image with dimensions `img`.
    /// Panics if the image is not evenly broken into blocks of dimension `dim`
    pub fn new(img: (u32, u32), dim: (u32, u32)) -> BlockQueue {
        if img.0 % dim.0 != 0 || img.1 % dim.1 != 0 {
            panic!("Image with dimension {} not evenly divided by dims of {}",
                   img, dim);
        }
        let num_blocks = (img.0 / dim.0, img.1 / dim.1);
        let mut blocks: Vec<(u32, u32)> = range(0, num_blocks.0 * num_blocks.1)
            .map(|i| (i % num_blocks.0, i / num_blocks.0)).collect();
        blocks.sort_by(|a, b| morton::morton2(a).cmp(&morton::morton2(b)));
        BlockQueue { blocks: blocks, dimensions: dim, next: AtomicUint::new(0) }
    }
    /// Get the dimensions of an individual block in the queue
    pub fn block_dim(&self) -> (u32, u32) { self.dimensions }
}

impl Iterator<(u32, u32)> for BlockQueue {
    fn next(&mut self) -> Option<(u32, u32)> {
        let i = self.next.fetch_add(1, AcqRel);
        if i >= self.blocks.len() {
            None
        } else {
            Some(self.blocks[i])
        }
    }
    fn size_hint(&self) -> (uint, Option<uint>) {
        let i = self.next.load(AcqRel);
        (i, Some(i))
    }
}

<commit_msg>Make it possible to iterate cleanly over the queue on multiple threads<commit_after>//! Provides a queue of block indices that the sampler positions can be initialized
//! from for the worker threads. The queue itself is not changed after creation
//! we simply work through it with an atomic counter to track the index of the next
//! block to work on

use std::vec::Vec;
use std::sync::atomic::{AtomicUint, AcqRel};
use sampler::morton;

/// The queue of blocks to be worked on shared immutably between worker threads.
pub struct BlockQueue {
    /// The block indices of blocks to work on for the image
    blocks: Vec<(u32, u32)>,
    /// Get the dimensions of an individual block
    dimensions: (u32, u32),
    /// Index of the next block to be worked on
    next: AtomicUint,
}

/// Iterator to work through the queue safely
pub struct BlockQueueIterator<'a> {
    queue: &'a BlockQueue,
}

impl BlockQueue {
    /// Create a block queue for the image with dimensions `img`.
    /// Panics if the image is not evenly broken into blocks of dimension `dim`
    pub fn new(img: (u32, u32), dim: (u32, u32)) -> BlockQueue {
        if img.0 % dim.0 != 0 || img.1 % dim.1 != 0 {
            panic!("Image with dimension {} not evenly divided by blocks of {}", img, dim);
        }
        let num_blocks = (img.0 / dim.0, img.1 / dim.1);
        let mut blocks: Vec<(u32, u32)> = range(0, num_blocks.0 * num_blocks.1)
            .map(|i| (i % num_blocks.0, i / num_blocks.0)).collect();
        blocks.sort_by(|a, b| morton::morton2(a).cmp(&morton::morton2(b)));
        BlockQueue { blocks: blocks, dimensions: dim, next: AtomicUint::new(0) }
    }
    /// Get the dimensions of an individual block in the queue
    pub fn block_dim(&self) -> (u32, u32) { self.dimensions }
    /// Get an iterator to work through the queue
    pub fn iter(&self) -> BlockQueueIterator { BlockQueueIterator { queue: self } }
    /// Get the next block in the queue or None if the queue is finished
    fn next(&self) -> Option<(u32, u32)> {
        let i = self.next.fetch_add(1, AcqRel);
        if i >= self.blocks.len() {
            None
        } else {
            Some(self.blocks[i])
        }
    }
    /// Get the length of the queue
    pub fn len(&self) -> uint { self.blocks.len() }
}

impl<'a> Iterator<(u32, u32)> for BlockQueueIterator<'a> {
    fn next(&mut self) -> Option<(u32, u32)> {
        self.queue.next()
    }
}

<|endoftext|>


<commit_before>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        println!("{}", pid);
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32::default();
        process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }
}
<commit_msg>fix: Remove println from windows code<commit_after>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32::default();
        process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }
}
<|endoftext|>