monty 0.0.21

//! Tests for bytecode operand overflow limits.
//!
//! These tests verify that the bytecode compiler handles cases where operands
//! exceed the u8/u16 limits of the bytecode encoding:
//!
//! - Local variable slots: Use wide instructions (u16), so up to 65535 locals work
//! - Function call arguments: Limited to 255 (u8 operand) - returns SyntaxError if exceeded
//! - Keyword argument counts: Limited to 255 (u8 operand) - returns SyntaxError if exceeded

use std::fmt::Write;

use monty::MontyRun;
use monty_types::{CompileOptions, ExcType, MontyException};

/// Generates Python code with N local variables in a function.
///
/// Creates: `def f(): v0=0; v1=1; ...; v{n-1}={n-1}; return v{n-1}`
fn generate_many_locals(count: usize) -> String {
    let mut code = String::from("def f():\n");
    for i in 0..count {
        writeln!(code, "    v{i} = {i}").unwrap();
    }
    writeln!(code, "    return v{}", count - 1).unwrap();
    code.push_str("f()");
    code
}

/// Generates Python code calling a function with N positional arguments.
///
/// Creates: `def f(*args): return len(args)\nf(0, 1, 2, ..., n-1)`
fn generate_many_positional_args(count: usize) -> String {
    let mut code = String::from("def f(*args): return len(args)\nf(");
    for i in 0..count {
        if i > 0 {
            code.push_str(", ");
        }
        code.push_str(&i.to_string());
    }
    code.push(')');
    code
}

/// Generates Python code calling a function with N keyword arguments.
///
/// Creates: `def f(**kw): return len(kw)\nf(k0=0, k1=1, ..., k{n-1}={n-1})`
fn generate_many_keyword_args(count: usize) -> String {
    let mut code = String::from("def f(**kw): return len(kw)\nf(");
    for i in 0..count {
        if i > 0 {
            code.push_str(", ");
        }
        write!(code, "k{i}={i}").unwrap();
    }
    code.push(')');
    code
}

/// Generates Python code with a function that has N parameters.
///
/// Creates: `def f(p0, p1, ..., p{n-1}): return p{n-1}\nf(0, 1, ..., n-1)`
fn generate_many_parameters(count: usize) -> String {
    let mut code = String::from("def f(");
    for i in 0..count {
        if i > 0 {
            code.push_str(", ");
        }
        write!(code, "p{i}").unwrap();
    }
    code.push_str("):\n");
    writeln!(code, "    return p{}", count - 1).unwrap();
    code.push_str("f(");
    for i in 0..count {
        if i > 0 {
            code.push_str(", ");
        }
        code.push_str(&i.to_string());
    }
    code.push(')');
    code
}

/// Asserts that a MontyRun result is a SyntaxError with a message containing the expected text.
fn assert_syntax_error(result: Result<MontyRun, MontyException>, expected_msg: &str) {
    let err = result.expect_err("expected SyntaxError");
    assert_eq!(
        err.exc_type(),
        ExcType::SyntaxError,
        "expected SyntaxError, got {:?}: {:?}",
        err.exc_type(),
        err.message()
    );
    let msg = err.message().expect("SyntaxError should have message");
    assert!(
        msg.contains(expected_msg),
        "expected message containing '{expected_msg}', got: {msg}"
    );
}

mod local_variable_limits {
    use super::*;

    #[test]
    fn locals_under_u8_limit_succeeds() {
        // 255 locals should work with u8 slots (0-254)
        let code = generate_many_locals(255);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "255 locals should compile successfully");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "255 locals should run successfully");
    }

    #[test]
    fn locals_at_u8_boundary_succeeds() {
        // 256 locals (slots 0-255) - uses wide instructions for slot 255+
        let code = generate_many_locals(256);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(
            result.is_ok(),
            "256 locals should compile successfully (wide instructions)"
        );

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "256 locals should run successfully");
    }

    #[test]
    fn locals_exceeding_u8_uses_wide_instructions() {
        // 257 locals requires LoadLocalW/StoreLocalW for slot 256
        let code = generate_many_locals(257);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "257 locals should compile (using wide instructions)");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "257 locals should run correctly with wide instructions");
    }

    #[test]
    fn locals_well_over_u8_limit() {
        // 300 locals - well into wide instruction territory
        let code = generate_many_locals(300);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "300 locals should compile successfully");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "300 locals should run successfully");
    }
}

mod function_argument_limits {
    use super::*;

    #[test]
    fn positional_args_under_u8_limit_succeeds() {
        // 255 positional args should work
        let code = generate_many_positional_args(255);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "255 positional args should compile successfully");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "255 positional args should run successfully");
    }

    #[test]
    fn positional_args_at_u8_boundary_returns_syntax_error() {
        // 256 positional args - exceeds u8 limit, should return SyntaxError
        let code = generate_many_positional_args(256);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 positional arguments");
    }

    #[test]
    fn positional_args_exceeding_u8_limit_returns_syntax_error() {
        // 257 positional args - clearly exceeds u8 capacity
        let code = generate_many_positional_args(257);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 positional arguments");
    }
}

mod keyword_argument_limits {
    use super::*;

    #[test]
    fn keyword_args_under_u8_limit_succeeds() {
        // 255 keyword args should work
        let code = generate_many_keyword_args(255);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "255 keyword args should compile successfully");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "255 keyword args should run successfully");
    }

    #[test]
    fn keyword_args_at_u8_boundary_returns_syntax_error() {
        // 256 keyword args - exceeds u8 limit, should return SyntaxError
        let code = generate_many_keyword_args(256);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 keyword arguments");
    }

    #[test]
    fn keyword_args_exceeding_u8_limit_returns_syntax_error() {
        // 257 keyword args - clearly exceeds u8 capacity
        let code = generate_many_keyword_args(257);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 keyword arguments");
    }
}

/// Generates a list comprehension with `count` `for` clauses.
///
/// Creates: `[0 for x0 in [0] for x1 in [0] ... for x{count-1} in [0]]`. The
/// compiler used to recurse once per generator with no up-front guard, so a
/// large `count` would overflow the Rust call stack during compilation,
/// before any runtime resource limits applied.
fn generate_comprehension_with_generators(count: usize) -> String {
    let mut code = String::from("x = [0");
    for i in 0..count {
        write!(code, " for x{i} in [0]").unwrap();
    }
    code.push(']');
    code
}

mod comprehension_generator_limits {
    use super::*;

    #[test]
    fn small_comprehension_succeeds() {
        // Each generator with a simple `for x in iter` target adds two items
        // to the operand stack (the iterator and the target leaf), so the
        // practical compile-time limit from the u8 `ListAppend` depth
        // operand is around 127. 50 is comfortably inside that window.
        let code = generate_comprehension_with_generators(50);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(
            result.is_ok(),
            "50 comprehension generators should compile successfully"
        );
    }

    #[test]
    fn generators_exceeding_max_returns_syntax_error() {
        // Above `MAX_COMP_GENERATORS` the compiler rejects with our
        // dedicated message before recursing into per-clause compilation.
        let code = generate_comprehension_with_generators(256);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "comprehension has too many nested clauses (256)");
    }

    #[test]
    fn many_generators_returns_syntax_error_not_stack_overflow() {
        // A comprehension with thousands of clauses used to crash the
        // compiler with a Rust stack overflow during MontyRun::new because
        // `compile_comprehension_generators` recursed once per clause with
        // no up-front guard.
        let code = generate_comprehension_with_generators(5000);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "comprehension has too many nested clauses (5000)");
    }
}

mod function_parameter_limits {
    use super::*;

    #[test]
    fn parameters_under_u8_limit_succeeds() {
        // 255 parameters should work - both definition and call
        let code = generate_many_parameters(255);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert!(result.is_ok(), "255 parameters should compile successfully");

        let run = result.unwrap();
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "255 parameters should run successfully");
    }

    #[test]
    fn parameters_at_u8_boundary_returns_syntax_error_for_call() {
        // 256 parameters - the function definition uses locals (wide instructions ok),
        // but the call site has 256 positional args which exceeds the limit
        let code = generate_many_parameters(256);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 positional arguments");
    }

    #[test]
    fn parameters_exceeding_u8_limit_returns_syntax_error_for_call() {
        // 257 parameters - same issue, call site has too many args
        let code = generate_many_parameters(257);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "more than 255 positional arguments");
    }
}

/// Generates a class with `count` simple `a{i} = {i}` class variables, plus an
/// assert reading the last one back through the class object.
///
/// Namespace assembly pushes two operand-stack entries (name const + value) per
/// member before `BuildDict` pops them all into the `type()` call's namespace
/// argument, so counts above 16383 overflow an `i16` stack-effect accumulator —
/// a regression guard for the i16→i32 widening.
fn generate_many_class_members(count: usize) -> String {
    let mut code = String::from("class C:\n");
    for i in 0..count {
        writeln!(code, "    a{i} = {i}").unwrap();
    }
    writeln!(code, "assert C.a{} == {}", count - 1, count - 1).unwrap();
    code
}

/// Generates `x = {0: 0, 1: 1, ...}` with `count` entries plus a read-back
/// assert — the dict-literal analogue of the class-member stack-effect case.
fn generate_large_dict_literal(count: usize) -> String {
    let mut code = String::from("x = {");
    for i in 0..count {
        if i > 0 {
            code.push_str(", ");
        }
        write!(code, "{i}: {i}").unwrap();
    }
    code.push('}');
    writeln!(code, "\nassert x[{}] == {}", count - 1, count - 1).unwrap();
    code
}

/// Generates one `try/finally` with many independent return sites.
fn generate_many_finally_return_sites(count: usize) -> String {
    let mut code = String::from("def f(x):\n    try:\n");
    for i in 0..count {
        writeln!(code, "        if x == {i}:\n            return {i}").unwrap();
    }
    code.push_str("    finally:\n        x = 0\n");
    code
}

/// Generates nested `try/finally` suites inside a return-path `finally`.
///
/// Each outer copy recompiles the nested suite, producing exponential copy
/// growth from source whose size is only linear in `depth`.
fn generate_nested_finally_suites(depth: usize) -> String {
    let mut code = String::from("def f():\n    try:\n        return 1\n    finally:\n");
    for level in 0..depth {
        let indent = "    ".repeat(level + 2);
        writeln!(code, "{indent}try:\n{indent}    pass\n{indent}finally:").unwrap();
    }
    writeln!(code, "{}pass\n\nassert f() == 1", "    ".repeat(depth + 2)).unwrap();
    code
}

mod stack_effect_limits {
    use super::*;

    #[test]
    fn class_members_above_i16_stack_effect() {
        // 16384 members -> 32768+ pushed operands, past i16::MAX (32767)
        let code = generate_many_class_members(16384);
        let run = MontyRun::new(code, "test.py", vec![], CompileOptions::default())
            .expect("16384 class members should compile");
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "16384 class members should run: {result:?}");
    }

    #[test]
    fn dict_literal_above_i16_stack_effect() {
        // 20000 entries -> 40000 pushed operands, past i16::MAX
        let code = generate_large_dict_literal(20000);
        let run = MontyRun::new(code, "test.py", vec![], CompileOptions::default())
            .expect("20000-entry dict literal should compile");
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "20000-entry dict literal should run: {result:?}");
    }
}

// `n` return sites emit `n + 2` finally copies: one per return site, plus the
// exception-path and fall-through copies every `try/finally` emits — hence
// the 1024-copy boundary sits at 1022/1023 sites.
mod finally_copy_limits {
    use super::*;

    /// The failure side of the boundary: 1023 sites are 1025 copies, one past
    /// `MAX_FINALLY_COPIES` (1024).
    #[test]
    fn excessive_inline_finally_copies_return_syntax_error() {
        let code = generate_many_finally_return_sites(1023);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "too many inline finally copies; maximum is 1024");
    }

    /// The success side of the boundary: 1022 sites are exactly
    /// `MAX_FINALLY_COPIES` (1024) copies, so an off-by-one in the limit
    /// guard or the constant cannot slip through the failure-only test above.
    #[test]
    fn max_inline_finally_copies_compile_and_run() {
        let code = generate_many_finally_return_sites(1022);
        let run = MontyRun::new(code, "test.py", vec![], CompileOptions::default())
            .expect("1024 inline finally copies should compile");
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "1024 inline finally copies should run: {result:?}");
    }

    /// Ten nested suites exceed the cap through repeated expansion even
    /// though the generated source contains only eleven `finally` statements.
    #[test]
    fn nested_finally_amplification_returns_syntax_error() {
        let code = generate_nested_finally_suites(10);
        let result = MontyRun::new(code, "test.py", vec![], CompileOptions::default());
        assert_syntax_error(result, "too many inline finally copies; maximum is 1024");
    }

    /// Nine nested suites produce 1,023 copies and exercise the return path,
    /// guarding the compact-source test against an overly conservative limit.
    #[test]
    fn nested_finally_amplification_below_limit_runs() {
        let code = generate_nested_finally_suites(9);
        let run = MontyRun::new(code, "test.py", vec![], CompileOptions::default())
            .expect("nested finally expansion below the copy limit should compile");
        let result = run.run_no_limits(vec![]);
        assert!(result.is_ok(), "nested finally expansion should run: {result:?}");
    }
}