win-auto-utils 0.2.1

Universal Windows automation utilities with memory, window, input, and color operations
//! Loop instruction handler
//!
//! Implements the `loop` instruction for creating loops.

use crate::script_engine::compiler::{GenericBlockMetadata, TerminatorMetadata};
use crate::script_engine::instruction::{
    BlockStructure, InstructionData, InstructionHandler, InstructionMetadata, ScriptError,
};
use crate::script_engine::VMContext;
use crate::scripts_builtin::terminator::TerminatorHandler;

/// Loop runtime state stored in VM execution state
#[derive(Debug)]
pub struct LoopRuntimeState {
    /// Remaining iterations (None for infinite loops)
    pub remaining: Option<u32>,
    /// Start IP of loop body (first instruction after 'loop')
    pub start_ip: usize,
    /// End IP of the loop block (the 'end' instruction)
    pub end_ip: usize,
}

/// Stack key for loop state management
pub const LOOP_STACK_KEY: &str = "__builtin_loop_stack";

/// Initialize loop terminator registration (call once at startup)
pub fn init_loop_terminator() {
    TerminatorHandler::register_handler("loop", |vm: &mut VMContext, _metadata: &TerminatorMetadata| {
        // Get loop stack from VM execution state
        let loop_stack = vm.get_or_create_execution_state::<Vec<LoopRuntimeState>>(LOOP_STACK_KEY);

        // Check the current loop state (peek without popping)
        if let Some(loop_state) = loop_stack.last_mut() {
            // Check if loop should continue
            let should_continue = match &loop_state.remaining {
                Some(count) => *count > 1,
                None => true, // Infinite loop
            };

            if should_continue {
                // Decrement counter
                if let Some(ref mut count) = loop_state.remaining {
                    *count -= 1;
                }

                // Jump back to loop body
                Ok(loop_state.start_ip + 1)
            } else {
                // Loop finished, pop and continue after 'end'
                loop_stack.pop();
                Ok(vm.ip + 1)
            }
        } else {
            Err(ScriptError::ExecutionError(
                "Loop end without active loop".into(),
            ))
        }
    });
}

/// Loop handler - start a loop block
///
/// Syntax: `loop [iterations]`
/// Creates a loop that runs a fixed number of times or infinitely.
///
/// Examples:
/// ```text
/// loop 5                 # Run 5 times
/// loop                   # Infinite loop (use break to exit)
/// ```
pub struct LoopHandler;

impl InstructionHandler for LoopHandler {
    fn name(&self) -> &str {
        "loop"
    }

    fn parse(&self, args: &[&str]) -> Result<InstructionData, ScriptError> {
        let iterations = if args.is_empty() {
            None // Infinite loop
        } else {
            let count = args[0].parse::<u32>().map_err(|e| {
                ScriptError::ParseError(format!("Invalid iteration count '{}': {}", args[0], e))
            })?;

            if count == 0 {
                return Err(ScriptError::ParseError(
                    "Loop iterations must be greater than 0".into(),
                ));
            }

            Some(count)
        };

        Ok(InstructionData::custom(iterations))
    }

    fn execute(
        &self,
        vm: &mut VMContext,
        data: &InstructionData,
        metadata: Option<&InstructionMetadata>,
    ) -> Result<(), ScriptError> {
        // Get iteration count from parsed data
        let iterations = *data.extract_custom::<Option<u32>>("Invalid loop data type")?;

        // Get pre-computed end_ip from metadata (auto-generated by compiler)
        let end_ip = metadata
            .and_then(|m| m.get::<GenericBlockMetadata>())
            .map(|m| m.end_ip)
            .ok_or_else(|| {
                ScriptError::ExecutionError("Loop instruction missing metadata".into())
            })?;

        let start_ip = vm.ip;

        // Push loop state to execution state stack
        let loop_stack = vm.get_or_create_execution_state::<Vec<LoopRuntimeState>>(LOOP_STACK_KEY);
        loop_stack.push(LoopRuntimeState {
            remaining: iterations,
            start_ip,
            end_ip,
        });

        Ok(())
    }

    /// Declare block structure for automatic pairing
    fn declare_block_structure(&self) -> Option<BlockStructure> {
        Some(BlockStructure::SimplePair {
            start_name: "loop",
            end_name: "end",
        })
    }
}

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

    #[test]
    fn test_parse_loop_with_iterations() {
        let handler = LoopHandler;
        let result = handler.parse(&["5"]).unwrap();

        match result {
            InstructionData::Custom(boxed) => {
                let iterations = boxed.downcast_ref::<Option<u32>>().unwrap();
                assert_eq!(*iterations, Some(5));
            }
            _ => panic!("Expected Custom data"),
        }
    }

    #[test]
    fn test_parse_loop_infinite() {
        let handler = LoopHandler;
        let result = handler.parse(&[]).unwrap();

        match result {
            InstructionData::Custom(boxed) => {
                let iterations = boxed.downcast_ref::<Option<u32>>().unwrap();
                assert_eq!(*iterations, None);
            }
            _ => panic!("Expected Custom data"),
        }
    }

    #[test]
    fn test_parse_loop_zero_iterations() {
        let handler = LoopHandler;
        assert!(handler.parse(&["0"]).is_err());
    }

    #[test]
    fn test_parse_loop_invalid_number() {
        let handler = LoopHandler;
        assert!(handler.parse(&["abc"]).is_err());
    }
}