타입 추론 개선

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<T>)
  3. More sophisticated IntoIterator pattern matching
  4. Iterator adapters (map, filter, etc.)
  5. Better handling of Option<T> return types from iterators