RustAst/tests/optimizer.rs

use myc::ast::environment::Environment;

// ── Block simplification ─────────────────────────────────────────────────────

#[test]
fn test_empty_block_eliminated() {
    // (do) is semantically Void/Nop — the optimizer must collapse it,
    // even though nothing inside the block changes (block_changed stays false).
    let env = Environment::new();
    let dump = env.dump_ast("(do)").unwrap();
    assert!(
        !dump.contains("Block"),
        "Empty block must be collapsed to Nop. Dump:\n{}",
        dump
    );
}

#[test]
fn test_single_expr_block_unwrapped() {
    // (do 42) carries exactly one expression — the Block wrapper is redundant.
    // The optimizer must unwrap it to a bare Constant node.
    let env = Environment::new();
    let result = env.run_script("(do 42)").unwrap();
    assert_eq!(format!("{}", result), "42");
    let dump = env.dump_ast("(do 42)").unwrap();
    assert!(
        !dump.contains("Block"),
        "Single-expression block must be unwrapped. Dump:\n{}",
        dump
    );
}

#[test]
fn test_inner_empty_block_eliminated() {
    // The inner (do) must be collapsed to Nop during recursive optimization.
    // Nop-pruning in the outer block then removes it, leaving just Constant(42).
    let env = Environment::new();
    let result = env.run_script("(do (do) 42)").unwrap();
    assert_eq!(format!("{}", result), "42");
    let dump = env.dump_ast("(do (do) 42)").unwrap();
    assert!(
        !dump.contains("Block"),
        "Inner empty block must be eliminated entirely. Dump:\n{}",
        dump
    );
}

#[test]
fn test_optimizer_destructuring_inlining_and_mutation() {
    let env = Environment::new();

    // 1. Destructuring definition should allow inlining if not mutated
    let source_inline = "(do (def [x y] [10 20]) (+ x y))";
    let res_inline = env.run_script(source_inline).unwrap();
    assert_eq!(format!("{}", res_inline), "30");

    // 2. Destructuring definition should NOT be inlined if mutated (regression test)
    let source_mutation = r#"
    (do
      (def [a b] [1 2])
      (def f (fn [] (assign a (+ a b))))
      (f)
      a)
    "#;
    let res_mutation = env.run_script(source_mutation).unwrap();
    assert_eq!(format!("{}", res_mutation), "3");
}

#[test]
fn test_inline_two_lambdas_same_slots() {
    let env = Environment::new();
    // Two lambdas each binding a local at slot 0 get inlined into the same scope.
    // Slot remapping must prevent collision between them.
    let source = r#"
    (do
        (def add1 (fn [x] (+ x 1)))
        (def add2 (fn [x] (+ x 2)))
        (+ (add1 10) (add2 20)))
    "#;
    assert_eq!(format!("{}", env.run_script(source).unwrap()), "33");
    let dump = env.dump_ast(source).unwrap();
    assert!(
        dump.contains("Constant: 33"),
        "Should be fully folded to Constant 33. Dump:\n{}",
        dump
    );
}

#[test]
fn test_multipass_const_propagation() {
    let env = Environment::new();
    // A constant chain where each step is only resolvable after the previous has been folded.
    // Requires multiple optimizer passes.
    let source = r#"
    (do
        (def a 1)
        (def b (+ a 1))
        (def c (+ b 1))
        (+ c 1))
    "#;
    assert_eq!(format!("{}", env.run_script(source).unwrap()), "4");
    let dump = env.dump_ast(source).unwrap();
    assert!(
        dump.contains("Constant: 4"),
        "Multi-pass chain should fold to Constant 4. Dump:\n{}",
        dump
    );
}

#[test]
fn test_dead_def_destructuring_correctness() {
    // A dead destructuring def must not change the result — with or without optimization.
    let source = "(do (def [x y] [1 2]) 42)";
    let mut env_opt = Environment::new();
    env_opt.optimization = true;
    let mut env_no = Environment::new();
    env_no.optimization = false;
    assert_eq!(
        format!("{}", env_opt.run_script(source).unwrap()),
        format!("{}", env_no.run_script(source).unwrap())
    );
}

#[test]
fn test_dead_destructuring_def_eliminated() {
    let env = Environment::new();
    // A dead destructuring def with a pure RHS should be eliminated by the optimizer.
    let dump = env.dump_ast("(do (def [x y] [1 2]) 42)").unwrap();
    assert!(
        !dump.contains("Def"),
        "Dead destructuring def should be eliminated from AST. Dump:\n{}",
        dump
    );
}

#[test]
fn test_optimizer_inlining_slot_clash_repro() {
    let mut env = Environment::new();
    env.optimization = true;

    let source = r#"
    (do
      (def outer (fn [length]
        (do
          (def history (series 100))
          (fn [val] length)
        )
      ))
      ((outer 20) 5)
    )
    "#;

    let res = env.run_script(source).unwrap();
    assert_eq!(format!("{}", res), "20");
}

#[test]
fn test_reproduce_inlining_slot_clash_crash() {
    let mut env = Environment::new();
    env.optimization = true;

    let source = r#"
(do
  (def MY_SMA
    (fn [length]
      (do
        (def history (series 100))
        (fn [val]
           (do
              (push history val)
              (len history))))))

  (def src (create-random-ohlc 42 100))
  (def s (.close src))

  [(pipe [s] (MY_SMA 5))]
  [(pipe [s] (MY_SMA 20))])
"#;

    let res = env.run_script(source);
    assert!(res.is_ok(), "Inlining slot clash triggered: {:?}", res.err());
}

// ── Slot reuse ────────────────────────────────────────────────────────────────

#[test]
fn test_slot_reuse_non_overlapping_scopes() {
    let mut env = Environment::new();
    env.optimization = true;

    // Two sequential inner blocks each bind a mutable local.
    // assign prevents the optimizer from inlining the variables away,
    // so both slots survive into the lowering pass.
    // x lives only in the first block; once that block ends its slot is free.
    // y is defined in the second block and must reuse x's slot.
    // Expected result: the function returns 3 (y = 2+1).
    let source = r#"
    (fn []
      (do
        (do (def x 1) (assign x (+ x 1)) x)
        (do (def y 2) (assign y (+ y 1)) y)))
    "#;

    // Correctness: calling the function must still produce the right value.
    let result = env
        .run_script("((fn [] (do (do (def x 1) (assign x (+ x 1)) x) (do (def y 2) (assign y (+ y 1)) y))))")
        .unwrap();
    assert_eq!(format!("{}", result), "3");

    // Stack efficiency: with slot reuse the Lambda's stack_size must be 1,
    // not 2. The dump format is "stack_size: N" in the Lambda metadata line.
    let dump = env.dump_ast(source).unwrap();
    assert!(
        dump.contains("stack_size: 1"),
        "Non-overlapping scopes must reuse the same slot (expected stack_size: 1). Dump:\n{}",
        dump
    );
}

// ── Tuple type recomputation after inlining ─────────────────────────────────

#[test]
fn test_inlined_tuple_has_concrete_type_homogeneous() {
    // A polymorphic function returning a tuple, called with all-Int args.
    // After inlining, the tuple type must be concrete (no unresolved TypeVars).
    let env = Environment::new();
    let source = r#"(do
        (def wrap (fn [x y z] [x y z]))
        (wrap 1 2 3)
    )"#;
    let dump = env.dump_ast_compact(source).unwrap();
    assert!(
        !dump.contains('?'),
        "Inlined tuple must not contain unresolved TypeVars. Dump:\n{dump}"
    );
}

#[test]
fn test_inlined_tuple_has_concrete_type_heterogeneous() {
    // Heterogeneous case: Int + String should produce Tuple([int, str]), not ?0/?1.
    let env = Environment::new();
    let source = r#"(do
        (def wrap (fn [x y] [x y]))
        (wrap 1 "hello")
    )"#;
    let dump = env.dump_ast_compact(source).unwrap();
    assert!(
        !dump.contains('?'),
        "Inlined tuple must not contain unresolved TypeVars. Dump:\n{dump}"
    );
}

#[test]
fn test_block_and_lambda_propagate_concrete_type_after_inlining() {
    // When a polymorphic function is inlined and the result tuple gets a concrete type,
    // the enclosing Block and outer Lambda return type must propagate that concrete type
    // upward — not retain stale TypeVars from before inlining.
    let env = Environment::new();
    let source = r#"(do
        (def last (fn [s] (s 0)))
        (def f (fn [n] n))
        (def r (series 5))
        (push r 4)
        [(last f) (last r)]
    )"#;
    let dump = env.dump_ast_compact(source).unwrap();
    // The outer Lambda and Block must carry concrete types, not stale TypeVars.
    // Inner polymorphic lambdas (like `last`) legitimately retain TypeVars.
    let first_line = dump.lines().next().unwrap();
    assert!(
        !first_line.contains('?'),
        "Outer Lambda return type must be concrete. First line:\n{first_line}"
    );
    let block_line = dump.lines().find(|l| l.contains("Block")).unwrap();
    assert!(
        !block_line.contains('?'),
        "Block type must be concrete. Line:\n{block_line}"
    );
}

// ── Record type recomputation after inlining ────────────────────────────────

#[test]
fn test_inlined_record_has_concrete_type() {
    // A polymorphic function returning a record, called with a concrete arg.
    // After inlining + folding, the record type must be concrete (no TypeVars).
    let env = Environment::new();
    let source = r#"(do
        (def wrap (fn [x] {:val x}))
        (wrap 42)
    )"#;
    let dump = env.dump_ast_compact(source).unwrap();
    assert!(
        !dump.contains('?'),
        "Inlined record must not contain unresolved TypeVars. Dump:\n{dump}"
    );
}

#[test]
fn test_inlined_record_has_concrete_type_multi_field() {
    // Heterogeneous case: Int + String fields should produce concrete record type.
    let env = Environment::new();
    let source = r#"(do
        (def mkr (fn [a b] {:x a :y b}))
        (mkr 1 "hello")
    )"#;
    let dump = env.dump_ast_compact(source).unwrap();
    assert!(
        !dump.contains('?'),
        "Inlined multi-field record must not contain unresolved TypeVars. Dump:\n{dump}"
    );
}