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::types::{Effect, StackType, Type};
9use crate::unification::{Subst, unify_stacks};
10use std::collections::HashMap;
11
12pub struct TypeChecker {
13    /// Environment mapping word names to their effects
14    env: HashMap<String, Effect>,
15    /// Counter for generating fresh type variables
16    fresh_counter: std::cell::Cell<usize>,
17    /// Quotation types tracked during type checking
18    /// Maps quotation ID (from AST) to inferred type (Quotation or Closure)
19    /// This type map is used by codegen to generate appropriate code
20    quotation_types: std::cell::RefCell<HashMap<usize, Type>>,
21    /// Expected quotation/closure type (from word signature, if any)
22    /// Used during type-driven capture inference
23    expected_quotation_type: std::cell::RefCell<Option<Type>>,
24}
25
26impl TypeChecker {
27    pub fn new() -> Self {
28        TypeChecker {
29            env: HashMap::new(),
30            fresh_counter: std::cell::Cell::new(0),
31            quotation_types: std::cell::RefCell::new(HashMap::new()),
32            expected_quotation_type: std::cell::RefCell::new(None),
33        }
34    }
35
36    /// Register external word effects (e.g., from included modules).
37    ///
38    /// Words with `Some(effect)` get their actual signature.
39    /// Words with `None` get a maximally polymorphic placeholder `( ..a -- ..b )`.
40    pub fn register_external_words(&mut self, words: &[(&str, Option<&Effect>)]) {
41        for (name, effect) in words {
42            if let Some(eff) = effect {
43                self.env.insert(name.to_string(), (*eff).clone());
44            } else {
45                // Maximally polymorphic placeholder
46                let placeholder = Effect::new(
47                    StackType::RowVar("ext_in".to_string()),
48                    StackType::RowVar("ext_out".to_string()),
49                );
50                self.env.insert(name.to_string(), placeholder);
51            }
52        }
53    }
54
55    /// Extract the type map (quotation ID -> inferred type)
56    ///
57    /// This should be called after check_program() to get the inferred types
58    /// for all quotations in the program. The map is used by codegen to generate
59    /// appropriate code for Quotations vs Closures.
60    pub fn take_quotation_types(&self) -> HashMap<usize, Type> {
61        self.quotation_types.replace(HashMap::new())
62    }
63
64    /// Generate a fresh variable name
65    fn fresh_var(&self, prefix: &str) -> String {
66        let n = self.fresh_counter.get();
67        self.fresh_counter.set(n + 1);
68        format!("{}${}", prefix, n)
69    }
70
71    /// Freshen all type and row variables in an effect
72    fn freshen_effect(&self, effect: &Effect) -> Effect {
73        let mut type_map = HashMap::new();
74        let mut row_map = HashMap::new();
75
76        let fresh_inputs = self.freshen_stack(&effect.inputs, &mut type_map, &mut row_map);
77        let fresh_outputs = self.freshen_stack(&effect.outputs, &mut type_map, &mut row_map);
78
79        Effect::new(fresh_inputs, fresh_outputs)
80    }
81
82    fn freshen_stack(
83        &self,
84        stack: &StackType,
85        type_map: &mut HashMap<String, String>,
86        row_map: &mut HashMap<String, String>,
87    ) -> StackType {
88        match stack {
89            StackType::Empty => StackType::Empty,
90            StackType::RowVar(name) => {
91                let fresh_name = row_map
92                    .entry(name.clone())
93                    .or_insert_with(|| self.fresh_var(name));
94                StackType::RowVar(fresh_name.clone())
95            }
96            StackType::Cons { rest, top } => {
97                let fresh_rest = self.freshen_stack(rest, type_map, row_map);
98                let fresh_top = self.freshen_type(top, type_map, row_map);
99                StackType::Cons {
100                    rest: Box::new(fresh_rest),
101                    top: fresh_top,
102                }
103            }
104        }
105    }
106
107    fn freshen_type(
108        &self,
109        ty: &Type,
110        type_map: &mut HashMap<String, String>,
111        row_map: &mut HashMap<String, String>,
112    ) -> Type {
113        match ty {
114            Type::Int | Type::Float | Type::Bool | Type::String => ty.clone(),
115            Type::Var(name) => {
116                let fresh_name = type_map
117                    .entry(name.clone())
118                    .or_insert_with(|| self.fresh_var(name));
119                Type::Var(fresh_name.clone())
120            }
121            Type::Quotation(effect) => {
122                let fresh_inputs = self.freshen_stack(&effect.inputs, type_map, row_map);
123                let fresh_outputs = self.freshen_stack(&effect.outputs, type_map, row_map);
124                Type::Quotation(Box::new(Effect::new(fresh_inputs, fresh_outputs)))
125            }
126            Type::Closure { effect, captures } => {
127                let fresh_inputs = self.freshen_stack(&effect.inputs, type_map, row_map);
128                let fresh_outputs = self.freshen_stack(&effect.outputs, type_map, row_map);
129                let fresh_captures = captures
130                    .iter()
131                    .map(|t| self.freshen_type(t, type_map, row_map))
132                    .collect();
133                Type::Closure {
134                    effect: Box::new(Effect::new(fresh_inputs, fresh_outputs)),
135                    captures: fresh_captures,
136                }
137            }
138        }
139    }
140
141    /// Type check a complete program
142    pub fn check_program(&mut self, program: &Program) -> Result<(), String> {
143        // First pass: collect all word signatures
144        // For words without explicit effects, use a maximally polymorphic placeholder
145        // This allows calls to work, and actual type safety comes from checking the body
146        for word in &program.words {
147            if let Some(effect) = &word.effect {
148                self.env.insert(word.name.clone(), effect.clone());
149            } else {
150                // Use placeholder effect: ( ..input -- ..output )
151                // This is maximally polymorphic and allows any usage
152                let placeholder = Effect::new(
153                    StackType::RowVar("input".to_string()),
154                    StackType::RowVar("output".to_string()),
155                );
156                self.env.insert(word.name.clone(), placeholder);
157            }
158        }
159
160        // Second pass: type check each word body
161        for word in &program.words {
162            self.check_word(word)?;
163        }
164
165        Ok(())
166    }
167
168    /// Type check a word definition
169    fn check_word(&self, word: &WordDef) -> Result<(), String> {
170        // If word has declared effect, verify body matches it
171        if let Some(declared_effect) = &word.effect {
172            // Check if the word's output type is a quotation or closure
173            // If so, store it as the expected type for capture inference
174            if let Some((_rest, top_type)) = declared_effect.outputs.clone().pop()
175                && matches!(top_type, Type::Quotation(_) | Type::Closure { .. })
176            {
177                *self.expected_quotation_type.borrow_mut() = Some(top_type);
178            }
179
180            // Infer the result stack starting from declared input
181            let (result_stack, _subst) =
182                self.infer_statements_from(&word.body, &declared_effect.inputs)?;
183
184            // Clear expected type after checking
185            *self.expected_quotation_type.borrow_mut() = None;
186
187            // Verify result matches declared output
188            unify_stacks(&declared_effect.outputs, &result_stack).map_err(|e| {
189                format!(
190                    "Word '{}': declared output stack ({:?}) doesn't match inferred ({:?}): {}",
191                    word.name, declared_effect.outputs, result_stack, e
192                )
193            })?;
194        } else {
195            // No declared effect - just verify body is well-typed
196            // Start from polymorphic input
197            self.infer_statements_from(&word.body, &StackType::RowVar("input".to_string()))?;
198        }
199
200        Ok(())
201    }
202
203    /// Infer the resulting stack type from a sequence of statements
204    /// starting from a given input stack
205    fn infer_statements_from(
206        &self,
207        statements: &[Statement],
208        start_stack: &StackType,
209    ) -> Result<(StackType, Subst), String> {
210        let mut current_stack = start_stack.clone();
211        let mut accumulated_subst = Subst::empty();
212
213        for (i, stmt) in statements.iter().enumerate() {
214            // Look ahead: if this is a quotation followed by a word that expects specific quotation type,
215            // set the expected type before checking the quotation
216            let saved_expected_type = if matches!(stmt, Statement::Quotation { .. }) {
217                // Save the current expected type
218                let saved = self.expected_quotation_type.borrow().clone();
219
220                // Try to set expected type based on lookahead
221                if let Some(Statement::WordCall(next_word)) = statements.get(i + 1) {
222                    // Check if the next word expects a specific quotation type
223                    if let Some(next_effect) = self.lookup_word_effect(next_word) {
224                        // Extract the quotation type expected by the next word
225                        // For operations like spawn: ( ..a Quotation(-- ) -- ..a Int )
226                        if let Some((_rest, quot_type)) = next_effect.inputs.clone().pop()
227                            && matches!(quot_type, Type::Quotation(_))
228                        {
229                            *self.expected_quotation_type.borrow_mut() = Some(quot_type);
230                        }
231                    }
232                }
233                Some(saved)
234            } else {
235                None
236            };
237
238            let (new_stack, subst) = self.infer_statement(stmt, current_stack)?;
239            current_stack = new_stack;
240            accumulated_subst = accumulated_subst.compose(&subst);
241
242            // Restore expected type after checking quotation
243            if let Some(saved) = saved_expected_type {
244                *self.expected_quotation_type.borrow_mut() = saved;
245            }
246        }
247
248        Ok((current_stack, accumulated_subst))
249    }
250
251    /// Infer the stack effect of a sequence of statements
252    /// Returns an Effect with both inputs and outputs normalized by applying discovered substitutions
253    fn infer_statements(&self, statements: &[Statement]) -> Result<Effect, String> {
254        let start = StackType::RowVar("input".to_string());
255        let (result, subst) = self.infer_statements_from(statements, &start)?;
256
257        // Apply the accumulated substitution to both start and result
258        // This ensures row variables are consistently named
259        let normalized_start = subst.apply_stack(&start);
260        let normalized_result = subst.apply_stack(&result);
261
262        Ok(Effect::new(normalized_start, normalized_result))
263    }
264
265    /// Infer the resulting stack type after a statement
266    /// Takes current stack, returns (new stack, substitution) after statement
267    fn infer_statement(
268        &self,
269        statement: &Statement,
270        current_stack: StackType,
271    ) -> Result<(StackType, Subst), String> {
272        match statement {
273            Statement::IntLiteral(_) => {
274                // Push Int onto stack
275                Ok((current_stack.push(Type::Int), Subst::empty()))
276            }
277
278            Statement::BoolLiteral(_) => {
279                // Push Bool onto stack (currently represented as Int at runtime)
280                // But we track it as Int in the type system
281                Ok((current_stack.push(Type::Int), Subst::empty()))
282            }
283
284            Statement::StringLiteral(_) => {
285                // Push String onto stack
286                Ok((current_stack.push(Type::String), Subst::empty()))
287            }
288
289            Statement::FloatLiteral(_) => {
290                // Push Float onto stack
291                Ok((current_stack.push(Type::Float), Subst::empty()))
292            }
293
294            Statement::WordCall(name) => {
295                // Look up word's effect
296                let effect = self
297                    .lookup_word_effect(name)
298                    .ok_or_else(|| format!("Unknown word: '{}'", name))?;
299
300                // Freshen the effect to avoid variable name clashes
301                let fresh_effect = self.freshen_effect(&effect);
302
303                // Special handling for spawn: auto-convert Quotation to Closure if needed
304                let adjusted_stack = if name == "spawn" {
305                    self.adjust_stack_for_spawn(current_stack, &fresh_effect)?
306                } else {
307                    current_stack
308                };
309
310                // Apply the freshened effect to current stack
311                self.apply_effect(&fresh_effect, adjusted_stack, name)
312            }
313
314            Statement::If {
315                then_branch,
316                else_branch,
317            } => {
318                // Pop condition (must be Int/Bool)
319                let (stack_after_cond, cond_type) =
320                    self.pop_type(&current_stack, "if condition")?;
321
322                // Condition must be Int (Forth-style: 0 = false, non-zero = true)
323                let cond_subst = unify_stacks(
324                    &StackType::singleton(Type::Int),
325                    &StackType::singleton(cond_type),
326                )
327                .map_err(|e| format!("if condition must be Int: {}", e))?;
328
329                let stack_after_cond = cond_subst.apply_stack(&stack_after_cond);
330
331                // Infer then branch
332                let then_effect = self.infer_statements(then_branch)?;
333                let (then_result, _then_subst) =
334                    self.apply_effect(&then_effect, stack_after_cond.clone(), "if then")?;
335
336                // Infer else branch (or use stack_after_cond if no else)
337                let (else_result, _else_subst) = if let Some(else_stmts) = else_branch {
338                    let else_effect = self.infer_statements(else_stmts)?;
339                    self.apply_effect(&else_effect, stack_after_cond, "if else")?
340                } else {
341                    (stack_after_cond, Subst::empty())
342                };
343
344                // Both branches must produce compatible stacks
345                let branch_subst = unify_stacks(&then_result, &else_result).map_err(|e| {
346                    format!(
347                        "if/else branches have incompatible stack effects:\n\
348                         \x20 then branch produces: {}\n\
349                         \x20 else branch produces: {}\n\
350                         \x20 Both branches of an if/else must produce the same stack shape.\n\
351                         \x20 Hint: Make sure both branches push/pop the same number of values.\n\
352                         \x20 Error: {}",
353                        then_result, else_result, e
354                    )
355                })?;
356
357                // Apply branch unification to get the final result
358                let result = branch_subst.apply_stack(&then_result);
359
360                // Propagate condition substitution composed with branch substitution
361                let total_subst = cond_subst.compose(&branch_subst);
362                Ok((result, total_subst))
363            }
364
365            Statement::Quotation { id, body } => {
366                // Type checking for quotations with automatic capture analysis
367                //
368                // A quotation is a block of deferred code.
369                //
370                // For stateless quotations:
371                //   Example: [ 1 add ]
372                //   Body effect: ( -- Int )  (pushes 1, needs Int from call site, adds)
373                //   Type: Quotation([Int -- Int])
374                //
375                // For closures (automatic capture):
376                //   Example: 5 [ add ]
377                //   Body effect: ( Int Int -- Int )  (needs 2 Ints)
378                //   One Int will be captured from creation site
379                //   Type: Closure { effect: [Int -- Int], captures: [Int] }
380
381                // Infer the effect of the quotation body
382                let body_effect = self.infer_statements(body)?;
383
384                // Perform capture analysis
385                let quot_type = self.analyze_captures(&body_effect, &current_stack)?;
386
387                // Record this quotation's type in the type map (for CodeGen to use later)
388                self.quotation_types
389                    .borrow_mut()
390                    .insert(*id, quot_type.clone());
391
392                // If this is a closure, we need to pop the captured values from the stack
393                let result_stack = match &quot_type {
394                    Type::Quotation(_) => {
395                        // Stateless - no captures, just push quotation onto stack
396                        current_stack.push(quot_type)
397                    }
398                    Type::Closure { captures, .. } => {
399                        // Pop captured values from stack, then push closure
400                        let mut stack = current_stack.clone();
401                        for _ in 0..captures.len() {
402                            let (new_stack, _value) = self.pop_type(&stack, "closure capture")?;
403                            stack = new_stack;
404                        }
405                        stack.push(quot_type)
406                    }
407                    _ => unreachable!("analyze_captures only returns Quotation or Closure"),
408                };
409
410                Ok((result_stack, Subst::empty()))
411            }
412        }
413    }
414
415    /// Look up the effect of a word (built-in or user-defined)
416    fn lookup_word_effect(&self, name: &str) -> Option<Effect> {
417        // First check built-ins
418        if let Some(effect) = builtin_signature(name) {
419            return Some(effect);
420        }
421
422        // Then check user-defined words
423        self.env.get(name).cloned()
424    }
425
426    /// Apply an effect to a stack
427    /// Effect: (inputs -- outputs)
428    /// Current stack must match inputs, result is outputs
429    /// Returns (result_stack, substitution)
430    fn apply_effect(
431        &self,
432        effect: &Effect,
433        current_stack: StackType,
434        operation: &str,
435    ) -> Result<(StackType, Subst), String> {
436        // Unify current stack with effect's input
437        let subst = unify_stacks(&effect.inputs, &current_stack).map_err(|e| {
438            format!(
439                "{}: stack type mismatch. Expected {:?}, got {:?}: {}",
440                operation, effect.inputs, current_stack, e
441            )
442        })?;
443
444        // Apply substitution to output
445        let result_stack = subst.apply_stack(&effect.outputs);
446
447        Ok((result_stack, subst))
448    }
449
450    /// Adjust stack for spawn operation by converting Quotation to Closure if needed
451    ///
452    /// spawn expects Quotation(Empty -- Empty), but if we have Quotation(T... -- U...)
453    /// with non-empty inputs, we auto-convert it to a Closure that captures those inputs.
454    fn adjust_stack_for_spawn(
455        &self,
456        current_stack: StackType,
457        spawn_effect: &Effect,
458    ) -> Result<StackType, String> {
459        // spawn expects: ( ..a Quotation(Empty -- Empty) -- ..a Int )
460        // Extract the expected quotation type from spawn's effect
461        let expected_quot_type = match &spawn_effect.inputs {
462            StackType::Cons { top, rest: _ } => {
463                if !matches!(top, Type::Quotation(_)) {
464                    return Ok(current_stack); // Not a quotation, don't adjust
465                }
466                top
467            }
468            _ => return Ok(current_stack),
469        };
470
471        // Check what's actually on the stack
472        let (rest_stack, actual_type) = match &current_stack {
473            StackType::Cons { rest, top } => (rest.as_ref().clone(), top),
474            _ => return Ok(current_stack), // Empty stack, nothing to adjust
475        };
476
477        // If top of stack is a Quotation with non-empty inputs, convert to Closure
478        if let Type::Quotation(actual_effect) = actual_type {
479            // Check if quotation needs inputs
480            if !matches!(actual_effect.inputs, StackType::Empty) {
481                // Extract expected effect from spawn's signature
482                let expected_effect = match expected_quot_type {
483                    Type::Quotation(eff) => eff.as_ref(),
484                    _ => return Ok(current_stack),
485                };
486
487                // Calculate what needs to be captured
488                let captures = self.calculate_captures(actual_effect, expected_effect)?;
489
490                // Create a Closure type
491                let closure_type = Type::Closure {
492                    effect: Box::new(expected_effect.clone()),
493                    captures: captures.clone(),
494                };
495
496                // Pop the captured values from the stack
497                // The values to capture are BELOW the quotation on the stack
498                let mut adjusted_stack = rest_stack;
499                for _ in &captures {
500                    adjusted_stack = match adjusted_stack {
501                        StackType::Cons { rest, .. } => rest.as_ref().clone(),
502                        _ => {
503                            return Err(format!(
504                                "spawn: not enough values on stack to capture. Need {} values",
505                                captures.len()
506                            ));
507                        }
508                    };
509                }
510
511                // Push the Closure onto the adjusted stack
512                return Ok(adjusted_stack.push(closure_type));
513            }
514        }
515
516        Ok(current_stack)
517    }
518
519    /// Analyze quotation captures
520    ///
521    /// Determines whether a quotation should be stateless (Type::Quotation)
522    /// or a closure (Type::Closure) based on the expected type from the word signature.
523    ///
524    /// Type-driven inference with automatic closure creation:
525    ///   - If expected type is Closure[effect], calculate what to capture
526    ///   - If expected type is Quotation[effect]:
527    ///     - If body needs more inputs than expected effect, auto-create Closure
528    ///     - Otherwise return stateless Quotation
529    ///   - If no expected type, default to stateless (conservative)
530    ///
531    /// Example 1 (auto-create closure):
532    ///   Expected: Quotation[-- ]          [spawn expects ( -- )]
533    ///   Body: [ handle-connection ]       [needs ( Int -- )]
534    ///   Body effect: ( Int -- )           [needs 1 Int]
535    ///   Expected effect: ( -- )           [provides 0 inputs]
536    ///   Result: Closure { effect: ( -- ), captures: [Int] }
537    ///
538    /// Example 2 (explicit closure):
539    ///   Signature: ( Int -- Closure[Int -- Int] )
540    ///   Body: [ add ]
541    ///   Body effect: ( Int Int -- Int )  [add needs 2 Ints]
542    ///   Expected effect: [Int -- Int]    [call site provides 1 Int]
543    ///   Result: Closure { effect: [Int -- Int], captures: [Int] }
544    fn analyze_captures(
545        &self,
546        body_effect: &Effect,
547        _current_stack: &StackType,
548    ) -> Result<Type, String> {
549        // Check if there's an expected type from the word signature
550        let expected = self.expected_quotation_type.borrow().clone();
551
552        match expected {
553            Some(Type::Closure { effect, .. }) => {
554                // User declared closure type - calculate captures
555                let captures = self.calculate_captures(body_effect, &effect)?;
556                Ok(Type::Closure { effect, captures })
557            }
558            Some(Type::Quotation(expected_effect)) => {
559                // User declared quotation type - check if we need to auto-create closure
560                // Auto-create closure only when:
561                // 1. Expected effect has empty inputs (like spawn's ( -- ))
562                // 2. Body effect has non-empty inputs (needs values to execute)
563
564                let expected_is_empty = matches!(expected_effect.inputs, StackType::Empty);
565                let body_needs_inputs = !matches!(body_effect.inputs, StackType::Empty);
566
567                if expected_is_empty && body_needs_inputs {
568                    // Body needs inputs but expected provides none
569                    // Auto-create closure to capture the inputs
570                    let captures = self.calculate_captures(body_effect, &expected_effect)?;
571                    Ok(Type::Closure {
572                        effect: expected_effect,
573                        captures,
574                    })
575                } else {
576                    // Body is compatible with expected effect - stateless quotation
577                    Ok(Type::Quotation(expected_effect))
578                }
579            }
580            _ => {
581                // No expected type - conservative default: stateless quotation
582                Ok(Type::Quotation(Box::new(body_effect.clone())))
583            }
584        }
585    }
586
587    /// Calculate capture types for a closure
588    ///
589    /// Given:
590    ///   - body_effect: what the quotation body needs (e.g., Int Int -- Int)
591    ///   - call_effect: what the call site will provide (e.g., Int -- Int)
592    ///
593    /// Calculate:
594    ///   - captures: types to capture from creation stack (e.g., [Int])
595    ///
596    /// Example:
597    ///   Body needs: (Int Int -- Int)  [2 inputs total]
598    ///   Call provides: (Int -- Int)   [1 input from call site]
599    ///   Captures: 2 - 1 = 1 Int       [1 input from creation site]
600    ///
601    /// Capture Ordering:
602    ///   Captures are returned bottom-to-top (deepest value first).
603    ///   This matches how push_closure pops from the stack:
604    ///   - Stack at creation: ( ...rest bottom top )
605    ///   - push_closure pops top-down: pop top, pop bottom
606    ///   - Stores as: env[0]=top, env[1]=bottom (reversed)
607    ///   - Closure function pushes: push env[0], push env[1]
608    ///   - Result: bottom is deeper, top is shallower (correct order)
609    fn calculate_captures(
610        &self,
611        body_effect: &Effect,
612        call_effect: &Effect,
613    ) -> Result<Vec<Type>, String> {
614        // Extract concrete types from stack types (bottom to top)
615        let body_inputs = self.extract_concrete_types(&body_effect.inputs);
616        let call_inputs = self.extract_concrete_types(&call_effect.inputs);
617
618        // Validate: call site shouldn't provide MORE than body needs
619        if call_inputs.len() > body_inputs.len() {
620            return Err(format!(
621                "Closure signature error: call site provides {} values but body only needs {}",
622                call_inputs.len(),
623                body_inputs.len()
624            ));
625        }
626
627        // Calculate how many to capture (from bottom of stack)
628        let capture_count = body_inputs.len() - call_inputs.len();
629
630        // Captures are the first N types (bottom of stack)
631        // Example: body needs [Int, String] (bottom to top), call provides [String]
632        // Captures: [Int] (the bottom type)
633        Ok(body_inputs[0..capture_count].to_vec())
634    }
635
636    /// Extract concrete types from a stack type (bottom to top order)
637    ///
638    /// Example:
639    ///   Input: Cons { rest: Cons { rest: Empty, top: Int }, top: String }
640    ///   Output: [Int, String]  (bottom to top)
641    ///
642    /// Row variables are not supported - this is only for concrete stacks
643    fn extract_concrete_types(&self, stack: &StackType) -> Vec<Type> {
644        let mut types = Vec::new();
645        let mut current = stack.clone();
646
647        // Pop types from top to bottom
648        while let Some((rest, top)) = current.pop() {
649            types.push(top);
650            current = rest;
651        }
652
653        // Reverse to get bottom-to-top order
654        types.reverse();
655        types
656    }
657
658    /// Pop a type from a stack type, returning (rest, top)
659    fn pop_type(&self, stack: &StackType, context: &str) -> Result<(StackType, Type), String> {
660        match stack {
661            StackType::Cons { rest, top } => Ok(((**rest).clone(), top.clone())),
662            StackType::Empty => Err(format!(
663                "{}: stack underflow - expected value on stack but stack is empty",
664                context
665            )),
666            StackType::RowVar(_) => {
667                // Can't statically determine if row variable is empty
668                // For now, assume it has at least one element
669                // This is conservative - real implementation would track constraints
670                Err(format!(
671                    "{}: cannot pop from polymorphic stack without more type information",
672                    context
673                ))
674            }
675        }
676    }
677}
678
679impl Default for TypeChecker {
680    fn default() -> Self {
681        Self::new()
682    }
683}
684
685#[cfg(test)]
686mod tests {
687    use super::*;
688
689    #[test]
690    fn test_simple_literal() {
691        let program = Program {
692            includes: vec![],
693            words: vec![WordDef {
694                name: "test".to_string(),
695                effect: Some(Effect::new(
696                    StackType::Empty,
697                    StackType::singleton(Type::Int),
698                )),
699                body: vec![Statement::IntLiteral(42)],
700                source: None,
701            }],
702        };
703
704        let mut checker = TypeChecker::new();
705        assert!(checker.check_program(&program).is_ok());
706    }
707
708    #[test]
709    fn test_simple_operation() {
710        // : test ( Int Int -- Int ) add ;
711        let program = Program {
712            includes: vec![],
713            words: vec![WordDef {
714                name: "test".to_string(),
715                effect: Some(Effect::new(
716                    StackType::Empty.push(Type::Int).push(Type::Int),
717                    StackType::singleton(Type::Int),
718                )),
719                body: vec![Statement::WordCall("add".to_string())],
720                source: None,
721            }],
722        };
723
724        let mut checker = TypeChecker::new();
725        assert!(checker.check_program(&program).is_ok());
726    }
727
728    #[test]
729    fn test_type_mismatch() {
730        // : test ( String -- ) write_line ;  with body: 42
731        let program = Program {
732            includes: vec![],
733            words: vec![WordDef {
734                name: "test".to_string(),
735                effect: Some(Effect::new(
736                    StackType::singleton(Type::String),
737                    StackType::Empty,
738                )),
739                body: vec![
740                    Statement::IntLiteral(42), // Pushes Int, not String!
741                    Statement::WordCall("write_line".to_string()),
742                ],
743                source: None,
744            }],
745        };
746
747        let mut checker = TypeChecker::new();
748        let result = checker.check_program(&program);
749        assert!(result.is_err());
750        assert!(result.unwrap_err().contains("Type mismatch"));
751    }
752
753    #[test]
754    fn test_polymorphic_dup() {
755        // : my-dup ( Int -- Int Int ) dup ;
756        let program = Program {
757            includes: vec![],
758            words: vec![WordDef {
759                name: "my-dup".to_string(),
760                effect: Some(Effect::new(
761                    StackType::singleton(Type::Int),
762                    StackType::Empty.push(Type::Int).push(Type::Int),
763                )),
764                body: vec![Statement::WordCall("dup".to_string())],
765                source: None,
766            }],
767        };
768
769        let mut checker = TypeChecker::new();
770        assert!(checker.check_program(&program).is_ok());
771    }
772
773    #[test]
774    fn test_conditional_branches() {
775        // : test ( Int Int -- String )
776        //   > if "greater" else "not greater" then ;
777        let program = Program {
778            includes: vec![],
779            words: vec![WordDef {
780                name: "test".to_string(),
781                effect: Some(Effect::new(
782                    StackType::Empty.push(Type::Int).push(Type::Int),
783                    StackType::singleton(Type::String),
784                )),
785                body: vec![
786                    Statement::WordCall(">".to_string()),
787                    Statement::If {
788                        then_branch: vec![Statement::StringLiteral("greater".to_string())],
789                        else_branch: Some(vec![Statement::StringLiteral(
790                            "not greater".to_string(),
791                        )]),
792                    },
793                ],
794                source: None,
795            }],
796        };
797
798        let mut checker = TypeChecker::new();
799        assert!(checker.check_program(&program).is_ok());
800    }
801
802    #[test]
803    fn test_mismatched_branches() {
804        // : test ( Int -- ? )
805        //   if 42 else "string" then ;  // ERROR: incompatible types
806        let program = Program {
807            includes: vec![],
808            words: vec![WordDef {
809                name: "test".to_string(),
810                effect: None,
811                body: vec![
812                    Statement::IntLiteral(1),
813                    Statement::If {
814                        then_branch: vec![Statement::IntLiteral(42)],
815                        else_branch: Some(vec![Statement::StringLiteral("string".to_string())]),
816                    },
817                ],
818                source: None,
819            }],
820        };
821
822        let mut checker = TypeChecker::new();
823        let result = checker.check_program(&program);
824        assert!(result.is_err());
825        assert!(result.unwrap_err().contains("incompatible"));
826    }
827
828    #[test]
829    fn test_user_defined_word_call() {
830        // : helper ( Int -- String ) int->string ;
831        // : main ( -- ) 42 helper write_line ;
832        let program = Program {
833            includes: vec![],
834            words: vec![
835                WordDef {
836                    name: "helper".to_string(),
837                    effect: Some(Effect::new(
838                        StackType::singleton(Type::Int),
839                        StackType::singleton(Type::String),
840                    )),
841                    body: vec![Statement::WordCall("int->string".to_string())],
842                    source: None,
843                },
844                WordDef {
845                    name: "main".to_string(),
846                    effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
847                    body: vec![
848                        Statement::IntLiteral(42),
849                        Statement::WordCall("helper".to_string()),
850                        Statement::WordCall("write_line".to_string()),
851                    ],
852                    source: None,
853                },
854            ],
855        };
856
857        let mut checker = TypeChecker::new();
858        assert!(checker.check_program(&program).is_ok());
859    }
860
861    #[test]
862    fn test_arithmetic_chain() {
863        // : test ( Int Int Int -- Int )
864        //   add multiply ;
865        let program = Program {
866            includes: vec![],
867            words: vec![WordDef {
868                name: "test".to_string(),
869                effect: Some(Effect::new(
870                    StackType::Empty
871                        .push(Type::Int)
872                        .push(Type::Int)
873                        .push(Type::Int),
874                    StackType::singleton(Type::Int),
875                )),
876                body: vec![
877                    Statement::WordCall("add".to_string()),
878                    Statement::WordCall("multiply".to_string()),
879                ],
880                source: None,
881            }],
882        };
883
884        let mut checker = TypeChecker::new();
885        assert!(checker.check_program(&program).is_ok());
886    }
887
888    #[test]
889    fn test_write_line_type_error() {
890        // : test ( Int -- ) write_line ;  // ERROR: write_line expects String
891        let program = Program {
892            includes: vec![],
893            words: vec![WordDef {
894                name: "test".to_string(),
895                effect: Some(Effect::new(
896                    StackType::singleton(Type::Int),
897                    StackType::Empty,
898                )),
899                body: vec![Statement::WordCall("write_line".to_string())],
900                source: None,
901            }],
902        };
903
904        let mut checker = TypeChecker::new();
905        let result = checker.check_program(&program);
906        assert!(result.is_err());
907        assert!(result.unwrap_err().contains("Type mismatch"));
908    }
909
910    #[test]
911    fn test_stack_underflow_drop() {
912        // : test ( -- ) drop ;  // ERROR: can't drop from empty stack
913        let program = Program {
914            includes: vec![],
915            words: vec![WordDef {
916                name: "test".to_string(),
917                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
918                body: vec![Statement::WordCall("drop".to_string())],
919                source: None,
920            }],
921        };
922
923        let mut checker = TypeChecker::new();
924        let result = checker.check_program(&program);
925        assert!(result.is_err());
926        assert!(result.unwrap_err().contains("mismatch"));
927    }
928
929    #[test]
930    fn test_stack_underflow_add() {
931        // : test ( Int -- Int ) add ;  // ERROR: add needs 2 values
932        let program = Program {
933            includes: vec![],
934            words: vec![WordDef {
935                name: "test".to_string(),
936                effect: Some(Effect::new(
937                    StackType::singleton(Type::Int),
938                    StackType::singleton(Type::Int),
939                )),
940                body: vec![Statement::WordCall("add".to_string())],
941                source: None,
942            }],
943        };
944
945        let mut checker = TypeChecker::new();
946        let result = checker.check_program(&program);
947        assert!(result.is_err());
948        assert!(result.unwrap_err().contains("mismatch"));
949    }
950
951    #[test]
952    fn test_csp_operations() {
953        // : test ( -- )
954        //   make-channel  # ( -- Int )
955        //   42 swap       # ( Int Int -- Int Int )
956        //   send          # ( Int Int -- )
957        // ;
958        let program = Program {
959            includes: vec![],
960            words: vec![WordDef {
961                name: "test".to_string(),
962                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
963                body: vec![
964                    Statement::WordCall("make-channel".to_string()),
965                    Statement::IntLiteral(42),
966                    Statement::WordCall("swap".to_string()),
967                    Statement::WordCall("send".to_string()),
968                ],
969                source: None,
970            }],
971        };
972
973        let mut checker = TypeChecker::new();
974        assert!(checker.check_program(&program).is_ok());
975    }
976
977    #[test]
978    fn test_complex_stack_shuffling() {
979        // : test ( Int Int Int -- Int )
980        //   rot add add ;
981        let program = Program {
982            includes: vec![],
983            words: vec![WordDef {
984                name: "test".to_string(),
985                effect: Some(Effect::new(
986                    StackType::Empty
987                        .push(Type::Int)
988                        .push(Type::Int)
989                        .push(Type::Int),
990                    StackType::singleton(Type::Int),
991                )),
992                body: vec![
993                    Statement::WordCall("rot".to_string()),
994                    Statement::WordCall("add".to_string()),
995                    Statement::WordCall("add".to_string()),
996                ],
997                source: None,
998            }],
999        };
1000
1001        let mut checker = TypeChecker::new();
1002        assert!(checker.check_program(&program).is_ok());
1003    }
1004
1005    #[test]
1006    fn test_empty_program() {
1007        // Program with no words should be valid
1008        let program = Program {
1009            includes: vec![],
1010            words: vec![],
1011        };
1012
1013        let mut checker = TypeChecker::new();
1014        assert!(checker.check_program(&program).is_ok());
1015    }
1016
1017    #[test]
1018    fn test_word_without_effect_declaration() {
1019        // : helper 42 ;  // No effect declaration
1020        let program = Program {
1021            includes: vec![],
1022            words: vec![WordDef {
1023                name: "helper".to_string(),
1024                effect: None,
1025                body: vec![Statement::IntLiteral(42)],
1026                source: None,
1027            }],
1028        };
1029
1030        let mut checker = TypeChecker::new();
1031        assert!(checker.check_program(&program).is_ok());
1032    }
1033
1034    #[test]
1035    fn test_nested_conditionals() {
1036        // : test ( Int Int Int Int -- String )
1037        //   > if
1038        //     > if "both true" else "first true" then
1039        //   else
1040        //     drop drop "first false"
1041        //   then ;
1042        // Note: Needs 4 Ints total (2 for each > comparison)
1043        // Else branch must drop unused Ints to match then branch's stack effect
1044        let program = Program {
1045            includes: vec![],
1046            words: vec![WordDef {
1047                name: "test".to_string(),
1048                effect: Some(Effect::new(
1049                    StackType::Empty
1050                        .push(Type::Int)
1051                        .push(Type::Int)
1052                        .push(Type::Int)
1053                        .push(Type::Int),
1054                    StackType::singleton(Type::String),
1055                )),
1056                body: vec![
1057                    Statement::WordCall(">".to_string()),
1058                    Statement::If {
1059                        then_branch: vec![
1060                            Statement::WordCall(">".to_string()),
1061                            Statement::If {
1062                                then_branch: vec![Statement::StringLiteral(
1063                                    "both true".to_string(),
1064                                )],
1065                                else_branch: Some(vec![Statement::StringLiteral(
1066                                    "first true".to_string(),
1067                                )]),
1068                            },
1069                        ],
1070                        else_branch: Some(vec![
1071                            Statement::WordCall("drop".to_string()),
1072                            Statement::WordCall("drop".to_string()),
1073                            Statement::StringLiteral("first false".to_string()),
1074                        ]),
1075                    },
1076                ],
1077                source: None,
1078            }],
1079        };
1080
1081        let mut checker = TypeChecker::new();
1082        match checker.check_program(&program) {
1083            Ok(_) => {}
1084            Err(e) => panic!("Type check failed: {}", e),
1085        }
1086    }
1087
1088    #[test]
1089    fn test_conditional_without_else() {
1090        // : test ( Int Int -- Int )
1091        //   > if 100 then ;
1092        // Both branches must leave same stack
1093        let program = Program {
1094            includes: vec![],
1095            words: vec![WordDef {
1096                name: "test".to_string(),
1097                effect: Some(Effect::new(
1098                    StackType::Empty.push(Type::Int).push(Type::Int),
1099                    StackType::singleton(Type::Int),
1100                )),
1101                body: vec![
1102                    Statement::WordCall(">".to_string()),
1103                    Statement::If {
1104                        then_branch: vec![Statement::IntLiteral(100)],
1105                        else_branch: None, // No else - should leave stack unchanged
1106                    },
1107                ],
1108                source: None,
1109            }],
1110        };
1111
1112        let mut checker = TypeChecker::new();
1113        let result = checker.check_program(&program);
1114        // This should fail because then pushes Int but else leaves stack empty
1115        assert!(result.is_err());
1116    }
1117
1118    #[test]
1119    fn test_multiple_word_chain() {
1120        // : helper1 ( Int -- String ) int->string ;
1121        // : helper2 ( String -- ) write_line ;
1122        // : main ( -- ) 42 helper1 helper2 ;
1123        let program = Program {
1124            includes: vec![],
1125            words: vec![
1126                WordDef {
1127                    name: "helper1".to_string(),
1128                    effect: Some(Effect::new(
1129                        StackType::singleton(Type::Int),
1130                        StackType::singleton(Type::String),
1131                    )),
1132                    body: vec![Statement::WordCall("int->string".to_string())],
1133                    source: None,
1134                },
1135                WordDef {
1136                    name: "helper2".to_string(),
1137                    effect: Some(Effect::new(
1138                        StackType::singleton(Type::String),
1139                        StackType::Empty,
1140                    )),
1141                    body: vec![Statement::WordCall("write_line".to_string())],
1142                    source: None,
1143                },
1144                WordDef {
1145                    name: "main".to_string(),
1146                    effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
1147                    body: vec![
1148                        Statement::IntLiteral(42),
1149                        Statement::WordCall("helper1".to_string()),
1150                        Statement::WordCall("helper2".to_string()),
1151                    ],
1152                    source: None,
1153                },
1154            ],
1155        };
1156
1157        let mut checker = TypeChecker::new();
1158        assert!(checker.check_program(&program).is_ok());
1159    }
1160
1161    #[test]
1162    fn test_all_stack_ops() {
1163        // : test ( Int Int Int -- Int Int Int Int )
1164        //   over nip tuck ;
1165        let program = Program {
1166            includes: vec![],
1167            words: vec![WordDef {
1168                name: "test".to_string(),
1169                effect: Some(Effect::new(
1170                    StackType::Empty
1171                        .push(Type::Int)
1172                        .push(Type::Int)
1173                        .push(Type::Int),
1174                    StackType::Empty
1175                        .push(Type::Int)
1176                        .push(Type::Int)
1177                        .push(Type::Int)
1178                        .push(Type::Int),
1179                )),
1180                body: vec![
1181                    Statement::WordCall("over".to_string()),
1182                    Statement::WordCall("nip".to_string()),
1183                    Statement::WordCall("tuck".to_string()),
1184                ],
1185                source: None,
1186            }],
1187        };
1188
1189        let mut checker = TypeChecker::new();
1190        assert!(checker.check_program(&program).is_ok());
1191    }
1192
1193    #[test]
1194    fn test_mixed_types_complex() {
1195        // : test ( -- )
1196        //   42 int->string      # ( -- String )
1197        //   100 200 >           # ( String -- String Int )
1198        //   if                  # ( String -- String )
1199        //     write_line        # ( String -- )
1200        //   else
1201        //     write_line
1202        //   then ;
1203        let program = Program {
1204            includes: vec![],
1205            words: vec![WordDef {
1206                name: "test".to_string(),
1207                effect: Some(Effect::new(StackType::Empty, StackType::Empty)),
1208                body: vec![
1209                    Statement::IntLiteral(42),
1210                    Statement::WordCall("int->string".to_string()),
1211                    Statement::IntLiteral(100),
1212                    Statement::IntLiteral(200),
1213                    Statement::WordCall(">".to_string()),
1214                    Statement::If {
1215                        then_branch: vec![Statement::WordCall("write_line".to_string())],
1216                        else_branch: Some(vec![Statement::WordCall("write_line".to_string())]),
1217                    },
1218                ],
1219                source: None,
1220            }],
1221        };
1222
1223        let mut checker = TypeChecker::new();
1224        assert!(checker.check_program(&program).is_ok());
1225    }
1226
1227    #[test]
1228    fn test_string_literal() {
1229        // : test ( -- String ) "hello" ;
1230        let program = Program {
1231            includes: vec![],
1232            words: vec![WordDef {
1233                name: "test".to_string(),
1234                effect: Some(Effect::new(
1235                    StackType::Empty,
1236                    StackType::singleton(Type::String),
1237                )),
1238                body: vec![Statement::StringLiteral("hello".to_string())],
1239                source: None,
1240            }],
1241        };
1242
1243        let mut checker = TypeChecker::new();
1244        assert!(checker.check_program(&program).is_ok());
1245    }
1246
1247    #[test]
1248    fn test_bool_literal() {
1249        // : test ( -- Int ) true ;
1250        // Booleans are represented as Int in the type system
1251        let program = Program {
1252            includes: vec![],
1253            words: vec![WordDef {
1254                name: "test".to_string(),
1255                effect: Some(Effect::new(
1256                    StackType::Empty,
1257                    StackType::singleton(Type::Int),
1258                )),
1259                body: vec![Statement::BoolLiteral(true)],
1260                source: None,
1261            }],
1262        };
1263
1264        let mut checker = TypeChecker::new();
1265        assert!(checker.check_program(&program).is_ok());
1266    }
1267
1268    #[test]
1269    fn test_type_error_in_nested_conditional() {
1270        // : test ( Int Int -- ? )
1271        //   > if
1272        //     42 write_line   # ERROR: write_line expects String, got Int
1273        //   else
1274        //     "ok" write_line
1275        //   then ;
1276        let program = Program {
1277            includes: vec![],
1278            words: vec![WordDef {
1279                name: "test".to_string(),
1280                effect: None,
1281                body: vec![
1282                    Statement::IntLiteral(10),
1283                    Statement::IntLiteral(20),
1284                    Statement::WordCall(">".to_string()),
1285                    Statement::If {
1286                        then_branch: vec![
1287                            Statement::IntLiteral(42),
1288                            Statement::WordCall("write_line".to_string()),
1289                        ],
1290                        else_branch: Some(vec![
1291                            Statement::StringLiteral("ok".to_string()),
1292                            Statement::WordCall("write_line".to_string()),
1293                        ]),
1294                    },
1295                ],
1296                source: None,
1297            }],
1298        };
1299
1300        let mut checker = TypeChecker::new();
1301        let result = checker.check_program(&program);
1302        assert!(result.is_err());
1303        assert!(result.unwrap_err().contains("Type mismatch"));
1304    }
1305
1306    #[test]
1307    fn test_read_line_operation() {
1308        // : test ( -- String ) read_line ;
1309        let program = Program {
1310            includes: vec![],
1311            words: vec![WordDef {
1312                name: "test".to_string(),
1313                effect: Some(Effect::new(
1314                    StackType::Empty,
1315                    StackType::singleton(Type::String),
1316                )),
1317                body: vec![Statement::WordCall("read_line".to_string())],
1318                source: None,
1319            }],
1320        };
1321
1322        let mut checker = TypeChecker::new();
1323        assert!(checker.check_program(&program).is_ok());
1324    }
1325
1326    #[test]
1327    fn test_comparison_operations() {
1328        // Test all comparison operators
1329        // : test ( Int Int -- Int Int Int Int Int Int )
1330        //   2dup = 2dup < 2dup > 2dup <= 2dup >= 2dup <> ;
1331        // Simplified: just test that comparisons work
1332        let program = Program {
1333            includes: vec![],
1334            words: vec![WordDef {
1335                name: "test".to_string(),
1336                effect: Some(Effect::new(
1337                    StackType::Empty.push(Type::Int).push(Type::Int),
1338                    StackType::singleton(Type::Int),
1339                )),
1340                body: vec![Statement::WordCall("<=".to_string())],
1341                source: None,
1342            }],
1343        };
1344
1345        let mut checker = TypeChecker::new();
1346        assert!(checker.check_program(&program).is_ok());
1347    }
1348
1349    #[test]
1350    fn test_recursive_word_definitions() {
1351        // Test mutually recursive words (type checking only, no runtime)
1352        // : is-even ( Int -- Int ) dup 0 = if drop 1 else 1 subtract is-odd then ;
1353        // : is-odd ( Int -- Int ) dup 0 = if drop 0 else 1 subtract is-even then ;
1354        //
1355        // Note: This tests that the checker can handle words that reference each other
1356        let program = Program {
1357            includes: vec![],
1358            words: vec![
1359                WordDef {
1360                    name: "is-even".to_string(),
1361                    effect: Some(Effect::new(
1362                        StackType::singleton(Type::Int),
1363                        StackType::singleton(Type::Int),
1364                    )),
1365                    body: vec![
1366                        Statement::WordCall("dup".to_string()),
1367                        Statement::IntLiteral(0),
1368                        Statement::WordCall("=".to_string()),
1369                        Statement::If {
1370                            then_branch: vec![
1371                                Statement::WordCall("drop".to_string()),
1372                                Statement::IntLiteral(1),
1373                            ],
1374                            else_branch: Some(vec![
1375                                Statement::IntLiteral(1),
1376                                Statement::WordCall("subtract".to_string()),
1377                                Statement::WordCall("is-odd".to_string()),
1378                            ]),
1379                        },
1380                    ],
1381                    source: None,
1382                },
1383                WordDef {
1384                    name: "is-odd".to_string(),
1385                    effect: Some(Effect::new(
1386                        StackType::singleton(Type::Int),
1387                        StackType::singleton(Type::Int),
1388                    )),
1389                    body: vec![
1390                        Statement::WordCall("dup".to_string()),
1391                        Statement::IntLiteral(0),
1392                        Statement::WordCall("=".to_string()),
1393                        Statement::If {
1394                            then_branch: vec![
1395                                Statement::WordCall("drop".to_string()),
1396                                Statement::IntLiteral(0),
1397                            ],
1398                            else_branch: Some(vec![
1399                                Statement::IntLiteral(1),
1400                                Statement::WordCall("subtract".to_string()),
1401                                Statement::WordCall("is-even".to_string()),
1402                            ]),
1403                        },
1404                    ],
1405                    source: None,
1406                },
1407            ],
1408        };
1409
1410        let mut checker = TypeChecker::new();
1411        assert!(checker.check_program(&program).is_ok());
1412    }
1413
1414    #[test]
1415    fn test_word_calling_word_with_row_polymorphism() {
1416        // Test that row variables unify correctly through word calls
1417        // : apply-twice ( Int -- Int ) dup add ;
1418        // : quad ( Int -- Int ) apply-twice apply-twice ;
1419        // Should work: both use row polymorphism correctly
1420        let program = Program {
1421            includes: vec![],
1422            words: vec![
1423                WordDef {
1424                    name: "apply-twice".to_string(),
1425                    effect: Some(Effect::new(
1426                        StackType::singleton(Type::Int),
1427                        StackType::singleton(Type::Int),
1428                    )),
1429                    body: vec![
1430                        Statement::WordCall("dup".to_string()),
1431                        Statement::WordCall("add".to_string()),
1432                    ],
1433                    source: None,
1434                },
1435                WordDef {
1436                    name: "quad".to_string(),
1437                    effect: Some(Effect::new(
1438                        StackType::singleton(Type::Int),
1439                        StackType::singleton(Type::Int),
1440                    )),
1441                    body: vec![
1442                        Statement::WordCall("apply-twice".to_string()),
1443                        Statement::WordCall("apply-twice".to_string()),
1444                    ],
1445                    source: None,
1446                },
1447            ],
1448        };
1449
1450        let mut checker = TypeChecker::new();
1451        assert!(checker.check_program(&program).is_ok());
1452    }
1453
1454    #[test]
1455    fn test_deep_stack_types() {
1456        // Test with many values on stack (10+ items)
1457        // : test ( Int Int Int Int Int Int Int Int Int Int -- Int )
1458        //   add add add add add add add add add ;
1459        let mut stack_type = StackType::Empty;
1460        for _ in 0..10 {
1461            stack_type = stack_type.push(Type::Int);
1462        }
1463
1464        let program = Program {
1465            includes: vec![],
1466            words: vec![WordDef {
1467                name: "test".to_string(),
1468                effect: Some(Effect::new(stack_type, StackType::singleton(Type::Int))),
1469                body: vec![
1470                    Statement::WordCall("add".to_string()),
1471                    Statement::WordCall("add".to_string()),
1472                    Statement::WordCall("add".to_string()),
1473                    Statement::WordCall("add".to_string()),
1474                    Statement::WordCall("add".to_string()),
1475                    Statement::WordCall("add".to_string()),
1476                    Statement::WordCall("add".to_string()),
1477                    Statement::WordCall("add".to_string()),
1478                    Statement::WordCall("add".to_string()),
1479                ],
1480                source: None,
1481            }],
1482        };
1483
1484        let mut checker = TypeChecker::new();
1485        assert!(checker.check_program(&program).is_ok());
1486    }
1487
1488    #[test]
1489    fn test_simple_quotation() {
1490        // : test ( -- Quot )
1491        //   [ 1 add ] ;
1492        // Quotation type should be [ ..input Int -- ..input Int ] (row polymorphic)
1493        let program = Program {
1494            includes: vec![],
1495            words: vec![WordDef {
1496                name: "test".to_string(),
1497                effect: Some(Effect::new(
1498                    StackType::Empty,
1499                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
1500                        StackType::RowVar("input".to_string()).push(Type::Int),
1501                        StackType::RowVar("input".to_string()).push(Type::Int),
1502                    )))),
1503                )),
1504                body: vec![Statement::Quotation {
1505                    id: 0,
1506                    body: vec![
1507                        Statement::IntLiteral(1),
1508                        Statement::WordCall("add".to_string()),
1509                    ],
1510                }],
1511                source: None,
1512            }],
1513        };
1514
1515        let mut checker = TypeChecker::new();
1516        match checker.check_program(&program) {
1517            Ok(_) => {}
1518            Err(e) => panic!("Type check failed: {}", e),
1519        }
1520    }
1521
1522    #[test]
1523    fn test_empty_quotation() {
1524        // : test ( -- Quot )
1525        //   [ ] ;
1526        // Empty quotation has effect ( ..input -- ..input ) (preserves stack)
1527        let program = Program {
1528            includes: vec![],
1529            words: vec![WordDef {
1530                name: "test".to_string(),
1531                effect: Some(Effect::new(
1532                    StackType::Empty,
1533                    StackType::singleton(Type::Quotation(Box::new(Effect::new(
1534                        StackType::RowVar("input".to_string()),
1535                        StackType::RowVar("input".to_string()),
1536                    )))),
1537                )),
1538                body: vec![Statement::Quotation {
1539                    id: 1,
1540                    body: vec![],
1541                }],
1542                source: None,
1543            }],
1544        };
1545
1546        let mut checker = TypeChecker::new();
1547        assert!(checker.check_program(&program).is_ok());
1548    }
1549
1550    // TODO: Re-enable once write_line is properly row-polymorphic
1551    // #[test]
1552    // fn test_quotation_with_string() {
1553    //     // : test ( -- Quot )
1554    //     //   [ "hello" write_line ] ;
1555    //     let program = Program { includes: vec![],
1556    //         words: vec![WordDef {
1557    //             name: "test".to_string(),
1558    //             effect: Some(Effect::new(
1559    //                 StackType::Empty,
1560    //                 StackType::singleton(Type::Quotation(Box::new(Effect::new(
1561    //                     StackType::RowVar("input".to_string()),
1562    //                     StackType::RowVar("input".to_string()),
1563    //                 )))),
1564    //             )),
1565    //             body: vec![Statement::Quotation(vec![
1566    //                 Statement::StringLiteral("hello".to_string()),
1567    //                 Statement::WordCall("write_line".to_string()),
1568    //             ])],
1569    //         }],
1570    //     };
1571    //
1572    //     let mut checker = TypeChecker::new();
1573    //     assert!(checker.check_program(&program).is_ok());
1574    // }
1575
1576    #[test]
1577    fn test_nested_quotation() {
1578        // : test ( -- Quot )
1579        //   [ [ 1 add ] ] ;
1580        // Outer quotation contains inner quotation (both row-polymorphic)
1581        let inner_quot_type = Type::Quotation(Box::new(Effect::new(
1582            StackType::RowVar("input".to_string()).push(Type::Int),
1583            StackType::RowVar("input".to_string()).push(Type::Int),
1584        )));
1585
1586        let outer_quot_type = Type::Quotation(Box::new(Effect::new(
1587            StackType::RowVar("input".to_string()),
1588            StackType::RowVar("input".to_string()).push(inner_quot_type.clone()),
1589        )));
1590
1591        let program = Program {
1592            includes: vec![],
1593            words: vec![WordDef {
1594                name: "test".to_string(),
1595                effect: Some(Effect::new(
1596                    StackType::Empty,
1597                    StackType::singleton(outer_quot_type),
1598                )),
1599                body: vec![Statement::Quotation {
1600                    id: 2,
1601                    body: vec![Statement::Quotation {
1602                        id: 3,
1603                        body: vec![
1604                            Statement::IntLiteral(1),
1605                            Statement::WordCall("add".to_string()),
1606                        ],
1607                    }],
1608                }],
1609                source: None,
1610            }],
1611        };
1612
1613        let mut checker = TypeChecker::new();
1614        assert!(checker.check_program(&program).is_ok());
1615    }
1616}