Custom Eloquent collection swapping

Author: AJenbo

docs/CHANGELOG.md

@@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **String-aware completion suppression.** Completion is now suppressed inside string literals (single-quoted strings, nowdocs, and plain text in double-quoted strings and heredocs) to reduce noise from accidental triggers. Interpolation contexts are detected and allowed through: `{$expr->}` brace interpolation and simple `"$var->"` property access both produce completions as expected. Array shape key completion (`$arr['`) is unaffected.
 - **Nested key completion for literal arrays.** When a variable is assigned a literal array with nested associative arrays (e.g. `$cfg = ['db' => ['host' => 'x', 'port' => 3306]]`), completing `$cfg['db']['` now offers the nested keys. Previously only a `@var array{...}` docblock produced nested completions. Works at arbitrary nesting depth, inside class methods, and at the top level.
 - **Generator yield type inference inside generator bodies.** When a function or method declares `@return Generator<TKey, TValue, TSend, TReturn>`, variables inside the generator body are now typed from the annotation. Variables that appear as the operand of a `yield` statement are inferred as TValue (the 2nd generic parameter), and variables assigned from a yield expression (`$var = yield $expr`) are inferred as TSend (the 3rd generic parameter). Works in class methods, top-level functions, key-value yields (`yield $k => $v`), cross-file resolution via PSR-4, and all four Generator type parameter arities. Explicit assignments still take priority over yield inference.
+- **Custom Eloquent collections.** Models that declare a custom collection via `#[CollectedBy(CustomCollection::class)]` or `/** @use HasCollection<CustomCollection> */ use HasCollection;` now resolve to the custom collection class instead of the standard `Illuminate\Database\Eloquent\Collection`. Custom collection methods appear in completions after `Model::where(...)->get()->`, after `Model::get()->`, and on relationship properties that return a collection of the model. The attribute takes priority over the trait when both are present. Works with short names, fully-qualified names, cross-file PSR-4 resolution, and same-file definitions. Models without a custom collection continue to use the standard Collection.
 
 ### Changed
 

example.php

@@ -889,6 +889,32 @@ public function demo(): void
 }
 
 
+// ── Custom Eloquent Collections ─────────────────────────────────────────────
+// Models with #[CollectedBy(CustomCollection::class)] or
+// /** @use HasCollection<CustomCollection> */ use HasCollection;
+// resolve to the custom collection class instead of the standard
+// Illuminate\Database\Eloquent\Collection. This means custom methods
+// like topRated() and averageRating() appear in completions after ->get().
+
+class CustomCollectionDemo
+{
+    public function demo(): void
+    {
+        // Builder chain → custom collection (via #[CollectedBy] on Review)
+        $reviews = Review::where('published', true)->get();
+        $reviews->topRated();        // custom method from ReviewCollection
+        $reviews->averageRating();   // custom method from ReviewCollection
+        $reviews->count();           // inherited from standard Collection
+        $reviews->first();           // inherited — returns Review|null
+
+        // Relationship properties also use the custom collection
+        $review = new Review();
+        $review->replies->topRated();       // HasMany<Review> → ReviewCollection
+        $review->replies->averageRating();  // ReviewCollection method
+    }
+}
+
+
 // ── Match Class-String Forwarding to Conditional Return Types ───────────────
 // When a variable holds a ::class value from a match expression and is then
 // passed to a function/method with @template T + @param class-string<T> +
@@ -2600,6 +2626,32 @@ class BlogTag extends \Illuminate\Database\Eloquent\Model
     public function getLabel(): string { return ''; }
 }
 
+// ── Custom Eloquent Collection Demo Models ──────────────────────────────────
+
+/**
+ * @template TKey of array-key
+ * @template TModel
+ * @extends \Illuminate\Database\Eloquent\Collection<TKey, TModel>
+ */
+class ReviewCollection extends \Illuminate\Database\Eloquent\Collection
+{
+    /** @return array<TKey, TModel> */
+    public function topRated(): array { return []; }
+
+    /** @return float */
+    public function averageRating(): float { return 0.0; }
+}
+
+#[\Illuminate\Database\Eloquent\Attributes\CollectedBy(ReviewCollection::class)]
+class Review extends \Illuminate\Database\Eloquent\Model
+{
+    public function getTitle(): string { return ''; }
+    public function getRating(): int { return 0; }
+
+    /** @return \Illuminate\Database\Eloquent\Relations\HasMany<Review, $this> */
+    public function replies(): mixed { return $this->hasMany(Review::class); }
+}
+
 } // end namespace Demo
 
 // ── Illuminate Stubs ────────────────────────────────────────────────────────
@@ -2656,6 +2708,15 @@ class MorphToMany {}
     class HasManyThrough {}
 }
 
+namespace Illuminate\Database\Eloquent\Attributes {
+    class CollectedBy {}
+}
+
+namespace Illuminate\Database\Eloquent {
+    /** @template TCollection */
+    trait HasCollection {}
+}
+
 namespace Illuminate\Database\Concerns {
 
     /**

src/completion/resolver.rs

@@ -1057,6 +1057,50 @@ impl Backend {
 
         match found {
             Some(cls) => {
+                // ── Custom Eloquent collection swapping ────────────────
+                // When the resolved class is the standard Eloquent
+                // Collection and one of the generic type args is a model
+                // with a `custom_collection` declared (via
+                // `#[CollectedBy]` or `@use HasCollection<X>`), swap to
+                // the custom collection class so that its own methods
+                // (e.g. `topRated()`) appear in completions.
+                //
+                // This handles the common chain pattern:
+                //   Model::where(...)->get()
+                // where Builder's `get()` returns
+                //   `\Illuminate\Database\Eloquent\Collection<int, TModel>`
+                // and TModel has been substituted to the concrete model.
+                //
+                // We compare against `base_clean` (the FQN extracted from
+                // the type hint) rather than `cls.file_namespace` because
+                // `file_namespace` is not always populated when classes are
+                // loaded via PSR-4 / classmap.
+                let is_eloquent_collection = {
+                    let bc = base_clean.strip_prefix('\\').unwrap_or(&base_clean);
+                    bc == crate::types::ELOQUENT_COLLECTION_FQN
+                };
+                let cls = if is_eloquent_collection && !generic_args.is_empty() {
+                    // The last generic arg is typically the model type.
+                    let model_arg = generic_args.last().unwrap();
+                    let model_clean = model_arg.strip_prefix('\\').unwrap_or(model_arg);
+                    let model_class = find_class_by_name(all_classes, model_clean)
+                        .cloned()
+                        .or_else(|| class_loader(model_clean));
+                    if let Some(ref mc) = model_class
+                        && let Some(ref coll_name) = mc.custom_collection
+                    {
+                        let coll_clean = coll_name.strip_prefix('\\').unwrap_or(coll_name);
+                        find_class_by_name(all_classes, coll_clean)
+                            .cloned()
+                            .or_else(|| class_loader(coll_clean))
+                            .unwrap_or(cls)
+                    } else {
+                        cls
+                    }
+                } else {
+                    cls
+                };
+
                 // Apply generic substitution if the type hint carried
                 // generic arguments and the class has template parameters.
                 // Resolve the class fully first (including trait methods,

src/parser/ast_update.rs

@@ -379,6 +379,12 @@ impl Backend {
                 .map(|m| Self::resolve_name(m, use_map, namespace))
                 .collect();
 
+            // Resolve custom collection class name to FQN
+            if let Some(ref coll) = class.custom_collection {
+                let resolved = Self::resolve_name(coll, use_map, namespace);
+                class.custom_collection = Some(resolved);
+            }
+
             // Resolve type arguments in @extends, @implements, and @use
             // generics so that after generic substitution, return types
             // and property types are fully-qualified and can be resolved

src/parser/classes.rs

@@ -1,10 +1,12 @@
 use std::collections::HashMap;
 
 use mago_span::HasSpan;
+use mago_syntax::ast::attribute::AttributeList;
 use mago_syntax::ast::class_like::method::MethodBody;
 use mago_syntax::ast::class_like::trait_use::{
     TraitUseAdaptation, TraitUseMethodReference, TraitUseSpecification,
 };
+use mago_syntax::ast::sequence::Sequence;
 /// Class, interface, trait, and enum extraction.
 ///
 /// Each class-like declaration is tagged with a [`ClassLikeKind`] so that
@@ -93,6 +95,65 @@ fn extract_class_docblock<'a>(
     }
 }
 
+/// Extract the custom collection class name from a `#[CollectedBy(X::class)]` attribute.
+///
+/// Scans the class's attribute lists for an attribute whose short name is
+/// `CollectedBy` and extracts the first argument's text with `::class` stripped.
+/// Returns `None` if no such attribute exists.
+fn extract_collected_by_attribute(
+    attribute_lists: &Sequence<'_, AttributeList<'_>>,
+    content: &str,
+) -> Option<String> {
+    for attr_list in attribute_lists.iter() {
+        for attr in attr_list.attributes.iter() {
+            let short = attr.name.last_segment();
+            if short != "CollectedBy" {
+                continue;
+            }
+            let arg_list = attr.argument_list.as_ref()?;
+            let first_arg = arg_list.arguments.first()?;
+            let span = first_arg.span();
+            let start = span.start.offset as usize;
+            let end = span.end.offset as usize;
+            let text = content.get(start..end)?;
+            let class_name = text.trim_end_matches("::class").trim();
+            if !class_name.is_empty() {
+                return Some(class_name.to_string());
+            }
+        }
+    }
+    None
+}
+
+/// Determine the custom collection class for an Eloquent model.
+///
+/// Checks two sources in priority order:
+///
+/// 1. `#[CollectedBy(CustomCollection::class)]` attribute on the class.
+/// 2. `/** @use HasCollection<CustomCollection> */` in `use_generics`.
+///
+/// The attribute takes priority because it is the newer Laravel API.
+fn extract_custom_collection(
+    attribute_lists: &Sequence<'_, AttributeList<'_>>,
+    use_generics: &[(String, Vec<String>)],
+    content: &str,
+) -> Option<String> {
+    // 1. Try the #[CollectedBy] attribute first.
+    if let Some(name) = extract_collected_by_attribute(attribute_lists, content) {
+        return Some(name);
+    }
+
+    // 2. Fall back to @use HasCollection<X>.
+    for (trait_name, args) in use_generics {
+        let short = trait_name.rsplit('\\').next().unwrap_or(trait_name);
+        if short == "HasCollection" && !args.is_empty() {
+            return Some(args[0].clone());
+        }
+    }
+
+    None
+}
+
 impl Backend {
     /// Recursively walk statements and extract class information.
     /// This handles classes at the top level as well as classes nested
@@ -141,6 +202,10 @@ impl Backend {
                     let start_offset = class.left_brace.start.offset;
                     let end_offset = class.right_brace.end.offset;
 
+                    let content = doc_ctx.map(|c| c.content).unwrap_or("");
+                    let custom_collection =
+                        extract_custom_collection(&class.attribute_lists, &use_generics, content);
+
                     classes.push(ClassInfo {
                         kind: ClassLikeKind::Class,
                         name: class_name,
@@ -166,6 +231,7 @@ impl Backend {
                         trait_aliases,
                         class_docblock: doc_info.raw_docblock,
                         file_namespace: None,
+                        custom_collection,
                     });
 
                     // Walk method bodies for anonymous classes.
@@ -236,6 +302,7 @@ impl Backend {
                         trait_aliases,
                         class_docblock: doc_info.raw_docblock,
                         file_namespace: None,
+                        custom_collection: None,
                     });
 
                     // Walk method bodies for anonymous classes.
@@ -288,6 +355,7 @@ impl Backend {
                         trait_aliases,
                         class_docblock: doc_info.raw_docblock,
                         file_namespace: None,
+                        custom_collection: None,
                     });
 
                     // Walk method bodies for anonymous classes.
@@ -358,6 +426,7 @@ impl Backend {
                         trait_aliases: vec![],
                         class_docblock: doc_info.raw_docblock,
                         file_namespace: None,
+                        custom_collection: None,
                     });
 
                     // Walk method bodies for anonymous classes.
@@ -440,6 +509,7 @@ impl Backend {
             trait_aliases,
             class_docblock: None,
             file_namespace: None,
+            custom_collection: None,
         }
     }
 

src/types.rs

@@ -520,6 +520,20 @@ pub struct ClassInfo {
     /// namespace blocks (e.g. `Illuminate\Database\Eloquent\Builder` vs
     /// `Illuminate\Database\Query\Builder`).
     pub file_namespace: Option<String>,
+    /// Custom collection class for Eloquent models.
+    ///
+    /// Detected from two Laravel mechanisms:
+    ///
+    /// 1. The `#[CollectedBy(CustomCollection::class)]` attribute on the
+    ///    model class.
+    /// 2. The `/** @use HasCollection<CustomCollection> */` docblock
+    ///    annotation on a `use HasCollection;` trait usage.
+    ///
+    /// When set, the `LaravelModelProvider` replaces
+    /// `\Illuminate\Database\Eloquent\Collection` with this class in
+    /// relationship property types and Builder-forwarded return types
+    /// (e.g. `get()`, `all()`).
+    pub custom_collection: Option<String>,
 }
 
 // ─── ClassInfo helpers ──────────────────────────────────────────────────────
@@ -563,6 +577,14 @@ pub(crate) struct FileContext {
     pub namespace: Option<String>,
 }
 
+// ─── Eloquent Constants ─────────────────────────────────────────────────────
+
+/// The fully-qualified name of the Eloquent Collection class.
+///
+/// Used by the `LaravelModelProvider` to detect and replace collection
+/// return types when a model declares a custom collection class.
+pub const ELOQUENT_COLLECTION_FQN: &str = "Illuminate\\Database\\Eloquent\\Collection";
+
 // ─── Recursion Depth Limits ─────────────────────────────────────────────────
 //
 // Centralised constants for the maximum recursion depth allowed when