Added tests, added Union Support, and code cleanup

crates/ide-diagnostics/src/handlers/unresolved_field.rs

@@ -1,4 +1,6 @@
-use hir::{db::ExpandDatabase, Adt, HasSource, HirDisplay, InFile};
+use std::iter;
+
+use hir::{db::ExpandDatabase, Adt, HasSource, HirDisplay, InFile, Struct, Union};
 use ide_db::{
     assists::{Assist, AssistId, AssistKind},
     base_db::FileRange,
@@ -7,8 +9,13 @@ use ide_db::{
     source_change::{SourceChange, SourceChangeBuilder},
 };
 use syntax::{
-    ast::{self, edit::IndentLevel, make},
-    AstNode, AstPtr, SyntaxKind,
+    algo,
+    ast::{self, edit::IndentLevel, make, FieldList, Name, Visibility},
+    AstNode, AstPtr, Direction, SyntaxKind, TextSize,
+};
+use syntax::{
+    ast::{edit::AstNodeEdit, Type},
+    SyntaxNode,
 };
 use text_edit::TextEdit;
 
@@ -52,12 +59,12 @@ pub(crate) fn unresolved_field(
 fn fixes(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedField) -> Option<Vec<Assist>> {
     let mut fixes = if d.method_with_same_name_exists { method_fix(ctx, &d.expr) } else { None };
     if let Some(fix) = add_field_fix(ctx, d) {
-        fixes.get_or_insert_with(Vec::new).push(fix);
+        fixes.get_or_insert_with(Vec::new).extend(fix);
     }
     fixes
 }
 
-fn add_field_fix(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedField) -> Option<Assist> {
+fn add_field_fix(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedField) -> Option<Vec<Assist>> {
     // Get the FileRange of the invalid field access
     let root = ctx.sema.db.parse_or_expand(d.expr.file_id);
     let expr = d.expr.value.to_node(&root);
@@ -65,10 +72,6 @@ fn add_field_fix(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedField) -> Opti
     let error_range = ctx.sema.original_range_opt(expr.syntax())?;
     // Convert the receiver to an ADT
     let adt = d.receiver.as_adt()?;
-    let Adt::Struct(adt) = adt else {
-        return None;
-    };
-
     let target_module = adt.module(ctx.sema.db);
 
     let suggested_type =
@@ -83,54 +86,161 @@ fn add_field_fix(ctx: &DiagnosticsContext<'_>, d: &hir::UnresolvedField) -> Opti
     if !is_editable_crate(target_module.krate(), ctx.sema.db) {
         return None;
     }
-    let adt_source = adt.source(ctx.sema.db)?;
+
+    // FIXME: Add Snippet Support
+    let field_name = d.name.as_str()?;
+
+    match adt {
+        Adt::Struct(adt_struct) => {
+            add_field_to_struct_fix(ctx, adt_struct, field_name, suggested_type, error_range)
+        }
+        Adt::Union(adt_union) => {
+            add_varient_to_union(ctx, adt_union, field_name, suggested_type, error_range)
+        }
+        _ => None,
+    }
+}
+fn add_varient_to_union(
+    ctx: &DiagnosticsContext<'_>,
+    adt_union: Union,
+    field_name: &str,
+    suggested_type: Type,
+    error_range: FileRange,
+) -> Option<Vec<Assist>> {
+    let adt_source = adt_union.source(ctx.sema.db)?;
     let adt_syntax = adt_source.syntax();
-    let range = adt_syntax.original_file_range(ctx.sema.db);
-
-    // Get range of final field in the struct
-    let (offset, needs_comma, indent) = match adt.fields(ctx.sema.db).last() {
-        Some(field) => {
-            let last_field = field.source(ctx.sema.db)?.value;
-            let hir::FieldSource::Named(record_field) = last_field else {
-                return None;
-            };
-            let last_field_syntax = record_field.syntax();
-            let last_field_imdent = IndentLevel::from_node(last_field_syntax);
-            (
-                last_field_syntax.text_range().end(),
-                !last_field_syntax.to_string().ends_with(','),
-                last_field_imdent,
-            )
-        }
-        None => {
-            // Empty Struct. Add a field right before the closing brace
-            let indent = IndentLevel::from_node(adt_syntax.value) + 1;
-            let record_field_list =
-                adt_syntax.value.children().find(|v| v.kind() == SyntaxKind::RECORD_FIELD_LIST)?;
-            let offset = record_field_list.first_token().map(|f| f.text_range().end())?;
-            (offset, false, indent)
-        }
+    let Some(field_list) = adt_source.value.record_field_list() else {
+        return None;
     };
+    let range = adt_syntax.original_file_range(ctx.sema.db);
+    let field_name = make::name(field_name);
 
-    let field_name = make::name(d.name.as_str()?);
+    let (offset, record_field) =
+        record_field_layout(None, field_name, suggested_type, field_list, adt_syntax.value)?;
 
-    // If the Type is in the same file. We don't need to add a visibility modifier. Otherwise make it pub(crate)
-    let visibility = if error_range.file_id == range.file_id { "" } else { "pub(crate)" };
     let mut src_change_builder = SourceChangeBuilder::new(range.file_id);
-    let comma = if needs_comma { "," } else { "" };
-    src_change_builder
-        .insert(offset, format!("{comma}\n{indent}{visibility}{field_name}: {suggested_type}\n"));
-
-    // FIXME: Add a Snippet for the new field type
-    let source_change = src_change_builder.finish();
-    Some(Assist {
-        id: AssistId("add-field-to-type", AssistKind::QuickFix),
-        label: Label::new("Add field to type".to_owned()),
+    src_change_builder.insert(offset, record_field);
+    Some(vec![Assist {
+        id: AssistId("add-varient-to-union", AssistKind::QuickFix),
+        label: Label::new("Add field to union".to_owned()),
         group: None,
         target: error_range.range,
-        source_change: Some(source_change),
+        source_change: Some(src_change_builder.finish()),
         trigger_signature_help: false,
-    })
+    }])
+}
+fn add_field_to_struct_fix(
+    ctx: &DiagnosticsContext<'_>,
+    adt_struct: Struct,
+    field_name: &str,
+    suggested_type: Type,
+    error_range: FileRange,
+) -> Option<Vec<Assist>> {
+    let struct_source = adt_struct.source(ctx.sema.db)?;
+    let struct_syntax = struct_source.syntax();
+    let struct_range = struct_syntax.original_file_range(ctx.sema.db);
+    let field_list = struct_source.value.field_list();
+    match field_list {
+        Some(FieldList::RecordFieldList(field_list)) => {
+            // Get range of final field in the struct
+            let visibility = if error_range.file_id == struct_range.file_id {
+                None
+            } else {
+                Some(make::visibility_pub_crate())
+            };
+            let field_name = make::name(field_name);
+
+            let (offset, record_field) = record_field_layout(
+                visibility,
+                field_name,
+                suggested_type,
+                field_list,
+                struct_syntax.value,
+            )?;
+
+            let mut src_change_builder = SourceChangeBuilder::new(struct_range.file_id);
+
+            // FIXME: Allow for choosing a visibility modifier see https://github.com/rust-lang/rust-analyzer/issues/11563
+            src_change_builder.insert(offset, record_field);
+            Some(vec![Assist {
+                id: AssistId("add-field-to-record-struct", AssistKind::QuickFix),
+                label: Label::new("Add field to Record Struct".to_owned()),
+                group: None,
+                target: error_range.range,
+                source_change: Some(src_change_builder.finish()),
+                trigger_signature_help: false,
+            }])
+        }
+        None => {
+            // Add a field list to the Unit Struct
+            let mut src_change_builder = SourceChangeBuilder::new(struct_range.file_id);
+            let field_name = make::name(field_name);
+            let visibility = if error_range.file_id == struct_range.file_id {
+                None
+            } else {
+                Some(make::visibility_pub_crate())
+            };
+            // FIXME: Allow for choosing a visibility modifier see https://github.com/rust-lang/rust-analyzer/issues/11563
+            let indent = IndentLevel::from_node(struct_syntax.value) + 1;
+
+            let field = make::record_field(visibility, field_name, suggested_type).indent(indent);
+            let record_field_list = make::record_field_list(iter::once(field));
+            // A Unit Struct with no `;` is invalid syntax. We should not suggest this fix.
+            let semi_colon =
+                algo::skip_trivia_token(struct_syntax.value.last_token()?, Direction::Prev)?;
+            if semi_colon.kind() != SyntaxKind::SEMICOLON {
+                return None;
+            }
+            src_change_builder.replace(semi_colon.text_range(), record_field_list.to_string());
+
+            Some(vec![Assist {
+                id: AssistId("convert-unit-struct-to-record-struct", AssistKind::QuickFix),
+                label: Label::new("Convert Unit Struct to Record Struct and add field".to_owned()),
+                group: None,
+                target: error_range.range,
+                source_change: Some(src_change_builder.finish()),
+                trigger_signature_help: false,
+            }])
+        }
+        Some(FieldList::TupleFieldList(_tuple)) => {
+            // FIXME: Add support for Tuple Structs. Tuple Structs are not sent to this diagnostic
+            None
+        }
+    }
+}
+/// Used to determine the layout of the record field in the struct.
+fn record_field_layout(
+    visibility: Option<Visibility>,
+    name: Name,
+    suggested_type: Type,
+    field_list: ast::RecordFieldList,
+    struct_syntax: &SyntaxNode,
+) -> Option<(TextSize, String)> {
+    let (offset, needs_comma, trailing_new_line, indent) = match field_list.fields().last() {
+        Some(record_field) => {
+            let syntax = algo::skip_trivia_token(field_list.r_curly_token()?, Direction::Prev)?;
+
+            let last_field_syntax = record_field.syntax();
+            let last_field_indent = IndentLevel::from_node(last_field_syntax);
+            (
+                last_field_syntax.text_range().end(),
+                syntax.kind() != SyntaxKind::COMMA,
+                false,
+                last_field_indent,
+            )
+        }
+        // Empty Struct. Add a field right before the closing brace
+        None => {
+            let indent = IndentLevel::from_node(struct_syntax) + 1;
+            let offset = field_list.r_curly_token()?.text_range().start();
+            (offset, false, true, indent)
+        }
+    };
+    let comma = if needs_comma { ",\n" } else { "" };
+    let trailing_new_line = if trailing_new_line { "\n" } else { "" };
+    let record_field = make::record_field(visibility, name, suggested_type);
+
+    Some((offset, format!("{comma}{indent}{record_field}{trailing_new_line}")))
 }
 // FIXME: We should fill out the call here, move the cursor and trigger signature help
 fn method_fix(
@@ -154,9 +264,11 @@ fn method_fix(
 }
 #[cfg(test)]
 mod tests {
+
     use crate::{
         tests::{
             check_diagnostics, check_diagnostics_with_config, check_diagnostics_with_disabled,
+            check_fix,
         },
         DiagnosticsConfig,
     };
@@ -245,4 +357,100 @@ fn foo() {
         config.disabled.insert("syntax-error".to_owned());
         check_diagnostics_with_config(config, "fn foo() { (). }");
     }
+
+    #[test]
+    fn unresolved_field_fix_on_unit() {
+        check_fix(
+            r#"
+                struct Foo;
+
+                fn foo() {
+                    Foo.bar$0;
+                }
+            "#,
+            r#"
+                struct Foo{ bar: () }
+
+                fn foo() {
+                    Foo.bar;
+                }
+            "#,
+        );
+    }
+    #[test]
+    fn unresolved_field_fix_on_empty() {
+        check_fix(
+            r#"
+                struct Foo{
+                }
+
+                fn foo() {
+                    let foo = Foo{};
+                    foo.bar$0;
+                }
+            "#,
+            r#"
+                struct Foo{
+                    bar: ()
+                }
+
+                fn foo() {
+                    let foo = Foo{};
+                    foo.bar;
+                }
+            "#,
+        );
+    }
+    #[test]
+    fn unresolved_field_fix_on_struct() {
+        check_fix(
+            r#"
+                struct Foo{
+                    a: i32
+                }
+
+                fn foo() {
+                    let foo = Foo{a: 0};
+                    foo.bar$0;
+                }
+            "#,
+            r#"
+                struct Foo{
+                    a: i32,
+                    bar: ()
+                }
+
+                fn foo() {
+                    let foo = Foo{a: 0};
+                    foo.bar;
+                }
+            "#,
+        );
+    }
+    #[test]
+    fn unresolved_field_fix_on_union() {
+        check_fix(
+            r#"
+                union Foo{
+                    a: i32
+                }
+
+                fn foo() {
+                    let foo = Foo{a: 0};
+                    foo.bar$0;
+                }
+            "#,
+            r#"
+                union Foo{
+                    a: i32,
+                    bar: ()
+                }
+
+                fn foo() {
+                    let foo = Foo{a: 0};
+                    foo.bar;
+                }
+            "#,
+        );
+    }
 }