infer: Add type inference support for Union types

This adds the type inference to Union types and add a small test case
for it, ensuring it keeps working in future.

Fixes: #5277
Signed-off-by: Otavio Salvador <[email protected]>

2 files changed, 32 insertions, 2 deletions

diff --git a/crates/ra_hir_ty/src/infer/expr.rs b/crates/ra_hir_ty/src/infer/expr.rs
index 06baac2a9..bd9a387f5 100644
--- a/crates/ra_hir_ty/src/infer/expr.rs
+++ b/crates/ra_hir_ty/src/infer/expr.rs
@@ -405,8 +405,15 @@ impl<'a> InferenceContext<'a> {
                                     .subst(&a_ty.parameters)
                             })
                         }
-                        // FIXME:
-                        TypeCtor::Adt(AdtId::UnionId(_)) => None,
+                        TypeCtor::Adt(AdtId::UnionId(u)) => {
+                            self.db.union_data(u).variant_data.field(name).map(|local_id| {
+                                let field = FieldId { parent: u.into(), local_id };
+                                self.write_field_resolution(tgt_expr, field);
+                                self.db.field_types(u.into())[field.local_id]
+                                    .clone()
+                                    .subst(&a_ty.parameters)
+                            })
+                        }
                         _ => None,
                     },
                     _ => None,
diff --git a/crates/ra_hir_ty/src/tests/simple.rs b/crates/ra_hir_ty/src/tests/simple.rs
index de63f4cce..6d3e264af 100644
--- a/crates/ra_hir_ty/src/tests/simple.rs
+++ b/crates/ra_hir_ty/src/tests/simple.rs
@@ -325,6 +325,29 @@ fn test() {
 }
 
 #[test]
+fn infer_union() {
+    assert_snapshot!(
+        infer(r#"
+union MyUnion {
+    foo: u32,
+    bar: f32,
+}
+
+unsafe fn baz(u: MyUnion) {
+    let inner = u.foo;
+}
+"#),
+        @r###"
+    61..62 'u': MyUnion
+    73..99 '{     ...foo; }': ()
+    83..88 'inner': u32
+    91..92 'u': MyUnion
+    91..96 'u.foo': u32
+    "###
+    );
+}
+
+#[test]
 fn infer_refs() {
     assert_snapshot!(
         infer(r#"