// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! PruneUnusedLValues (PruneTemporaryLValues)
//!
//! Nulls out lvalues for temporary variables that are never accessed later.
//!
//! Corresponds to `src/ReactiveScopes/PruneTemporaryLValues.ts`.

use rustc_hash::FxHashSet;

use crate::react_compiler_hir::{
    DeclarationId, EvaluationOrder, Place, ReactiveFunction, ReactiveInstruction,
    ReactiveStatement, ReactiveTerminal, ReactiveValue, environment::Environment,
};

use crate::react_compiler_reactive_scopes::visitors::{self, ReactiveFunctionVisitor};

/// Nulls out lvalues for unnamed temporaries that are never used.
/// TS: `pruneUnusedLValues`
///
/// Uses ReactiveFunctionVisitor to collect unnamed lvalue DeclarationIds,
/// removing them when referenced as operands. After the visitor pass,
/// a second pass nulls out the remaining unused lvalues.
///
/// This uses a two-phase approach because Rust's ReactiveFunctionVisitor
/// takes immutable references, so we cannot modify lvalues during the visit.
/// The TS version stores mutable instruction references and modifies them
/// after the visitor completes.
pub fn prune_unused_lvalues<'a>(func: &mut ReactiveFunction<'a>, env: &Environment<'a>) {
    // Phase 1: Use ReactiveFunctionVisitor to identify unused unnamed lvalues.
    // When we see an unnamed lvalue on an instruction, we add its DeclarationId.
    // When we see a place reference (operand), we remove its DeclarationId.
    let visitor = Visitor { env };
    let mut lvalues: FxHashSet<DeclarationId> = FxHashSet::default();
    visitors::visit_reactive_function(func, &visitor, &mut lvalues);

    // Phase 2: Null out lvalues whose DeclarationId remains in the map.
    // In the TS, this is done by iterating the stored instruction references.
    // In Rust, we walk the tree to find instructions with matching DeclarationIds.
    if !lvalues.is_empty() {
        null_unused_lvalues(&mut func.body, env, &lvalues);
    }
}

/// TS: `type LValues = Map<DeclarationId, ReactiveInstruction>`
/// In Rust, we only need the set of DeclarationIds (not the instruction refs)
/// because we apply changes in a separate pass.
type LValues = FxHashSet<DeclarationId>;

/// TS: `class Visitor extends ReactiveFunctionVisitor<LValues>`
struct Visitor<'a, 'e> {
    env: &'e Environment<'a>,
}

impl<'a, 'e> ReactiveFunctionVisitor<'a> for Visitor<'a, 'e> {
    type State = LValues;

    fn env(&self) -> &Environment<'a> {
        self.env
    }

    /// TS: `visitPlace(_id, place, state) { state.delete(place.identifier.declarationId) }`
    fn visit_place(&self, _id: EvaluationOrder, place: &Place, state: &mut LValues) {
        let ident = &self.env.identifiers[place.identifier];
        state.remove(&ident.declaration_id);
    }

    /// TS: `visitInstruction(instruction, state)`
    /// Calls traverseInstruction first (visits operands via visitPlace),
    /// then checks if the lvalue is unnamed and adds to map.
    fn visit_instruction(&self, instruction: &ReactiveInstruction<'a>, state: &mut LValues) {
        self.traverse_instruction(instruction, state);
        if let Some(lv) = &instruction.lvalue {
            let ident = &self.env.identifiers[lv.identifier];
            if ident.name.is_none() {
                state.insert(ident.declaration_id);
            }
        }
    }
}

/// Phase 2: Walk the tree and null out lvalues whose DeclarationId is unused.
/// This is necessary because Rust's visitor takes immutable references.
fn null_unused_lvalues<'a>(
    block: &mut Vec<ReactiveStatement<'a>>,
    env: &Environment<'a>,
    unused: &FxHashSet<DeclarationId>,
) {
    for stmt in block.iter_mut() {
        match stmt {
            ReactiveStatement::Instruction(instr) => {
                null_unused_in_instruction(instr, env, unused);
            }
            ReactiveStatement::Scope(scope) => {
                null_unused_lvalues(&mut scope.instructions, env, unused);
            }
            ReactiveStatement::PrunedScope(scope) => {
                null_unused_lvalues(&mut scope.instructions, env, unused);
            }
            ReactiveStatement::Terminal(stmt) => {
                null_unused_in_terminal(&mut stmt.terminal, env, unused);
            }
        }
    }
}

fn null_unused_in_instruction<'a>(
    instr: &mut ReactiveInstruction<'a>,
    env: &Environment<'a>,
    unused: &FxHashSet<DeclarationId>,
) {
    if let Some(lv) = &instr.lvalue {
        let ident = &env.identifiers[lv.identifier];
        if unused.contains(&ident.declaration_id) {
            instr.lvalue = None;
        }
    }
    null_unused_in_value(&mut instr.value, env, unused);
}

fn null_unused_in_value<'a>(
    value: &mut ReactiveValue<'a>,
    env: &Environment<'a>,
    unused: &FxHashSet<DeclarationId>,
) {
    match value {
        ReactiveValue::SequenceExpression { instructions, value: inner, .. } => {
            for instr in instructions.iter_mut() {
                null_unused_in_instruction(instr, env, unused);
            }
            null_unused_in_value(inner, env, unused);
        }
        ReactiveValue::LogicalExpression { left, right, .. } => {
            null_unused_in_value(left, env, unused);
            null_unused_in_value(right, env, unused);
        }
        ReactiveValue::ConditionalExpression { test, consequent, alternate, .. } => {
            null_unused_in_value(test, env, unused);
            null_unused_in_value(consequent, env, unused);
            null_unused_in_value(alternate, env, unused);
        }
        ReactiveValue::OptionalExpression { value: inner, .. } => {
            null_unused_in_value(inner, env, unused);
        }
        ReactiveValue::Instruction(_) => {}
    }
}

fn null_unused_in_terminal<'a>(
    terminal: &mut ReactiveTerminal<'a>,
    env: &Environment<'a>,
    unused: &FxHashSet<DeclarationId>,
) {
    match terminal {
        ReactiveTerminal::Break { .. } | ReactiveTerminal::Continue { .. } => {}
        ReactiveTerminal::Return { .. } | ReactiveTerminal::Throw { .. } => {}
        ReactiveTerminal::For { init, test, update, loop_block, .. } => {
            null_unused_in_value(init, env, unused);
            null_unused_in_value(test, env, unused);
            null_unused_lvalues(loop_block, env, unused);
            if let Some(update) = update {
                null_unused_in_value(update, env, unused);
            }
        }
        ReactiveTerminal::ForOf { init, test, loop_block, .. } => {
            null_unused_in_value(init, env, unused);
            null_unused_in_value(test, env, unused);
            null_unused_lvalues(loop_block, env, unused);
        }
        ReactiveTerminal::ForIn { init, loop_block, .. } => {
            null_unused_in_value(init, env, unused);
            null_unused_lvalues(loop_block, env, unused);
        }
        ReactiveTerminal::DoWhile { loop_block, test, .. } => {
            null_unused_lvalues(loop_block, env, unused);
            null_unused_in_value(test, env, unused);
        }
        ReactiveTerminal::While { test, loop_block, .. } => {
            null_unused_in_value(test, env, unused);
            null_unused_lvalues(loop_block, env, unused);
        }
        ReactiveTerminal::If { consequent, alternate, .. } => {
            null_unused_lvalues(consequent, env, unused);
            if let Some(alt) = alternate {
                null_unused_lvalues(alt, env, unused);
            }
        }
        ReactiveTerminal::Switch { cases, .. } => {
            for case in cases.iter_mut() {
                if let Some(block) = &mut case.block {
                    null_unused_lvalues(block, env, unused);
                }
            }
        }
        ReactiveTerminal::Label { block, .. } => {
            null_unused_lvalues(block, env, unused);
        }
        ReactiveTerminal::Try { block, handler, .. } => {
            null_unused_lvalues(block, env, unused);
            null_unused_lvalues(handler, env, unused);
        }
    }
}
