roc/compiler/unify/src/unify.rs

use roc_collections::all::{default_hasher, get_shared, relative_complement, union, MutMap};
use roc_module::ident::{Lowercase, TagName};
use roc_module::symbol::Symbol;
use roc_types::subs::Content::{self, *};
use roc_types::subs::{Descriptor, FlatType, Mark, OptVariable, RecordFields, Subs, Variable};
use roc_types::types::{gather_fields_ref, ErrorType, Mismatch, RecordField, RecordStructure};

macro_rules! mismatch {
    () => {{
        if cfg!(debug_assertions) {
            println!(
                "Mismatch in {} Line {} Column {}",
                file!(),
                line!(),
                column!()
            );
        }

        vec![Mismatch::TypeMismatch]
    }};
    ($msg:expr) => {{
        if cfg!(debug_assertions) {
            println!(
                "Mismatch in {} Line {} Column {}",
                file!(),
                line!(),
                column!()
            );
            println!($msg);
            println!("");
        }


        vec![Mismatch::TypeMismatch]
    }};
    ($msg:expr,) => {{
        mismatch!($msg)
    }};
    ($msg:expr, $($arg:tt)*) => {{
        if cfg!(debug_assertions) {
            println!(
                "Mismatch in {} Line {} Column {}",
                file!(),
                line!(),
                column!()
            );
            println!($msg, $($arg)*);
            println!("");
        }

        vec![Mismatch::TypeMismatch]
    }};
}

type Pool = Vec<Variable>;

#[derive(Debug)]
pub struct Context {
    first: Variable,
    first_desc: Descriptor,
    second: Variable,
    second_desc: Descriptor,
}

#[derive(Debug)]
pub enum Unified {
    Success(Pool),
    Failure(Pool, ErrorType, ErrorType),
    BadType(Pool, roc_types::types::Problem),
}

#[derive(Debug)]
struct TagUnionStructure {
    tags: MutMap<TagName, Vec<Variable>>,
    ext: Variable,
}

type Outcome = Vec<Mismatch>;

#[inline(always)]
pub fn unify(subs: &mut Subs, var1: Variable, var2: Variable) -> Unified {
    let mut vars = Vec::new();
    let mismatches = unify_pool(subs, &mut vars, var1, var2);

    if mismatches.is_empty() {
        Unified::Success(vars)
    } else {
        let (type1, mut problems) = subs.var_to_error_type(var1);
        let (type2, problems2) = subs.var_to_error_type(var2);

        problems.extend(problems2);

        subs.union(var1, var2, Content::Error.into());

        if !problems.is_empty() {
            Unified::BadType(vars, problems.remove(0))
        } else {
            Unified::Failure(vars, type1, type2)
        }
    }
}

#[inline(always)]
pub fn unify_pool(subs: &mut Subs, pool: &mut Pool, var1: Variable, var2: Variable) -> Outcome {
    if subs.equivalent(var1, var2) {
        Vec::new()
    } else {
        let ctx = Context {
            first: var1,
            first_desc: subs.get(var1),
            second: var2,
            second_desc: subs.get(var2),
        };

        unify_context(subs, pool, ctx)
    }
}

fn unify_context(subs: &mut Subs, pool: &mut Pool, ctx: Context) -> Outcome {
    if false {
        // if true, print the types that are unified.
        //
        // NOTE: names are generated here (when creating an error type) and that modifies names
        // generated by pretty_print.rs. So many test will fail with changes in variable names when
        // this block runs.
        //        let (type1, _problems1) = subs.var_to_error_type(ctx.first);
        //        let (type2, _problems2) = subs.var_to_error_type(ctx.second);
        //        println!("\n --------------- \n");
        //        dbg!(ctx.first, type1);
        //        println!("\n --- \n");
        //        dbg!(ctx.second, type2);
        //        println!("\n --------------- \n");
        println!(
            "{:?} {:?} ~ {:?} {:?}",
            ctx.first,
            subs.get(ctx.first).content,
            ctx.second,
            subs.get(ctx.second).content
        );
    }
    match &ctx.first_desc.content {
        FlexVar(opt_name) => unify_flex(subs, &ctx, opt_name, &ctx.second_desc.content),
        RecursionVar {
            opt_name,
            structure,
        } => unify_recursion(
            subs,
            pool,
            &ctx,
            opt_name,
            *structure,
            &ctx.second_desc.content,
        ),
        RigidVar(name) => unify_rigid(subs, &ctx, name, &ctx.second_desc.content),
        Structure(flat_type) => {
            unify_structure(subs, pool, &ctx, flat_type, &ctx.second_desc.content)
        }
        Alias(symbol, args, real_var) => unify_alias(subs, pool, &ctx, *symbol, args, *real_var),
        Error => {
            // Error propagates. Whatever we're comparing it to doesn't matter!
            merge(subs, &ctx, Error)
        }
    }
}

#[inline(always)]
fn unify_alias(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    symbol: Symbol,
    args: &[(Lowercase, Variable)],
    real_var: Variable,
) -> Outcome {
    let other_content = &ctx.second_desc.content;

    match other_content {
        FlexVar(_) => {
            // Alias wins
            merge(subs, ctx, Alias(symbol, args.to_owned(), real_var))
        }
        RecursionVar { structure, .. } => unify_pool(subs, pool, real_var, *structure),
        RigidVar(_) => unify_pool(subs, pool, real_var, ctx.second),
        Alias(other_symbol, other_args, other_real_var) => {
            if symbol == *other_symbol {
                if args.len() == other_args.len() {
                    let mut problems = Vec::new();
                    for ((_, l_var), (_, r_var)) in args.iter().zip(other_args.iter()) {
                        problems.extend(unify_pool(subs, pool, *l_var, *r_var));
                    }

                    if problems.is_empty() {
                        problems.extend(merge(subs, ctx, other_content.clone()));
                    }

                    if problems.is_empty() {
                        problems.extend(unify_pool(subs, pool, real_var, *other_real_var));
                    }

                    problems
                } else {
                    mismatch!("{}", symbol)
                }
            } else {
                unify_pool(subs, pool, real_var, *other_real_var)
            }
        }
        Structure(_) => unify_pool(subs, pool, real_var, ctx.second),
        Error => merge(subs, ctx, Error),
    }
}

#[inline(always)]
fn unify_structure(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    flat_type: &FlatType,
    other: &Content,
) -> Outcome {
    match other {
        FlexVar(_) => {
            // If the other is flex, Structure wins!
            merge(subs, ctx, Structure(flat_type.clone()))
        }
        RigidVar(name) => {
            // Type mismatch! Rigid can only unify with flex.
            mismatch!("trying to unify {:?} with rigid var {:?}", &flat_type, name)
        }
        RecursionVar { structure, .. } => match flat_type {
            FlatType::TagUnion(_, _) => {
                // unify the structure with this unrecursive tag union
                unify_pool(subs, pool, ctx.first, *structure)
            }
            FlatType::RecursiveTagUnion(rec, _, _) => {
                debug_assert!(is_recursion_var(subs, *rec));
                // unify the structure with this recursive tag union
                unify_pool(subs, pool, ctx.first, *structure)
            }
            FlatType::FunctionOrTagUnion(_, _, _) => {
                // unify the structure with this unrecursive tag union
                unify_pool(subs, pool, ctx.first, *structure)
            }
            _ => todo!("rec structure {:?}", &flat_type),
        },

        Structure(ref other_flat_type) => {
            // Unify the two flat types
            unify_flat_type(subs, pool, ctx, flat_type, other_flat_type)
        }
        Alias(_, _, real_var) => unify_pool(subs, pool, ctx.first, *real_var),
        Error => merge(subs, ctx, Error),
    }
}

fn unify_record(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    rec1: RecordStructure,
    rec2: RecordStructure,
) -> Outcome {
    let fields1 = rec1.fields;
    let fields2 = rec2.fields;

    let separate = RecordFields::separate(fields1, fields2);

    let shared_fields = separate.in_both;

    if separate.only_in_1.is_empty() {
        if separate.only_in_2.is_empty() {
            let ext_problems = unify_pool(subs, pool, rec1.ext, rec2.ext);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut field_problems =
                unify_shared_fields(subs, pool, ctx, shared_fields, OtherFields::None, rec1.ext);

            field_problems.extend(ext_problems);

            field_problems
        } else {
            let flat_type = FlatType::Record(separate.only_in_2, rec2.ext);
            let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
            let ext_problems = unify_pool(subs, pool, rec1.ext, sub_record);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut field_problems = unify_shared_fields(
                subs,
                pool,
                ctx,
                shared_fields,
                OtherFields::None,
                sub_record,
            );

            field_problems.extend(ext_problems);

            field_problems
        }
    } else if separate.only_in_2.is_empty() {
        let flat_type = FlatType::Record(separate.only_in_1, rec1.ext);
        let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
        let ext_problems = unify_pool(subs, pool, sub_record, rec2.ext);

        if !ext_problems.is_empty() {
            return ext_problems;
        }

        let mut field_problems = unify_shared_fields(
            subs,
            pool,
            ctx,
            shared_fields,
            OtherFields::None,
            sub_record,
        );

        field_problems.extend(ext_problems);

        field_problems
    } else {
        let it = (&separate.only_in_1)
            .into_iter()
            .chain((&separate.only_in_2).into_iter());
        let other: RecordFields = it.collect();

        let other_fields = OtherFields::Other(other);

        let ext = fresh(subs, pool, ctx, Content::FlexVar(None));
        let flat_type1 = FlatType::Record(separate.only_in_1, ext);
        let flat_type2 = FlatType::Record(separate.only_in_2, ext);

        let sub1 = fresh(subs, pool, ctx, Structure(flat_type1));
        let sub2 = fresh(subs, pool, ctx, Structure(flat_type2));

        let rec1_problems = unify_pool(subs, pool, rec1.ext, sub2);
        if !rec1_problems.is_empty() {
            return rec1_problems;
        }

        let rec2_problems = unify_pool(subs, pool, sub1, rec2.ext);
        if !rec2_problems.is_empty() {
            return rec2_problems;
        }

        let mut field_problems =
            unify_shared_fields(subs, pool, ctx, shared_fields, other_fields, ext);

        field_problems.reserve(rec1_problems.len() + rec2_problems.len());
        field_problems.extend(rec1_problems);
        field_problems.extend(rec2_problems);

        field_problems
    }
}

enum OtherFields {
    None,
    Other(RecordFields),
}

fn unify_shared_fields(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    shared_fields: Vec<(Lowercase, RecordField<Variable>, RecordField<Variable>)>,
    other_fields: OtherFields,
    ext: Variable,
) -> Outcome {
    let mut matching_fields = Vec::with_capacity(shared_fields.len());
    let num_shared_fields = shared_fields.len();

    for (name, actual, expected) in shared_fields {
        let local_problems = unify_pool(subs, pool, actual.into_inner(), expected.into_inner());

        if local_problems.is_empty() {
            use RecordField::*;

            // Unification of optional fields
            //
            // Demanded does not unify with Optional
            // Unifying Required with Demanded => Demanded
            // Unifying Optional with Required => Required
            // Unifying X with X => X
            let actual = match (actual, expected) {
                (Demanded(_), Optional(_)) | (Optional(_), Demanded(_)) => {
                    // this is an error, but we continue to give better error messages
                    continue;
                }
                (Demanded(val), Required(_))
                | (Required(val), Demanded(_))
                | (Demanded(val), Demanded(_)) => Demanded(val),
                (Required(val), Required(_)) => Required(val),
                (Required(val), Optional(_)) => Required(val),
                (Optional(val), Required(_)) => Required(val),
                (Optional(val), Optional(_)) => Optional(val),
            };

            matching_fields.push((name, actual));
        }
    }

    if num_shared_fields == matching_fields.len() {
        // pull fields in from the ext_var

        let mut ext_fields = MutMap::default();
        let new_ext_var =
            match roc_types::pretty_print::chase_ext_record(subs, ext, &mut ext_fields) {
                Ok(()) => Variable::EMPTY_RECORD,
                Err((new, _)) => new,
            };

        let fields: RecordFields = match other_fields {
            OtherFields::None => {
                if ext_fields.is_empty() {
                    RecordFields::from_sorted_vec(matching_fields)
                } else {
                    matching_fields
                        .into_iter()
                        .chain(ext_fields.into_iter())
                        .collect()
                }
            }
            OtherFields::Other(other_fields) => matching_fields
                .into_iter()
                .chain(other_fields.into_iter())
                .chain(ext_fields.into_iter())
                .collect(),
        };

        let flat_type = FlatType::Record(fields, new_ext_var);

        merge(subs, ctx, Structure(flat_type))
    } else {
        mismatch!("in unify_shared_fields")
    }
}

struct Separate<K, V> {
    only_in_1: MutMap<K, V>,
    only_in_2: MutMap<K, V>,
    in_both: MutMap<K, (V, V)>,
}

fn separate<K, V>(tags1: MutMap<K, V>, mut tags2: MutMap<K, V>) -> Separate<K, V>
where
    K: Ord + std::hash::Hash,
{
    let mut only_in_1 = MutMap::with_capacity_and_hasher(tags1.len(), default_hasher());

    let max_common = tags1.len().min(tags2.len());
    let mut in_both = MutMap::with_capacity_and_hasher(max_common, default_hasher());

    for (k, v1) in tags1.into_iter() {
        match tags2.remove(&k) {
            Some(v2) => {
                in_both.insert(k, (v1, v2));
            }
            None => {
                only_in_1.insert(k, v1);
            }
        }
    }

    Separate {
        only_in_1,
        only_in_2: tags2,
        in_both,
    }
}

fn unify_tag_union(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    rec1: TagUnionStructure,
    rec2: TagUnionStructure,
    recursion: (Option<Variable>, Option<Variable>),
) -> Outcome {
    let tags1 = rec1.tags;
    let tags2 = rec2.tags;

    let recursion_var = match recursion {
        (None, None) => None,
        (Some(v), None) | (None, Some(v)) => Some(v),
        (Some(v1), Some(_v2)) => Some(v1),
    };

    // heuristic: our closure defunctionalization scheme generates a bunch of one-tag unions
    // also our number types fall in this category too.
    if tags1.len() == 1
        && tags2.len() == 1
        && tags1 == tags2
        && subs.get_root_key_without_compacting(rec1.ext)
            == subs.get_root_key_without_compacting(rec2.ext)
    {
        return unify_shared_tags_merge(subs, ctx, tags1, rec1.ext, recursion_var);
    }

    let Separate {
        only_in_1: unique_tags1,
        only_in_2: unique_tags2,
        in_both: shared_tags,
    } = separate(tags1, tags2);

    if unique_tags1.is_empty() {
        if unique_tags2.is_empty() {
            let ext_problems = unify_pool(subs, pool, rec1.ext, rec2.ext);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut tag_problems = unify_shared_tags(
                subs,
                pool,
                ctx,
                shared_tags,
                OtherTags::Empty,
                rec1.ext,
                recursion_var,
            );

            tag_problems.extend(ext_problems);

            tag_problems
        } else {
            let flat_type = FlatType::TagUnion(unique_tags2, rec2.ext);
            let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
            let ext_problems = unify_pool(subs, pool, rec1.ext, sub_record);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut tag_problems = unify_shared_tags(
                subs,
                pool,
                ctx,
                shared_tags,
                OtherTags::Empty,
                sub_record,
                recursion_var,
            );

            tag_problems.extend(ext_problems);

            tag_problems
        }
    } else if unique_tags2.is_empty() {
        let flat_type = FlatType::TagUnion(unique_tags1, rec1.ext);
        let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
        let ext_problems = unify_pool(subs, pool, sub_record, rec2.ext);

        if !ext_problems.is_empty() {
            return ext_problems;
        }

        let mut tag_problems = unify_shared_tags(
            subs,
            pool,
            ctx,
            shared_tags,
            OtherTags::Empty,
            sub_record,
            recursion_var,
        );

        tag_problems.extend(ext_problems);

        tag_problems
    } else {
        let other_tags = OtherTags::Union {
            tags1: unique_tags1.clone(),
            tags2: unique_tags2.clone(),
        };

        let ext = fresh(subs, pool, ctx, Content::FlexVar(None));
        let flat_type1 = FlatType::TagUnion(unique_tags1, ext);
        let flat_type2 = FlatType::TagUnion(unique_tags2, ext);

        let sub1 = fresh(subs, pool, ctx, Structure(flat_type1));
        let sub2 = fresh(subs, pool, ctx, Structure(flat_type2));

        // NOTE: for clearer error messages, we rollback unification of the ext vars when either fails
        //
        // This is inspired by
        //
        //
        //      f : [ Red, Green ] -> Bool
        //      f = \_ -> True
        //
        //      f Blue
        //
        //  In this case, we want the mismatch to be between `[ Blue ]a` and `[ Red, Green ]`, but
        //  without rolling back, the mismatch is between `[ Blue, Red, Green ]a` and `[ Red, Green ]`.
        //  TODO is this also required for the other cases?

        let snapshot = subs.snapshot();

        let ext1_problems = unify_pool(subs, pool, rec1.ext, sub2);
        if !ext1_problems.is_empty() {
            subs.rollback_to(snapshot);
            return ext1_problems;
        }

        let ext2_problems = unify_pool(subs, pool, sub1, rec2.ext);
        if !ext2_problems.is_empty() {
            subs.rollback_to(snapshot);
            return ext2_problems;
        }

        subs.commit_snapshot(snapshot);

        let mut tag_problems =
            unify_shared_tags(subs, pool, ctx, shared_tags, other_tags, ext, recursion_var);

        tag_problems.reserve(ext1_problems.len() + ext2_problems.len());
        tag_problems.extend(ext1_problems);
        tag_problems.extend(ext2_problems);

        tag_problems
    }
}

fn unify_tag_union_not_recursive_recursive(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    rec1: TagUnionStructure,
    rec2: TagUnionStructure,
    recursion_var: Variable,
) -> Outcome {
    let tags1 = rec1.tags;
    let tags2 = rec2.tags;
    let shared_tags = get_shared(&tags1, &tags2);
    // NOTE: don't use `difference` here. In contrast to Haskell, im's `difference` is symmetric
    let unique_tags1 = relative_complement(&tags1, &tags2);
    let unique_tags2 = relative_complement(&tags2, &tags1);

    if unique_tags1.is_empty() {
        if unique_tags2.is_empty() {
            let ext_problems = unify_pool(subs, pool, rec1.ext, rec2.ext);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut tag_problems = unify_shared_tags_recursive_not_recursive(
                subs,
                pool,
                ctx,
                shared_tags,
                MutMap::default(),
                rec1.ext,
                recursion_var,
            );

            tag_problems.extend(ext_problems);

            tag_problems
        } else {
            let flat_type = FlatType::TagUnion(unique_tags2, rec2.ext);
            let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
            let ext_problems = unify_pool(subs, pool, rec1.ext, sub_record);

            if !ext_problems.is_empty() {
                return ext_problems;
            }

            let mut tag_problems = unify_shared_tags_recursive_not_recursive(
                subs,
                pool,
                ctx,
                shared_tags,
                MutMap::default(),
                sub_record,
                recursion_var,
            );

            tag_problems.extend(ext_problems);

            tag_problems
        }
    } else if unique_tags2.is_empty() {
        let flat_type = FlatType::TagUnion(unique_tags1, rec1.ext);
        let sub_record = fresh(subs, pool, ctx, Structure(flat_type));
        let ext_problems = unify_pool(subs, pool, sub_record, rec2.ext);

        if !ext_problems.is_empty() {
            return ext_problems;
        }

        let mut tag_problems = unify_shared_tags_recursive_not_recursive(
            subs,
            pool,
            ctx,
            shared_tags,
            MutMap::default(),
            sub_record,
            recursion_var,
        );

        tag_problems.extend(ext_problems);

        tag_problems
    } else {
        let other_tags = union(unique_tags1.clone(), &unique_tags2);

        let ext = fresh(subs, pool, ctx, Content::FlexVar(None));
        let flat_type1 = FlatType::TagUnion(unique_tags1, ext);
        let flat_type2 = FlatType::TagUnion(unique_tags2, ext);

        let sub1 = fresh(subs, pool, ctx, Structure(flat_type1));
        let sub2 = fresh(subs, pool, ctx, Structure(flat_type2));

        // NOTE: for clearer error messages, we rollback unification of the ext vars when either fails
        //
        // This is inspired by
        //
        //
        //      f : [ Red, Green ] -> Bool
        //      f = \_ -> True
        //
        //      f Blue
        //
        //  In this case, we want the mismatch to be between `[ Blue ]a` and `[ Red, Green ]`, but
        //  without rolling back, the mismatch is between `[ Blue, Red, Green ]a` and `[ Red, Green ]`.
        //  TODO is this also required for the other cases?

        let snapshot = subs.snapshot();

        let ext1_problems = unify_pool(subs, pool, rec1.ext, sub2);
        if !ext1_problems.is_empty() {
            subs.rollback_to(snapshot);
            return ext1_problems;
        }

        let ext2_problems = unify_pool(subs, pool, sub1, rec2.ext);
        if !ext2_problems.is_empty() {
            subs.rollback_to(snapshot);
            return ext2_problems;
        }

        subs.commit_snapshot(snapshot);

        let mut tag_problems = unify_shared_tags_recursive_not_recursive(
            subs,
            pool,
            ctx,
            shared_tags,
            other_tags,
            ext,
            recursion_var,
        );

        tag_problems.reserve(ext1_problems.len() + ext2_problems.len());
        tag_problems.extend(ext1_problems);
        tag_problems.extend(ext2_problems);

        tag_problems
    }
}

/// Is the given variable a structure. Does not consider Attr itself a structure, and instead looks
/// into it.
#[allow(dead_code)]
fn is_structure(var: Variable, subs: &mut Subs) -> bool {
    match subs.get_content_without_compacting(var) {
        Content::Alias(_, _, actual) => is_structure(*actual, subs),
        Content::Structure(_) => true,
        _ => false,
    }
}

fn unify_shared_tags_recursive_not_recursive(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    shared_tags: MutMap<TagName, (Vec<Variable>, Vec<Variable>)>,
    other_tags: MutMap<TagName, Vec<Variable>>,
    ext: Variable,
    recursion_var: Variable,
) -> Outcome {
    let mut matching_tags = MutMap::default();
    let num_shared_tags = shared_tags.len();

    for (name, (actual_vars, expected_vars)) in shared_tags {
        let mut matching_vars = Vec::with_capacity(actual_vars.len());

        let actual_len = actual_vars.len();
        let expected_len = expected_vars.len();

        for (actual, expected) in actual_vars.into_iter().zip(expected_vars.into_iter()) {
            // NOTE the arguments of a tag can be recursive. For instance in the expression
            //
            //  Cons 1 (Cons "foo" Nil)
            //
            // We need to not just check the outer layer (inferring ConsList Int)
            // but also the inner layer (finding a type error, as desired)
            //
            // This correction introduces the same issue as in https://github.com/elm/compiler/issues/1964
            // Polymorphic recursion is now a type error.
            //
            // The strategy is to expand the recursive tag union as deeply as the non-recursive one
            // is.
            //
            // > RecursiveTagUnion(rvar, [ Cons a rvar, Nil ], ext)
            //
            // Conceptually becomes
            //
            // > RecursiveTagUnion(rvar, [ Cons a [ Cons a rvar, Nil ], Nil ], ext)
            //
            // and so on until the whole non-recursive tag union can be unified with it.
            let mut problems = Vec::new();

            {
                // we always unify NonRecursive with Recursive, so this should never happen
                //debug_assert_ne!(Some(actual), recursion_var);

                problems.extend(unify_pool(subs, pool, actual, expected));
            }

            if problems.is_empty() {
                matching_vars.push(expected);
            }
        }

        // only do this check after unification so the error message has more info
        if actual_len == expected_len && actual_len == matching_vars.len() {
            matching_tags.insert(name, matching_vars);
        }
    }

    if num_shared_tags == matching_tags.len() {
        // merge fields from the ext_var into this tag union
        let mut fields = Vec::new();
        let new_ext_var = match roc_types::pretty_print::chase_ext_tag_union(subs, ext, &mut fields)
        {
            Ok(()) => Variable::EMPTY_TAG_UNION,
            Err((new, _)) => new,
        };

        let mut new_tags = union(matching_tags, &other_tags);
        new_tags.extend(fields.into_iter());

        let flat_type = FlatType::RecursiveTagUnion(recursion_var, new_tags, new_ext_var);

        merge(subs, ctx, Structure(flat_type))
    } else {
        mismatch!("Problem with Tag Union")
    }
}

enum OtherTags {
    Empty,
    Union {
        tags1: MutMap<TagName, Vec<Variable>>,
        tags2: MutMap<TagName, Vec<Variable>>,
    },
}

fn unify_shared_tags(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    shared_tags: MutMap<TagName, (Vec<Variable>, Vec<Variable>)>,
    other_tags: OtherTags,
    ext: Variable,
    recursion_var: Option<Variable>,
) -> Outcome {
    let mut matching_tags = MutMap::default();
    let num_shared_tags = shared_tags.len();

    for (name, (actual_vars, expected_vars)) in shared_tags {
        let mut matching_vars = Vec::with_capacity(actual_vars.len());

        let actual_len = actual_vars.len();
        let expected_len = expected_vars.len();

        for (actual, expected) in actual_vars.into_iter().zip(expected_vars.into_iter()) {
            // NOTE the arguments of a tag can be recursive. For instance in the expression
            //
            //  Cons 1 (Cons "foo" Nil)
            //
            // We need to not just check the outer layer (inferring ConsList Int)
            // but also the inner layer (finding a type error, as desired)
            //
            // This correction introduces the same issue as in https://github.com/elm/compiler/issues/1964
            // Polymorphic recursion is now a type error.
            //
            // The strategy is to expand the recursive tag union as deeply as the non-recursive one
            // is.
            //
            // > RecursiveTagUnion(rvar, [ Cons a rvar, Nil ], ext)
            //
            // Conceptually becomes
            //
            // > RecursiveTagUnion(rvar, [ Cons a [ Cons a rvar, Nil ], Nil ], ext)
            //
            // and so on until the whole non-recursive tag union can be unified with it.
            let mut problems = Vec::new();

            {
                problems.extend(unify_pool(subs, pool, actual, expected));
            }

            if problems.is_empty() {
                matching_vars.push(actual);
            }
        }

        // only do this check after unification so the error message has more info
        if actual_len == expected_len && actual_len == matching_vars.len() {
            matching_tags.insert(name, matching_vars);
        }
    }

    if num_shared_tags == matching_tags.len() {
        let mut new_tags = matching_tags;

        // merge fields from the ext_var into this tag union
        let mut fields = Vec::new();
        let new_ext_var = match roc_types::pretty_print::chase_ext_tag_union(subs, ext, &mut fields)
        {
            Ok(()) => Variable::EMPTY_TAG_UNION,
            Err((new, _)) => new,
        };
        new_tags.extend(fields.into_iter());

        match other_tags {
            OtherTags::Empty => {}
            OtherTags::Union { tags1, tags2 } => {
                new_tags.reserve(tags1.len() + tags2.len());
                new_tags.extend(tags1);
                new_tags.extend(tags2);
            }
        }

        unify_shared_tags_merge(subs, ctx, new_tags, new_ext_var, recursion_var)
    } else {
        mismatch!(
            "Problem with Tag Union\nThere should be {:?} matching tags, but I only got \n{:?}",
            num_shared_tags,
            &matching_tags
        )
    }
}

fn unify_shared_tags_merge(
    subs: &mut Subs,
    ctx: &Context,
    new_tags: MutMap<TagName, Vec<Variable>>,
    new_ext_var: Variable,
    recursion_var: Option<Variable>,
) -> Outcome {
    let flat_type = if let Some(rec) = recursion_var {
        debug_assert!(is_recursion_var(subs, rec));
        FlatType::RecursiveTagUnion(rec, new_tags, new_ext_var)
    } else {
        FlatType::TagUnion(new_tags, new_ext_var)
    };

    merge(subs, ctx, Structure(flat_type))
}

#[inline(always)]
fn unify_flat_type(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    left: &FlatType,
    right: &FlatType,
) -> Outcome {
    use roc_types::subs::FlatType::*;

    match (left, right) {
        (EmptyRecord, EmptyRecord) => merge(subs, ctx, Structure(left.clone())),

        (Record(fields, ext), EmptyRecord) if fields.has_only_optional_fields() => {
            unify_pool(subs, pool, *ext, ctx.second)
        }

        (EmptyRecord, Record(fields, ext)) if fields.has_only_optional_fields() => {
            unify_pool(subs, pool, ctx.first, *ext)
        }

        (Record(fields1, ext1), Record(fields2, ext2)) => {
            let rec1 = gather_fields_ref(subs, fields1, *ext1);
            let rec2 = gather_fields_ref(subs, fields2, *ext2);

            unify_record(subs, pool, ctx, rec1, rec2)
        }

        (EmptyTagUnion, EmptyTagUnion) => merge(subs, ctx, Structure(left.clone())),

        (TagUnion(tags, ext), EmptyTagUnion) if tags.is_empty() => {
            unify_pool(subs, pool, *ext, ctx.second)
        }

        (EmptyTagUnion, TagUnion(tags, ext)) if tags.is_empty() => {
            unify_pool(subs, pool, ctx.first, *ext)
        }

        (TagUnion(tags1, ext1), TagUnion(tags2, ext2)) => {
            let union1 = gather_tags(subs, tags1.clone(), *ext1);
            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            unify_tag_union(subs, pool, ctx, union1, union2, (None, None))
        }

        (RecursiveTagUnion(recursion_var, tags1, ext1), TagUnion(tags2, ext2)) => {
            debug_assert!(is_recursion_var(subs, *recursion_var));
            // this never happens in type-correct programs, but may happen if there is a type error
            let union1 = gather_tags(subs, tags1.clone(), *ext1);
            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            unify_tag_union(
                subs,
                pool,
                ctx,
                union1,
                union2,
                (Some(*recursion_var), None),
            )
        }

        (TagUnion(tags1, ext1), RecursiveTagUnion(recursion_var, tags2, ext2)) => {
            debug_assert!(is_recursion_var(subs, *recursion_var));
            let union1 = gather_tags(subs, tags1.clone(), *ext1);
            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            unify_tag_union_not_recursive_recursive(subs, pool, ctx, union1, union2, *recursion_var)
        }

        (RecursiveTagUnion(rec1, tags1, ext1), RecursiveTagUnion(rec2, tags2, ext2)) => {
            debug_assert!(is_recursion_var(subs, *rec1));
            debug_assert!(is_recursion_var(subs, *rec2));
            let union1 = gather_tags(subs, tags1.clone(), *ext1);
            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            let mut problems =
                unify_tag_union(subs, pool, ctx, union1, union2, (Some(*rec1), Some(*rec2)));
            problems.extend(unify_pool(subs, pool, *rec1, *rec2));

            problems
        }

        (Apply(l_symbol, l_args), Apply(r_symbol, r_args)) if l_symbol == r_symbol => {
            let problems = unify_zip(subs, pool, l_args.iter(), r_args.iter());

            if problems.is_empty() {
                merge(subs, ctx, Structure(Apply(*r_symbol, (*r_args).clone())))
            } else {
                problems
            }
        }
        (Func(l_args, l_closure, l_ret), Func(r_args, r_closure, r_ret))
            if l_args.len() == r_args.len() =>
        {
            let arg_problems = unify_zip(subs, pool, l_args.iter(), r_args.iter());
            let ret_problems = unify_pool(subs, pool, *l_ret, *r_ret);
            let closure_problems = unify_pool(subs, pool, *l_closure, *r_closure);

            if arg_problems.is_empty() && closure_problems.is_empty() && ret_problems.is_empty() {
                merge(
                    subs,
                    ctx,
                    Structure(Func((*r_args).clone(), *r_closure, *r_ret)),
                )
            } else {
                let mut problems = ret_problems;

                problems.extend(closure_problems);
                problems.extend(arg_problems);

                problems
            }
        }
        (FunctionOrTagUnion(tag_name, tag_symbol, ext), Func(args, closure, ret)) => {
            unify_function_or_tag_union_and_func(
                subs,
                pool,
                ctx,
                tag_name,
                *tag_symbol,
                *ext,
                args,
                *ret,
                *closure,
                true,
            )
        }
        (Func(args, closure, ret), FunctionOrTagUnion(tag_name, tag_symbol, ext)) => {
            unify_function_or_tag_union_and_func(
                subs,
                pool,
                ctx,
                tag_name,
                *tag_symbol,
                *ext,
                args,
                *ret,
                *closure,
                false,
            )
        }
        (FunctionOrTagUnion(tag_name_1, _, ext_1), FunctionOrTagUnion(tag_name_2, _, ext_2)) => {
            if tag_name_1 == tag_name_2 {
                let problems = unify_pool(subs, pool, *ext_1, *ext_2);
                if problems.is_empty() {
                    let content = subs.get_content_without_compacting(ctx.second).clone();
                    merge(subs, ctx, content)
                } else {
                    problems
                }
            } else {
                let mut tags1 = MutMap::default();
                tags1.insert(*tag_name_1.clone(), vec![]);
                let union1 = gather_tags(subs, tags1, *ext_1);

                let mut tags2 = MutMap::default();
                tags2.insert(*tag_name_2.clone(), vec![]);
                let union2 = gather_tags(subs, tags2, *ext_2);

                unify_tag_union(subs, pool, ctx, union1, union2, (None, None))
            }
        }
        (TagUnion(tags1, ext1), FunctionOrTagUnion(tag_name, _, ext2)) => {
            let union1 = gather_tags(subs, tags1.clone(), *ext1);

            let mut tags2 = MutMap::default();
            tags2.insert(*tag_name.clone(), vec![]);
            let union2 = gather_tags(subs, tags2, *ext2);

            unify_tag_union(subs, pool, ctx, union1, union2, (None, None))
        }
        (FunctionOrTagUnion(tag_name, _, ext1), TagUnion(tags2, ext2)) => {
            let mut tags1 = MutMap::default();
            tags1.insert(*tag_name.clone(), vec![]);
            let union1 = gather_tags(subs, tags1, *ext1);

            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            unify_tag_union(subs, pool, ctx, union1, union2, (None, None))
        }

        (RecursiveTagUnion(recursion_var, tags1, ext1), FunctionOrTagUnion(tag_name, _, ext2)) => {
            // this never happens in type-correct programs, but may happen if there is a type error
            debug_assert!(is_recursion_var(subs, *recursion_var));

            let mut tags2 = MutMap::default();
            tags2.insert(*tag_name.clone(), vec![]);

            let union1 = gather_tags(subs, tags1.clone(), *ext1);
            let union2 = gather_tags(subs, tags2, *ext2);

            unify_tag_union(
                subs,
                pool,
                ctx,
                union1,
                union2,
                (Some(*recursion_var), None),
            )
        }

        (FunctionOrTagUnion(tag_name, _, ext1), RecursiveTagUnion(recursion_var, tags2, ext2)) => {
            debug_assert!(is_recursion_var(subs, *recursion_var));

            let mut tags1 = MutMap::default();
            tags1.insert(*tag_name.clone(), vec![]);

            let union1 = gather_tags(subs, tags1, *ext1);
            let union2 = gather_tags(subs, tags2.clone(), *ext2);

            unify_tag_union_not_recursive_recursive(subs, pool, ctx, union1, union2, *recursion_var)
        }

        (other1, other2) => mismatch!(
            "Trying to unify two flat types that are incompatible: {:?} ~ {:?}",
            other1,
            other2
        ),
    }
}

fn unify_zip<'a, I>(subs: &mut Subs, pool: &mut Pool, left_iter: I, right_iter: I) -> Outcome
where
    I: Iterator<Item = &'a Variable>,
{
    let mut problems = Vec::new();

    let it = left_iter.zip(right_iter);

    for (&l_var, &r_var) in it {
        problems.extend(unify_pool(subs, pool, l_var, r_var));
    }

    problems
}

#[inline(always)]
fn unify_rigid(subs: &mut Subs, ctx: &Context, name: &Lowercase, other: &Content) -> Outcome {
    match other {
        FlexVar(_) => {
            // If the other is flex, rigid wins!
            merge(subs, ctx, RigidVar(name.clone()))
        }
        RigidVar(_) | RecursionVar { .. } | Structure(_) | Alias(_, _, _) => {
            // Type mismatch! Rigid can only unify with flex, even if the
            // rigid names are the same.
            mismatch!("Rigid {:?} with {:?}", ctx.first, &other)
        }
        Error => {
            // Error propagates.
            merge(subs, ctx, Error)
        }
    }
}

#[inline(always)]
fn unify_flex(
    subs: &mut Subs,
    ctx: &Context,
    opt_name: &Option<Lowercase>,
    other: &Content,
) -> Outcome {
    match other {
        FlexVar(None) => {
            // If both are flex, and only left has a name, keep the name around.
            merge(subs, ctx, FlexVar(opt_name.clone()))
        }

        FlexVar(Some(_)) | RigidVar(_) | RecursionVar { .. } | Structure(_) | Alias(_, _, _) => {
            // TODO special-case boolean here
            // In all other cases, if left is flex, defer to right.
            // (This includes using right's name if both are flex and named.)
            merge(subs, ctx, other.clone())
        }

        Error => merge(subs, ctx, Error),
    }
}

#[inline(always)]
fn unify_recursion(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    opt_name: &Option<Lowercase>,
    structure: Variable,
    other: &Content,
) -> Outcome {
    match other {
        RecursionVar {
            opt_name: other_opt_name,
            structure: _other_structure,
        } => {
            // NOTE: structure and other_structure may not be unified yet, but will be
            // we should not do that here, it would create an infinite loop!
            let name = opt_name.clone().or_else(|| other_opt_name.clone());
            merge(
                subs,
                ctx,
                RecursionVar {
                    opt_name: name,
                    structure,
                },
            )
        }

        Structure(_) => {
            // unify the structure variable with this Structure
            unify_pool(subs, pool, structure, ctx.second)
        }
        RigidVar(_) => mismatch!("RecursionVar {:?} with rigid {:?}", ctx.first, &other),

        FlexVar(_) => merge(
            subs,
            ctx,
            RecursionVar {
                structure,
                opt_name: opt_name.clone(),
            },
        ),

        Alias(_, _, actual) => {
            // look at the type the alias stands for

            unify_pool(subs, pool, ctx.first, *actual)
        }

        Error => merge(subs, ctx, Error),
    }
}

pub fn merge(subs: &mut Subs, ctx: &Context, content: Content) -> Outcome {
    let rank = ctx.first_desc.rank.min(ctx.second_desc.rank);
    let desc = Descriptor {
        content,
        rank,
        mark: Mark::NONE,
        copy: OptVariable::NONE,
    };

    subs.union(ctx.first, ctx.second, desc);

    Vec::new()
}

fn register(subs: &mut Subs, desc: Descriptor, pool: &mut Pool) -> Variable {
    let var = subs.fresh(desc);

    pool.push(var);

    var
}

fn fresh(subs: &mut Subs, pool: &mut Pool, ctx: &Context, content: Content) -> Variable {
    register(
        subs,
        Descriptor {
            content,
            rank: ctx.first_desc.rank.min(ctx.second_desc.rank),
            mark: Mark::NONE,
            copy: OptVariable::NONE,
        },
        pool,
    )
}

fn gather_tags(
    subs: &mut Subs,
    mut tags: MutMap<TagName, Vec<Variable>>,
    var: Variable,
) -> TagUnionStructure {
    use roc_types::subs::Content::*;
    use roc_types::subs::FlatType::*;

    match subs.get_content_without_compacting(var) {
        Structure(TagUnion(sub_tags, sub_ext)) => {
            for (k, v) in sub_tags {
                tags.insert(k.clone(), v.clone());
            }

            let sub_ext = *sub_ext;
            gather_tags(subs, tags, sub_ext)
        }

        Alias(_, _, var) => {
            // TODO according to elm/compiler: "TODO may be dropping useful alias info here"
            let var = *var;
            gather_tags(subs, tags, var)
        }

        _ => TagUnionStructure { tags, ext: var },
    }
}

fn is_recursion_var(subs: &Subs, var: Variable) -> bool {
    matches!(
        subs.get_content_without_compacting(var),
        Content::RecursionVar { .. }
    )
}

#[allow(clippy::too_many_arguments)]
fn unify_function_or_tag_union_and_func(
    subs: &mut Subs,
    pool: &mut Pool,
    ctx: &Context,
    tag_name: &TagName,
    tag_symbol: Symbol,
    tag_ext: Variable,
    function_arguments: &[Variable],
    function_return: Variable,
    function_lambda_set: Variable,
    left: bool,
) -> Outcome {
    use FlatType::*;

    let mut new_tags = MutMap::with_capacity_and_hasher(1, default_hasher());

    new_tags.insert(tag_name.clone(), function_arguments.to_owned());

    let content = Structure(TagUnion(new_tags, tag_ext));

    let new_tag_union_var = fresh(subs, pool, ctx, content);

    let mut problems = if left {
        unify_pool(subs, pool, new_tag_union_var, function_return)
    } else {
        unify_pool(subs, pool, function_return, new_tag_union_var)
    };

    {
        let lambda_set_ext = subs.fresh_unnamed_flex_var();

        let mut closure_tags = MutMap::with_capacity_and_hasher(1, default_hasher());
        closure_tags.insert(TagName::Closure(tag_symbol), vec![]);

        let lambda_set_content = Structure(TagUnion(closure_tags, lambda_set_ext));

        let tag_lambda_set = register(
            subs,
            Descriptor {
                content: lambda_set_content,
                rank: ctx.first_desc.rank.min(ctx.second_desc.rank),
                mark: Mark::NONE,
                copy: OptVariable::NONE,
            },
            pool,
        );

        let closure_problems = if left {
            unify_pool(subs, pool, tag_lambda_set, function_lambda_set)
        } else {
            unify_pool(subs, pool, function_lambda_set, tag_lambda_set)
        };

        problems.extend(closure_problems);
    }

    if problems.is_empty() {
        let desc = if left {
            subs.get(ctx.second)
        } else {
            subs.get(ctx.first)
        };

        subs.union(ctx.first, ctx.second, desc);
    }

    problems
}