mago-analyzer 1.21.1

use std::cell::RefCell;
use std::rc::Rc;

use foldhash::HashSet;

use mago_algebra::assertion_set::AssertionSet;
use mago_algebra::clause::Clause;
use mago_atom::Atom;
use mago_atom::AtomMap;
use mago_atom::AtomSet;
use mago_atom::atom;
use mago_codex::assertion::Assertion;
use mago_codex::context::ScopeContext;
use mago_codex::ttype::TType;
use mago_codex::ttype::add_optional_union_type;
use mago_codex::ttype::atomic::TAtomic;
use mago_codex::ttype::atomic::scalar::TScalar;
use mago_codex::ttype::comparator::ComparisonResult;
use mago_codex::ttype::comparator::union_comparator;
use mago_codex::ttype::get_mixed;
use mago_codex::ttype::get_never;
use mago_codex::ttype::union::TUnion;
use mago_span::Span;

use crate::common::global::get_super_globals;
use crate::context::Context;
use crate::context::block_flags::BlockContextFlags;
use crate::context::scope::control_action::ControlActionSet;
use crate::context::scope::finally_scope::FinallyScope;
use crate::context::scope::var_has_root;
use crate::reconciler::assertion_reconciler;
use crate::reconciler::negated_assertion_reconciler;

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum BreakContext {
    Switch,
    Loop,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReferenceConstraintSource {
    Global,
    Static,
    Parameter,
    Argument,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReferenceConstraint {
    pub constraint_span: Span,
    pub source: ReferenceConstraintSource,
    pub constraint_type: Option<Rc<TUnion>>,
}

#[derive(Clone, Debug)]
pub struct BlockContext<'ctx> {
    pub scope: ScopeContext<'ctx>,
    pub locals: AtomMap<Rc<TUnion>>,
    pub static_locals: AtomSet,
    pub variables_possibly_in_scope: AtomSet,
    pub conditionally_referenced_variable_ids: AtomSet,
    pub assigned_variable_ids: AtomMap<u32>,
    pub possibly_assigned_variable_ids: AtomSet,

    /// Maps variable names to the number of times they have been referenced in the current scope.
    ///
    /// This might not contain all variables in `locals`, as it is only updated when a variable is referenced.
    pub referenced_counts: AtomMap<u32>,

    /// Maps variable names to the local variable that they reference.
    ///
    /// All keys and values in this map are guaranteed to be set in `locals`.
    pub references_in_scope: AtomMap<Atom>,

    /// Set of references to variables in another scope. These references will be marked as used if they are assigned to.
    pub references_to_external_scope: AtomSet,

    /// A set of references that might still be in scope from a scope likely to cause confusion. This applies
    /// to references set inside a loop or if statement, since it's easy to forget about PHP's weird scope
    /// rules, and assigning to a reference will change the referenced variable rather than shadowing it.
    pub references_possibly_from_confusing_scope: AtomSet,

    /// A map of variable names to their reference constraints,
    /// where the key is the variable name and the value is the reference constraint.
    pub by_reference_constraints: AtomMap<ReferenceConstraint>,

    /// Bitflags for various context states (inside_conditional, inside_isset, etc.)
    pub flags: BlockContextFlags,

    pub clauses: Vec<Rc<Clause>>,
    pub reconciled_expression_clauses: Vec<Rc<Clause>>,
    pub known_functions: AtomSet,
    pub known_constants: AtomSet,
    pub break_types: Vec<BreakContext>,
    pub finally_scope: Option<Rc<RefCell<FinallyScope>>>,
    pub parent_conflicting_clause_variables: AtomSet,
    pub loop_bounds: (u32, u32),
    pub if_body_context: Option<Rc<RefCell<Self>>>,
    pub control_actions: ControlActionSet,
    pub possibly_thrown_exceptions: AtomMap<HashSet<Span>>,

    /// Properties that are DEFINITELY initialized in ALL code paths.
    /// Uses intersection semantics - a property must be initialized in ALL branches.
    pub definitely_initialized_properties: AtomSet,

    /// Properties that are POSSIBLY initialized (in at least one path).
    /// Uses union semantics - a property initialized in ANY branch is included.
    pub possibly_initialized_properties: AtomSet,

    /// Methods called on $this that are DEFINITELY called in ALL code paths.
    /// Uses intersection semantics for tracking transitive initialization.
    pub definitely_called_methods: HashSet<Atom>,

    /// Methods called on $this in at least one path.
    /// Uses union semantics.
    pub called_methods: HashSet<Atom>,

    /// If this method calls `parent::`<method>() where <method> is a class initializer,
    /// this holds the initializer method name.
    pub calls_parent_initializer: Option<Atom>,

    /// Active method call assertions from guard methods.
    /// When inside a conditional like `if ($obj->isValid())`, this tracks assertions
    /// about method return types that should be applied.
    ///
    /// Key: method call expression id (e.g., "$statements->first()")
    /// Value: assertion set to apply to the method's return type
    pub active_method_call_assertions: AtomMap<AssertionSet>,
}

impl BreakContext {
    #[inline]
    pub const fn is_switch(&self) -> bool {
        matches!(self, BreakContext::Switch)
    }
}

impl ReferenceConstraint {
    pub fn new(constraint_span: Span, source: ReferenceConstraintSource, constraint_type: Option<Rc<TUnion>>) -> Self {
        let constraint_type = constraint_type.map(|mut constraint_type| {
            if constraint_type.has_literal_string()
                || constraint_type.has_literal_int()
                || constraint_type.has_literal_float()
            {
                let union = Rc::make_mut(&mut constraint_type);

                if union.has_literal_string() {
                    union.types.to_mut().push(TAtomic::Scalar(TScalar::string()));
                }

                if union.has_literal_int() {
                    union.types.to_mut().push(TAtomic::Scalar(TScalar::int()));
                }

                if union.has_literal_float() {
                    union.types.to_mut().push(TAtomic::Scalar(TScalar::float()));
                }
            }

            constraint_type
        });

        Self { constraint_span, source, constraint_type }
    }
}

impl<'ctx> BlockContext<'ctx> {
    pub fn new(scope: ScopeContext<'ctx>, register_super_globals: bool) -> Self {
        let mut block_context = Self {
            scope,
            locals: AtomMap::default(),
            static_locals: AtomSet::default(),
            variables_possibly_in_scope: AtomSet::default(),
            conditionally_referenced_variable_ids: AtomSet::default(),
            assigned_variable_ids: AtomMap::default(),
            possibly_assigned_variable_ids: AtomSet::default(),
            referenced_counts: AtomMap::default(),
            references_in_scope: AtomMap::default(),
            references_to_external_scope: AtomSet::default(),
            references_possibly_from_confusing_scope: AtomSet::default(),
            by_reference_constraints: AtomMap::default(),
            flags: BlockContextFlags::new(),
            clauses: Vec::new(),
            reconciled_expression_clauses: Vec::new(),
            known_functions: AtomSet::default(),
            known_constants: AtomSet::default(),
            break_types: Vec::new(),
            finally_scope: None,
            parent_conflicting_clause_variables: AtomSet::default(),
            loop_bounds: (0, 0),
            if_body_context: None,
            control_actions: ControlActionSet::new(),
            possibly_thrown_exceptions: AtomMap::default(),
            definitely_initialized_properties: AtomSet::default(),
            possibly_initialized_properties: AtomSet::default(),
            definitely_called_methods: HashSet::default(),
            called_methods: HashSet::default(),
            calls_parent_initializer: None,
            active_method_call_assertions: AtomMap::default(),
        };

        if register_super_globals {
            for (var_name, var_type) in get_super_globals() {
                block_context.locals.insert(atom(var_name), var_type);
            }
        }

        block_context
    }

    pub fn is_global_scope(&self) -> bool {
        self.scope.is_global()
    }

    pub fn update_references_possibly_from_confusing_scope(&mut self, confusing_scope_context: &BlockContext<'ctx>) {
        let references = confusing_scope_context
            .references_in_scope
            .keys()
            .chain(confusing_scope_context.references_to_external_scope.iter());

        for reference_id in references {
            if !self.references_in_scope.contains_key(reference_id)
                && !self.references_to_external_scope.contains(reference_id)
            {
                self.references_possibly_from_confusing_scope.insert(*reference_id);
            }
        }

        self.references_possibly_from_confusing_scope
            .extend(confusing_scope_context.references_possibly_from_confusing_scope.iter().copied());
    }

    pub fn get_redefined_locals(
        &self,
        new_locals: &AtomMap<Rc<TUnion>>,
        include_new_vars: bool,
        removed_vars: &mut AtomSet,
    ) -> AtomMap<Rc<TUnion>> {
        let mut redefined_vars = AtomMap::default();

        let mut var_ids = self.locals.keys().collect::<Vec<_>>();
        var_ids.extend(new_locals.keys());

        for var_id in var_ids {
            if let Some(this_type) = self.locals.get(var_id) {
                if let Some(new_type) = new_locals.get(var_id) {
                    if new_type != this_type {
                        redefined_vars.insert(*var_id, this_type.clone());
                    }
                } else if include_new_vars {
                    redefined_vars.insert(*var_id, this_type.clone());
                }
            } else {
                removed_vars.insert(*var_id);
            }
        }

        redefined_vars
    }

    pub fn get_new_or_updated_locals(original_context: &Self, new_context: &Self) -> AtomSet {
        let mut redefined_var_ids = AtomSet::default();

        for (var_id, new_type) in &new_context.locals {
            if let Some(original_type) = original_context.locals.get(var_id) {
                if original_context.assigned_variable_ids.get(var_id).unwrap_or(&0)
                    != new_context.assigned_variable_ids.get(var_id).unwrap_or(&0)
                    || original_type != new_type
                {
                    redefined_var_ids.insert(*var_id);
                }
            } else {
                redefined_var_ids.insert(*var_id);
            }
        }

        redefined_var_ids
    }

    pub fn remove_reconciled_clause_refs(
        clauses: &Vec<Rc<Clause>>,
        changed_var_ids: &AtomSet,
    ) -> (Vec<Rc<Clause>>, Vec<Rc<Clause>>) {
        let mut included_clauses = Vec::new();
        let mut rejected_clauses = Vec::new();

        'outer: for c in clauses {
            if c.wedge {
                included_clauses.push(c.clone());
                continue;
            }

            for key in c.possibilities.keys() {
                for changed_var_id in changed_var_ids {
                    if changed_var_id == key || var_has_root(*key, *changed_var_id) {
                        rejected_clauses.push(c.clone());
                        continue 'outer;
                    }
                }
            }

            included_clauses.push(c.clone());
        }

        (included_clauses, rejected_clauses)
    }

    pub fn remove_reconciled_clauses(clauses: &Vec<Clause>, changed_var_ids: &AtomSet) -> (Vec<Clause>, Vec<Clause>) {
        let mut included_clauses = Vec::new();
        let mut rejected_clauses = Vec::new();

        'outer: for c in clauses {
            if c.wedge {
                included_clauses.push(c.clone());
                continue;
            }

            for key in c.possibilities.keys() {
                if changed_var_ids.contains(key) {
                    rejected_clauses.push(c.clone());
                    continue 'outer;
                }
            }

            included_clauses.push(c.clone());
        }

        (included_clauses, rejected_clauses)
    }

    pub(crate) fn filter_clauses<'arena>(
        context: &mut Context<'ctx, 'arena>,
        remove_var_id: Atom,
        clauses: Vec<Rc<Clause>>,
        new_type: Option<&TUnion>,
    ) -> Vec<Rc<Clause>> {
        let mut clauses_to_keep = Vec::new();
        let mut other_clauses = Vec::new();

        'outer: for clause in clauses {
            for var_id in clause.possibilities.keys() {
                if var_has_root(*var_id, remove_var_id) {
                    continue 'outer;
                }
            }

            let keep_clause = should_keep_clause(&clause, remove_var_id, new_type);

            if keep_clause {
                clauses_to_keep.push(clause.clone());
            } else {
                other_clauses.push(clause);
            }
        }

        if let Some(new_type) = new_type
            && !new_type.is_mixed()
        {
            for clause in other_clauses {
                let mut type_changed = false;
                let Some(possibilities) = clause.possibilities.get(&remove_var_id) else {
                    clauses_to_keep.push(clause.clone());

                    continue;
                };

                for (_, assertion) in possibilities {
                    if assertion.is_negation() {
                        type_changed = true;
                        break;
                    }

                    let result_type = assertion_reconciler::reconcile(
                        context,
                        assertion,
                        Some(&new_type.clone()),
                        None,
                        false,
                        None,
                        false,
                        false,
                    );

                    if result_type != *new_type {
                        type_changed = true;
                        break;
                    }
                }

                if !type_changed {
                    clauses_to_keep.push(clause.clone());
                }
            }
        }

        clauses_to_keep
    }

    pub(crate) fn remove_variable_from_conflicting_clauses<'arena>(
        &mut self,
        context: &mut Context<'ctx, 'arena>,
        remove_var_id: Atom,
        new_type: Option<&TUnion>,
    ) {
        self.clauses = BlockContext::filter_clauses(context, remove_var_id, self.clauses.clone(), new_type);

        self.parent_conflicting_clause_variables.insert(remove_var_id);
    }

    pub(crate) fn remove_descendants<'arena>(
        &mut self,
        context: &mut Context<'ctx, 'arena>,
        remove_var_id: Atom,
        existing_type: &TUnion,
        new_type: Option<&TUnion>,
    ) {
        self.remove_variable_from_conflicting_clauses(
            context,
            remove_var_id,
            if existing_type.is_mixed() {
                None
            } else if let Some(new_type) = new_type {
                Some(new_type)
            } else {
                None
            },
        );

        let keys = self.locals.keys().copied().collect::<Vec<_>>();

        for var_id in keys {
            if var_has_root(var_id, remove_var_id) {
                self.locals.remove(&var_id);
            }
        }
    }

    /// Registers a variable that is referenced conditionally, like in a property
    /// or array access (`$foo->bar`, `$foo[0]`).
    pub fn add_conditionally_referenced_variable(&mut self, var_name: &str) {
        /// Strips an accessor suffix (from the first `->` or `[`) from a variable name.
        /// Returns the original slice if no accessor is found.
        fn strip_accessor_suffix(var_name: &str) -> &str {
            let first_separator_pos = var_name
                .find("->")
                .map(|pos| {
                    // If we find '->', see if '[' comes before it.
                    var_name.find('[').map_or(pos, |bracket_pos| pos.min(bracket_pos))
                })
                .or_else(|| {
                    // If '->' wasn't found, just look for '['.
                    var_name.find('[')
                });

            if let Some(pos) = first_separator_pos { &var_name[..pos] } else { var_name }
        }

        let stripped_var = strip_accessor_suffix(var_name);

        // A variable is conditionally referenced if it's part of an access chain
        // (i.e., its suffix was stripped) and the base variable is not `$this`.
        if stripped_var != "$this" || stripped_var != var_name {
            self.conditionally_referenced_variable_ids.insert(atom(var_name));
        }
    }

    /// Checks if a variable exists in the local scope, while also registering it
    /// as a conditionally referenced variable if it's part of an access chain.
    #[must_use]
    pub fn has_variable(&mut self, var_name: &str) -> bool {
        self.add_conditionally_referenced_variable(var_name);
        self.locals.contains_key(&atom(var_name))
    }

    pub(crate) fn remove_variable<'arena>(
        &mut self,
        var_name: &str,
        remove_descendants: bool,
        context: &mut Context<'ctx, 'arena>,
    ) {
        let var_atom = atom(var_name);
        if let Some(existing_type) = self.locals.remove(&var_atom)
            && remove_descendants
        {
            self.remove_descendants(context, var_atom, &existing_type, None);
        }

        self.assigned_variable_ids.remove(&var_atom);
        self.possibly_assigned_variable_ids.remove(&var_atom);
        self.conditionally_referenced_variable_ids.remove(&var_atom);
    }

    pub fn remove_possible_reference(&mut self, remove_var_id: &str) {
        let remove_var_atom = atom(remove_var_id);
        if let Some(reference_count) = self.referenced_counts.get(&remove_var_atom)
            && *reference_count > 0
        {
            // If a referenced variable goes out of scope, we need to update the references.
            // All of the references to this variable are still references to the same value,
            // so we pick the first one and make the rest of the references point to it.
            let mut references = vec![];
            for (reference, referenced) in &self.references_in_scope {
                if *referenced == remove_var_atom {
                    references.push(*reference);
                }
            }

            for reference in &references {
                self.references_in_scope.remove(reference);
            }

            debug_assert!(
                !references.is_empty(),
                "No references found for variable {remove_var_id}, even though it has a reference count of {reference_count}"
            );

            if !references.is_empty() {
                self.referenced_counts.remove(&remove_var_atom);

                let first_reference = references.remove(0);
                if !references.is_empty() {
                    let reassigned_references = references.len() as u32;
                    if let Some(existing_count) = self.referenced_counts.get_mut(&first_reference) {
                        *existing_count += reassigned_references;
                    } else {
                        self.referenced_counts.insert(first_reference, reassigned_references);
                    }

                    for reference in references {
                        self.references_in_scope.insert(reference, first_reference);
                    }
                }
            }
        }

        if self.references_in_scope.contains_key(&remove_var_atom) {
            self.decrement_reference_count(remove_var_id);
        }

        self.locals.remove(&remove_var_atom);
        self.variables_possibly_in_scope.remove(&remove_var_atom);
        self.assigned_variable_ids.remove(&remove_var_atom);
        self.possibly_assigned_variable_ids.remove(&remove_var_atom);
        self.conditionally_referenced_variable_ids.remove(&remove_var_atom);
    }

    pub fn update(
        &mut self,
        context: &mut Context<'ctx, '_>,
        start_block_context: &Self,
        end_block_context: &mut Self,
        has_leaving_statements: bool,
        vars_to_update: &AtomSet,
        updated_vars: &mut AtomSet,
    ) {
        if vars_to_update.is_empty() {
            return;
        }

        for (variable_id, old_type) in &start_block_context.locals {
            if !vars_to_update.contains(variable_id) {
                continue;
            }

            let new_type = if !has_leaving_statements && end_block_context.has_variable(variable_id.as_str()) {
                end_block_context.locals.get(variable_id).cloned()
            } else {
                None
            };

            let Some(existing_type) = self.locals.get(variable_id).map(std::convert::AsRef::as_ref).cloned() else {
                if let Some(new_type) = new_type {
                    self.locals.insert(*variable_id, new_type);
                    updated_vars.insert(*variable_id);
                }

                continue;
            };

            let old_type = old_type.as_ref().clone();

            let should_substitute = match &new_type {
                Some(new_type) => !old_type.eq(new_type),
                None => existing_type.types.len() > 1,
            };

            let resulting_type = if should_substitute {
                updated_vars.insert(*variable_id);

                substitute_types(context, existing_type, old_type, new_type.as_deref())
            } else {
                existing_type
            };

            self.locals.insert(*variable_id, Rc::new(resulting_type));
        }
    }

    /// Decrement the reference count of the variable that $`ref_id` is referring to. This needs to
    /// be done before $`ref_id` is changed to no longer reference its currently referenced variable,
    /// for example by unsetting, reassigning to another reference, or being shadowed by a global.
    pub fn decrement_reference_count(&mut self, ref_id: &str) -> bool {
        let ref_atom = atom(ref_id);
        let Some(ref_target) = self.references_in_scope.get(&ref_atom) else {
            return false;
        };

        let Some(reference_count) = self.referenced_counts.get_mut(ref_target) else {
            return false;
        };

        if *reference_count < 1 {
            return false;
        }

        *reference_count -= 1;

        true
    }
}

impl std::fmt::Display for ReferenceConstraintSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ReferenceConstraintSource::Global => write!(f, "global"),
            ReferenceConstraintSource::Static => write!(f, "static"),
            ReferenceConstraintSource::Parameter => write!(f, "parameter"),
            ReferenceConstraintSource::Argument => write!(f, "argument"),
        }
    }
}

fn substitute_types(
    context: &mut Context<'_, '_>,
    existing_type: TUnion,
    old_type: TUnion,
    new_type: Option<&TUnion>,
) -> TUnion {
    if existing_type.is_mixed() || existing_type.is_never() {
        return existing_type;
    }

    let updated_type =
        if existing_type.eq(&old_type) { get_mixed() } else { subtract_union_types(context, existing_type, old_type) };

    add_optional_union_type(updated_type, new_type, context.codebase)
}

/// Subtracts the types in `type_to_remove` from `existing_type`.
///
/// This function iterates through each atomic type in `existing_type`. For each of these,
/// it iteratively applies the logic of "is not `atomic_from_remove_set`" for every
/// atomic type in `type_to_remove`. This effectively refines each part of `existing_type`
/// to exclude any possibilities covered by `type_to_remove`.
///
/// This is primarily useful for determining remaining possible types for a match subject
/// after some conditional arms have been considered.
///
/// # Arguments
///
/// * `context` - The reconciliation context, and providing access to codebase.
/// * `existing_type` - The initial `TUnion` type (the minuend).
/// * `type_to_remove` - The `TUnion` type whose components should be subtracted from `existing_type`.
///
/// # Returns
///
/// A new `TUnion` representing `existing_type - type_to_remove`.
pub fn subtract_union_types(context: &mut Context<'_, '_>, existing_type: TUnion, type_to_remove: TUnion) -> TUnion {
    if existing_type == type_to_remove {
        return get_never();
    }

    if !(existing_type.has_literal_value() && type_to_remove.has_literal_value())
        && union_comparator::is_contained_by(
            context.codebase,
            &existing_type,
            &type_to_remove,
            false,
            false,
            true,
            &mut ComparisonResult::new(),
        )
    {
        return existing_type;
    }

    let mut result = existing_type;
    for atomic in type_to_remove.types.into_owned() {
        let assertion = Assertion::IsNotType(atomic);
        let key = result.get_id();

        result = negated_assertion_reconciler::reconcile(context, &assertion, &result, None, key, None, true);
        if result.is_never() {
            break;
        }
    }

    result
}

fn should_keep_clause(clause: &Rc<Clause>, remove_var_id: Atom, new_type: Option<&TUnion>) -> bool {
    if let Some(possibilities) = clause.possibilities.get(&remove_var_id) {
        if possibilities.len() == 1
            && let Some((_, Assertion::IsType(assertion_type))) = possibilities.first()
            && let Some(new_type) = new_type
            && new_type.is_single()
        {
            return new_type.get_single() == assertion_type;
        }

        false
    } else {
        true
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use mago_atom::atom;
    use mago_codex::context::ScopeContext;
    use mago_codex::ttype::get_mixed;

    use super::BlockContext;

    fn block_context_with_locals(vars: &[&str]) -> BlockContext<'static> {
        let mut block_context = BlockContext::new(ScopeContext::new(), false);
        for variable in vars {
            let variable_atom = atom(variable);
            block_context.locals.insert(variable_atom, Rc::new(get_mixed()));
            block_context.variables_possibly_in_scope.insert(variable_atom);
        }

        block_context
    }

    #[test]
    fn remove_possible_reference_preserves_promoted_reference_counts() {
        let mut block_context = block_context_with_locals(&["$a", "$b", "$c"]);

        // Model: $b =& $a; $c =& $b.
        block_context.references_in_scope.insert(atom("$b"), atom("$a"));
        block_context.references_in_scope.insert(atom("$c"), atom("$b"));
        block_context.referenced_counts.insert(atom("$a"), 1);
        block_context.referenced_counts.insert(atom("$b"), 1);

        block_context.remove_possible_reference("$a");

        // $b is promoted to root, but still has $c pointing to it.
        assert_eq!(block_context.references_in_scope.get(&atom("$c")).copied(), Some(atom("$b")));
        assert_eq!(block_context.referenced_counts.get(&atom("$b")).copied(), Some(1));

        // Removing $b should re-root $c and clear dangling reference links.
        block_context.remove_possible_reference("$b");
        assert!(!block_context.references_in_scope.contains_key(&atom("$c")));
    }

    #[test]
    fn remove_possible_reference_adds_to_existing_promoted_count() {
        let mut block_context = block_context_with_locals(&["$a", "$b", "$c", "$d"]);

        // Model before removal:
        // $b =& $a, $c =& $a, and $d =& $b.
        block_context.references_in_scope.insert(atom("$b"), atom("$a"));
        block_context.references_in_scope.insert(atom("$c"), atom("$a"));
        block_context.references_in_scope.insert(atom("$d"), atom("$b"));
        block_context.referenced_counts.insert(atom("$a"), 2);
        block_context.referenced_counts.insert(atom("$b"), 1);

        block_context.remove_possible_reference("$a");

        // $b is promoted and keeps its prior 1, plus 1 reassigned reference from $c.
        assert_eq!(block_context.referenced_counts.get(&atom("$b")).copied(), Some(2));
        assert_eq!(block_context.references_in_scope.get(&atom("$c")).copied(), Some(atom("$b")));
        assert_eq!(block_context.references_in_scope.get(&atom("$d")).copied(), Some(atom("$b")));
    }

    #[test]
    fn remove_possible_reference_sets_promoted_count_when_missing() {
        let mut block_context = block_context_with_locals(&["$a", "$b", "$c"]);

        // Model before removal:
        // $b =& $a, $c =& $a with no prior count entry for $b.
        block_context.references_in_scope.insert(atom("$b"), atom("$a"));
        block_context.references_in_scope.insert(atom("$c"), atom("$a"));
        block_context.referenced_counts.insert(atom("$a"), 2);

        block_context.remove_possible_reference("$a");

        // $b becomes primary and should have count 1 (the reassigned $c).
        assert_eq!(block_context.referenced_counts.get(&atom("$b")).copied(), Some(1));
        assert_eq!(block_context.references_in_scope.get(&atom("$c")).copied(), Some(atom("$b")));
    }

    #[test]
    #[cfg(debug_assertions)]
    #[should_panic(expected = "No references found for variable")]
    fn remove_possible_reference_panics_on_stale_counts_in_debug() {
        let mut block_context = block_context_with_locals(&["$stale"]);
        block_context.referenced_counts.insert(atom("$stale"), 1);

        // No references_in_scope entry for $stale -> intentional debug panic.
        block_context.remove_possible_reference("$stale");
    }
}