//! MIR lowering for patterns

use hir_def::{
    hir::{ExprId, RecordFieldPat},
    signatures::VariantFields,
};
use rustc_type_ir::inherent::{IntoKind, Ty as _};

use crate::{
    BindingMode, ByRef,
    mir::{
        FieldIndex, LocalId, MutBorrowKind, Operand, OperandKind, PlaceRef, Projection,
        lower::{
            BasicBlockId, BinOp, BindingId, BorrowKind, Expr, Idx, MemoryMap, MirLowerCtx,
            MirLowerError, MirSpan, Pat, PatId, PlaceElem, ProjectionElem, ResolveValueResult,
            Result, Rvalue, SwitchTargets, TerminatorKind, Ty, TyKind, ValueNs, VariantId,
        },
    },
};
use crate::{method_resolution::CandidateId, next_solver::GenericArgs};

macro_rules! not_supported {
    ($x: expr) => {
        return Err(MirLowerError::NotSupported(format!($x)))
    };
}

pub(super) enum AdtPatternShape<'a> {
    Tuple { args: &'a [PatId], ellipsis: Option<u32> },
    Record { args: &'a [RecordFieldPat] },
    Unit,
}

/// We need to do pattern matching in two phases: One to check if the pattern matches, and one to fill the bindings
/// of patterns. This is necessary to prevent double moves and similar problems. For example:
/// ```ignore
/// struct X;
/// match (X, 3) {
///     (b, 2) | (b, 3) => {},
///     _ => {}
/// }
/// ```
/// If we do everything in one pass, we will move `X` to the first `b`, then we see that the second field of tuple
/// doesn't match and we should move the `X` to the second `b` (which here is the same thing, but doesn't need to be) and
/// it might even doesn't match the second pattern and we may want to not move `X` at all.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MatchingMode {
    /// Check that if this pattern matches
    Check,
    /// Assume that this pattern matches, fill bindings
    Bind,
    /// Assume that this pattern matches, assign to existing variables.
    Assign,
}

impl<'db> MirLowerCtx<'_, 'db> {
    /// It gets a `current` unterminated block, appends some statements and possibly a terminator to it to check if
    /// the pattern matches and write bindings, and returns two unterminated blocks, one for the matched path (which
    /// can be the `current` block) and one for the mismatched path. If the input pattern is irrefutable, the
    /// mismatched path block is `None`.
    ///
    /// By default, it will create a new block for mismatched path. If you already have one, you can provide it with
    /// `current_else` argument to save an unnecessary jump. If `current_else` isn't `None`, the result mismatched path
    /// wouldn't be `None` as well. Note that this function will add jumps to the beginning of the `current_else` block,
    /// so it should be an empty block.
    pub(super) fn pattern_match(
        &mut self,
        current: BasicBlockId,
        current_else: Option<BasicBlockId>,
        cond_place: PlaceRef<'db>,
        pattern: PatId,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let (current, current_else) = self.pattern_match_inner(
            current,
            current_else,
            cond_place,
            pattern,
            MatchingMode::Check,
        )?;
        let (current, current_else) = self.pattern_match_inner(
            current,
            current_else,
            cond_place,
            pattern,
            MatchingMode::Bind,
        )?;
        Ok((current, current_else))
    }

    pub(super) fn pattern_match_assignment(
        &mut self,
        current: BasicBlockId,
        value: PlaceRef<'db>,
        pattern: PatId,
    ) -> Result<'db, BasicBlockId> {
        let (current, _) =
            self.pattern_match_inner(current, None, value, pattern, MatchingMode::Assign)?;
        Ok(current)
    }

    pub(super) fn match_self_param(
        &mut self,
        id: BindingId,
        current: BasicBlockId,
        local: LocalId,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        self.pattern_match_binding(
            id,
            BindingMode(ByRef::No, rustc_ast_ir::Mutability::Not),
            local.into(),
            MirSpan::SelfParam,
            current,
            None,
        )
    }

    fn pattern_match_inner(
        &mut self,
        mut current: BasicBlockId,
        mut current_else: Option<BasicBlockId>,
        mut cond_place: PlaceRef<'db>,
        pattern: PatId,
        mode: MatchingMode,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let cnt = self.infer.pat_adjustments.get(&pattern).map(|x| x.len()).unwrap_or_default();
        cond_place.projection = Projection::new_from_iter(
            cond_place
                .projection
                .as_slice()
                .iter()
                .cloned()
                .chain((0..cnt).map(|_| ProjectionElem::Deref)),
        );
        Ok(match &self.store[pattern] {
            Pat::Missing | Pat::Rest | Pat::NotNull => {
                return Err(MirLowerError::IncompletePattern);
            }
            Pat::Wild => (current, current_else),
            Pat::Tuple { args, ellipsis } => {
                let place_ty = cond_place.ty(&self.result, &self.infcx, self.env).ty;
                let subst = match place_ty.kind() {
                    TyKind::Tuple(s) => s,
                    _ => {
                        return Err(MirLowerError::TypeError(
                            "non tuple type matched with tuple pattern",
                        ));
                    }
                };
                self.pattern_match_tuple_like(
                    current,
                    current_else,
                    args,
                    *ellipsis,
                    (0..subst.len()).map(|i| PlaceElem::Field(FieldIndex(i as u32))),
                    &cond_place,
                    mode,
                )?
            }
            Pat::Or(pats) => {
                let then_target = self.new_basic_block();
                let mut finished = false;
                for pat in &**pats {
                    let (mut next, next_else) = self.pattern_match_inner(
                        current,
                        None,
                        cond_place,
                        *pat,
                        MatchingMode::Check,
                    )?;
                    if mode != MatchingMode::Check {
                        (next, _) = self.pattern_match_inner(next, None, cond_place, *pat, mode)?;
                    }
                    self.set_goto(next, then_target, pattern.into());
                    match next_else {
                        Some(t) => {
                            current = t;
                        }
                        None => {
                            finished = true;
                            break;
                        }
                    }
                }
                if !finished {
                    if mode == MatchingMode::Check {
                        let ce = *current_else.get_or_insert_with(|| self.new_basic_block());
                        self.set_goto(current, ce, pattern.into());
                    } else {
                        self.set_terminator(current, TerminatorKind::Unreachable, pattern.into());
                    }
                }
                (then_target, current_else)
            }
            Pat::Record { args, .. } => {
                let Some(variant) = self.infer.variant_resolution_for_pat(pattern) else {
                    not_supported!("unresolved variant for record");
                };
                self.pattern_matching_variant(
                    cond_place,
                    variant,
                    current,
                    pattern.into(),
                    current_else,
                    AdtPatternShape::Record { args },
                    mode,
                )?
            }
            Pat::Range { start, end, range_type: _ } => {
                let mut add_check = |l: &ExprId, binop| -> Result<'db, ()> {
                    let lv =
                        self.lower_literal_or_const_to_operand(self.infer.pat_ty(pattern), l)?;
                    let else_target = *current_else.get_or_insert_with(|| self.new_basic_block());
                    let next = self.new_basic_block();
                    let discr =
                        self.temp(Ty::new_bool(self.interner()), current, pattern.into())?.into();
                    self.push_assignment(
                        current,
                        discr,
                        Rvalue::CheckedBinaryOp(
                            binop,
                            lv,
                            Operand { kind: OperandKind::Copy(cond_place.store()), span: None },
                        ),
                        pattern.into(),
                    );
                    let discr = Operand { kind: OperandKind::Copy(discr.store()), span: None };
                    self.set_terminator(
                        current,
                        TerminatorKind::SwitchInt {
                            discr,
                            targets: SwitchTargets::static_if(1, next, else_target),
                        },
                        pattern.into(),
                    );
                    current = next;
                    Ok(())
                };
                if mode == MatchingMode::Check {
                    if let Some(start) = start {
                        add_check(start, BinOp::Le)?;
                    }
                    if let Some(end) = end {
                        add_check(end, BinOp::Ge)?;
                    }
                }
                (current, current_else)
            }
            Pat::Slice { prefix, slice, suffix } => {
                let pat_ty = self.infer.pat_ty(pattern);
                // FIXME: MIR lowering should be skipped for bodies with inference errors. Once
                // that happens, this recovery for invalid slice patterns can be removed.
                if !matches!(pat_ty.kind(), TyKind::Array(..) | TyKind::Slice(_)) {
                    return Err(MirLowerError::TypeError(
                        "non array or slice type matched with slice pattern",
                    ));
                }

                if mode == MatchingMode::Check {
                    // emit runtime length check for slice
                    if let TyKind::Slice(_) = pat_ty.kind() {
                        let pattern_len = prefix.len() + suffix.len();
                        let place_len = self
                            .temp(Ty::new_usize(self.interner()), current, pattern.into())?
                            .into();
                        self.push_assignment(
                            current,
                            place_len,
                            Rvalue::Len(cond_place.store()),
                            pattern.into(),
                        );
                        let else_target =
                            *current_else.get_or_insert_with(|| self.new_basic_block());
                        let next = self.new_basic_block();
                        if slice.is_none() {
                            self.set_terminator(
                                current,
                                TerminatorKind::SwitchInt {
                                    discr: Operand {
                                        kind: OperandKind::Copy(place_len.store()),
                                        span: None,
                                    },
                                    targets: SwitchTargets::static_if(
                                        pattern_len as u128,
                                        next,
                                        else_target,
                                    ),
                                },
                                pattern.into(),
                            );
                        } else {
                            let c = Operand::from_concrete_const(
                                pattern_len.to_le_bytes().into(),
                                MemoryMap::default(),
                                Ty::new_usize(self.interner()),
                            );
                            let discr = self
                                .temp(Ty::new_bool(self.interner()), current, pattern.into())?
                                .into();
                            self.push_assignment(
                                current,
                                discr,
                                Rvalue::CheckedBinaryOp(
                                    BinOp::Le,
                                    c,
                                    Operand {
                                        kind: OperandKind::Copy(place_len.store()),
                                        span: None,
                                    },
                                ),
                                pattern.into(),
                            );
                            let discr =
                                Operand { kind: OperandKind::Copy(discr.store()), span: None };
                            self.set_terminator(
                                current,
                                TerminatorKind::SwitchInt {
                                    discr,
                                    targets: SwitchTargets::static_if(1, next, else_target),
                                },
                                pattern.into(),
                            );
                        }
                        current = next;
                    }
                }
                for (i, &pat) in prefix.iter().enumerate() {
                    let next_place = cond_place.project(ProjectionElem::ConstantIndex {
                        offset: i as u64,
                        from_end: false,
                    });
                    (current, current_else) =
                        self.pattern_match_inner(current, current_else, next_place, pat, mode)?;
                }
                if let &Some(slice) = slice
                    && mode != MatchingMode::Check
                    && let Pat::Bind { id, subpat: _ } = self.store[slice]
                {
                    let next_place = cond_place.project(ProjectionElem::Subslice {
                        from: prefix.len() as u64,
                        to: suffix.len() as u64,
                    });
                    let mode = self.infer.binding_modes[slice];
                    (current, current_else) = self.pattern_match_binding(
                        id,
                        mode,
                        next_place,
                        (slice).into(),
                        current,
                        current_else,
                    )?;
                }
                for (i, &pat) in suffix.iter().enumerate() {
                    let next_place = cond_place.project(ProjectionElem::ConstantIndex {
                        offset: i as u64,
                        from_end: true,
                    });
                    (current, current_else) =
                        self.pattern_match_inner(current, current_else, next_place, pat, mode)?;
                }
                (current, current_else)
            }
            Pat::Path(p) => match self.infer.variant_resolution_for_pat(pattern) {
                Some(variant) => self.pattern_matching_variant(
                    cond_place,
                    variant,
                    current,
                    pattern.into(),
                    current_else,
                    AdtPatternShape::Unit,
                    mode,
                )?,
                None => {
                    let unresolved_name = || {
                        MirLowerError::unresolved_path(
                            self.db,
                            p,
                            self.display_target(),
                            self.owner.expression_store_owner(self.db),
                            self.store,
                        )
                    };
                    let hygiene = self.store.pat_path_hygiene(pattern);
                    let pr = self
                        .resolver
                        .resolve_path_in_value_ns(self.db, p, hygiene)
                        .ok_or_else(unresolved_name)?;

                    if let (
                        MatchingMode::Assign,
                        ResolveValueResult::ValueNs(ValueNs::LocalBinding(binding)),
                    ) = (mode, &pr)
                    {
                        let local = self.binding_local(*binding)?;
                        self.push_match_assignment(
                            current,
                            local,
                            BindingMode(ByRef::No, rustc_ast_ir::Mutability::Not),
                            cond_place,
                            pattern.into(),
                        );
                        return Ok((current, current_else));
                    }

                    // The path is not a variant or a local, so it is a const
                    if mode != MatchingMode::Check {
                        // A const don't bind anything. Only needs check.
                        return Ok((current, current_else));
                    }
                    let (c, subst) = 'b: {
                        if let Some(x) = self.infer.assoc_resolutions_for_pat(pattern)
                            && let CandidateId::ConstId(c) = x.0
                        {
                            break 'b (c, x.1);
                        }
                        if let ResolveValueResult::ValueNs(ValueNs::ConstId(c)) = pr {
                            break 'b (c, GenericArgs::empty(self.interner()));
                        }
                        not_supported!("path in pattern position that is not const or variant")
                    };
                    let tmp =
                        self.temp(self.infer.pat_ty(pattern), current, pattern.into())?.into();
                    let span = pattern.into();
                    self.lower_const(c.into(), current, tmp, subst, span)?;
                    let tmp2 =
                        self.temp(Ty::new_bool(self.interner()), current, pattern.into())?.into();
                    self.push_assignment(
                        current,
                        tmp2,
                        Rvalue::CheckedBinaryOp(
                            BinOp::Eq,
                            Operand { kind: OperandKind::Copy(tmp.store()), span: None },
                            Operand { kind: OperandKind::Copy(cond_place.store()), span: None },
                        ),
                        span,
                    );
                    let next = self.new_basic_block();
                    let else_target = current_else.unwrap_or_else(|| self.new_basic_block());
                    self.set_terminator(
                        current,
                        TerminatorKind::SwitchInt {
                            discr: Operand { kind: OperandKind::Copy(tmp2.store()), span: None },
                            targets: SwitchTargets::static_if(1, next, else_target),
                        },
                        span,
                    );
                    (next, Some(else_target))
                }
            },
            Pat::Lit(l) => match &self.store[*l] {
                Expr::Literal(l) => {
                    if mode == MatchingMode::Check {
                        let c = self.lower_literal_to_operand(self.infer.pat_ty(pattern), l)?;
                        self.pattern_match_const(current_else, current, c, cond_place, pattern)?
                    } else {
                        (current, current_else)
                    }
                }
                _ => not_supported!("expression path literal"),
            },
            Pat::Bind { id, subpat } => {
                if let Some(subpat) = subpat {
                    (current, current_else) =
                        self.pattern_match_inner(current, current_else, cond_place, *subpat, mode)?
                }
                if mode != MatchingMode::Check {
                    let mode = self.infer.binding_modes[pattern];
                    self.pattern_match_binding(
                        *id,
                        mode,
                        cond_place,
                        pattern.into(),
                        current,
                        current_else,
                    )?
                } else {
                    (current, current_else)
                }
            }
            Pat::TupleStruct { path: _, args, ellipsis } => {
                let Some(variant) = self.infer.variant_resolution_for_pat(pattern) else {
                    not_supported!("unresolved variant");
                };
                self.pattern_matching_variant(
                    cond_place,
                    variant,
                    current,
                    pattern.into(),
                    current_else,
                    AdtPatternShape::Tuple { args, ellipsis: *ellipsis },
                    mode,
                )?
            }
            Pat::Ref { pat, mutability: _ } => {
                let ty = cond_place.ty(&self.result, &self.infcx, self.env).ty;
                if !ty.is_ref() {
                    return Err(MirLowerError::TypeError(
                        "non reference type matched with reference pattern",
                    ));
                }
                let cond_place = cond_place.project(ProjectionElem::Deref);
                self.pattern_match_inner(current, current_else, cond_place, *pat, mode)?
            }
            &Pat::Expr(expr) => {
                stdx::always!(
                    mode == MatchingMode::Assign,
                    "Pat::Expr can only come in destructuring assignments"
                );
                let Some((lhs_place, current)) = self.lower_expr_as_place(current, expr, false)?
                else {
                    return Ok((current, current_else));
                };
                self.push_assignment(
                    current,
                    lhs_place,
                    Operand { kind: OperandKind::Copy(cond_place.store()), span: None }.into(),
                    expr.into(),
                );
                (current, current_else)
            }
            Pat::Box { .. } => not_supported!("box pattern"),
            Pat::Deref { .. } => not_supported!("deref pattern"),
            Pat::ConstBlock(_) => not_supported!("const block pattern"),
        })
    }

    fn pattern_match_binding(
        &mut self,
        id: BindingId,
        mode: BindingMode,
        cond_place: PlaceRef<'db>,
        span: MirSpan,
        current: BasicBlockId,
        current_else: Option<BasicBlockId>,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let target_place = self.binding_local(id)?;
        self.push_storage_live(id, current)?;
        self.push_match_assignment(current, target_place, mode, cond_place, span);
        Ok((current, current_else))
    }

    fn push_match_assignment(
        &mut self,
        current: BasicBlockId,
        target_place: LocalId,
        mode: BindingMode,
        cond_place: PlaceRef<'db>,
        span: MirSpan,
    ) {
        self.push_assignment(
            current,
            target_place.into(),
            match mode {
                BindingMode(ByRef::No, _) => {
                    Operand { kind: OperandKind::Copy(cond_place.store()), span: None }.into()
                }
                BindingMode(ByRef::Yes(rustc_ast_ir::Mutability::Not), _) => {
                    Rvalue::Ref(BorrowKind::Shared, cond_place.store())
                }
                BindingMode(ByRef::Yes(rustc_ast_ir::Mutability::Mut), _) => Rvalue::Ref(
                    BorrowKind::Mut { kind: MutBorrowKind::Default },
                    cond_place.store(),
                ),
            },
            span,
        );
    }

    fn pattern_match_const(
        &mut self,
        current_else: Option<BasicBlockId>,
        current: BasicBlockId,
        c: Operand,
        cond_place: PlaceRef<'db>,
        pattern: Idx<Pat>,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let then_target = self.new_basic_block();
        let else_target = current_else.unwrap_or_else(|| self.new_basic_block());
        let discr = self.temp(Ty::new_bool(self.interner()), current, pattern.into())?.into();
        self.push_assignment(
            current,
            discr,
            Rvalue::CheckedBinaryOp(
                BinOp::Eq,
                c,
                Operand { kind: OperandKind::Copy(cond_place.store()), span: None },
            ),
            pattern.into(),
        );
        let discr = Operand { kind: OperandKind::Copy(discr.store()), span: None };
        self.set_terminator(
            current,
            TerminatorKind::SwitchInt {
                discr,
                targets: SwitchTargets::static_if(1, then_target, else_target),
            },
            pattern.into(),
        );
        Ok((then_target, Some(else_target)))
    }

    fn pattern_matching_variant(
        &mut self,
        cond_place: PlaceRef<'db>,
        variant: VariantId,
        mut current: BasicBlockId,
        span: MirSpan,
        mut current_else: Option<BasicBlockId>,
        shape: AdtPatternShape<'_>,
        mode: MatchingMode,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let place_ty = cond_place.ty(&self.result, &self.infcx, self.env).ty;
        let Some((place_adt, _)) = place_ty.as_adt() else {
            return Err(MirLowerError::TypeError("non ADT type matched with ADT pattern"));
        };
        if place_adt != variant.adt_id(self.db) {
            return Err(MirLowerError::TypeError("ADT pattern does not match place type"));
        }

        Ok(match variant {
            VariantId::EnumVariantId(v) => {
                if mode == MatchingMode::Check {
                    let e = self.const_eval_discriminant(v)? as u128;
                    let tmp = self.discr_temp_place(current);
                    self.push_assignment(
                        current,
                        tmp,
                        Rvalue::Discriminant(cond_place.store()),
                        span,
                    );
                    let next = self.new_basic_block();
                    let else_target = current_else.get_or_insert_with(|| self.new_basic_block());
                    self.set_terminator(
                        current,
                        TerminatorKind::SwitchInt {
                            discr: Operand { kind: OperandKind::Copy(tmp.store()), span: None },
                            targets: SwitchTargets::static_if(e, next, *else_target),
                        },
                        span,
                    );
                    current = next;
                }
                self.pattern_matching_variant_fields(
                    shape,
                    v.fields(self.db),
                    variant,
                    current,
                    current_else,
                    &cond_place,
                    mode,
                )?
            }
            VariantId::StructId(s) => self.pattern_matching_variant_fields(
                shape,
                s.fields(self.db),
                variant,
                current,
                current_else,
                &cond_place,
                mode,
            )?,
            VariantId::UnionId(_) => {
                return Err(MirLowerError::TypeError("pattern matching on union"));
            }
        })
    }

    fn pattern_matching_variant_fields(
        &mut self,
        shape: AdtPatternShape<'_>,
        variant_data: &VariantFields,
        v: VariantId,
        current: BasicBlockId,
        current_else: Option<BasicBlockId>,
        cond_place: &PlaceRef<'db>,
        mode: MatchingMode,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let downcast_place = if matches!(v, VariantId::EnumVariantId(_)) {
            cond_place.project(ProjectionElem::Downcast(v))
        } else {
            *cond_place
        };
        Ok(match shape {
            AdtPatternShape::Record { args } => {
                let it = args
                    .iter()
                    .map(|x| {
                        let field_id =
                            variant_data.field(&x.name).ok_or(MirLowerError::UnresolvedField)?;
                        Ok((PlaceElem::Field(field_id.into()), x.pat))
                    })
                    .collect::<Result<'db, Vec<_>>>()?;
                self.pattern_match_adt(
                    current,
                    current_else,
                    it.into_iter(),
                    &downcast_place,
                    mode,
                )?
            }
            AdtPatternShape::Tuple { args, ellipsis } => {
                let fields = variant_data.fields().iter().map(|(x, _)| PlaceElem::Field(x.into()));
                self.pattern_match_tuple_like(
                    current,
                    current_else,
                    args,
                    ellipsis,
                    fields,
                    &downcast_place,
                    mode,
                )?
            }
            AdtPatternShape::Unit => (current, current_else),
        })
    }

    fn pattern_match_adt(
        &mut self,
        mut current: BasicBlockId,
        mut current_else: Option<BasicBlockId>,
        args: impl Iterator<Item = (PlaceElem, PatId)>,
        cond_place: &PlaceRef<'db>,
        mode: MatchingMode,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        for (proj, arg) in args {
            let cond_place = cond_place.project(proj);
            (current, current_else) =
                self.pattern_match_inner(current, current_else, cond_place, arg, mode)?;
        }
        Ok((current, current_else))
    }

    fn pattern_match_tuple_like(
        &mut self,
        current: BasicBlockId,
        current_else: Option<BasicBlockId>,
        args: &[PatId],
        ellipsis: Option<u32>,
        fields: impl DoubleEndedIterator<Item = PlaceElem> + Clone,
        cond_place: &PlaceRef<'db>,
        mode: MatchingMode,
    ) -> Result<'db, (BasicBlockId, Option<BasicBlockId>)> {
        let (al, ar) = args.split_at(ellipsis.map_or(args.len(), |it| it as usize));
        let it = al
            .iter()
            .zip(fields.clone())
            .chain(ar.iter().rev().zip(fields.rev()))
            .map(|(x, y)| (y, *x));
        self.pattern_match_adt(current, current_else, it, cond_place, mode)
    }
}
