Skip to main content

seqc/
typechecker.rs

1//! Enhanced type checker for Seq with full type tracking
2//!
3//! Uses row polymorphism and unification to verify stack effects.
4//! Based on cem2's type checker but simplified for Phase 8.5.
5
6use crate::ast::{Program, Statement, WordDef};
7use crate::builtins::builtin_signature;
8use crate::call_graph::CallGraph;
9use crate::capture_analysis::{calculate_captures, extract_concrete_types};
10use crate::types::{
11    Effect, SideEffect, StackType, Type, UnionTypeInfo, VariantFieldInfo, VariantInfo,
12};
13use crate::unification::{Subst, unify_stacks, unify_types};
14use std::collections::HashMap;
15
16/// Format a line number as an error message prefix (e.g., "at line 42: ").
17/// Line numbers are 0-indexed internally, so we add 1 for display.
18fn format_line_prefix(line: usize) -> String {
19    format!("at line {}: ", line + 1)
20}
21
22/// Validate that `main` has an allowed signature (Issue #355).
23///
24/// Only `( -- )` and `( -- Int )` are accepted. The first is "void main"
25/// (process exits with 0). The second is "int main" (return value is the
26/// process exit code).
27///
28/// Any other shape — extra inputs, multiple outputs, non-Int output —
29/// is rejected with an actionable error.
30fn validate_main_effect(effect: &Effect) -> Result<(), String> {
31    // Inputs must be empty (just the row var, no concrete types)
32    let inputs_ok = matches!(&effect.inputs, StackType::Empty | StackType::RowVar(_));
33
34    // Outputs: either empty (void main) or exactly one Int (int main)
35    let outputs_ok = match &effect.outputs {
36        StackType::Empty | StackType::RowVar(_) => true,
37        StackType::Cons {
38            rest,
39            top: Type::Int,
40        } if matches!(**rest, StackType::Empty | StackType::RowVar(_)) => true,
41        _ => false,
42    };
43
44    if inputs_ok && outputs_ok {
45        return Ok(());
46    }
47
48    Err(format!(
49        "Word 'main' has an invalid stack effect: ( {} -- {} ).\n\
50         `main` must be declared with one of:\n\
51           ( -- )       — void main, process exits with code 0\n\
52           ( -- Int )   — exit code is the returned Int\n\
53         Other shapes are not allowed.",
54        effect.inputs, effect.outputs
55    ))
56}
57
58pub struct TypeChecker {
59    /// Environment mapping word names to their effects
60    env: HashMap<String, Effect>,
61    /// Union type registry - maps union names to their type information
62    /// Contains variant names and field types for each union
63    unions: HashMap<String, UnionTypeInfo>,
64    /// Counter for generating fresh type variables
65    fresh_counter: std::cell::Cell<usize>,
66    /// Quotation types tracked during type checking
67    /// Maps quotation ID (from AST) to inferred type (Quotation or Closure)
68    /// This type map is used by codegen to generate appropriate code
69    quotation_types: std::cell::RefCell<HashMap<usize, Type>>,
70    /// Expected quotation/closure type (from word signature, if any)
71    /// Used during type-driven capture inference
72    expected_quotation_type: std::cell::RefCell<Option<Type>>,
73    /// Current word being type-checked (for detecting recursive tail calls)
74    /// Used to identify divergent branches in if/else expressions
75    /// Stores (name, line_number) for better error messages
76    current_word: std::cell::RefCell<Option<(String, Option<usize>)>>,
77    /// Per-statement type info for codegen optimization (Issue #186)
78    /// Maps (word_name, statement_index) -> concrete top-of-stack type before statement
79    /// Only stores trivially-copyable types (Int, Float, Bool) to enable optimizations
80    statement_top_types: std::cell::RefCell<HashMap<(String, usize), Type>>,
81    /// Call graph for detecting mutual recursion (Issue #229)
82    /// Used to improve divergent branch detection beyond direct recursion
83    call_graph: Option<CallGraph>,
84    /// Current aux stack type during word body checking (Issue #350)
85    /// Tracked per-word; reset to Empty at each word boundary.
86    current_aux_stack: std::cell::RefCell<StackType>,
87    /// Maximum aux stack depth per word, for codegen alloca sizing (Issue #350)
88    /// Maps word_name -> max_depth (number of %Value allocas needed)
89    aux_max_depths: std::cell::RefCell<HashMap<String, usize>>,
90    /// Maximum aux stack depth per quotation, for codegen alloca sizing (Issue #393)
91    /// Maps quotation_id -> max_depth (number of %Value allocas needed)
92    /// Quotation IDs are program-wide unique, assigned by the parser.
93    quotation_aux_depths: std::cell::RefCell<HashMap<usize, usize>>,
94    /// Stack of currently-active quotation IDs during type checking (Issue #393).
95    /// Pushed when entering `infer_quotation`, popped on exit. The top of the
96    /// stack is the innermost quotation. Empty means we're in word body scope.
97    /// Replaces the old `in_quotation_scope` boolean.
98    quotation_id_stack: std::cell::RefCell<Vec<usize>>,
99    /// Resolved arithmetic sugar: maps (line, column) -> concrete op name.
100    /// Keyed by source location, which is unique per occurrence and available
101    /// to both the typechecker and codegen via the AST span.
102    resolved_sugar: std::cell::RefCell<HashMap<(usize, usize), String>>,
103}
104
105impl TypeChecker {
106    pub fn new() -> Self {
107        TypeChecker {
108            env: HashMap::new(),
109            unions: HashMap::new(),
110            fresh_counter: std::cell::Cell::new(0),
111            quotation_types: std::cell::RefCell::new(HashMap::new()),
112            expected_quotation_type: std::cell::RefCell::new(None),
113            current_word: std::cell::RefCell::new(None),
114            statement_top_types: std::cell::RefCell::new(HashMap::new()),
115            call_graph: None,
116            current_aux_stack: std::cell::RefCell::new(StackType::Empty),
117            aux_max_depths: std::cell::RefCell::new(HashMap::new()),
118            quotation_aux_depths: std::cell::RefCell::new(HashMap::new()),
119            quotation_id_stack: std::cell::RefCell::new(Vec::new()),
120            resolved_sugar: std::cell::RefCell::new(HashMap::new()),
121        }
122    }
123
124    /// Set the call graph for mutual recursion detection.
125    ///
126    /// When set, the type checker can detect divergent branches caused by
127    /// mutual recursion (e.g., even/odd pattern) in addition to direct recursion.
128    pub fn set_call_graph(&mut self, call_graph: CallGraph) {
129        self.call_graph = Some(call_graph);
130    }
131
132    /// Get line info prefix for error messages (e.g., "at line 42: " or "")
133    fn line_prefix(&self) -> String {
134        self.current_word
135            .borrow()
136            .as_ref()
137            .and_then(|(_, line)| line.map(format_line_prefix))
138            .unwrap_or_default()
139    }
140
141    /// Look up a union type by name
142    pub fn get_union(&self, name: &str) -> Option<&UnionTypeInfo> {
143        self.unions.get(name)
144    }
145
146    /// Get all registered union types
147    pub fn get_unions(&self) -> &HashMap<String, UnionTypeInfo> {
148        &self.unions
149    }
150
151    /// Find variant info by name across all unions
152    ///
153    /// Returns (union_name, variant_info) for the variant
154    fn find_variant(&self, variant_name: &str) -> Option<(&str, &VariantInfo)> {
155        for (union_name, union_info) in &self.unions {
156            for variant in &union_info.variants {
157                if variant.name == variant_name {
158                    return Some((union_name.as_str(), variant));
159                }
160            }
161        }
162        None
163    }
164
165    /// Register external word effects (e.g., from included modules or FFI).
166    ///
167    /// All external words must have explicit stack effects for type safety.
168    pub fn register_external_words(&mut self, words: &[(&str, &Effect)]) {
169        for (name, effect) in words {
170            self.env.insert(name.to_string(), (*effect).clone());
171        }
172    }
173
174    /// Register external union type names (e.g., from included modules).
175    ///
176    /// This allows field types in union definitions to reference types from includes.
177    /// We only register the name as a valid type; we don't need full variant info
178    /// since the actual union definition lives in the included file.
179    pub fn register_external_unions(&mut self, union_names: &[&str]) {
180        for name in union_names {
181            // Insert a placeholder union with no variants
182            // This makes is_valid_type_name() return true for this type
183            self.unions.insert(
184                name.to_string(),
185                UnionTypeInfo {
186                    name: name.to_string(),
187                    variants: vec![],
188                },
189            );
190        }
191    }
192
193    /// Extract the type map (quotation ID -> inferred type)
194    ///
195    /// This should be called after check_program() to get the inferred types
196    /// for all quotations in the program. The map is used by codegen to generate
197    /// appropriate code for Quotations vs Closures.
198    pub fn take_quotation_types(&self) -> HashMap<usize, Type> {
199        self.quotation_types.replace(HashMap::new())
200    }
201
202    /// Extract per-statement type info for codegen optimization (Issue #186)
203    /// Returns map of (word_name, statement_index) -> top-of-stack type
204    pub fn take_statement_top_types(&self) -> HashMap<(String, usize), Type> {
205        self.statement_top_types.replace(HashMap::new())
206    }
207
208    /// Extract resolved arithmetic sugar for codegen
209    /// Maps (line, column) -> concrete operation name
210    pub fn take_resolved_sugar(&self) -> HashMap<(usize, usize), String> {
211        self.resolved_sugar.replace(HashMap::new())
212    }
213
214    /// Extract per-word aux stack max depths for codegen alloca sizing (Issue #350)
215    pub fn take_aux_max_depths(&self) -> HashMap<String, usize> {
216        self.aux_max_depths.replace(HashMap::new())
217    }
218
219    /// Extract per-quotation aux stack max depths for codegen alloca sizing (Issue #393)
220    /// Maps quotation_id -> max_depth
221    pub fn take_quotation_aux_depths(&self) -> HashMap<usize, usize> {
222        self.quotation_aux_depths.replace(HashMap::new())
223    }
224
225    /// Count the number of concrete types in a StackType (for aux depth tracking)
226    fn stack_depth(stack: &StackType) -> usize {
227        let mut depth = 0;
228        let mut current = stack;
229        while let StackType::Cons { rest, .. } = current {
230            depth += 1;
231            current = rest;
232        }
233        depth
234    }
235
236    /// Check if the top of the stack is a trivially-copyable type (Int, Float, Bool)
237    /// These types have no heap references and can be memcpy'd in codegen.
238    fn get_trivially_copyable_top(stack: &StackType) -> Option<Type> {
239        match stack {
240            StackType::Cons { top, .. } => match top {
241                Type::Int | Type::Float | Type::Bool => Some(top.clone()),
242                _ => None,
243            },
244            _ => None,
245        }
246    }
247
248    /// Record the top-of-stack type for a statement if it's trivially copyable (Issue #186)
249    fn capture_statement_type(&self, word_name: &str, stmt_index: usize, stack: &StackType) {
250        if let Some(top_type) = Self::get_trivially_copyable_top(stack) {
251            self.statement_top_types
252                .borrow_mut()
253                .insert((word_name.to_string(), stmt_index), top_type);
254        }
255    }
256
257    /// Generate a fresh variable name
258    fn fresh_var(&self, prefix: &str) -> String {
259        let n = self.fresh_counter.get();
260        self.fresh_counter.set(n + 1);
261        format!("{}${}", prefix, n)
262    }
263
264    /// Freshen all type and row variables in an effect
265    fn freshen_effect(&self, effect: &Effect) -> Effect {
266        let mut type_map = HashMap::new();
267        let mut row_map = HashMap::new();
268
269        let fresh_inputs = self.freshen_stack(&effect.inputs, &mut type_map, &mut row_map);
270        let fresh_outputs = self.freshen_stack(&effect.outputs, &mut type_map, &mut row_map);
271
272        // Freshen the side effects too
273        let fresh_effects = effect
274            .effects
275            .iter()
276            .map(|e| self.freshen_side_effect(e, &mut type_map, &mut row_map))
277            .collect();
278
279        Effect::with_effects(fresh_inputs, fresh_outputs, fresh_effects)
280    }
281
282    fn freshen_side_effect(
283        &self,
284        effect: &SideEffect,
285        type_map: &mut HashMap<String, String>,
286        row_map: &mut HashMap<String, String>,
287    ) -> SideEffect {
288        match effect {
289            SideEffect::Yield(ty) => {
290                SideEffect::Yield(Box::new(self.freshen_type(ty, type_map, row_map)))
291            }
292        }
293    }
294
295    fn freshen_stack(
296        &self,
297        stack: &StackType,
298        type_map: &mut HashMap<String, String>,
299        row_map: &mut HashMap<String, String>,
300    ) -> StackType {
301        match stack {
302            StackType::Empty => StackType::Empty,
303            StackType::RowVar(name) => {
304                let fresh_name = row_map
305                    .entry(name.clone())
306                    .or_insert_with(|| self.fresh_var(name));
307                StackType::RowVar(fresh_name.clone())
308            }
309            StackType::Cons { rest, top } => {
310                let fresh_rest = self.freshen_stack(rest, type_map, row_map);
311                let fresh_top = self.freshen_type(top, type_map, row_map);
312                StackType::Cons {
313                    rest: Box::new(fresh_rest),
314                    top: fresh_top,
315                }
316            }
317        }
318    }
319
320    fn freshen_type(
321        &self,
322        ty: &Type,
323        type_map: &mut HashMap<String, String>,
324        row_map: &mut HashMap<String, String>,
325    ) -> Type {
326        match ty {
327            Type::Int | Type::Float | Type::Bool | Type::String | Type::Symbol | Type::Channel => {
328                ty.clone()
329            }
330            Type::Var(name) => {
331                let fresh_name = type_map
332                    .entry(name.clone())
333                    .or_insert_with(|| self.fresh_var(name));
334                Type::Var(fresh_name.clone())
335            }
336            Type::Quotation(effect) => {
337                let fresh_inputs = self.freshen_stack(&effect.inputs, type_map, row_map);
338                let fresh_outputs = self.freshen_stack(&effect.outputs, type_map, row_map);
339                Type::Quotation(Box::new(Effect::new(fresh_inputs, fresh_outputs)))
340            }
341            Type::Closure { effect, captures } => {
342                let fresh_inputs = self.freshen_stack(&effect.inputs, type_map, row_map);
343                let fresh_outputs = self.freshen_stack(&effect.outputs, type_map, row_map);
344                let fresh_captures = captures
345                    .iter()
346                    .map(|t| self.freshen_type(t, type_map, row_map))
347                    .collect();
348                Type::Closure {
349                    effect: Box::new(Effect::new(fresh_inputs, fresh_outputs)),
350                    captures: fresh_captures,
351                }
352            }
353            // Union types are concrete named types - no freshening needed
354            Type::Union(name) => Type::Union(name.clone()),
355        }
356    }
357
358    /// Parse a type name string into a Type
359    ///
360    /// Supports: Int, Float, Bool, String, Channel, and union types
361    fn parse_type_name(&self, name: &str) -> Type {
362        match name {
363            "Int" => Type::Int,
364            "Float" => Type::Float,
365            "Bool" => Type::Bool,
366            "String" => Type::String,
367            "Channel" => Type::Channel,
368            // Any other name is assumed to be a union type reference
369            other => Type::Union(other.to_string()),
370        }
371    }
372
373    /// Check if a type name is a known valid type
374    ///
375    /// Returns true for built-in types (Int, Float, Bool, String, Channel) and
376    /// registered union type names
377    fn is_valid_type_name(&self, name: &str) -> bool {
378        matches!(name, "Int" | "Float" | "Bool" | "String" | "Channel")
379            || self.unions.contains_key(name)
380    }
381
382    /// Validate that all field types in union definitions reference known types
383    ///
384    /// Note: Field count validation happens earlier in generate_constructors()
385    fn validate_union_field_types(&self, program: &Program) -> Result<(), String> {
386        for union_def in &program.unions {
387            for variant in &union_def.variants {
388                for field in &variant.fields {
389                    if !self.is_valid_type_name(&field.type_name) {
390                        return Err(format!(
391                            "Unknown type '{}' in field '{}' of variant '{}' in union '{}'. \
392                             Valid types are: Int, Float, Bool, String, Channel, or a defined union name.",
393                            field.type_name, field.name, variant.name, union_def.name
394                        ));
395                    }
396                }
397            }
398        }
399        Ok(())
400    }
401
402    /// Validate that all types in a stack effect are known types
403    ///
404    /// RFC #345: This catches cases where an uppercase identifier was parsed as a type
405    /// variable but should have been a union type (e.g., from an include that wasn't
406    /// available when parsing). Type variables should be single uppercase letters (T, U, V).
407    fn validate_effect_types(&self, effect: &Effect, word_name: &str) -> Result<(), String> {
408        self.validate_stack_types(&effect.inputs, word_name)?;
409        self.validate_stack_types(&effect.outputs, word_name)?;
410        Ok(())
411    }
412
413    /// Validate types in a stack type
414    fn validate_stack_types(&self, stack: &StackType, word_name: &str) -> Result<(), String> {
415        match stack {
416            StackType::Empty | StackType::RowVar(_) => Ok(()),
417            StackType::Cons { rest, top } => {
418                self.validate_type(top, word_name)?;
419                self.validate_stack_types(rest, word_name)
420            }
421        }
422    }
423
424    /// Validate a single type
425    ///
426    /// Type variables are allowed - they're used for polymorphism.
427    /// Only Type::Union types are validated to ensure they're registered.
428    fn validate_type(&self, ty: &Type, word_name: &str) -> Result<(), String> {
429        match ty {
430            Type::Var(_) => {
431                // Type variables are always valid - they represent polymorphic types
432                // Examples: T, U, V, Ctx, Handle, Acc, etc.
433                // After fixup_union_types(), any union name that was mistakenly parsed
434                // as a type variable will have been converted to Type::Union
435                Ok(())
436            }
437            Type::Quotation(effect) => self.validate_effect_types(effect, word_name),
438            Type::Closure { effect, captures } => {
439                self.validate_effect_types(effect, word_name)?;
440                for cap in captures {
441                    self.validate_type(cap, word_name)?;
442                }
443                Ok(())
444            }
445            // Concrete types are always valid
446            Type::Int | Type::Float | Type::Bool | Type::String | Type::Symbol | Type::Channel => {
447                Ok(())
448            }
449            // Union types are valid if they're registered
450            Type::Union(name) => {
451                if !self.unions.contains_key(name) {
452                    return Err(format!(
453                        "In word '{}': Unknown union type '{}' in stack effect.\n\
454                         Make sure this union is defined before the word that uses it.",
455                        word_name, name
456                    ));
457                }
458                Ok(())
459            }
460        }
461    }
462
463    /// Type check a complete program
464    pub fn check_program(&mut self, program: &Program) -> Result<(), String> {
465        // First pass: register all union definitions
466        for union_def in &program.unions {
467            let variants = union_def
468                .variants
469                .iter()
470                .map(|v| VariantInfo {
471                    name: v.name.clone(),
472                    fields: v
473                        .fields
474                        .iter()
475                        .map(|f| VariantFieldInfo {
476                            name: f.name.clone(),
477                            field_type: self.parse_type_name(&f.type_name),
478                        })
479                        .collect(),
480                })
481                .collect();
482
483            self.unions.insert(
484                union_def.name.clone(),
485                UnionTypeInfo {
486                    name: union_def.name.clone(),
487                    variants,
488                },
489            );
490        }
491
492        // Validate field types in unions reference known types
493        self.validate_union_field_types(program)?;
494
495        // Second pass: collect all word signatures
496        // All words must have explicit stack effect declarations (v2.0 requirement)
497        for word in &program.words {
498            if let Some(effect) = &word.effect {
499                // RFC #345: Validate that all types in the effect are known
500                // This catches cases where an uppercase identifier was parsed as a type variable
501                // but should have been a union type (e.g., from an include)
502                self.validate_effect_types(effect, &word.name)?;
503                self.env.insert(word.name.clone(), effect.clone());
504            } else {
505                return Err(format!(
506                    "Word '{}' is missing a stack effect declaration.\n\
507                     All words must declare their stack effect, e.g.: : {} ( -- ) ... ;",
508                    word.name, word.name
509                ));
510            }
511        }
512
513        // Validate main's signature (Issue #355).
514        // Only `( -- )` and `( -- Int )` are allowed.
515        if let Some(main_effect) = self.env.get("main") {
516            validate_main_effect(main_effect)?;
517        }
518
519        // Third pass: type check each word body
520        for word in &program.words {
521            self.check_word(word)?;
522        }
523
524        Ok(())
525    }
526
527    /// Type check a word definition
528    fn check_word(&self, word: &WordDef) -> Result<(), String> {
529        // Track current word for detecting recursive tail calls (divergent branches)
530        let line = word.source.as_ref().map(|s| s.start_line);
531        *self.current_word.borrow_mut() = Some((word.name.clone(), line));
532
533        // Reset aux stack for this word (Issue #350)
534        *self.current_aux_stack.borrow_mut() = StackType::Empty;
535
536        // All words must have declared effects (enforced in check_program)
537        let declared_effect = word.effect.as_ref().expect("word must have effect");
538
539        // Check if the word's output type is a quotation or closure
540        // If so, store it as the expected type for capture inference
541        if let Some((_rest, top_type)) = declared_effect.outputs.clone().pop()
542            && matches!(top_type, Type::Quotation(_) | Type::Closure { .. })
543        {
544            *self.expected_quotation_type.borrow_mut() = Some(top_type);
545        }
546
547        // Infer the result stack and effects starting from declared input
548        let (result_stack, _subst, inferred_effects) =
549            self.infer_statements_from(&word.body, &declared_effect.inputs, true)?;
550
551        // Clear expected type after checking
552        *self.expected_quotation_type.borrow_mut() = None;
553
554        // Verify result matches declared output
555        let line_info = line.map(format_line_prefix).unwrap_or_default();
556        unify_stacks(&declared_effect.outputs, &result_stack).map_err(|e| {
557            format!(
558                "{}Word '{}': declared output stack ({}) doesn't match inferred ({}): {}",
559                line_info, word.name, declared_effect.outputs, result_stack, e
560            )
561        })?;
562
563        // Verify computational effects match (bidirectional)
564        // 1. Check that each inferred effect has a matching declared effect (by kind)
565        // Type variables in effects are matched by kind (Yield matches Yield)
566        for inferred in &inferred_effects {
567            if !self.effect_matches_any(inferred, &declared_effect.effects) {
568                return Err(format!(
569                    "{}Word '{}': body produces effect '{}' but no matching effect is declared.\n\
570                     Hint: Add '| Yield <type>' to the word's stack effect declaration.",
571                    line_info, word.name, inferred
572                ));
573            }
574        }
575
576        // 2. Check that each declared effect is actually produced (effect soundness)
577        // This prevents declaring effects that don't occur
578        for declared in &declared_effect.effects {
579            if !self.effect_matches_any(declared, &inferred_effects) {
580                return Err(format!(
581                    "{}Word '{}': declares effect '{}' but body doesn't produce it.\n\
582                     Hint: Remove the effect declaration or ensure the body uses yield.",
583                    line_info, word.name, declared
584                ));
585            }
586        }
587
588        // Verify aux stack is empty at word boundary (Issue #350)
589        let aux_stack = self.current_aux_stack.borrow().clone();
590        if aux_stack != StackType::Empty {
591            return Err(format!(
592                "{}Word '{}': aux stack is not empty at word return.\n\
593                 Remaining aux stack: {}\n\
594                 Every >aux must be matched by a corresponding aux> before the word returns.",
595                line_info, word.name, aux_stack
596            ));
597        }
598
599        // Clear current word
600        *self.current_word.borrow_mut() = None;
601
602        Ok(())
603    }
604
605    /// Infer the resulting stack type from a sequence of statements
606    /// starting from a given input stack
607    /// Returns (final_stack, substitution, accumulated_effects)
608    ///
609    /// `capture_stmt_types`: If true, capture statement type info for codegen optimization.
610    /// Should only be true for top-level word bodies, not for nested branches/loops.
611    fn infer_statements_from(
612        &self,
613        statements: &[Statement],
614        start_stack: &StackType,
615        capture_stmt_types: bool,
616    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
617        let mut current_stack = start_stack.clone();
618        let mut accumulated_subst = Subst::empty();
619        let mut accumulated_effects: Vec<SideEffect> = Vec::new();
620        let mut skip_next = false;
621
622        for (i, stmt) in statements.iter().enumerate() {
623            // Skip this statement if we already handled it (e.g., pick/roll after literal)
624            if skip_next {
625                skip_next = false;
626                continue;
627            }
628
629            // Special case: IntLiteral followed by pick or roll
630            // Handle them as a fused operation with correct type semantics
631            if let Statement::IntLiteral(n) = stmt
632                && let Some(Statement::WordCall {
633                    name: next_word, ..
634                }) = statements.get(i + 1)
635            {
636                if next_word == "pick" {
637                    let (new_stack, subst) = self.handle_literal_pick(*n, current_stack.clone())?;
638                    current_stack = new_stack;
639                    accumulated_subst = accumulated_subst.compose(&subst);
640                    skip_next = true; // Skip the "pick" word
641                    continue;
642                } else if next_word == "roll" {
643                    let (new_stack, subst) = self.handle_literal_roll(*n, current_stack.clone())?;
644                    current_stack = new_stack;
645                    accumulated_subst = accumulated_subst.compose(&subst);
646                    skip_next = true; // Skip the "roll" word
647                    continue;
648                }
649            }
650
651            // Look ahead: if this is a quotation followed by a word that expects specific quotation type,
652            // set the expected type before checking the quotation
653            let saved_expected_type = if matches!(stmt, Statement::Quotation { .. }) {
654                // Save the current expected type
655                let saved = self.expected_quotation_type.borrow().clone();
656
657                // Try to set expected type based on lookahead
658                if let Some(Statement::WordCall {
659                    name: next_word, ..
660                }) = statements.get(i + 1)
661                {
662                    // Check if the next word expects a specific quotation type
663                    if let Some(next_effect) = self.lookup_word_effect(next_word) {
664                        // Extract the quotation type expected by the next word
665                        // For operations like spawn: ( ..a Quotation(-- ) -- ..a Int )
666                        if let Some((_rest, quot_type)) = next_effect.inputs.clone().pop()
667                            && matches!(quot_type, Type::Quotation(_))
668                        {
669                            *self.expected_quotation_type.borrow_mut() = Some(quot_type);
670                        }
671                    }
672                }
673                Some(saved)
674            } else {
675                None
676            };
677
678            // Capture statement type info for codegen optimization (Issue #186)
679            // Record the top-of-stack type BEFORE this statement for operations like dup
680            // Only capture for top-level word bodies, not nested branches/loops
681            if capture_stmt_types && let Some((word_name, _)) = self.current_word.borrow().as_ref()
682            {
683                self.capture_statement_type(word_name, i, &current_stack);
684            }
685
686            let (new_stack, subst, effects) = self.infer_statement(stmt, current_stack)?;
687            current_stack = new_stack;
688            accumulated_subst = accumulated_subst.compose(&subst);
689
690            // Accumulate side effects from this statement
691            for effect in effects {
692                if !accumulated_effects.contains(&effect) {
693                    accumulated_effects.push(effect);
694                }
695            }
696
697            // Restore expected type after checking quotation
698            if let Some(saved) = saved_expected_type {
699                *self.expected_quotation_type.borrow_mut() = saved;
700            }
701        }
702
703        Ok((current_stack, accumulated_subst, accumulated_effects))
704    }
705
706    /// Handle `n pick` where n is a literal integer
707    ///
708    /// pick(n) copies the value at position n to the top of the stack.
709    /// Position 0 is the top, 1 is below top, etc.
710    ///
711    /// Example: `2 pick` on stack ( A B C ) produces ( A B C A )
712    /// - Position 0: C (top)
713    /// - Position 1: B
714    /// - Position 2: A
715    /// - Result: copy A to top
716    fn handle_literal_pick(
717        &self,
718        n: i64,
719        current_stack: StackType,
720    ) -> Result<(StackType, Subst), String> {
721        if n < 0 {
722            return Err(format!("pick: index must be non-negative, got {}", n));
723        }
724
725        // Get the type at position n
726        let type_at_n = self.get_type_at_position(&current_stack, n as usize, "pick")?;
727
728        // Push a copy of that type onto the stack
729        Ok((current_stack.push(type_at_n), Subst::empty()))
730    }
731
732    /// Handle `n roll` where n is a literal integer
733    ///
734    /// roll(n) moves the value at position n to the top of the stack,
735    /// shifting all items above it down by one position.
736    ///
737    /// Example: `2 roll` on stack ( A B C ) produces ( B C A )
738    /// - Position 0: C (top)
739    /// - Position 1: B
740    /// - Position 2: A
741    /// - Result: move A to top, B and C shift down
742    fn handle_literal_roll(
743        &self,
744        n: i64,
745        current_stack: StackType,
746    ) -> Result<(StackType, Subst), String> {
747        if n < 0 {
748            return Err(format!("roll: index must be non-negative, got {}", n));
749        }
750
751        // For roll, we need to:
752        // 1. Extract the type at position n
753        // 2. Remove it from that position
754        // 3. Push it on top
755        self.rotate_type_to_top(current_stack, n as usize)
756    }
757
758    /// Get the type at position n in the stack (0 = top)
759    fn get_type_at_position(&self, stack: &StackType, n: usize, op: &str) -> Result<Type, String> {
760        let mut current = stack;
761        let mut pos = 0;
762
763        loop {
764            match current {
765                StackType::Cons { rest, top } => {
766                    if pos == n {
767                        return Ok(top.clone());
768                    }
769                    pos += 1;
770                    current = rest;
771                }
772                StackType::RowVar(name) => {
773                    // We've hit a row variable before reaching position n
774                    // This means the type at position n is unknown statically.
775                    // Generate a fresh type variable to represent it.
776                    // This allows the code to type-check, with the actual type
777                    // determined by unification with how the value is used.
778                    //
779                    // Note: This works correctly even in conditional branches because
780                    // branches are now inferred from the actual stack (not abstractly),
781                    // so row variables only appear when the word itself has polymorphic inputs.
782                    let fresh_type = Type::Var(self.fresh_var(&format!("{}_{}", op, name)));
783                    return Ok(fresh_type);
784                }
785                StackType::Empty => {
786                    return Err(format!(
787                        "{}{}: stack underflow - position {} requested but stack has only {} concrete items",
788                        self.line_prefix(),
789                        op,
790                        n,
791                        pos
792                    ));
793                }
794            }
795        }
796    }
797
798    /// Remove the type at position n and push it on top (for roll)
799    fn rotate_type_to_top(&self, stack: StackType, n: usize) -> Result<(StackType, Subst), String> {
800        if n == 0 {
801            // roll(0) is a no-op
802            return Ok((stack, Subst::empty()));
803        }
804
805        // Collect all types from top to the target position
806        let mut types_above: Vec<Type> = Vec::new();
807        let mut current = stack;
808        let mut pos = 0;
809
810        // Pop items until we reach position n
811        loop {
812            match current {
813                StackType::Cons { rest, top } => {
814                    if pos == n {
815                        // Found the target - 'top' is what we want to move to the top
816                        // Rebuild the stack: rest, then types_above (reversed), then top
817                        let mut result = *rest;
818                        // Push types_above back in reverse order (bottom to top)
819                        for ty in types_above.into_iter().rev() {
820                            result = result.push(ty);
821                        }
822                        // Push the rotated type on top
823                        result = result.push(top);
824                        return Ok((result, Subst::empty()));
825                    }
826                    types_above.push(top);
827                    pos += 1;
828                    current = *rest;
829                }
830                StackType::RowVar(name) => {
831                    // Reached a row variable before position n
832                    // The type at position n is in the row variable.
833                    // Generate a fresh type variable to represent the moved value.
834                    //
835                    // Note: This preserves stack size correctly because we're moving
836                    // (not copying) a value. The row variable conceptually "loses"
837                    // an item which appears on top. Since we can't express "row minus one",
838                    // we generate a fresh type and trust unification to constrain it.
839                    //
840                    // This works correctly in conditional branches because branches are
841                    // now inferred from the actual stack (not abstractly), so row variables
842                    // only appear when the word itself has polymorphic inputs.
843                    let fresh_type = Type::Var(self.fresh_var(&format!("roll_{}", name)));
844
845                    // Reconstruct the stack with the rolled type on top
846                    let mut result = StackType::RowVar(name.clone());
847                    for ty in types_above.into_iter().rev() {
848                        result = result.push(ty);
849                    }
850                    result = result.push(fresh_type);
851                    return Ok((result, Subst::empty()));
852                }
853                StackType::Empty => {
854                    return Err(format!(
855                        "{}roll: stack underflow - position {} requested but stack has only {} items",
856                        self.line_prefix(),
857                        n,
858                        pos
859                    ));
860                }
861            }
862        }
863    }
864
865    /// Infer the stack effect of a sequence of statements
866    /// Returns an Effect with both inputs and outputs normalized by applying discovered substitutions
867    /// Also includes any computational side effects (Yield, etc.)
868    fn infer_statements(&self, statements: &[Statement]) -> Result<Effect, String> {
869        let start = StackType::RowVar("input".to_string());
870        // Don't capture statement types for quotation bodies - only top-level word bodies
871        let (result, subst, effects) = self.infer_statements_from(statements, &start, false)?;
872
873        // Apply the accumulated substitution to both start and result
874        // This ensures row variables are consistently named
875        let normalized_start = subst.apply_stack(&start);
876        let normalized_result = subst.apply_stack(&result);
877
878        Ok(Effect::with_effects(
879            normalized_start,
880            normalized_result,
881            effects,
882        ))
883    }
884
885    /// Infer the stack effect of a match expression
886    fn infer_match(
887        &self,
888        arms: &[crate::ast::MatchArm],
889        match_span: &Option<crate::ast::Span>,
890        current_stack: StackType,
891    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
892        if arms.is_empty() {
893            return Err("match expression must have at least one arm".to_string());
894        }
895
896        // Pop the matched value from the stack
897        let (stack_after_match, _matched_type) =
898            self.pop_type(&current_stack, "match expression")?;
899
900        // Track all arm results for unification
901        let mut arm_results: Vec<StackType> = Vec::new();
902        let mut combined_subst = Subst::empty();
903        let mut merged_effects: Vec<SideEffect> = Vec::new();
904
905        // Save aux stack before match arms (Issue #350)
906        let aux_before_match = self.current_aux_stack.borrow().clone();
907        let mut aux_after_arms: Vec<StackType> = Vec::new();
908
909        for arm in arms {
910            // Restore aux stack before each arm (Issue #350)
911            *self.current_aux_stack.borrow_mut() = aux_before_match.clone();
912
913            // Get variant name from pattern
914            let variant_name = match &arm.pattern {
915                crate::ast::Pattern::Variant(name) => name.as_str(),
916                crate::ast::Pattern::VariantWithBindings { name, .. } => name.as_str(),
917            };
918
919            // Look up variant info
920            let (_union_name, variant_info) = self
921                .find_variant(variant_name)
922                .ok_or_else(|| format!("Unknown variant '{}' in match pattern", variant_name))?;
923
924            // Push fields onto the stack based on pattern type
925            let arm_stack = self.push_variant_fields(
926                &stack_after_match,
927                &arm.pattern,
928                variant_info,
929                variant_name,
930            )?;
931
932            // Type check the arm body directly from the actual stack
933            // Don't capture statement types for match arms - only top-level word bodies
934            let (arm_result, arm_subst, arm_effects) =
935                self.infer_statements_from(&arm.body, &arm_stack, false)?;
936
937            combined_subst = combined_subst.compose(&arm_subst);
938            arm_results.push(arm_result);
939            aux_after_arms.push(self.current_aux_stack.borrow().clone());
940
941            // Merge effects from this arm
942            for effect in arm_effects {
943                if !merged_effects.contains(&effect) {
944                    merged_effects.push(effect);
945                }
946            }
947        }
948
949        // Verify all arms produce the same aux stack (Issue #350)
950        if aux_after_arms.len() > 1 {
951            let first_aux = &aux_after_arms[0];
952            for (i, arm_aux) in aux_after_arms.iter().enumerate().skip(1) {
953                if arm_aux != first_aux {
954                    let match_line = match_span.as_ref().map(|s| s.line + 1).unwrap_or(0);
955                    return Err(format!(
956                        "at line {}: match arms have incompatible aux stack effects:\n\
957                         \x20 arm 0 aux: {}\n\
958                         \x20 arm {} aux: {}\n\
959                         \x20 All match arms must leave the aux stack in the same state.",
960                        match_line, first_aux, i, arm_aux
961                    ));
962                }
963            }
964        }
965        // Set aux to the first arm's result (all are verified equal)
966        if let Some(aux) = aux_after_arms.into_iter().next() {
967            *self.current_aux_stack.borrow_mut() = aux;
968        }
969
970        // Unify all arm results to ensure they're compatible
971        let mut final_result = arm_results[0].clone();
972        for (i, arm_result) in arm_results.iter().enumerate().skip(1) {
973            // Get line info for error reporting
974            let match_line = match_span.as_ref().map(|s| s.line + 1).unwrap_or(0);
975            let arm0_line = arms[0].span.as_ref().map(|s| s.line + 1).unwrap_or(0);
976            let arm_i_line = arms[i].span.as_ref().map(|s| s.line + 1).unwrap_or(0);
977
978            let arm_subst = unify_stacks(&final_result, arm_result).map_err(|e| {
979                if match_line > 0 && arm0_line > 0 && arm_i_line > 0 {
980                    format!(
981                        "at line {}: match arms have incompatible stack effects:\n\
982                         \x20 arm 0 (line {}) produces: {}\n\
983                         \x20 arm {} (line {}) produces: {}\n\
984                         \x20 All match arms must produce the same stack shape.\n\
985                         \x20 Error: {}",
986                        match_line, arm0_line, final_result, i, arm_i_line, arm_result, e
987                    )
988                } else {
989                    format!(
990                        "match arms have incompatible stack effects:\n\
991                         \x20 arm 0 produces: {}\n\
992                         \x20 arm {} produces: {}\n\
993                         \x20 All match arms must produce the same stack shape.\n\
994                         \x20 Error: {}",
995                        final_result, i, arm_result, e
996                    )
997                }
998            })?;
999            combined_subst = combined_subst.compose(&arm_subst);
1000            final_result = arm_subst.apply_stack(&final_result);
1001        }
1002
1003        Ok((final_result, combined_subst, merged_effects))
1004    }
1005
1006    /// Push variant fields onto the stack based on the match pattern
1007    fn push_variant_fields(
1008        &self,
1009        stack: &StackType,
1010        pattern: &crate::ast::Pattern,
1011        variant_info: &VariantInfo,
1012        variant_name: &str,
1013    ) -> Result<StackType, String> {
1014        let mut arm_stack = stack.clone();
1015        match pattern {
1016            crate::ast::Pattern::Variant(_) => {
1017                // Stack-based: push all fields in declaration order
1018                for field in &variant_info.fields {
1019                    arm_stack = arm_stack.push(field.field_type.clone());
1020                }
1021            }
1022            crate::ast::Pattern::VariantWithBindings { bindings, .. } => {
1023                // Named bindings: validate and push only bound fields
1024                for binding in bindings {
1025                    let field = variant_info
1026                        .fields
1027                        .iter()
1028                        .find(|f| &f.name == binding)
1029                        .ok_or_else(|| {
1030                            let available: Vec<_> = variant_info
1031                                .fields
1032                                .iter()
1033                                .map(|f| f.name.as_str())
1034                                .collect();
1035                            format!(
1036                                "Unknown field '{}' in pattern for variant '{}'.\n\
1037                                 Available fields: {}",
1038                                binding,
1039                                variant_name,
1040                                available.join(", ")
1041                            )
1042                        })?;
1043                    arm_stack = arm_stack.push(field.field_type.clone());
1044                }
1045            }
1046        }
1047        Ok(arm_stack)
1048    }
1049
1050    /// Check if a branch ends with a recursive tail call to the current word
1051    /// or to a mutually recursive word.
1052    ///
1053    /// Such branches are "divergent" - they never return to the if/else,
1054    /// so their stack effect shouldn't constrain the other branch.
1055    ///
1056    /// # Detection Capabilities
1057    ///
1058    /// - Direct recursion: word calls itself
1059    /// - Mutual recursion: word calls another word in the same SCC (when call graph is available)
1060    ///
1061    /// # Limitations
1062    ///
1063    /// This detection does NOT detect:
1064    /// - Calls to known non-returning functions (panic, exit, infinite loops)
1065    /// - Nested control flow with tail calls (if ... if ... recurse then then)
1066    ///
1067    /// These patterns will still require branch unification. Future enhancements
1068    /// could track known non-returning functions or support explicit divergence
1069    /// annotations (similar to Rust's `!` type).
1070    fn is_divergent_branch(&self, statements: &[Statement]) -> bool {
1071        let Some((current_word_name, _)) = self.current_word.borrow().as_ref().cloned() else {
1072            return false;
1073        };
1074        let Some(Statement::WordCall { name, .. }) = statements.last() else {
1075            return false;
1076        };
1077
1078        // Direct recursion: word calls itself
1079        if name == &current_word_name {
1080            return true;
1081        }
1082
1083        // Mutual recursion: word calls another word in the same SCC
1084        if let Some(ref graph) = self.call_graph
1085            && graph.are_mutually_recursive(&current_word_name, name)
1086        {
1087            return true;
1088        }
1089
1090        false
1091    }
1092
1093    /// Infer the stack effect of an if/else expression
1094    fn infer_if(
1095        &self,
1096        then_branch: &[Statement],
1097        else_branch: &Option<Vec<Statement>>,
1098        if_span: &Option<crate::ast::Span>,
1099        current_stack: StackType,
1100    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1101        // Pop condition (must be Bool)
1102        let (stack_after_cond, cond_type) = self.pop_type(&current_stack, "if condition")?;
1103
1104        // Condition must be Bool
1105        let cond_subst = unify_stacks(
1106            &StackType::singleton(Type::Bool),
1107            &StackType::singleton(cond_type),
1108        )
1109        .map_err(|e| format!("if condition must be Bool: {}", e))?;
1110
1111        let stack_after_cond = cond_subst.apply_stack(&stack_after_cond);
1112
1113        // Check for divergent branches (recursive tail calls)
1114        let then_diverges = self.is_divergent_branch(then_branch);
1115        let else_diverges = else_branch
1116            .as_ref()
1117            .map(|stmts| self.is_divergent_branch(stmts))
1118            .unwrap_or(false);
1119
1120        // Save aux stack before branching (Issue #350)
1121        let aux_before_branches = self.current_aux_stack.borrow().clone();
1122
1123        // Infer branches directly from the actual stack
1124        // Don't capture statement types for if branches - only top-level word bodies
1125        let (then_result, then_subst, then_effects) =
1126            self.infer_statements_from(then_branch, &stack_after_cond, false)?;
1127        let aux_after_then = self.current_aux_stack.borrow().clone();
1128
1129        // Restore aux stack before checking else branch (Issue #350)
1130        *self.current_aux_stack.borrow_mut() = aux_before_branches.clone();
1131
1132        // Infer else branch (or use stack_after_cond if no else)
1133        let (else_result, else_subst, else_effects) = if let Some(else_stmts) = else_branch {
1134            self.infer_statements_from(else_stmts, &stack_after_cond, false)?
1135        } else {
1136            (stack_after_cond.clone(), Subst::empty(), vec![])
1137        };
1138        let aux_after_else = self.current_aux_stack.borrow().clone();
1139
1140        // Verify aux stacks match between branches (Issue #350)
1141        // Skip check if one branch diverges (never returns)
1142        if !then_diverges && !else_diverges && aux_after_then != aux_after_else {
1143            let if_line = if_span.as_ref().map(|s| s.line + 1).unwrap_or(0);
1144            return Err(format!(
1145                "at line {}: if/else branches have incompatible aux stack effects:\n\
1146                 \x20 then branch aux: {}\n\
1147                 \x20 else branch aux: {}\n\
1148                 \x20 Both branches must leave the aux stack in the same state.",
1149                if_line, aux_after_then, aux_after_else
1150            ));
1151        }
1152
1153        // Set aux to the non-divergent branch's result (or then if neither diverges)
1154        if then_diverges && !else_diverges {
1155            *self.current_aux_stack.borrow_mut() = aux_after_else;
1156        } else {
1157            *self.current_aux_stack.borrow_mut() = aux_after_then;
1158        }
1159
1160        // Merge effects from both branches (if either yields, the whole if yields)
1161        let mut merged_effects = then_effects;
1162        for effect in else_effects {
1163            if !merged_effects.contains(&effect) {
1164                merged_effects.push(effect);
1165            }
1166        }
1167
1168        // Handle divergent branches: if one branch diverges (never returns),
1169        // use the other branch's stack type without requiring unification.
1170        // This supports patterns like:
1171        //   chan.receive not if drop store-loop then
1172        // where the then branch recurses and the else branch continues.
1173        let (result, branch_subst) = if then_diverges && !else_diverges {
1174            // Then branch diverges, use else branch's type
1175            (else_result, Subst::empty())
1176        } else if else_diverges && !then_diverges {
1177            // Else branch diverges, use then branch's type
1178            (then_result, Subst::empty())
1179        } else {
1180            // Both branches must produce compatible stacks (normal case)
1181            let if_line = if_span.as_ref().map(|s| s.line + 1).unwrap_or(0);
1182            let branch_subst = unify_stacks(&then_result, &else_result).map_err(|e| {
1183                if if_line > 0 {
1184                    format!(
1185                        "at line {}: if/else branches have incompatible stack effects:\n\
1186                         \x20 then branch produces: {}\n\
1187                         \x20 else branch produces: {}\n\
1188                         \x20 Both branches of an if/else must produce the same stack shape.\n\
1189                         \x20 Hint: Make sure both branches push/pop the same number of values.\n\
1190                         \x20 Error: {}",
1191                        if_line, then_result, else_result, e
1192                    )
1193                } else {
1194                    format!(
1195                        "if/else branches have incompatible stack effects:\n\
1196                         \x20 then branch produces: {}\n\
1197                         \x20 else branch produces: {}\n\
1198                         \x20 Both branches of an if/else must produce the same stack shape.\n\
1199                         \x20 Hint: Make sure both branches push/pop the same number of values.\n\
1200                         \x20 Error: {}",
1201                        then_result, else_result, e
1202                    )
1203                }
1204            })?;
1205            (branch_subst.apply_stack(&then_result), branch_subst)
1206        };
1207
1208        // Propagate all substitutions
1209        let total_subst = cond_subst
1210            .compose(&then_subst)
1211            .compose(&else_subst)
1212            .compose(&branch_subst);
1213        Ok((result, total_subst, merged_effects))
1214    }
1215
1216    /// Infer the stack effect of a quotation
1217    /// Quotations capture effects in their type - they don't propagate effects to the outer scope
1218    fn infer_quotation(
1219        &self,
1220        id: usize,
1221        body: &[Statement],
1222        current_stack: StackType,
1223    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1224        // Save and clear expected type so nested quotations don't inherit it.
1225        // The expected type applies only to THIS quotation, not inner ones.
1226        let expected_for_this_quotation = self.expected_quotation_type.borrow().clone();
1227        *self.expected_quotation_type.borrow_mut() = None;
1228
1229        // Save enclosing aux stack and enter quotation scope (Issue #350, #393).
1230        // Quotations are compiled as separate LLVM functions; each gets its own
1231        // aux slot table. The save/restore here means the enclosing word's aux
1232        // state is undisturbed by the quotation, and the quotation's aux usage
1233        // is tracked independently in `quotation_aux_depths` (Issue #393).
1234        let saved_aux = self.current_aux_stack.borrow().clone();
1235        *self.current_aux_stack.borrow_mut() = StackType::Empty;
1236        self.quotation_id_stack.borrow_mut().push(id);
1237
1238        // Run the body inference and balance check inside an immediately-invoked
1239        // closure so we can restore scope state on every exit path — including
1240        // errors. Without this, an error in body inference or the balance check
1241        // would leave the typechecker with a corrupt scope stack and a polluted
1242        // aux stack, which matters for callers that inspect errors and continue.
1243        let body_result: Result<Effect, String> = (|| {
1244            // Infer the effect of the quotation body.
1245            //
1246            // If we have an expected quotation type from a combinator's signature
1247            // (e.g., list.fold expects [..b Acc T -- ..b Acc]), seed the body
1248            // inference with that input stack. Without this, the body inference
1249            // starts from a polymorphic row variable, and operations like >aux
1250            // can't pop because they don't know the type. Issue #393.
1251            let body_effect = if let Some(expected) = &expected_for_this_quotation {
1252                let expected_effect = match expected {
1253                    Type::Quotation(eff) => Some((**eff).clone()),
1254                    Type::Closure { effect, .. } => Some((**effect).clone()),
1255                    _ => None,
1256                };
1257                if let Some(eff) = expected_effect {
1258                    // Freshen to avoid row-variable name clashes with the
1259                    // enclosing scope.
1260                    let fresh = self.freshen_effect(&eff);
1261                    let (result, subst, effects) =
1262                        self.infer_statements_from(body, &fresh.inputs, false)?;
1263                    let normalized_start = subst.apply_stack(&fresh.inputs);
1264                    let normalized_result = subst.apply_stack(&result);
1265                    Effect::with_effects(normalized_start, normalized_result, effects)
1266                } else {
1267                    self.infer_statements(body)?
1268                }
1269            } else {
1270                self.infer_statements(body)?
1271            };
1272
1273            // Verify quotation's aux stack is balanced (Issue #350).
1274            // Lexical scoping: every >aux inside the quotation must have a
1275            // matching aux> inside the same quotation.
1276            let quot_aux = self.current_aux_stack.borrow().clone();
1277            if quot_aux != StackType::Empty {
1278                return Err(format!(
1279                    "Quotation has unbalanced aux stack.\n\
1280                     Remaining aux stack: {}\n\
1281                     Every >aux must be matched by a corresponding aux> within the quotation.",
1282                    quot_aux
1283                ));
1284            }
1285
1286            Ok(body_effect)
1287        })();
1288
1289        // Always restore scope state, regardless of whether the body inference
1290        // succeeded or failed.
1291        *self.current_aux_stack.borrow_mut() = saved_aux;
1292        self.quotation_id_stack.borrow_mut().pop();
1293        *self.expected_quotation_type.borrow_mut() = expected_for_this_quotation.clone();
1294
1295        let body_effect = body_result?;
1296
1297        // Perform capture analysis
1298        let quot_type = self.analyze_captures(&body_effect, &current_stack)?;
1299
1300        // If this is a closure, we need to pop the captured values from the stack
1301        // and correct the capture types from the caller's actual stack.
1302        let result_stack = match &quot_type {
1303            Type::Quotation(_) => {
1304                // Stateless - no captures. Record in type map for codegen.
1305                self.quotation_types
1306                    .borrow_mut()
1307                    .insert(id, quot_type.clone());
1308                current_stack.push(quot_type)
1309            }
1310            Type::Closure {
1311                captures, effect, ..
1312            } => {
1313                // Pop captured values from the caller's stack.
1314                // The capture COUNT comes from analyze_captures (based on
1315                // body vs expected input comparison), but the capture TYPES
1316                // come from the caller's stack — not from the body's inference.
1317                //
1318                // We intentionally do NOT call unify_types on the popped types.
1319                // The body's inference may have constrained a type variable to
1320                // Int/Float via its operations (e.g., i.+), even when the actual
1321                // stack value is a Variant. unify_types(Var("V$nn"), Int) would
1322                // succeed and propagate the wrong type to codegen, which would
1323                // then emit env_get_int for a Variant value — a runtime crash.
1324                // Using the caller's actual types directly ensures codegen emits
1325                // the correct getter for the runtime Value type.
1326                let mut stack = current_stack.clone();
1327                let mut actual_captures: Vec<Type> = Vec::new();
1328                for _ in (0..captures.len()).rev() {
1329                    let (new_stack, actual_type) = self.pop_type(&stack, "closure capture")?;
1330                    actual_captures.push(actual_type);
1331                    stack = new_stack;
1332                }
1333                // actual_captures is in pop order (top-down), reverse to
1334                // get bottom-to-top (matching calculate_captures convention)
1335                actual_captures.reverse();
1336
1337                // Rebuild the closure type with the actual capture types
1338                let corrected_quot_type = Type::Closure {
1339                    effect: effect.clone(),
1340                    captures: actual_captures,
1341                };
1342
1343                // Update the type map so codegen sees the corrected types
1344                self.quotation_types
1345                    .borrow_mut()
1346                    .insert(id, corrected_quot_type.clone());
1347
1348                stack.push(corrected_quot_type)
1349            }
1350            _ => unreachable!("analyze_captures only returns Quotation or Closure"),
1351        };
1352
1353        // Quotations don't propagate effects - they capture them in the quotation type
1354        // The effect annotation on the quotation type (e.g., [ ..a -- ..b | Yield Int ])
1355        // indicates what effects the quotation may produce when called
1356        Ok((result_stack, Subst::empty(), vec![]))
1357    }
1358
1359    /// Infer the stack effect of a word call
1360    fn infer_word_call(
1361        &self,
1362        name: &str,
1363        span: &Option<crate::ast::Span>,
1364        current_stack: StackType,
1365    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1366        // Arithmetic sugar resolution: resolve +, -, *, / etc. to concrete ops
1367        // based on the types currently on the stack.
1368        let is_sugar = matches!(
1369            name,
1370            "+" | "-" | "*" | "/" | "%" | "=" | "<" | ">" | "<=" | ">=" | "<>"
1371        );
1372        if is_sugar {
1373            if let Some(resolved) = self.resolve_arithmetic_sugar(name, &current_stack) {
1374                // Record the resolution for codegen, keyed by source location (line, column)
1375                if let Some(s) = span {
1376                    self.resolved_sugar
1377                        .borrow_mut()
1378                        .insert((s.line, s.column), resolved.clone());
1379                }
1380                // Proceed as if the user wrote the resolved name
1381                return self.infer_word_call(&resolved, span, current_stack);
1382            }
1383            // Sugar op but types don't match — give a helpful error
1384            let line_prefix = self.line_prefix();
1385            let (top_desc, second_desc) = {
1386                let top = current_stack.clone().pop().map(|(_, t)| format!("{}", t));
1387                let second = current_stack
1388                    .clone()
1389                    .pop()
1390                    .and_then(|(r, _)| r.pop().map(|(_, t)| format!("{}", t)));
1391                (
1392                    top.unwrap_or_else(|| "empty".to_string()),
1393                    second.unwrap_or_else(|| "empty".to_string()),
1394                )
1395            };
1396            let (type_options, suggestion) = match name {
1397                "+" => (
1398                    "Int+Int, Float+Float, or String+String",
1399                    "Use `i.+`, `f.+`, or `string.concat`.",
1400                ),
1401                "=" => (
1402                    "Int+Int, Float+Float, or String+String (equality)",
1403                    "Use `i.=`, `f.=`, or `string.equal?`.",
1404                ),
1405                "%" => (
1406                    "Int+Int only — float modulo is not supported",
1407                    "Use `i.%` for integer modulo.",
1408                ),
1409                _ => (
1410                    "Int+Int or Float+Float",
1411                    "Use the `i.` or `f.` prefixed variant.",
1412                ),
1413            };
1414            return Err(format!(
1415                "{}`{}` requires matching types ({}), got ({}, {}). {}",
1416                line_prefix, name, type_options, second_desc, top_desc, suggestion,
1417            ));
1418        }
1419
1420        // Special handling for aux stack operations (Issue #350)
1421        if name == ">aux" {
1422            return self.infer_to_aux(span, current_stack);
1423        }
1424        if name == "aux>" {
1425            return self.infer_from_aux(span, current_stack);
1426        }
1427
1428        // Special handling for `call`: extract and apply the quotation's actual effect
1429        // This ensures stack pollution through quotations is caught (Issue #228)
1430        if name == "call" {
1431            return self.infer_call(span, current_stack);
1432        }
1433
1434        // Special handling for dataflow combinators
1435        if name == "dip" {
1436            return self.infer_dip(span, current_stack);
1437        }
1438        if name == "keep" {
1439            return self.infer_keep(span, current_stack);
1440        }
1441        if name == "bi" {
1442            return self.infer_bi(span, current_stack);
1443        }
1444
1445        // Look up word's effect
1446        let effect = self
1447            .lookup_word_effect(name)
1448            .ok_or_else(|| format!("Unknown word: '{}'", name))?;
1449
1450        // Freshen the effect to avoid variable name clashes
1451        let fresh_effect = self.freshen_effect(&effect);
1452
1453        // Special handling for strand.spawn: auto-convert Quotation to Closure if needed
1454        let adjusted_stack = if name == "strand.spawn" {
1455            self.adjust_stack_for_spawn(current_stack, &fresh_effect)?
1456        } else {
1457            current_stack
1458        };
1459
1460        // Apply the freshened effect to current stack
1461        let (result_stack, subst) = self.apply_effect(&fresh_effect, adjusted_stack, name, span)?;
1462
1463        // Propagate side effects from the called word
1464        // Note: strand.weave "handles" Yield effects (consumes them from the quotation)
1465        // strand.spawn requires pure quotations (checked separately)
1466        let propagated_effects = fresh_effect.effects.clone();
1467
1468        Ok((result_stack, subst, propagated_effects))
1469    }
1470
1471    /// Handle >aux: pop from main stack, push onto scope-local aux stack
1472    /// (Issue #350, Issue #393).
1473    ///
1474    /// In word-body scope, depth is tracked per word in `aux_max_depths`.
1475    /// In quotation-body scope, depth is tracked per quotation ID in
1476    /// `quotation_aux_depths`. Each quotation gets its own slot table at
1477    /// codegen time.
1478    fn infer_to_aux(
1479        &self,
1480        _span: &Option<crate::ast::Span>,
1481        current_stack: StackType,
1482    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1483        let (rest, top_type) = self.pop_type(&current_stack, ">aux")?;
1484
1485        // Push onto aux stack
1486        let mut aux = self.current_aux_stack.borrow_mut();
1487        *aux = aux.clone().push(top_type);
1488
1489        // Track max depth for codegen alloca sizing.
1490        // If we're inside a quotation, key the depth by quotation ID.
1491        // Otherwise, key by the enclosing word name.
1492        let depth = Self::stack_depth(&aux);
1493        let quot_stack = self.quotation_id_stack.borrow();
1494        if let Some(&quot_id) = quot_stack.last() {
1495            let mut depths = self.quotation_aux_depths.borrow_mut();
1496            let entry = depths.entry(quot_id).or_insert(0);
1497            if depth > *entry {
1498                *entry = depth;
1499            }
1500        } else if let Some((word_name, _)) = self.current_word.borrow().as_ref() {
1501            let mut depths = self.aux_max_depths.borrow_mut();
1502            let entry = depths.entry(word_name.clone()).or_insert(0);
1503            if depth > *entry {
1504                *entry = depth;
1505            }
1506        }
1507
1508        Ok((rest, Subst::empty(), vec![]))
1509    }
1510
1511    /// Handle aux>: pop from aux stack, push onto main stack (Issue #350, #393).
1512    fn infer_from_aux(
1513        &self,
1514        _span: &Option<crate::ast::Span>,
1515        current_stack: StackType,
1516    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1517        let mut aux = self.current_aux_stack.borrow_mut();
1518        match aux.clone().pop() {
1519            Some((rest, top_type)) => {
1520                *aux = rest;
1521                Ok((current_stack.push(top_type), Subst::empty(), vec![]))
1522            }
1523            None => {
1524                let line_info = self.line_prefix();
1525                Err(format!(
1526                    "{}aux>: aux stack is empty. Every aux> must be paired with a preceding >aux.",
1527                    line_info
1528                ))
1529            }
1530        }
1531    }
1532
1533    /// Special handling for `call` to properly propagate quotation effects (Issue #228)
1534    ///
1535    /// The generic `call` signature `( ..a Q -- ..b )` has independent row variables,
1536    /// which doesn't constrain the output based on the quotation's actual effect.
1537    /// This function extracts the quotation's effect and applies it properly.
1538    fn infer_call(
1539        &self,
1540        span: &Option<crate::ast::Span>,
1541        current_stack: StackType,
1542    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1543        // Pop the quotation from the stack
1544        let line_prefix = self.line_prefix();
1545        let (remaining_stack, quot_type) = current_stack.clone().pop().ok_or_else(|| {
1546            format!(
1547                "{}call: stack underflow - expected quotation on stack",
1548                line_prefix
1549            )
1550        })?;
1551
1552        // Extract the quotation's effect
1553        let quot_effect = match &quot_type {
1554            Type::Quotation(effect) => (**effect).clone(),
1555            Type::Closure { effect, .. } => (**effect).clone(),
1556            Type::Var(_) => {
1557                // Type variable - fall back to polymorphic behavior
1558                // This happens when the quotation type isn't known yet
1559                let effect = self
1560                    .lookup_word_effect("call")
1561                    .ok_or_else(|| "Unknown word: 'call'".to_string())?;
1562                let fresh_effect = self.freshen_effect(&effect);
1563                let (result_stack, subst) =
1564                    self.apply_effect(&fresh_effect, current_stack, "call", span)?;
1565                return Ok((result_stack, subst, vec![]));
1566            }
1567            _ => {
1568                return Err(format!(
1569                    "call: expected quotation or closure on stack, got {}",
1570                    quot_type
1571                ));
1572            }
1573        };
1574
1575        // Check for Yield effects - quotations with Yield must use strand.weave
1576        if quot_effect.has_yield() {
1577            return Err("Cannot call quotation with Yield effect directly.\n\
1578                 Quotations that yield values must be wrapped with `strand.weave`.\n\
1579                 Example: `[ yielding-code ] strand.weave` instead of `[ yielding-code ] call`"
1580                .to_string());
1581        }
1582
1583        // Freshen the quotation's effect to avoid variable clashes
1584        let fresh_effect = self.freshen_effect(&quot_effect);
1585
1586        // Apply the quotation's effect to the remaining stack
1587        let (result_stack, subst) =
1588            self.apply_effect(&fresh_effect, remaining_stack, "call", span)?;
1589
1590        // Propagate side effects from the quotation
1591        let propagated_effects = fresh_effect.effects.clone();
1592
1593        Ok((result_stack, subst, propagated_effects))
1594    }
1595
1596    /// Resolve arithmetic sugar operators to concrete operations based on
1597    /// the types on the stack. Returns `None` if the name is not a sugar op.
1598    fn resolve_arithmetic_sugar(&self, name: &str, stack: &StackType) -> Option<String> {
1599        // Only handle known sugar operators
1600        let is_binary = matches!(
1601            name,
1602            "+" | "-" | "*" | "/" | "%" | "=" | "<" | ">" | "<=" | ">=" | "<>"
1603        );
1604        if !is_binary {
1605            return None;
1606        }
1607
1608        // Peek at the top two types on the stack
1609        let (rest, top) = stack.clone().pop()?;
1610        let (_, second) = rest.pop()?;
1611
1612        match (name, &second, &top) {
1613            // Int × Int operations
1614            ("+", Type::Int, Type::Int) => Some("i.+".to_string()),
1615            ("-", Type::Int, Type::Int) => Some("i.-".to_string()),
1616            ("*", Type::Int, Type::Int) => Some("i.*".to_string()),
1617            ("/", Type::Int, Type::Int) => Some("i./".to_string()),
1618            ("%", Type::Int, Type::Int) => Some("i.%".to_string()),
1619            ("=", Type::Int, Type::Int) => Some("i.=".to_string()),
1620            ("<", Type::Int, Type::Int) => Some("i.<".to_string()),
1621            (">", Type::Int, Type::Int) => Some("i.>".to_string()),
1622            ("<=", Type::Int, Type::Int) => Some("i.<=".to_string()),
1623            (">=", Type::Int, Type::Int) => Some("i.>=".to_string()),
1624            ("<>", Type::Int, Type::Int) => Some("i.<>".to_string()),
1625
1626            // Float × Float operations
1627            ("+", Type::Float, Type::Float) => Some("f.+".to_string()),
1628            ("-", Type::Float, Type::Float) => Some("f.-".to_string()),
1629            ("*", Type::Float, Type::Float) => Some("f.*".to_string()),
1630            ("/", Type::Float, Type::Float) => Some("f./".to_string()),
1631            ("=", Type::Float, Type::Float) => Some("f.=".to_string()),
1632            ("<", Type::Float, Type::Float) => Some("f.<".to_string()),
1633            (">", Type::Float, Type::Float) => Some("f.>".to_string()),
1634            ("<=", Type::Float, Type::Float) => Some("f.<=".to_string()),
1635            (">=", Type::Float, Type::Float) => Some("f.>=".to_string()),
1636            ("<>", Type::Float, Type::Float) => Some("f.<>".to_string()),
1637
1638            // String operations (only + for concat, = for equality)
1639            ("+", Type::String, Type::String) => Some("string.concat".to_string()),
1640            ("=", Type::String, Type::String) => Some("string.equal?".to_string()),
1641
1642            // No match — not a sugar op for these types (will fall through
1643            // to normal lookup, which will fail with "Unknown word: '+'" —
1644            // giving the user a clear error that they need explicit types)
1645            _ => None,
1646        }
1647    }
1648
1649    /// Infer the stack effect of `dip`: ( ..a x quot -- ..b x )
1650    ///
1651    /// Hide the value below the quotation, run the quotation on the rest
1652    /// of the stack, then restore the hidden value on top.
1653    fn infer_dip(
1654        &self,
1655        span: &Option<crate::ast::Span>,
1656        current_stack: StackType,
1657    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1658        let line_prefix = self.line_prefix();
1659
1660        // Pop the quotation
1661        let (stack_after_quot, quot_type) = current_stack.clone().pop().ok_or_else(|| {
1662            format!(
1663                "{}dip: stack underflow - expected quotation on stack",
1664                line_prefix
1665            )
1666        })?;
1667
1668        // Extract the quotation's effect
1669        let quot_effect = match &quot_type {
1670            Type::Quotation(effect) => (**effect).clone(),
1671            Type::Closure { effect, .. } => (**effect).clone(),
1672            Type::Var(_) => {
1673                // Unknown quotation type — fall back to generic builtin signature
1674                let effect = self
1675                    .lookup_word_effect("dip")
1676                    .ok_or_else(|| "Unknown word: 'dip'".to_string())?;
1677                let fresh_effect = self.freshen_effect(&effect);
1678                let (result_stack, subst) =
1679                    self.apply_effect(&fresh_effect, current_stack, "dip", span)?;
1680                return Ok((result_stack, subst, vec![]));
1681            }
1682            _ => {
1683                return Err(format!(
1684                    "{}dip: expected quotation or closure on top of stack, got {}",
1685                    line_prefix, quot_type
1686                ));
1687            }
1688        };
1689
1690        if quot_effect.has_yield() {
1691            return Err("dip: quotation must not have Yield effects.\n\
1692                 Use strand.weave for quotations that yield."
1693                .to_string());
1694        }
1695
1696        // Pop the preserved value (below the quotation)
1697        let (rest_stack, preserved_type) = stack_after_quot.clone().pop().ok_or_else(|| {
1698            format!(
1699                "{}dip: stack underflow - expected a value below the quotation",
1700                line_prefix
1701            )
1702        })?;
1703
1704        // Freshen and apply the quotation's effect to the stack below the preserved value
1705        let fresh_effect = self.freshen_effect(&quot_effect);
1706        let (result_stack, subst) =
1707            self.apply_effect(&fresh_effect, rest_stack, "dip (quotation)", span)?;
1708
1709        // Push the preserved value back on top, applying substitution in case
1710        // preserved_type contains type variables resolved during unification
1711        let resolved_preserved = subst.apply_type(&preserved_type);
1712        let result_stack = result_stack.push(resolved_preserved);
1713
1714        let propagated_effects = fresh_effect.effects.clone();
1715        Ok((result_stack, subst, propagated_effects))
1716    }
1717
1718    /// Infer the stack effect of `keep`: ( ..a x quot -- ..b x )
1719    ///
1720    /// Run the quotation on the value (quotation receives x), then
1721    /// restore the original value on top. Like `over >aux call aux>`.
1722    fn infer_keep(
1723        &self,
1724        span: &Option<crate::ast::Span>,
1725        current_stack: StackType,
1726    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1727        let line_prefix = self.line_prefix();
1728
1729        // Pop the quotation
1730        let (stack_after_quot, quot_type) = current_stack.clone().pop().ok_or_else(|| {
1731            format!(
1732                "{}keep: stack underflow - expected quotation on stack",
1733                line_prefix
1734            )
1735        })?;
1736
1737        // Extract the quotation's effect
1738        let quot_effect = match &quot_type {
1739            Type::Quotation(effect) => (**effect).clone(),
1740            Type::Closure { effect, .. } => (**effect).clone(),
1741            Type::Var(_) => {
1742                let effect = self
1743                    .lookup_word_effect("keep")
1744                    .ok_or_else(|| "Unknown word: 'keep'".to_string())?;
1745                let fresh_effect = self.freshen_effect(&effect);
1746                let (result_stack, subst) =
1747                    self.apply_effect(&fresh_effect, current_stack, "keep", span)?;
1748                return Ok((result_stack, subst, vec![]));
1749            }
1750            _ => {
1751                return Err(format!(
1752                    "{}keep: expected quotation or closure on top of stack, got {}",
1753                    line_prefix, quot_type
1754                ));
1755            }
1756        };
1757
1758        if quot_effect.has_yield() {
1759            return Err("keep: quotation must not have Yield effects.\n\
1760                 Use strand.weave for quotations that yield."
1761                .to_string());
1762        }
1763
1764        // Peek at the preserved value type (it stays, we just need its type)
1765        let (_rest_stack, preserved_type) = stack_after_quot.clone().pop().ok_or_else(|| {
1766            format!(
1767                "{}keep: stack underflow - expected a value below the quotation",
1768                line_prefix
1769            )
1770        })?;
1771
1772        // The quotation receives x on the stack (stack_after_quot still has x on top).
1773        // Apply the quotation's effect to the stack INCLUDING x.
1774        let fresh_effect = self.freshen_effect(&quot_effect);
1775        let (result_stack, subst) =
1776            self.apply_effect(&fresh_effect, stack_after_quot, "keep (quotation)", span)?;
1777
1778        // Push the preserved value back on top, applying substitution in case
1779        // preserved_type contains type variables resolved during unification
1780        let resolved_preserved = subst.apply_type(&preserved_type);
1781        let result_stack = result_stack.push(resolved_preserved);
1782
1783        let propagated_effects = fresh_effect.effects.clone();
1784        Ok((result_stack, subst, propagated_effects))
1785    }
1786
1787    /// Infer the stack effect of `bi`: ( ..a x quot1 quot2 -- ..c )
1788    ///
1789    /// Apply two quotations to the same value. First quotation receives x,
1790    /// then second quotation receives x on top of the first quotation's results.
1791    fn infer_bi(
1792        &self,
1793        span: &Option<crate::ast::Span>,
1794        current_stack: StackType,
1795    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1796        let line_prefix = self.line_prefix();
1797
1798        // Pop quot2 (top)
1799        let (stack1, quot2_type) = current_stack.clone().pop().ok_or_else(|| {
1800            format!(
1801                "{}bi: stack underflow - expected second quotation on stack",
1802                line_prefix
1803            )
1804        })?;
1805
1806        // Pop quot1
1807        let (stack2, quot1_type) = stack1.clone().pop().ok_or_else(|| {
1808            format!(
1809                "{}bi: stack underflow - expected first quotation on stack",
1810                line_prefix
1811            )
1812        })?;
1813
1814        // Extract both quotation effects
1815        let quot1_effect = match &quot1_type {
1816            Type::Quotation(effect) => (**effect).clone(),
1817            Type::Closure { effect, .. } => (**effect).clone(),
1818            Type::Var(_) => {
1819                let effect = self
1820                    .lookup_word_effect("bi")
1821                    .ok_or_else(|| "Unknown word: 'bi'".to_string())?;
1822                let fresh_effect = self.freshen_effect(&effect);
1823                let (result_stack, subst) =
1824                    self.apply_effect(&fresh_effect, current_stack, "bi", span)?;
1825                return Ok((result_stack, subst, vec![]));
1826            }
1827            _ => {
1828                return Err(format!(
1829                    "{}bi: expected quotation or closure as first quotation, got {}",
1830                    line_prefix, quot1_type
1831                ));
1832            }
1833        };
1834
1835        let quot2_effect = match &quot2_type {
1836            Type::Quotation(effect) => (**effect).clone(),
1837            Type::Closure { effect, .. } => (**effect).clone(),
1838            Type::Var(_) => {
1839                let effect = self
1840                    .lookup_word_effect("bi")
1841                    .ok_or_else(|| "Unknown word: 'bi'".to_string())?;
1842                let fresh_effect = self.freshen_effect(&effect);
1843                let (result_stack, subst) =
1844                    self.apply_effect(&fresh_effect, current_stack, "bi", span)?;
1845                return Ok((result_stack, subst, vec![]));
1846            }
1847            _ => {
1848                return Err(format!(
1849                    "{}bi: expected quotation or closure as second quotation, got {}",
1850                    line_prefix, quot2_type
1851                ));
1852            }
1853        };
1854
1855        if quot1_effect.has_yield() || quot2_effect.has_yield() {
1856            return Err("bi: quotations must not have Yield effects.\n\
1857                 Use strand.weave for quotations that yield."
1858                .to_string());
1859        }
1860
1861        // stack2 has x on top (the value both quotations operate on)
1862        // Peek at x's type for the second application
1863        let (_rest, preserved_type) = stack2.clone().pop().ok_or_else(|| {
1864            format!(
1865                "{}bi: stack underflow - expected a value below the quotations",
1866                line_prefix
1867            )
1868        })?;
1869
1870        // Apply quot1 to stack including x
1871        let fresh_effect1 = self.freshen_effect(&quot1_effect);
1872        let (after_quot1, subst1) =
1873            self.apply_effect(&fresh_effect1, stack2, "bi (first quotation)", span)?;
1874
1875        // Push x again for quot2, applying subst1 in case preserved_type
1876        // contains type variables that were resolved during quot1's unification
1877        let resolved_preserved = subst1.apply_type(&preserved_type);
1878        let with_x = after_quot1.push(resolved_preserved);
1879
1880        // Apply quot2
1881        let fresh_effect2 = self.freshen_effect(&quot2_effect);
1882        let (result_stack, subst2) =
1883            self.apply_effect(&fresh_effect2, with_x, "bi (second quotation)", span)?;
1884
1885        let subst = subst1.compose(&subst2);
1886
1887        let mut effects = fresh_effect1.effects.clone();
1888        for e in fresh_effect2.effects.clone() {
1889            if !effects.contains(&e) {
1890                effects.push(e);
1891            }
1892        }
1893
1894        Ok((result_stack, subst, effects))
1895    }
1896
1897    /// Infer the resulting stack type after a statement
1898    /// Takes current stack, returns (new stack, substitution, side effects) after statement
1899    fn infer_statement(
1900        &self,
1901        statement: &Statement,
1902        current_stack: StackType,
1903    ) -> Result<(StackType, Subst, Vec<SideEffect>), String> {
1904        match statement {
1905            Statement::IntLiteral(_) => Ok((current_stack.push(Type::Int), Subst::empty(), vec![])),
1906            Statement::BoolLiteral(_) => {
1907                Ok((current_stack.push(Type::Bool), Subst::empty(), vec![]))
1908            }
1909            Statement::StringLiteral(_) => {
1910                Ok((current_stack.push(Type::String), Subst::empty(), vec![]))
1911            }
1912            Statement::FloatLiteral(_) => {
1913                Ok((current_stack.push(Type::Float), Subst::empty(), vec![]))
1914            }
1915            Statement::Symbol(_) => Ok((current_stack.push(Type::Symbol), Subst::empty(), vec![])),
1916            Statement::Match { arms, span } => self.infer_match(arms, span, current_stack),
1917            Statement::WordCall { name, span } => self.infer_word_call(name, span, current_stack),
1918            Statement::If {
1919                then_branch,
1920                else_branch,
1921                span,
1922            } => self.infer_if(then_branch, else_branch, span, current_stack),
1923            Statement::Quotation { id, body, .. } => self.infer_quotation(*id, body, current_stack),
1924        }
1925    }
1926
1927    /// Look up the effect of a word (built-in or user-defined)
1928    fn lookup_word_effect(&self, name: &str) -> Option<Effect> {
1929        // First check built-ins
1930        if let Some(effect) = builtin_signature(name) {
1931            return Some(effect);
1932        }
1933
1934        // Then check user-defined words
1935        self.env.get(name).cloned()
1936    }
1937
1938    /// Apply an effect to a stack
1939    /// Effect: (inputs -- outputs)
1940    /// Current stack must match inputs, result is outputs
1941    /// Returns (result_stack, substitution)
1942    fn apply_effect(
1943        &self,
1944        effect: &Effect,
1945        current_stack: StackType,
1946        operation: &str,
1947        span: &Option<crate::ast::Span>,
1948    ) -> Result<(StackType, Subst), String> {
1949        // Check for stack underflow: if the effect needs more concrete values than
1950        // the current stack provides, and the stack has a "rigid" row variable at its base,
1951        // this would be unsound (the row var could be Empty at runtime).
1952        // Bug #169: "phantom stack entries"
1953        //
1954        // We only check for "rigid" row variables (named "rest" from declared effects).
1955        // Row variables named "input" are from inference and CAN grow to discover requirements.
1956        let effect_concrete = Self::count_concrete_types(&effect.inputs);
1957        let stack_concrete = Self::count_concrete_types(&current_stack);
1958
1959        if let Some(row_var_name) = Self::get_row_var_base(&current_stack) {
1960            // Only check "rigid" row variables (from declared effects, not inference).
1961            //
1962            // Row variable naming convention (established in parser.rs:build_stack_type):
1963            // - "rest": Created by the parser for declared stack effects. When a word declares
1964            //   `( String Int -- String )`, the parser creates `( ..rest String Int -- ..rest String )`.
1965            //   This "rest" is rigid because the caller guarantees exactly these concrete types.
1966            // - "rest$N": Freshened versions created during type checking when calling other words.
1967            //   These represent the callee's stack context and can grow during unification.
1968            // - "input": Created for words without declared effects during inference.
1969            //   These are flexible and grow to discover the word's actual requirements.
1970            //
1971            // Only the original "rest" (exact match) should trigger underflow checking.
1972            let is_rigid = row_var_name == "rest";
1973
1974            if is_rigid && effect_concrete > stack_concrete {
1975                let word_name = self
1976                    .current_word
1977                    .borrow()
1978                    .as_ref()
1979                    .map(|(n, _)| n.clone())
1980                    .unwrap_or_else(|| "unknown".to_string());
1981                return Err(format!(
1982                    "{}In '{}': {}: stack underflow - requires {} value(s), only {} provided",
1983                    self.line_prefix(),
1984                    word_name,
1985                    operation,
1986                    effect_concrete,
1987                    stack_concrete
1988                ));
1989            }
1990        }
1991
1992        // Unify current stack with effect's input
1993        let subst = unify_stacks(&effect.inputs, &current_stack).map_err(|e| {
1994            let line_info = span
1995                .as_ref()
1996                .map(|s| format_line_prefix(s.line))
1997                .unwrap_or_default();
1998            format!(
1999                "{}{}: stack type mismatch. Expected {}, got {}: {}",
2000                line_info, operation, effect.inputs, current_stack, e
2001            )
2002        })?;
2003
2004        // Apply substitution to output
2005        let result_stack = subst.apply_stack(&effect.outputs);
2006
2007        Ok((result_stack, subst))
2008    }
2009
2010    /// Count the number of concrete (non-row-variable) types in a stack
2011    fn count_concrete_types(stack: &StackType) -> usize {
2012        let mut count = 0;
2013        let mut current = stack;
2014        while let StackType::Cons { rest, top: _ } = current {
2015            count += 1;
2016            current = rest;
2017        }
2018        count
2019    }
2020
2021    /// Get the row variable name at the base of a stack, if any
2022    fn get_row_var_base(stack: &StackType) -> Option<String> {
2023        let mut current = stack;
2024        while let StackType::Cons { rest, top: _ } = current {
2025            current = rest;
2026        }
2027        match current {
2028            StackType::RowVar(name) => Some(name.clone()),
2029            _ => None,
2030        }
2031    }
2032
2033    /// Adjust stack for strand.spawn operation by converting Quotation to Closure if needed
2034    ///
2035    /// strand.spawn expects Quotation(Empty -- Empty), but if we have Quotation(T... -- U...)
2036    /// with non-empty inputs, we auto-convert it to a Closure that captures those inputs.
2037    fn adjust_stack_for_spawn(
2038        &self,
2039        current_stack: StackType,
2040        spawn_effect: &Effect,
2041    ) -> Result<StackType, String> {
2042        // strand.spawn expects: ( ..a Quotation(Empty -- Empty) -- ..a Int )
2043        // Extract the expected quotation type from strand.spawn's effect
2044        let expected_quot_type = match &spawn_effect.inputs {
2045            StackType::Cons { top, rest: _ } => {
2046                if !matches!(top, Type::Quotation(_)) {
2047                    return Ok(current_stack); // Not a quotation, don't adjust
2048                }
2049                top
2050            }
2051            _ => return Ok(current_stack),
2052        };
2053
2054        // Check what's actually on the stack
2055        let (rest_stack, actual_type) = match &current_stack {
2056            StackType::Cons { rest, top } => (rest.as_ref().clone(), top),
2057            _ => return Ok(current_stack), // Empty stack, nothing to adjust
2058        };
2059
2060        // If top of stack is a Quotation with non-empty inputs, convert to Closure
2061        if let Type::Quotation(actual_effect) = actual_type {
2062            // Check if quotation needs inputs
2063            if !matches!(actual_effect.inputs, StackType::Empty) {
2064                // Extract expected effect from spawn's signature
2065                let expected_effect = match expected_quot_type {
2066                    Type::Quotation(eff) => eff.as_ref(),
2067                    _ => return Ok(current_stack),
2068                };
2069
2070                // Calculate what needs to be captured
2071                let captures = calculate_captures(actual_effect, expected_effect)?;
2072
2073                // Create a Closure type
2074                let closure_type = Type::Closure {
2075                    effect: Box::new(expected_effect.clone()),
2076                    captures: captures.clone(),
2077                };
2078
2079                // Pop the captured values from the stack
2080                // The values to capture are BELOW the quotation on the stack
2081                let mut adjusted_stack = rest_stack;
2082                for _ in &captures {
2083                    adjusted_stack = match adjusted_stack {
2084                        StackType::Cons { rest, .. } => rest.as_ref().clone(),
2085                        _ => {
2086                            return Err(format!(
2087                                "strand.spawn: not enough values on stack to capture. Need {} values",
2088                                captures.len()
2089                            ));
2090                        }
2091                    };
2092                }
2093
2094                // Push the Closure onto the adjusted stack
2095                return Ok(adjusted_stack.push(closure_type));
2096            }
2097        }
2098
2099        Ok(current_stack)
2100    }
2101
2102    /// Analyze quotation captures
2103    ///
2104    /// Determines whether a quotation should be stateless (Type::Quotation)
2105    /// or a closure (Type::Closure) based on the expected type from the word signature.
2106    ///
2107    /// Type-driven inference with automatic closure creation:
2108    ///   - If expected type is Closure[effect], calculate what to capture
2109    ///   - If expected type is Quotation[effect]:
2110    ///     - If body needs more inputs than expected effect, auto-create Closure
2111    ///     - Otherwise return stateless Quotation
2112    ///   - If no expected type, default to stateless (conservative)
2113    ///
2114    /// Example 1 (auto-create closure):
2115    ///   Expected: Quotation[-- ]          [spawn expects ( -- )]
2116    ///   Body: [ handle-connection ]       [needs ( Int -- )]
2117    ///   Body effect: ( Int -- )           [needs 1 Int]
2118    ///   Expected effect: ( -- )           [provides 0 inputs]
2119    ///   Result: Closure { effect: ( -- ), captures: [Int] }
2120    ///
2121    /// Example 2 (explicit closure):
2122    ///   Signature: ( Int -- Closure[Int -- Int] )
2123    ///   Body: [ add ]
2124    ///   Body effect: ( Int Int -- Int )  [add needs 2 Ints]
2125    ///   Expected effect: [Int -- Int]    [call site provides 1 Int]
2126    ///   Result: Closure { effect: [Int -- Int], captures: [Int] }
2127    fn analyze_captures(
2128        &self,
2129        body_effect: &Effect,
2130        _current_stack: &StackType,
2131    ) -> Result<Type, String> {
2132        // Check if there's an expected type from the word signature
2133        let expected = self.expected_quotation_type.borrow().clone();
2134
2135        match expected {
2136            Some(Type::Closure { effect, .. }) => {
2137                // User declared closure type - calculate captures
2138                let captures = calculate_captures(body_effect, &effect)?;
2139                Ok(Type::Closure { effect, captures })
2140            }
2141            Some(Type::Quotation(expected_effect)) => {
2142                // Check if we need to auto-create a closure by comparing the
2143                // body's concrete input count against what the combinator provides.
2144                let body_inputs = extract_concrete_types(&body_effect.inputs);
2145                let expected_inputs = extract_concrete_types(&expected_effect.inputs);
2146
2147                // Auto-capture triggers when the body needs more concrete inputs
2148                // than the expected provides. Three branches:
2149                // (a) Expected is empty (strand.spawn): body needs any inputs → capture all.
2150                // (b) Expected has concrete inputs (list.fold): body has MORE → capture excess.
2151                // (c) Expected has ONLY a row variable and no concrete inputs
2152                //     (strand.weave): don't capture, fall through to unification.
2153                let expected_is_empty = matches!(expected_effect.inputs, StackType::Empty);
2154                let should_capture = if expected_is_empty {
2155                    !body_inputs.is_empty()
2156                } else if !expected_inputs.is_empty() {
2157                    body_inputs.len() > expected_inputs.len()
2158                } else {
2159                    false // row-variable-only expected — don't capture, unify instead
2160                };
2161
2162                if should_capture {
2163                    // Body needs more inputs than the combinator provides.
2164                    // The excess (bottommost) become captures; the topmost must
2165                    // align with what the combinator provides.
2166                    //
2167                    // Example: list.fold expects ( ..b Acc T -- ..b Acc ).
2168                    // Body inferred as ( ..b X Acc T -- ..b Acc ).
2169                    // expected_inputs = [Acc, T], body_inputs = [X, Acc, T].
2170                    // Captures = [X]. Topmost 2 of body must match expected's 2.
2171                    //
2172                    // Issue #395: this extends the empty-input auto-capture
2173                    // (used by strand.spawn) to the non-empty case.
2174                    let captures = calculate_captures(body_effect, &expected_effect)?;
2175                    Ok(Type::Closure {
2176                        effect: expected_effect,
2177                        captures,
2178                    })
2179                } else {
2180                    // Body has same or fewer inputs — standard unification path.
2181                    // This catches:
2182                    // - Stack pollution: body pushes values when expected is stack-neutral
2183                    // - Stack underflow: body consumes values when expected is stack-neutral
2184                    // - Wrong return type: body returns Int when Bool expected
2185                    let body_quot = Type::Quotation(Box::new(body_effect.clone()));
2186                    let expected_quot = Type::Quotation(expected_effect.clone());
2187                    unify_types(&body_quot, &expected_quot).map_err(|e| {
2188                        format!(
2189                            "quotation effect mismatch: expected {}, got {}: {}",
2190                            expected_effect, body_effect, e
2191                        )
2192                    })?;
2193
2194                    // Body is compatible with expected effect - stateless quotation
2195                    Ok(Type::Quotation(expected_effect))
2196                }
2197            }
2198            _ => {
2199                // No expected type - conservative default: stateless quotation
2200                Ok(Type::Quotation(Box::new(body_effect.clone())))
2201            }
2202        }
2203    }
2204
2205    /// Check if an inferred effect matches any of the declared effects
2206    /// Effects match by kind (e.g., Yield matches Yield, regardless of type parameters)
2207    /// Type parameters should unify, but for now we just check the effect kind
2208    fn effect_matches_any(&self, inferred: &SideEffect, declared: &[SideEffect]) -> bool {
2209        declared.iter().any(|decl| match (inferred, decl) {
2210            (SideEffect::Yield(_), SideEffect::Yield(_)) => true,
2211        })
2212    }
2213
2214    /// Pop a type from a stack type, returning (rest, top)
2215    fn pop_type(&self, stack: &StackType, context: &str) -> Result<(StackType, Type), String> {
2216        match stack {
2217            StackType::Cons { rest, top } => Ok(((**rest).clone(), top.clone())),
2218            StackType::Empty => Err(format!(
2219                "{}: stack underflow - expected value on stack but stack is empty",
2220                context
2221            )),
2222            StackType::RowVar(_) => {
2223                // Can't statically determine if row variable is empty
2224                // For now, assume it has at least one element
2225                // This is conservative - real implementation would track constraints
2226                Err(format!(
2227                    "{}: cannot pop from polymorphic stack without more type information",
2228                    context
2229                ))
2230            }
2231        }
2232    }
2233}
2234
2235impl Default for TypeChecker {
2236    fn default() -> Self {
2237        Self::new()
2238    }
2239}
2240
2241#[cfg(test)]
2242mod tests {
2243    use super::*;
2244
2245    #[test]
2246    fn test_simple_literal() {
2247        let program = Program {
2248            includes: vec![],
2249            unions: vec![],
2250            words: vec![WordDef {
2251                name: "test".to_string(),
2252                effect: Some(Effect::new(
2253                    StackType::Empty,
2254                    StackType::singleton(Type::Int),
2255                )),
2256                body: vec![Statement::IntLiteral(42)],
2257                source: None,
2258                allowed_lints: vec![],
2259            }],
2260        };
2261
2262        let mut checker = TypeChecker::new();
2263        assert!(checker.check_program(&program).is_ok());
2264    }
2265
2266    #[test]
2267    fn test_simple_operation() {
2268        // : test ( Int Int -- Int ) add ;
2269        let program = Program {
2270            includes: vec![],
2271            unions: vec![],
2272            words: vec![WordDef {
2273                name: "test".to_string(),
2274                effect: Some(Effect::new(
2275                    StackType::Empty.push(Type::Int).push(Type::Int),
2276                    StackType::singleton(Type::Int),
2277                )),
2278                body: vec![Statement::WordCall {
2279                    name: "i.add".to_string(),
2280                    span: None,
2281                }],
2282                source: None,
2283                allowed_lints: vec![],
2284            }],
2285        };
2286
2287        let mut checker = TypeChecker::new();
2288        assert!(checker.check_program(&program).is_ok());
2289    }
2290
2291    #[test]
2292    fn test_type_mismatch() {
2293        // : test ( String -- ) io.write-line ;  with body: 42
2294        let program = Program {
2295            includes: vec![],
2296            unions: vec![],
2297            words: vec![WordDef {
2298                name: "test".to_string(),
2299                effect: Some(Effect::new(
2300                    StackType::singleton(Type::String),
2301                    StackType::Empty,
2302                )),
2303                body: vec![
2304                    Statement::IntLiteral(42), // Pushes Int, not String!
2305                    Statement::WordCall {
2306                        name: "io.write-line".to_string(),
2307                        span: None,
2308                    },
2309                ],
2310                source: None,
2311                allowed_lints: vec![],
2312            }],
2313        };
2314
2315        let mut checker = TypeChecker::new();
2316        let result = checker.check_program(&program);
2317        assert!(result.is_err());
2318        assert!(result.unwrap_err().contains("Type mismatch"));
2319    }
2320
2321    #[test]
2322    fn test_polymorphic_dup() {
2323        // : my-dup ( Int -- Int Int ) dup ;
2324        let program = Program {
2325            includes: vec![],
2326            unions: vec![],
2327            words: vec![WordDef {
2328                name: "my-dup".to_string(),
2329                effect: Some(Effect::new(
2330                    StackType::singleton(Type::Int),
2331                    StackType::Empty.push(Type::Int).push(Type::Int),
2332                )),
2333                body: vec![Statement::WordCall {
2334                    name: "dup".to_string(),
2335                    span: None,
2336                }],
2337                source: None,
2338                allowed_lints: vec![],
2339            }],
2340        };
2341
2342        let mut checker = TypeChecker::new();
2343        assert!(checker.check_program(&program).is_ok());
2344    }
2345
2346    #[test]
2347    fn test_conditional_branches() {
2348        // : test ( Int Int -- String )
2349        //   > if "greater" else "not greater" then ;
2350        let program = Program {
2351            includes: vec![],
2352            unions: vec![],
2353            words: vec![WordDef {
2354                name: "test".to_string(),
2355                effect: Some(Effect::new(
2356                    StackType::Empty.push(Type::Int).push(Type::Int),
2357                    StackType::singleton(Type::String),
2358                )),
2359                body: vec![
2360                    Statement::WordCall {
2361                        name: "i.>".to_string(),
2362                        span: None,
2363                    },
2364                    Statement::If {
2365                        then_branch: vec![Statement::StringLiteral("greater".to_string())],
2366                        else_branch: Some(vec![Statement::StringLiteral(
2367                            "not greater".to_string(),
2368                        )]),
2369                        span: None,
2370                    },
2371                ],
2372                source: None,
2373                allowed_lints: vec![],
2374            }],
2375        };
2376
2377        let mut checker = TypeChecker::new();
2378        assert!(checker.check_program(&program).is_ok());
2379    }
2380
2381    #[test]
2382    fn test_mismatched_branches() {
2383        // : test ( -- Int )
2384        //   true if 42 else "string" then ;  // ERROR: incompatible types
2385        let program = Program {
2386            includes: vec![],
2387            unions: vec![],
2388            words: vec![WordDef {
2389                name: "test".to_string(),
2390                effect: Some(Effect::new(
2391                    StackType::Empty,
2392                    StackType::singleton(Type::Int),
2393                )),
2394                body: vec![
2395                    Statement::BoolLiteral(true),
2396                    Statement::If {
2397                        then_branch: vec![Statement::IntLiteral(42)],
2398                        else_branch: Some(vec![Statement::StringLiteral("string".to_string())]),
2399                        span: None,
2400                    },
2401                ],
2402                source: None,
2403                allowed_lints: vec![],
2404            }],
2405        };
2406
2407        let mut checker = TypeChecker::new();
2408        let result = checker.check_program(&program);
2409        assert!(result.is_err());
2410        assert!(result.unwrap_err().contains("incompatible"));
2411    }
2412
2413    #[test]
2414    fn test_user_defined_word_call() {
2415        // : helper ( Int -- String ) int->string ;
2416        // : main ( -- ) 42 helper io.write-line ;
2417        let program = Program {
2418            includes: vec![],
2419            unions: vec![],
2420            words: vec![
2421                WordDef {
2422                    name: "helper".to_string(),
2423                    effect: Some(Effect::new(
2424                        StackType::singleton(Type::Int),
2425                        StackType::singleton(Type::String),
2426                    )),
2427                    body: vec![Statement::WordCall {
2428                        name: "int->string".to_string(),
2429                        span: None,
2430                    }],
2431                    source: None,
2432                    allowed_lints: vec![],
2433                },
2434                WordDef {
2435                    name: "main".to_string(),
2436                    effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
2437                    body: vec![
2438                        Statement::IntLiteral(42),
2439                        Statement::WordCall {
2440                            name: "helper".to_string(),
2441                            span: None,
2442                        },
2443                        Statement::WordCall {
2444                            name: "io.write-line".to_string(),
2445                            span: None,
2446                        },
2447                    ],
2448                    source: None,
2449                    allowed_lints: vec![],
2450                },
2451            ],
2452        };
2453
2454        let mut checker = TypeChecker::new();
2455        assert!(checker.check_program(&program).is_ok());
2456    }
2457
2458    #[test]
2459    fn test_arithmetic_chain() {
2460        // : test ( Int Int Int -- Int )
2461        //   add multiply ;
2462        let program = Program {
2463            includes: vec![],
2464            unions: vec![],
2465            words: vec![WordDef {
2466                name: "test".to_string(),
2467                effect: Some(Effect::new(
2468                    StackType::Empty
2469                        .push(Type::Int)
2470                        .push(Type::Int)
2471                        .push(Type::Int),
2472                    StackType::singleton(Type::Int),
2473                )),
2474                body: vec![
2475                    Statement::WordCall {
2476                        name: "i.add".to_string(),
2477                        span: None,
2478                    },
2479                    Statement::WordCall {
2480                        name: "i.multiply".to_string(),
2481                        span: None,
2482                    },
2483                ],
2484                source: None,
2485                allowed_lints: vec![],
2486            }],
2487        };
2488
2489        let mut checker = TypeChecker::new();
2490        assert!(checker.check_program(&program).is_ok());
2491    }
2492
2493    #[test]
2494    fn test_write_line_type_error() {
2495        // : test ( Int -- ) io.write-line ;  // ERROR: io.write-line expects String
2496        let program = Program {
2497            includes: vec![],
2498            unions: vec![],
2499            words: vec![WordDef {
2500                name: "test".to_string(),
2501                effect: Some(Effect::new(
2502                    StackType::singleton(Type::Int),
2503                    StackType::Empty,
2504                )),
2505                body: vec![Statement::WordCall {
2506                    name: "io.write-line".to_string(),
2507                    span: None,
2508                }],
2509                source: None,
2510                allowed_lints: vec![],
2511            }],
2512        };
2513
2514        let mut checker = TypeChecker::new();
2515        let result = checker.check_program(&program);
2516        assert!(result.is_err());
2517        assert!(result.unwrap_err().contains("Type mismatch"));
2518    }
2519
2520    #[test]
2521    fn test_stack_underflow_drop() {
2522        // : test ( -- ) drop ;  // ERROR: can't drop from empty stack
2523        let program = Program {
2524            includes: vec![],
2525            unions: vec![],
2526            words: vec![WordDef {
2527                name: "test".to_string(),
2528                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
2529                body: vec![Statement::WordCall {
2530                    name: "drop".to_string(),
2531                    span: None,
2532                }],
2533                source: None,
2534                allowed_lints: vec![],
2535            }],
2536        };
2537
2538        let mut checker = TypeChecker::new();
2539        let result = checker.check_program(&program);
2540        assert!(result.is_err());
2541        assert!(result.unwrap_err().contains("mismatch"));
2542    }
2543
2544    #[test]
2545    fn test_stack_underflow_add() {
2546        // : test ( Int -- Int ) add ;  // ERROR: add needs 2 values
2547        let program = Program {
2548            includes: vec![],
2549            unions: vec![],
2550            words: vec![WordDef {
2551                name: "test".to_string(),
2552                effect: Some(Effect::new(
2553                    StackType::singleton(Type::Int),
2554                    StackType::singleton(Type::Int),
2555                )),
2556                body: vec![Statement::WordCall {
2557                    name: "i.add".to_string(),
2558                    span: None,
2559                }],
2560                source: None,
2561                allowed_lints: vec![],
2562            }],
2563        };
2564
2565        let mut checker = TypeChecker::new();
2566        let result = checker.check_program(&program);
2567        assert!(result.is_err());
2568        assert!(result.unwrap_err().contains("mismatch"));
2569    }
2570
2571    /// Issue #169: rot with only 2 values should fail at compile time
2572    /// Previously this was silently accepted due to implicit row polymorphism
2573    #[test]
2574    fn test_stack_underflow_rot_issue_169() {
2575        // : test ( -- ) 3 4 rot ;  // ERROR: rot needs 3 values, only 2 provided
2576        // Note: The parser generates `( ..rest -- ..rest )` for `( -- )`, so we use RowVar("rest")
2577        // to match the actual parsing behavior. The "rest" row variable is rigid.
2578        let program = Program {
2579            includes: vec![],
2580            unions: vec![],
2581            words: vec![WordDef {
2582                name: "test".to_string(),
2583                effect: Some(Effect::new(
2584                    StackType::RowVar("rest".to_string()),
2585                    StackType::RowVar("rest".to_string()),
2586                )),
2587                body: vec![
2588                    Statement::IntLiteral(3),
2589                    Statement::IntLiteral(4),
2590                    Statement::WordCall {
2591                        name: "rot".to_string(),
2592                        span: None,
2593                    },
2594                ],
2595                source: None,
2596                allowed_lints: vec![],
2597            }],
2598        };
2599
2600        let mut checker = TypeChecker::new();
2601        let result = checker.check_program(&program);
2602        assert!(result.is_err(), "rot with 2 values should fail");
2603        let err = result.unwrap_err();
2604        assert!(
2605            err.contains("stack underflow") || err.contains("requires 3"),
2606            "Error should mention underflow: {}",
2607            err
2608        );
2609    }
2610
2611    #[test]
2612    fn test_csp_operations() {
2613        // : test ( -- )
2614        //   chan.make     # ( -- Channel )
2615        //   42 swap       # ( Channel Int -- Int Channel )
2616        //   chan.send     # ( Int Channel -- Bool )
2617        //   drop          # ( Bool -- )
2618        // ;
2619        let program = Program {
2620            includes: vec![],
2621            unions: vec![],
2622            words: vec![WordDef {
2623                name: "test".to_string(),
2624                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
2625                body: vec![
2626                    Statement::WordCall {
2627                        name: "chan.make".to_string(),
2628                        span: None,
2629                    },
2630                    Statement::IntLiteral(42),
2631                    Statement::WordCall {
2632                        name: "swap".to_string(),
2633                        span: None,
2634                    },
2635                    Statement::WordCall {
2636                        name: "chan.send".to_string(),
2637                        span: None,
2638                    },
2639                    Statement::WordCall {
2640                        name: "drop".to_string(),
2641                        span: None,
2642                    },
2643                ],
2644                source: None,
2645                allowed_lints: vec![],
2646            }],
2647        };
2648
2649        let mut checker = TypeChecker::new();
2650        assert!(checker.check_program(&program).is_ok());
2651    }
2652
2653    #[test]
2654    fn test_complex_stack_shuffling() {
2655        // : test ( Int Int Int -- Int )
2656        //   rot add add ;
2657        let program = Program {
2658            includes: vec![],
2659            unions: vec![],
2660            words: vec![WordDef {
2661                name: "test".to_string(),
2662                effect: Some(Effect::new(
2663                    StackType::Empty
2664                        .push(Type::Int)
2665                        .push(Type::Int)
2666                        .push(Type::Int),
2667                    StackType::singleton(Type::Int),
2668                )),
2669                body: vec![
2670                    Statement::WordCall {
2671                        name: "rot".to_string(),
2672                        span: None,
2673                    },
2674                    Statement::WordCall {
2675                        name: "i.add".to_string(),
2676                        span: None,
2677                    },
2678                    Statement::WordCall {
2679                        name: "i.add".to_string(),
2680                        span: None,
2681                    },
2682                ],
2683                source: None,
2684                allowed_lints: vec![],
2685            }],
2686        };
2687
2688        let mut checker = TypeChecker::new();
2689        assert!(checker.check_program(&program).is_ok());
2690    }
2691
2692    #[test]
2693    fn test_empty_program() {
2694        // Program with no words should be valid
2695        let program = Program {
2696            includes: vec![],
2697            unions: vec![],
2698            words: vec![],
2699        };
2700
2701        let mut checker = TypeChecker::new();
2702        assert!(checker.check_program(&program).is_ok());
2703    }
2704
2705    #[test]
2706    fn test_word_without_effect_declaration() {
2707        // : helper 42 ;  // No effect declaration - should error
2708        let program = Program {
2709            includes: vec![],
2710            unions: vec![],
2711            words: vec![WordDef {
2712                name: "helper".to_string(),
2713                effect: None,
2714                body: vec![Statement::IntLiteral(42)],
2715                source: None,
2716                allowed_lints: vec![],
2717            }],
2718        };
2719
2720        let mut checker = TypeChecker::new();
2721        let result = checker.check_program(&program);
2722        assert!(result.is_err());
2723        assert!(
2724            result
2725                .unwrap_err()
2726                .contains("missing a stack effect declaration")
2727        );
2728    }
2729
2730    #[test]
2731    fn test_nested_conditionals() {
2732        // : test ( Int Int Int Int -- String )
2733        //   > if
2734        //     > if "both true" else "first true" then
2735        //   else
2736        //     drop drop "first false"
2737        //   then ;
2738        // Note: Needs 4 Ints total (2 for each > comparison)
2739        // Else branch must drop unused Ints to match then branch's stack effect
2740        let program = Program {
2741            includes: vec![],
2742            unions: vec![],
2743            words: vec![WordDef {
2744                name: "test".to_string(),
2745                effect: Some(Effect::new(
2746                    StackType::Empty
2747                        .push(Type::Int)
2748                        .push(Type::Int)
2749                        .push(Type::Int)
2750                        .push(Type::Int),
2751                    StackType::singleton(Type::String),
2752                )),
2753                body: vec![
2754                    Statement::WordCall {
2755                        name: "i.>".to_string(),
2756                        span: None,
2757                    },
2758                    Statement::If {
2759                        then_branch: vec![
2760                            Statement::WordCall {
2761                                name: "i.>".to_string(),
2762                                span: None,
2763                            },
2764                            Statement::If {
2765                                then_branch: vec![Statement::StringLiteral(
2766                                    "both true".to_string(),
2767                                )],
2768                                else_branch: Some(vec![Statement::StringLiteral(
2769                                    "first true".to_string(),
2770                                )]),
2771                                span: None,
2772                            },
2773                        ],
2774                        else_branch: Some(vec![
2775                            Statement::WordCall {
2776                                name: "drop".to_string(),
2777                                span: None,
2778                            },
2779                            Statement::WordCall {
2780                                name: "drop".to_string(),
2781                                span: None,
2782                            },
2783                            Statement::StringLiteral("first false".to_string()),
2784                        ]),
2785                        span: None,
2786                    },
2787                ],
2788                source: None,
2789                allowed_lints: vec![],
2790            }],
2791        };
2792
2793        let mut checker = TypeChecker::new();
2794        match checker.check_program(&program) {
2795            Ok(_) => {}
2796            Err(e) => panic!("Type check failed: {}", e),
2797        }
2798    }
2799
2800    #[test]
2801    fn test_conditional_without_else() {
2802        // : test ( Int Int -- Int )
2803        //   > if 100 then ;
2804        // Both branches must leave same stack
2805        let program = Program {
2806            includes: vec![],
2807            unions: vec![],
2808            words: vec![WordDef {
2809                name: "test".to_string(),
2810                effect: Some(Effect::new(
2811                    StackType::Empty.push(Type::Int).push(Type::Int),
2812                    StackType::singleton(Type::Int),
2813                )),
2814                body: vec![
2815                    Statement::WordCall {
2816                        name: "i.>".to_string(),
2817                        span: None,
2818                    },
2819                    Statement::If {
2820                        then_branch: vec![Statement::IntLiteral(100)],
2821                        else_branch: None, // No else - should leave stack unchanged
2822                        span: None,
2823                    },
2824                ],
2825                source: None,
2826                allowed_lints: vec![],
2827            }],
2828        };
2829
2830        let mut checker = TypeChecker::new();
2831        let result = checker.check_program(&program);
2832        // This should fail because then pushes Int but else leaves stack empty
2833        assert!(result.is_err());
2834    }
2835
2836    #[test]
2837    fn test_multiple_word_chain() {
2838        // : helper1 ( Int -- String ) int->string ;
2839        // : helper2 ( String -- ) io.write-line ;
2840        // : main ( -- ) 42 helper1 helper2 ;
2841        let program = Program {
2842            includes: vec![],
2843            unions: vec![],
2844            words: vec![
2845                WordDef {
2846                    name: "helper1".to_string(),
2847                    effect: Some(Effect::new(
2848                        StackType::singleton(Type::Int),
2849                        StackType::singleton(Type::String),
2850                    )),
2851                    body: vec![Statement::WordCall {
2852                        name: "int->string".to_string(),
2853                        span: None,
2854                    }],
2855                    source: None,
2856                    allowed_lints: vec![],
2857                },
2858                WordDef {
2859                    name: "helper2".to_string(),
2860                    effect: Some(Effect::new(
2861                        StackType::singleton(Type::String),
2862                        StackType::Empty,
2863                    )),
2864                    body: vec![Statement::WordCall {
2865                        name: "io.write-line".to_string(),
2866                        span: None,
2867                    }],
2868                    source: None,
2869                    allowed_lints: vec![],
2870                },
2871                WordDef {
2872                    name: "main".to_string(),
2873                    effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
2874                    body: vec![
2875                        Statement::IntLiteral(42),
2876                        Statement::WordCall {
2877                            name: "helper1".to_string(),
2878                            span: None,
2879                        },
2880                        Statement::WordCall {
2881                            name: "helper2".to_string(),
2882                            span: None,
2883                        },
2884                    ],
2885                    source: None,
2886                    allowed_lints: vec![],
2887                },
2888            ],
2889        };
2890
2891        let mut checker = TypeChecker::new();
2892        assert!(checker.check_program(&program).is_ok());
2893    }
2894
2895    #[test]
2896    fn test_all_stack_ops() {
2897        // : test ( Int Int Int -- Int Int Int Int )
2898        //   over nip tuck ;
2899        let program = Program {
2900            includes: vec![],
2901            unions: vec![],
2902            words: vec![WordDef {
2903                name: "test".to_string(),
2904                effect: Some(Effect::new(
2905                    StackType::Empty
2906                        .push(Type::Int)
2907                        .push(Type::Int)
2908                        .push(Type::Int),
2909                    StackType::Empty
2910                        .push(Type::Int)
2911                        .push(Type::Int)
2912                        .push(Type::Int)
2913                        .push(Type::Int),
2914                )),
2915                body: vec![
2916                    Statement::WordCall {
2917                        name: "over".to_string(),
2918                        span: None,
2919                    },
2920                    Statement::WordCall {
2921                        name: "nip".to_string(),
2922                        span: None,
2923                    },
2924                    Statement::WordCall {
2925                        name: "tuck".to_string(),
2926                        span: None,
2927                    },
2928                ],
2929                source: None,
2930                allowed_lints: vec![],
2931            }],
2932        };
2933
2934        let mut checker = TypeChecker::new();
2935        assert!(checker.check_program(&program).is_ok());
2936    }
2937
2938    #[test]
2939    fn test_mixed_types_complex() {
2940        // : test ( -- )
2941        //   42 int->string      # ( -- String )
2942        //   100 200 >           # ( String -- String Int )
2943        //   if                  # ( String -- String )
2944        //     io.write-line     # ( String -- )
2945        //   else
2946        //     io.write-line
2947        //   then ;
2948        let program = Program {
2949            includes: vec![],
2950            unions: vec![],
2951            words: vec![WordDef {
2952                name: "test".to_string(),
2953                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
2954                body: vec![
2955                    Statement::IntLiteral(42),
2956                    Statement::WordCall {
2957                        name: "int->string".to_string(),
2958                        span: None,
2959                    },
2960                    Statement::IntLiteral(100),
2961                    Statement::IntLiteral(200),
2962                    Statement::WordCall {
2963                        name: "i.>".to_string(),
2964                        span: None,
2965                    },
2966                    Statement::If {
2967                        then_branch: vec![Statement::WordCall {
2968                            name: "io.write-line".to_string(),
2969                            span: None,
2970                        }],
2971                        else_branch: Some(vec![Statement::WordCall {
2972                            name: "io.write-line".to_string(),
2973                            span: None,
2974                        }]),
2975                        span: None,
2976                    },
2977                ],
2978                source: None,
2979                allowed_lints: vec![],
2980            }],
2981        };
2982
2983        let mut checker = TypeChecker::new();
2984        assert!(checker.check_program(&program).is_ok());
2985    }
2986
2987    #[test]
2988    fn test_string_literal() {
2989        // : test ( -- String ) "hello" ;
2990        let program = Program {
2991            includes: vec![],
2992            unions: vec![],
2993            words: vec![WordDef {
2994                name: "test".to_string(),
2995                effect: Some(Effect::new(
2996                    StackType::Empty,
2997                    StackType::singleton(Type::String),
2998                )),
2999                body: vec![Statement::StringLiteral("hello".to_string())],
3000                source: None,
3001                allowed_lints: vec![],
3002            }],
3003        };
3004
3005        let mut checker = TypeChecker::new();
3006        assert!(checker.check_program(&program).is_ok());
3007    }
3008
3009    #[test]
3010    fn test_bool_literal() {
3011        // : test ( -- Bool ) true ;
3012        // Booleans are now properly typed as Bool
3013        let program = Program {
3014            includes: vec![],
3015            unions: vec![],
3016            words: vec![WordDef {
3017                name: "test".to_string(),
3018                effect: Some(Effect::new(
3019                    StackType::Empty,
3020                    StackType::singleton(Type::Bool),
3021                )),
3022                body: vec![Statement::BoolLiteral(true)],
3023                source: None,
3024                allowed_lints: vec![],
3025            }],
3026        };
3027
3028        let mut checker = TypeChecker::new();
3029        assert!(checker.check_program(&program).is_ok());
3030    }
3031
3032    #[test]
3033    fn test_type_error_in_nested_conditional() {
3034        // : test ( -- )
3035        //   10 20 i.> if
3036        //     42 io.write-line   # ERROR: io.write-line expects String, got Int
3037        //   else
3038        //     "ok" io.write-line
3039        //   then ;
3040        let program = Program {
3041            includes: vec![],
3042            unions: vec![],
3043            words: vec![WordDef {
3044                name: "test".to_string(),
3045                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
3046                body: vec![
3047                    Statement::IntLiteral(10),
3048                    Statement::IntLiteral(20),
3049                    Statement::WordCall {
3050                        name: "i.>".to_string(),
3051                        span: None,
3052                    },
3053                    Statement::If {
3054                        then_branch: vec![
3055                            Statement::IntLiteral(42),
3056                            Statement::WordCall {
3057                                name: "io.write-line".to_string(),
3058                                span: None,
3059                            },
3060                        ],
3061                        else_branch: Some(vec![
3062                            Statement::StringLiteral("ok".to_string()),
3063                            Statement::WordCall {
3064                                name: "io.write-line".to_string(),
3065                                span: None,
3066                            },
3067                        ]),
3068                        span: None,
3069                    },
3070                ],
3071                source: None,
3072                allowed_lints: vec![],
3073            }],
3074        };
3075
3076        let mut checker = TypeChecker::new();
3077        let result = checker.check_program(&program);
3078        assert!(result.is_err());
3079        assert!(result.unwrap_err().contains("Type mismatch"));
3080    }
3081
3082    #[test]
3083    fn test_read_line_operation() {
3084        // : test ( -- String Bool ) io.read-line ;
3085        // io.read-line now returns ( -- String Bool ) for error handling
3086        let program = Program {
3087            includes: vec![],
3088            unions: vec![],
3089            words: vec![WordDef {
3090                name: "test".to_string(),
3091                effect: Some(Effect::new(
3092                    StackType::Empty,
3093                    StackType::from_vec(vec![Type::String, Type::Bool]),
3094                )),
3095                body: vec![Statement::WordCall {
3096                    name: "io.read-line".to_string(),
3097                    span: None,
3098                }],
3099                source: None,
3100                allowed_lints: vec![],
3101            }],
3102        };
3103
3104        let mut checker = TypeChecker::new();
3105        assert!(checker.check_program(&program).is_ok());
3106    }
3107
3108    #[test]
3109    fn test_comparison_operations() {
3110        // Test all comparison operators
3111        // : test ( Int Int -- Bool )
3112        //   i.<= ;
3113        // Simplified: just test that comparisons work and return Bool
3114        let program = Program {
3115            includes: vec![],
3116            unions: vec![],
3117            words: vec![WordDef {
3118                name: "test".to_string(),
3119                effect: Some(Effect::new(
3120                    StackType::Empty.push(Type::Int).push(Type::Int),
3121                    StackType::singleton(Type::Bool),
3122                )),
3123                body: vec![Statement::WordCall {
3124                    name: "i.<=".to_string(),
3125                    span: None,
3126                }],
3127                source: None,
3128                allowed_lints: vec![],
3129            }],
3130        };
3131
3132        let mut checker = TypeChecker::new();
3133        assert!(checker.check_program(&program).is_ok());
3134    }
3135
3136    #[test]
3137    fn test_recursive_word_definitions() {
3138        // Test mutually recursive words (type checking only, no runtime)
3139        // : is-even ( Int -- Int ) dup 0 = if drop 1 else 1 subtract is-odd then ;
3140        // : is-odd ( Int -- Int ) dup 0 = if drop 0 else 1 subtract is-even then ;
3141        //
3142        // Note: This tests that the checker can handle words that reference each other
3143        let program = Program {
3144            includes: vec![],
3145            unions: vec![],
3146            words: vec![
3147                WordDef {
3148                    name: "is-even".to_string(),
3149                    effect: Some(Effect::new(
3150                        StackType::singleton(Type::Int),
3151                        StackType::singleton(Type::Int),
3152                    )),
3153                    body: vec![
3154                        Statement::WordCall {
3155                            name: "dup".to_string(),
3156                            span: None,
3157                        },
3158                        Statement::IntLiteral(0),
3159                        Statement::WordCall {
3160                            name: "i.=".to_string(),
3161                            span: None,
3162                        },
3163                        Statement::If {
3164                            then_branch: vec![
3165                                Statement::WordCall {
3166                                    name: "drop".to_string(),
3167                                    span: None,
3168                                },
3169                                Statement::IntLiteral(1),
3170                            ],
3171                            else_branch: Some(vec![
3172                                Statement::IntLiteral(1),
3173                                Statement::WordCall {
3174                                    name: "i.subtract".to_string(),
3175                                    span: None,
3176                                },
3177                                Statement::WordCall {
3178                                    name: "is-odd".to_string(),
3179                                    span: None,
3180                                },
3181                            ]),
3182                            span: None,
3183                        },
3184                    ],
3185                    source: None,
3186                    allowed_lints: vec![],
3187                },
3188                WordDef {
3189                    name: "is-odd".to_string(),
3190                    effect: Some(Effect::new(
3191                        StackType::singleton(Type::Int),
3192                        StackType::singleton(Type::Int),
3193                    )),
3194                    body: vec![
3195                        Statement::WordCall {
3196                            name: "dup".to_string(),
3197                            span: None,
3198                        },
3199                        Statement::IntLiteral(0),
3200                        Statement::WordCall {
3201                            name: "i.=".to_string(),
3202                            span: None,
3203                        },
3204                        Statement::If {
3205                            then_branch: vec![
3206                                Statement::WordCall {
3207                                    name: "drop".to_string(),
3208                                    span: None,
3209                                },
3210                                Statement::IntLiteral(0),
3211                            ],
3212                            else_branch: Some(vec![
3213                                Statement::IntLiteral(1),
3214                                Statement::WordCall {
3215                                    name: "i.subtract".to_string(),
3216                                    span: None,
3217                                },
3218                                Statement::WordCall {
3219                                    name: "is-even".to_string(),
3220                                    span: None,
3221                                },
3222                            ]),
3223                            span: None,
3224                        },
3225                    ],
3226                    source: None,
3227                    allowed_lints: vec![],
3228                },
3229            ],
3230        };
3231
3232        let mut checker = TypeChecker::new();
3233        assert!(checker.check_program(&program).is_ok());
3234    }
3235
3236    #[test]
3237    fn test_word_calling_word_with_row_polymorphism() {
3238        // Test that row variables unify correctly through word calls
3239        // : apply-twice ( Int -- Int ) dup add ;
3240        // : quad ( Int -- Int ) apply-twice apply-twice ;
3241        // Should work: both use row polymorphism correctly
3242        let program = Program {
3243            includes: vec![],
3244            unions: vec![],
3245            words: vec![
3246                WordDef {
3247                    name: "apply-twice".to_string(),
3248                    effect: Some(Effect::new(
3249                        StackType::singleton(Type::Int),
3250                        StackType::singleton(Type::Int),
3251                    )),
3252                    body: vec![
3253                        Statement::WordCall {
3254                            name: "dup".to_string(),
3255                            span: None,
3256                        },
3257                        Statement::WordCall {
3258                            name: "i.add".to_string(),
3259                            span: None,
3260                        },
3261                    ],
3262                    source: None,
3263                    allowed_lints: vec![],
3264                },
3265                WordDef {
3266                    name: "quad".to_string(),
3267                    effect: Some(Effect::new(
3268                        StackType::singleton(Type::Int),
3269                        StackType::singleton(Type::Int),
3270                    )),
3271                    body: vec![
3272                        Statement::WordCall {
3273                            name: "apply-twice".to_string(),
3274                            span: None,
3275                        },
3276                        Statement::WordCall {
3277                            name: "apply-twice".to_string(),
3278                            span: None,
3279                        },
3280                    ],
3281                    source: None,
3282                    allowed_lints: vec![],
3283                },
3284            ],
3285        };
3286
3287        let mut checker = TypeChecker::new();
3288        assert!(checker.check_program(&program).is_ok());
3289    }
3290
3291    #[test]
3292    fn test_deep_stack_types() {
3293        // Test with many values on stack (10+ items)
3294        // : test ( Int Int Int Int Int Int Int Int Int Int -- Int )
3295        //   add add add add add add add add add ;
3296        let mut stack_type = StackType::Empty;
3297        for _ in 0..10 {
3298            stack_type = stack_type.push(Type::Int);
3299        }
3300
3301        let program = Program {
3302            includes: vec![],
3303            unions: vec![],
3304            words: vec![WordDef {
3305                name: "test".to_string(),
3306                effect: Some(Effect::new(stack_type, StackType::singleton(Type::Int))),
3307                body: vec![
3308                    Statement::WordCall {
3309                        name: "i.add".to_string(),
3310                        span: None,
3311                    },
3312                    Statement::WordCall {
3313                        name: "i.add".to_string(),
3314                        span: None,
3315                    },
3316                    Statement::WordCall {
3317                        name: "i.add".to_string(),
3318                        span: None,
3319                    },
3320                    Statement::WordCall {
3321                        name: "i.add".to_string(),
3322                        span: None,
3323                    },
3324                    Statement::WordCall {
3325                        name: "i.add".to_string(),
3326                        span: None,
3327                    },
3328                    Statement::WordCall {
3329                        name: "i.add".to_string(),
3330                        span: None,
3331                    },
3332                    Statement::WordCall {
3333                        name: "i.add".to_string(),
3334                        span: None,
3335                    },
3336                    Statement::WordCall {
3337                        name: "i.add".to_string(),
3338                        span: None,
3339                    },
3340                    Statement::WordCall {
3341                        name: "i.add".to_string(),
3342                        span: None,
3343                    },
3344                ],
3345                source: None,
3346                allowed_lints: vec![],
3347            }],
3348        };
3349
3350        let mut checker = TypeChecker::new();
3351        assert!(checker.check_program(&program).is_ok());
3352    }
3353
3354    #[test]
3355    fn test_simple_quotation() {
3356        // : test ( -- Quot )
3357        //   [ 1 add ] ;
3358        // Quotation type should be [ ..input Int -- ..input Int ] (row polymorphic)
3359        let program = Program {
3360            includes: vec![],
3361            unions: vec![],
3362            words: vec![WordDef {
3363                name: "test".to_string(),
3364                effect: Some(Effect::new(
3365                    StackType::Empty,
3366                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
3367                        StackType::RowVar("input".to_string()).push(Type::Int),
3368                        StackType::RowVar("input".to_string()).push(Type::Int),
3369                    )))),
3370                )),
3371                body: vec![Statement::Quotation {
3372                    span: None,
3373                    id: 0,
3374                    body: vec![
3375                        Statement::IntLiteral(1),
3376                        Statement::WordCall {
3377                            name: "i.add".to_string(),
3378                            span: None,
3379                        },
3380                    ],
3381                }],
3382                source: None,
3383                allowed_lints: vec![],
3384            }],
3385        };
3386
3387        let mut checker = TypeChecker::new();
3388        match checker.check_program(&program) {
3389            Ok(_) => {}
3390            Err(e) => panic!("Type check failed: {}", e),
3391        }
3392    }
3393
3394    #[test]
3395    fn test_empty_quotation() {
3396        // : test ( -- Quot )
3397        //   [ ] ;
3398        // Empty quotation has effect ( ..input -- ..input ) (preserves stack)
3399        let program = Program {
3400            includes: vec![],
3401            unions: vec![],
3402            words: vec![WordDef {
3403                name: "test".to_string(),
3404                effect: Some(Effect::new(
3405                    StackType::Empty,
3406                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
3407                        StackType::RowVar("input".to_string()),
3408                        StackType::RowVar("input".to_string()),
3409                    )))),
3410                )),
3411                body: vec![Statement::Quotation {
3412                    span: None,
3413                    id: 1,
3414                    body: vec![],
3415                }],
3416                source: None,
3417                allowed_lints: vec![],
3418            }],
3419        };
3420
3421        let mut checker = TypeChecker::new();
3422        assert!(checker.check_program(&program).is_ok());
3423    }
3424
3425    #[test]
3426    fn test_nested_quotation() {
3427        // : test ( -- Quot )
3428        //   [ [ 1 add ] ] ;
3429        // Outer quotation contains inner quotation (both row-polymorphic)
3430        let inner_quot_type = Type::Quotation(Box::new(Effect::new(
3431            StackType::RowVar("input".to_string()).push(Type::Int),
3432            StackType::RowVar("input".to_string()).push(Type::Int),
3433        )));
3434
3435        let outer_quot_type = Type::Quotation(Box::new(Effect::new(
3436            StackType::RowVar("input".to_string()),
3437            StackType::RowVar("input".to_string()).push(inner_quot_type.clone()),
3438        )));
3439
3440        let program = Program {
3441            includes: vec![],
3442            unions: vec![],
3443            words: vec![WordDef {
3444                name: "test".to_string(),
3445                effect: Some(Effect::new(
3446                    StackType::Empty,
3447                    StackType::singleton(outer_quot_type),
3448                )),
3449                body: vec![Statement::Quotation {
3450                    span: None,
3451                    id: 2,
3452                    body: vec![Statement::Quotation {
3453                        span: None,
3454                        id: 3,
3455                        body: vec![
3456                            Statement::IntLiteral(1),
3457                            Statement::WordCall {
3458                                name: "i.add".to_string(),
3459                                span: None,
3460                            },
3461                        ],
3462                    }],
3463                }],
3464                source: None,
3465                allowed_lints: vec![],
3466            }],
3467        };
3468
3469        let mut checker = TypeChecker::new();
3470        assert!(checker.check_program(&program).is_ok());
3471    }
3472
3473    #[test]
3474    fn test_invalid_field_type_error() {
3475        use crate::ast::{UnionDef, UnionField, UnionVariant};
3476
3477        let program = Program {
3478            includes: vec![],
3479            unions: vec![UnionDef {
3480                name: "Message".to_string(),
3481                variants: vec![UnionVariant {
3482                    name: "Get".to_string(),
3483                    fields: vec![UnionField {
3484                        name: "chan".to_string(),
3485                        type_name: "InvalidType".to_string(),
3486                    }],
3487                    source: None,
3488                }],
3489                source: None,
3490            }],
3491            words: vec![],
3492        };
3493
3494        let mut checker = TypeChecker::new();
3495        let result = checker.check_program(&program);
3496        assert!(result.is_err());
3497        let err = result.unwrap_err();
3498        assert!(err.contains("Unknown type 'InvalidType'"));
3499        assert!(err.contains("chan"));
3500        assert!(err.contains("Get"));
3501        assert!(err.contains("Message"));
3502    }
3503
3504    #[test]
3505    fn test_roll_inside_conditional_with_concrete_stack() {
3506        // Bug #93: n roll inside if/else should work when stack has enough concrete items
3507        // : test ( Int Int Int Int -- Int Int Int Int )
3508        //   dup 0 > if
3509        //     3 roll    # Works: 4 concrete items available
3510        //   else
3511        //     rot rot   # Alternative that also works
3512        //   then ;
3513        let program = Program {
3514            includes: vec![],
3515            unions: vec![],
3516            words: vec![WordDef {
3517                name: "test".to_string(),
3518                effect: Some(Effect::new(
3519                    StackType::Empty
3520                        .push(Type::Int)
3521                        .push(Type::Int)
3522                        .push(Type::Int)
3523                        .push(Type::Int),
3524                    StackType::Empty
3525                        .push(Type::Int)
3526                        .push(Type::Int)
3527                        .push(Type::Int)
3528                        .push(Type::Int),
3529                )),
3530                body: vec![
3531                    Statement::WordCall {
3532                        name: "dup".to_string(),
3533                        span: None,
3534                    },
3535                    Statement::IntLiteral(0),
3536                    Statement::WordCall {
3537                        name: "i.>".to_string(),
3538                        span: None,
3539                    },
3540                    Statement::If {
3541                        then_branch: vec![
3542                            Statement::IntLiteral(3),
3543                            Statement::WordCall {
3544                                name: "roll".to_string(),
3545                                span: None,
3546                            },
3547                        ],
3548                        else_branch: Some(vec![
3549                            Statement::WordCall {
3550                                name: "rot".to_string(),
3551                                span: None,
3552                            },
3553                            Statement::WordCall {
3554                                name: "rot".to_string(),
3555                                span: None,
3556                            },
3557                        ]),
3558                        span: None,
3559                    },
3560                ],
3561                source: None,
3562                allowed_lints: vec![],
3563            }],
3564        };
3565
3566        let mut checker = TypeChecker::new();
3567        // This should now work because both branches have 4 concrete items
3568        match checker.check_program(&program) {
3569            Ok(_) => {}
3570            Err(e) => panic!("Type check failed: {}", e),
3571        }
3572    }
3573
3574    #[test]
3575    fn test_roll_inside_match_arm_with_concrete_stack() {
3576        // Similar to bug #93 but for match arms: n roll inside match should work
3577        // when stack has enough concrete items from the match context
3578        use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};
3579
3580        // Define a simple union: union Result = Ok | Err
3581        let union_def = UnionDef {
3582            name: "Result".to_string(),
3583            variants: vec![
3584                UnionVariant {
3585                    name: "Ok".to_string(),
3586                    fields: vec![],
3587                    source: None,
3588                },
3589                UnionVariant {
3590                    name: "Err".to_string(),
3591                    fields: vec![],
3592                    source: None,
3593                },
3594            ],
3595            source: None,
3596        };
3597
3598        // : test ( Int Int Int Int Result -- Int Int Int Int )
3599        //   match
3600        //     Ok => 3 roll
3601        //     Err => rot rot
3602        //   end ;
3603        let program = Program {
3604            includes: vec![],
3605            unions: vec![union_def],
3606            words: vec![WordDef {
3607                name: "test".to_string(),
3608                effect: Some(Effect::new(
3609                    StackType::Empty
3610                        .push(Type::Int)
3611                        .push(Type::Int)
3612                        .push(Type::Int)
3613                        .push(Type::Int)
3614                        .push(Type::Union("Result".to_string())),
3615                    StackType::Empty
3616                        .push(Type::Int)
3617                        .push(Type::Int)
3618                        .push(Type::Int)
3619                        .push(Type::Int),
3620                )),
3621                body: vec![Statement::Match {
3622                    arms: vec![
3623                        MatchArm {
3624                            pattern: Pattern::Variant("Ok".to_string()),
3625                            body: vec![
3626                                Statement::IntLiteral(3),
3627                                Statement::WordCall {
3628                                    name: "roll".to_string(),
3629                                    span: None,
3630                                },
3631                            ],
3632                            span: None,
3633                        },
3634                        MatchArm {
3635                            pattern: Pattern::Variant("Err".to_string()),
3636                            body: vec![
3637                                Statement::WordCall {
3638                                    name: "rot".to_string(),
3639                                    span: None,
3640                                },
3641                                Statement::WordCall {
3642                                    name: "rot".to_string(),
3643                                    span: None,
3644                                },
3645                            ],
3646                            span: None,
3647                        },
3648                    ],
3649                    span: None,
3650                }],
3651                source: None,
3652                allowed_lints: vec![],
3653            }],
3654        };
3655
3656        let mut checker = TypeChecker::new();
3657        match checker.check_program(&program) {
3658            Ok(_) => {}
3659            Err(e) => panic!("Type check failed: {}", e),
3660        }
3661    }
3662
3663    #[test]
3664    fn test_roll_with_row_polymorphic_input() {
3665        // roll reaching into row variable should work (needed for stdlib)
3666        // : test ( T U V W -- U V W T )
3667        //   3 roll ;   # Rotates: brings position 3 to top
3668        let program = Program {
3669            includes: vec![],
3670            unions: vec![],
3671            words: vec![WordDef {
3672                name: "test".to_string(),
3673                effect: Some(Effect::new(
3674                    StackType::Empty
3675                        .push(Type::Var("T".to_string()))
3676                        .push(Type::Var("U".to_string()))
3677                        .push(Type::Var("V".to_string()))
3678                        .push(Type::Var("W".to_string())),
3679                    StackType::Empty
3680                        .push(Type::Var("U".to_string()))
3681                        .push(Type::Var("V".to_string()))
3682                        .push(Type::Var("W".to_string()))
3683                        .push(Type::Var("T".to_string())),
3684                )),
3685                body: vec![
3686                    Statement::IntLiteral(3),
3687                    Statement::WordCall {
3688                        name: "roll".to_string(),
3689                        span: None,
3690                    },
3691                ],
3692                source: None,
3693                allowed_lints: vec![],
3694            }],
3695        };
3696
3697        let mut checker = TypeChecker::new();
3698        let result = checker.check_program(&program);
3699        assert!(result.is_ok(), "roll test failed: {:?}", result.err());
3700    }
3701
3702    #[test]
3703    fn test_pick_with_row_polymorphic_input() {
3704        // pick reaching into row variable should work (needed for stdlib)
3705        // : test ( T U V -- T U V T )
3706        //   2 pick ;   # Copies element at index 2 (0-indexed from top)
3707        let program = Program {
3708            includes: vec![],
3709            unions: vec![],
3710            words: vec![WordDef {
3711                name: "test".to_string(),
3712                effect: Some(Effect::new(
3713                    StackType::Empty
3714                        .push(Type::Var("T".to_string()))
3715                        .push(Type::Var("U".to_string()))
3716                        .push(Type::Var("V".to_string())),
3717                    StackType::Empty
3718                        .push(Type::Var("T".to_string()))
3719                        .push(Type::Var("U".to_string()))
3720                        .push(Type::Var("V".to_string()))
3721                        .push(Type::Var("T".to_string())),
3722                )),
3723                body: vec![
3724                    Statement::IntLiteral(2),
3725                    Statement::WordCall {
3726                        name: "pick".to_string(),
3727                        span: None,
3728                    },
3729                ],
3730                source: None,
3731                allowed_lints: vec![],
3732            }],
3733        };
3734
3735        let mut checker = TypeChecker::new();
3736        assert!(checker.check_program(&program).is_ok());
3737    }
3738
3739    #[test]
3740    fn test_valid_union_reference_in_field() {
3741        use crate::ast::{UnionDef, UnionField, UnionVariant};
3742
3743        let program = Program {
3744            includes: vec![],
3745            unions: vec![
3746                UnionDef {
3747                    name: "Inner".to_string(),
3748                    variants: vec![UnionVariant {
3749                        name: "Val".to_string(),
3750                        fields: vec![UnionField {
3751                            name: "x".to_string(),
3752                            type_name: "Int".to_string(),
3753                        }],
3754                        source: None,
3755                    }],
3756                    source: None,
3757                },
3758                UnionDef {
3759                    name: "Outer".to_string(),
3760                    variants: vec![UnionVariant {
3761                        name: "Wrap".to_string(),
3762                        fields: vec![UnionField {
3763                            name: "inner".to_string(),
3764                            type_name: "Inner".to_string(), // Reference to other union
3765                        }],
3766                        source: None,
3767                    }],
3768                    source: None,
3769                },
3770            ],
3771            words: vec![],
3772        };
3773
3774        let mut checker = TypeChecker::new();
3775        assert!(
3776            checker.check_program(&program).is_ok(),
3777            "Union reference in field should be valid"
3778        );
3779    }
3780
3781    #[test]
3782    fn test_divergent_recursive_tail_call() {
3783        // Test that recursive tail calls in if/else branches are recognized as divergent.
3784        // This pattern is common in actor loops:
3785        //
3786        // : store-loop ( Channel -- )
3787        //   dup           # ( chan chan )
3788        //   chan.receive  # ( chan value Bool )
3789        //   not if        # ( chan value )
3790        //     drop        # ( chan ) - drop value, keep chan for recursion
3791        //     store-loop  # diverges - never returns
3792        //   then
3793        //   # else: ( chan value ) - process msg normally
3794        //   drop drop     # ( )
3795        // ;
3796        //
3797        // The then branch ends with a recursive call (store-loop), so it diverges.
3798        // The else branch (implicit empty) continues with the stack after the if.
3799        // Without divergent branch detection, this would fail because:
3800        //   - then branch produces: () (after drop store-loop)
3801        //   - else branch produces: (chan value)
3802        // But since then diverges, we should use else's type.
3803
3804        let program = Program {
3805            includes: vec![],
3806            unions: vec![],
3807            words: vec![WordDef {
3808                name: "store-loop".to_string(),
3809                effect: Some(Effect::new(
3810                    StackType::singleton(Type::Channel), // ( Channel -- )
3811                    StackType::Empty,
3812                )),
3813                body: vec![
3814                    // dup -> ( chan chan )
3815                    Statement::WordCall {
3816                        name: "dup".to_string(),
3817                        span: None,
3818                    },
3819                    // chan.receive -> ( chan value Bool )
3820                    Statement::WordCall {
3821                        name: "chan.receive".to_string(),
3822                        span: None,
3823                    },
3824                    // not -> ( chan value Bool )
3825                    Statement::WordCall {
3826                        name: "not".to_string(),
3827                        span: None,
3828                    },
3829                    // if drop store-loop then
3830                    Statement::If {
3831                        then_branch: vec![
3832                            // drop value -> ( chan )
3833                            Statement::WordCall {
3834                                name: "drop".to_string(),
3835                                span: None,
3836                            },
3837                            // store-loop -> diverges
3838                            Statement::WordCall {
3839                                name: "store-loop".to_string(), // recursive tail call
3840                                span: None,
3841                            },
3842                        ],
3843                        else_branch: None, // implicit else continues with ( chan value )
3844                        span: None,
3845                    },
3846                    // After if: ( chan value ) - drop both
3847                    Statement::WordCall {
3848                        name: "drop".to_string(),
3849                        span: None,
3850                    },
3851                    Statement::WordCall {
3852                        name: "drop".to_string(),
3853                        span: None,
3854                    },
3855                ],
3856                source: None,
3857                allowed_lints: vec![],
3858            }],
3859        };
3860
3861        let mut checker = TypeChecker::new();
3862        let result = checker.check_program(&program);
3863        assert!(
3864            result.is_ok(),
3865            "Divergent recursive tail call should be accepted: {:?}",
3866            result.err()
3867        );
3868    }
3869
3870    #[test]
3871    fn test_divergent_else_branch() {
3872        // Test that divergence detection works for else branches too.
3873        //
3874        // : process-loop ( Channel -- )
3875        //   dup chan.receive   # ( chan value Bool )
3876        //   if                 # ( chan value )
3877        //     drop drop        # normal exit: ( )
3878        //   else
3879        //     drop             # ( chan )
3880        //     process-loop     # diverges - retry on failure
3881        //   then
3882        // ;
3883
3884        let program = Program {
3885            includes: vec![],
3886            unions: vec![],
3887            words: vec![WordDef {
3888                name: "process-loop".to_string(),
3889                effect: Some(Effect::new(
3890                    StackType::singleton(Type::Channel), // ( Channel -- )
3891                    StackType::Empty,
3892                )),
3893                body: vec![
3894                    Statement::WordCall {
3895                        name: "dup".to_string(),
3896                        span: None,
3897                    },
3898                    Statement::WordCall {
3899                        name: "chan.receive".to_string(),
3900                        span: None,
3901                    },
3902                    Statement::If {
3903                        then_branch: vec![
3904                            // success: drop value and chan
3905                            Statement::WordCall {
3906                                name: "drop".to_string(),
3907                                span: None,
3908                            },
3909                            Statement::WordCall {
3910                                name: "drop".to_string(),
3911                                span: None,
3912                            },
3913                        ],
3914                        else_branch: Some(vec![
3915                            // failure: drop value, keep chan, recurse
3916                            Statement::WordCall {
3917                                name: "drop".to_string(),
3918                                span: None,
3919                            },
3920                            Statement::WordCall {
3921                                name: "process-loop".to_string(), // recursive tail call
3922                                span: None,
3923                            },
3924                        ]),
3925                        span: None,
3926                    },
3927                ],
3928                source: None,
3929                allowed_lints: vec![],
3930            }],
3931        };
3932
3933        let mut checker = TypeChecker::new();
3934        let result = checker.check_program(&program);
3935        assert!(
3936            result.is_ok(),
3937            "Divergent else branch should be accepted: {:?}",
3938            result.err()
3939        );
3940    }
3941
3942    #[test]
3943    fn test_non_tail_call_recursion_not_divergent() {
3944        // Test that recursion NOT in tail position is not treated as divergent.
3945        // This should fail type checking because after the recursive call,
3946        // there's more code that changes the stack.
3947        //
3948        // : bad-loop ( Int -- Int )
3949        //   dup 0 i.> if
3950        //     1 i.subtract bad-loop  # recursive call
3951        //     1 i.add                # more code after - not tail position!
3952        //   then
3953        // ;
3954        //
3955        // This should fail because:
3956        // - then branch: recurse then add 1 -> stack changes after recursion
3957        // - else branch (implicit): stack is ( Int )
3958        // Without proper handling, this could incorrectly pass.
3959
3960        let program = Program {
3961            includes: vec![],
3962            unions: vec![],
3963            words: vec![WordDef {
3964                name: "bad-loop".to_string(),
3965                effect: Some(Effect::new(
3966                    StackType::singleton(Type::Int),
3967                    StackType::singleton(Type::Int),
3968                )),
3969                body: vec![
3970                    Statement::WordCall {
3971                        name: "dup".to_string(),
3972                        span: None,
3973                    },
3974                    Statement::IntLiteral(0),
3975                    Statement::WordCall {
3976                        name: "i.>".to_string(),
3977                        span: None,
3978                    },
3979                    Statement::If {
3980                        then_branch: vec![
3981                            Statement::IntLiteral(1),
3982                            Statement::WordCall {
3983                                name: "i.subtract".to_string(),
3984                                span: None,
3985                            },
3986                            Statement::WordCall {
3987                                name: "bad-loop".to_string(), // NOT in tail position
3988                                span: None,
3989                            },
3990                            Statement::IntLiteral(1),
3991                            Statement::WordCall {
3992                                name: "i.add".to_string(), // code after recursion
3993                                span: None,
3994                            },
3995                        ],
3996                        else_branch: None,
3997                        span: None,
3998                    },
3999                ],
4000                source: None,
4001                allowed_lints: vec![],
4002            }],
4003        };
4004
4005        let mut checker = TypeChecker::new();
4006        // This should pass because the branches ARE compatible:
4007        // - then: produces Int (after bad-loop returns Int, then add 1)
4008        // - else: produces Int (from the dup at start)
4009        // The key is that bad-loop is NOT in tail position, so it's not divergent.
4010        let result = checker.check_program(&program);
4011        assert!(
4012            result.is_ok(),
4013            "Non-tail recursion should type check normally: {:?}",
4014            result.err()
4015        );
4016    }
4017
4018    #[test]
4019    fn test_call_yield_quotation_error() {
4020        // Phase 2c: Calling a quotation with Yield effect directly should error.
4021        // : bad ( Ctx -- Ctx ) [ yield ] call ;
4022        // This is wrong because yield quotations must be wrapped with strand.weave.
4023        let program = Program {
4024            includes: vec![],
4025            unions: vec![],
4026            words: vec![WordDef {
4027                name: "bad".to_string(),
4028                effect: Some(Effect::new(
4029                    StackType::singleton(Type::Var("Ctx".to_string())),
4030                    StackType::singleton(Type::Var("Ctx".to_string())),
4031                )),
4032                body: vec![
4033                    // Push a dummy value that will be yielded
4034                    Statement::IntLiteral(42),
4035                    Statement::Quotation {
4036                        span: None,
4037                        id: 0,
4038                        body: vec![Statement::WordCall {
4039                            name: "yield".to_string(),
4040                            span: None,
4041                        }],
4042                    },
4043                    Statement::WordCall {
4044                        name: "call".to_string(),
4045                        span: None,
4046                    },
4047                ],
4048                source: None,
4049                allowed_lints: vec![],
4050            }],
4051        };
4052
4053        let mut checker = TypeChecker::new();
4054        let result = checker.check_program(&program);
4055        assert!(
4056            result.is_err(),
4057            "Calling yield quotation directly should fail"
4058        );
4059        let err = result.unwrap_err();
4060        assert!(
4061            err.contains("Yield") || err.contains("strand.weave"),
4062            "Error should mention Yield or strand.weave: {}",
4063            err
4064        );
4065    }
4066
4067    #[test]
4068    fn test_strand_weave_yield_quotation_ok() {
4069        // Phase 2c: Using strand.weave on a Yield quotation is correct.
4070        // : good ( -- Int Handle ) 42 [ yield ] strand.weave ;
4071        let program = Program {
4072            includes: vec![],
4073            unions: vec![],
4074            words: vec![WordDef {
4075                name: "good".to_string(),
4076                effect: Some(Effect::new(
4077                    StackType::Empty,
4078                    StackType::Empty
4079                        .push(Type::Int)
4080                        .push(Type::Var("Handle".to_string())),
4081                )),
4082                body: vec![
4083                    Statement::IntLiteral(42),
4084                    Statement::Quotation {
4085                        span: None,
4086                        id: 0,
4087                        body: vec![Statement::WordCall {
4088                            name: "yield".to_string(),
4089                            span: None,
4090                        }],
4091                    },
4092                    Statement::WordCall {
4093                        name: "strand.weave".to_string(),
4094                        span: None,
4095                    },
4096                ],
4097                source: None,
4098                allowed_lints: vec![],
4099            }],
4100        };
4101
4102        let mut checker = TypeChecker::new();
4103        let result = checker.check_program(&program);
4104        assert!(
4105            result.is_ok(),
4106            "strand.weave on yield quotation should pass: {:?}",
4107            result.err()
4108        );
4109    }
4110
4111    #[test]
4112    fn test_call_pure_quotation_ok() {
4113        // Phase 2c: Calling a pure quotation (no Yield) is fine.
4114        // : ok ( Int -- Int ) [ 1 i.add ] call ;
4115        let program = Program {
4116            includes: vec![],
4117            unions: vec![],
4118            words: vec![WordDef {
4119                name: "ok".to_string(),
4120                effect: Some(Effect::new(
4121                    StackType::singleton(Type::Int),
4122                    StackType::singleton(Type::Int),
4123                )),
4124                body: vec![
4125                    Statement::Quotation {
4126                        span: None,
4127                        id: 0,
4128                        body: vec![
4129                            Statement::IntLiteral(1),
4130                            Statement::WordCall {
4131                                name: "i.add".to_string(),
4132                                span: None,
4133                            },
4134                        ],
4135                    },
4136                    Statement::WordCall {
4137                        name: "call".to_string(),
4138                        span: None,
4139                    },
4140                ],
4141                source: None,
4142                allowed_lints: vec![],
4143            }],
4144        };
4145
4146        let mut checker = TypeChecker::new();
4147        let result = checker.check_program(&program);
4148        assert!(
4149            result.is_ok(),
4150            "Calling pure quotation should pass: {:?}",
4151            result.err()
4152        );
4153    }
4154
4155    // ==========================================================================
4156    // Stack Pollution Detection Tests (Issue #228)
4157    // These tests verify the type checker catches stack effect mismatches
4158    // ==========================================================================
4159
4160    #[test]
4161    fn test_pollution_extra_push() {
4162        // : test ( Int -- Int ) 42 ;
4163        // Declares consuming 1 Int, producing 1 Int
4164        // But body pushes 42 on top of input, leaving 2 values
4165        let program = Program {
4166            includes: vec![],
4167            unions: vec![],
4168            words: vec![WordDef {
4169                name: "test".to_string(),
4170                effect: Some(Effect::new(
4171                    StackType::singleton(Type::Int),
4172                    StackType::singleton(Type::Int),
4173                )),
4174                body: vec![Statement::IntLiteral(42)],
4175                source: None,
4176                allowed_lints: vec![],
4177            }],
4178        };
4179
4180        let mut checker = TypeChecker::new();
4181        let result = checker.check_program(&program);
4182        assert!(
4183            result.is_err(),
4184            "Should reject: declares ( Int -- Int ) but leaves 2 values on stack"
4185        );
4186    }
4187
4188    #[test]
4189    fn test_pollution_extra_dup() {
4190        // : test ( Int -- Int ) dup ;
4191        // Declares producing 1 Int, but dup produces 2
4192        let program = Program {
4193            includes: vec![],
4194            unions: vec![],
4195            words: vec![WordDef {
4196                name: "test".to_string(),
4197                effect: Some(Effect::new(
4198                    StackType::singleton(Type::Int),
4199                    StackType::singleton(Type::Int),
4200                )),
4201                body: vec![Statement::WordCall {
4202                    name: "dup".to_string(),
4203                    span: None,
4204                }],
4205                source: None,
4206                allowed_lints: vec![],
4207            }],
4208        };
4209
4210        let mut checker = TypeChecker::new();
4211        let result = checker.check_program(&program);
4212        assert!(
4213            result.is_err(),
4214            "Should reject: declares ( Int -- Int ) but dup produces 2 values"
4215        );
4216    }
4217
4218    #[test]
4219    fn test_pollution_consumes_extra() {
4220        // : test ( Int -- Int ) drop drop 42 ;
4221        // Declares consuming 1 Int, but body drops twice
4222        let program = Program {
4223            includes: vec![],
4224            unions: vec![],
4225            words: vec![WordDef {
4226                name: "test".to_string(),
4227                effect: Some(Effect::new(
4228                    StackType::singleton(Type::Int),
4229                    StackType::singleton(Type::Int),
4230                )),
4231                body: vec![
4232                    Statement::WordCall {
4233                        name: "drop".to_string(),
4234                        span: None,
4235                    },
4236                    Statement::WordCall {
4237                        name: "drop".to_string(),
4238                        span: None,
4239                    },
4240                    Statement::IntLiteral(42),
4241                ],
4242                source: None,
4243                allowed_lints: vec![],
4244            }],
4245        };
4246
4247        let mut checker = TypeChecker::new();
4248        let result = checker.check_program(&program);
4249        assert!(
4250            result.is_err(),
4251            "Should reject: declares ( Int -- Int ) but consumes 2 values"
4252        );
4253    }
4254
4255    #[test]
4256    fn test_pollution_in_then_branch() {
4257        // : test ( Bool -- Int )
4258        //   if 1 2 else 3 then ;
4259        // Then branch pushes 2 values, else pushes 1
4260        let program = Program {
4261            includes: vec![],
4262            unions: vec![],
4263            words: vec![WordDef {
4264                name: "test".to_string(),
4265                effect: Some(Effect::new(
4266                    StackType::singleton(Type::Bool),
4267                    StackType::singleton(Type::Int),
4268                )),
4269                body: vec![Statement::If {
4270                    then_branch: vec![
4271                        Statement::IntLiteral(1),
4272                        Statement::IntLiteral(2), // Extra value!
4273                    ],
4274                    else_branch: Some(vec![Statement::IntLiteral(3)]),
4275                    span: None,
4276                }],
4277                source: None,
4278                allowed_lints: vec![],
4279            }],
4280        };
4281
4282        let mut checker = TypeChecker::new();
4283        let result = checker.check_program(&program);
4284        assert!(
4285            result.is_err(),
4286            "Should reject: then branch pushes 2 values, else pushes 1"
4287        );
4288    }
4289
4290    #[test]
4291    fn test_pollution_in_else_branch() {
4292        // : test ( Bool -- Int )
4293        //   if 1 else 2 3 then ;
4294        // Then branch pushes 1, else pushes 2 values
4295        let program = Program {
4296            includes: vec![],
4297            unions: vec![],
4298            words: vec![WordDef {
4299                name: "test".to_string(),
4300                effect: Some(Effect::new(
4301                    StackType::singleton(Type::Bool),
4302                    StackType::singleton(Type::Int),
4303                )),
4304                body: vec![Statement::If {
4305                    then_branch: vec![Statement::IntLiteral(1)],
4306                    else_branch: Some(vec![
4307                        Statement::IntLiteral(2),
4308                        Statement::IntLiteral(3), // Extra value!
4309                    ]),
4310                    span: None,
4311                }],
4312                source: None,
4313                allowed_lints: vec![],
4314            }],
4315        };
4316
4317        let mut checker = TypeChecker::new();
4318        let result = checker.check_program(&program);
4319        assert!(
4320            result.is_err(),
4321            "Should reject: then branch pushes 1 value, else pushes 2"
4322        );
4323    }
4324
4325    #[test]
4326    fn test_pollution_both_branches_extra() {
4327        // : test ( Bool -- Int )
4328        //   if 1 2 else 3 4 then ;
4329        // Both branches push 2 values but declared output is 1
4330        let program = Program {
4331            includes: vec![],
4332            unions: vec![],
4333            words: vec![WordDef {
4334                name: "test".to_string(),
4335                effect: Some(Effect::new(
4336                    StackType::singleton(Type::Bool),
4337                    StackType::singleton(Type::Int),
4338                )),
4339                body: vec![Statement::If {
4340                    then_branch: vec![Statement::IntLiteral(1), Statement::IntLiteral(2)],
4341                    else_branch: Some(vec![Statement::IntLiteral(3), Statement::IntLiteral(4)]),
4342                    span: None,
4343                }],
4344                source: None,
4345                allowed_lints: vec![],
4346            }],
4347        };
4348
4349        let mut checker = TypeChecker::new();
4350        let result = checker.check_program(&program);
4351        assert!(
4352            result.is_err(),
4353            "Should reject: both branches push 2 values, but declared output is 1"
4354        );
4355    }
4356
4357    #[test]
4358    fn test_pollution_branch_consumes_extra() {
4359        // : test ( Bool Int -- Int )
4360        //   if drop drop 1 else then ;
4361        // Then branch consumes more than available from declared inputs
4362        let program = Program {
4363            includes: vec![],
4364            unions: vec![],
4365            words: vec![WordDef {
4366                name: "test".to_string(),
4367                effect: Some(Effect::new(
4368                    StackType::Empty.push(Type::Bool).push(Type::Int),
4369                    StackType::singleton(Type::Int),
4370                )),
4371                body: vec![Statement::If {
4372                    then_branch: vec![
4373                        Statement::WordCall {
4374                            name: "drop".to_string(),
4375                            span: None,
4376                        },
4377                        Statement::WordCall {
4378                            name: "drop".to_string(),
4379                            span: None,
4380                        },
4381                        Statement::IntLiteral(1),
4382                    ],
4383                    else_branch: Some(vec![]),
4384                    span: None,
4385                }],
4386                source: None,
4387                allowed_lints: vec![],
4388            }],
4389        };
4390
4391        let mut checker = TypeChecker::new();
4392        let result = checker.check_program(&program);
4393        assert!(
4394            result.is_err(),
4395            "Should reject: then branch consumes Bool (should only have Int after if)"
4396        );
4397    }
4398
4399    #[test]
4400    fn test_pollution_quotation_wrong_arity_output() {
4401        // : test ( Int -- Int )
4402        //   [ dup ] call ;
4403        // Quotation produces 2 values, but word declares 1 output
4404        let program = Program {
4405            includes: vec![],
4406            unions: vec![],
4407            words: vec![WordDef {
4408                name: "test".to_string(),
4409                effect: Some(Effect::new(
4410                    StackType::singleton(Type::Int),
4411                    StackType::singleton(Type::Int),
4412                )),
4413                body: vec![
4414                    Statement::Quotation {
4415                        span: None,
4416                        id: 0,
4417                        body: vec![Statement::WordCall {
4418                            name: "dup".to_string(),
4419                            span: None,
4420                        }],
4421                    },
4422                    Statement::WordCall {
4423                        name: "call".to_string(),
4424                        span: None,
4425                    },
4426                ],
4427                source: None,
4428                allowed_lints: vec![],
4429            }],
4430        };
4431
4432        let mut checker = TypeChecker::new();
4433        let result = checker.check_program(&program);
4434        assert!(
4435            result.is_err(),
4436            "Should reject: quotation [dup] produces 2 values, declared output is 1"
4437        );
4438    }
4439
4440    #[test]
4441    fn test_pollution_quotation_wrong_arity_input() {
4442        // : test ( Int -- Int )
4443        //   [ drop drop 42 ] call ;
4444        // Quotation consumes 2 values, but only 1 available
4445        let program = Program {
4446            includes: vec![],
4447            unions: vec![],
4448            words: vec![WordDef {
4449                name: "test".to_string(),
4450                effect: Some(Effect::new(
4451                    StackType::singleton(Type::Int),
4452                    StackType::singleton(Type::Int),
4453                )),
4454                body: vec![
4455                    Statement::Quotation {
4456                        span: None,
4457                        id: 0,
4458                        body: vec![
4459                            Statement::WordCall {
4460                                name: "drop".to_string(),
4461                                span: None,
4462                            },
4463                            Statement::WordCall {
4464                                name: "drop".to_string(),
4465                                span: None,
4466                            },
4467                            Statement::IntLiteral(42),
4468                        ],
4469                    },
4470                    Statement::WordCall {
4471                        name: "call".to_string(),
4472                        span: None,
4473                    },
4474                ],
4475                source: None,
4476                allowed_lints: vec![],
4477            }],
4478        };
4479
4480        let mut checker = TypeChecker::new();
4481        let result = checker.check_program(&program);
4482        assert!(
4483            result.is_err(),
4484            "Should reject: quotation [drop drop 42] consumes 2 values, only 1 available"
4485        );
4486    }
4487
4488    #[test]
4489    fn test_missing_effect_provides_helpful_error() {
4490        // : myword 42 ;
4491        // No effect annotation - should error with helpful message including word name
4492        let program = Program {
4493            includes: vec![],
4494            unions: vec![],
4495            words: vec![WordDef {
4496                name: "myword".to_string(),
4497                effect: None, // No annotation
4498                body: vec![Statement::IntLiteral(42)],
4499                source: None,
4500                allowed_lints: vec![],
4501            }],
4502        };
4503
4504        let mut checker = TypeChecker::new();
4505        let result = checker.check_program(&program);
4506        assert!(result.is_err());
4507        let err = result.unwrap_err();
4508        assert!(err.contains("myword"), "Error should mention word name");
4509        assert!(
4510            err.contains("stack effect"),
4511            "Error should mention stack effect"
4512        );
4513    }
4514
4515    #[test]
4516    fn test_valid_effect_exact_match() {
4517        // : test ( Int Int -- Int ) i.+ ;
4518        // Exact match - consumes 2, produces 1
4519        let program = Program {
4520            includes: vec![],
4521            unions: vec![],
4522            words: vec![WordDef {
4523                name: "test".to_string(),
4524                effect: Some(Effect::new(
4525                    StackType::Empty.push(Type::Int).push(Type::Int),
4526                    StackType::singleton(Type::Int),
4527                )),
4528                body: vec![Statement::WordCall {
4529                    name: "i.add".to_string(),
4530                    span: None,
4531                }],
4532                source: None,
4533                allowed_lints: vec![],
4534            }],
4535        };
4536
4537        let mut checker = TypeChecker::new();
4538        let result = checker.check_program(&program);
4539        assert!(result.is_ok(), "Should accept: effect matches exactly");
4540    }
4541
4542    #[test]
4543    fn test_valid_polymorphic_passthrough() {
4544        // : test ( a -- a ) ;
4545        // Identity function - row polymorphism allows this
4546        let program = Program {
4547            includes: vec![],
4548            unions: vec![],
4549            words: vec![WordDef {
4550                name: "test".to_string(),
4551                effect: Some(Effect::new(
4552                    StackType::Cons {
4553                        rest: Box::new(StackType::RowVar("rest".to_string())),
4554                        top: Type::Var("a".to_string()),
4555                    },
4556                    StackType::Cons {
4557                        rest: Box::new(StackType::RowVar("rest".to_string())),
4558                        top: Type::Var("a".to_string()),
4559                    },
4560                )),
4561                body: vec![], // Empty body - just pass through
4562                source: None,
4563                allowed_lints: vec![],
4564            }],
4565        };
4566
4567        let mut checker = TypeChecker::new();
4568        let result = checker.check_program(&program);
4569        assert!(result.is_ok(), "Should accept: polymorphic identity");
4570    }
4571
4572    // ==========================================================================
4573    // Closure Nesting Tests (Issue #230)
4574    // Tests for deep closure nesting, transitive captures, and edge cases
4575    // ==========================================================================
4576
4577    #[test]
4578    fn test_closure_basic_capture() {
4579        // : make-adder ( Int -- Closure )
4580        //   [ i.+ ] ;
4581        // The quotation needs 2 Ints (for i.+) but caller will only provide 1
4582        // So it captures 1 Int from the creation site
4583        // Must declare as Closure type to trigger capture analysis
4584        let program = Program {
4585            includes: vec![],
4586            unions: vec![],
4587            words: vec![WordDef {
4588                name: "make-adder".to_string(),
4589                effect: Some(Effect::new(
4590                    StackType::singleton(Type::Int),
4591                    StackType::singleton(Type::Closure {
4592                        effect: Box::new(Effect::new(
4593                            StackType::RowVar("r".to_string()).push(Type::Int),
4594                            StackType::RowVar("r".to_string()).push(Type::Int),
4595                        )),
4596                        captures: vec![Type::Int], // Captures 1 Int
4597                    }),
4598                )),
4599                body: vec![Statement::Quotation {
4600                    span: None,
4601                    id: 0,
4602                    body: vec![Statement::WordCall {
4603                        name: "i.add".to_string(),
4604                        span: None,
4605                    }],
4606                }],
4607                source: None,
4608                allowed_lints: vec![],
4609            }],
4610        };
4611
4612        let mut checker = TypeChecker::new();
4613        let result = checker.check_program(&program);
4614        assert!(
4615            result.is_ok(),
4616            "Basic closure capture should work: {:?}",
4617            result.err()
4618        );
4619    }
4620
4621    #[test]
4622    fn test_closure_nested_two_levels() {
4623        // : outer ( -- Quot )
4624        //   [ [ 1 i.+ ] ] ;
4625        // Outer quotation: no inputs, just returns inner quotation
4626        // Inner quotation: pushes 1 then adds (needs 1 Int from caller)
4627        let program = Program {
4628            includes: vec![],
4629            unions: vec![],
4630            words: vec![WordDef {
4631                name: "outer".to_string(),
4632                effect: Some(Effect::new(
4633                    StackType::Empty,
4634                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
4635                        StackType::RowVar("r".to_string()),
4636                        StackType::RowVar("r".to_string()).push(Type::Quotation(Box::new(
4637                            Effect::new(
4638                                StackType::RowVar("s".to_string()).push(Type::Int),
4639                                StackType::RowVar("s".to_string()).push(Type::Int),
4640                            ),
4641                        ))),
4642                    )))),
4643                )),
4644                body: vec![Statement::Quotation {
4645                    span: None,
4646                    id: 0,
4647                    body: vec![Statement::Quotation {
4648                        span: None,
4649                        id: 1,
4650                        body: vec![
4651                            Statement::IntLiteral(1),
4652                            Statement::WordCall {
4653                                name: "i.add".to_string(),
4654                                span: None,
4655                            },
4656                        ],
4657                    }],
4658                }],
4659                source: None,
4660                allowed_lints: vec![],
4661            }],
4662        };
4663
4664        let mut checker = TypeChecker::new();
4665        let result = checker.check_program(&program);
4666        assert!(
4667            result.is_ok(),
4668            "Two-level nested quotations should work: {:?}",
4669            result.err()
4670        );
4671    }
4672
4673    #[test]
4674    fn test_closure_nested_three_levels() {
4675        // : deep ( -- Quot )
4676        //   [ [ [ 1 i.+ ] ] ] ;
4677        // Three levels of nesting, innermost does actual work
4678        let inner_effect = Effect::new(
4679            StackType::RowVar("a".to_string()).push(Type::Int),
4680            StackType::RowVar("a".to_string()).push(Type::Int),
4681        );
4682        let middle_effect = Effect::new(
4683            StackType::RowVar("b".to_string()),
4684            StackType::RowVar("b".to_string()).push(Type::Quotation(Box::new(inner_effect))),
4685        );
4686        let outer_effect = Effect::new(
4687            StackType::RowVar("c".to_string()),
4688            StackType::RowVar("c".to_string()).push(Type::Quotation(Box::new(middle_effect))),
4689        );
4690
4691        let program = Program {
4692            includes: vec![],
4693            unions: vec![],
4694            words: vec![WordDef {
4695                name: "deep".to_string(),
4696                effect: Some(Effect::new(
4697                    StackType::Empty,
4698                    StackType::singleton(Type::Quotation(Box::new(outer_effect))),
4699                )),
4700                body: vec![Statement::Quotation {
4701                    span: None,
4702                    id: 0,
4703                    body: vec![Statement::Quotation {
4704                        span: None,
4705                        id: 1,
4706                        body: vec![Statement::Quotation {
4707                            span: None,
4708                            id: 2,
4709                            body: vec![
4710                                Statement::IntLiteral(1),
4711                                Statement::WordCall {
4712                                    name: "i.add".to_string(),
4713                                    span: None,
4714                                },
4715                            ],
4716                        }],
4717                    }],
4718                }],
4719                source: None,
4720                allowed_lints: vec![],
4721            }],
4722        };
4723
4724        let mut checker = TypeChecker::new();
4725        let result = checker.check_program(&program);
4726        assert!(
4727            result.is_ok(),
4728            "Three-level nested quotations should work: {:?}",
4729            result.err()
4730        );
4731    }
4732
4733    #[test]
4734    fn test_closure_use_after_creation() {
4735        // : use-adder ( -- Int )
4736        //   5 make-adder   // Creates closure capturing 5
4737        //   10 swap call ; // Calls closure with 10, should return 15
4738        //
4739        // Tests that closure is properly typed when called later
4740        let adder_type = Type::Closure {
4741            effect: Box::new(Effect::new(
4742                StackType::RowVar("r".to_string()).push(Type::Int),
4743                StackType::RowVar("r".to_string()).push(Type::Int),
4744            )),
4745            captures: vec![Type::Int],
4746        };
4747
4748        let program = Program {
4749            includes: vec![],
4750            unions: vec![],
4751            words: vec![
4752                WordDef {
4753                    name: "make-adder".to_string(),
4754                    effect: Some(Effect::new(
4755                        StackType::singleton(Type::Int),
4756                        StackType::singleton(adder_type.clone()),
4757                    )),
4758                    body: vec![Statement::Quotation {
4759                        span: None,
4760                        id: 0,
4761                        body: vec![Statement::WordCall {
4762                            name: "i.add".to_string(),
4763                            span: None,
4764                        }],
4765                    }],
4766                    source: None,
4767                    allowed_lints: vec![],
4768                },
4769                WordDef {
4770                    name: "use-adder".to_string(),
4771                    effect: Some(Effect::new(
4772                        StackType::Empty,
4773                        StackType::singleton(Type::Int),
4774                    )),
4775                    body: vec![
4776                        Statement::IntLiteral(5),
4777                        Statement::WordCall {
4778                            name: "make-adder".to_string(),
4779                            span: None,
4780                        },
4781                        Statement::IntLiteral(10),
4782                        Statement::WordCall {
4783                            name: "swap".to_string(),
4784                            span: None,
4785                        },
4786                        Statement::WordCall {
4787                            name: "call".to_string(),
4788                            span: None,
4789                        },
4790                    ],
4791                    source: None,
4792                    allowed_lints: vec![],
4793                },
4794            ],
4795        };
4796
4797        let mut checker = TypeChecker::new();
4798        let result = checker.check_program(&program);
4799        assert!(
4800            result.is_ok(),
4801            "Closure usage after creation should work: {:?}",
4802            result.err()
4803        );
4804    }
4805
4806    #[test]
4807    fn test_closure_wrong_call_type() {
4808        // : bad-use ( -- Int )
4809        //   5 make-adder   // Creates Int -> Int closure
4810        //   "hello" swap call ; // Tries to call with String - should fail!
4811        let adder_type = Type::Closure {
4812            effect: Box::new(Effect::new(
4813                StackType::RowVar("r".to_string()).push(Type::Int),
4814                StackType::RowVar("r".to_string()).push(Type::Int),
4815            )),
4816            captures: vec![Type::Int],
4817        };
4818
4819        let program = Program {
4820            includes: vec![],
4821            unions: vec![],
4822            words: vec![
4823                WordDef {
4824                    name: "make-adder".to_string(),
4825                    effect: Some(Effect::new(
4826                        StackType::singleton(Type::Int),
4827                        StackType::singleton(adder_type.clone()),
4828                    )),
4829                    body: vec![Statement::Quotation {
4830                        span: None,
4831                        id: 0,
4832                        body: vec![Statement::WordCall {
4833                            name: "i.add".to_string(),
4834                            span: None,
4835                        }],
4836                    }],
4837                    source: None,
4838                    allowed_lints: vec![],
4839                },
4840                WordDef {
4841                    name: "bad-use".to_string(),
4842                    effect: Some(Effect::new(
4843                        StackType::Empty,
4844                        StackType::singleton(Type::Int),
4845                    )),
4846                    body: vec![
4847                        Statement::IntLiteral(5),
4848                        Statement::WordCall {
4849                            name: "make-adder".to_string(),
4850                            span: None,
4851                        },
4852                        Statement::StringLiteral("hello".to_string()), // Wrong type!
4853                        Statement::WordCall {
4854                            name: "swap".to_string(),
4855                            span: None,
4856                        },
4857                        Statement::WordCall {
4858                            name: "call".to_string(),
4859                            span: None,
4860                        },
4861                    ],
4862                    source: None,
4863                    allowed_lints: vec![],
4864                },
4865            ],
4866        };
4867
4868        let mut checker = TypeChecker::new();
4869        let result = checker.check_program(&program);
4870        assert!(
4871            result.is_err(),
4872            "Calling Int closure with String should fail"
4873        );
4874    }
4875
4876    #[test]
4877    fn test_closure_multiple_captures() {
4878        // : make-between ( Int Int -- Quot )
4879        //   [ dup rot i.>= swap rot i.<= and ] ;
4880        // Captures both min and max, checks if value is between them
4881        // Body needs: value min max (3 Ints)
4882        // Caller provides: value (1 Int)
4883        // Captures: min max (2 Ints)
4884        let program = Program {
4885            includes: vec![],
4886            unions: vec![],
4887            words: vec![WordDef {
4888                name: "make-between".to_string(),
4889                effect: Some(Effect::new(
4890                    StackType::Empty.push(Type::Int).push(Type::Int),
4891                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
4892                        StackType::RowVar("r".to_string()).push(Type::Int),
4893                        StackType::RowVar("r".to_string()).push(Type::Bool),
4894                    )))),
4895                )),
4896                body: vec![Statement::Quotation {
4897                    span: None,
4898                    id: 0,
4899                    body: vec![
4900                        // Simplified: just do a comparison that uses all 3 values
4901                        Statement::WordCall {
4902                            name: "i.>=".to_string(),
4903                            span: None,
4904                        },
4905                        // Note: This doesn't match the comment but tests multi-capture
4906                    ],
4907                }],
4908                source: None,
4909                allowed_lints: vec![],
4910            }],
4911        };
4912
4913        let mut checker = TypeChecker::new();
4914        let result = checker.check_program(&program);
4915        // This should work - the quotation body uses values from stack
4916        // The exact behavior depends on how captures are inferred
4917        // For now, we're testing that it doesn't crash
4918        assert!(
4919            result.is_ok() || result.is_err(),
4920            "Multiple captures should be handled (pass or fail gracefully)"
4921        );
4922    }
4923
4924    #[test]
4925    fn test_quotation_type_preserved_through_word() {
4926        // : identity-quot ( Quot -- Quot ) ;
4927        // Tests that quotation types are preserved when passed through words
4928        let quot_type = Type::Quotation(Box::new(Effect::new(
4929            StackType::RowVar("r".to_string()).push(Type::Int),
4930            StackType::RowVar("r".to_string()).push(Type::Int),
4931        )));
4932
4933        let program = Program {
4934            includes: vec![],
4935            unions: vec![],
4936            words: vec![WordDef {
4937                name: "identity-quot".to_string(),
4938                effect: Some(Effect::new(
4939                    StackType::singleton(quot_type.clone()),
4940                    StackType::singleton(quot_type.clone()),
4941                )),
4942                body: vec![], // Identity - just return what's on stack
4943                source: None,
4944                allowed_lints: vec![],
4945            }],
4946        };
4947
4948        let mut checker = TypeChecker::new();
4949        let result = checker.check_program(&program);
4950        assert!(
4951            result.is_ok(),
4952            "Quotation type should be preserved through identity word: {:?}",
4953            result.err()
4954        );
4955    }
4956
4957    #[test]
4958    fn test_closure_captures_value_for_inner_quotation() {
4959        // : make-inner-adder ( Int -- Closure )
4960        //   [ [ i.+ ] swap call ] ;
4961        // The closure captures an Int
4962        // When called, it creates an inner quotation and calls it with the captured value
4963        // This tests that closures can work with nested quotations
4964        let closure_effect = Effect::new(
4965            StackType::RowVar("r".to_string()).push(Type::Int),
4966            StackType::RowVar("r".to_string()).push(Type::Int),
4967        );
4968
4969        let program = Program {
4970            includes: vec![],
4971            unions: vec![],
4972            words: vec![WordDef {
4973                name: "make-inner-adder".to_string(),
4974                effect: Some(Effect::new(
4975                    StackType::singleton(Type::Int),
4976                    StackType::singleton(Type::Closure {
4977                        effect: Box::new(closure_effect),
4978                        captures: vec![Type::Int],
4979                    }),
4980                )),
4981                body: vec![Statement::Quotation {
4982                    span: None,
4983                    id: 0,
4984                    body: vec![
4985                        // The captured Int and the caller's Int are on stack
4986                        Statement::WordCall {
4987                            name: "i.add".to_string(),
4988                            span: None,
4989                        },
4990                    ],
4991                }],
4992                source: None,
4993                allowed_lints: vec![],
4994            }],
4995        };
4996
4997        let mut checker = TypeChecker::new();
4998        let result = checker.check_program(&program);
4999        assert!(
5000            result.is_ok(),
5001            "Closure with capture for inner work should pass: {:?}",
5002            result.err()
5003        );
5004    }
5005
5006    #[test]
5007    fn test_union_type_mismatch_should_fail() {
5008        // RFC #345: Different union types should not unify
5009        // This tests that passing union B to a function expecting union A fails
5010        use crate::ast::{UnionDef, UnionField, UnionVariant};
5011
5012        let mut program = Program {
5013            includes: vec![],
5014            unions: vec![
5015                UnionDef {
5016                    name: "UnionA".to_string(),
5017                    variants: vec![UnionVariant {
5018                        name: "AVal".to_string(),
5019                        fields: vec![UnionField {
5020                            name: "x".to_string(),
5021                            type_name: "Int".to_string(),
5022                        }],
5023                        source: None,
5024                    }],
5025                    source: None,
5026                },
5027                UnionDef {
5028                    name: "UnionB".to_string(),
5029                    variants: vec![UnionVariant {
5030                        name: "BVal".to_string(),
5031                        fields: vec![UnionField {
5032                            name: "y".to_string(),
5033                            type_name: "Int".to_string(),
5034                        }],
5035                        source: None,
5036                    }],
5037                    source: None,
5038                },
5039            ],
5040            words: vec![
5041                // : takes-a ( UnionA -- ) drop ;
5042                WordDef {
5043                    name: "takes-a".to_string(),
5044                    effect: Some(Effect::new(
5045                        StackType::RowVar("rest".to_string())
5046                            .push(Type::Union("UnionA".to_string())),
5047                        StackType::RowVar("rest".to_string()),
5048                    )),
5049                    body: vec![Statement::WordCall {
5050                        name: "drop".to_string(),
5051                        span: None,
5052                    }],
5053                    source: None,
5054                    allowed_lints: vec![],
5055                },
5056                // : main ( -- ) 99 Make-BVal takes-a ;
5057                // This should FAIL: Make-BVal returns UnionB, takes-a expects UnionA
5058                WordDef {
5059                    name: "main".to_string(),
5060                    effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
5061                    body: vec![
5062                        Statement::IntLiteral(99),
5063                        Statement::WordCall {
5064                            name: "Make-BVal".to_string(),
5065                            span: None,
5066                        },
5067                        Statement::WordCall {
5068                            name: "takes-a".to_string(),
5069                            span: None,
5070                        },
5071                    ],
5072                    source: None,
5073                    allowed_lints: vec![],
5074                },
5075            ],
5076        };
5077
5078        // Generate constructors (Make-AVal, Make-BVal)
5079        program.generate_constructors().unwrap();
5080
5081        let mut checker = TypeChecker::new();
5082        let result = checker.check_program(&program);
5083
5084        // This SHOULD fail because UnionB != UnionA
5085        assert!(
5086            result.is_err(),
5087            "Passing UnionB to function expecting UnionA should fail, but got: {:?}",
5088            result
5089        );
5090        let err = result.unwrap_err();
5091        assert!(
5092            err.contains("Union") || err.contains("mismatch"),
5093            "Error should mention union type mismatch, got: {}",
5094            err
5095        );
5096    }
5097
5098    // =========================================================================
5099    // Aux stack tests (Issue #350)
5100    // =========================================================================
5101
5102    fn make_word_call(name: &str) -> Statement {
5103        Statement::WordCall {
5104            name: name.to_string(),
5105            span: None,
5106        }
5107    }
5108
5109    #[test]
5110    fn test_aux_basic_round_trip() {
5111        // : test ( Int -- Int ) >aux aux> ;
5112        let program = Program {
5113            includes: vec![],
5114            unions: vec![],
5115            words: vec![WordDef {
5116                name: "test".to_string(),
5117                effect: Some(Effect::new(
5118                    StackType::singleton(Type::Int),
5119                    StackType::singleton(Type::Int),
5120                )),
5121                body: vec![make_word_call(">aux"), make_word_call("aux>")],
5122                source: None,
5123                allowed_lints: vec![],
5124            }],
5125        };
5126
5127        let mut checker = TypeChecker::new();
5128        assert!(checker.check_program(&program).is_ok());
5129    }
5130
5131    #[test]
5132    fn test_aux_preserves_type() {
5133        // : test ( String -- String ) >aux aux> ;
5134        let program = Program {
5135            includes: vec![],
5136            unions: vec![],
5137            words: vec![WordDef {
5138                name: "test".to_string(),
5139                effect: Some(Effect::new(
5140                    StackType::singleton(Type::String),
5141                    StackType::singleton(Type::String),
5142                )),
5143                body: vec![make_word_call(">aux"), make_word_call("aux>")],
5144                source: None,
5145                allowed_lints: vec![],
5146            }],
5147        };
5148
5149        let mut checker = TypeChecker::new();
5150        assert!(checker.check_program(&program).is_ok());
5151    }
5152
5153    #[test]
5154    fn test_aux_unbalanced_error() {
5155        // : test ( Int -- ) >aux ;  -- ERROR: aux not empty at return
5156        let program = Program {
5157            includes: vec![],
5158            unions: vec![],
5159            words: vec![WordDef {
5160                name: "test".to_string(),
5161                effect: Some(Effect::new(
5162                    StackType::singleton(Type::Int),
5163                    StackType::Empty,
5164                )),
5165                body: vec![make_word_call(">aux")],
5166                source: None,
5167                allowed_lints: vec![],
5168            }],
5169        };
5170
5171        let mut checker = TypeChecker::new();
5172        let result = checker.check_program(&program);
5173        assert!(result.is_err());
5174        let err = result.unwrap_err();
5175        assert!(
5176            err.contains("aux stack is not empty"),
5177            "Expected aux stack balance error, got: {}",
5178            err
5179        );
5180    }
5181
5182    #[test]
5183    fn test_aux_pop_empty_error() {
5184        // : test ( -- Int ) aux> ;  -- ERROR: aux is empty
5185        let program = Program {
5186            includes: vec![],
5187            unions: vec![],
5188            words: vec![WordDef {
5189                name: "test".to_string(),
5190                effect: Some(Effect::new(
5191                    StackType::Empty,
5192                    StackType::singleton(Type::Int),
5193                )),
5194                body: vec![make_word_call("aux>")],
5195                source: None,
5196                allowed_lints: vec![],
5197            }],
5198        };
5199
5200        let mut checker = TypeChecker::new();
5201        let result = checker.check_program(&program);
5202        assert!(result.is_err());
5203        let err = result.unwrap_err();
5204        assert!(
5205            err.contains("aux stack is empty"),
5206            "Expected aux empty error, got: {}",
5207            err
5208        );
5209    }
5210
5211    #[test]
5212    fn test_aux_multiple_values() {
5213        // >aux >aux aux> aux> is identity (LIFO preserves order)
5214        // Input: ( Int String ), >aux pops String, >aux pops Int
5215        // aux> pushes Int, aux> pushes String → output: ( Int String )
5216        let program = Program {
5217            includes: vec![],
5218            unions: vec![],
5219            words: vec![WordDef {
5220                name: "test".to_string(),
5221                effect: Some(Effect::new(
5222                    StackType::Empty.push(Type::Int).push(Type::String),
5223                    StackType::Empty.push(Type::Int).push(Type::String),
5224                )),
5225                body: vec![
5226                    make_word_call(">aux"),
5227                    make_word_call(">aux"),
5228                    make_word_call("aux>"),
5229                    make_word_call("aux>"),
5230                ],
5231                source: None,
5232                allowed_lints: vec![],
5233            }],
5234        };
5235
5236        let mut checker = TypeChecker::new();
5237        assert!(checker.check_program(&program).is_ok());
5238    }
5239
5240    #[test]
5241    fn test_aux_max_depths_tracked() {
5242        // : test ( Int -- Int ) >aux aux> ;
5243        let program = Program {
5244            includes: vec![],
5245            unions: vec![],
5246            words: vec![WordDef {
5247                name: "test".to_string(),
5248                effect: Some(Effect::new(
5249                    StackType::singleton(Type::Int),
5250                    StackType::singleton(Type::Int),
5251                )),
5252                body: vec![make_word_call(">aux"), make_word_call("aux>")],
5253                source: None,
5254                allowed_lints: vec![],
5255            }],
5256        };
5257
5258        let mut checker = TypeChecker::new();
5259        checker.check_program(&program).unwrap();
5260        let depths = checker.take_aux_max_depths();
5261        assert_eq!(depths.get("test"), Some(&1));
5262    }
5263
5264    #[test]
5265    fn test_aux_in_match_balanced() {
5266        // Aux used symmetrically in match arms: each arm does >aux aux>
5267        use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};
5268
5269        let union_def = UnionDef {
5270            name: "Choice".to_string(),
5271            variants: vec![
5272                UnionVariant {
5273                    name: "Left".to_string(),
5274                    fields: vec![],
5275                    source: None,
5276                },
5277                UnionVariant {
5278                    name: "Right".to_string(),
5279                    fields: vec![],
5280                    source: None,
5281                },
5282            ],
5283            source: None,
5284        };
5285
5286        // : test ( Int Choice -- Int )
5287        //   swap >aux match Left => aux> end Right => aux> end ;
5288        let program = Program {
5289            includes: vec![],
5290            unions: vec![union_def],
5291            words: vec![WordDef {
5292                name: "test".to_string(),
5293                effect: Some(Effect::new(
5294                    StackType::Empty
5295                        .push(Type::Int)
5296                        .push(Type::Union("Choice".to_string())),
5297                    StackType::singleton(Type::Int),
5298                )),
5299                body: vec![
5300                    make_word_call("swap"),
5301                    make_word_call(">aux"),
5302                    Statement::Match {
5303                        arms: vec![
5304                            MatchArm {
5305                                pattern: Pattern::Variant("Left".to_string()),
5306                                body: vec![make_word_call("aux>")],
5307                                span: None,
5308                            },
5309                            MatchArm {
5310                                pattern: Pattern::Variant("Right".to_string()),
5311                                body: vec![make_word_call("aux>")],
5312                                span: None,
5313                            },
5314                        ],
5315                        span: None,
5316                    },
5317                ],
5318                source: None,
5319                allowed_lints: vec![],
5320            }],
5321        };
5322
5323        let mut checker = TypeChecker::new();
5324        assert!(checker.check_program(&program).is_ok());
5325    }
5326
5327    #[test]
5328    fn test_aux_in_match_unbalanced_error() {
5329        // Match arms with different aux states should error
5330        use crate::ast::{MatchArm, Pattern, UnionDef, UnionVariant};
5331
5332        let union_def = UnionDef {
5333            name: "Choice".to_string(),
5334            variants: vec![
5335                UnionVariant {
5336                    name: "Left".to_string(),
5337                    fields: vec![],
5338                    source: None,
5339                },
5340                UnionVariant {
5341                    name: "Right".to_string(),
5342                    fields: vec![],
5343                    source: None,
5344                },
5345            ],
5346            source: None,
5347        };
5348
5349        // : test ( Int Choice -- Int )
5350        //   swap >aux match Left => aux> end Right => end ;  -- ERROR: unbalanced
5351        let program = Program {
5352            includes: vec![],
5353            unions: vec![union_def],
5354            words: vec![WordDef {
5355                name: "test".to_string(),
5356                effect: Some(Effect::new(
5357                    StackType::Empty
5358                        .push(Type::Int)
5359                        .push(Type::Union("Choice".to_string())),
5360                    StackType::singleton(Type::Int),
5361                )),
5362                body: vec![
5363                    make_word_call("swap"),
5364                    make_word_call(">aux"),
5365                    Statement::Match {
5366                        arms: vec![
5367                            MatchArm {
5368                                pattern: Pattern::Variant("Left".to_string()),
5369                                body: vec![make_word_call("aux>")],
5370                                span: None,
5371                            },
5372                            MatchArm {
5373                                pattern: Pattern::Variant("Right".to_string()),
5374                                body: vec![],
5375                                span: None,
5376                            },
5377                        ],
5378                        span: None,
5379                    },
5380                ],
5381                source: None,
5382                allowed_lints: vec![],
5383            }],
5384        };
5385
5386        let mut checker = TypeChecker::new();
5387        let result = checker.check_program(&program);
5388        assert!(result.is_err());
5389        let err = result.unwrap_err();
5390        assert!(
5391            err.contains("aux stack"),
5392            "Expected aux stack mismatch error, got: {}",
5393            err
5394        );
5395    }
5396
5397    #[test]
5398    fn test_aux_in_quotation_balanced_accepted() {
5399        // Issue #393: balanced >aux/aux> inside a quotation is now allowed.
5400        // The word produces a quotation [ Int -- Int ]. The quotation body
5401        // uses >aux/aux> to round-trip the input Int. Lexical scoping is
5402        // preserved because both ops are inside the same quotation.
5403        let program = Program {
5404            includes: vec![],
5405            unions: vec![],
5406            words: vec![WordDef {
5407                name: "test".to_string(),
5408                effect: Some(Effect::new(
5409                    StackType::Empty,
5410                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
5411                        StackType::singleton(Type::Int),
5412                        StackType::singleton(Type::Int),
5413                    )))),
5414                )),
5415                body: vec![Statement::Quotation {
5416                    span: None,
5417                    id: 0,
5418                    body: vec![make_word_call(">aux"), make_word_call("aux>")],
5419                }],
5420                source: None,
5421                allowed_lints: vec![],
5422            }],
5423        };
5424
5425        let mut checker = TypeChecker::new();
5426        let result = checker.check_program(&program);
5427        assert!(
5428            result.is_ok(),
5429            "should accept balanced aux in quotation: {:?}",
5430            result.err()
5431        );
5432    }
5433
5434    #[test]
5435    fn test_aux_in_quotation_unbalanced_rejected() {
5436        // Issue #393: a quotation that pushes onto aux without popping
5437        // must still be rejected. Lexical scoping requires every >aux to
5438        // be matched by an aux> within the same quotation.
5439        let program = Program {
5440            includes: vec![],
5441            unions: vec![],
5442            words: vec![WordDef {
5443                name: "test".to_string(),
5444                effect: Some(Effect::new(
5445                    StackType::Empty,
5446                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
5447                        StackType::singleton(Type::Int),
5448                        StackType::Empty,
5449                    )))),
5450                )),
5451                body: vec![Statement::Quotation {
5452                    span: None,
5453                    id: 0,
5454                    body: vec![make_word_call(">aux")],
5455                }],
5456                source: None,
5457                allowed_lints: vec![],
5458            }],
5459        };
5460
5461        let mut checker = TypeChecker::new();
5462        let result = checker.check_program(&program);
5463        assert!(result.is_err());
5464        let err = result.unwrap_err();
5465        assert!(
5466            err.contains("unbalanced aux stack"),
5467            "Expected unbalanced aux error, got: {}",
5468            err
5469        );
5470    }
5471
5472    // =========================================================================
5473    // Dataflow combinator tests
5474    // =========================================================================
5475
5476    #[test]
5477    fn test_dip_basic() {
5478        // : test ( Int Int -- Int Int )  [ 1 i.+ ] dip ;
5479        // Increments value below top, preserving top
5480        let program = Program {
5481            includes: vec![],
5482            unions: vec![],
5483            words: vec![WordDef {
5484                name: "test".to_string(),
5485                effect: Some(Effect::new(
5486                    StackType::Empty.push(Type::Int).push(Type::Int),
5487                    StackType::Empty.push(Type::Int).push(Type::Int),
5488                )),
5489                body: vec![
5490                    Statement::Quotation {
5491                        id: 0,
5492                        body: vec![
5493                            Statement::IntLiteral(1),
5494                            Statement::WordCall {
5495                                name: "i.+".to_string(),
5496                                span: None,
5497                            },
5498                        ],
5499                        span: None,
5500                    },
5501                    Statement::WordCall {
5502                        name: "dip".to_string(),
5503                        span: None,
5504                    },
5505                ],
5506                source: None,
5507                allowed_lints: vec![],
5508            }],
5509        };
5510
5511        let mut checker = TypeChecker::new();
5512        assert!(checker.check_program(&program).is_ok());
5513    }
5514
5515    #[test]
5516    fn test_dip_type_mismatch() {
5517        // : test ( String Int -- ?? )  [ 1 i.+ ] dip ;
5518        // Should fail: quotation expects Int but gets String
5519        let program = Program {
5520            includes: vec![],
5521            unions: vec![],
5522            words: vec![WordDef {
5523                name: "test".to_string(),
5524                effect: Some(Effect::new(
5525                    StackType::Empty.push(Type::String).push(Type::Int),
5526                    StackType::Empty.push(Type::Int).push(Type::Int),
5527                )),
5528                body: vec![
5529                    Statement::Quotation {
5530                        id: 0,
5531                        body: vec![
5532                            Statement::IntLiteral(1),
5533                            Statement::WordCall {
5534                                name: "i.+".to_string(),
5535                                span: None,
5536                            },
5537                        ],
5538                        span: None,
5539                    },
5540                    Statement::WordCall {
5541                        name: "dip".to_string(),
5542                        span: None,
5543                    },
5544                ],
5545                source: None,
5546                allowed_lints: vec![],
5547            }],
5548        };
5549
5550        let mut checker = TypeChecker::new();
5551        assert!(checker.check_program(&program).is_err());
5552    }
5553
5554    #[test]
5555    fn test_keep_basic() {
5556        // : test ( Int -- Int Int )  [ dup i.* ] keep ;
5557        // Squares and keeps original
5558        let program = Program {
5559            includes: vec![],
5560            unions: vec![],
5561            words: vec![WordDef {
5562                name: "test".to_string(),
5563                effect: Some(Effect::new(
5564                    StackType::singleton(Type::Int),
5565                    StackType::Empty.push(Type::Int).push(Type::Int),
5566                )),
5567                body: vec![
5568                    Statement::Quotation {
5569                        id: 0,
5570                        body: vec![
5571                            Statement::WordCall {
5572                                name: "dup".to_string(),
5573                                span: None,
5574                            },
5575                            Statement::WordCall {
5576                                name: "i.*".to_string(),
5577                                span: None,
5578                            },
5579                        ],
5580                        span: None,
5581                    },
5582                    Statement::WordCall {
5583                        name: "keep".to_string(),
5584                        span: None,
5585                    },
5586                ],
5587                source: None,
5588                allowed_lints: vec![],
5589            }],
5590        };
5591
5592        let mut checker = TypeChecker::new();
5593        assert!(checker.check_program(&program).is_ok());
5594    }
5595
5596    #[test]
5597    fn test_bi_basic() {
5598        // : test ( Int -- Int Int )  [ 2 i.* ] [ 3 i.* ] bi ;
5599        // Double and triple
5600        let program = Program {
5601            includes: vec![],
5602            unions: vec![],
5603            words: vec![WordDef {
5604                name: "test".to_string(),
5605                effect: Some(Effect::new(
5606                    StackType::singleton(Type::Int),
5607                    StackType::Empty.push(Type::Int).push(Type::Int),
5608                )),
5609                body: vec![
5610                    Statement::Quotation {
5611                        id: 0,
5612                        body: vec![
5613                            Statement::IntLiteral(2),
5614                            Statement::WordCall {
5615                                name: "i.*".to_string(),
5616                                span: None,
5617                            },
5618                        ],
5619                        span: None,
5620                    },
5621                    Statement::Quotation {
5622                        id: 1,
5623                        body: vec![
5624                            Statement::IntLiteral(3),
5625                            Statement::WordCall {
5626                                name: "i.*".to_string(),
5627                                span: None,
5628                            },
5629                        ],
5630                        span: None,
5631                    },
5632                    Statement::WordCall {
5633                        name: "bi".to_string(),
5634                        span: None,
5635                    },
5636                ],
5637                source: None,
5638                allowed_lints: vec![],
5639            }],
5640        };
5641
5642        let mut checker = TypeChecker::new();
5643        assert!(checker.check_program(&program).is_ok());
5644    }
5645
5646    #[test]
5647    fn test_keep_type_mismatch() {
5648        // : test ( String -- ?? )  [ 1 i.+ ] keep ;
5649        // Should fail: quotation expects Int but gets String
5650        let program = Program {
5651            includes: vec![],
5652            unions: vec![],
5653            words: vec![WordDef {
5654                name: "test".to_string(),
5655                effect: Some(Effect::new(
5656                    StackType::singleton(Type::String),
5657                    StackType::Empty.push(Type::Int).push(Type::String),
5658                )),
5659                body: vec![
5660                    Statement::Quotation {
5661                        id: 0,
5662                        body: vec![
5663                            Statement::IntLiteral(1),
5664                            Statement::WordCall {
5665                                name: "i.+".to_string(),
5666                                span: None,
5667                            },
5668                        ],
5669                        span: None,
5670                    },
5671                    Statement::WordCall {
5672                        name: "keep".to_string(),
5673                        span: None,
5674                    },
5675                ],
5676                source: None,
5677                allowed_lints: vec![],
5678            }],
5679        };
5680
5681        let mut checker = TypeChecker::new();
5682        assert!(checker.check_program(&program).is_err());
5683    }
5684
5685    #[test]
5686    fn test_bi_type_mismatch() {
5687        // : test ( String -- ?? )  [ string.length ] [ 1 i.+ ] bi ;
5688        // Should fail: second quotation expects Int but value is String
5689        let program = Program {
5690            includes: vec![],
5691            unions: vec![],
5692            words: vec![WordDef {
5693                name: "test".to_string(),
5694                effect: Some(Effect::new(
5695                    StackType::singleton(Type::String),
5696                    StackType::Empty.push(Type::Int).push(Type::Int),
5697                )),
5698                body: vec![
5699                    Statement::Quotation {
5700                        id: 0,
5701                        body: vec![Statement::WordCall {
5702                            name: "string.length".to_string(),
5703                            span: None,
5704                        }],
5705                        span: None,
5706                    },
5707                    Statement::Quotation {
5708                        id: 1,
5709                        body: vec![
5710                            Statement::IntLiteral(1),
5711                            Statement::WordCall {
5712                                name: "i.+".to_string(),
5713                                span: None,
5714                            },
5715                        ],
5716                        span: None,
5717                    },
5718                    Statement::WordCall {
5719                        name: "bi".to_string(),
5720                        span: None,
5721                    },
5722                ],
5723                source: None,
5724                allowed_lints: vec![],
5725            }],
5726        };
5727
5728        let mut checker = TypeChecker::new();
5729        assert!(checker.check_program(&program).is_err());
5730    }
5731
5732    #[test]
5733    fn test_dip_underflow() {
5734        // : test ( -- ?? )  [ 1 ] dip ;
5735        // Should fail: dip needs a value below the quotation
5736        let program = Program {
5737            includes: vec![],
5738            unions: vec![],
5739            words: vec![WordDef {
5740                name: "test".to_string(),
5741                effect: Some(Effect::new(
5742                    StackType::Empty,
5743                    StackType::singleton(Type::Int),
5744                )),
5745                body: vec![
5746                    Statement::Quotation {
5747                        id: 0,
5748                        body: vec![Statement::IntLiteral(1)],
5749                        span: None,
5750                    },
5751                    Statement::WordCall {
5752                        name: "dip".to_string(),
5753                        span: None,
5754                    },
5755                ],
5756                source: None,
5757                allowed_lints: vec![],
5758            }],
5759        };
5760
5761        let mut checker = TypeChecker::new();
5762        let result = checker.check_program(&program);
5763        assert!(result.is_err());
5764        let err = result.unwrap_err();
5765        assert!(
5766            err.contains("stack underflow"),
5767            "Expected underflow error, got: {}",
5768            err
5769        );
5770    }
5771
5772    #[test]
5773    fn test_dip_preserves_type() {
5774        // : test ( Int String -- Int String )  [ 1 i.+ ] dip ;
5775        // The String on top is preserved, Int below is incremented
5776        let program = Program {
5777            includes: vec![],
5778            unions: vec![],
5779            words: vec![WordDef {
5780                name: "test".to_string(),
5781                effect: Some(Effect::new(
5782                    StackType::Empty.push(Type::Int).push(Type::String),
5783                    StackType::Empty.push(Type::Int).push(Type::String),
5784                )),
5785                body: vec![
5786                    Statement::Quotation {
5787                        id: 0,
5788                        body: vec![
5789                            Statement::IntLiteral(1),
5790                            Statement::WordCall {
5791                                name: "i.+".to_string(),
5792                                span: None,
5793                            },
5794                        ],
5795                        span: None,
5796                    },
5797                    Statement::WordCall {
5798                        name: "dip".to_string(),
5799                        span: None,
5800                    },
5801                ],
5802                source: None,
5803                allowed_lints: vec![],
5804            }],
5805        };
5806
5807        let mut checker = TypeChecker::new();
5808        assert!(checker.check_program(&program).is_ok());
5809    }
5810
5811    #[test]
5812    fn test_keep_underflow() {
5813        // : test ( -- ?? )  [ drop ] keep ;
5814        // Should fail: keep needs a value below the quotation
5815        let program = Program {
5816            includes: vec![],
5817            unions: vec![],
5818            words: vec![WordDef {
5819                name: "test".to_string(),
5820                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
5821                body: vec![
5822                    Statement::Quotation {
5823                        id: 0,
5824                        body: vec![Statement::WordCall {
5825                            name: "drop".to_string(),
5826                            span: None,
5827                        }],
5828                        span: None,
5829                    },
5830                    Statement::WordCall {
5831                        name: "keep".to_string(),
5832                        span: None,
5833                    },
5834                ],
5835                source: None,
5836                allowed_lints: vec![],
5837            }],
5838        };
5839
5840        let mut checker = TypeChecker::new();
5841        let result = checker.check_program(&program);
5842        assert!(result.is_err());
5843        let err = result.unwrap_err();
5844        assert!(
5845            err.contains("stack underflow") || err.contains("underflow"),
5846            "Expected underflow error, got: {}",
5847            err
5848        );
5849    }
5850
5851    #[test]
5852    fn test_bi_underflow() {
5853        // : test ( -- ?? )  [ 1 ] [ 2 ] bi ;
5854        // Should fail: bi needs a value below the two quotations
5855        let program = Program {
5856            includes: vec![],
5857            unions: vec![],
5858            words: vec![WordDef {
5859                name: "test".to_string(),
5860                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
5861                body: vec![
5862                    Statement::Quotation {
5863                        id: 0,
5864                        body: vec![Statement::IntLiteral(1)],
5865                        span: None,
5866                    },
5867                    Statement::Quotation {
5868                        id: 1,
5869                        body: vec![Statement::IntLiteral(2)],
5870                        span: None,
5871                    },
5872                    Statement::WordCall {
5873                        name: "bi".to_string(),
5874                        span: None,
5875                    },
5876                ],
5877                source: None,
5878                allowed_lints: vec![],
5879            }],
5880        };
5881
5882        let mut checker = TypeChecker::new();
5883        let result = checker.check_program(&program);
5884        assert!(result.is_err());
5885        let err = result.unwrap_err();
5886        assert!(
5887            err.contains("stack underflow") || err.contains("underflow"),
5888            "Expected underflow error, got: {}",
5889            err
5890        );
5891    }
5892
5893    #[test]
5894    fn test_bi_polymorphic_quotations() {
5895        // : test ( Int -- Int String )  [ 2 i.* ] [ int->string ] bi ;
5896        // Two quotations with different output types — verifies independent typing
5897        let program = Program {
5898            includes: vec![],
5899            unions: vec![],
5900            words: vec![WordDef {
5901                name: "test".to_string(),
5902                effect: Some(Effect::new(
5903                    StackType::singleton(Type::Int),
5904                    StackType::Empty.push(Type::Int).push(Type::String),
5905                )),
5906                body: vec![
5907                    Statement::Quotation {
5908                        id: 0,
5909                        body: vec![
5910                            Statement::IntLiteral(2),
5911                            Statement::WordCall {
5912                                name: "i.*".to_string(),
5913                                span: None,
5914                            },
5915                        ],
5916                        span: None,
5917                    },
5918                    Statement::Quotation {
5919                        id: 1,
5920                        body: vec![Statement::WordCall {
5921                            name: "int->string".to_string(),
5922                            span: None,
5923                        }],
5924                        span: None,
5925                    },
5926                    Statement::WordCall {
5927                        name: "bi".to_string(),
5928                        span: None,
5929                    },
5930                ],
5931                source: None,
5932                allowed_lints: vec![],
5933            }],
5934        };
5935
5936        let mut checker = TypeChecker::new();
5937        assert!(checker.check_program(&program).is_ok());
5938    }
5939}