Type Assists 4/n: mismatched atom, fix spec

Summary: Finish the mismatched return atom assist, by adding the option to update the spec to match the return value.

Reviewed By: robertoaloi

Differential Revision: D55019579

fbshipit-source-id: 9396e25381a62fbeacba8081dd2c60881ae39de5

crates/hir/src/body.rs

@@ -916,6 +916,10 @@ impl BodySourceMap {
             .copied()
     }
 
+    pub fn type_expr(&self, type_expr_id: TypeExprId) -> Option<ExprSource> {
+        self.type_expr_map_back.get(type_expr_id).copied()
+    }
+
     pub fn term_id(&self, expr: InFile<&ast::Expr>) -> Option<TermId> {
         self.term_map.get(&InFileAstPtr::from_infile(expr)).copied()
     }

crates/ide/src/diagnostics/eqwalizer_assists/expected_type.rs

@@ -8,20 +8,25 @@
  */
 
 use elp_ide_db::elp_base_db::FileId;
+use elp_ide_db::elp_base_db::FilePosition;
+use elp_ide_db::find_best_token;
 use elp_ide_db::source_change::SourceChange;
 use elp_ide_db::EqwalizerDiagnostic;
 use elp_types_db::eqwalizer::tc_diagnostics::ExpectedSubtype;
 use elp_types_db::eqwalizer::tc_diagnostics::TypeError;
 use elp_types_db::eqwalizer::types::Type;
 use elp_types_db::eqwalizer::StructuredDiagnostic;
+use hir::InFile;
+use hir::Literal;
 use hir::Semantic;
+use hir::TypeExpr;
 use text_edit::TextEdit;
 
 use crate::diagnostics::Diagnostic;
 use crate::fix;
 
 pub fn expected_type(
-    _sema: &Semantic,
+    sema: &Semantic,
     file_id: FileId,
     d: &EqwalizerDiagnostic,
     diagnostic: &mut Diagnostic,
@@ -43,16 +48,60 @@ pub fn expected_type(
                     SourceChange::from_text_edit(file_id, edit),
                     d.range,
                 ));
+                add_spec_fix(sema, file_id, got, diagnostic);
             }
             _ => {}
         }
     }
 }
 
+fn add_spec_fix(
+    sema: &Semantic,
+    file_id: FileId,
+    got: &Type,
+    diagnostic: &mut Diagnostic,
+) -> Option<()> {
+    let token = find_best_token(
+        sema,
+        FilePosition {
+            file_id,
+            offset: diagnostic.range.start(),
+        },
+    )?;
+    let function_id = sema.find_enclosing_function(file_id, &token.value.parent()?)?;
+    let def_map = sema.def_map(file_id);
+    let function_def = def_map.get_by_function_id(&InFile::new(file_id, function_id))?;
+    let spec_id = InFile::new(file_id, function_def.spec.clone()?.spec_id);
+    let spec_body = sema.db.spec_body(spec_id);
+    // We are looking for a single spec signature with a single atom return type.
+    match &spec_body.sigs[..] {
+        [sig] => {
+            match spec_body.body[sig.result] {
+                TypeExpr::Literal(Literal::Atom(_)) => {
+                    // We have a single atom. Make an edit to replace it.
+                    let (_, body_map) = sema.db.spec_body_with_source(spec_id);
+                    let source = body_map.type_expr(sig.result)?;
+                    let range = source.range();
+                    let edit = TextEdit::replace(range, format!("{got}").to_string());
+                    diagnostic.add_fix(fix(
+                        "fix_expected_type",
+                        format!("Change spec atom to '{got}'").as_str(),
+                        SourceChange::from_text_edit(file_id, edit),
+                        diagnostic.range,
+                    ));
+                }
+                _ => {}
+            }
+        }
+        _ => {}
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::check_diagnostics;
-    use crate::tests::check_fix;
+    use crate::tests::check_specific_fix;
 
     #[test]
     fn mismatched_atom() {
@@ -70,8 +119,9 @@ mod tests {
     }
 
     #[test]
-    fn mismatched_atom_fix() {
-        check_fix(
+    fn mismatched_atom_fix_return() {
+        check_specific_fix(
+            "Change returned atom to 'spec_atom'",
             r#"
             //- eqwalizer
             //- /play/src/bar.erl app:play
@@ -89,4 +139,26 @@ mod tests {
          "#,
         )
     }
+
+    #[test]
+    fn mismatched_atom_fix_spec() {
+        check_specific_fix(
+            "Change spec atom to 'something_else'",
+            r#"
+            //- eqwalizer
+            //- /play/src/bar.erl app:play
+            -module(bar).
+
+            -spec baz() -> spec_atom.
+            baz() -> somethin~g_else.
+                  %% ^^^^^^^^^^^^^^ 💡 error: eqwalizer: incompatible_types
+            "#,
+            r#"
+            -module(bar).
+
+            -spec baz() -> something_else.
+            baz() -> something_else.
+         "#,
+        )
+    }
 }