rust-analyzer - Unnamed repository; edit this file 'description' to name the repository.


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

//! Post-nameres attribute resolution.

use base_db::CrateId;
use hir_expand::{
    attrs::{Attr, AttrId, AttrInput},
    inert_attr_macro::find_builtin_attr_idx,
    MacroCallId, MacroCallKind, MacroDefId,
};
use span::SyntaxContextId;
use syntax::ast;
use triomphe::Arc;

use crate::{
    db::DefDatabase,
    item_scope::BuiltinShadowMode,
    nameres::path_resolution::ResolveMode,
    path::{self, ModPath, PathKind},
    AstIdWithPath, LocalModuleId, MacroId, UnresolvedMacro,
};

use super::{DefMap, MacroSubNs};

pub enum ResolvedAttr {
    /// Attribute resolved to an attribute macro.
    Macro(MacroCallId),
    /// Attribute resolved to something else that does not require expansion.
    Other,
}

impl DefMap {
    pub(crate) fn resolve_attr_macro(
        &self,
        db: &dyn DefDatabase,
        original_module: LocalModuleId,
        ast_id: AstIdWithPath<ast::Item>,
        attr: &Attr,
    ) -> Result<ResolvedAttr, UnresolvedMacro> {
        // NB: does not currently work for derive helpers as they aren't recorded in the `DefMap`

        if self.is_builtin_or_registered_attr(&ast_id.path) {
            return Ok(ResolvedAttr::Other);
        }

        let resolved_res = self.resolve_path_fp_with_macro(
            db,
            ResolveMode::Other,
            original_module,
            &ast_id.path,
            BuiltinShadowMode::Module,
            Some(MacroSubNs::Attr),
        );
        let def = match resolved_res.resolved_def.take_macros() {
            Some(def) => {
                // `MacroSubNs` is just a hint, so the path may still resolve to a custom derive
                // macro, or even function-like macro when the path is qualified.
                if def.is_attribute(db) {
                    def
                } else {
                    return Ok(ResolvedAttr::Other);
                }
            }
            None => return Err(UnresolvedMacro { path: ast_id.path.as_ref().clone() }),
        };

        Ok(ResolvedAttr::Macro(attr_macro_as_call_id(
            db,
            &ast_id,
            attr,
            self.krate,
            db.macro_def(def),
        )))
    }

    pub(crate) fn is_builtin_or_registered_attr(&self, path: &ModPath) -> bool {
        if path.kind != PathKind::Plain {
            return false;
        }

        let segments = path.segments();

        if let Some(name) = segments.first() {
            let name = name.symbol();
            let pred = |n: &_| *n == *name;

            let is_tool = self.data.registered_tools.iter().any(pred);
            // FIXME: tool modules can be shadowed by actual modules
            if is_tool {
                return true;
            }

            if segments.len() == 1 {
                if find_builtin_attr_idx(name).is_some() {
                    return true;
                }
                if self.data.registered_attrs.iter().any(pred) {
                    return true;
                }
            }
        }
        false
    }
}

pub(super) fn attr_macro_as_call_id(
    db: &dyn DefDatabase,
    item_attr: &AstIdWithPath<ast::Item>,
    macro_attr: &Attr,
    krate: CrateId,
    def: MacroDefId,
) -> MacroCallId {
    let arg = match macro_attr.input.as_deref() {
        Some(AttrInput::TokenTree(tt)) => {
            let mut tt = tt.as_ref().clone();
            tt.delimiter.kind = tt::DelimiterKind::Invisible;
            Some(tt)
        }

        _ => None,
    };

    def.make_call(
        db.upcast(),
        krate,
        MacroCallKind::Attr {
            ast_id: item_attr.ast_id,
            attr_args: arg.map(Arc::new),
            invoc_attr_index: macro_attr.id,
        },
        macro_attr.ctxt,
    )
}

pub(super) fn derive_macro_as_call_id(
    db: &dyn DefDatabase,
    item_attr: &AstIdWithPath<ast::Adt>,
    derive_attr_index: AttrId,
    derive_pos: u32,
    call_site: SyntaxContextId,
    krate: CrateId,
    resolver: impl Fn(&path::ModPath) -> Option<(MacroId, MacroDefId)>,
    derive_macro_id: MacroCallId,
) -> Result<(MacroId, MacroDefId, MacroCallId), UnresolvedMacro> {
    let (macro_id, def_id) = resolver(&item_attr.path)
        .filter(|(_, def_id)| def_id.is_derive())
        .ok_or_else(|| UnresolvedMacro { path: item_attr.path.as_ref().clone() })?;
    let call_id = def_id.make_call(
        db.upcast(),
        krate,
        MacroCallKind::Derive {
            ast_id: item_attr.ast_id,
            derive_index: derive_pos,
            derive_attr_index,
            derive_macro_id,
        },
        call_site,
    );
    Ok((macro_id, def_id, call_id))
}
use hir::{Name, sym};
use ide_db::{famous_defs::FamousDefs, syntax_helpers::suggest_name};
use syntax::{
    AstNode,
    ast::{self, HasAttrs, HasLoopBody, edit::IndentLevel},
    syntax_editor::Position,
};

use crate::{AssistContext, AssistId, Assists};

// Assist: convert_for_loop_to_while_let
//
// Converts a for loop into a while let on the Iterator.
//
// ```
// fn main() {
//     let x = vec![1, 2, 3];
//     for$0 v in x {
//         let y = v * 2;
//     };
// }
// ```
// ->
// ```
// fn main() {
//     let x = vec![1, 2, 3];
//     let mut iter = x.into_iter();
//     while let Some(v) = iter.next() {
//         let y = v * 2;
//     };
// }
// ```
pub(crate) fn convert_for_loop_to_while_let(
    acc: &mut Assists,
    ctx: &AssistContext<'_>,
) -> Option<()> {
    let for_loop = ctx.find_node_at_offset::<ast::ForExpr>()?;
    let iterable = for_loop.iterable()?;
    let pat = for_loop.pat()?;
    let body = for_loop.loop_body()?;
    if body.syntax().text_range().start() < ctx.offset() {
        cov_mark::hit!(not_available_in_body);
        return None;
    }

    acc.add(
        AssistId::refactor_rewrite("convert_for_loop_to_while_let"),
        "Replace this for loop with `while let`",
        for_loop.syntax().text_range(),
        |builder| {
            let editor = builder.make_editor(for_loop.syntax());
            let make = editor.make();

            let (iterable, method) = if impls_core_iter(&ctx.sema, &iterable) {
                (iterable, None)
            } else if let Some((expr, method)) = is_ref_and_impls_iter_method(&ctx.sema, &iterable)
            {
                (expr, Some(make.name_ref(method.as_str())))
            } else if let ast::Expr::RefExpr(_) = iterable {
                (make.expr_paren(iterable).into(), Some(make.name_ref("into_iter")))
            } else {
                (iterable, Some(make.name_ref("into_iter")))
            };

            let iterable = if let Some(method) = method {
                make.expr_method_call(iterable, method, make.arg_list([])).into()
            } else {
                iterable
            };

            let mut new_name = suggest_name::NameGenerator::new_from_scope_locals(
                ctx.sema.scope(for_loop.syntax()),
            );
            let tmp_var = new_name.suggest_name("iter");

            let mut_expr = make.let_stmt(
                make.ident_pat(false, true, make.name(&tmp_var)).into(),
                None,
                Some(iterable),
            );
            let indent = IndentLevel::from_node(for_loop.syntax());

            if let Some(label) = for_loop.label() {
                let label = label.syntax();
                editor.insert(Position::before(for_loop.syntax()), make.whitespace(" "));
                editor.insert(Position::before(for_loop.syntax()), label);
            }
            crate::utils::insert_attributes(for_loop.syntax(), &editor, for_loop.attrs());

            editor.insert(
                Position::before(for_loop.syntax()),
                make.whitespace(format!("\n{indent}").as_str()),
            );
            editor.insert(Position::before(for_loop.syntax()), mut_expr.syntax());

            let opt_pat = make.tuple_struct_pat(make.ident_path("Some"), [pat]);
            let iter_next_expr = make.expr_method_call(
                make.expr_path(make.ident_path(&tmp_var)),
                make.name_ref("next"),
                make.arg_list([]),
            );
            let cond = make.expr_let(opt_pat.into(), iter_next_expr.into());

            let while_loop = make.expr_while_loop(cond.into(), body);

            editor.replace(for_loop.syntax(), while_loop.syntax());

            builder.add_file_edits(ctx.vfs_file_id(), editor);
        },
    )
}

/// If iterable is a reference where the expression behind the reference implements a method
/// returning an Iterator called iter or iter_mut (depending on the type of reference) then return
/// the expression behind the reference and the method name
fn is_ref_and_impls_iter_method(
    sema: &hir::Semantics<'_, ide_db::RootDatabase>,
    iterable: &ast::Expr,
) -> Option<(ast::Expr, hir::Name)> {
    let ref_expr = match iterable {
        ast::Expr::RefExpr(r) => r,
        _ => return None,
    };
    let wanted_method = Name::new_symbol_root(if ref_expr.mut_token().is_some() {
        sym::iter_mut
    } else {
        sym::iter
    });
    let expr_behind_ref = ref_expr.expr()?;
    let ty = sema.type_of_expr(&expr_behind_ref)?.adjusted();
    let scope = sema.scope(iterable.syntax())?;
    let krate = scope.krate();
    let iter_trait = FamousDefs(sema, krate).core_iter_Iterator()?;

    let has_wanted_method = ty
        .iterate_method_candidates(sema.db, &scope, Some(&wanted_method), |func| {
            if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) {
                return Some(());
            }
            None
        })
        .is_some();
    if !has_wanted_method {
        return None;
    }

    Some((expr_behind_ref, wanted_method))
}

/// Whether iterable implements core::Iterator
fn impls_core_iter(sema: &hir::Semantics<'_, ide_db::RootDatabase>, iterable: &ast::Expr) -> bool {
    (|| {
        let it_typ = sema.type_of_expr(iterable)?.adjusted();

        let module = sema.scope(iterable.syntax())?.module();

        let krate = module.krate(sema.db);
        let iter_trait = FamousDefs(sema, krate).core_iter_Iterator()?;
        cov_mark::hit!(test_already_impls_iterator);
        Some(it_typ.impls_trait(sema.db, iter_trait, &[]))
    })()
    .unwrap_or(false)
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn each_to_for_simple_for() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    for $0v in x {
        v *= 2;
    };
}",
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    let mut iter = x.into_iter();
    while let Some(v) = iter.next() {
        v *= 2;
    };
}",
        )
    }

    #[test]
    fn each_to_for_with_label() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    'a: for $0v in x {
        v *= 2;
        break 'a;
    };
}",
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    let mut iter = x.into_iter();
    'a: while let Some(v) = iter.next() {
        v *= 2;
        break 'a;
    };
}",
        )
    }

    #[test]
    fn each_to_for_with_attributes() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    #[allow(unused)]
    #[deny(unsafe_code)]
    for $0v in x {
        v *= 2;
    };
}",
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    let mut iter = x.into_iter();
    #[allow(unused)]
    #[deny(unsafe_code)]
    while let Some(v) = iter.next() {
        v *= 2;
    };
}",
        )
    }

    #[test]
    fn each_to_for_for_in_range() {
        check_assist(
            convert_for_loop_to_while_let,
            r#"
//- minicore: range, iterators
impl<T> core::iter::Iterator for core::ops::Range<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

fn main() {
    for $0x in 0..92 {
        print!("{}", x);
    }
}"#,
            r#"
impl<T> core::iter::Iterator for core::ops::Range<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

fn main() {
    let mut iter = 0..92;
    while let Some(x) = iter.next() {
        print!("{}", x);
    }
}"#,
        )
    }

    #[test]
    fn each_to_for_not_available_in_body() {
        cov_mark::check!(not_available_in_body);
        check_assist_not_applicable(
            convert_for_loop_to_while_let,
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    for v in x {
        $0v *= 2;
    }
}",
        )
    }

    #[test]
    fn each_to_for_for_borrowed() {
        check_assist(
            convert_for_loop_to_while_let,
            r#"
//- minicore: iterators
use core::iter::{Repeat, repeat};

struct S;
impl S {
    fn iter(&self) -> Repeat<i32> { repeat(92) }
    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
}

fn main() {
    let x = S;
    for $0v in &x {
        let a = v * 2;
    }
}
"#,
            r#"
use core::iter::{Repeat, repeat};

struct S;
impl S {
    fn iter(&self) -> Repeat<i32> { repeat(92) }
    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
}

fn main() {
    let x = S;
    let mut iter = x.iter();
    while let Some(v) = iter.next() {
        let a = v * 2;
    }
}
"#,
        )
    }

    #[test]
    fn each_to_for_for_borrowed_no_iter_method() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
struct NoIterMethod;
fn main() {
    let x = NoIterMethod;
    for $0v in &x {
        let a = v * 2;
    }
}
",
            r"
struct NoIterMethod;
fn main() {
    let x = NoIterMethod;
    let mut iter = (&x).into_iter();
    while let Some(v) = iter.next() {
        let a = v * 2;
    }
}
",
        )
    }

    #[test]
    fn each_to_for_for_borrowed_no_iter_method_mut() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
struct NoIterMethod;
fn main() {
    let x = NoIterMethod;
    for $0v in &mut x {
        let a = v * 2;
    }
}
",
            r"
struct NoIterMethod;
fn main() {
    let x = NoIterMethod;
    let mut iter = (&mut x).into_iter();
    while let Some(v) = iter.next() {
        let a = v * 2;
    }
}
",
        )
    }

    #[test]
    fn each_to_for_for_borrowed_mut() {
        check_assist(
            convert_for_loop_to_while_let,
            r#"
//- minicore: iterators
use core::iter::{Repeat, repeat};

struct S;
impl S {
    fn iter(&self) -> Repeat<i32> { repeat(92) }
    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
}

fn main() {
    let x = S;
    for $0v in &mut x {
        let a = v * 2;
    }
}
"#,
            r#"
use core::iter::{Repeat, repeat};

struct S;
impl S {
    fn iter(&self) -> Repeat<i32> { repeat(92) }
    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
}

fn main() {
    let x = S;
    let mut iter = x.iter_mut();
    while let Some(v) = iter.next() {
        let a = v * 2;
    }
}
"#,
        )
    }

    #[test]
    fn each_to_for_for_borrowed_mut_behind_var() {
        check_assist(
            convert_for_loop_to_while_let,
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    let y = &mut x;
    for $0v in y {
        *v *= 2;
    }
}",
            r"
fn main() {
    let mut x = vec![1, 2, 3];
    let y = &mut x;
    let mut iter = y.into_iter();
    while let Some(v) = iter.next() {
        *v *= 2;
    }
}",
        )
    }

    #[test]
    fn each_to_for_already_impls_iterator() {
        cov_mark::check!(test_already_impls_iterator);
        check_assist(
            convert_for_loop_to_while_let,
            r#"
//- minicore: iterators
fn main() {
    for$0 a in core::iter::repeat(92).take(1) {
        println!("{}", a);
    }
}
"#,
            r#"
fn main() {
    let mut iter = core::iter::repeat(92).take(1);
    while let Some(a) = iter.next() {
        println!("{}", a);
    }
}
"#,
        );
    }
}