cliard24 0.1.2

//! # Iterative Expression Solver
//!
//! This module provides an iterative (explicit stack) solver for finding arithmetic expressions
//! that evaluate to a target value using given operands. It is a non-recursive alternative to
//! the recursive `ExprSolver`, maintaining the same search logic but using manual stack management.
//!
//! The algorithm performs a depth-first search over all possible binary expression trees
//! constructible from the input operands and operators.

use crate::expression::{
    ast::{BinaryOpr, EXPR_EPSILON, ExprNode, ExprNum, ExprTree},
    error::AstError,
};
use std::collections::HashSet;

/// Represents a single frame in the explicit call stack used for iterative search.
///
/// Each frame corresponds to a state in the depth-first search where we are attempting
/// to combine a specific pair of operands/sub-expressions with a specific operator.
struct StackFrame {
    /// Index of the left operand within `node_indices`.
    l_idx: usize,
    /// Index of the right operand within `node_indices`.
    r_idx: usize,
    /// Current index within the `oprs` slice being tried.
    opr_idx: usize,
    /// Whether the order of operands (`l_idx` and `r_idx`) has been swapped for the current operator.
    /// Relevant only for non-commutative operators.
    has_swapped: bool,
    /// List of expression node indices (`usize` referencing `ExprTree`) available for combination in this frame.
    node_indices: Vec<usize>,
}

impl StackFrame {
    /// Creates a new `StackFrame` with its `node_indices` vector pre-allocated to the given capacity.
    ///
    /// The other fields (`l_idx`, `r_idx`, `opr_idx`, `has_swapped`) are initialized to their starting states.
    fn with_capacity(capacity: usize) -> Self {
        StackFrame {
            l_idx: 0,
            r_idx: 1,
            opr_idx: 0,
            has_swapped: false,
            node_indices: Vec::with_capacity(capacity),
        }
    }
}

/// Solves for expressions that evaluate to the target value using an iterative (stack-based) algorithm.
///
/// # Arguments
/// * `nums` - Slice of numerical operands to use in expressions.
/// * `oprs` - Slice of binary operators to consider for combinations.
/// * `target` - The numerical value that valid expressions must equal.
/// * `limit` - Maximum number of solutions to find (0 means no limit).
///
/// # Returns
/// * `Ok(HashSet<String>)` - A set of unique string representations of valid expressions.
/// * `Err(AstError<'a>)` - If an error occurs during expression tree construction or evaluation.
///
/// # Errors
/// Returns an `AstError` if `ExprTree::add_node` or `ExprTree::add_root` fails.
///
/// # Algorithm Overview
/// 1.  Initialize the search stack: one frame per operand, each containing a literal node.
/// 2.  Use a loop to simulate recursion, managed by `cur_frame_idx` which points to the current active frame.
/// 3.  In each iteration:
///     - If the solution `limit` is reached, break.
///     - If `cur_frame_idx` points to the top (single node), check if it's a solution.
///     - Otherwise, try to combine the current pair of nodes (`l_idx`, `r_idx`) with the current operator (`opr_idx`).
///     - If successful, create a new frame (`cur_frame_idx + 1`) with the combined node and remaining nodes.
///     - Update state (`l_idx`, `r_idx`, `opr_idx`, `has_swapped`) to explore the next combination, backtracking by moving `cur_frame_idx` as needed.
pub fn solve24<'a>(
    nums: &[ExprNum],
    oprs: &[BinaryOpr],
    target: ExprNum,
    limit: usize,
) -> Result<HashSet<String>, AstError<'a>> {
    let mut stack = Vec::with_capacity(nums.len());
    let mut tree = ExprTree::new();
    let mut solutions = HashSet::new();

    // Initialize stack: one frame per operand, each containing a literal expression node.
    for (i, n) in nums.iter().enumerate() {
        // Capacity decreases for higher frames as operands get combined.
        stack.push(StackFrame::with_capacity(nums.len() - i));
        // Add a literal node for this operand to the first frame's list.
        stack[0]
            .node_indices
            .push(tree.add_node(ExprNode::Literal(*n))?);
    }

    // Main loop simulating depth-first search. `cur_frame_idx` is our "program counter".
    let mut cur_frame_idx: usize = 0;
    let top_frame_idx = stack.len() - 1; // Index of the final frame (single combined expression).
    loop {
        // Early termination if solution limit is reached
        if limit != 0 && limit <= solutions.len() {
            break;
        }

        // --- Case 1: We are at the top frame (only one node remains -> a complete expression) ---
        if cur_frame_idx >= top_frame_idx {
            // Evaluate the single remaining expression node.
            if let Ok(Some(final_result)) = tree.get_pre_value(stack[cur_frame_idx].node_indices[0])
            {
                // Check if the expression evaluates to the target (with floating-point tolerance)
                if (target - final_result).abs() < EXPR_EPSILON {
                    solutions.insert(tree.to_string());
                }
            }
            // This frame's exploration is done. Backtrack.
            // The node created to reach this frame is popped.
            cur_frame_idx -= 1;
            tree.pop_node();
            continue;
        } else {
            let cur_frame = &mut stack[cur_frame_idx];
            // --- Case 2: Current frame has exhausted all operand pairs (l_idx, r_idx) ---
            if cur_frame.l_idx >= (cur_frame.node_indices.len() - 1) {
                // Recursion base case: the bottom-most frame is also exhausted.
                if cur_frame_idx == 0 {
                    break; // Search complete.
                }
                // Current frame exhausted. Backtrack.
                cur_frame_idx -= 1;
                tree.pop_node();
                // Reset the frame's state for potential future reuse (though not in this algorithm path).
                cur_frame.l_idx = 0;
                cur_frame.r_idx = 1;
                continue;
            }

            // --- Case 3 & 4: There is a valid pair (l_idx, r_idx) to combine ---
            // Determine the actual node indices, considering if operands have been swapped.
            let (l_idx, r_idx) = if cur_frame.has_swapped {
                (
                    cur_frame.node_indices[cur_frame.r_idx],
                    cur_frame.node_indices[cur_frame.l_idx],
                )
            } else {
                (
                    cur_frame.node_indices[cur_frame.l_idx],
                    cur_frame.node_indices[cur_frame.r_idx],
                )
            };
            // Create a new expression node combining l_idx and r_idx with the current operator.
            let new_root_idx = tree.add_root(oprs[cur_frame.opr_idx], l_idx, r_idx)?;

            // Only proceed if the combination evaluates successfully (non-division-by-zero, etc.)
            if let Ok(Some(_)) = tree.get_pre_value(new_root_idx) {
                // Split the stack to get mutable references to the current and next frame.
                let (cur_stack, next_stack) = stack.split_at_mut(cur_frame_idx + 1);
                let (cur_frame, next_frame) = (&mut cur_stack[cur_frame_idx], &mut next_stack[0]);
                // Build the node list for the next recursion level: the new combined node plus all others.
                next_frame.node_indices.clear();
                next_frame.node_indices.push(new_root_idx);
                for &idx in cur_frame.node_indices.iter() {
                    if idx != l_idx && idx != r_idx {
                        next_frame.node_indices.push(idx);
                    }
                }
            }
            // If evaluation failed, we simply skip creating the next frame and proceed to state update.
        }

        // --- State Update: Determine the next combination to try in the current frame ---
        let cur_frame = &mut stack[cur_frame_idx];
        if !cur_frame.has_swapped && !oprs[cur_frame.opr_idx].is_commutation() {
            // The operator is non-commutative and we haven't tried swapping operand order yet.
            // Next iteration should try the swapped version.
            cur_frame.has_swapped = true;
            cur_frame_idx += 1; // Stay in/advance to the "try swapped" state.
            continue;
        }
        // Reset swap flag and move to the next operator.
        cur_frame.has_swapped = false;
        cur_frame.opr_idx += 1;
        if cur_frame.opr_idx < oprs.len() {
            // More operators to try for the current (l_idx, r_idx) pair.
            cur_frame_idx += 1;
            continue;
        }
        // All operators tried for current pair. Move to next right operand (r_idx).
        cur_frame.opr_idx = 0;
        cur_frame.r_idx += 1;
        if cur_frame.r_idx < cur_frame.node_indices.len() {
            cur_frame_idx += 1;
            continue;
        }
        // All right operands tried for current left operand (l_idx). Move to next left operand.
        cur_frame.l_idx += 1;
        cur_frame.r_idx = cur_frame.l_idx + 1;
        cur_frame_idx += 1;
        // Loop continues. If l_idx is now out of bounds, it will be caught at the start of the next iteration (Case 2).
    }
    Ok(solutions)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_no_solution() {
        let nums = [1.0, 1.0, 1.0, 1.0];
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        assert_eq!(solutions.len(), 0, "1,1,1,1不能解出24");
    }

    #[test]
    fn test_2numbers_solution() {
        let nums = [4.0, 6.0];
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        assert!(solutions.len() > 0, "4,6能解出24");
    }

    #[test]
    fn test_3numbers_solution() {
        let nums = [3.0, 8.0, 1.0];
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        assert!(solutions.len() > 0, "3,8,1能解出24");
    }

    #[test]
    fn test_4numbers_solution() {
        let nums = [5.0, 5.0, 4.0, 1.0];
        let start = std::time::Instant::now();
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        let duration = start.elapsed();
        println!("runtime-new: {}", duration.as_nanos());
        let count = solutions.len();
        assert!(count > 0, "5,5,4,1能解出24");
        println!("找到 {} 个解法", count);
    }

    #[test]
    fn test_5numbers_solution() {
        let nums = [1.0, 2.0, 3.0, 4.0, 5.0];
        let start = std::time::Instant::now();
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        let duration = start.elapsed();
        println!("runtime-new: {}", duration.as_nanos());
        let count = solutions.len();
        assert!(count > 0, "1,2,3,4,5能解出24");
        println!("找到 {} 个解法", count);
    }

    #[test]
    fn test_6numbers_solution() {
        let nums = [1.0, 1.0, 1.0, 1.0, 1.0, 5.0];
        let start = std::time::Instant::now();
        let solutions = solve24(&nums, &BinaryOpr::get_all_operators(), 24.0, 0).unwrap();
        let duration = start.elapsed();
        println!("runtime-new: {}", duration.as_millis());
        let count = solutions.len();
        assert!(count > 0, "1,1,1,1,1,5能解出24");
        println!("找到 {} 个解法", count);
    }
}