<commit_before>use rustc_hash::FxHashSet;

use syntax::{
    ast::{self, AstNode, AstToken, VisibilityOwner},
    Direction, NodeOrToken, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, TextRange, TextSize,
};

#[derive(Debug, PartialEq, Eq)]
pub enum FoldKind {
    Comment,
    Imports,
    Mods,
    Block,
    ArgList,
    Region,
    Consts,
    Statics,
    Array,
    WhereClause,
    ReturnType,
}

#[derive(Debug)]
pub struct Fold {
    pub range: TextRange,
    pub kind: FoldKind,
}

// Feature: Folding
//
// Defines folding regions for curly braced blocks, runs of consecutive import
// statements, and `region` / `endregion` comment markers.
pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
    let mut res = vec![];
    let mut visited_comments = FxHashSet::default();
    let mut visited_imports = FxHashSet::default();
    let mut visited_mods = FxHashSet::default();
    let mut visited_consts = FxHashSet::default();
    let mut visited_statics = FxHashSet::default();
    // regions can be nested, here is a LIFO buffer
    let mut regions_starts: Vec<TextSize> = vec![];

    for element in file.syntax().descendants_with_tokens() {
        // Fold items that span multiple lines
        if let Some(kind) = fold_kind(element.kind()) {
            let is_multiline = match &element {
                NodeOrToken::Node(node) => node.text().contains_char('\n'),
                NodeOrToken::Token(token) => token.text().contains('\n'),
            };
            if is_multiline {
                res.push(Fold { range: element.text_range(), kind });
                continue;
            }
        }

        match element {
            NodeOrToken::Token(token) => {
                // Fold groups of comments
                if let Some(comment) = ast::Comment::cast(token) {
                    if !visited_comments.contains(&comment) {
                        // regions are not real comments
                        if comment.text().trim().starts_with("// region:") {
                            regions_starts.push(comment.syntax().text_range().start());
                        } else if comment.text().trim().starts_with("// endregion") {
                            if let Some(region) = regions_starts.pop() {
                                res.push(Fold {
                                    range: TextRange::new(
                                        region,
                                        comment.syntax().text_range().end(),
                                    ),
                                    kind: FoldKind::Region,
                                })
                            }
                        } else {
                            if let Some(range) =
                                contiguous_range_for_comment(comment, &mut visited_comments)
                            {
                                res.push(Fold { range, kind: FoldKind::Comment })
                            }
                        }
                    }
                }
            }
            NodeOrToken::Node(node) => {
                // Fold groups of imports
                if node.kind() == USE && !visited_imports.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_imports) {
                        res.push(Fold { range, kind: FoldKind::Imports })
                    }
                }

                // Fold groups of mods
                if node.kind() == MODULE && !has_visibility(&node) && !visited_mods.contains(&node)
                {
                    if let Some(range) =
                        contiguous_range_for_group_unless(&node, has_visibility, &mut visited_mods)
                    {
                        res.push(Fold { range, kind: FoldKind::Mods })
                    }
                }

                // Fold groups of consts
                if node.kind() == CONST && !visited_consts.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_consts) {
                        res.push(Fold { range, kind: FoldKind::Consts })
                    }
                }
                // Fold groups of consts
                if node.kind() == STATIC && !visited_statics.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_statics) {
                        res.push(Fold { range, kind: FoldKind::Statics })
                    }
                }

                // Fold where clause
                if node.kind() == WHERE_CLAUSE {
                    if let Some(range) = fold_range_for_where_clause(&node) {
                        res.push(Fold { range, kind: FoldKind::WhereClause })
                    }
                }
            }
        }
    }

    res
}

fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
    match kind {
        COMMENT => Some(FoldKind::Comment),
        ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
        ARRAY_EXPR => Some(FoldKind::Array),
        RET_TYPE => Some(FoldKind::ReturnType),
        ASSOC_ITEM_LIST
        | RECORD_FIELD_LIST
        | RECORD_PAT_FIELD_LIST
        | RECORD_EXPR_FIELD_LIST
        | ITEM_LIST
        | EXTERN_ITEM_LIST
        | USE_TREE_LIST
        | BLOCK_EXPR
        | MATCH_ARM_LIST
        | VARIANT_LIST
        | TOKEN_TREE => Some(FoldKind::Block),
        _ => None,
    }
}

fn has_visibility(node: &SyntaxNode) -> bool {
    ast::Module::cast(node.clone()).and_then(|m| m.visibility()).is_some()
}

fn contiguous_range_for_group(
    first: &SyntaxNode,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    contiguous_range_for_group_unless(first, |_| false, visited)
}

fn contiguous_range_for_group_unless(
    first: &SyntaxNode,
    unless: impl Fn(&SyntaxNode) -> bool,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    let mut last = first.clone();
    for element in first.siblings_with_tokens(Direction::Next) {
        let node = match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                // There is a blank line or another token, which means that the
                // group ends here
                break;
            }
            NodeOrToken::Node(node) => node,
        };

        // Stop if we find a node that doesn't belong to the group
        if node.kind() != first.kind() || unless(&node) {
            break;
        }

        visited.insert(node.clone());
        last = node;
    }

    if first != &last {
        Some(TextRange::new(first.text_range().start(), last.text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn contiguous_range_for_comment(
    first: ast::Comment,
    visited: &mut FxHashSet<ast::Comment>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    // Only fold comments of the same flavor
    let group_kind = first.kind();
    if !group_kind.shape.is_line() {
        return None;
    }

    let mut last = first.clone();
    for element in first.syntax().siblings_with_tokens(Direction::Next) {
        match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token.clone()) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                if let Some(c) = ast::Comment::cast(token) {
                    if c.kind() == group_kind {
                        // regions are not real comments
                        if c.text().trim().starts_with("// region:")
                            || c.text().trim().starts_with("// endregion")
                        {
                            break;
                        } else {
                            visited.insert(c.clone());
                            last = c;
                            continue;
                        }
                    }
                }
                // The comment group ends because either:
                // * An element of a different kind was reached
                // * A comment of a different flavor was reached
                break;
            }
            NodeOrToken::Node(_) => break,
        };
    }

    if first != last {
        Some(TextRange::new(first.syntax().text_range().start(), last.syntax().text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn fold_range_for_where_clause(node: &SyntaxNode) -> Option<TextRange> {
    let first_where_pred = node.first_child();
    let last_where_pred = node.last_child();

    if first_where_pred != last_where_pred {
        let mut it = node.descendants_with_tokens();
        if let (Some(_where_clause), Some(where_kw), Some(last_comma)) =
            (it.next(), it.next(), it.last())
        {
            let start = where_kw.text_range().end();
            let end = last_comma.text_range().end();
            return Some(TextRange::new(start, end));
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use test_utils::extract_tags;

    use super::*;

    fn check(ra_fixture: &str) {
        let (ranges, text) = extract_tags(ra_fixture, "fold");

        let parse = SourceFile::parse(&text);
        let folds = folding_ranges(&parse.tree());
        assert_eq!(
            folds.len(),
            ranges.len(),
            "The amount of folds is different than the expected amount"
        );

        for (fold, (range, attr)) in folds.iter().zip(ranges.into_iter()) {
            assert_eq!(fold.range.start(), range.start());
            assert_eq!(fold.range.end(), range.end());

            let kind = match fold.kind {
                FoldKind::Comment => "comment",
                FoldKind::Imports => "imports",
                FoldKind::Mods => "mods",
                FoldKind::Block => "block",
                FoldKind::ArgList => "arglist",
                FoldKind::Region => "region",
                FoldKind::Consts => "consts",
                FoldKind::Statics => "statics",
                FoldKind::Array => "array",
                FoldKind::WhereClause => "whereclause",
                FoldKind::ReturnType => "returntype",
            };
            assert_eq!(kind, &attr.unwrap());
        }
    }

    #[test]
    fn test_fold_comments() {
        check(
            r#"
<fold comment>// Hello
// this is a multiline
// comment
//</fold>

// But this is not

fn main() <fold block>{
    <fold comment>// We should
    // also
    // fold
    // this one.</fold>
    <fold comment>//! But this one is different
    //! because it has another flavor</fold>
    <fold comment>/* As does this
    multiline comment */</fold>
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_imports() {
        check(
            r#"
use std::<fold block>{
    str,
    vec,
    io as iop
}</fold>;

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_mods() {
        check(
            r#"

pub mod foo;
<fold mods>mod after_pub;
mod after_pub_next;</fold>

<fold mods>mod before_pub;
mod before_pub_next;</fold>
pub mod bar;

mod not_folding_single;
pub mod foobar;
pub not_folding_single_next;

<fold mods>#[cfg(test)]
mod with_attribute;
mod with_attribute_next;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

<fold imports>use std::collections::HashMap;
// Some random comment
use std::collections::VecDeque;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_and_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

use std::collections::<fold block>{
    HashMap,
    VecDeque,
}</fold>;
// Some random comment

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_folds_structs() {
        check(
            r#"
struct Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_traits() {
        check(
            r#"
trait Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_macros() {
        check(
            r#"
macro_rules! foo <fold block>{
    ($($tt:tt)*) => { $($tt)* }
}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_match_arms() {
        check(
            r#"
fn main() <fold block>{
    match 0 <fold block>{
        0 => 0,
        _ => 1,
    }</fold>
}</fold>
"#,
        );
    }

    #[test]
    fn fold_big_calls() {
        check(
            r#"
fn main() <fold block>{
    frobnicate<fold arglist>(
        1,
        2,
        3,
    )</fold>
}</fold>
"#,
        )
    }

    #[test]
    fn fold_record_literals() {
        check(
            r#"
const _: S = S <fold block>{

}</fold>;
"#,
        )
    }

    #[test]
    fn fold_multiline_params() {
        check(
            r#"
fn foo<fold arglist>(
    x: i32,
    y: String,
)</fold> {}
"#,
        )
    }

    #[test]
    fn fold_multiline_array() {
        check(
            r#"
const FOO: [usize; 4] = <fold array>[
    1,
    2,
    3,
    4,
]</fold>;
"#,
        )
    }

    #[test]
    fn fold_region() {
        check(
            r#"
// 1. some normal comment
<fold region>// region: test
// 2. some normal comment
calling_function(x,y);
// endregion: test</fold>
"#,
        )
    }

    #[test]
    fn fold_consecutive_const() {
        check(
            r#"
<fold consts>const FIRST_CONST: &str = "first";
const SECOND_CONST: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_consecutive_static() {
        check(
            r#"
<fold statics>static FIRST_STATIC: &str = "first";
static SECOND_STATIC: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_where_clause() {
        // fold multi-line and don't fold single line.
        check(
            r#"
fn foo()
where<fold whereclause>
    A: Foo,
    B: Foo,
    C: Foo,
    D: Foo,</fold> {}

fn bar()
where
    A: Bar, {}
"#,
        )
    }

    #[test]
    fn fold_return_type() {
        check(
            r#"
fn foo()<fold returntype>-> (
    bool,
    bool,
)</fold> { (true, true) }
            "#,
        )
    }
}
<commit_msg>Test for correct amount of emitted folds<commit_after>use rustc_hash::FxHashSet;

use syntax::{
    ast::{self, AstNode, AstToken, VisibilityOwner},
    Direction, NodeOrToken, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, TextRange, TextSize,
};

#[derive(Debug, PartialEq, Eq)]
pub enum FoldKind {
    Comment,
    Imports,
    Mods,
    Block,
    ArgList,
    Region,
    Consts,
    Statics,
    Array,
    WhereClause,
    ReturnType,
}

#[derive(Debug)]
pub struct Fold {
    pub range: TextRange,
    pub kind: FoldKind,
}

// Feature: Folding
//
// Defines folding regions for curly braced blocks, runs of consecutive import
// statements, and `region` / `endregion` comment markers.
pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
    let mut res = vec![];
    let mut visited_comments = FxHashSet::default();
    let mut visited_imports = FxHashSet::default();
    let mut visited_mods = FxHashSet::default();
    let mut visited_consts = FxHashSet::default();
    let mut visited_statics = FxHashSet::default();
    // regions can be nested, here is a LIFO buffer
    let mut regions_starts: Vec<TextSize> = vec![];

    for element in file.syntax().descendants_with_tokens() {
        // Fold items that span multiple lines
        if let Some(kind) = fold_kind(element.kind()) {
            let is_multiline = match &element {
                NodeOrToken::Node(node) => node.text().contains_char('\n'),
                NodeOrToken::Token(token) => token.text().contains('\n'),
            };
            if is_multiline {
                res.push(Fold { range: element.text_range(), kind });
                continue;
            }
        }

        match element {
            NodeOrToken::Token(token) => {
                // Fold groups of comments
                if let Some(comment) = ast::Comment::cast(token) {
                    if !visited_comments.contains(&comment) {
                        // regions are not real comments
                        if comment.text().trim().starts_with("// region:") {
                            regions_starts.push(comment.syntax().text_range().start());
                        } else if comment.text().trim().starts_with("// endregion") {
                            if let Some(region) = regions_starts.pop() {
                                res.push(Fold {
                                    range: TextRange::new(
                                        region,
                                        comment.syntax().text_range().end(),
                                    ),
                                    kind: FoldKind::Region,
                                })
                            }
                        } else {
                            if let Some(range) =
                                contiguous_range_for_comment(comment, &mut visited_comments)
                            {
                                res.push(Fold { range, kind: FoldKind::Comment })
                            }
                        }
                    }
                }
            }
            NodeOrToken::Node(node) => {
                // Fold groups of imports
                if node.kind() == USE && !visited_imports.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_imports) {
                        res.push(Fold { range, kind: FoldKind::Imports })
                    }
                }

                // Fold groups of mods
                if node.kind() == MODULE && !has_visibility(&node) && !visited_mods.contains(&node)
                {
                    if let Some(range) =
                        contiguous_range_for_group_unless(&node, has_visibility, &mut visited_mods)
                    {
                        res.push(Fold { range, kind: FoldKind::Mods })
                    }
                }

                // Fold groups of consts
                if node.kind() == CONST && !visited_consts.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_consts) {
                        res.push(Fold { range, kind: FoldKind::Consts })
                    }
                }
                // Fold groups of consts
                if node.kind() == STATIC && !visited_statics.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_statics) {
                        res.push(Fold { range, kind: FoldKind::Statics })
                    }
                }

                // Fold where clause
                if node.kind() == WHERE_CLAUSE {
                    if let Some(range) = fold_range_for_where_clause(&node) {
                        res.push(Fold { range, kind: FoldKind::WhereClause })
                    }
                }
            }
        }
    }

    res
}

fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
    match kind {
        COMMENT => Some(FoldKind::Comment),
        ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
        ARRAY_EXPR => Some(FoldKind::Array),
        RET_TYPE => Some(FoldKind::ReturnType),
        ASSOC_ITEM_LIST
        | RECORD_FIELD_LIST
        | RECORD_PAT_FIELD_LIST
        | RECORD_EXPR_FIELD_LIST
        | ITEM_LIST
        | EXTERN_ITEM_LIST
        | USE_TREE_LIST
        | BLOCK_EXPR
        | MATCH_ARM_LIST
        | VARIANT_LIST
        | TOKEN_TREE => Some(FoldKind::Block),
        _ => None,
    }
}

fn has_visibility(node: &SyntaxNode) -> bool {
    ast::Module::cast(node.clone()).and_then(|m| m.visibility()).is_some()
}

fn contiguous_range_for_group(
    first: &SyntaxNode,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    contiguous_range_for_group_unless(first, |_| false, visited)
}

fn contiguous_range_for_group_unless(
    first: &SyntaxNode,
    unless: impl Fn(&SyntaxNode) -> bool,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    let mut last = first.clone();
    for element in first.siblings_with_tokens(Direction::Next) {
        let node = match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                // There is a blank line or another token, which means that the
                // group ends here
                break;
            }
            NodeOrToken::Node(node) => node,
        };

        // Stop if we find a node that doesn't belong to the group
        if node.kind() != first.kind() || unless(&node) {
            break;
        }

        visited.insert(node.clone());
        last = node;
    }

    if first != &last {
        Some(TextRange::new(first.text_range().start(), last.text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn contiguous_range_for_comment(
    first: ast::Comment,
    visited: &mut FxHashSet<ast::Comment>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    // Only fold comments of the same flavor
    let group_kind = first.kind();
    if !group_kind.shape.is_line() {
        return None;
    }

    let mut last = first.clone();
    for element in first.syntax().siblings_with_tokens(Direction::Next) {
        match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token.clone()) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                if let Some(c) = ast::Comment::cast(token) {
                    if c.kind() == group_kind {
                        // regions are not real comments
                        if c.text().trim().starts_with("// region:")
                            || c.text().trim().starts_with("// endregion")
                        {
                            break;
                        } else {
                            visited.insert(c.clone());
                            last = c;
                            continue;
                        }
                    }
                }
                // The comment group ends because either:
                // * An element of a different kind was reached
                // * A comment of a different flavor was reached
                break;
            }
            NodeOrToken::Node(_) => break,
        };
    }

    if first != last {
        Some(TextRange::new(first.syntax().text_range().start(), last.syntax().text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn fold_range_for_where_clause(node: &SyntaxNode) -> Option<TextRange> {
    let first_where_pred = node.first_child();
    let last_where_pred = node.last_child();

    if first_where_pred != last_where_pred {
        let mut it = node.descendants_with_tokens();
        if let (Some(_where_clause), Some(where_kw), Some(last_comma)) =
            (it.next(), it.next(), it.last())
        {
            let start = where_kw.text_range().end();
            let end = last_comma.text_range().end();
            return Some(TextRange::new(start, end));
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use test_utils::extract_tags;

    use super::*;

    fn check(ra_fixture: &str) {
        let (ranges, text) = extract_tags(ra_fixture, "fold");

        let parse = SourceFile::parse(&text);
        let folds = folding_ranges(&parse.tree());
        assert_eq!(
            folds.len(),
            ranges.len(),
            "The amount of folds is different than the expected amount"
        );

        for (fold, (range, attr)) in folds.iter().zip(ranges.into_iter()) {
            assert_eq!(fold.range.start(), range.start());
            assert_eq!(fold.range.end(), range.end());

            let kind = match fold.kind {
                FoldKind::Comment => "comment",
                FoldKind::Imports => "imports",
                FoldKind::Mods => "mods",
                FoldKind::Block => "block",
                FoldKind::ArgList => "arglist",
                FoldKind::Region => "region",
                FoldKind::Consts => "consts",
                FoldKind::Statics => "statics",
                FoldKind::Array => "array",
                FoldKind::WhereClause => "whereclause",
                FoldKind::ReturnType => "returntype",
            };
            assert_eq!(kind, &attr.unwrap());
        }
    }

    #[test]
    fn test_fold_comments() {
        check(
            r#"
<fold comment>// Hello
// this is a multiline
// comment
//</fold>

// But this is not

fn main() <fold block>{
    <fold comment>// We should
    // also
    // fold
    // this one.</fold>
    <fold comment>//! But this one is different
    //! because it has another flavor</fold>
    <fold comment>/* As does this
    multiline comment */</fold>
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_imports() {
        check(
            r#"
use std::<fold block>{
    str,
    vec,
    io as iop
}</fold>;

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_mods() {
        check(
            r#"

pub mod foo;
<fold mods>mod after_pub;
mod after_pub_next;</fold>

<fold mods>mod before_pub;
mod before_pub_next;</fold>
pub mod bar;

mod not_folding_single;
pub mod foobar;
pub not_folding_single_next;

<fold mods>#[cfg(test)]
mod with_attribute;
mod with_attribute_next;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

<fold imports>use std::collections::HashMap;
// Some random comment
use std::collections::VecDeque;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_and_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

use std::collections::<fold block>{
    HashMap,
    VecDeque,
}</fold>;
// Some random comment

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_folds_structs() {
        check(
            r#"
struct Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_traits() {
        check(
            r#"
trait Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_macros() {
        check(
            r#"
macro_rules! foo <fold block>{
    ($($tt:tt)*) => { $($tt)* }
}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_match_arms() {
        check(
            r#"
fn main() <fold block>{
    match 0 <fold block>{
        0 => 0,
        _ => 1,
    }</fold>
}</fold>
"#,
        );
    }

    #[test]
    fn fold_big_calls() {
        check(
            r#"
fn main() <fold block>{
    frobnicate<fold arglist>(
        1,
        2,
        3,
    )</fold>
}</fold>
"#,
        )
    }

    #[test]
    fn fold_record_literals() {
        check(
            r#"
const _: S = S <fold block>{

}</fold>;
"#,
        )
    }

    #[test]
    fn fold_multiline_params() {
        check(
            r#"
fn foo<fold arglist>(
    x: i32,
    y: String,
)</fold> {}
"#,
        )
    }

    #[test]
    fn fold_multiline_array() {
        check(
            r#"
const FOO: [usize; 4] = <fold array>[
    1,
    2,
    3,
    4,
]</fold>;
"#,
        )
    }

    #[test]
    fn fold_region() {
        check(
            r#"
// 1. some normal comment
<fold region>// region: test
// 2. some normal comment
calling_function(x,y);
// endregion: test</fold>
"#,
        )
    }

    #[test]
    fn fold_consecutive_const() {
        check(
            r#"
<fold consts>const FIRST_CONST: &str = "first";
const SECOND_CONST: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_consecutive_static() {
        check(
            r#"
<fold statics>static FIRST_STATIC: &str = "first";
static SECOND_STATIC: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_where_clause() {
        // fold multi-line and don't fold single line.
        check(
            r#"
fn foo()
where<fold whereclause>
    A: Foo,
    B: Foo,
    C: Foo,
    D: Foo,</fold> {}

fn bar()
where
    A: Bar, {}
"#,
        )
    }

    #[test]
    fn fold_return_type() {
        check(
            r#"
fn foo()<fold returntype>-> (
    bool,
    bool,
)</fold> { (true, true) }

fn bar() -> (bool, bool) { (true, true) }
            "#,
        )
    }
}
<|endoftext|>


<commit_before>use rustc_hash::FxHashSet;

use syntax::{
    ast::{self, AstNode, AstToken, VisibilityOwner},
    Direction, NodeOrToken, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, TextRange, TextSize,
};

#[derive(Debug, PartialEq, Eq)]
pub enum FoldKind {
    Comment,
    Imports,
    Mods,
    Block,
    ArgList,
    Region,
    Consts,
    Statics,
    Array,
    WhereClause,
    ReturnType,
}

#[derive(Debug)]
pub struct Fold {
    pub range: TextRange,
    pub kind: FoldKind,
}

// Feature: Folding
//
// Defines folding regions for curly braced blocks, runs of consecutive import
// statements, and `region` / `endregion` comment markers.
pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
    let mut res = vec![];
    let mut visited_comments = FxHashSet::default();
    let mut visited_imports = FxHashSet::default();
    let mut visited_mods = FxHashSet::default();
    let mut visited_consts = FxHashSet::default();
    let mut visited_statics = FxHashSet::default();
    // regions can be nested, here is a LIFO buffer
    let mut regions_starts: Vec<TextSize> = vec![];

    for element in file.syntax().descendants_with_tokens() {
        // Fold items that span multiple lines
        if let Some(kind) = fold_kind(element.kind()) {
            let is_multiline = match &element {
                NodeOrToken::Node(node) => node.text().contains_char('\n'),
                NodeOrToken::Token(token) => token.text().contains('\n'),
            };
            if is_multiline {
                res.push(Fold { range: element.text_range(), kind });
                continue;
            }
        }

        match element {
            NodeOrToken::Token(token) => {
                // Fold groups of comments
                if let Some(comment) = ast::Comment::cast(token) {
                    if !visited_comments.contains(&comment) {
                        // regions are not real comments
                        if comment.text().trim().starts_with("// region:") {
                            regions_starts.push(comment.syntax().text_range().start());
                        } else if comment.text().trim().starts_with("// endregion") {
                            if let Some(region) = regions_starts.pop() {
                                res.push(Fold {
                                    range: TextRange::new(
                                        region,
                                        comment.syntax().text_range().end(),
                                    ),
                                    kind: FoldKind::Region,
                                })
                            }
                        } else {
                            if let Some(range) =
                                contiguous_range_for_comment(comment, &mut visited_comments)
                            {
                                res.push(Fold { range, kind: FoldKind::Comment })
                            }
                        }
                    }
                }
            }
            NodeOrToken::Node(node) => {
                // Fold groups of imports
                if node.kind() == USE && !visited_imports.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_imports) {
                        res.push(Fold { range, kind: FoldKind::Imports })
                    }
                }

                // Fold groups of mods
                if node.kind() == MODULE && !has_visibility(&node) && !visited_mods.contains(&node)
                {
                    if let Some(range) =
                        contiguous_range_for_group_unless(&node, has_visibility, &mut visited_mods)
                    {
                        res.push(Fold { range, kind: FoldKind::Mods })
                    }
                }

                // Fold groups of consts
                if node.kind() == CONST && !visited_consts.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_consts) {
                        res.push(Fold { range, kind: FoldKind::Consts })
                    }
                }
                // Fold groups of consts
                if node.kind() == STATIC && !visited_statics.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_statics) {
                        res.push(Fold { range, kind: FoldKind::Statics })
                    }
                }

                // Fold where clause
                if node.kind() == WHERE_CLAUSE {
                    if let Some(range) = fold_range_for_where_clause(&node) {
                        res.push(Fold { range, kind: FoldKind::WhereClause })
                    }
                }
            }
        }
    }

    res
}

fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
    match kind {
        COMMENT => Some(FoldKind::Comment),
        ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
        ARRAY_EXPR => Some(FoldKind::Array),
        RET_TYPE => Some(FoldKind::ReturnType),
        ASSOC_ITEM_LIST
        | RECORD_FIELD_LIST
        | RECORD_PAT_FIELD_LIST
        | RECORD_EXPR_FIELD_LIST
        | ITEM_LIST
        | EXTERN_ITEM_LIST
        | USE_TREE_LIST
        | BLOCK_EXPR
        | MATCH_ARM_LIST
        | VARIANT_LIST
        | TOKEN_TREE => Some(FoldKind::Block),
        _ => None,
    }
}

fn has_visibility(node: &SyntaxNode) -> bool {
    ast::Module::cast(node.clone()).and_then(|m| m.visibility()).is_some()
}

fn contiguous_range_for_group(
    first: &SyntaxNode,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    contiguous_range_for_group_unless(first, |_| false, visited)
}

fn contiguous_range_for_group_unless(
    first: &SyntaxNode,
    unless: impl Fn(&SyntaxNode) -> bool,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    let mut last = first.clone();
    for element in first.siblings_with_tokens(Direction::Next) {
        let node = match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                // There is a blank line or another token, which means that the
                // group ends here
                break;
            }
            NodeOrToken::Node(node) => node,
        };

        // Stop if we find a node that doesn't belong to the group
        if node.kind() != first.kind() || unless(&node) {
            break;
        }

        visited.insert(node.clone());
        last = node;
    }

    if first != &last {
        Some(TextRange::new(first.text_range().start(), last.text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn contiguous_range_for_comment(
    first: ast::Comment,
    visited: &mut FxHashSet<ast::Comment>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    // Only fold comments of the same flavor
    let group_kind = first.kind();
    if !group_kind.shape.is_line() {
        return None;
    }

    let mut last = first.clone();
    for element in first.syntax().siblings_with_tokens(Direction::Next) {
        match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token.clone()) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                if let Some(c) = ast::Comment::cast(token) {
                    if c.kind() == group_kind {
                        // regions are not real comments
                        if c.text().trim().starts_with("// region:")
                            || c.text().trim().starts_with("// endregion")
                        {
                            break;
                        } else {
                            visited.insert(c.clone());
                            last = c;
                            continue;
                        }
                    }
                }
                // The comment group ends because either:
                // * An element of a different kind was reached
                // * A comment of a different flavor was reached
                break;
            }
            NodeOrToken::Node(_) => break,
        };
    }

    if first != last {
        Some(TextRange::new(first.syntax().text_range().start(), last.syntax().text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn fold_range_for_where_clause(node: &SyntaxNode) -> Option<TextRange> {
    let first_where_pred = node.first_child();
    let last_where_pred = node.last_child();

    if first_where_pred != last_where_pred {
        let mut it = node.descendants_with_tokens();
        if let (Some(_where_clause), Some(where_kw), Some(last_comma)) =
            (it.next(), it.next(), it.last())
        {
            let start = where_kw.text_range().end();
            let end = last_comma.text_range().end();
            return Some(TextRange::new(start, end));
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use test_utils::extract_tags;

    use super::*;

    fn check(ra_fixture: &str) {
        let (ranges, text) = extract_tags(ra_fixture, "fold");

        let parse = SourceFile::parse(&text);
        let folds = folding_ranges(&parse.tree());
        assert_eq!(
            folds.len(),
            ranges.len(),
            "The amount of folds is different than the expected amount"
        );

        for (fold, (range, attr)) in folds.iter().zip(ranges.into_iter()) {
            assert_eq!(fold.range.start(), range.start());
            assert_eq!(fold.range.end(), range.end());

            let kind = match fold.kind {
                FoldKind::Comment => "comment",
                FoldKind::Imports => "imports",
                FoldKind::Mods => "mods",
                FoldKind::Block => "block",
                FoldKind::ArgList => "arglist",
                FoldKind::Region => "region",
                FoldKind::Consts => "consts",
                FoldKind::Statics => "statics",
                FoldKind::Array => "array",
                FoldKind::WhereClause => "whereclause",
                FoldKind::ReturnType => "returntype",
            };
            assert_eq!(kind, &attr.unwrap());
        }
    }

    #[test]
    fn test_fold_comments() {
        check(
            r#"
<fold comment>// Hello
// this is a multiline
// comment
//</fold>

// But this is not

fn main() <fold block>{
    <fold comment>// We should
    // also
    // fold
    // this one.</fold>
    <fold comment>//! But this one is different
    //! because it has another flavor</fold>
    <fold comment>/* As does this
    multiline comment */</fold>
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_imports() {
        check(
            r#"
use std::<fold block>{
    str,
    vec,
    io as iop
}</fold>;

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_mods() {
        check(
            r#"

pub mod foo;
<fold mods>mod after_pub;
mod after_pub_next;</fold>

<fold mods>mod before_pub;
mod before_pub_next;</fold>
pub mod bar;

mod not_folding_single;
pub mod foobar;
pub not_folding_single_next;

<fold mods>#[cfg(test)]
mod with_attribute;
mod with_attribute_next;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

<fold imports>use std::collections::HashMap;
// Some random comment
use std::collections::VecDeque;</fold>

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_fold_import_and_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

use std::collections::<fold block>{
    HashMap,
    VecDeque,
}</fold>;
// Some random comment

fn main() <fold block>{
}</fold>"#,
        );
    }

    #[test]
    fn test_folds_structs() {
        check(
            r#"
struct Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_traits() {
        check(
            r#"
trait Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_macros() {
        check(
            r#"
macro_rules! foo <fold block>{
    ($($tt:tt)*) => { $($tt)* }
}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_match_arms() {
        check(
            r#"
fn main() <fold block>{
    match 0 <fold block>{
        0 => 0,
        _ => 1,
    }</fold>
}</fold>
"#,
        );
    }

    #[test]
    fn fold_big_calls() {
        check(
            r#"
fn main() <fold block>{
    frobnicate<fold arglist>(
        1,
        2,
        3,
    )</fold>
}</fold>
"#,
        )
    }

    #[test]
    fn fold_record_literals() {
        check(
            r#"
const _: S = S <fold block>{

}</fold>;
"#,
        )
    }

    #[test]
    fn fold_multiline_params() {
        check(
            r#"
fn foo<fold arglist>(
    x: i32,
    y: String,
)</fold> {}
"#,
        )
    }

    #[test]
    fn fold_multiline_array() {
        check(
            r#"
const FOO: [usize; 4] = <fold array>[
    1,
    2,
    3,
    4,
]</fold>;
"#,
        )
    }

    #[test]
    fn fold_region() {
        check(
            r#"
// 1. some normal comment
<fold region>// region: test
// 2. some normal comment
calling_function(x,y);
// endregion: test</fold>
"#,
        )
    }

    #[test]
    fn fold_consecutive_const() {
        check(
            r#"
<fold consts>const FIRST_CONST: &str = "first";
const SECOND_CONST: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_consecutive_static() {
        check(
            r#"
<fold statics>static FIRST_STATIC: &str = "first";
static SECOND_STATIC: &str = "second";</fold>
            "#,
        )
    }

    #[test]
    fn fold_where_clause() {
        // fold multi-line and don't fold single line.
        check(
            r#"
fn foo()
where<fold whereclause>
    A: Foo,
    B: Foo,
    C: Foo,
    D: Foo,</fold> {}

fn bar()
where
    A: Bar, {}
"#,
        )
    }

    #[test]
    fn fold_return_type() {
        check(
            r#"
fn foo()<fold returntype>-> (
    bool,
    bool,
)</fold> { (true, true) }

fn bar() -> (bool, bool) { (true, true) }
            "#,
        )
    }
}
<commit_msg>Don't fold consecutive modules with item lists<commit_after>use rustc_hash::FxHashSet;

use syntax::{
    ast::{self, AstNode, AstToken, VisibilityOwner},
    Direction, NodeOrToken, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, TextRange, TextSize,
};

#[derive(Debug, PartialEq, Eq)]
pub enum FoldKind {
    Comment,
    Imports,
    Mods,
    Block,
    ArgList,
    Region,
    Consts,
    Statics,
    Array,
    WhereClause,
    ReturnType,
}

#[derive(Debug)]
pub struct Fold {
    pub range: TextRange,
    pub kind: FoldKind,
}

// Feature: Folding
//
// Defines folding regions for curly braced blocks, runs of consecutive import
// statements, and `region` / `endregion` comment markers.
pub(crate) fn folding_ranges(file: &SourceFile) -> Vec<Fold> {
    let mut res = vec![];
    let mut visited_comments = FxHashSet::default();
    let mut visited_imports = FxHashSet::default();
    let mut visited_mods = FxHashSet::default();
    let mut visited_consts = FxHashSet::default();
    let mut visited_statics = FxHashSet::default();
    // regions can be nested, here is a LIFO buffer
    let mut regions_starts: Vec<TextSize> = vec![];

    for element in file.syntax().descendants_with_tokens() {
        // Fold items that span multiple lines
        if let Some(kind) = fold_kind(element.kind()) {
            let is_multiline = match &element {
                NodeOrToken::Node(node) => node.text().contains_char('\n'),
                NodeOrToken::Token(token) => token.text().contains('\n'),
            };
            if is_multiline {
                res.push(Fold { range: element.text_range(), kind });
                continue;
            }
        }

        match element {
            NodeOrToken::Token(token) => {
                // Fold groups of comments
                if let Some(comment) = ast::Comment::cast(token) {
                    if !visited_comments.contains(&comment) {
                        // regions are not real comments
                        if comment.text().trim().starts_with("// region:") {
                            regions_starts.push(comment.syntax().text_range().start());
                        } else if comment.text().trim().starts_with("// endregion") {
                            if let Some(region) = regions_starts.pop() {
                                res.push(Fold {
                                    range: TextRange::new(
                                        region,
                                        comment.syntax().text_range().end(),
                                    ),
                                    kind: FoldKind::Region,
                                })
                            }
                        } else {
                            if let Some(range) =
                                contiguous_range_for_comment(comment, &mut visited_comments)
                            {
                                res.push(Fold { range, kind: FoldKind::Comment })
                            }
                        }
                    }
                }
            }
            NodeOrToken::Node(node) => {
                // Fold groups of imports
                if node.kind() == USE && !visited_imports.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_imports) {
                        res.push(Fold { range, kind: FoldKind::Imports })
                    }
                }

                // Fold groups of mods
                if let Some(module) = ast::Module::cast(node.clone()) {
                    if !has_visibility(&node)
                        && !visited_mods.contains(&node)
                        && module.item_list().is_none()
                    {
                        if let Some(range) = contiguous_range_for_group_unless(
                            &node,
                            has_visibility,
                            &mut visited_mods,
                        ) {
                            res.push(Fold { range, kind: FoldKind::Mods })
                        }
                    }
                }

                // Fold groups of consts
                if node.kind() == CONST && !visited_consts.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_consts) {
                        res.push(Fold { range, kind: FoldKind::Consts })
                    }
                }
                // Fold groups of consts
                if node.kind() == STATIC && !visited_statics.contains(&node) {
                    if let Some(range) = contiguous_range_for_group(&node, &mut visited_statics) {
                        res.push(Fold { range, kind: FoldKind::Statics })
                    }
                }

                // Fold where clause
                if node.kind() == WHERE_CLAUSE {
                    if let Some(range) = fold_range_for_where_clause(&node) {
                        res.push(Fold { range, kind: FoldKind::WhereClause })
                    }
                }
            }
        }
    }

    res
}

fn fold_kind(kind: SyntaxKind) -> Option<FoldKind> {
    match kind {
        COMMENT => Some(FoldKind::Comment),
        ARG_LIST | PARAM_LIST => Some(FoldKind::ArgList),
        ARRAY_EXPR => Some(FoldKind::Array),
        RET_TYPE => Some(FoldKind::ReturnType),
        ASSOC_ITEM_LIST
        | RECORD_FIELD_LIST
        | RECORD_PAT_FIELD_LIST
        | RECORD_EXPR_FIELD_LIST
        | ITEM_LIST
        | EXTERN_ITEM_LIST
        | USE_TREE_LIST
        | BLOCK_EXPR
        | MATCH_ARM_LIST
        | VARIANT_LIST
        | TOKEN_TREE => Some(FoldKind::Block),
        _ => None,
    }
}

fn has_visibility(node: &SyntaxNode) -> bool {
    ast::Module::cast(node.clone()).and_then(|m| m.visibility()).is_some()
}

fn contiguous_range_for_group(
    first: &SyntaxNode,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    contiguous_range_for_group_unless(first, |_| false, visited)
}

fn contiguous_range_for_group_unless(
    first: &SyntaxNode,
    unless: impl Fn(&SyntaxNode) -> bool,
    visited: &mut FxHashSet<SyntaxNode>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    let mut last = first.clone();
    for element in first.siblings_with_tokens(Direction::Next) {
        let node = match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                // There is a blank line or another token, which means that the
                // group ends here
                break;
            }
            NodeOrToken::Node(node) => node,
        };

        // Stop if we find a node that doesn't belong to the group
        if node.kind() != first.kind() || unless(&node) {
            break;
        }

        visited.insert(node.clone());
        last = node;
    }

    if first != &last {
        Some(TextRange::new(first.text_range().start(), last.text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn contiguous_range_for_comment(
    first: ast::Comment,
    visited: &mut FxHashSet<ast::Comment>,
) -> Option<TextRange> {
    visited.insert(first.clone());

    // Only fold comments of the same flavor
    let group_kind = first.kind();
    if !group_kind.shape.is_line() {
        return None;
    }

    let mut last = first.clone();
    for element in first.syntax().siblings_with_tokens(Direction::Next) {
        match element {
            NodeOrToken::Token(token) => {
                if let Some(ws) = ast::Whitespace::cast(token.clone()) {
                    if !ws.spans_multiple_lines() {
                        // Ignore whitespace without blank lines
                        continue;
                    }
                }
                if let Some(c) = ast::Comment::cast(token) {
                    if c.kind() == group_kind {
                        // regions are not real comments
                        if c.text().trim().starts_with("// region:")
                            || c.text().trim().starts_with("// endregion")
                        {
                            break;
                        } else {
                            visited.insert(c.clone());
                            last = c;
                            continue;
                        }
                    }
                }
                // The comment group ends because either:
                // * An element of a different kind was reached
                // * A comment of a different flavor was reached
                break;
            }
            NodeOrToken::Node(_) => break,
        };
    }

    if first != last {
        Some(TextRange::new(first.syntax().text_range().start(), last.syntax().text_range().end()))
    } else {
        // The group consists of only one element, therefore it cannot be folded
        None
    }
}

fn fold_range_for_where_clause(node: &SyntaxNode) -> Option<TextRange> {
    let first_where_pred = node.first_child();
    let last_where_pred = node.last_child();

    if first_where_pred != last_where_pred {
        let mut it = node.descendants_with_tokens();
        if let (Some(_where_clause), Some(where_kw), Some(last_comma)) =
            (it.next(), it.next(), it.last())
        {
            let start = where_kw.text_range().end();
            let end = last_comma.text_range().end();
            return Some(TextRange::new(start, end));
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use test_utils::extract_tags;

    use super::*;

    fn check(ra_fixture: &str) {
        let (ranges, text) = extract_tags(ra_fixture, "fold");

        let parse = SourceFile::parse(&text);
        let folds = folding_ranges(&parse.tree());
        assert_eq!(
            folds.len(),
            ranges.len(),
            "The amount of folds is different than the expected amount"
        );

        for (fold, (range, attr)) in folds.iter().zip(ranges.into_iter()) {
            assert_eq!(fold.range.start(), range.start());
            assert_eq!(fold.range.end(), range.end());

            let kind = match fold.kind {
                FoldKind::Comment => "comment",
                FoldKind::Imports => "imports",
                FoldKind::Mods => "mods",
                FoldKind::Block => "block",
                FoldKind::ArgList => "arglist",
                FoldKind::Region => "region",
                FoldKind::Consts => "consts",
                FoldKind::Statics => "statics",
                FoldKind::Array => "array",
                FoldKind::WhereClause => "whereclause",
                FoldKind::ReturnType => "returntype",
            };
            assert_eq!(kind, &attr.unwrap());
        }
    }

    #[test]
    fn test_fold_comments() {
        check(
            r#"
<fold comment>// Hello
// this is a multiline
// comment
//</fold>

// But this is not

fn main() <fold block>{
    <fold comment>// We should
    // also
    // fold
    // this one.</fold>
    <fold comment>//! But this one is different
    //! because it has another flavor</fold>
    <fold comment>/* As does this
    multiline comment */</fold>
}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_imports() {
        check(
            r#"
use std::<fold block>{
    str,
    vec,
    io as iop
}</fold>;
"#,
        );
    }

    #[test]
    fn test_fold_mods() {
        check(
            r#"

pub mod foo;
<fold mods>mod after_pub;
mod after_pub_next;</fold>

<fold mods>mod before_pub;
mod before_pub_next;</fold>
pub mod bar;

mod not_folding_single;
pub mod foobar;
pub not_folding_single_next;

<fold mods>#[cfg(test)]
mod with_attribute;
mod with_attribute_next;</fold>

mod inline0 {}
mod inline1 {}

mod inline2 <fold block>{

}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_import_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

<fold imports>use std::collections::HashMap;
// Some random comment
use std::collections::VecDeque;</fold>
"#,
        );
    }

    #[test]
    fn test_fold_import_and_groups() {
        check(
            r#"
<fold imports>use std::str;
use std::vec;
use std::io as iop;</fold>

<fold imports>use std::mem;
use std::f64;</fold>

use std::collections::<fold block>{
    HashMap,
    VecDeque,
}</fold>;
// Some random comment
"#,
        );
    }

    #[test]
    fn test_folds_structs() {
        check(
            r#"
struct Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_traits() {
        check(
            r#"
trait Foo <fold block>{
}</fold>
"#,
        );
    }

    #[test]
    fn test_folds_macros() {
        check(
            r#"
macro_rules! foo <fold block>{
    ($($tt:tt)*) => { $($tt)* }
}</fold>
"#,
        );
    }

    #[test]
    fn test_fold_match_arms() {
        check(
            r#"
fn main() <fold block>{
    match 0 <fold block>{
        0 => 0,
        _ => 1,
    }</fold>
}</fold>
"#,
        );
    }

    #[test]
    fn fold_big_calls() {
        check(
            r#"
fn main() <fold block>{
    frobnicate<fold arglist>(
        1,
        2,
        3,
    )</fold>
}</fold>
"#,
        )
    }

    #[test]
    fn fold_record_literals() {
        check(
            r#"
const _: S = S <fold block>{

}</fold>;
"#,
        )
    }

    #[test]
    fn fold_multiline_params() {
        check(
            r#"
fn foo<fold arglist>(
    x: i32,
    y: String,
)</fold> {}
"#,
        )
    }

    #[test]
    fn fold_multiline_array() {
        check(
            r#"
const FOO: [usize; 4] = <fold array>[
    1,
    2,
    3,
    4,
]</fold>;
"#,
        )
    }

    #[test]
    fn fold_region() {
        check(
            r#"
// 1. some normal comment
<fold region>// region: test
// 2. some normal comment
calling_function(x,y);
// endregion: test</fold>
"#,
        )
    }

    #[test]
    fn fold_consecutive_const() {
        check(
            r#"
<fold consts>const FIRST_CONST: &str = "first";
const SECOND_CONST: &str = "second";</fold>
"#,
        )
    }

    #[test]
    fn fold_consecutive_static() {
        check(
            r#"
<fold statics>static FIRST_STATIC: &str = "first";
static SECOND_STATIC: &str = "second";</fold>
"#,
        )
    }

    #[test]
    fn fold_where_clause() {
        // fold multi-line and don't fold single line.
        check(
            r#"
fn foo()
where<fold whereclause>
    A: Foo,
    B: Foo,
    C: Foo,
    D: Foo,</fold> {}

fn bar()
where
    A: Bar, {}
"#,
        )
    }

    #[test]
    fn fold_return_type() {
        check(
            r#"
fn foo()<fold returntype>-> (
    bool,
    bool,
)</fold> { (true, true) }

fn bar() -> (bool, bool) { (true, true) }
"#,
        )
    }
}
<|endoftext|>


<commit_before>//! Advertizes the capabilities of the LSP Server.
use std::env;

use lsp_types::{
    CallHierarchyServerCapability, ClientCapabilities, CodeActionKind, CodeActionOptions,
    CodeActionProviderCapability, CodeLensOptions, CompletionOptions,
    DocumentOnTypeFormattingOptions, FoldingRangeProviderCapability, HoverProviderCapability,
    ImplementationProviderCapability, RenameOptions, RenameProviderCapability, SaveOptions,
    SelectionRangeProviderCapability, SemanticTokensDocumentProvider, SemanticTokensLegend,
    SemanticTokensOptions, ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability,
    TextDocumentSyncKind, TextDocumentSyncOptions, TypeDefinitionProviderCapability,
    WorkDoneProgressOptions,
};
use serde_json::json;

use crate::semantic_tokens;

pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabilities {
    let code_action_provider = code_action_capabilities(client_caps);

    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_NO_INCREMENTAL_SYNC").is_ok() {
                TextDocumentSyncKind::Full
            } else {
                TextDocumentSyncKind::Incremental
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default().into()),
        })),
        hover_provider: Some(HoverProviderCapability::Simple(true)),
        completion_provider: Some(CompletionOptions {
            resolve_provider: None,
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(true),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(true),
        document_highlight_provider: Some(true),
        document_symbol_provider: Some(true),
        workspace_symbol_provider: Some(true),
        code_action_provider: Some(code_action_provider),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(true),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(RenameProviderCapability::Options(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: None,
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider: Some(
            SemanticTokensOptions {
                legend: SemanticTokensLegend {
                    token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                    token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
                },

                document_provider: Some(SemanticTokensDocumentProvider::Bool(true)),
                range_provider: Some(true),
                work_done_progress_options: Default::default(),
            }
            .into(),
        ),
        experimental: Some(json!({
            "joinLines": true,
            "ssr": true,
            "onEnter": true,
            "parentModule": true,
            "runnables": {
                "kinds": [ "cargo" ],
            },
        })),
    }
}

fn code_action_capabilities(client_caps: &ClientCapabilities) -> CodeActionProviderCapability {
    client_caps
        .text_document
        .as_ref()
        .and_then(|it| it.code_action.as_ref())
        .and_then(|it| it.code_action_literal_support.as_ref())
        .map_or(CodeActionProviderCapability::Simple(true), |_| {
            CodeActionProviderCapability::Options(CodeActionOptions {
                // Advertise support for all built-in CodeActionKinds.
                // Ideally we would base this off of the client capabilities
                // but the client is supposed to fall back gracefully for unknown values.
                code_action_kinds: Some(vec![
                    CodeActionKind::EMPTY,
                    CodeActionKind::QUICKFIX,
                    CodeActionKind::REFACTOR,
                    CodeActionKind::REFACTOR_EXTRACT,
                    CodeActionKind::REFACTOR_INLINE,
                    CodeActionKind::REFACTOR_REWRITE,
                ]),
                work_done_progress_options: Default::default(),
            })
        })
}
<commit_msg>Prepare server capabilities based on client ones.<commit_after>//! Advertizes the capabilities of the LSP Server.
use std::env;

use lsp_types::{
    CallHierarchyServerCapability, ClientCapabilities, CodeActionKind, CodeActionOptions,
    CodeActionProviderCapability, CodeLensOptions, CompletionOptions,
    DocumentOnTypeFormattingOptions, FoldingRangeProviderCapability, HoverProviderCapability,
    ImplementationProviderCapability, RenameOptions, RenameProviderCapability, SaveOptions,
    SelectionRangeProviderCapability, SemanticTokensDocumentProvider, SemanticTokensLegend,
    SemanticTokensOptions, SemanticTokensServerCapabilities, ServerCapabilities,
    SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
    TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
};
use serde_json::{json, Value};

use crate::semantic_tokens;

pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabilities {
    let code_action_provider = code_action_capabilities(client_caps);
    let semantic_tokens_provider = semantic_tokens_capabilities(client_caps);
    let experimental = experimental_capabilities(client_caps);

    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_NO_INCREMENTAL_SYNC").is_ok() {
                TextDocumentSyncKind::Full
            } else {
                TextDocumentSyncKind::Incremental
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default().into()),
        })),
        hover_provider: Some(HoverProviderCapability::Simple(true)),
        completion_provider: Some(CompletionOptions {
            resolve_provider: None,
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(true),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(true),
        document_highlight_provider: Some(true),
        document_symbol_provider: Some(true),
        workspace_symbol_provider: Some(true),
        code_action_provider: Some(code_action_provider),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(true),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(RenameProviderCapability::Options(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: None,
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider,
        experimental,
    }
}

fn experimental_capabilities(client_caps: &ClientCapabilities) -> Option<Value> {
    client_caps.experimental.as_ref().and_then(|it| {
        it.as_object().map(|map| {
            let mut obj = json!({});
            let result = obj.as_object_mut().unwrap();

            if map.contains_key("joinLines") {
                result.insert("joinLines".into(), true.into());
            }

            if map.contains_key("ssr") {
                result.insert("ssr".into(), true.into());
            }

            if map.contains_key("onEnter") {
                result.insert("onEnter".into(), true.into());
            }

            if map.contains_key("parentModule") {
                result.insert("parentModule".into(), true.into());
            }

            if map.contains_key("runnables") {
                result.insert("runnables".into(), json!({ "kinds": [ "cargo" ] }));
            }

            obj
        })
    })
}

fn semantic_tokens_capabilities(
    client_caps: &ClientCapabilities,
) -> Option<SemanticTokensServerCapabilities> {
    client_caps.text_document.as_ref().and_then(|it| it.semantic_tokens.as_ref()).map(|_|
            // client supports semanticTokens
            SemanticTokensOptions {
            legend: SemanticTokensLegend {
                token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
            },

            document_provider: Some(SemanticTokensDocumentProvider::Bool(true)),
            range_provider: Some(true),
            work_done_progress_options: Default::default(),
        }
        .into())
}

fn code_action_capabilities(client_caps: &ClientCapabilities) -> CodeActionProviderCapability {
    client_caps
        .text_document
        .as_ref()
        .and_then(|it| it.code_action.as_ref())
        .and_then(|it| it.code_action_literal_support.as_ref())
        .map_or(CodeActionProviderCapability::Simple(true), |caps| {
            let mut action_kinds = vec![
                CodeActionKind::EMPTY,
                CodeActionKind::QUICKFIX,
                CodeActionKind::REFACTOR,
                CodeActionKind::REFACTOR_EXTRACT,
                CodeActionKind::REFACTOR_INLINE,
                CodeActionKind::REFACTOR_REWRITE,
            ];

            // Not all clients can fall back gracefully for unknown values.
            // Microsoft.VisualStudio.LanguageServer.Protocol.CodeActionKind does not support CodeActionKind::EMPTY
            // So have to filter out.
            action_kinds
                .retain(|it| caps.code_action_kind.value_set.contains(&it.as_str().to_owned()));

            CodeActionProviderCapability::Options(CodeActionOptions {
                code_action_kinds: Some(action_kinds),
                work_done_progress_options: Default::default(),
            })
        })
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "cargo-clippy", allow(unreadable_literal))]
#[macro_use]
extern crate nom;

use nom::{Err, Needed};
#[cfg(feature = "alloc")]
use nom::number::streaming::be_u64;

// Parser definition

// We request a length that would trigger an overflow if computing consumed + requested
named!(parser01<&[u8],()>,
    do_parse!(
        hdr: take!(1) >>
        data: take!(18446744073709551615) >>
        ({
          let _ = hdr;
          let _ = data;
          ()
        })
    )
);

// We request a length that would trigger an overflow if computing consumed + requested
named!(parser02<&[u8],(&[u8],&[u8])>,
    tuple!(take!(1),take!(18446744073709551615))
);

#[test]
fn overflow_incomplete_do_parse() {
  assert_eq!(
    parser01(&b"3"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
fn overflow_incomplete_tuple() {
  assert_eq!(
    parser02(&b"3"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_bytes() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  // Trigger an overflow in length_data
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xff\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many0() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  // Trigger an overflow in many0
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many1() {
  named!(multi<&[u8], Vec<&[u8]> >, many1!( length_data!(be_u64) ) );

  // Trigger an overflow in many1
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many_till() {
  named!(multi<&[u8], (Vec<&[u8]>, &[u8]) >, many_till!( length_data!(be_u64), tag!("abc") ) );

  // Trigger an overflow in many_till
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many_m_n() {
  named!(multi<&[u8], Vec<&[u8]> >, many_m_n!(2, 4, length_data!(be_u64) ) );

  // Trigger an overflow in many_m_n
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_count() {
  named!(counter<&[u8], Vec<&[u8]> >, count!( length_data!(be_u64), 2 ) );

  assert_eq!(
    counter(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_count() {
  use nom::number::streaming::be_u8;
  named!(multi<&[u8], Vec<&[u8]> >, length_count!( be_u8, length_data!(be_u64) ) );

  assert_eq!(
    multi(&b"\x04\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xee\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551598)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_data() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xff\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}
<commit_msg>only build overflow tests on 64 bits arch<commit_after>#![cfg_attr(feature = "cargo-clippy", allow(unreadable_literal))]
#![cfg(target_pointer_width = "64")]

#[macro_use]
extern crate nom;

use nom::{Err, Needed};
#[cfg(feature = "alloc")]
use nom::number::streaming::be_u64;

// Parser definition

// We request a length that would trigger an overflow if computing consumed + requested
named!(parser01<&[u8],()>,
    do_parse!(
        hdr: take!(1) >>
        data: take!(18446744073709551615) >>
        ({
          let _ = hdr;
          let _ = data;
          ()
        })
    )
);

// We request a length that would trigger an overflow if computing consumed + requested
named!(parser02<&[u8],(&[u8],&[u8])>,
    tuple!(take!(1),take!(18446744073709551615))
);

#[test]
fn overflow_incomplete_do_parse() {
  assert_eq!(
    parser01(&b"3"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
fn overflow_incomplete_tuple() {
  assert_eq!(
    parser02(&b"3"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_bytes() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  // Trigger an overflow in length_data
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xff\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many0() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  // Trigger an overflow in many0
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many1() {
  named!(multi<&[u8], Vec<&[u8]> >, many1!( length_data!(be_u64) ) );

  // Trigger an overflow in many1
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many_till() {
  named!(multi<&[u8], (Vec<&[u8]>, &[u8]) >, many_till!( length_data!(be_u64), tag!("abc") ) );

  // Trigger an overflow in many_till
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_many_m_n() {
  named!(multi<&[u8], Vec<&[u8]> >, many_m_n!(2, 4, length_data!(be_u64) ) );

  // Trigger an overflow in many_m_n
  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_count() {
  named!(counter<&[u8], Vec<&[u8]> >, count!( length_data!(be_u64), 2 ) );

  assert_eq!(
    counter(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xef\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551599)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_count() {
  use nom::number::streaming::be_u8;
  named!(multi<&[u8], Vec<&[u8]> >, length_count!( be_u8, length_data!(be_u64) ) );

  assert_eq!(
    multi(&b"\x04\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xee\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551598)))
  );
}

#[test]
#[cfg(feature = "alloc")]
fn overflow_incomplete_length_data() {
  named!(multi<&[u8], Vec<&[u8]> >, many0!( length_data!(be_u64) ) );

  assert_eq!(
    multi(&b"\x00\x00\x00\x00\x00\x00\x00\x01\xaa\xff\xff\xff\xff\xff\xff\xff\xff\xaa"[..]),
    Err(Err::Incomplete(Needed::Size(18446744073709551615)))
  );
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::HashSet;

use serde_json;
use failure::Context;

use errors::prelude::*;
use services::ledger::merkletree::merkletree::MerkleTree;
use services::pool::merkle_tree_factory;
use services::pool::rust_base58::{FromBase58, ToBase58};
use services::pool::types::{CatchupReq, Message};

pub enum CatchupProgress {
    ShouldBeStarted(
        Vec<u8>, //target_mt_root
        usize, //target_mt_size
        MerkleTree,
    ),
    NotNeeded(MerkleTree),
    Restart(MerkleTree),
    InProgress,
}

pub fn build_catchup_req(merkle: &MerkleTree, target_mt_size: usize) -> IndyResult<Option<(String, String)>> {
    let txns_cnt = target_mt_size - merkle.count();

    if txns_cnt <= 0 {
        warn!("No transactions to catch up!");
        return Ok(None);
    }
    let seq_no_start = merkle.count() + 1;
    let seq_no_end = target_mt_size;

    let cr = CatchupReq {
        ledgerId: 0,
        seqNoStart: seq_no_start.clone(),
        seqNoEnd: seq_no_end.clone(),
        catchupTill: target_mt_size,
    };

    let req_id = format!("{}{}", seq_no_start, seq_no_end);

    let req_json = serde_json::to_string(&Message::CatchupReq(cr))
        .to_indy(IndyErrorKind::InvalidState, "Cannot serialize CatchupRequest")?;

    trace!("catchup_req msg: {:?}", req_json);
    Ok(Some((req_id, req_json)))
}

pub fn check_nodes_responses_on_status(nodes_votes: &HashMap<(String, usize, Option<Vec<String>>), HashSet<String>>,
                                       merkle_tree: &MerkleTree,
                                       node_cnt: usize,
                                       f: usize,
                                       pool_name: &str) -> IndyResult<CatchupProgress> {
    //FIXME fetch not just most popular vote but most popular not a timeout
    if let Some((most_popular_vote, votes_cnt)) = nodes_votes.iter().map(|(key, val)| (key, val.len())).max_by_key(|entry| entry.1) {
        let is_consensus_reached =
            votes_cnt == node_cnt - f
                || votes_cnt == f + 1 && most_popular_vote.0.neq("timeout");
        if is_consensus_reached {
            if most_popular_vote.0.eq("timeout") {
                return _try_to_restart_catch_up(pool_name, err_msg(IndyErrorKind::PoolTimeout, "Pool timeout"))
            }

            return _try_to_catch_up(most_popular_vote, merkle_tree).or_else(|err| {
                if merkle_tree_factory::drop_cache(pool_name).is_ok() {
                    let merkle_tree = merkle_tree_factory::create(pool_name)?;
                    _try_to_catch_up(most_popular_vote, &merkle_tree)
                } else {
                    Err(err)
                }
            });
        } else {
            let reps_cnt: usize = nodes_votes.values().map(HashSet::len).sum();
            let positive_votes_cnt = votes_cnt + (node_cnt - reps_cnt);
            let is_consensus_reachable = positive_votes_cnt < node_cnt - f;
            if is_consensus_reachable {
                //TODO: maybe we should change the error, but it was made to escape changing of ErrorCode returned to client
               return  _try_to_restart_catch_up(pool_name, err_msg(IndyErrorKind::PoolTimeout, "No consensus possible"))
            }
        }
    }
    Ok(CatchupProgress::InProgress)
}

fn _try_to_restart_catch_up(pool_name: &str, err: IndyError) -> IndyResult<CatchupProgress>{
    if merkle_tree_factory::drop_cache(pool_name).is_ok() {
        let merkle_tree = merkle_tree_factory::create(pool_name)?;
        return Ok(CatchupProgress::Restart(merkle_tree))
    } else {
        return Err(err);
    }
}

fn _try_to_catch_up(ledger_status: &(String, usize, Option<Vec<String>>), merkle_tree: &MerkleTree) -> IndyResult<CatchupProgress> {
    let &(ref target_mt_root, target_mt_size, ref hashes) = ledger_status;
    let cur_mt_size = merkle_tree.count();
    let cur_mt_hash = merkle_tree.root_hash().to_base58();

    if target_mt_size == cur_mt_size {
        if cur_mt_hash.eq(target_mt_root) {
            return Ok(CatchupProgress::NotNeeded(merkle_tree.clone()));
        } else {
            return Err(err_msg(IndyErrorKind::InvalidState,
                               "Ledger merkle tree is not acceptable for current tree."));
        }
    } else if target_mt_size > cur_mt_size {
        let target_mt_root = target_mt_root
            .from_base58()
            .map_err(|err| Context::new(err))
            .to_indy(IndyErrorKind::InvalidStructure, "Can't parse target MerkleTree hash from nodes responses")?; // FIXME: review kind

        match hashes {
            &None => (),
            &Some(ref hashes) => check_cons_proofs(merkle_tree, hashes, &target_mt_root, target_mt_size)?,
        };

        return Ok(CatchupProgress::ShouldBeStarted(target_mt_root, target_mt_size, merkle_tree.clone()));
    } else {
        return Err(err_msg(IndyErrorKind::InvalidState, "Local merkle tree greater than mt from ledger"));
    }
}

pub fn check_cons_proofs(mt: &MerkleTree, cons_proofs: &Vec<String>, target_mt_root: &Vec<u8>, target_mt_size: usize) -> IndyResult<()> {
    let mut bytes_proofs: Vec<Vec<u8>> = Vec::new();

    for cons_proof in cons_proofs {
        let cons_proof: &String = cons_proof;

        bytes_proofs.push(
            cons_proof
                .from_base58()
                .map_err(|err| Context::new(err))
                .to_indy(IndyErrorKind::InvalidStructure, "Can't decode node consistency proof")? // FIXME: review kind
        );
    }

    if !mt.consistency_proof(target_mt_root, target_mt_size, &bytes_proofs)? {
        return Err(err_msg(IndyErrorKind::InvalidState, "Consistency proof verification failed")); // FIXME: review kind
    }

    Ok(())
}<commit_msg>Resolved FIXME<commit_after>use std::collections::HashMap;
use std::collections::HashSet;

use serde_json;
use failure::Context;

use errors::prelude::*;
use services::ledger::merkletree::merkletree::MerkleTree;
use services::pool::merkle_tree_factory;
use services::pool::rust_base58::{FromBase58, ToBase58};
use services::pool::types::{CatchupReq, Message};

pub enum CatchupProgress {
    ShouldBeStarted(
        Vec<u8>, //target_mt_root
        usize, //target_mt_size
        MerkleTree,
    ),
    NotNeeded(MerkleTree),
    Restart(MerkleTree),
    InProgress,
}

pub fn build_catchup_req(merkle: &MerkleTree, target_mt_size: usize) -> IndyResult<Option<(String, String)>> {
    let txns_cnt = target_mt_size - merkle.count();

    if txns_cnt <= 0 {
        warn!("No transactions to catch up!");
        return Ok(None);
    }
    let seq_no_start = merkle.count() + 1;
    let seq_no_end = target_mt_size;

    let cr = CatchupReq {
        ledgerId: 0,
        seqNoStart: seq_no_start.clone(),
        seqNoEnd: seq_no_end.clone(),
        catchupTill: target_mt_size,
    };

    let req_id = format!("{}{}", seq_no_start, seq_no_end);

    let req_json = serde_json::to_string(&Message::CatchupReq(cr))
        .to_indy(IndyErrorKind::InvalidState, "Cannot serialize CatchupRequest")?;

    trace!("catchup_req msg: {:?}", req_json);
    Ok(Some((req_id, req_json)))
}

pub fn check_nodes_responses_on_status(nodes_votes: &HashMap<(String, usize, Option<Vec<String>>), HashSet<String>>,
                                       merkle_tree: &MerkleTree,
                                       node_cnt: usize,
                                       f: usize,
                                       pool_name: &str) -> IndyResult<CatchupProgress> {
    let (votes, timeout_votes): (HashMap<&(String, usize, Option<Vec<String>>), usize>, HashMap<&(String, usize, Option<Vec<String>>), usize>) =
        nodes_votes
            .iter()
            .map(|(key, val)| (key, val.len()))
            .partition(|((key, _, _), _)| key != "timeout");

    let most_popular_not_timeout =
        votes
            .iter()
            .max_by_key(|entry| entry.1);

    let timeout_votes = timeout_votes.iter().last();

    if let Some((most_popular_not_timeout_vote, vote_cnt)) = most_popular_not_timeout {
        if *vote_cnt == f + 1 {
            return _try_to_catch_up(most_popular_not_timeout_vote, merkle_tree).or_else(|err| {
                if merkle_tree_factory::drop_cache(pool_name).is_ok() {
                    let merkle_tree = merkle_tree_factory::create(pool_name)?;
                    _try_to_catch_up(most_popular_not_timeout_vote, &merkle_tree)
                } else {
                    Err(err)
                }
            });
        }
    } else if let Some((_, votes_cnt)) = timeout_votes {
        if *votes_cnt == node_cnt - f {
            return _try_to_restart_catch_up(pool_name, err_msg(IndyErrorKind::PoolTimeout, "Pool timeout"));
        } else {
            let reps_cnt: usize = nodes_votes.values().map(HashSet::len).sum();
            let positive_votes_cnt = votes_cnt + (node_cnt - reps_cnt);
            let is_consensus_reachable = positive_votes_cnt < node_cnt - f;
            if is_consensus_reachable {
                //TODO: maybe we should change the error, but it was made to escape changing of ErrorCode returned to client
                return _try_to_restart_catch_up(pool_name, err_msg(IndyErrorKind::PoolTimeout, "No consensus possible"));
            }
        }
    }
    Ok(CatchupProgress::InProgress)
}

fn _try_to_restart_catch_up(pool_name: &str, err: IndyError) -> IndyResult<CatchupProgress> {
    if merkle_tree_factory::drop_cache(pool_name).is_ok() {
        let merkle_tree = merkle_tree_factory::create(pool_name)?;
        return Ok(CatchupProgress::Restart(merkle_tree));
    } else {
        return Err(err);
    }
}

fn _try_to_catch_up(ledger_status: &(String, usize, Option<Vec<String>>), merkle_tree: &MerkleTree) -> IndyResult<CatchupProgress> {
    let &(ref target_mt_root, target_mt_size, ref hashes) = ledger_status;
    let cur_mt_size = merkle_tree.count();
    let cur_mt_hash = merkle_tree.root_hash().to_base58();

    if target_mt_size == cur_mt_size {
        if cur_mt_hash.eq(target_mt_root) {
            return Ok(CatchupProgress::NotNeeded(merkle_tree.clone()));
        } else {
            return Err(err_msg(IndyErrorKind::InvalidState,
                               "Ledger merkle tree is not acceptable for current tree."));
        }
    } else if target_mt_size > cur_mt_size {
        let target_mt_root = target_mt_root
            .from_base58()
            .map_err(|err| Context::new(err))
            .to_indy(IndyErrorKind::InvalidStructure, "Can't parse target MerkleTree hash from nodes responses")?; // FIXME: review kind

        match hashes {
            &None => (),
            &Some(ref hashes) => check_cons_proofs(merkle_tree, hashes, &target_mt_root, target_mt_size)?,
        };

        return Ok(CatchupProgress::ShouldBeStarted(target_mt_root, target_mt_size, merkle_tree.clone()));
    } else {
        return Err(err_msg(IndyErrorKind::InvalidState, "Local merkle tree greater than mt from ledger"));
    }
}

pub fn check_cons_proofs(mt: &MerkleTree, cons_proofs: &Vec<String>, target_mt_root: &Vec<u8>, target_mt_size: usize) -> IndyResult<()> {
    let mut bytes_proofs: Vec<Vec<u8>> = Vec::new();

    for cons_proof in cons_proofs {
        let cons_proof: &String = cons_proof;

        bytes_proofs.push(
            cons_proof
                .from_base58()
                .map_err(|err| Context::new(err))
                .to_indy(IndyErrorKind::InvalidStructure, "Can't decode node consistency proof")? // FIXME: review kind
        );
    }

    if !mt.consistency_proof(target_mt_root, target_mt_size, &bytes_proofs)? {
        return Err(err_msg(IndyErrorKind::InvalidState, "Consistency proof verification failed")); // FIXME: review kind
    }

    Ok(())
}<|endoftext|>