seq-compiler 7.4.3

use super::*;
use crate::ast::{Program, Statement, WordDef};
use crate::types::{Effect, StackType, Type};

#[test]
fn test_simple_literal() {
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::IntLiteral(42)],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_simple_operation() {
    // : test ( Int Int -- Int ) add ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::WordCall {
                name: "i.add".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_type_mismatch() {
    // : test ( String -- ) io.write-line ;  with body: 42
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::String),
                StackType::Empty,
            )),
            body: vec![
                Statement::IntLiteral(42), // Pushes Int, not String!
                Statement::WordCall {
                    name: "io.write-line".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("Type mismatch"));
}

#[test]
fn test_polymorphic_dup() {
    // : my-dup ( Int -- Int Int ) dup ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "my-dup".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![Statement::WordCall {
                name: "dup".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_conditional_branches() {
    // : test ( Int Int -- String )
    //   > if "greater" else "not greater" then ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::String),
            )),
            body: vec![
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![Statement::StringLiteral(b"greater".to_vec())],
                    else_branch: Some(vec![Statement::StringLiteral(b"not greater".to_vec())]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_mismatched_branches() {
    // : test ( -- Int )
    //   true if 42 else "string" then ;  // ERROR: incompatible types
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::BoolLiteral(true),
                Statement::If {
                    then_branch: vec![Statement::IntLiteral(42)],
                    else_branch: Some(vec![Statement::StringLiteral(b"string".to_vec())]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("incompatible"));
}

#[test]
fn test_user_defined_word_call() {
    // : helper ( Int -- String ) int->string ;
    // : main ( -- ) 42 helper io.write-line ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "helper".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::String),
                )),
                body: vec![Statement::WordCall {
                    name: "int->string".to_string(),
                    span: None,
                }],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "main".to_string(),
                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
                body: vec![
                    Statement::IntLiteral(42),
                    Statement::WordCall {
                        name: "helper".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "io.write-line".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_arithmetic_chain() {
    // : test ( Int Int Int -- Int )
    //   add multiply ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.multiply".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_write_line_type_error() {
    // : test ( Int -- ) io.write-line ;  // ERROR: io.write-line expects String
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty,
            )),
            body: vec![Statement::WordCall {
                name: "io.write-line".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("Type mismatch"));
}

#[test]
fn test_stack_underflow_drop() {
    // : test ( -- ) drop ;  // ERROR: can't drop from empty stack
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![Statement::WordCall {
                name: "drop".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("mismatch"));
}

#[test]
fn test_stack_underflow_add() {
    // : test ( Int -- Int ) add ;  // ERROR: add needs 2 values
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::WordCall {
                name: "i.add".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("mismatch"));
}

/// Issue #169: rot with only 2 values should fail at compile time
/// Previously this was silently accepted due to implicit row polymorphism
#[test]
fn test_stack_underflow_rot_issue_169() {
    // : test ( -- ) 3 4 rot ;  // ERROR: rot needs 3 values, only 2 provided
    // Note: The parser generates `( ..rest -- ..rest )` for `( -- )`, so we use RowVar("rest")
    // to match the actual parsing behavior. The "rest" row variable is rigid.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::RowVar("rest".to_string()),
                StackType::RowVar("rest".to_string()),
            )),
            body: vec![
                Statement::IntLiteral(3),
                Statement::IntLiteral(4),
                Statement::WordCall {
                    name: "rot".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err(), "rot with 2 values should fail");
    let err = result.unwrap_err();
    assert!(
        err.contains("stack underflow") || err.contains("requires 3"),
        "Error should mention underflow: {}",
        err
    );
}

#[test]
fn test_csp_operations() {
    // : test ( -- )
    //   chan.make     # ( -- Channel )
    //   42 swap       # ( Channel Int -- Int Channel )
    //   chan.send     # ( Int Channel -- Bool )
    //   drop          # ( Bool -- )
    // ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::WordCall {
                    name: "chan.make".to_string(),
                    span: None,
                },
                Statement::IntLiteral(42),
                Statement::WordCall {
                    name: "swap".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "chan.send".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_complex_stack_shuffling() {
    // : test ( Int Int Int -- Int )
    //   rot add add ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "rot".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_empty_program() {
    // Program with no words should be valid
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_word_without_effect_declaration() {
    // : helper 42 ;  // No effect declaration - should error
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "helper".to_string(),
            effect: None,
            body: vec![Statement::IntLiteral(42)],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(
        result
            .unwrap_err()
            .contains("missing a stack effect declaration")
    );
}

#[test]
fn test_nested_conditionals() {
    // : test ( Int Int Int Int -- String )
    //   > if
    //     > if "both true" else "first true" then
    //   else
    //     drop drop "first false"
    //   then ;
    // Note: Needs 4 Ints total (2 for each > comparison)
    // Else branch must drop unused Ints to match then branch's stack effect
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::singleton(Type::String),
            )),
            body: vec![
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![
                        Statement::WordCall {
                            name: "i.>".to_string(),
                            span: None,
                        },
                        Statement::If {
                            then_branch: vec![Statement::StringLiteral(b"both true".to_vec())],
                            else_branch: Some(vec![Statement::StringLiteral(
                                b"first true".to_vec(),
                            )]),
                            span: None,
                        },
                    ],
                    else_branch: Some(vec![
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::StringLiteral(b"first false".to_vec()),
                    ]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    match checker.check_program(&program) {
        Ok(_) => {}
        Err(e) => panic!("Type check failed: {}", e),
    }
}

#[test]
fn test_conditional_without_else() {
    // : test ( Int Int -- Int )
    //   > if 100 then ;
    // Both branches must leave same stack
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![Statement::IntLiteral(100)],
                    else_branch: None, // No else - should leave stack unchanged
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    // This should fail because then pushes Int but else leaves stack empty
    assert!(result.is_err());
}

#[test]
fn test_multiple_word_chain() {
    // : helper1 ( Int -- String ) int->string ;
    // : helper2 ( String -- ) io.write-line ;
    // : main ( -- ) 42 helper1 helper2 ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "helper1".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::String),
                )),
                body: vec![Statement::WordCall {
                    name: "int->string".to_string(),
                    span: None,
                }],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "helper2".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::String),
                    StackType::Empty,
                )),
                body: vec![Statement::WordCall {
                    name: "io.write-line".to_string(),
                    span: None,
                }],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "main".to_string(),
                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
                body: vec![
                    Statement::IntLiteral(42),
                    Statement::WordCall {
                        name: "helper1".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "helper2".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_all_stack_ops() {
    // : test ( Int Int Int -- Int Int Int Int )
    //   over nip tuck ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "over".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "nip".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "tuck".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_mixed_types_complex() {
    // : test ( -- )
    //   42 int->string      # ( -- String )
    //   100 200 >           # ( String -- String Int )
    //   if                  # ( String -- String )
    //     io.write-line     # ( String -- )
    //   else
    //     io.write-line
    //   then ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::IntLiteral(42),
                Statement::WordCall {
                    name: "int->string".to_string(),
                    span: None,
                },
                Statement::IntLiteral(100),
                Statement::IntLiteral(200),
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![Statement::WordCall {
                        name: "io.write-line".to_string(),
                        span: None,
                    }],
                    else_branch: Some(vec![Statement::WordCall {
                        name: "io.write-line".to_string(),
                        span: None,
                    }]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_string_literal() {
    // : test ( -- String ) "hello" ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::String),
            )),
            body: vec![Statement::StringLiteral(b"hello".to_vec())],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_bool_literal() {
    // : test ( -- Bool ) true ;
    // Booleans are now properly typed as Bool
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Bool),
            )),
            body: vec![Statement::BoolLiteral(true)],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_type_error_in_nested_conditional() {
    // : test ( -- )
    //   10 20 i.> if
    //     42 io.write-line   # ERROR: io.write-line expects String, got Int
    //   else
    //     "ok" io.write-line
    //   then ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::IntLiteral(10),
                Statement::IntLiteral(20),
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![
                        Statement::IntLiteral(42),
                        Statement::WordCall {
                            name: "io.write-line".to_string(),
                            span: None,
                        },
                    ],
                    else_branch: Some(vec![
                        Statement::StringLiteral(b"ok".to_vec()),
                        Statement::WordCall {
                            name: "io.write-line".to_string(),
                            span: None,
                        },
                    ]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    assert!(result.unwrap_err().contains("Type mismatch"));
}

#[test]
fn test_read_line_operation() {
    // : test ( -- String Bool ) io.read-line ;
    // io.read-line now returns ( -- String Bool ) for error handling
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::from_vec(vec![Type::String, Type::Bool]),
            )),
            body: vec![Statement::WordCall {
                name: "io.read-line".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_comparison_operations() {
    // Test all comparison operators
    // : test ( Int Int -- Bool )
    //   i.<= ;
    // Simplified: just test that comparisons work and return Bool
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::Bool),
            )),
            body: vec![Statement::WordCall {
                name: "i.<=".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_recursive_word_definitions() {
    // Test mutually recursive words (type checking only, no runtime)
    // : is-even ( Int -- Int ) dup 0 = if drop 1 else 1 subtract is-odd then ;
    // : is-odd ( Int -- Int ) dup 0 = if drop 0 else 1 subtract is-even then ;
    //
    // Note: This tests that the checker can handle words that reference each other
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "is-even".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::WordCall {
                        name: "dup".to_string(),
                        span: None,
                    },
                    Statement::IntLiteral(0),
                    Statement::WordCall {
                        name: "i.=".to_string(),
                        span: None,
                    },
                    Statement::If {
                        then_branch: vec![
                            Statement::WordCall {
                                name: "drop".to_string(),
                                span: None,
                            },
                            Statement::IntLiteral(1),
                        ],
                        else_branch: Some(vec![
                            Statement::IntLiteral(1),
                            Statement::WordCall {
                                name: "i.subtract".to_string(),
                                span: None,
                            },
                            Statement::WordCall {
                                name: "is-odd".to_string(),
                                span: None,
                            },
                        ]),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "is-odd".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::WordCall {
                        name: "dup".to_string(),
                        span: None,
                    },
                    Statement::IntLiteral(0),
                    Statement::WordCall {
                        name: "i.=".to_string(),
                        span: None,
                    },
                    Statement::If {
                        then_branch: vec![
                            Statement::WordCall {
                                name: "drop".to_string(),
                                span: None,
                            },
                            Statement::IntLiteral(0),
                        ],
                        else_branch: Some(vec![
                            Statement::IntLiteral(1),
                            Statement::WordCall {
                                name: "i.subtract".to_string(),
                                span: None,
                            },
                            Statement::WordCall {
                                name: "is-even".to_string(),
                                span: None,
                            },
                        ]),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_word_calling_word_with_row_polymorphism() {
    // Test that row variables unify correctly through word calls
    // : apply-twice ( Int -- Int ) dup add ;
    // : quad ( Int -- Int ) apply-twice apply-twice ;
    // Should work: both use row polymorphism correctly
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "apply-twice".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::WordCall {
                        name: "dup".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "i.add".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "quad".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::WordCall {
                        name: "apply-twice".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "apply-twice".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_deep_stack_types() {
    // Test with many values on stack (10+ items)
    // : test ( Int Int Int Int Int Int Int Int Int Int -- Int )
    //   add add add add add add add add add ;
    let mut stack_type = StackType::Empty;
    for _ in 0..10 {
        stack_type = stack_type.push(Type::Int);
    }

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(stack_type, StackType::singleton(Type::Int))),
            body: vec![
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_simple_quotation() {
    // : test ( -- Quot )
    //   [ 1 add ] ;
    // Quotation type should be [ ..input Int -- ..input Int ] (row polymorphic)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::RowVar("input".to_string()).push(Type::Int),
                    StackType::RowVar("input".to_string()).push(Type::Int),
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![
                    Statement::IntLiteral(1),
                    Statement::WordCall {
                        name: "i.add".to_string(),
                        span: None,
                    },
                ],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    match checker.check_program(&program) {
        Ok(_) => {}
        Err(e) => panic!("Type check failed: {}", e),
    }
}

#[test]
fn test_empty_quotation() {
    // : test ( -- Quot )
    //   [ ] ;
    // Empty quotation has effect ( ..input -- ..input ) (preserves stack)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::RowVar("input".to_string()),
                    StackType::RowVar("input".to_string()),
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 1,
                body: vec![],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_nested_quotation() {
    // : test ( -- Quot )
    //   [ [ 1 add ] ] ;
    // Outer quotation contains inner quotation (both row-polymorphic)
    let inner_quot_type = Type::Quotation(Box::new(Effect::new(
        StackType::RowVar("input".to_string()).push(Type::Int),
        StackType::RowVar("input".to_string()).push(Type::Int),
    )));

    let outer_quot_type = Type::Quotation(Box::new(Effect::new(
        StackType::RowVar("input".to_string()),
        StackType::RowVar("input".to_string()).push(inner_quot_type.clone()),
    )));

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(outer_quot_type),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 2,
                body: vec![Statement::Quotation {
                    span: None,
                    id: 3,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.add".to_string(),
                            span: None,
                        },
                    ],
                }],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_invalid_field_type_error() {
    use crate::ast::{UnionDef, UnionField, UnionVariant};

    let program = Program {
        includes: vec![],
        unions: vec![UnionDef {
            name: "Message".to_string(),
            variants: vec![UnionVariant {
                name: "Get".to_string(),
                fields: vec![UnionField {
                    name: "chan".to_string(),
                    type_name: "InvalidType".to_string(),
                }],
                source: None,
            }],
            source: None,
        }],
        words: vec![],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(err.contains("Unknown type 'InvalidType'"));
    assert!(err.contains("chan"));
    assert!(err.contains("Get"));
    assert!(err.contains("Message"));
}

#[test]
fn test_roll_inside_conditional_with_concrete_stack() {
    // Bug #93: n roll inside if/else should work when stack has enough concrete items
    // : test ( Int Int Int Int -- Int Int Int Int )
    //   dup 0 > if
    //     3 roll    # Works: 4 concrete items available
    //   else
    //     rot rot   # Alternative that also works
    //   then ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "dup".to_string(),
                    span: None,
                },
                Statement::IntLiteral(0),
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![
                        Statement::IntLiteral(3),
                        Statement::WordCall {
                            name: "roll".to_string(),
                            span: None,
                        },
                    ],
                    else_branch: Some(vec![
                        Statement::WordCall {
                            name: "rot".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "rot".to_string(),
                            span: None,
                        },
                    ]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    // This should now work because both branches have 4 concrete items
    match checker.check_program(&program) {
        Ok(_) => {}
        Err(e) => panic!("Type check failed: {}", e),
    }
}

#[test]
fn test_roll_inside_match_arm_with_concrete_stack() {
    // Similar to bug #93 but for match arms: n roll inside match should work
    // when stack has enough concrete items from the match context
    use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};

    // Define a simple union: union Result = Ok | Err
    let union_def = UnionDef {
        name: "Result".to_string(),
        variants: vec![
            UnionVariant {
                name: "Ok".to_string(),
                fields: vec![],
                source: None,
            },
            UnionVariant {
                name: "Err".to_string(),
                fields: vec![],
                source: None,
            },
        ],
        source: None,
    };

    // : test ( Int Int Int Int Result -- Int Int Int Int )
    //   match
    //     Ok => 3 roll
    //     Err => rot rot
    //   end ;
    let program = Program {
        includes: vec![],
        unions: vec![union_def],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Union("Result".to_string())),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![Statement::Match {
                arms: vec![
                    MatchArm {
                        pattern: Pattern::Variant("Ok".to_string()),
                        body: vec![
                            Statement::IntLiteral(3),
                            Statement::WordCall {
                                name: "roll".to_string(),
                                span: None,
                            },
                        ],
                        span: None,
                    },
                    MatchArm {
                        pattern: Pattern::Variant("Err".to_string()),
                        body: vec![
                            Statement::WordCall {
                                name: "rot".to_string(),
                                span: None,
                            },
                            Statement::WordCall {
                                name: "rot".to_string(),
                                span: None,
                            },
                        ],
                        span: None,
                    },
                ],
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    match checker.check_program(&program) {
        Ok(_) => {}
        Err(e) => panic!("Type check failed: {}", e),
    }
}

#[test]
fn test_roll_with_row_polymorphic_input() {
    // roll reaching into row variable should work (needed for stdlib)
    // : test ( T U V W -- U V W T )
    //   3 roll ;   # Rotates: brings position 3 to top
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Var("T".to_string()))
                    .push(Type::Var("U".to_string()))
                    .push(Type::Var("V".to_string()))
                    .push(Type::Var("W".to_string())),
                StackType::Empty
                    .push(Type::Var("U".to_string()))
                    .push(Type::Var("V".to_string()))
                    .push(Type::Var("W".to_string()))
                    .push(Type::Var("T".to_string())),
            )),
            body: vec![
                Statement::IntLiteral(3),
                Statement::WordCall {
                    name: "roll".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_ok(), "roll test failed: {:?}", result.err());
}

#[test]
fn test_pick_with_row_polymorphic_input() {
    // pick reaching into row variable should work (needed for stdlib)
    // : test ( T U V -- T U V T )
    //   2 pick ;   # Copies element at index 2 (0-indexed from top)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Var("T".to_string()))
                    .push(Type::Var("U".to_string()))
                    .push(Type::Var("V".to_string())),
                StackType::Empty
                    .push(Type::Var("T".to_string()))
                    .push(Type::Var("U".to_string()))
                    .push(Type::Var("V".to_string()))
                    .push(Type::Var("T".to_string())),
            )),
            body: vec![
                Statement::IntLiteral(2),
                Statement::WordCall {
                    name: "pick".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_valid_union_reference_in_field() {
    use crate::ast::{UnionDef, UnionField, UnionVariant};

    let program = Program {
        includes: vec![],
        unions: vec![
            UnionDef {
                name: "Inner".to_string(),
                variants: vec![UnionVariant {
                    name: "Val".to_string(),
                    fields: vec![UnionField {
                        name: "x".to_string(),
                        type_name: "Int".to_string(),
                    }],
                    source: None,
                }],
                source: None,
            },
            UnionDef {
                name: "Outer".to_string(),
                variants: vec![UnionVariant {
                    name: "Wrap".to_string(),
                    fields: vec![UnionField {
                        name: "inner".to_string(),
                        type_name: "Inner".to_string(), // Reference to other union
                    }],
                    source: None,
                }],
                source: None,
            },
        ],
        words: vec![],
    };

    let mut checker = TypeChecker::new();
    assert!(
        checker.check_program(&program).is_ok(),
        "Union reference in field should be valid"
    );
}

#[test]
fn test_divergent_recursive_tail_call() {
    // Test that recursive tail calls in if/else branches are recognized as divergent.
    // This pattern is common in actor loops:
    //
    // : store-loop ( Channel -- )
    //   dup           # ( chan chan )
    //   chan.receive  # ( chan value Bool )
    //   not if        # ( chan value )
    //     drop        # ( chan ) - drop value, keep chan for recursion
    //     store-loop  # diverges - never returns
    //   then
    //   # else: ( chan value ) - process msg normally
    //   drop drop     # ( )
    // ;
    //
    // The then branch ends with a recursive call (store-loop), so it diverges.
    // The else branch (implicit empty) continues with the stack after the if.
    // Without divergent branch detection, this would fail because:
    //   - then branch produces: () (after drop store-loop)
    //   - else branch produces: (chan value)
    // But since then diverges, we should use else's type.

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "store-loop".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Channel), // ( Channel -- )
                StackType::Empty,
            )),
            body: vec![
                // dup -> ( chan chan )
                Statement::WordCall {
                    name: "dup".to_string(),
                    span: None,
                },
                // chan.receive -> ( chan value Bool )
                Statement::WordCall {
                    name: "chan.receive".to_string(),
                    span: None,
                },
                // not -> ( chan value Bool )
                Statement::WordCall {
                    name: "not".to_string(),
                    span: None,
                },
                // if drop store-loop then
                Statement::If {
                    then_branch: vec![
                        // drop value -> ( chan )
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        // store-loop -> diverges
                        Statement::WordCall {
                            name: "store-loop".to_string(), // recursive tail call
                            span: None,
                        },
                    ],
                    else_branch: None, // implicit else continues with ( chan value )
                    span: None,
                },
                // After if: ( chan value ) - drop both
                Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Divergent recursive tail call should be accepted: {:?}",
        result.err()
    );
}

#[test]
fn test_divergent_else_branch() {
    // Test that divergence detection works for else branches too.
    //
    // : process-loop ( Channel -- )
    //   dup chan.receive   # ( chan value Bool )
    //   if                 # ( chan value )
    //     drop drop        # normal exit: ( )
    //   else
    //     drop             # ( chan )
    //     process-loop     # diverges - retry on failure
    //   then
    // ;

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "process-loop".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Channel), // ( Channel -- )
                StackType::Empty,
            )),
            body: vec![
                Statement::WordCall {
                    name: "dup".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "chan.receive".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![
                        // success: drop value and chan
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                    ],
                    else_branch: Some(vec![
                        // failure: drop value, keep chan, recurse
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "process-loop".to_string(), // recursive tail call
                            span: None,
                        },
                    ]),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Divergent else branch should be accepted: {:?}",
        result.err()
    );
}

#[test]
fn test_non_tail_call_recursion_not_divergent() {
    // Test that recursion NOT in tail position is not treated as divergent.
    // This should fail type checking because after the recursive call,
    // there's more code that changes the stack.
    //
    // : bad-loop ( Int -- Int )
    //   dup 0 i.> if
    //     1 i.subtract bad-loop  # recursive call
    //     1 i.add                # more code after - not tail position!
    //   then
    // ;
    //
    // This should fail because:
    // - then branch: recurse then add 1 -> stack changes after recursion
    // - else branch (implicit): stack is ( Int )
    // Without proper handling, this could incorrectly pass.

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "bad-loop".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "dup".to_string(),
                    span: None,
                },
                Statement::IntLiteral(0),
                Statement::WordCall {
                    name: "i.>".to_string(),
                    span: None,
                },
                Statement::If {
                    then_branch: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.subtract".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "bad-loop".to_string(), // NOT in tail position
                            span: None,
                        },
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.add".to_string(), // code after recursion
                            span: None,
                        },
                    ],
                    else_branch: None,
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    // This should pass because the branches ARE compatible:
    // - then: produces Int (after bad-loop returns Int, then add 1)
    // - else: produces Int (from the dup at start)
    // The key is that bad-loop is NOT in tail position, so it's not divergent.
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Non-tail recursion should type check normally: {:?}",
        result.err()
    );
}

#[test]
fn test_call_yield_quotation_error() {
    // Phase 2c: Calling a quotation with Yield effect directly should error.
    // : bad ( W -- W ) [ yield ] call ;
    // This is wrong because yield quotations must be wrapped with strand.weave.
    // (W is a single-letter type variable — multi-character names like `Ctx`
    // are no longer accepted as ad-hoc type vars under the strict rule.)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "bad".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Var("W".to_string())),
                StackType::singleton(Type::Var("W".to_string())),
            )),
            body: vec![
                // Push a dummy value that will be yielded
                Statement::IntLiteral(42),
                Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "yield".to_string(),
                        span: None,
                    }],
                },
                Statement::WordCall {
                    name: "call".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Calling yield quotation directly should fail"
    );
    let err = result.unwrap_err();
    assert!(
        err.contains("Yield") || err.contains("strand.weave"),
        "Error should mention Yield or strand.weave: {}",
        err
    );
}

#[test]
fn test_strand_weave_yield_quotation_ok() {
    // Phase 2c: Using strand.weave on a Yield quotation is correct.
    // : good ( -- Int H ) 42 [ yield ] strand.weave ;
    // (single-letter type variable — multi-character names like `Handle`
    // are no longer accepted as ad-hoc type vars under the strict rule.)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "good".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Var("H".to_string())),
            )),
            body: vec![
                Statement::IntLiteral(42),
                Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "yield".to_string(),
                        span: None,
                    }],
                },
                Statement::WordCall {
                    name: "strand.weave".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "strand.weave on yield quotation should pass: {:?}",
        result.err()
    );
}

#[test]
fn test_call_pure_quotation_ok() {
    // Phase 2c: Calling a pure quotation (no Yield) is fine.
    // : ok ( Int -- Int ) [ 1 i.add ] call ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "ok".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.add".to_string(),
                            span: None,
                        },
                    ],
                },
                Statement::WordCall {
                    name: "call".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Calling pure quotation should pass: {:?}",
        result.err()
    );
}

// ==========================================================================
// Stack Pollution Detection Tests (Issue #228)
// These tests verify the type checker catches stack effect mismatches
// ==========================================================================

#[test]
fn test_pollution_extra_push() {
    // : test ( Int -- Int ) 42 ;
    // Declares consuming 1 Int, producing 1 Int
    // But body pushes 42 on top of input, leaving 2 values
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::IntLiteral(42)],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: declares ( Int -- Int ) but leaves 2 values on stack"
    );
}

#[test]
fn test_pollution_extra_dup() {
    // : test ( Int -- Int ) dup ;
    // Declares producing 1 Int, but dup produces 2
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::WordCall {
                name: "dup".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: declares ( Int -- Int ) but dup produces 2 values"
    );
}

#[test]
fn test_pollution_consumes_extra() {
    // : test ( Int -- Int ) drop drop 42 ;
    // Declares consuming 1 Int, but body drops twice
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                },
                Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                },
                Statement::IntLiteral(42),
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: declares ( Int -- Int ) but consumes 2 values"
    );
}

#[test]
fn test_pollution_in_then_branch() {
    // : test ( Bool -- Int )
    //   if 1 2 else 3 then ;
    // Then branch pushes 2 values, else pushes 1
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Bool),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::If {
                then_branch: vec![
                    Statement::IntLiteral(1),
                    Statement::IntLiteral(2), // Extra value!
                ],
                else_branch: Some(vec![Statement::IntLiteral(3)]),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: then branch pushes 2 values, else pushes 1"
    );
}

#[test]
fn test_pollution_in_else_branch() {
    // : test ( Bool -- Int )
    //   if 1 else 2 3 then ;
    // Then branch pushes 1, else pushes 2 values
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Bool),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::If {
                then_branch: vec![Statement::IntLiteral(1)],
                else_branch: Some(vec![
                    Statement::IntLiteral(2),
                    Statement::IntLiteral(3), // Extra value!
                ]),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: then branch pushes 1 value, else pushes 2"
    );
}

#[test]
fn test_pollution_both_branches_extra() {
    // : test ( Bool -- Int )
    //   if 1 2 else 3 4 then ;
    // Both branches push 2 values but declared output is 1
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Bool),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::If {
                then_branch: vec![Statement::IntLiteral(1), Statement::IntLiteral(2)],
                else_branch: Some(vec![Statement::IntLiteral(3), Statement::IntLiteral(4)]),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: both branches push 2 values, but declared output is 1"
    );
}

#[test]
fn test_pollution_branch_consumes_extra() {
    // : test ( Bool Int -- Int )
    //   if drop drop 1 else then ;
    // Then branch consumes more than available from declared inputs
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Bool).push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::If {
                then_branch: vec![
                    Statement::WordCall {
                        name: "drop".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "drop".to_string(),
                        span: None,
                    },
                    Statement::IntLiteral(1),
                ],
                else_branch: Some(vec![]),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: then branch consumes Bool (should only have Int after if)"
    );
}

#[test]
fn test_pollution_quotation_wrong_arity_output() {
    // : test ( Int -- Int )
    //   [ dup ] call ;
    // Quotation produces 2 values, but word declares 1 output
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "dup".to_string(),
                        span: None,
                    }],
                },
                Statement::WordCall {
                    name: "call".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: quotation [dup] produces 2 values, declared output is 1"
    );
}

#[test]
fn test_pollution_quotation_wrong_arity_input() {
    // : test ( Int -- Int )
    //   [ drop drop 42 ] call ;
    // Quotation consumes 2 values, but only 1 available
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "drop".to_string(),
                            span: None,
                        },
                        Statement::IntLiteral(42),
                    ],
                },
                Statement::WordCall {
                    name: "call".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Should reject: quotation [drop drop 42] consumes 2 values, only 1 available"
    );
}

#[test]
fn test_missing_effect_provides_helpful_error() {
    // : myword 42 ;
    // No effect annotation - should error with helpful message including word name
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "myword".to_string(),
            effect: None, // No annotation
            body: vec![Statement::IntLiteral(42)],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(err.contains("myword"), "Error should mention word name");
    assert!(
        err.contains("stack effect"),
        "Error should mention stack effect"
    );
}

#[test]
fn test_valid_effect_exact_match() {
    // : test ( Int Int -- Int ) i.+ ;
    // Exact match - consumes 2, produces 1
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::WordCall {
                name: "i.add".to_string(),
                span: None,
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_ok(), "Should accept: effect matches exactly");
}

#[test]
fn test_valid_polymorphic_passthrough() {
    // : test ( a -- a ) ;
    // Identity function - row polymorphism allows this
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Cons {
                    rest: Box::new(StackType::RowVar("rest".to_string())),
                    top: Type::Var("a".to_string()),
                },
                StackType::Cons {
                    rest: Box::new(StackType::RowVar("rest".to_string())),
                    top: Type::Var("a".to_string()),
                },
            )),
            body: vec![], // Empty body - just pass through
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_ok(), "Should accept: polymorphic identity");
}

// ==========================================================================
// Closure Nesting Tests (Issue #230)
// Tests for deep closure nesting, transitive captures, and edge cases
// ==========================================================================

#[test]
fn test_closure_basic_capture() {
    // : make-adder ( Int -- Closure )
    //   [ i.+ ] ;
    // The quotation needs 2 Ints (for i.+) but caller will only provide 1
    // So it captures 1 Int from the creation site
    // Must declare as Closure type to trigger capture analysis
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "make-adder".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Closure {
                    effect: Box::new(Effect::new(
                        StackType::RowVar("r".to_string()).push(Type::Int),
                        StackType::RowVar("r".to_string()).push(Type::Int),
                    )),
                    captures: vec![Type::Int], // Captures 1 Int
                }),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![Statement::WordCall {
                    name: "i.add".to_string(),
                    span: None,
                }],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Basic closure capture should work: {:?}",
        result.err()
    );
}

#[test]
fn test_closure_nested_two_levels() {
    // : outer ( -- Quot )
    //   [ [ 1 i.+ ] ] ;
    // Outer quotation: no inputs, just returns inner quotation
    // Inner quotation: pushes 1 then adds (needs 1 Int from caller)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "outer".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::RowVar("r".to_string()),
                    StackType::RowVar("r".to_string()).push(Type::Quotation(Box::new(
                        Effect::new(
                            StackType::RowVar("s".to_string()).push(Type::Int),
                            StackType::RowVar("s".to_string()).push(Type::Int),
                        ),
                    ))),
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![Statement::Quotation {
                    span: None,
                    id: 1,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.add".to_string(),
                            span: None,
                        },
                    ],
                }],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Two-level nested quotations should work: {:?}",
        result.err()
    );
}

#[test]
fn test_closure_nested_three_levels() {
    // : deep ( -- Quot )
    //   [ [ [ 1 i.+ ] ] ] ;
    // Three levels of nesting, innermost does actual work
    let inner_effect = Effect::new(
        StackType::RowVar("a".to_string()).push(Type::Int),
        StackType::RowVar("a".to_string()).push(Type::Int),
    );
    let middle_effect = Effect::new(
        StackType::RowVar("b".to_string()),
        StackType::RowVar("b".to_string()).push(Type::Quotation(Box::new(inner_effect))),
    );
    let outer_effect = Effect::new(
        StackType::RowVar("c".to_string()),
        StackType::RowVar("c".to_string()).push(Type::Quotation(Box::new(middle_effect))),
    );

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "deep".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(outer_effect))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![Statement::Quotation {
                    span: None,
                    id: 1,
                    body: vec![Statement::Quotation {
                        span: None,
                        id: 2,
                        body: vec![
                            Statement::IntLiteral(1),
                            Statement::WordCall {
                                name: "i.add".to_string(),
                                span: None,
                            },
                        ],
                    }],
                }],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Three-level nested quotations should work: {:?}",
        result.err()
    );
}

#[test]
fn test_closure_use_after_creation() {
    // : use-adder ( -- Int )
    //   5 make-adder   // Creates closure capturing 5
    //   10 swap call ; // Calls closure with 10, should return 15
    //
    // Tests that closure is properly typed when called later
    let adder_type = Type::Closure {
        effect: Box::new(Effect::new(
            StackType::RowVar("r".to_string()).push(Type::Int),
            StackType::RowVar("r".to_string()).push(Type::Int),
        )),
        captures: vec![Type::Int],
    };

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "make-adder".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(adder_type.clone()),
                )),
                body: vec![Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "i.add".to_string(),
                        span: None,
                    }],
                }],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "use-adder".to_string(),
                effect: Some(Effect::new(
                    StackType::Empty,
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::IntLiteral(5),
                    Statement::WordCall {
                        name: "make-adder".to_string(),
                        span: None,
                    },
                    Statement::IntLiteral(10),
                    Statement::WordCall {
                        name: "swap".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "call".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Closure usage after creation should work: {:?}",
        result.err()
    );
}

#[test]
fn test_closure_wrong_call_type() {
    // : bad-use ( -- Int )
    //   5 make-adder   // Creates Int -> Int closure
    //   "hello" swap call ; // Tries to call with String - should fail!
    let adder_type = Type::Closure {
        effect: Box::new(Effect::new(
            StackType::RowVar("r".to_string()).push(Type::Int),
            StackType::RowVar("r".to_string()).push(Type::Int),
        )),
        captures: vec![Type::Int],
    };

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![
            WordDef {
                name: "make-adder".to_string(),
                effect: Some(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(adder_type.clone()),
                )),
                body: vec![Statement::Quotation {
                    span: None,
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "i.add".to_string(),
                        span: None,
                    }],
                }],
                source: None,
                allowed_lints: vec![],
            },
            WordDef {
                name: "bad-use".to_string(),
                effect: Some(Effect::new(
                    StackType::Empty,
                    StackType::singleton(Type::Int),
                )),
                body: vec![
                    Statement::IntLiteral(5),
                    Statement::WordCall {
                        name: "make-adder".to_string(),
                        span: None,
                    },
                    Statement::StringLiteral(b"hello".to_vec()), // Wrong type!
                    Statement::WordCall {
                        name: "swap".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "call".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_err(),
        "Calling Int closure with String should fail"
    );
}

#[test]
fn test_closure_multiple_captures() {
    // : make-between ( Int Int -- Quot )
    //   [ dup rot i.>= swap rot i.<= and ] ;
    // Captures both min and max, checks if value is between them
    // Body needs: value min max (3 Ints)
    // Caller provides: value (1 Int)
    // Captures: min max (2 Ints)
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "make-between".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::RowVar("r".to_string()).push(Type::Int),
                    StackType::RowVar("r".to_string()).push(Type::Bool),
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![
                    // Simplified: just do a comparison that uses all 3 values
                    Statement::WordCall {
                        name: "i.>=".to_string(),
                        span: None,
                    },
                    // Note: This doesn't match the comment but tests multi-capture
                ],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    // This should work - the quotation body uses values from stack
    // The exact behavior depends on how captures are inferred
    // For now, we're testing that it doesn't crash
    assert!(
        result.is_ok() || result.is_err(),
        "Multiple captures should be handled (pass or fail gracefully)"
    );
}

#[test]
fn test_quotation_type_preserved_through_word() {
    // : identity-quot ( Quot -- Quot ) ;
    // Tests that quotation types are preserved when passed through words
    let quot_type = Type::Quotation(Box::new(Effect::new(
        StackType::RowVar("r".to_string()).push(Type::Int),
        StackType::RowVar("r".to_string()).push(Type::Int),
    )));

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "identity-quot".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(quot_type.clone()),
                StackType::singleton(quot_type.clone()),
            )),
            body: vec![], // Identity - just return what's on stack
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Quotation type should be preserved through identity word: {:?}",
        result.err()
    );
}

#[test]
fn test_closure_captures_value_for_inner_quotation() {
    // : make-inner-adder ( Int -- Closure )
    //   [ [ i.+ ] swap call ] ;
    // The closure captures an Int
    // When called, it creates an inner quotation and calls it with the captured value
    // This tests that closures can work with nested quotations
    let closure_effect = Effect::new(
        StackType::RowVar("r".to_string()).push(Type::Int),
        StackType::RowVar("r".to_string()).push(Type::Int),
    );

    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "make-inner-adder".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Closure {
                    effect: Box::new(closure_effect),
                    captures: vec![Type::Int],
                }),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![
                    // The captured Int and the caller's Int are on stack
                    Statement::WordCall {
                        name: "i.add".to_string(),
                        span: None,
                    },
                ],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "Closure with capture for inner work should pass: {:?}",
        result.err()
    );
}

#[test]
fn test_union_type_mismatch_should_fail() {
    // RFC #345: Different union types should not unify
    // This tests that passing union B to a function expecting union A fails
    use crate::ast::{UnionDef, UnionField, UnionVariant};

    let mut program = Program {
        includes: vec![],
        unions: vec![
            UnionDef {
                name: "UnionA".to_string(),
                variants: vec![UnionVariant {
                    name: "AVal".to_string(),
                    fields: vec![UnionField {
                        name: "x".to_string(),
                        type_name: "Int".to_string(),
                    }],
                    source: None,
                }],
                source: None,
            },
            UnionDef {
                name: "UnionB".to_string(),
                variants: vec![UnionVariant {
                    name: "BVal".to_string(),
                    fields: vec![UnionField {
                        name: "y".to_string(),
                        type_name: "Int".to_string(),
                    }],
                    source: None,
                }],
                source: None,
            },
        ],
        words: vec![
            // : takes-a ( UnionA -- ) drop ;
            WordDef {
                name: "takes-a".to_string(),
                effect: Some(Effect::new(
                    StackType::RowVar("rest".to_string()).push(Type::Union("UnionA".to_string())),
                    StackType::RowVar("rest".to_string()),
                )),
                body: vec![Statement::WordCall {
                    name: "drop".to_string(),
                    span: None,
                }],
                source: None,
                allowed_lints: vec![],
            },
            // : main ( -- ) 99 Make-BVal takes-a ;
            // This should FAIL: Make-BVal returns UnionB, takes-a expects UnionA
            WordDef {
                name: "main".to_string(),
                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
                body: vec![
                    Statement::IntLiteral(99),
                    Statement::WordCall {
                        name: "Make-BVal".to_string(),
                        span: None,
                    },
                    Statement::WordCall {
                        name: "takes-a".to_string(),
                        span: None,
                    },
                ],
                source: None,
                allowed_lints: vec![],
            },
        ],
    };

    // Generate constructors (Make-AVal, Make-BVal)
    program.generate_constructors().unwrap();

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);

    // This SHOULD fail because UnionB != UnionA
    assert!(
        result.is_err(),
        "Passing UnionB to function expecting UnionA should fail, but got: {:?}",
        result
    );
    let err = result.unwrap_err();
    assert!(
        err.contains("Union") || err.contains("mismatch"),
        "Error should mention union type mismatch, got: {}",
        err
    );
}

// =========================================================================
// Aux stack tests (Issue #350)
// =========================================================================

fn make_word_call(name: &str) -> Statement {
    Statement::WordCall {
        name: name.to_string(),
        span: None,
    }
}

#[test]
fn test_aux_basic_round_trip() {
    // : test ( Int -- Int ) >aux aux> ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![make_word_call(">aux"), make_word_call("aux>")],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_aux_preserves_type() {
    // : test ( String -- String ) >aux aux> ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::String),
                StackType::singleton(Type::String),
            )),
            body: vec![make_word_call(">aux"), make_word_call("aux>")],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_aux_unbalanced_error() {
    // : test ( Int -- ) >aux ;  -- ERROR: aux not empty at return
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty,
            )),
            body: vec![make_word_call(">aux")],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("aux stack is not empty"),
        "Expected aux stack balance error, got: {}",
        err
    );
}

#[test]
fn test_aux_pop_empty_error() {
    // : test ( -- Int ) aux> ;  -- ERROR: aux is empty
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Int),
            )),
            body: vec![make_word_call("aux>")],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("aux stack is empty"),
        "Expected aux empty error, got: {}",
        err
    );
}

#[test]
fn test_aux_multiple_values() {
    // >aux >aux aux> aux> is identity (LIFO preserves order)
    // Input: ( Int String ), >aux pops String, >aux pops Int
    // aux> pushes Int, aux> pushes String → output: ( Int String )
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::String),
                StackType::Empty.push(Type::Int).push(Type::String),
            )),
            body: vec![
                make_word_call(">aux"),
                make_word_call(">aux"),
                make_word_call("aux>"),
                make_word_call("aux>"),
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_aux_max_depths_tracked() {
    // : test ( Int -- Int ) >aux aux> ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![make_word_call(">aux"), make_word_call("aux>")],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    checker.check_program(&program).unwrap();
    let depths = checker.take_aux_max_depths();
    assert_eq!(depths.get("test"), Some(&1));
}

#[test]
fn test_aux_in_match_balanced() {
    // Aux used symmetrically in match arms: each arm does >aux aux>
    use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};

    let union_def = UnionDef {
        name: "Choice".to_string(),
        variants: vec![
            UnionVariant {
                name: "Left".to_string(),
                fields: vec![],
                source: None,
            },
            UnionVariant {
                name: "Right".to_string(),
                fields: vec![],
                source: None,
            },
        ],
        source: None,
    };

    // : test ( Int Choice -- Int )
    //   swap >aux match Left => aux> end Right => aux> end ;
    let program = Program {
        includes: vec![],
        unions: vec![union_def],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Union("Choice".to_string())),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                make_word_call("swap"),
                make_word_call(">aux"),
                Statement::Match {
                    arms: vec![
                        MatchArm {
                            pattern: Pattern::Variant("Left".to_string()),
                            body: vec![make_word_call("aux>")],
                            span: None,
                        },
                        MatchArm {
                            pattern: Pattern::Variant("Right".to_string()),
                            body: vec![make_word_call("aux>")],
                            span: None,
                        },
                    ],
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_aux_in_match_unbalanced_error() {
    // Match arms with different aux states should error
    use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};

    let union_def = UnionDef {
        name: "Choice".to_string(),
        variants: vec![
            UnionVariant {
                name: "Left".to_string(),
                fields: vec![],
                source: None,
            },
            UnionVariant {
                name: "Right".to_string(),
                fields: vec![],
                source: None,
            },
        ],
        source: None,
    };

    // : test ( Int Choice -- Int )
    //   swap >aux match Left => aux> end Right => end ;  -- ERROR: unbalanced
    let program = Program {
        includes: vec![],
        unions: vec![union_def],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Union("Choice".to_string())),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                make_word_call("swap"),
                make_word_call(">aux"),
                Statement::Match {
                    arms: vec![
                        MatchArm {
                            pattern: Pattern::Variant("Left".to_string()),
                            body: vec![make_word_call("aux>")],
                            span: None,
                        },
                        MatchArm {
                            pattern: Pattern::Variant("Right".to_string()),
                            body: vec![],
                            span: None,
                        },
                    ],
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("aux stack"),
        "Expected aux stack mismatch error, got: {}",
        err
    );
}

#[test]
fn test_aux_in_quotation_balanced_accepted() {
    // Issue #393: balanced >aux/aux> inside a quotation is now allowed.
    // The word produces a quotation [ Int -- Int ]. The quotation body
    // uses >aux/aux> to round-trip the input Int. Lexical scoping is
    // preserved because both ops are inside the same quotation.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::singleton(Type::Int),
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![make_word_call(">aux"), make_word_call("aux>")],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(
        result.is_ok(),
        "should accept balanced aux in quotation: {:?}",
        result.err()
    );
}

#[test]
fn test_aux_in_quotation_unbalanced_rejected() {
    // Issue #393: a quotation that pushes onto aux without popping
    // must still be rejected. Lexical scoping requires every >aux to
    // be matched by an aux> within the same quotation.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Quotation(Box::new(Effect::new(
                    StackType::singleton(Type::Int),
                    StackType::Empty,
                )))),
            )),
            body: vec![Statement::Quotation {
                span: None,
                id: 0,
                body: vec![make_word_call(">aux")],
            }],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("unbalanced aux stack"),
        "Expected unbalanced aux error, got: {}",
        err
    );
}

// =========================================================================
// Dataflow combinator tests
// =========================================================================

#[test]
fn test_dip_basic() {
    // : test ( Int Int -- Int Int )  [ 1 i.+ ] dip ;
    // Increments value below top, preserving top
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.+".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_dip_type_mismatch() {
    // : test ( String Int -- ?? )  [ 1 i.+ ] dip ;
    // Should fail: quotation expects Int but gets String
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::String).push(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.+".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_err());
}

#[test]
fn test_keep_basic() {
    // : test ( Int -- Int Int )  [ dup i.* ] keep ;
    // Squares and keeps original
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::WordCall {
                            name: "dup".to_string(),
                            span: None,
                        },
                        Statement::WordCall {
                            name: "i.*".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "keep".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_bi_basic() {
    // : test ( Int -- Int Int )  [ 2 i.* ] [ 3 i.* ] bi ;
    // Double and triple
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(2),
                        Statement::WordCall {
                            name: "i.*".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![
                        Statement::IntLiteral(3),
                        Statement::WordCall {
                            name: "i.*".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "bi".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_keep_type_mismatch() {
    // : test ( String -- ?? )  [ 1 i.+ ] keep ;
    // Should fail: quotation expects Int but gets String
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::String),
                StackType::Empty.push(Type::Int).push(Type::String),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.+".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "keep".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_err());
}

#[test]
fn test_bi_type_mismatch() {
    // : test ( String -- ?? )  [ string.length ] [ 1 i.+ ] bi ;
    // Should fail: second quotation expects Int but value is String
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::String),
                StackType::Empty.push(Type::Int).push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "string.length".to_string(),
                        span: None,
                    }],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.+".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "bi".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_err());
}

#[test]
fn test_dip_underflow() {
    // : test ( -- ?? )  [ 1 ] dip ;
    // Should fail: dip needs a value below the quotation
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty,
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::IntLiteral(1)],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("stack underflow"),
        "Expected underflow error, got: {}",
        err
    );
}

#[test]
fn test_dip_preserves_type() {
    // : test ( Int String -- Int String )  [ 1 i.+ ] dip ;
    // The String on top is preserved, Int below is incremented
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::String),
                StackType::Empty.push(Type::Int).push(Type::String),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(1),
                        Statement::WordCall {
                            name: "i.+".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_keep_underflow() {
    // : test ( -- ?? )  [ drop ] keep ;
    // Should fail: keep needs a value below the quotation
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::WordCall {
                        name: "drop".to_string(),
                        span: None,
                    }],
                    span: None,
                },
                Statement::WordCall {
                    name: "keep".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("stack underflow") || err.contains("underflow"),
        "Expected underflow error, got: {}",
        err
    );
}

#[test]
fn test_bi_underflow() {
    // : test ( -- ?? )  [ 1 ] [ 2 ] bi ;
    // Should fail: bi needs a value below the two quotations
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::IntLiteral(1)],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![Statement::IntLiteral(2)],
                    span: None,
                },
                Statement::WordCall {
                    name: "bi".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("stack underflow") || err.contains("underflow"),
        "Expected underflow error, got: {}",
        err
    );
}

// =========================================================================
// variant.* type safety
//
// Regression coverage for the `variant.*` builtins requiring Type::Variant
// (or a Union via the Union <: Variant relaxation) on the stack.
// =========================================================================

fn check_word_with_body(
    name: &str,
    inputs: StackType,
    outputs: StackType,
    body: Vec<Statement>,
) -> Result<(), String> {
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: name.to_string(),
            effect: Some(Effect::new(inputs, outputs)),
            body,
            source: None,
            allowed_lints: vec![],
        }],
    };
    let mut checker = TypeChecker::new();
    checker.check_program(&program).map(|_| ())
}

fn variant_call(name: &str) -> Statement {
    Statement::WordCall {
        name: name.to_string(),
        span: None,
    }
}

#[test]
fn test_variant_field_at_rejects_string() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::singleton(Type::Int),
        vec![
            Statement::StringLiteral(b"alpha beta".to_vec()),
            Statement::IntLiteral(0),
            variant_call("variant.field-at"),
        ],
    );
    let err = result.expect_err("expected type error for String -> variant.field-at");
    assert!(
        err.contains("variant.field-at") && err.contains("String"),
        "error should name variant.field-at and String, got: {err}"
    );
}

#[test]
fn test_variant_tag_rejects_int() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::singleton(Type::Symbol),
        vec![Statement::IntLiteral(7), variant_call("variant.tag")],
    );
    let err = result.expect_err("expected type error for Int -> variant.tag");
    assert!(err.contains("variant.tag"));
}

#[test]
fn test_variant_field_count_rejects_string() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::singleton(Type::Int),
        vec![
            Statement::StringLiteral(b"x".to_vec()),
            variant_call("variant.field-count"),
        ],
    );
    let err = result.expect_err("expected type error for String -> variant.field-count");
    assert!(
        err.contains("variant.field-count"),
        "error should name variant.field-count, got: {err}"
    );
}

#[test]
fn test_variant_init_rejects_string() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::singleton(Type::Variant),
        vec![
            Statement::StringLiteral(b"x".to_vec()),
            variant_call("variant.init"),
        ],
    );
    let err = result.expect_err("expected type error for String -> variant.init");
    assert!(
        err.contains("variant.init"),
        "error should name variant.init, got: {err}"
    );
}

#[test]
fn test_variant_append_rejects_string_base() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::singleton(Type::Variant),
        vec![
            Statement::StringLiteral(b"x".to_vec()),
            Statement::IntLiteral(1),
            variant_call("variant.append"),
        ],
    );
    let err = result.expect_err("expected type error for String -> variant.append");
    assert!(
        err.contains("variant.append"),
        "error should name variant.append, got: {err}"
    );
}

#[test]
fn test_variant_last_rejects_string() {
    let result = check_word_with_body(
        "bad",
        StackType::Empty,
        StackType::Empty.push(Type::Var("T".to_string())),
        vec![
            Statement::StringLiteral(b"x".to_vec()),
            variant_call("variant.last"),
        ],
    );
    let err = result.expect_err("expected type error for String -> variant.last");
    assert!(
        err.contains("variant.last"),
        "error should name variant.last, got: {err}"
    );
}

#[test]
fn test_union_value_accepted_by_variant_field_at() {
    // Direct exercise of the Union(_) <: Variant relaxation rule:
    // a value typed `Union("Box")` should be accepted by variant.field-at
    // (which is signed against `Variant`).
    let union_def = crate::ast::UnionDef {
        name: "Box".to_string(),
        variants: vec![crate::ast::UnionVariant {
            name: "Cell".to_string(),
            fields: vec![crate::ast::UnionField {
                name: "x".to_string(),
                type_name: "Int".to_string(),
            }],
            source: None,
        }],
        source: None,
    };
    let program = Program {
        includes: vec![],
        unions: vec![union_def],
        words: vec![WordDef {
            name: "first".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Union("Box".to_string())),
                StackType::singleton(Type::Int),
            )),
            body: vec![Statement::IntLiteral(0), variant_call("variant.field-at")],
            source: None,
            allowed_lints: vec![],
        }],
    };
    let mut checker = TypeChecker::new();
    checker
        .check_program(&program)
        .expect("Union(Box) should be accepted by variant.field-at");
}

#[test]
fn test_variant_make_then_field_at_typechecks() {
    let result = check_word_with_body(
        "ok",
        StackType::Empty,
        StackType::singleton(Type::Int),
        vec![
            Statement::IntLiteral(42),
            Statement::Symbol("Foo".to_string()),
            variant_call("variant.make-1"),
            Statement::IntLiteral(0),
            variant_call("variant.field-at"),
        ],
    );
    assert!(
        result.is_ok(),
        "variant.make-1 -> variant.field-at should typecheck, got: {:?}",
        result
    );
}

#[test]
fn test_bi_polymorphic_quotations() {
    // : test ( Int -- Int String )  [ 2 i.* ] [ int->string ] bi ;
    // Two quotations with different output types — verifies independent typing
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::Empty.push(Type::Int).push(Type::String),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::IntLiteral(2),
                        Statement::WordCall {
                            name: "i.*".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![Statement::WordCall {
                        name: "int->string".to_string(),
                        span: None,
                    }],
                    span: None,
                },
                Statement::WordCall {
                    name: "bi".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

// `if`-combinator tests. These exercise `infer_if_combinator` directly
// (the dynamic-dispatch path) rather than going through the AST
// normalization pass — `check_program` doesn't run normalization, so a
// `[Quotation, Quotation, WordCall("if")]` AST reaches the combinator
// inference path the same way runtime-built quotations would.

#[test]
fn test_if_combinator_basic() {
    // : test ( Bool -- Int )  [ 1 ] [ 2 ] if ;
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Bool),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::IntLiteral(1)],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![Statement::IntLiteral(2)],
                    span: None,
                },
                Statement::WordCall {
                    name: "if".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_if_combinator_branch_mismatch() {
    // : test ( Bool -- ?? )  [ 1 ] [ "string" ] if ;
    // Should fail: branches produce different output types.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Bool),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::IntLiteral(1)],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![Statement::StringLiteral(b"string".to_vec())],
                    span: None,
                },
                Statement::WordCall {
                    name: "if".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("incompatible") || err.contains("unify"),
        "Expected branch-mismatch error, got: {}",
        err
    );
}

// =============================================================================
// Issue #471: row-polymorphic refinement inside quotation bodies.
//
// Before the polymorphic_pop fix, a combinator (dip/keep/bi/if) that needs to
// pop a value from below a quotation would fail when the surrounding stack was
// a bare row variable — which is the case inside any quotation body whose
// incoming stack is unconstrained. These tests pin the new behavior down: the
// flexible row variable refines to "row + fresh T", letting nested combinators
// type-check, while the rigid `..rest` row variable from a declared signature
// stays rigid so genuine underflow is still rejected.
// =============================================================================

#[test]
fn test_nested_dip_in_dip_body() {
    // : test ( Int Int Int -- Int Int Int )  [ [ 1 i.+ ] dip ] dip ;
    // The outer dip preserves the top Int; the inner dip (inside the
    // outer quotation body, where the stack starts as a row variable)
    // must polymorphic-pop its preserved value out of that row.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::Quotation {
                            id: 1,
                            body: vec![
                                Statement::IntLiteral(1),
                                Statement::WordCall {
                                    name: "i.+".to_string(),
                                    span: None,
                                },
                            ],
                            span: None,
                        },
                        Statement::WordCall {
                            name: "dip".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_keep_inside_dip_body() {
    // : test ( Int Int -- Int Int Int )  [ [ 1 i.+ ] keep ] dip ;
    // Outer dip preserves top Int; body uses `keep`, which both passes
    // its value to the quotation and preserves it. Body effect:
    // ( ..r Int -- ..r Int Int ).
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::Quotation {
                            id: 1,
                            body: vec![
                                Statement::IntLiteral(1),
                                Statement::WordCall {
                                    name: "i.+".to_string(),
                                    span: None,
                                },
                            ],
                            span: None,
                        },
                        Statement::WordCall {
                            name: "keep".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_dip_inside_keep_body() {
    // : test ( Int Int -- Int Int Int )  [ [ 1 i.+ ] dip ] keep ;
    // Keep's quotation runs on a stack that still has the preserved
    // value on top; the inner dip preserves that value while running
    // [ 1 i.+ ] on what's below.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::Quotation {
                            id: 1,
                            body: vec![
                                Statement::IntLiteral(1),
                                Statement::WordCall {
                                    name: "i.+".to_string(),
                                    span: None,
                                },
                            ],
                            span: None,
                        },
                        Statement::WordCall {
                            name: "dip".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "keep".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_bi_inside_dip_body() {
    // : test ( Int Int -- Int Int Int )  [ [ 1 i.+ ] [ 2 i.+ ] bi ] dip ;
    // Body uses `bi` on the row-variable input the outer dip exposed.
    // `bi` peeks a preserved value out of that row, just like dip/keep.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::Empty.push(Type::Int).push(Type::Int),
                StackType::Empty
                    .push(Type::Int)
                    .push(Type::Int)
                    .push(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![
                        Statement::Quotation {
                            id: 1,
                            body: vec![
                                Statement::IntLiteral(1),
                                Statement::WordCall {
                                    name: "i.+".to_string(),
                                    span: None,
                                },
                            ],
                            span: None,
                        },
                        Statement::Quotation {
                            id: 2,
                            body: vec![
                                Statement::IntLiteral(2),
                                Statement::WordCall {
                                    name: "i.+".to_string(),
                                    span: None,
                                },
                            ],
                            span: None,
                        },
                        Statement::WordCall {
                            name: "bi".to_string(),
                            span: None,
                        },
                    ],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    assert!(checker.check_program(&program).is_ok());
}

#[test]
fn test_dip_rigid_rest_underflow_still_rejected() {
    // : test ( -- )  [ ] dip ;
    // The outer signature is `( -- )` — `..rest` is rigid (user
    // declared no inputs). `polymorphic_pop` must refuse to refine
    // a rigid row variable; this is the soundness boundary. Without
    // the rigidity check the prototype silently accepted this, so
    // the test pins the contract down.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![],
                    span: None,
                },
                Statement::WordCall {
                    name: "dip".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("stack underflow"),
        "Expected underflow error at rigid signature boundary, got: {}",
        err
    );
}

#[test]
fn test_if_combinator_non_bool_condition() {
    // : test ( Int -- ?? )  [ 1 ] [ 2 ] if ;
    // Should fail: condition is Int, not Bool.
    let program = Program {
        includes: vec![],
        unions: vec![],
        words: vec![WordDef {
            name: "test".to_string(),
            effect: Some(Effect::new(
                StackType::singleton(Type::Int),
                StackType::singleton(Type::Int),
            )),
            body: vec![
                Statement::Quotation {
                    id: 0,
                    body: vec![Statement::IntLiteral(1)],
                    span: None,
                },
                Statement::Quotation {
                    id: 1,
                    body: vec![Statement::IntLiteral(2)],
                    span: None,
                },
                Statement::WordCall {
                    name: "if".to_string(),
                    span: None,
                },
            ],
            source: None,
            allowed_lints: vec![],
        }],
    };

    let mut checker = TypeChecker::new();
    let result = checker.check_program(&program);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        err.contains("Bool") || err.contains("condition"),
        "Expected non-Bool condition error, got: {}",
        err
    );
}