타입 추론 개선

Iterator Type Inference Improvements

Overview

This document describes the improvements made to iterator type inference in the VAIS type system (vais-types v0.0.1).

Problem

Previously, the type checker had limited support for iterator type inference:

  • Method calls only resolved methods defined directly on structs, not from trait implementations
  • Loop expressions could only infer element types for built-in types (Array, Range)
  • Custom iterator types implementing the Iterator trait were not properly handled
  • The TODO comment // For other types, try to iterate (could be a custom iterator trait later) indicated incomplete implementation

Solution

1. Trait Method Resolution (find_trait_method)

Added a new helper method that resolves methods from trait implementations:

fn find_trait_method(&self, receiver_type: &ResolvedType, method_name: &str) -> Option<TraitMethodSig>

This method:

  • Takes a receiver type and method name
  • Searches through all trait implementations for that type
  • Returns the method signature if found in any implemented trait

2. Iterator Item Type Inference (get_iterator_item_type)

Added a comprehensive method to infer the element type from iterators:

fn get_iterator_item_type(&self, iter_type: &ResolvedType) -> Option<ResolvedType>

This method handles:

  • Built-in iterable types (Array, Range)
  • Types implementing the Iterator trait
  • Types implementing the IntoIterator trait (for conversion to iterators)
  • Recursive resolution for nested iterator types

3. Enhanced MethodCall Expression Handling

Updated the MethodCall expression handler to:

  1. First check for methods defined directly on the struct
  2. If not found, search for methods in trait implementations using find_trait_method
  3. Properly type-check arguments and return the correct return type

4. Improved Loop Expression Handling

Replaced the limited pattern matching for loop iterators with:

  • Call to get_iterator_item_type for comprehensive type inference
  • Proper binding of loop variables with inferred types
  • Warning generation when iterator item type cannot be inferred (instead of silent failure)

Benefits

  1. Full Trait Support: Methods from trait implementations are now properly resolved
  2. Custom Iterators: User-defined types implementing Iterator trait work correctly
  3. Better Type Safety: Loop variables have proper type information
  4. IntoIterator Pattern: Support for types that convert to iterators
  5. Better Error Messages: Warnings when type inference fails

Examples

Trait Method Resolution

W Iterator {
    F next(&self) -> i64
}

S Counter { value: i64 }

X Counter: Iterator {
    F next(&self) -> i64 {
        val := self.value
        self.value = self.value + 1
        val
    }
}

F main() -> i64 {
    c := Counter { value: 0 }
    # Method call now resolves to Iterator trait implementation
    val := c.next()  # Properly inferred as i64
    0
}

Iterator Type Inference in Loops

# Built-in types still work
L i:0..10 { }        # i is i64

# Arrays work
arr := [1, 2, 3]
L x:arr { }          # x is i64

# Custom iterators with trait implementations now work
counter := Counter { value: 0 }
# Note: Direct loop iteration would require IntoIterator
# Manual iteration works with trait methods

Testing

All existing tests pass, and new test file demonstrates the improvements:

  • examples/trait_method_inference_test.vais - Tests trait method resolution
  • All existing iterator tests continue to work correctly

Implementation Details

The implementation is in /Users/sswoo/study/projects/vais/crates/vais-types/src/lib.rs:

  • Lines 2157-2177: find_trait_method helper
  • Lines 2199-2263: get_iterator_item_type helper
  • Lines 1519-1537: Enhanced MethodCall handling
  • Lines 1383-1411: Improved Loop expression handling

Future Enhancements

Potential future improvements:

  1. Support for associated types in trait definitions
  2. Generic iterator implementations (Iterator)
  3. More sophisticated IntoIterator pattern matching
  4. Iterator adapters (map, filter, etc.)
  5. Better handling of Option return types from iterators