Patch SPL iterator stubs for generic type propagation

Author: AJenbo

docs/CHANGELOG.md

@@ -26,9 +26,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **Array shape keys with special characters.** Array literal keys containing backslashes, newlines, or other non-identifier characters are now properly quoted and escaped in type display (e.g. `array{'foo\\bar\nbaz': string}` instead of a truncated type string).
 - **Static method calls on class-string unions.** `$variable::method()` where `$variable` holds a union of class-strings now resolves through all possible classes and unions their return types.
 - **Trait `self` return type resolution through inheritance.** When a trait method has return type `self`, calling it on a subclass now correctly resolves to the class that uses the trait (the declaring class), not the calling subclass. `static` and `$this` continue to resolve to the runtime class.
-- **`NoRewindIterator` generic resolution.** Added stub patch so `new NoRewindIterator($generator)` propagates the wrapped iterator's type parameters.
+- **SPL iterator generic type propagation.** `CachingIterator`, `InfiniteIterator`, `LimitIterator`, `NoRewindIterator`, `CallbackFilterIterator`, and `FilterIterator` now propagate the wrapped iterator's generic type parameters. Calling `->current()` or `->key()` on any of these decorators returns the correct element type instead of `mixed`.
+- **`ArrayIterator` constructor generic inference.** `new ArrayIterator($typedArray)` now infers `TKey` and `TValue` from the array argument's generic type, so the resulting iterator preserves key and value types through the chain.
 - **Template param inference from type bounds.** When a class template parameter has a bound like `TIterator as Iterator<TKey, TValue>` and `TIterator` is resolved to a concrete generic type, the nested template params (`TKey`, `TValue`) are now inferred from the concrete type's generic arguments.
-- **`@psalm-param`/`@phpstan-param` priority over `@param`.** When a docblock contained both `@param array $x` and `@psalm-param array<TKey, T> $x`, the parameter type was resolved from whichever tag appeared first in the docblock. Now `@phpstan-param` always takes precedence over `@psalm-param`, which always takes precedence over `@param`, matching PHPStan and Psalm behaviour. This fixes generic constructor inference when stubs use paired `@param`/`@psalm-param` tags.
+- **`@psalm-param`/`@phpstan-param` priority over `@param`.** When a docblock contained both `@param array $x` and `@psalm-param array<TKey, T> $x`, the parameter type was resolved from whichever tag appeared first in the docblock. Now `@phpstan-param` always takes precedence over `@psalm-param`, which always takes precedence over `@param`, matching PHPStan and Psalm behaviour.
 - **`static` return type through first-class callables.** `self::method(...)()`, `static::method(...)()`, `parent::method(...)()`, and `$this->method(...)()` now preserve `static` in the return type when the underlying method declares `@return static`.
 - **ArrayAccess array-access assignment.** `$obj[$key] = $val` on objects implementing `ArrayAccess` no longer overwrites the variable's generic type annotation with an array type.
 - **Method-level `@template` with `key-of` bound.** When a method declares `@template K as key-of<TData>` and returns `TData[K]`, passing a string literal argument (e.g. `$config->get('name')`) now infers K from the literal and resolves the return type to the specific array shape value type.

docs/todo/bugs.md

@@ -20,11 +20,12 @@ to second-guess upstream output.
 Remaining failures have multiple root causes (the original
 multi-namespace theory was incorrect for most of them):
 
-- **Lines 16, 29, 41, 56, 68:** SPL iterator stubs
-  (`CachingIterator`, `InfiniteIterator`, `LimitIterator`,
-  `CallbackFilterIterator`, `NoRewindIterator`) lack `@template`
-  annotations, so generic type propagation through iterator
-  decorator constructors is impossible with current stubs.
+- **Lines 16, 29, 41, 56, 68:** `range()` returns bare `array` in
+  phpstorm-stubs (no generic args), so the type information is lost
+  before it reaches the iterator constructors.  The iterator
+  decorator stub patches themselves work correctly when the input
+  array has generic type annotations (see
+  `tests/psalm_assertions/spl_iterator_patches.php`).
 - **Lines 602, 788:** Union generic method resolution and static
   method template inference work correctly in single-namespace
   files. The failures are caused by the multi-namespace test
@@ -42,6 +43,10 @@ multi-namespace theory was incorrect for most of them):
   `extract_param_raw_type_from_info` returned the first match in
   document order instead of respecting `@phpstan-param` >
   `@psalm-param` > `@param` priority.
+- SPL iterator decorator stubs (`CachingIterator`,
+  `InfiniteIterator`, `LimitIterator`, `CallbackFilterIterator`,
+  `FilterIterator`) and `ArrayIterator` constructor patched with
+  `@template` annotations matching PHPStan's stubs.
 
 **Tests:** SKIPs in `tests/psalm_assertions/template_class_template.php`
 (lines 16, 29, 41, 56, 68, 602, 788).
@@ -72,10 +77,10 @@ tests that may now pass. Run
 `cargo nextest run --test assert_type_runner --no-fail-fast` with
 the SKIP removed to verify.
 
-Remaining SKIPs (11) are:
-- `template_class_template.php` (7) — B14: SPL stubs lack
-  @template annotations (5), multi-namespace test runner
-  limitation (2)
+Remaining SKIPs (7) are:
+- `template_class_template.php` (5) — `range()` returns bare
+  `array`, no generic info to propagate through iterator chain;
+  (2) — multi-namespace test runner limitation
 - `magic_method_annotation.php` (3) — B14 cross-namespace
   resolution in single-file test runner
 - `mixin_annotation.php` (1) — `IteratorIterator` not in fixture

src/completion/variable/rhs_resolution.rs

@@ -1450,7 +1450,16 @@ fn resolve_generic_wrapper_template(
         wrapper_name,
         "array" | "list" | "non-empty-array" | "non-empty-list"
     ) {
-        return resolve_array_literal_generic(tpl_position, arg_text, rctx);
+        // Try to infer from array literal first.
+        if let Some(result) = resolve_array_literal_generic(tpl_position, arg_text, rctx) {
+            return Some(result);
+        }
+        // If the argument is not a literal (e.g. a variable), resolve its
+        // type and extract the generic arg at the given position.
+        if let Some(resolved) = Backend::resolve_arg_text_to_type(arg_text, rctx) {
+            return extract_generic_arg_at_position(&resolved, tpl_position);
+        }
+        return None;
     }
 
     // Load the wrapper class.
@@ -1488,7 +1497,7 @@ fn resolve_generic_wrapper_template(
     wrapper_subs.get(wrapper_tpl.as_str()).cloned()
 }
 
-/// Infer a generic type argument from an array literal.
+/// Extract a generic type argument from an array literal.
 ///
 /// For `@param array<TKey, TValue> $kv` with argument `["a" => 1]`:
 /// - `tpl_position == 0` → key type (`string`)
@@ -1557,6 +1566,19 @@ fn resolve_array_literal_generic(
     }
 }
 
+/// Extract the generic type argument at a given position from a resolved type.
+///
+/// For `array<int, string>` with position 0 → `int`, position 1 → `string`.
+/// For `list<User>` with position 0 → `User`.
+/// Also handles `PhpType::Array(inner)` as a single-arg generic.
+fn extract_generic_arg_at_position(ty: &PhpType, position: usize) -> Option<PhpType> {
+    match ty {
+        PhpType::Generic(_, args) => args.get(position).cloned(),
+        PhpType::Array(inner) if position == 0 => Some(inner.as_ref().clone()),
+        _ => None,
+    }
+}
+
 /// Resolve `$arr[0]` / `$arr[$key]` by extracting the generic element
 /// type from the base array's annotation or assignment.
 fn resolve_rhs_array_access<'b>(

src/stub_patches.rs

@@ -41,6 +41,18 @@
 //!    iterator available on the wrapper.
 //!    PHPStan ref: `stubs/iterable.stub`
 //!
+//! 3. **`FilterIterator`** -- extends `IteratorIterator` but stubs lack
+//!    `@template` params.  PHPStan adds the same three template params
+//!    and `@template-extends IteratorIterator<TKey, TValue, TIterator>`.
+//!
+//! 4. **`NoRewindIterator`**, **`CachingIterator`**, **`InfiniteIterator`**,
+//!    **`LimitIterator`** -- all extend `IteratorIterator`.  Same template
+//!    params + `@extends` generics + constructor binding `TIterator → $iterator`.
+//!
+//! 5. **`CallbackFilterIterator`** -- extends `FilterIterator`.
+//!    Same template params + `@extends FilterIterator<TKey, TValue, TIterator>`
+//!    + constructor binding.
+//!
 //! ## Removing patches
 //!
 //! When phpstorm-stubs gains proper annotations for a patched symbol,
@@ -82,7 +94,13 @@ pub fn apply_class_stub_patches(class: &mut ClassInfo) {
     match class.name.as_str() {
         "WeakMap" => patch_weak_map(class),
         "IteratorIterator" => patch_iterator_iterator(class),
+        "FilterIterator" => patch_filter_iterator(class),
         "NoRewindIterator" => patch_no_rewind_iterator(class),
+        "CachingIterator" => patch_caching_iterator(class),
+        "InfiniteIterator" => patch_infinite_iterator(class),
+        "LimitIterator" => patch_limit_iterator(class),
+        "CallbackFilterIterator" => patch_callback_filter_iterator(class),
+        "ArrayIterator" => patch_array_iterator(class),
         _ => {}
     }
 }
@@ -131,6 +149,13 @@ fn patch_iterator_iterator(class: &mut ClassInfo) {
     if !class.mixins.contains(&t_iter) {
         class.mixins.push(t_iter);
     }
+
+    // Patch current() → TValue and key() → TKey.
+    // phpstorm-stubs declare `current(): mixed` and `key(): mixed` which
+    // hides the generic type.  PHPStan's stubs override these.
+    patch_method_return_type(class, "current", PhpType::Named("TValue".to_string()));
+    patch_method_return_type(class, "key", PhpType::Named("TKey".to_string()));
+
     // Patch the constructor: add template binding TIterator → $iterator
     // so that `new IteratorIterator(new Subject())` infers TIterator = Subject.
     if let Some(ctor_idx) = class
@@ -153,16 +178,125 @@ fn patch_iterator_iterator(class: &mut ClassInfo) {
 }
 
 /// Add `@template TKey`, `@template TValue`,
-/// `@template TIterator of Iterator<TKey, TValue>`,
-/// `@extends IteratorIterator<TKey, TValue, TIterator>`,
-/// and constructor template binding `TIterator → $iterator`.
+/// `@template TIterator of Traversable<TKey, TValue>`,
+/// `@extends IteratorIterator<TKey, TValue, TIterator>`.
+///
+/// `FilterIterator` is abstract and extends `IteratorIterator`.
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_filter_iterator(class: &mut ClassInfo) {
+    if !class.template_params.is_empty() {
+        return;
+    }
+    patch_iterator_iterator_subclass(class, "IteratorIterator");
+    patch_method_return_type(class, "current", PhpType::Named("TValue".to_string()));
+    patch_method_return_type(class, "key", PhpType::Named("TKey".to_string()));
+}
+
+/// Patch `NoRewindIterator` with template params inherited from `IteratorIterator`.
 ///
 /// Without this patch, `new NoRewindIterator(generator())` resolves as
 /// bare `NoRewindIterator` without propagating the generator's type params.
+/// PHPStan ref: `stubs/iterable.stub`
 fn patch_no_rewind_iterator(class: &mut ClassInfo) {
     if !class.template_params.is_empty() {
         return;
     }
+    patch_iterator_iterator_subclass(class, "IteratorIterator");
+    patch_constructor_iterator_binding(class);
+}
+
+/// Patch `CachingIterator` with template params inherited from `IteratorIterator`.
+///
+/// `CachingIterator` extends `IteratorIterator` and wraps an iterator.
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_caching_iterator(class: &mut ClassInfo) {
+    if !class.template_params.is_empty() {
+        return;
+    }
+    patch_iterator_iterator_subclass(class, "IteratorIterator");
+    patch_method_return_type(class, "current", PhpType::Named("TValue".to_string()));
+    patch_method_return_type(class, "key", PhpType::Named("TKey".to_string()));
+    patch_constructor_iterator_binding(class);
+}
+
+/// Patch `InfiniteIterator` with template params inherited from `IteratorIterator`.
+///
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_infinite_iterator(class: &mut ClassInfo) {
+    if !class.template_params.is_empty() {
+        return;
+    }
+    patch_iterator_iterator_subclass(class, "IteratorIterator");
+    patch_constructor_iterator_binding(class);
+}
+
+/// Patch `LimitIterator` with template params inherited from `IteratorIterator`.
+///
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_limit_iterator(class: &mut ClassInfo) {
+    if !class.template_params.is_empty() {
+        return;
+    }
+    patch_iterator_iterator_subclass(class, "IteratorIterator");
+    patch_method_return_type(class, "current", PhpType::Named("TValue".to_string()));
+    patch_method_return_type(class, "key", PhpType::Named("TKey".to_string()));
+    patch_constructor_iterator_binding(class);
+}
+
+/// Patch `CallbackFilterIterator` with template params inherited from `FilterIterator`.
+///
+/// `CallbackFilterIterator` extends `FilterIterator` (not `IteratorIterator` directly).
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_callback_filter_iterator(class: &mut ClassInfo) {
+    if !class.template_params.is_empty() {
+        return;
+    }
+    patch_iterator_iterator_subclass(class, "FilterIterator");
+    patch_constructor_iterator_binding(class);
+}
+
+/// Patch `ArrayIterator` constructor to bind template params from the `$array` arg.
+///
+/// phpstorm-stubs declare `@template TKey of array-key` and `@template TValue`
+/// on the class, but the constructor's `@param` is just `object|array` with no
+/// generics.  PHPStan's stubs use `@param array<TKey, TValue> $array`.
+/// PHPStan ref: `stubs/iterable.stub`
+fn patch_array_iterator(class: &mut ClassInfo) {
+    if let Some(ctor_idx) = class
+        .methods
+        .iter()
+        .position(|m| m.name.as_str() == "__construct")
+    {
+        let mut ctor = (*class.methods[ctor_idx]).clone();
+
+        // Add template bindings: TKey → $array, TValue → $array
+        for tpl_name in ["TKey", "TValue"] {
+            let binding = (atom(tpl_name), atom("$array"));
+            if !ctor.template_bindings.iter().any(|(t, _)| t == &binding.0) {
+                ctor.template_bindings.push(binding);
+            }
+        }
+
+        // Set the parameter type hint to array<TKey, TValue> so that
+        // classify_template_binding can determine the GenericWrapper mode.
+        if let Some(param) = ctor.parameters.iter_mut().find(|p| p.name == "$array") {
+            param.type_hint = Some(PhpType::Generic(
+                "array".to_string(),
+                vec![
+                    PhpType::Named("TKey".to_string()),
+                    PhpType::Named("TValue".to_string()),
+                ],
+            ));
+        }
+
+        class.methods.make_mut()[ctor_idx] = std::sync::Arc::new(ctor);
+    }
+}
+
+/// Shared helper: add `@template TKey, TValue, TIterator` and
+/// `@extends <parent><TKey, TValue, TIterator>` to an `IteratorIterator`
+/// subclass (or sub-subclass like `CallbackFilterIterator`).
+fn patch_iterator_iterator_subclass(class: &mut ClassInfo, parent: &str) {
     add_templates(class, &[("TKey", None), ("TValue", None)]);
     let t_iter = atom("TIterator");
     if !class.template_params.contains(&t_iter) {
@@ -173,30 +307,33 @@ fn patch_no_rewind_iterator(class: &mut ClassInfo) {
         .entry(atom("TIterator"))
         .or_insert_with(|| {
             PhpType::Generic(
-                "Iterator".to_string(),
+                "Traversable".to_string(),
                 vec![
                     PhpType::Named("TKey".to_string()),
                     PhpType::Named("TValue".to_string()),
                 ],
             )
         });
-    // Add @extends generics so inheritance resolves TKey/TValue.
-    let iter_iter_atom = atom("IteratorIterator");
+    let parent_atom = atom(parent);
     if !class
         .extends_generics
         .iter()
-        .any(|(n, _)| *n == iter_iter_atom)
+        .any(|(n, _)| *n == parent_atom)
     {
         class.extends_generics.push((
-            iter_iter_atom,
+            parent_atom,
             vec![
                 PhpType::Named("TKey".to_string()),
                 PhpType::Named("TValue".to_string()),
                 PhpType::Named("TIterator".to_string()),
             ],
         ));
     }
-    // Patch constructor: bind TIterator from the $iterator parameter.
+}
+
+/// Shared helper: patch the constructor to bind `TIterator` from the
+/// `$iterator` parameter.
+fn patch_constructor_iterator_binding(class: &mut ClassInfo) {
     if let Some(ctor_idx) = class
         .methods
         .iter()
@@ -214,6 +351,22 @@ fn patch_no_rewind_iterator(class: &mut ClassInfo) {
     }
 }
 
+/// Override a method's return type on a class.
+///
+/// If the method exists, replaces its `return_type` with the given type.
+/// Used to patch stub methods like `current(): mixed` → `current(): TValue`.
+fn patch_method_return_type(class: &mut ClassInfo, method_name: &str, return_type: PhpType) {
+    if let Some(idx) = class
+        .methods
+        .iter()
+        .position(|m| m.name.as_str() == method_name)
+    {
+        let mut method = (*class.methods[idx]).clone();
+        method.return_type = Some(return_type);
+        class.methods.make_mut()[idx] = std::sync::Arc::new(method);
+    }
+}
+
 // ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 // Helpers
 // ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

tests/psalm_assertions/spl_iterator_patches.php

@@ -0,0 +1,30 @@
+<?php
+// Test that iterator decorator patches propagate generic types.
+
+/** @var ArrayIterator<int, string> $inner */
+
+$cachingIter = new CachingIterator($inner);
+$cachingValue = $cachingIter->current();
+assertType('string', $cachingValue);
+
+$infiniteIter = new InfiniteIterator($inner);
+$infiniteValue = $infiniteIter->current();
+assertType('string', $infiniteValue);
+
+$limitIter = new LimitIterator($inner, 0, 10);
+$limitValue = $limitIter->current();
+assertType('string', $limitValue);
+
+$noRewindIter = new NoRewindIterator($inner);
+$noRewindValue = $noRewindIter->current();
+assertType('string', $noRewindValue);
+
+$cbFilterIter = new CallbackFilterIterator($inner, function ($v) { return true; });
+$cbFilterValue = $cbFilterIter->current();
+assertType('string', $cbFilterValue);
+
+/** @var array<string, int> $typedArray */
+$arrayIter = new ArrayIterator($typedArray);
+assertType('ArrayIterator<string, int>', $arrayIter);
+$arrayIterValue = $arrayIter->current();
+assertType('int', $arrayIterValue);