Construct exhaustiveness branches with condition, not branch, variable

Previously we would construct the shapes of unions used in the pattern
tree for exhaustiveness checking using the type of the branch patterns,
rather than the type of the condition variable. Clearly we want to
always use the condition variable, otherwise some branches will be
seen as exhaustive, when they are not!

To do this, we now index into the condition variable while refying the
patterns to build the tree for exhaustiveness checking.

Closes #4068

crates/compiler/can/src/exhaustive.rs

@@ -1,6 +1,7 @@
 use crate::expr::{self, IntValue, WhenBranch};
 use crate::pattern::DestructType;
 use roc_collections::all::HumanIndex;
+use roc_collections::VecMap;
 use roc_error_macros::internal_error;
 use roc_exhaustive::{
     is_useful, Ctor, CtorName, Error, Guard, Literal, Pattern, RenderAs, TagId, Union,
@@ -23,6 +24,7 @@ pub struct ExhaustiveSummary {
 
 pub fn check(
     subs: &Subs,
+    real_var: Variable,
     sketched_rows: SketchedRows,
     context: ExhaustiveContext,
 ) -> ExhaustiveSummary {
@@ -33,7 +35,7 @@ pub fn check(
         non_redundant_rows,
         errors,
         redundancies,
-    } = sketched_rows.reify_to_non_redundant(subs);
+    } = sketched_rows.reify_to_non_redundant(subs, real_var);
     all_errors.extend(errors);
 
     let exhaustive = match roc_exhaustive::check(overall_region, context, non_redundant_rows) {
@@ -55,27 +57,155 @@ pub fn check(
 enum SketchedPattern {
     Anything,
     Literal(Literal),
-    Ctor(Variable, TagName, Vec<SketchedPattern>),
-    KnownCtor(Union, TagId, Vec<SketchedPattern>),
+    /// A constructor whose expected union is not yet known.
+    /// We'll know the whole union when reifying the sketched pattern against an expected case type.
+    Ctor(TagName, Vec<SketchedPattern>),
+    KnownCtor(Union, IndexCtor<'static>, TagId, Vec<SketchedPattern>),
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum IndexCtor<'a> {
+    /// Index an opaque type. There should be one argument.
+    Opaque,
+    /// Index a record type. The arguments are the types of the record fields.
+    Record,
+    /// Index a guard constructor. The arguments are a faux guard pattern, and then the real
+    /// pattern being guarded. E.g. `A B if g` becomes Guard { [True, (A B)] }.
+    Guard,
+    /// Index a tag union with the given tag constructor.
+    Tag(&'a TagName),
+}
+
+/// Index a variable as a certain constructor, to get the expected argument types of that constructor.
+fn index_var(
+    subs: &Subs,
+    mut var: Variable,
+    ctor: IndexCtor,
+    render_as: &RenderAs,
+) -> Vec<Variable> {
+    if matches!(ctor, IndexCtor::Guard) {
+        // `A B if g` becomes Guard { [True, (A B)] }, so the arguments are a bool, and the type
+        // of the pattern.
+        return vec![Variable::BOOL, var];
+    }
+    loop {
+        match subs.get_content_without_compacting(var) {
+            Content::FlexVar(_)
+            | Content::RigidVar(_)
+            | Content::FlexAbleVar(_, _)
+            | Content::RigidAbleVar(_, _)
+            | Content::LambdaSet(_)
+            | Content::RangedNumber(..) => internal_error!("not a indexable constructor"),
+            Content::Error => {
+                internal_error!("errors should not be reachable during exhautiveness checking")
+            }
+            Content::RecursionVar {
+                structure,
+                opt_name: _,
+            } => {
+                var = *structure;
+            }
+            Content::Structure(structure) => match structure {
+                FlatType::Apply(_, _)
+                | FlatType::Func(_, _, _)
+                | FlatType::FunctionOrTagUnion(_, _, _) => {
+                    internal_error!("not an indexable constructor")
+                }
+                FlatType::Erroneous(_) => {
+                    internal_error!("errors should not be reachable during exhautiveness checking")
+                }
+                FlatType::Record(fields, ext) => {
+                    let fields_order = match render_as {
+                        RenderAs::Record(fields) => fields,
+                        _ => internal_error!(
+                            "record constructors must always be rendered as records"
+                        ),
+                    };
+                    let iter = fields
+                        .unsorted_iterator(subs, *ext)
+                        .expect("should not have errors if performing exhautiveness checking");
+
+                    let map: VecMap<_, _> = iter
+                        .map(|(name, field)| (name, *field.as_inner()))
+                        .collect();
+
+                    let field_types = fields_order
+                        .iter()
+                        .map(|field| {
+                            *map.get(&field)
+                                .expect("field must be present during exhautiveness checking")
+                        })
+                        .collect();
+
+                    return field_types;
+                }
+                FlatType::TagUnion(tags, ext) | FlatType::RecursiveTagUnion(_, tags, ext) => {
+                    let tag_ctor = match ctor {
+                        IndexCtor::Tag(name) => name,
+                        _ => {
+                            internal_error!("constructor in a tag union must be tag")
+                        }
+                    };
+                    let mut iter = tags.unsorted_iterator(subs, *ext);
+                    let opt_vars = iter.find_map(|(tag, vars)| {
+                        if tag == tag_ctor {
+                            Some(vars.to_vec())
+                        } else {
+                            None
+                        }
+                    });
+                    let vars = opt_vars.expect("constructor must be known in the indexable type if we are exhautiveness checking");
+                    return vars;
+                }
+                FlatType::EmptyRecord => {
+                    debug_assert!(matches!(ctor, IndexCtor::Record));
+                    return vec![];
+                }
+                FlatType::EmptyTagUnion => {
+                    internal_error!("empty tag unions are not indexable")
+                }
+            },
+            Content::Alias(_, _, var, AliasKind::Opaque) => {
+                debug_assert!(matches!(ctor, IndexCtor::Opaque));
+                return vec![*var];
+            }
+            Content::Alias(_, _, inner, AliasKind::Structural) => {
+                var = *inner;
+            }
+        }
+    }
 }
 
 impl SketchedPattern {
-    fn reify(self, subs: &Subs) -> Pattern {
+    fn reify(self, subs: &Subs, real_var: Variable) -> Pattern {
         match self {
             Self::Anything => Pattern::Anything,
             Self::Literal(lit) => Pattern::Literal(lit),
-            Self::KnownCtor(union, tag_id, patterns) => Pattern::Ctor(
-                union,
-                tag_id,
-                patterns.into_iter().map(|pat| pat.reify(subs)).collect(),
-            ),
-            Self::Ctor(var, tag_name, patterns) => {
-                let (union, tag_id) = convert_tag(subs, var, &tag_name);
-                Pattern::Ctor(
-                    union,
-                    tag_id,
-                    patterns.into_iter().map(|pat| pat.reify(subs)).collect(),
-                )
+            Self::KnownCtor(union, index_ctor, tag_id, patterns) => {
+                let arg_vars = index_var(subs, real_var, index_ctor, &union.render_as);
+
+                debug_assert!(arg_vars.len() == patterns.len());
+                let args = (patterns.into_iter())
+                    .zip(arg_vars)
+                    .map(|(pat, var)| {
+                        // FIXME
+                        pat.reify(subs, var)
+                    })
+                    .collect();
+
+                Pattern::Ctor(union, tag_id, args)
+            }
+            Self::Ctor(tag_name, patterns) => {
+                let arg_vars = index_var(subs, real_var, IndexCtor::Tag(&tag_name), &RenderAs::Tag);
+                let (union, tag_id) = convert_tag(subs, real_var, &tag_name);
+
+                debug_assert!(arg_vars.len() == patterns.len());
+                let args = (patterns.into_iter())
+                    .zip(arg_vars)
+                    .map(|(pat, var)| pat.reify(subs, var))
+                    .collect();
+
+                Pattern::Ctor(union, tag_id, args)
             }
         }
     }
@@ -96,12 +226,12 @@ pub struct SketchedRows {
 }
 
 impl SketchedRows {
-    fn reify_to_non_redundant(self, subs: &Subs) -> NonRedundantSummary {
-        to_nonredundant_rows(subs, self)
+    fn reify_to_non_redundant(self, subs: &Subs, real_var: Variable) -> NonRedundantSummary {
+        to_nonredundant_rows(subs, real_var, self)
     }
 }
 
-fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedPattern {
+fn sketch_pattern(pattern: &crate::pattern::Pattern) -> SketchedPattern {
     use crate::pattern::Pattern::*;
     use SketchedPattern as SP;
 
@@ -131,9 +261,7 @@ fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedP
                     DestructType::Required | DestructType::Optional(..) => {
                         patterns.push(SP::Anything)
                     }
-                    DestructType::Guard(_, guard) => {
-                        patterns.push(sketch_pattern(destruct.var, &guard.value))
-                    }
+                    DestructType::Guard(_, guard) => patterns.push(sketch_pattern(&guard.value)),
                 }
             }
 
@@ -146,7 +274,7 @@ fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedP
                 }],
             };
 
-            SP::KnownCtor(union, tag_id, patterns)
+            SP::KnownCtor(union, IndexCtor::Record, tag_id, patterns)
         }
 
         AppliedTag {
@@ -156,16 +284,16 @@ fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedP
         } => {
             let simplified_args: std::vec::Vec<_> = arguments
                 .iter()
-                .map(|(var, arg)| sketch_pattern(*var, &arg.value))
+                .map(|(_, arg)| sketch_pattern(&arg.value))
                 .collect();
 
-            SP::Ctor(var, tag_name.clone(), simplified_args)
+            SP::Ctor(tag_name.clone(), simplified_args)
         }
 
         UnwrappedOpaque {
             opaque, argument, ..
         } => {
-            let (arg_var, argument) = &(**argument);
+            let (_, argument) = &(**argument);
 
             let tag_id = TagId(0);
 
@@ -180,8 +308,9 @@ fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedP
 
             SP::KnownCtor(
                 union,
+                IndexCtor::Opaque,
                 tag_id,
-                vec![sketch_pattern(*arg_var, &argument.value)],
+                vec![sketch_pattern(&argument.value)],
             )
         }
 
@@ -197,11 +326,7 @@ fn sketch_pattern(var: Variable, pattern: &crate::pattern::Pattern) -> SketchedP
     }
 }
 
-pub fn sketch_when_branches(
-    target_var: Variable,
-    region: Region,
-    patterns: &[expr::WhenBranch],
-) -> SketchedRows {
+pub fn sketch_when_branches(region: Region, patterns: &[expr::WhenBranch]) -> SketchedRows {
     let mut rows: Vec<SketchedRow> = Vec::with_capacity(patterns.len());
 
     // If any of the branches has a guard, e.g.
@@ -256,18 +381,16 @@ pub fn sketch_when_branches(
 
                 vec![SP::KnownCtor(
                     union,
+                    IndexCtor::Guard,
                     tag_id,
                     // NB: ordering the guard pattern first seems to be better at catching
                     // non-exhaustive constructors in the second argument; see the paper to see if
                     // there is a way to improve this in general.
-                    vec![
-                        guard_pattern,
-                        sketch_pattern(target_var, &loc_pat.pattern.value),
-                    ],
+                    vec![guard_pattern, sketch_pattern(&loc_pat.pattern.value)],
                 )]
             } else {
                 // Simple case
-                vec![sketch_pattern(target_var, &loc_pat.pattern.value)]
+                vec![sketch_pattern(&loc_pat.pattern.value)]
             };
 
             let row = SketchedRow {
@@ -286,13 +409,9 @@ pub fn sketch_when_branches(
     }
 }
 
-pub fn sketch_pattern_to_rows(
-    target_var: Variable,
-    region: Region,
-    pattern: &crate::pattern::Pattern,
-) -> SketchedRows {
+pub fn sketch_pattern_to_rows(region: Region, pattern: &crate::pattern::Pattern) -> SketchedRows {
     let row = SketchedRow {
-        patterns: vec![sketch_pattern(target_var, pattern)],
+        patterns: vec![sketch_pattern(pattern)],
         region,
         // A single row cannot be redundant!
         redundant_mark: RedundantMark::known_non_redundant(),
@@ -313,7 +432,11 @@ struct NonRedundantSummary {
 }
 
 /// INVARIANT: Produces a list of rows where (forall row. length row == 1)
-fn to_nonredundant_rows(subs: &Subs, rows: SketchedRows) -> NonRedundantSummary {
+fn to_nonredundant_rows(
+    subs: &Subs,
+    real_var: Variable,
+    rows: SketchedRows,
+) -> NonRedundantSummary {
     let SketchedRows {
         rows,
         overall_region,
@@ -335,7 +458,7 @@ fn to_nonredundant_rows(subs: &Subs, rows: SketchedRows) -> NonRedundantSummary
     {
         let next_row: Vec<Pattern> = patterns
             .into_iter()
-            .map(|pattern| pattern.reify(subs))
+            .map(|pattern| pattern.reify(subs, real_var))
             .collect();
 
         let redundant_err = if !is_inhabited_row(&next_row) {
@@ -425,15 +548,17 @@ fn convert_tag(subs: &Subs, whole_var: Variable, this_tag: &TagName) -> (Union,
             let mut alternatives = Vec::with_capacity(num_tags);
             let alternatives_iter = sorted_tags.into_iter().chain(opt_openness_tag.into_iter());
 
-            for (index, (tag, args)) in alternatives_iter.enumerate() {
+            let mut index = 0;
+            for (tag, args) in alternatives_iter {
                 let is_inhabited = args.iter().all(|v| subs.is_inhabited(*v));
-
                 if !is_inhabited {
                     // This constructor is not material; we don't need to match over it!
                     continue;
                 }
 
                 let tag_id = TagId(index as TagIdIntType);
+                index += 1;
+
                 if this_tag == &tag {
                     my_tag_id = tag_id;
                 }