fix(core): single match properly handles comments

crates/cairo-lint-core/src/fix.rs

@@ -8,6 +8,7 @@ use cairo_lang_syntax::node::{SyntaxNode, TypedSyntaxNode};
 use cairo_lang_utils::Upcast;
 
 use crate::db::AnalysisDatabase;
+use crate::lints::single_match::is_expr_unit;
 use crate::plugin::{diagnostic_kind_from_message, CairoLintKind};
 
 #[derive(Default)]
@@ -36,24 +37,27 @@ impl Fixer {
         let second_arm = &arms[1];
         let (pattern, first_expr) =
             match (&first_arm.patterns(db).elements(db)[0], &second_arm.patterns(db).elements(db)[0]) {
-                (Pattern::Underscore(_), Pattern::Enum(pat)) => {
-                    (&pat.as_syntax_node().get_text_without_trivia(db), second_arm)
-                }
-                (Pattern::Enum(pat), Pattern::Underscore(_)) => {
-                    (&pat.as_syntax_node().get_text_without_trivia(db), first_arm)
-                }
-                (Pattern::Underscore(_), Pattern::Struct(pat)) => {
-                    (&pat.as_syntax_node().get_text_without_trivia(db), second_arm)
-                }
-                (Pattern::Struct(pat), Pattern::Underscore(_)) => {
-                    (&pat.as_syntax_node().get_text_without_trivia(db), first_arm)
+                (Pattern::Underscore(_), Pattern::Enum(pat)) => (pat.as_syntax_node(), second_arm),
+                (Pattern::Enum(pat), Pattern::Underscore(_)) => (pat.as_syntax_node(), first_arm),
+                (Pattern::Underscore(_), Pattern::Struct(pat)) => (pat.as_syntax_node(), second_arm),
+                (Pattern::Struct(pat), Pattern::Underscore(_)) => (pat.as_syntax_node(), first_arm),
+                (Pattern::Enum(pat1), Pattern::Enum(pat2)) => {
+                    if is_expr_unit(second_arm.expression(db), db) {
+                        (pat1.as_syntax_node(), first_arm)
+                    } else {
+                        (pat2.as_syntax_node(), second_arm)
+                    }
                 }
                 (_, _) => panic!("Incorrect diagnostic"),
             };
+        let mut pattern_span = pattern.span(db);
+        pattern_span.end = pattern.span_start_without_trivia(db);
+        let indent = node.get_text(db).chars().take_while(|c| c.is_whitespace()).collect::<String>();
+        let trivia = pattern.clone().get_text_of_span(db, pattern_span).trim().to_string();
+        let trivia = if trivia.is_empty() { trivia } else { format!("{indent}{trivia}\n") };
         format!(
-            "{}if let {} = {} {{ {} }}",
-            node.get_text(db).chars().take_while(|c| c.is_whitespace()).collect::<String>(),
-            pattern,
+            "{trivia}{indent}if let {} = {} {{ {} }}",
+            pattern.get_text_without_trivia(db),
             match_expr.expr(db).as_syntax_node().get_text_without_trivia(db),
             first_expr.expression(db).as_syntax_node().get_text_without_trivia(db),
         )

crates/cairo-lint-core/src/lib.rs

@@ -1,4 +1,4 @@
-#![feature(let_chains)]
+#![feature(let_chains, is_none_or)]
 pub mod db;
 pub mod fix;
 pub mod lints;

crates/cairo-lint-core/src/lints/single_match.rs

@@ -1,55 +1,100 @@
+use cairo_lang_defs::ids::{FileIndex, ModuleFileId, ModuleId};
 use cairo_lang_defs::plugin::PluginDiagnostic;
-use cairo_lang_diagnostics::Severity;
+use cairo_lang_diagnostics::{DiagnosticsBuilder, Severity};
+use cairo_lang_semantic::db::SemanticGroup;
+use cairo_lang_semantic::diagnostic::NotFoundItemType;
+use cairo_lang_semantic::expr::inference::InferenceId;
+use cairo_lang_semantic::resolve::{AsSegments, ResolvedGenericItem, Resolver};
 use cairo_lang_syntax::node::ast::{Expr, ExprBlock, ExprListParenthesized, ExprMatch, Pattern, Statement};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::{TypedStablePtr, TypedSyntaxNode};
+use cairo_lang_utils::try_extract_matches;
 
 pub const DESTRUCT_MATCH: &str =
     "you seem to be trying to use `match` for destructuring a single pattern. Consider using `if let`";
 pub const MATCH_FOR_EQUALITY: &str = "you seem to be trying to use `match` for an equality check. Consider using `if`";
 
-fn tuple_expr_in_block_expr(
-    db: &dyn SyntaxGroup,
-    block_expr: &ExprBlock,
-    is_single_armed: &mut bool,
-) -> Option<ExprListParenthesized> {
+fn is_expr_list_parenthesised_unit(expr: &ExprListParenthesized, db: &dyn SyntaxGroup) -> bool {
+    expr.expressions(db).elements(db).is_empty()
+}
+
+fn is_block_expr_unit_without_comment(block_expr: &ExprBlock, db: &dyn SyntaxGroup) -> bool {
     let statements = block_expr.statements(db).elements(db);
     if statements.is_empty() {
-        *is_single_armed = true;
+        return true;
     }
     if statements.len() == 1
         && let Statement::Expr(statement_expr) = &statements[0]
         && let Expr::Tuple(tuple_expr) = statement_expr.expr(db)
     {
-        Some(tuple_expr)
+        let tuple_node = tuple_expr.as_syntax_node();
+        if tuple_node.span(db).start != tuple_node.span_start_without_trivia(db) {
+            return false;
+        }
+        is_expr_list_parenthesised_unit(&tuple_expr, db)
     } else {
-        None
+        false
     }
 }
 
-pub fn check_single_match(db: &dyn SyntaxGroup, match_expr: &ExprMatch, diagnostics: &mut Vec<PluginDiagnostic>) {
-    let arms = match_expr.arms(db).elements(db);
+pub fn is_expr_unit(expr: Expr, db: &dyn SyntaxGroup) -> bool {
+    match expr {
+        Expr::Block(block_expr) => is_block_expr_unit_without_comment(&block_expr, db),
+        Expr::Tuple(tuple_expr) => is_expr_list_parenthesised_unit(&tuple_expr, db),
+        _ => false,
+    }
+}
+pub fn check_single_match(
+    db: &dyn SemanticGroup,
+    match_expr: &ExprMatch,
+    diagnostics: &mut Vec<PluginDiagnostic>,
+    module_id: &ModuleId,
+) {
+    let syntax_db = db.upcast();
+    let arms = match_expr.arms(syntax_db).elements(syntax_db);
     let mut is_single_armed = false;
+    let mut is_complete = false;
     let mut is_destructuring = false;
     if arms.len() == 2 {
-        for arm in arms {
-            let patterns = arm.patterns(db).elements(db);
-            match &patterns[0] {
-                Pattern::Underscore(_) => {
-                    let tuple_expr = match arm.expression(db) {
-                        Expr::Block(block_expr) => tuple_expr_in_block_expr(db, &block_expr, &mut is_single_armed),
-                        Expr::Tuple(tuple_expr) => Some(tuple_expr),
-                        _ => None,
-                    };
-                    is_single_armed =
-                        tuple_expr.is_some_and(|list| list.expressions(db).elements(db).is_empty()) || is_single_armed;
+        let first_arm = &arms[0];
+        let second_arm = &arms[1];
+        let mut enum_len = None;
+        let mut is_first_arm_unit = false;
+        if let Some(pattern) = first_arm.patterns(syntax_db).elements(syntax_db).first() {
+            match pattern {
+                Pattern::Underscore(_) => return,
+                Pattern::Enum(pat) => {
+                    let mut diagnostics = DiagnosticsBuilder::default();
+                    let path = pat.path(syntax_db).to_segments(syntax_db);
+                    let item = Resolver::new(db, ModuleFileId(*module_id, FileIndex(0)), InferenceId::NoContext)
+                        .resolve_generic_path(&mut diagnostics, path, NotFoundItemType::Identifier)
+                        .unwrap();
+                    let generic_variant = try_extract_matches!(item, ResolvedGenericItem::Variant).unwrap();
+                    enum_len = Some(db.enum_variants(generic_variant.enum_id).unwrap().len());
+                    is_destructuring = true;
                 }
-                Pattern::Enum(_) | Pattern::Struct(_) => {
+                Pattern::Struct(_) => {
                     is_destructuring = true;
                 }
                 _ => (),
-            };
-        }
+            }
+            is_first_arm_unit = is_expr_unit(first_arm.expression(syntax_db), syntax_db)
+        };
+        if let Some(pattern) = second_arm.patterns(syntax_db).elements(syntax_db).first() {
+            match pattern {
+                Pattern::Underscore(_) => {
+                    is_complete = true;
+                }
+                Pattern::Enum(_) => {
+                    if enum_len == Some(2) {
+                        is_complete = true;
+                    }
+                }
+                _ => (),
+            }
+            is_single_armed =
+                is_expr_unit(second_arm.expression(syntax_db), syntax_db) && is_complete || is_first_arm_unit;
+        };
     };
     match (is_single_armed, is_destructuring) {
         (true, false) => diagnostics.push(PluginDiagnostic {

crates/cairo-lint-core/src/plugin.rs

@@ -49,6 +49,7 @@ impl AnalyzerPlugin for CairoLint {
                                 db.upcast(),
                                 &ExprMatch::from_syntax_node(db.upcast(), descendant),
                                 &mut diags,
+                                &module_id,
                             ),
                             SyntaxKind::ItemExternFunction => (),
                             _ => (),