Add #[ArrayShape] attribute support for array shape inference

Functions and methods annotated with #[ArrayShape(["key" => "type",
...])]
now produce array shape key completions, hover type info, and correct
type
resolution. Includes tests for function, method, and union return types.
Removes completed todo item from docs/todo.md and
docs/todo/completion.md.
Changelog updated.

Author: AJenbo

docs/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **`#[ArrayShape]` attribute support.** Functions and methods annotated with `#[ArrayShape(["key" => "type", ...])]` (used by ~84 phpstorm-stubs entries) now produce array shape key completions, hover type info, and correct type resolution. Affects commonly used functions like `parse_url`, `stat`, `pathinfo`, `gc_status`, `getimagesize`, and `session_get_cookie_params`.
 - **Convert to arrow function.** A new `refactor.rewrite` code action converts single-expression closures to arrow functions (`function($x) { return $x * 2; }` to `fn($x) => $x * 2`). The action is only offered when the conversion is safe: single return statement, no by-reference `use` captures, no `void`/`never` return type, and PHP >= 7.4.
 - **Convert switch to match.** A new `refactor.rewrite` code action converts `switch` statements to `match` expressions when all arms are single-expression returns or assignments to the same variable. Handles fall-through cases (merged with commas), trailing `break` removal, and `throw` arms. Requires PHP >= 8.0.
 - **Extract interface.** A new `refactor.extract` code action generates an interface from a concrete class. All public method signatures (excluding the constructor) are extracted into a new `{ClassName}Interface.php` file in the same directory, and the class is updated with `implements {ClassName}Interface`. Class-level and method-level `@template` tags are preserved when referenced by extracted methods.

docs/todo.md

@@ -25,7 +25,6 @@ within the same impact tier.
 
 | #   | Item                                                                                                                  | Impact     | Effort |
 | --- | --------------------------------------------------------------------------------------------------------------------- | ---------- | ------ |
-| C2  | [`#[ArrayShape]` return shapes on stub functions](todo/completion.md#c2-arrayshape-return-shapes-on-stub-functions)   | Medium     | Medium |
 | D10 | [PHPMD diagnostic proxy](todo/diagnostics.md#d10-phpmd-diagnostic-proxy)                                              | Low        | Medium |
 
 ## Sprint 7 — 1.0 release & IDE extensions

docs/todo/completion.md

@@ -46,44 +46,6 @@ These functions have return type semantics that don't fit into either
 
 ---
 
-## C2. `#[ArrayShape]` return shapes on stub functions
-
-**Impact: Medium · Effort: Medium**
-
-phpstorm-stubs annotate ~84 functions and methods with
-`#[ArrayShape(["key" => "type", ...])]` to declare the structure of
-their array return values. Almost none of these have a companion
-`@return array{...}` docblock, so the shape information is invisible
-to PHPantom. This affects commonly used functions like `parse_url`,
-`stat`, `pathinfo`, `gc_status`, `getimagesize`,
-`session_get_cookie_params`, `stream_get_meta_data`, and
-`password_get_info`.
-
-```php
-#[ArrayShape(["lifetime" => "int", "path" => "string", "domain" => "string",
-              "secure" => "bool", "httponly" => "bool", "samesite" => "string"])]
-function session_get_cookie_params(): array {}
-
-#[ArrayShape(["runs" => "int", "collected" => "int", "threshold" => "int", "roots" => "int"])]
-function gc_status(): array {}
-```
-
-**Attribute FQN:** `JetBrains\PhpStorm\ArrayShape`. Stub files import
-it via `use JetBrains\PhpStorm\ArrayShape;`. No aliases are used.
-Match by resolving through the `DocblockCtx` use-map and comparing the
-last segment of the resolved FQN (same pattern as `Deprecated` and
-`PhpStormStubsElementAvailable`).
-
-**Implementation:** During function/method extraction, scan for the
-`ArrayShape` attribute. Parse the associative array literal in its
-argument to build an `array{key: type, ...}` string, and use it as
-the effective return type (or parameter type when applied to a
-parameter). This complements the existing docblock `array{...}`
-parsing and should feed into the same `return_type` field on
-`FunctionInfo` / `MethodInfo`.
-
----
-
 ## C3. Go-to-definition for array shape keys via bracket access
 
 **Impact: Low-Medium · Effort: Medium**

src/completion/array_shape.rs

@@ -532,14 +532,18 @@ impl Backend {
                 &dummy_class
             }
         };
+        let function_loader = self.function_loader(file_ctx);
         let resolved = crate::completion::variable::resolution::resolve_variable_types(
             var_name,
             effective_class,
             &file_ctx.classes,
             content,
             cursor_offset as u32,
             &class_loader,
-            Loaders::default(),
+            Loaders {
+                function_loader: Some(&function_loader),
+                ..Loaders::default()
+            },
         );
         if resolved.is_empty() {
             None

src/parser/classes.rs

@@ -1959,6 +1959,15 @@ impl Backend {
                             parsed_doc_type.as_ref(),
                         );
 
+                        // Apply #[ArrayShape] override if present.
+                        let effective = if let Some(ctx) = doc_ctx {
+                            effective.map(|ty| {
+                                super::apply_array_shape_override(ty, &method.attribute_lists, ctx)
+                            })
+                        } else {
+                            effective
+                        };
+
                         let conditional = docblock::extract_conditional_return_type_from_info(info);
 
                         // Extract method-level @template params, their bounds,
@@ -2064,8 +2073,18 @@ impl Backend {
                         // #[Deprecated] attribute on the method itself.
                         let depr_info =
                             merge_deprecation_info(None, &method.attribute_lists, doc_ctx);
+
+                        // Apply #[ArrayShape] override even without a docblock.
+                        let effective_ret = if let Some(ctx) = doc_ctx {
+                            native_return_type.clone().map(|ty| {
+                                super::apply_array_shape_override(ty, &method.attribute_lists, ctx)
+                            })
+                        } else {
+                            native_return_type.clone()
+                        };
+
                         (
-                            native_return_type.clone(),
+                            effective_ret,
                             None,
                             depr_info.message,
                             depr_info.replacement,

src/parser/functions.rs

@@ -167,6 +167,14 @@ impl Backend {
                             parsed_doc_return.as_ref(),
                         );
 
+                        // If the effective return type is bare `array` (or
+                        // nullable/union containing array) and the function
+                        // has a `#[ArrayShape]` attribute, replace it with
+                        // the more specific shape type.
+                        let effective = effective.map(|ty| {
+                            super::apply_array_shape_override(ty, &func.attribute_lists, ctx)
+                        });
+
                         let conditional = info
                             .as_ref()
                             .and_then(docblock::extract_conditional_return_type_from_info);

src/parser/mod.rs

@@ -68,6 +68,9 @@ const ATTR_ELEMENT_AVAILABLE: &str = "PhpStormStubsElementAvailable";
 /// Last segment of the `LanguageLevelTypeAware` attribute FQN.
 const ATTR_LANGUAGE_LEVEL_TYPE_AWARE: &str = "LanguageLevelTypeAware";
 
+/// Last segment of the `ArrayShape` attribute FQN.
+const ATTR_ARRAY_SHAPE: &str = "ArrayShape";
+
 /// Fully-qualified names (without leading `\`) that we recognise as
 /// deprecation attributes.  Only the native PHP 8.4 `\Deprecated` and
 /// the JetBrains stubs `\JetBrains\PhpStorm\Deprecated` should match.
@@ -102,6 +105,13 @@ impl DocblockCtx<'_> {
         let canonical = self.resolve_attr_last_segment(name_str).unwrap_or(name_str);
         canonical == ATTR_LANGUAGE_LEVEL_TYPE_AWARE
     }
+
+    /// Check whether `attr_short_name` resolves to `ArrayShape`.
+    fn is_array_shape_attr(&self, attr_short_name: &[u8]) -> bool {
+        let name_str = bytes_to_str(attr_short_name);
+        let canonical = self.resolve_attr_last_segment(name_str).unwrap_or(name_str);
+        canonical == ATTR_ARRAY_SHAPE
+    }
 }
 
 // ─── PhpStormStubsElementAvailable Attribute Parsing ────────────────────────
@@ -325,6 +335,103 @@ pub(crate) fn extract_language_level_type_for_param(
     extract_language_level_type(&param.attribute_lists, ctx, php_version)
 }
 
+/// Extract an `array{key: type, ...}` shape from a `#[ArrayShape]` attribute.
+///
+/// phpstorm-stubs annotate functions and methods with
+/// `#[ArrayShape(["key" => "type", ...])]` to declare the structure of
+/// their array return values.  This function extracts that information
+/// and returns it as a `PhpType::ArrayShape`.
+pub(crate) fn extract_array_shape_type(
+    attribute_lists: &Sequence<'_, attribute::AttributeList<'_>>,
+    ctx: &DocblockCtx<'_>,
+) -> Option<PhpType> {
+    use crate::php_type::ShapeEntry;
+
+    for attr_list in attribute_lists.iter() {
+        for attr in attr_list.attributes.iter() {
+            if !ctx.is_array_shape_attr(last_segment(attr.name.value())) {
+                continue;
+            }
+
+            let arg_list = attr.argument_list.as_ref()?;
+            // The ArrayShape attribute takes a single positional argument:
+            // an associative array literal ["key" => "type", ...].
+            let first_arg = arg_list.arguments.first()?;
+            let expr = match first_arg {
+                argument::Argument::Positional(p) => p.value,
+                argument::Argument::Named(n) => n.value,
+            };
+
+            let elements: Box<dyn Iterator<Item = &ArrayElement<'_>>> = match expr {
+                Expression::Array(arr) => Box::new(arr.elements.iter()),
+                Expression::LegacyArray(arr) => Box::new(arr.elements.iter()),
+                _ => return None,
+            };
+
+            let mut entries = Vec::new();
+            for elem in elements {
+                if let ArrayElement::KeyValue(kv) = elem {
+                    let key = extract_string_literal_value(kv.key, ctx.content)?;
+                    let value_str = extract_string_literal_value(kv.value, ctx.content)?;
+                    let value_type = PhpType::parse(&value_str);
+                    entries.push(ShapeEntry {
+                        key: Some(key),
+                        value_type,
+                        optional: false,
+                    });
+                }
+            }
+
+            if !entries.is_empty() {
+                return Some(PhpType::ArrayShape(entries));
+            }
+        }
+    }
+
+    None
+}
+
+/// Replace bare `array` components in `ty` with the shape from an
+/// `#[ArrayShape]` attribute, if present.  Handles `array`, `?array`,
+/// and `array|false` patterns.
+pub(crate) fn apply_array_shape_override(
+    ty: PhpType,
+    attribute_lists: &Sequence<'_, attribute::AttributeList<'_>>,
+    ctx: &DocblockCtx<'_>,
+) -> PhpType {
+    let Some(shape) = extract_array_shape_type(attribute_lists, ctx) else {
+        return ty;
+    };
+
+    match &ty {
+        // Exact `array` → replace with shape.
+        PhpType::Named(n) if n == "array" => shape,
+        // `?array` → `?array{...}`
+        PhpType::Nullable(inner) => {
+            if matches!(inner.as_ref(), PhpType::Named(n) if n == "array") {
+                PhpType::Nullable(Box::new(shape))
+            } else {
+                ty
+            }
+        }
+        // `array|false` or other unions containing bare `array`.
+        PhpType::Union(members) => {
+            let new_members: Vec<PhpType> = members
+                .iter()
+                .map(|m| {
+                    if matches!(m, PhpType::Named(n) if n == "array") {
+                        shape.clone()
+                    } else {
+                        m.clone()
+                    }
+                })
+                .collect();
+            PhpType::Union(new_members)
+        }
+        _ => ty,
+    }
+}
+
 /// Parse the version → type array inside `LanguageLevelTypeAware`.
 ///
 /// Handles both `['8.0' => 'string']` (short array) and

src/php_type.rs

@@ -1438,6 +1438,10 @@ impl PhpType {
         match self {
             PhpType::ArrayShape(entries) | PhpType::ObjectShape(entries) => Some(entries),
             PhpType::Nullable(inner) => inner.shape_entries(),
+            PhpType::Union(members) => {
+                // Find the first array/object shape member in the union.
+                members.iter().find_map(|m| m.shape_entries())
+            }
             _ => None,
         }
     }