Rollup merge of #90575 - m-ou-se:compatible-variant-improvements, r=estebank

Improve suggestions for compatible variants on type mismatch.

Fixes #90553.

Before:
![image](https://user-images.githubusercontent.com/783247/140385675-6ff41090-eca2-41bc-b161-99c5dabfec61.png)

After:
![image](https://user-images.githubusercontent.com/783247/140385748-20cf26b5-ea96-4e56-8af2-5fe1ab16fd3b.png)

r? `````@estebank`````

compiler/rustc_span/src/source_map.rs

@@ -593,14 +593,19 @@ impl SourceMap {
     }
 
     pub fn span_to_margin(&self, sp: Span) -> Option<usize> {
-        match self.span_to_prev_source(sp) {
-            Err(_) => None,
-            Ok(source) => {
-                let last_line = source.rsplit_once('\n').unwrap_or(("", &source)).1;
+        Some(self.indentation_before(sp)?.len())
+    }
 
-                Some(last_line.len() - last_line.trim_start().len())
-            }
-        }
+    pub fn indentation_before(&self, sp: Span) -> Option<String> {
+        self.span_to_source(sp, |src, start_index, _| {
+            let before = &src[..start_index];
+            let last_line = before.rsplit_once('\n').map_or(before, |(_, last)| last);
+            Ok(last_line
+                .split_once(|c: char| !c.is_whitespace())
+                .map_or(last_line, |(indent, _)| indent)
+                .to_string())
+        })
+        .ok()
     }
 
     /// Returns the source snippet as `String` before the given `Span`.

compiler/rustc_typeck/src/check/demand.rs

@@ -199,7 +199,50 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 return;
             }
 
-            let mut compatible_variants = expected_adt
+            // If the expression is of type () and it's the return expression of a block,
+            // we suggest adding a separate return expression instead.
+            // (To avoid things like suggesting `Ok(while .. { .. })`.)
+            if expr_ty.is_unit() {
+                if let Some(hir::Node::Block(&hir::Block {
+                    span: block_span, expr: Some(e), ..
+                })) = self.tcx.hir().find(self.tcx.hir().get_parent_node(expr.hir_id))
+                {
+                    if e.hir_id == expr.hir_id {
+                        if let Some(span) = expr.span.find_ancestor_inside(block_span) {
+                            let return_suggestions =
+                                if self.tcx.is_diagnostic_item(sym::Result, expected_adt.did) {
+                                    vec!["Ok(())".to_string()]
+                                } else if self.tcx.is_diagnostic_item(sym::Option, expected_adt.did)
+                                {
+                                    vec!["None".to_string(), "Some(())".to_string()]
+                                } else {
+                                    return;
+                                };
+                            if let Some(indent) =
+                                self.tcx.sess.source_map().indentation_before(span.shrink_to_lo())
+                            {
+                                // Add a semicolon, except after `}`.
+                                let semicolon =
+                                    match self.tcx.sess.source_map().span_to_snippet(span) {
+                                        Ok(s) if s.ends_with('}') => "",
+                                        _ => ";",
+                                    };
+                                err.span_suggestions(
+                                    span.shrink_to_hi(),
+                                    "try adding an expression at the end of the block",
+                                    return_suggestions
+                                        .into_iter()
+                                        .map(|r| format!("{}\n{}{}", semicolon, indent, r)),
+                                    Applicability::MaybeIncorrect,
+                                );
+                            }
+                            return;
+                        }
+                    }
+                }
+            }
+
+            let compatible_variants: Vec<String> = expected_adt
                 .variants
                 .iter()
                 .filter(|variant| variant.fields.len() == 1)
@@ -220,19 +263,33 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         None
                     }
                 })
-                .peekable();
+                .collect();
 
-            if compatible_variants.peek().is_some() {
-                if let Ok(expr_text) = self.tcx.sess.source_map().span_to_snippet(expr.span) {
-                    let suggestions = compatible_variants.map(|v| format!("{}({})", v, expr_text));
-                    let msg = "try using a variant of the expected enum";
-                    err.span_suggestions(
-                        expr.span,
-                        msg,
-                        suggestions,
-                        Applicability::MaybeIncorrect,
-                    );
-                }
+            if let [variant] = &compatible_variants[..] {
+                // Just a single matching variant.
+                err.multipart_suggestion(
+                    &format!("try wrapping the expression in `{}`", variant),
+                    vec![
+                        (expr.span.shrink_to_lo(), format!("{}(", variant)),
+                        (expr.span.shrink_to_hi(), ")".to_string()),
+                    ],
+                    Applicability::MaybeIncorrect,
+                );
+            } else if compatible_variants.len() > 1 {
+                // More than one matching variant.
+                err.multipart_suggestions(
+                    &format!(
+                        "try wrapping the expression in a variant of `{}`",
+                        self.tcx.def_path_str(expected_adt.did)
+                    ),
+                    compatible_variants.into_iter().map(|variant| {
+                        vec![
+                            (expr.span.shrink_to_lo(), format!("{}(", variant)),
+                            (expr.span.shrink_to_hi(), ")".to_string()),
+                        ]
+                    }),
+                    Applicability::MaybeIncorrect,
+                );
             }
         }
     }