rust/crates/ide_assists/src/handlers/merge_match_arms.rs

use itertools::Itertools;
use std::iter::successors;
use syntax::{
    algo::neighbor,
    ast::{self, AstNode},
    Direction,
};

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

// Assist: merge_match_arms
//
// Merges the current match arm with the following if their bodies are identical.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         $0Action::Move(..) => foo(),
//         Action::Stop => foo(),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move(..) | Action::Stop => foo(),
//     }
// }
// ```
pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
    // Don't try to handle arms with guards for now - can add support for this later
    if current_arm.guard().is_some() {
        return None;
    }
    let current_expr = current_arm.expr()?;
    let current_text_range = current_arm.syntax().text_range();
    let current_arm_types = get_arm_types(&ctx, &current_arm);

    // We check if the following match arms match this one. We could, but don't,
    // compare to the previous match arm as well.
    let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
        .take_while(|arm| match arm.expr() {
            Some(expr) if arm.guard().is_none() => {
                let same_text = expr.syntax().text() == current_expr.syntax().text();
                if !same_text {
                    return false;
                }

                return are_same_types(&current_arm_types, arm, ctx);
            }
            _ => false,
        })
        .collect::<Vec<_>>();

    if arms_to_merge.len() <= 1 {
        return None;
    }

    acc.add(
        AssistId("merge_match_arms", AssistKind::RefactorRewrite),
        "Merge match arms",
        current_text_range,
        |edit| {
            let pats = if arms_to_merge.iter().any(contains_placeholder) {
                "_".into()
            } else {
                arms_to_merge
                    .iter()
                    .filter_map(ast::MatchArm::pat)
                    .map(|x| x.syntax().to_string())
                    .collect::<Vec<String>>()
                    .join(" | ")
            };

            let arm = format!("{} => {},", pats, current_expr.syntax().text());

            if let [first, .., last] = &*arms_to_merge {
                let start = first.syntax().text_range().start();
                let end = last.syntax().text_range().end();

                edit.replace(TextRange::new(start, end), arm);
            }
        },
    )
}

fn contains_placeholder(a: &ast::MatchArm) -> bool {
    matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
}

fn are_same_types(
    current_arm_types: &Vec<Option<hir::TypeInfo>>,
    arm: &ast::MatchArm,
    ctx: &AssistContext,
) -> bool {
    let arm_types = get_arm_types(&ctx, &arm);
    for i in 0..arm_types.len() {
        let other_arm_type = &arm_types[i];
        let current_arm_type = &current_arm_types[i];
        if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
            return &other_arm_type.original == &current_arm_type.original;
        }
    }

    return true;
}

fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
    match arm.pat() {
        Some(ast::Pat::TupleStructPat(tp)) => {
            tp.fields().into_iter().map(|field| ctx.sema.type_of_pat(&field)).collect_vec()
        }
        _ => Vec::new(),
    }
}

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

    use super::*;

    #[test]
    fn merge_match_arms_single_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32$0 }
        X::B => { 1i32 }
        X::C => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_multiple_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => {$0 1i32 },
        X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B | X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_placeholder_pattern() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        X::B => { 2i$032 },
        _ => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        _ => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merges_all_subsequent_arms() {
        check_assist(
            merge_match_arms,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A$0 => 92,
        X::B => 92,
        X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A | X::B | X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
        )
    }

    #[test]
    fn merge_match_arms_rejects_guards() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X {
    A(i32),
    B,
    C
}

fn main() {
    let x = X::A;
    let y = match x {
        X::A(a) if a > 5 => { $01i32 },
        X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<f64, f32>::Ok(0f64) {
        Ok(x) => $0x.classify(),
        Err(x) => x.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type_multiple_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_multiple_fields() {
        check_assist(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
            r#"
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) | Err(x) => x.1.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f64),
    OptionC(f32)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_different_number_of_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64, f64), (f64, f64)>::Ok((0f64, 0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_same_destructuring_different_types() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
struct Point {
    x: i32,
    y: i32,
}

fn func() {
    let p = Point { x: 0, y: 7 };

    match p {
        Point { x, y: 0 } => $0"",
        Point { x: 0, y } => "",
        Point { x, y } => "",
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_range() {
        check_assist(
            merge_match_arms,
            r#"
let x = 'c';

    match x {
        'a'..='j' => $0"",
        'c'..='z' => "",
        _ => "other",
    };
"#,
            r#"
let x = 'c';

    match x {
        'a'..='j' | 'c'..='z' => "",
        _ => "other",
    };
"#,
        );
    }

    #[test]
    fn merge_match_arms_enum_without_field() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    NoField,
    AField(u8)
}

fn func(x: MyEnum) {
    match x {
        MyEnum::NoField => $0"",
        MyEnum::AField(x) => ""
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_different_types() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Write(String),
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Write(text) => "",
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_same_types() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { x: i32, y: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Crawl { x, y } => "",
    };
}
        "#,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { x: i32, y: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } | MyEnum::Crawl { x, y } => "",
    };
}
        "#,        
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_same_types_different_name() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { a: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Crawl { a, b } => "",
    };
}
        "#       
        )
    }

    #[test]
    fn merge_match_arms_enum_nested_pattern_different_names() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(r, g, b)) => $0"",
        Message::ChangeColor(Color::Hsv(h, s, v)) => "",
        _ => "other"
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_nested_pattern_same_names() {
        check_assist(
            merge_match_arms,
            r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(a, b, c)) => $0"",
        Message::ChangeColor(Color::Hsv(a, b, c)) => "",
        _ => "other"
    };
}
        "#,
        r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(a, b, c)) | Message::ChangeColor(Color::Hsv(a, b, c)) => "",
        _ => "other"
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_with_ignore() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, a: i32 },
    Crawl { x: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, .. } => $0"",
        MyEnum::Crawl { x, .. } => "",
    };
}
        "#,
            r#"
enum MyEnum {
    Move { x: i32, a: i32 },
    Crawl { x: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, .. } | MyEnum::Crawl { x, .. } => "",
    };
}
        "#,        
        )
    }
}