mago-analyzer 1.47.1

use indexmap::IndexMap;
use mago_allocator::Arena;

use itertools::Itertools;
use mago_algebra::AlgebraThresholds;
use mago_algebra::assertion_set::AssertionSet;
use mago_algebra::clause::Clause;
use mago_algebra::disjoin_clauses;
use mago_algebra::negate_formula;
use mago_codex::assertion::Assertion;
use mago_codex::ttype::atomic::TAtomic;
use mago_codex::ttype::atomic::scalar::TScalar;
use mago_span::HasSpan;
use mago_span::Span;
use mago_syntax::cst::*;
use mago_word::Word;
use mago_word::WordMap;

use crate::artifacts::AnalysisArtifacts;
use crate::assertion::scrape_assertions;
use crate::context::assertion::AssertionContext;
use crate::context::scope::var_has_root;
use crate::utils::misc::unwrap_expression;

/// Recursively traverses a conditional expression to generate a corresponding logical formula.
///
/// This function serves as the primary entry point for converting control flow conditions
/// (e.g., from `if`, `while`, or ternary expressions) into a set of logical clauses.
/// The resulting formula is represented in Disjunctive Normal Form (a vector of clauses
/// where each clause is a conjunction of assertions).
///
/// The function breaks down the expression by handling logical operators:
/// - **Binary `&&` and `||`**: Delegates to specialized handlers (`handle_binary_and_operation`,
///   `handle_binary_or_operation`) to correctly combine the formulas from the left and
///   right-hand sides.
/// - **Unary `!` (Not)**: Applies negation to the operand's formula. It includes an
///   optimization for De Morgan's laws, transforming `!(A || B)` into `!A && !B` and
///   `!(A && B)` into `!A || !B` before processing.
///
/// For any other expression (the base case of the recursion), it scrapes atomic
/// assertions and converts them into a set of clauses.
///
/// # Parameters
///
/// * `conditional_object_id`: The span of the overall conditional statement (e.g., the `if` keyword).
/// * `creating_object_id`: The span of the specific part of the expression currently being analyzed.
/// * `conditional`: The conditional expression to convert into a formula.
/// * `assertion_context`: The context required for generating assertions.
/// * `artifacts`: A mutable reference to the analysis artifacts.
/// * `algebra_thresholds`: Thresholds for controlling algebra operations complexity.
/// * `formula_size_threshold`: The maximum allowed formula size before returning `None`.
///
/// # Returns
///
/// Returns `Some(Vec<Clause>)` representing the logical formula. Returns `None` if the
/// formula's complexity exceeds the provided `formula_size_threshold` at
/// any point during the recursive process.
#[allow(clippy::too_many_arguments)]
fn get_boolean_literal_comparison_formula<A>(
    conditional_object_id: Span,
    creating_object_id: Span,
    other_side: &Expression,
    literal_is_true: bool,
    is_identical: bool,
    assertion_context: AssertionContext<'_, '_, A>,
    artifacts: &AnalysisArtifacts,
    algebra_thresholds: &AlgebraThresholds,
    formula_size_threshold: u16,
) -> Option<Vec<Clause>>
where
    A: Arena,
{
    if let Some(var_name) = assertion_context.get_expression_id(other_side) {
        let literal_atomic =
            if literal_is_true { TAtomic::Scalar(TScalar::r#true()) } else { TAtomic::Scalar(TScalar::r#false()) };
        let assertion =
            if is_identical { Assertion::IsType(literal_atomic) } else { Assertion::IsNotType(literal_atomic) };

        let mut clause_map = IndexMap::new();
        let mut type_map = IndexMap::new();
        type_map.insert(assertion.to_hash(), assertion);
        clause_map.insert(var_name, type_map);

        return Some(vec![Clause::new(
            clause_map,
            conditional_object_id,
            creating_object_id,
            Some(false),
            Some(true),
            Some(false),
        )]);
    }

    let formula = get_formula(
        conditional_object_id,
        creating_object_id,
        other_side,
        assertion_context,
        artifacts,
        algebra_thresholds,
        formula_size_threshold,
    )?;

    let should_negate = if is_identical { !literal_is_true } else { literal_is_true };
    if should_negate { negate_formula(formula, algebra_thresholds) } else { Some(formula) }
}

pub fn get_formula<A>(
    conditional_object_id: Span,
    creating_object_id: Span,
    conditional: &Expression,
    assertion_context: AssertionContext<'_, '_, A>,
    artifacts: &AnalysisArtifacts,
    algebra_thresholds: &AlgebraThresholds,
    formula_size_threshold: u16,
) -> Option<Vec<Clause>>
where
    A: Arena,
{
    let expression = unwrap_expression(conditional);

    if let Expression::Binary(binary) = expression {
        if matches!(binary.operator, BinaryOperator::And(_) | BinaryOperator::LowAnd(_)) {
            return handle_binary_and_operation(
                conditional_object_id,
                binary.lhs,
                binary.rhs,
                assertion_context,
                artifacts,
                algebra_thresholds,
                formula_size_threshold,
            );
        }

        if matches!(binary.operator, BinaryOperator::Or(_) | BinaryOperator::LowOr(_)) {
            return handle_binary_or_operation(
                conditional_object_id,
                binary.lhs,
                binary.rhs,
                assertion_context,
                artifacts,
                algebra_thresholds,
                formula_size_threshold,
            );
        }

        if let BinaryOperator::Identical(_) | BinaryOperator::NotIdentical(_) = binary.operator {
            let check_boolean = |expr: &Expression| -> (bool, bool) {
                if expr.is_true() {
                    return (true, false);
                }

                if expr.is_false() {
                    return (false, true);
                }

                artifacts.get_expression_type(expr).map_or((false, false), |t| {
                    if t.is_true() {
                        (true, false)
                    } else if t.is_false() {
                        (false, true)
                    } else {
                        (false, false)
                    }
                })
            };

            let is_identical = matches!(binary.operator, BinaryOperator::Identical(_));
            let (left_is_true, left_is_false) = check_boolean(binary.lhs);
            let (right_is_true, right_is_false) = check_boolean(binary.rhs);

            match (left_is_true || left_is_false, right_is_true || right_is_false) {
                (true, _) => {
                    return get_boolean_literal_comparison_formula(
                        conditional_object_id,
                        creating_object_id,
                        binary.rhs,
                        left_is_true,
                        is_identical,
                        assertion_context,
                        artifacts,
                        algebra_thresholds,
                        formula_size_threshold,
                    );
                }
                (_, true) => {
                    return get_boolean_literal_comparison_formula(
                        conditional_object_id,
                        creating_object_id,
                        binary.lhs,
                        right_is_true,
                        is_identical,
                        assertion_context,
                        artifacts,
                        algebra_thresholds,
                        formula_size_threshold,
                    );
                }
                _ => {}
            }
        }
    }

    if let Expression::UnaryPrefix(unary_prefix) = expression
        && unary_prefix.operator.is_not()
    {
        if let Expression::Construct(Construct::Isset(isset_construct)) = unary_prefix.operand
            && isset_construct.values.len() > 1
        {
            let scraped_assertions = scrape_assertions(unary_prefix.operand, artifacts, assertion_context);

            let mut clauses = Vec::new();

            for assertions in scraped_assertions {
                for (var, anded_types) in assertions {
                    let var = if let Some(stripped) = var.as_bytes().strip_prefix(b"=") {
                        mago_word::word(stripped)
                    } else {
                        var
                    };

                    for orred_types in anded_types {
                        let has_equality =
                            orred_types.first().is_some_and(mago_codex::assertion::Assertion::has_equality);
                        let mapped_orred_types = orred_types
                            .into_iter()
                            .map(|orred_type| (orred_type.to_hash(), orred_type))
                            .collect::<IndexMap<_, _>>();

                        clauses.push(Clause::new(
                            {
                                let mut map = IndexMap::new();
                                map.insert(var, mapped_orred_types);
                                map
                            },
                            conditional_object_id,
                            creating_object_id,
                            Some(false),
                            Some(true),
                            Some(has_equality),
                        ));

                        if clauses.len() > usize::from(formula_size_threshold) {
                            return None;
                        }
                    }
                }
            }

            return negate_formula(clauses, algebra_thresholds);
        }

        if let Expression::Binary(binary_expression) = unwrap_expression(unary_prefix.operand) {
            if matches!(binary_expression.operator, BinaryOperator::Or(_) | BinaryOperator::LowOr(_)) {
                return handle_binary_and_operation(
                    conditional_object_id,
                    &Expression::UnaryPrefix(UnaryPrefix {
                        operator: unary_prefix.operator.clone(),
                        operand: assertion_context.arena.alloc(binary_expression.lhs.clone()),
                    }),
                    &Expression::UnaryPrefix(UnaryPrefix {
                        operator: unary_prefix.operator.clone(),
                        operand: assertion_context.arena.alloc(binary_expression.rhs.clone()),
                    }),
                    assertion_context,
                    artifacts,
                    algebra_thresholds,
                    formula_size_threshold,
                );
            }

            if matches!(binary_expression.operator, BinaryOperator::And(_) | BinaryOperator::LowAnd(_)) {
                return handle_binary_or_operation(
                    conditional_object_id,
                    &Expression::UnaryPrefix(UnaryPrefix {
                        operator: unary_prefix.operator.clone(),
                        operand: assertion_context.arena.alloc(binary_expression.lhs.clone()),
                    }),
                    &Expression::UnaryPrefix(UnaryPrefix {
                        operator: unary_prefix.operator.clone(),
                        operand: assertion_context.arena.alloc(binary_expression.rhs.clone()),
                    }),
                    assertion_context,
                    artifacts,
                    algebra_thresholds,
                    formula_size_threshold,
                );
            }
        }

        let unary_operand_span = unary_prefix.operand.span();
        let negated = negate_formula(
            get_formula(
                conditional_object_id,
                unary_operand_span,
                unary_prefix.operand,
                assertion_context,
                artifacts,
                algebra_thresholds,
                formula_size_threshold,
            )?,
            algebra_thresholds,
        )?;

        return if negated.len() > usize::from(formula_size_threshold) { None } else { Some(negated) };
    }

    if let Expression::Conditional(conditional_expr) = expression
        && let Some(then) = conditional_expr.then
        && artifacts.get_expression_type(conditional_expr.r#else).is_some_and(|t| t.is_always_falsy())
    {
        return handle_binary_and_operation(
            conditional_object_id,
            conditional_expr.condition,
            then,
            assertion_context,
            artifacts,
            algebra_thresholds,
            formula_size_threshold,
        );
    }

    let mut formula = get_formula_from_assertions(
        conditional_object_id,
        creating_object_id,
        expression,
        scrape_assertions(expression, artifacts, assertion_context),
        formula_size_threshold,
    )?;

    add_nullsafe_base_clauses(expression, &mut formula, conditional_object_id, creating_object_id, assertion_context);

    if formula.len() > usize::from(formula_size_threshold) { None } else { Some(formula) }
}

fn add_nullsafe_base_clauses<A>(
    expression: &Expression,
    formula: &mut Vec<Clause>,
    conditional_object_id: Span,
    creating_object_id: Span,
    assertion_context: AssertionContext<'_, '_, A>,
) where
    A: Arena,
{
    let mut current = Some(unwrap_expression(expression));
    while let Some(expr) = current {
        match expr {
            Expression::Access(Access::NullSafeProperty(access)) => {
                push_not_null_clause(
                    access.object,
                    formula,
                    conditional_object_id,
                    creating_object_id,
                    assertion_context,
                );
                current = Some(unwrap_expression(access.object));
            }
            Expression::Access(Access::Property(access)) => {
                current = Some(unwrap_expression(access.object));
            }
            Expression::Call(Call::NullSafeMethod(call)) => {
                push_not_null_clause(
                    call.object,
                    formula,
                    conditional_object_id,
                    creating_object_id,
                    assertion_context,
                );
                current = Some(unwrap_expression(call.object));
            }
            Expression::Call(Call::Method(call)) => {
                current = Some(unwrap_expression(call.object));
            }
            _ => {
                current = None;
            }
        }
    }
}

fn push_not_null_clause<A>(
    base: &Expression,
    formula: &mut Vec<Clause>,
    conditional_object_id: Span,
    creating_object_id: Span,
    assertion_context: AssertionContext<'_, '_, A>,
) where
    A: Arena,
{
    let Some(base_id) = assertion_context.get_expression_id(base) else {
        return;
    };

    let assertion = Assertion::IsNotType(TAtomic::Null);
    let assertion_hash = assertion.to_hash();

    if formula.iter().any(|clause| {
        clause.possibilities.len() == 1
            && clause
                .possibilities
                .get(&base_id)
                .is_some_and(|types| types.len() == 1 && types.contains_key(&assertion_hash))
    }) {
        return;
    }

    let mut clause_map = IndexMap::new();
    let mut type_map = IndexMap::new();
    type_map.insert(assertion_hash, assertion);
    clause_map.insert(base_id, type_map);

    formula.push(Clause::new(
        clause_map,
        conditional_object_id,
        creating_object_id,
        Some(false),
        Some(true),
        Some(false),
    ));
}

fn get_formula_from_assertions(
    conditional_object_id: Span,
    creating_object_id: Span,
    conditional: &Expression,
    anded_assertions: Vec<WordMap<AssertionSet>>,
    formula_size_threshold: u16,
) -> Option<Vec<Clause>> {
    let mut clauses = Vec::new();
    for assertions in anded_assertions {
        for (var_id, anded_types) in assertions {
            for orred_types in anded_types {
                let Some(first_type) = orred_types.first() else {
                    continue; // should not happen
                };

                let has_equality = first_type.has_equality();
                clauses.push(Clause::new(
                    {
                        let mut map = IndexMap::new();
                        map.insert(
                            var_id,
                            orred_types.into_iter().map(|a| (a.to_hash(), a)).collect::<IndexMap<_, _>>(),
                        );
                        map
                    },
                    conditional_object_id,
                    creating_object_id,
                    Some(false),
                    Some(true),
                    Some(has_equality),
                ));
            }
        }
    }

    if !clauses.is_empty() {
        return if clauses.len() > usize::from(formula_size_threshold) { None } else { Some(clauses) };
    }

    let conditional_span = conditional.span();
    let conditional_ref =
        Word::from(format!("*{}-{}", conditional_span.start.offset, conditional_span.end.offset).as_str());

    Some(vec![Clause::new(
        {
            let mut map = IndexMap::new();
            map.insert(conditional_ref, IndexMap::from([(Assertion::Truthy.to_hash(), Assertion::Truthy)]));
            map
        },
        conditional_object_id,
        creating_object_id,
        None,
        None,
        None,
    )])
}

pub fn negate_or_synthesize<A>(
    clauses: Vec<Clause>,
    conditional: &Expression,
    assertion_context: AssertionContext<'_, '_, A>,
    artifacts: &AnalysisArtifacts,
    algebra_thresholds: &AlgebraThresholds,
    formula_size_threshold: u16,
) -> Vec<Clause>
where
    A: Arena,
{
    match negate_formula(clauses, algebra_thresholds) {
        Some(negated_clauses) => negated_clauses,
        None => match get_formula(
            conditional.span(),
            conditional.span(),
            &Expression::UnaryPrefix(UnaryPrefix {
                operator: UnaryPrefixOperator::Not(conditional.span()),
                operand: assertion_context.arena.alloc(conditional.clone()),
            }),
            assertion_context,
            artifacts,
            algebra_thresholds,
            formula_size_threshold,
        ) {
            Some(synthesized_clauses) => synthesized_clauses,
            None => {
                // If we cannot negate the formula, we return an empty vector
                // This is a fallback, and it should not happen in normal cases
                vec![Clause::new(IndexMap::new(), conditional.span(), conditional.span(), Some(true), None, None)]
            }
        },
    }
}

#[inline]
fn handle_binary_or_operation<A>(
    conditional_object_id: Span,
    left: &Expression,
    right: &Expression,
    assertion_context: AssertionContext<'_, '_, A>,
    artifacts: &AnalysisArtifacts,
    algebra_thresholds: &AlgebraThresholds,
    formula_size_threshold: u16,
) -> Option<Vec<Clause>>
where
    A: Arena,
{
    let left_clauses = get_formula(
        conditional_object_id,
        left.span(),
        left,
        assertion_context,
        artifacts,
        algebra_thresholds,
        formula_size_threshold,
    )?;
    let right_clauses = get_formula(
        conditional_object_id,
        right.span(),
        right,
        assertion_context,
        artifacts,
        algebra_thresholds,
        formula_size_threshold,
    )?;
    let clauses = disjoin_clauses(left_clauses, right_clauses, conditional_object_id, algebra_thresholds);

    if clauses.len() > usize::from(formula_size_threshold) { None } else { Some(clauses) }
}

#[inline]
fn handle_binary_and_operation<A>(
    conditional_object_id: Span,
    left: &Expression,
    right: &Expression,
    assertion_context: AssertionContext<'_, '_, A>,
    artifacts: &AnalysisArtifacts,
    algebra_thresholds: &AlgebraThresholds,
    formula_size_threshold: u16,
) -> Option<Vec<Clause>>
where
    A: Arena,
{
    let mut clauses = get_formula(
        conditional_object_id,
        left.span(),
        left,
        assertion_context,
        artifacts,
        algebra_thresholds,
        formula_size_threshold,
    )?;
    clauses.extend(get_formula(
        conditional_object_id,
        right.span(),
        right,
        assertion_context,
        artifacts,
        algebra_thresholds,
        formula_size_threshold,
    )?);

    if clauses.len() > usize::from(formula_size_threshold) { None } else { Some(clauses) }
}

pub fn remove_clauses_with_mixed_variables(
    clauses: Vec<Clause>,
    mut mixed_var_ids: Vec<Word>,
    cond_object_id: Span,
) -> Vec<Clause> {
    clauses
        .into_iter()
        .map(|c| {
            mixed_var_ids.retain(|id| !c.possibilities.contains_key(id));

            if c.possibilities.keys().cartesian_product(&mixed_var_ids).any(|(key, id)| var_has_root(*key, *id)) {
                return Clause::new(IndexMap::new(), cond_object_id, cond_object_id, Some(true), None, None);
            }

            c
        })
        .collect::<Vec<Clause>>()
}