ide_db/syntax_helpers/

node_ext.rs

//! Various helper functions to work with SyntaxNodes.
use std::ops::ControlFlow;

use either::Either;
use itertools::Itertools;
use parser::T;
use span::Edition;
use syntax::{
    AstNode, AstToken, Direction, Preorder, RustLanguage, SyntaxToken, WalkEvent,
    algo::non_trivia_sibling,
    ast::{self, HasLoopBody, MacroCall, PathSegmentKind, VisibilityKind},
    syntax_editor::Element,
};

pub fn expr_as_name_ref(expr: &ast::Expr) -> Option<ast::NameRef> {
    if let ast::Expr::PathExpr(expr) = expr {
        let path = expr.path()?;
        path.as_single_name_ref()
    } else {
        None
    }
}

pub fn full_path_of_name_ref(name_ref: &ast::NameRef) -> Option<ast::Path> {
    let mut ancestors = name_ref.syntax().ancestors();
    let _ = ancestors.next()?; // skip self
    let _ = ancestors.next().filter(|it| ast::PathSegment::can_cast(it.kind()))?; // skip self
    ancestors.take_while(|it| ast::Path::can_cast(it.kind())).last().and_then(ast::Path::cast)
}

pub fn block_as_lone_tail(block: &ast::BlockExpr) -> Option<ast::Expr> {
    block.statements().next().is_none().then(|| block.tail_expr()).flatten()
}

/// Preorder walk all the expression's child expressions.
pub fn walk_expr(expr: &ast::Expr, cb: &mut dyn FnMut(ast::Expr)) {
    preorder_expr(expr, &mut |ev| {
        if let WalkEvent::Enter(expr) = ev {
            cb(expr);
        }
        false
    })
}

pub fn is_closure_or_blk_with_modif(expr: &ast::Expr) -> bool {
    match expr {
        ast::Expr::BlockExpr(block_expr) => {
            matches!(
                block_expr.modifier(),
                Some(
                    ast::BlockModifier::Async(_)
                        | ast::BlockModifier::Try(_)
                        | ast::BlockModifier::Const(_)
                )
            )
        }
        ast::Expr::ClosureExpr(_) => true,
        _ => false,
    }
}

/// Preorder walk all the expression's child expressions preserving events.
/// If the callback returns true on an [`WalkEvent::Enter`], the subtree of the expression will be skipped.
/// Note that the subtree may already be skipped due to the context analysis this function does.
pub fn preorder_expr(start: &ast::Expr, cb: &mut dyn FnMut(WalkEvent<ast::Expr>) -> bool) {
    preorder_expr_with_ctx_checker(start, &is_closure_or_blk_with_modif, cb);
}

pub fn preorder_expr_with_ctx_checker(
    start: &ast::Expr,
    check_ctx: &dyn Fn(&ast::Expr) -> bool,
    cb: &mut dyn FnMut(WalkEvent<ast::Expr>) -> bool,
) {
    let mut preorder = start.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(node) => {
                if let Some(expr) = ast::Expr::cast(node) {
                    cb(WalkEvent::Leave(expr));
                }
                continue;
            }
        };
        if let Some(let_stmt) = node.parent().and_then(ast::LetStmt::cast)
            && let_stmt.initializer().map(|it| it.syntax() != &node).unwrap_or(true)
            && let_stmt.let_else().map(|it| it.syntax() != &node).unwrap_or(true)
        {
            // skipping potential const pat expressions in  let statements
            preorder.skip_subtree();
            continue;
        }

        match ast::Stmt::cast(node.clone()) {
            // Don't skip subtree since we want to process the expression child next
            Some(ast::Stmt::ExprStmt(_)) | Some(ast::Stmt::LetStmt(_)) => (),
            // skip inner items which might have their own expressions
            Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
            None => {
                // skip const args, those expressions are a different context
                if ast::GenericArg::can_cast(node.kind()) {
                    preorder.skip_subtree();
                } else if let Some(expr) = ast::Expr::cast(node) {
                    let is_different_context = check_ctx(&expr) && expr.syntax() != start.syntax();
                    let skip = cb(WalkEvent::Enter(expr));
                    if skip || is_different_context {
                        preorder.skip_subtree();
                    }
                }
            }
        }
    }
}

/// Preorder walk all the expression's child patterns.
pub fn walk_patterns_in_expr(start: &ast::Expr, cb: &mut dyn FnMut(ast::Pat)) {
    let mut preorder = start.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        match ast::Stmt::cast(node.clone()) {
            Some(ast::Stmt::LetStmt(l)) => {
                if let Some(pat) = l.pat() {
                    _ = walk_pat(&pat, &mut |pat| {
                        cb(pat);
                        ControlFlow::<(), ()>::Continue(())
                    });
                }
                if let Some(expr) = l.initializer() {
                    walk_patterns_in_expr(&expr, cb);
                }
                preorder.skip_subtree();
            }
            // Don't skip subtree since we want to process the expression child next
            Some(ast::Stmt::ExprStmt(_)) => (),
            // skip inner items which might have their own patterns
            Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
            None => {
                // skip const args, those are a different context
                if ast::GenericArg::can_cast(node.kind()) {
                    preorder.skip_subtree();
                } else if let Some(expr) = ast::Expr::cast(node.clone()) {
                    let is_different_context = match &expr {
                        ast::Expr::BlockExpr(block_expr) => {
                            matches!(
                                block_expr.modifier(),
                                Some(
                                    ast::BlockModifier::Async(_)
                                        | ast::BlockModifier::Try(_)
                                        | ast::BlockModifier::Const(_)
                                )
                            )
                        }
                        ast::Expr::ClosureExpr(_) => true,
                        _ => false,
                    } && expr.syntax() != start.syntax();
                    if is_different_context {
                        preorder.skip_subtree();
                    }
                } else if let Some(pat) = ast::Pat::cast(node) {
                    preorder.skip_subtree();
                    _ = walk_pat(&pat, &mut |pat| {
                        cb(pat);
                        ControlFlow::<(), ()>::Continue(())
                    });
                }
            }
        }
    }
}

/// Preorder walk all the pattern's sub patterns.
pub fn walk_pat<T>(
    pat: &ast::Pat,
    cb: &mut dyn FnMut(ast::Pat) -> ControlFlow<T>,
) -> ControlFlow<T> {
    let mut preorder = pat.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        let kind = node.kind();
        match ast::Pat::cast(node) {
            Some(pat @ ast::Pat::ConstBlockPat(_)) => {
                preorder.skip_subtree();
                cb(pat)?;
            }
            Some(pat) => {
                cb(pat)?;
            }
            // skip const args
            None if ast::GenericArg::can_cast(kind) => {
                preorder.skip_subtree();
            }
            None => (),
        }
    }
    ControlFlow::Continue(())
}

/// Preorder walk all the type's sub types.
// FIXME: Make the control flow more proper
pub fn walk_ty(ty: &ast::Type, cb: &mut dyn FnMut(ast::Type) -> bool) {
    let mut preorder = ty.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        let kind = node.kind();
        match ast::Type::cast(node) {
            Some(ty @ ast::Type::MacroType(_)) => {
                preorder.skip_subtree();
                cb(ty);
            }
            Some(ty) => {
                if cb(ty) {
                    preorder.skip_subtree();
                }
            }
            // skip const args
            None if ast::ConstArg::can_cast(kind) => {
                preorder.skip_subtree();
            }
            None => (),
        }
    }
}

pub fn vis_eq(this: &ast::Visibility, other: &ast::Visibility) -> bool {
    match (this.kind(), other.kind()) {
        (VisibilityKind::In(this), VisibilityKind::In(other)) => {
            stdx::iter_eq_by(this.segments(), other.segments(), |lhs, rhs| {
                lhs.kind().zip(rhs.kind()).is_some_and(|it| match it {
                    (PathSegmentKind::CrateKw, PathSegmentKind::CrateKw)
                    | (PathSegmentKind::SelfKw, PathSegmentKind::SelfKw)
                    | (PathSegmentKind::SuperKw, PathSegmentKind::SuperKw) => true,
                    (PathSegmentKind::Name(lhs), PathSegmentKind::Name(rhs)) => {
                        lhs.text() == rhs.text()
                    }
                    _ => false,
                })
            })
        }
        (VisibilityKind::PubSelf, VisibilityKind::PubSelf)
        | (VisibilityKind::PubSuper, VisibilityKind::PubSuper)
        | (VisibilityKind::PubCrate, VisibilityKind::PubCrate)
        | (VisibilityKind::Pub, VisibilityKind::Pub) => true,
        _ => false,
    }
}

/// Returns the `let` only if there is exactly one (that is, `let pat = expr`
/// or `((let pat = expr))`, but not `let pat = expr && expr` or `non_let_expr`).
pub fn single_let(expr: ast::Expr) -> Option<ast::LetExpr> {
    match expr {
        ast::Expr::ParenExpr(expr) => expr.expr().and_then(single_let),
        ast::Expr::LetExpr(expr) => Some(expr),
        _ => None,
    }
}

pub fn is_pattern_cond(expr: ast::Expr) -> bool {
    match expr {
        ast::Expr::BinExpr(expr)
            if expr.op_kind() == Some(ast::BinaryOp::LogicOp(ast::LogicOp::And)) =>
        {
            expr.lhs().map_or(false, is_pattern_cond) || expr.rhs().map_or(false, is_pattern_cond)
        }
        ast::Expr::ParenExpr(expr) => expr.expr().is_some_and(is_pattern_cond),
        ast::Expr::LetExpr(_) => true,
        _ => false,
    }
}

/// Calls `cb` on each expression inside `expr` that is at "tail position".
/// Does not walk into `break` or `return` expressions.
/// Note that modifying the tree while iterating it will cause undefined iteration which might
/// potentially results in an out of bounds panic.
pub fn for_each_tail_expr(expr: &ast::Expr, cb: &mut dyn FnMut(&ast::Expr)) {
    let walk_loop = |cb: &mut dyn FnMut(&ast::Expr), label, body: Option<ast::BlockExpr>| {
        for_each_break_expr(label, body.and_then(|it| it.stmt_list()), &mut |b| {
            cb(&ast::Expr::BreakExpr(b))
        })
    };
    match expr {
        ast::Expr::BlockExpr(b) => {
            match b.modifier() {
                Some(
                    ast::BlockModifier::Async(_)
                    | ast::BlockModifier::Try(_)
                    | ast::BlockModifier::Const(_),
                ) => return cb(expr),

                Some(ast::BlockModifier::Label(label)) => {
                    for_each_break_expr(Some(label), b.stmt_list(), &mut |b| {
                        cb(&ast::Expr::BreakExpr(b))
                    });
                }
                Some(ast::BlockModifier::Unsafe(_)) => (),
                Some(ast::BlockModifier::Gen(_)) => (),
                Some(ast::BlockModifier::AsyncGen(_)) => (),
                None => (),
            }
            if let Some(stmt_list) = b.stmt_list()
                && let Some(e) = stmt_list.tail_expr()
            {
                for_each_tail_expr(&e, cb);
            }
        }
        ast::Expr::IfExpr(if_) => {
            let mut if_ = if_.clone();
            loop {
                if let Some(block) = if_.then_branch() {
                    for_each_tail_expr(&ast::Expr::BlockExpr(block), cb);
                }
                match if_.else_branch() {
                    Some(ast::ElseBranch::IfExpr(it)) => if_ = it,
                    Some(ast::ElseBranch::Block(block)) => {
                        for_each_tail_expr(&ast::Expr::BlockExpr(block), cb);
                        break;
                    }
                    None => break,
                }
            }
        }
        ast::Expr::LoopExpr(l) => walk_loop(cb, l.label(), l.loop_body()),
        ast::Expr::WhileExpr(w) => walk_loop(cb, w.label(), w.loop_body()),
        ast::Expr::ForExpr(f) => walk_loop(cb, f.label(), f.loop_body()),
        ast::Expr::MatchExpr(m) => {
            if let Some(arms) = m.match_arm_list() {
                arms.arms().filter_map(|arm| arm.expr()).for_each(|e| for_each_tail_expr(&e, cb));
            }
        }
        ast::Expr::ArrayExpr(_)
        | ast::Expr::AwaitExpr(_)
        | ast::Expr::BinExpr(_)
        | ast::Expr::BreakExpr(_)
        | ast::Expr::CallExpr(_)
        | ast::Expr::CastExpr(_)
        | ast::Expr::ClosureExpr(_)
        | ast::Expr::ContinueExpr(_)
        | ast::Expr::FieldExpr(_)
        | ast::Expr::IndexExpr(_)
        | ast::Expr::Literal(_)
        | ast::Expr::MacroExpr(_)
        | ast::Expr::MethodCallExpr(_)
        | ast::Expr::ParenExpr(_)
        | ast::Expr::PathExpr(_)
        | ast::Expr::PrefixExpr(_)
        | ast::Expr::RangeExpr(_)
        | ast::Expr::RecordExpr(_)
        | ast::Expr::RefExpr(_)
        | ast::Expr::ReturnExpr(_)
        | ast::Expr::BecomeExpr(_)
        | ast::Expr::TryExpr(_)
        | ast::Expr::TupleExpr(_)
        | ast::Expr::LetExpr(_)
        | ast::Expr::UnderscoreExpr(_)
        | ast::Expr::YieldExpr(_)
        | ast::Expr::YeetExpr(_)
        | ast::Expr::OffsetOfExpr(_)
        | ast::Expr::FormatArgsExpr(_)
        | ast::Expr::AsmExpr(_) => cb(expr),
    }
}

pub fn for_each_break_and_continue_expr(
    label: Option<ast::Label>,
    body: Option<ast::StmtList>,
    cb: &mut dyn FnMut(ast::Expr),
) {
    let label = label.and_then(|lbl| lbl.lifetime());
    if let Some(b) = body {
        let tree_depth_iterator = TreeWithDepthIterator::new(b);
        for (expr, depth) in tree_depth_iterator {
            match expr {
                ast::Expr::BreakExpr(b)
                    if (depth == 0 && b.lifetime().is_none())
                        || eq_label_lt(&label, &b.lifetime()) =>
                {
                    cb(ast::Expr::BreakExpr(b));
                }
                ast::Expr::ContinueExpr(c)
                    if (depth == 0 && c.lifetime().is_none())
                        || eq_label_lt(&label, &c.lifetime()) =>
                {
                    cb(ast::Expr::ContinueExpr(c));
                }
                _ => (),
            }
        }
    }
}

fn for_each_break_expr(
    label: Option<ast::Label>,
    body: Option<ast::StmtList>,
    cb: &mut dyn FnMut(ast::BreakExpr),
) {
    let label = label.and_then(|lbl| lbl.lifetime());
    if let Some(b) = body {
        let tree_depth_iterator = TreeWithDepthIterator::new(b);
        for (expr, depth) in tree_depth_iterator {
            match expr {
                ast::Expr::BreakExpr(b)
                    if (depth == 0 && b.lifetime().is_none())
                        || eq_label_lt(&label, &b.lifetime()) =>
                {
                    cb(b);
                }
                _ => (),
            }
        }
    }
}

pub fn eq_label_lt(lt1: &Option<ast::Lifetime>, lt2: &Option<ast::Lifetime>) -> bool {
    lt1.as_ref().zip(lt2.as_ref()).is_some_and(|(lt, lbl)| lt.text() == lbl.text())
}

/// Find the loop or block to break or continue, multiple results may be caused by macros.
pub fn find_loops(
    sema: &hir::Semantics<'_, crate::RootDatabase>,
    token: &syntax::SyntaxToken,
) -> Option<impl Iterator<Item = ast::Expr>> {
    let parent = token.parent()?;
    let lbl = syntax::match_ast! {
        match parent {
            ast::BreakExpr(break_) => break_.lifetime(),
            ast::ContinueExpr(continue_) => continue_.lifetime(),
            _ => None,
        }
    };
    let label_matches =
        move |it: Option<ast::Label>| match (lbl.as_ref(), it.and_then(|it| it.lifetime())) {
            (Some(lbl), Some(it)) => lbl.text() == it.text(),
            (None, _) => true,
            (Some(_), None) => false,
        };

    let find_ancestors = move |token| {
        for anc in sema.token_ancestors_with_macros(token).filter_map(ast::Expr::cast) {
            let node = match &anc {
                ast::Expr::LoopExpr(loop_) if label_matches(loop_.label()) => anc,
                ast::Expr::WhileExpr(while_) if label_matches(while_.label()) => anc,
                ast::Expr::ForExpr(for_) if label_matches(for_.label()) => anc,
                ast::Expr::BlockExpr(blk)
                    if blk.label().is_some() && label_matches(blk.label()) =>
                {
                    anc
                }
                _ => continue,
            };

            return Some(node);
        }
        None
    };

    sema.descend_into_macros(token.clone()).into_iter().filter_map(find_ancestors).into()
}

struct TreeWithDepthIterator {
    preorder: Preorder<RustLanguage>,
    depth: u32,
}

impl TreeWithDepthIterator {
    fn new(body: ast::StmtList) -> Self {
        let preorder = body.syntax().preorder();
        Self { preorder, depth: 0 }
    }
}

impl Iterator for TreeWithDepthIterator {
    type Item = (ast::Expr, u32);

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(event) = self.preorder.find_map(|ev| match ev {
            WalkEvent::Enter(it) => ast::Expr::cast(it).map(WalkEvent::Enter),
            WalkEvent::Leave(it) => ast::Expr::cast(it).map(WalkEvent::Leave),
        }) {
            match event {
                WalkEvent::Enter(
                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_),
                ) => {
                    self.depth += 1;
                }
                WalkEvent::Leave(
                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_),
                ) => {
                    self.depth -= 1;
                }
                WalkEvent::Enter(ast::Expr::BlockExpr(e)) if e.label().is_some() => {
                    self.depth += 1;
                }
                WalkEvent::Leave(ast::Expr::BlockExpr(e)) if e.label().is_some() => {
                    self.depth -= 1;
                }
                WalkEvent::Enter(expr) => return Some((expr, self.depth)),
                _ => (),
            }
        }
        None
    }
}

/// Parses the input token tree as comma separated plain paths.
pub fn parse_tt_as_comma_sep_paths(
    input: ast::TokenTree,
    edition: Edition,
) -> Option<Vec<ast::Path>> {
    let r_paren = input.r_paren_token();
    let tokens =
        input.syntax().children_with_tokens().skip(1).map_while(|it| match it.into_token() {
            // seeing a keyword means the attribute is unclosed so stop parsing here
            Some(tok) if tok.kind().is_keyword(edition) => None,
            // don't include the right token tree parenthesis if it exists
            tok @ Some(_) if tok == r_paren => None,
            // only nodes that we can find are other TokenTrees, those are unexpected in this parse though
            None => None,
            Some(tok) => Some(tok),
        });
    let input_expressions = tokens.chunk_by(|tok| tok.kind() == T![,]);
    let paths = input_expressions
        .into_iter()
        .filter_map(|(is_sep, group)| (!is_sep).then_some(group))
        .filter_map(|mut tokens| {
            syntax::hacks::parse_expr_from_str(&tokens.join(""), Edition::CURRENT).and_then(
                |expr| match expr {
                    ast::Expr::PathExpr(it) => it.path(),
                    _ => None,
                },
            )
        })
        .collect();
    Some(paths)
}

pub fn macro_call_for_string_token(string: &ast::String) -> Option<MacroCall> {
    let macro_call = string.syntax().parent_ancestors().find_map(ast::MacroCall::cast)?;
    Some(macro_call)
}

pub fn is_in_macro_matcher(token: &SyntaxToken) -> bool {
    let Some(macro_def) = token
        .parent_ancestors()
        .map_while(Either::<ast::TokenTree, ast::Macro>::cast)
        .find_map(Either::right)
    else {
        return false;
    };
    let range = token.text_range();
    let Some(body) = (match macro_def {
        ast::Macro::MacroDef(macro_def) => {
            if let Some(args) = macro_def.args() {
                return args.syntax().text_range().contains_range(range);
            }
            macro_def.body()
        }
        ast::Macro::MacroRules(macro_rules) => macro_rules.token_tree(),
    }) else {
        return false;
    };
    if !body.syntax().text_range().contains_range(range) {
        return false;
    }
    body.token_trees_and_tokens().filter_map(|tt| tt.into_node()).any(|tt| {
        let Some(next) = non_trivia_sibling(tt.syntax().syntax_element(), Direction::Next) else {
            return false;
        };
        let Some(next_next) = next.next_sibling_or_token() else { return false };
        next.kind() == T![=]
            && next_next.kind() == T![>]
            && tt.syntax().text_range().contains_range(range)
    })
}