Generate per-model Relation types

This a more accurate represenation of what ActiveRecord actually does at
runtime, and therefore gives better code suggestions when dealing with
associations and model classes.

I've settled on this as being the best compromise after trying a few
different approaches. There are two "challenges" that I believe can't be
met any other way (at this time).

1. It is not possible to write a return type annotation for the methods
   of various ActiveRecord mixins that will be correct for both a model
   class (e.g. Person.where) and relation (e.g. people.where). [0], [1]
2. It is not possible to represent ActiveRecords "class methods are also
   relation methods" behaviour without model-specific relation types.

It's conceivable that Solargraph could change it's interpretation of [self]
on class methods to solve challenge iftheshoefritz#1, but the second challenge really
forces our hand. In order to represent this correctly, Solargraph would
need support for method-missing delegation, **and** delegating those
missing methods to an associated/generic type.

Given @castwide is working on RBS support and the Ruby ecosystem is
likely to move that way in the future, it seems pragmatic to eat the
cost repetition / manual labour in this gem rather than try to push YARD
types into supporting that degree of type-level programming. 😅

[0]: castwide/solargraph#592
[1]: lsegal/yard#1257

lib/solargraph/rails/annotations.rb

@@ -42,13 +42,44 @@
 #     # @yieldself [ActionDispatch::Routing::Mapper]
 #     def draw; end
 #   end
+#
+#   # this module doesn't really exist, it's here to avoid repeating these mixins
+#   module ActiveRecord::RelationMethods
+#     include Enumerable
+#     include ActiveRecord::QueryMethods
+#     include ActiveRecord::FinderMethods
+#     include ActiveRecord::Calculations
+#     include ActiveRecord::Batches
+#   end
+#
+#   class ActiveRecord::Relation
+#     include ActiveRecord::RelationMethods
+#   end
+#
 #   class ActiveRecord::Base
-#     extend ActiveRecord::QueryMethods
-#     extend ActiveRecord::FinderMethods
 #     extend ActiveRecord::Associations::ClassMethods
 #     extend ActiveRecord::Inheritance::ClassMethods
 #     extend ActiveRecord::ModelSchema::ClassMethods
 #     extend ActiveRecord::Transactions::ClassMethods
 #     extend ActiveRecord::Scoping::Named::ClassMethods
+#     extend ActiveRecord::RelationMethods
 #     include ActiveRecord::Persistence
 #   end
+
+# @!override ActiveRecord::Batches#find_each
+#   @yieldparam_single_parameter
+
+# @!override ActiveRecord::Calculations#count
+#   @return [Integer, Hash]
+# @!override ActiveRecord::Calculations#pluck
+#   @overload pluck(one)
+#     @return [Array]
+#   @overload pluck(one, two, *more)
+#     @return [Array<Array>]
+
+# @!override ActiveRecord::QueryMethods::WhereChain#not
+#   @return_single_parameter
+# @!override ActiveRecord::QueryMethods::WhereChain#missing
+#   @return_single_parameter
+# @!override ActiveRecord::QueryMethods::WhereChain#associated
+#   @return_single_parameter

lib/solargraph/rails/model.rb

@@ -9,11 +9,35 @@ def self.valid_filename?(filename)
         filename.include?('app/models')
       end
 
+      # @param source_map [Solargraph::SourceMap]
+      # @param ns [Solargraph::Pin::Namespace]
       def process(source_map, ns)
         return [] unless self.class.valid_filename?(source_map.filename)
 
-        walker = Walker.from_source(source_map.source)
         pins = []
+        abstract = false
+
+        # ActiveRecord defines a hidden subclass of ActiveRecord::Relation for
+        # each model class that inherits from ActiveRecord::Base.
+        pins << relation = Solargraph::Pin::Namespace.new(
+          name: 'ActiveRecord_Relation',
+          type: :class,
+          visibility: :private,
+          closure: ns,
+        )
+        pins << Solargraph::Pin::Reference::Superclass.new(
+          name: "ActiveRecord::Relation",
+          closure: relation,
+        )
+
+        pins << Solargraph::Pin::Method.new(
+          name: 'model',
+          scope: :instance,
+          closure: relation,
+          comments: "@return [Class<#{ns.name}>]"
+        )
+
+        walker = Walker.from_source(source_map.source)
 
         walker.on :send, [nil, :belongs_to] do |ast|
           pins << singular_association(ns, ast)
@@ -31,60 +55,92 @@ def process(source_map, ns)
           pins << plural_association(ns, ast)
         end
 
+        walker.on :send, [:self, :abstract_class=, :true] do |ast|
+          abstract = true
+        end
+
         walker.on :send, [nil, :scope] do |ast|
           next if ast.children[2].nil?
           name = ast.children[2].children.last
 
-          method_pin =
-            Util.build_public_method(
-              ns,
-              name.to_s,
-              types: ns.return_type.map(&:tag),
-              scope: :class,
-              location: Util.build_location(ast, ns.filename)
-            )
+          parameters = []
 
           if ast.children.last.type == :block
             location = ast.children.last.location
-            block_pin =
-              source_map.locate_block_pin(location.line, location.column)
+            block_pin = source_map.locate_block_pin(location.line, location.column)
+            parameters.concat(block_pin.parameters.clone)
+            # hacky
             block_pin.instance_variable_set(:@binder, relation)
-            method_pin.parameters.concat(block_pin.parameters.clone)
           end
-          pins << method_pin
+
+          location = Util.build_location(ast, ns.filename)
+          # define scopes as a class methods on the model, and instance methods
+          # on the hidden relation class
+          pins << Util.build_public_method(
+            ns,
+            name.to_s,
+            scope: :class,
+            parameters: parameters,
+            types: [relation_type(ns.name)],
+            location: location
+          )
+          pins << Util.build_public_method(
+            relation,
+            name.to_s,
+            scope: :instance,
+            parameters: parameters,
+            types: [relation_type(ns.name)],
+            location: location
+          )
         end
 
         walker.walk
-        if pins.any?
-          Solargraph.logger.debug(
-            "[Rails][Model] added #{pins.map(&:name)} to #{ns.path}"
-          )
+
+        # Class methods on the model are exposed as *instance* methods on the
+        # hidden ActiveRecord_Relation class.
+        #
+        # Uses DelegatedMethod pins (instead of build_public_method) so Solargraph
+        # will show the "real" method pin for type inference, probing, docs etc.
+        source_map.pins.each do |pin|
+          next unless pin.is_a?(Solargraph::Pin::Method) && pin.scope == :class && pin.closure == ns
+
+          pins << Solargraph::Pin::DelegatedMethod.new(closure: relation, scope: :instance, method: pin)
+        end
+
+
+        unless abstract
+          pins += relation_method_pins(ns, :class, ns.path)
+          pins += relation_method_pins(relation, :instance, ns.path)
         end
+
+        Solargraph.logger.debug("[Rails][Model] added #{pins.map(&:name)} to #{ns.path}")
+
         pins
       end
 
+
       def plural_association(ns, ast)
-        relation_name = ast.children[2].children.first
+        association_name = ast.children[2].children.first
         class_name =
           extract_custom_class_name(ast) ||
-            relation_name.to_s.singularize.camelize
+            association_name.to_s.singularize.camelize
 
         Util.build_public_method(
           ns,
-          relation_name.to_s,
-          types: ["ActiveRecord::Associations::CollectionProxy<#{class_name}>"],
+          association_name.to_s,
+          types: [relation_type(class_name)],
           location: Util.build_location(ast, ns.filename)
         )
       end
 
       def singular_association(ns, ast)
-        relation_name = ast.children[2].children.first
+        association_name = ast.children[2].children.first
         class_name =
-          extract_custom_class_name(ast) || relation_name.to_s.camelize
+          extract_custom_class_name(ast) || association_name.to_s.camelize
 
         Util.build_public_method(
           ns,
-          relation_name.to_s,
+          association_name.to_s,
           types: [class_name],
           location: Util.build_location(ast, ns.filename)
         )
@@ -96,6 +152,122 @@ def extract_custom_class_name(ast)
 
         node.children.last
       end
+
+      # Generate method pins for ActiveRecord methods in the given namespace/scope, where the
+      # the return types will be templated with the provided model class.
+      #
+      # These method pins don't need to include any documentation, as Solargraph will merge
+      # documentation from Rails when it resolves the "method stack" for each pin.
+      #
+      # @param ns [Solargraph::Pin::Namespace] the namespace (model or relation class) in which to define methods.
+      # @param scope [:instance, :class] the method scope (:class for the model and :instance for the relation).
+      # @param model_class [String] the model class (e.g. "Person") that should be used in return types.
+      # @return [Array<Solargraph::Pin::Method>]
+      def relation_method_pins(namespace, scope, model_class)
+        pins = []
+        RETURNS_RELATION.each do |method|
+          pins << Util.build_public_method(namespace, method, scope: scope, types: [relation_type(model_class)])
+        end
+        RETURNS_INSTANCE.each do |method|
+          pins << Util.build_public_method(namespace, method, scope: scope, types: [model_class])
+        end
+        OVERLOADED.each do |method, overloads|
+          comments = overloads.map do |args, lines|
+            lines = ["@return [#{lines}]"] if lines.is_a?(String)
+            lines = ["@overload #{method}#{args}"] + lines
+            lines.map { |line| line.gsub '$T', model_class }.join("\n  ")
+          end
+          pins << Util.build_public_method(namespace, method, scope: scope, comments: comments.join("\n"))
+        end
+        pins
+      end
+
+      # construct the type name for the models hidden relation class.
+      # the additional type parameter is _not_ redundant, it makes enumerable methods work.
+      def relation_type(model_path)
+        "#{model_path}::ActiveRecord_Relation"
+      end
+
+      RETURNS_RELATION = %w[
+        all
+        and
+        annotate
+        distinct
+        eager_load
+        excluding
+        from
+        group
+        having
+        in_order_of
+        includes
+        invert_where
+        joins
+        left_joins
+        left_outer_joins
+        limit
+        lock
+        none
+        offset
+        or
+        order
+        preload
+        readonly
+        references
+        reorder
+        reselect
+        reverse_order
+        rewhere
+        select
+        strict_loading
+        unscope
+        where
+        without
+      ]
+
+      RETURNS_INSTANCE = %w[
+        find
+        find_by find_by!
+        take
+        take!
+        sole find_sole_by
+        first  second  third  fourth  fifth  third_to_last  second_to_last  last
+        first! second! third! fourth! fifth! third_to_last! second_to_last! last!
+        forty_two
+        forty_two!
+      ]
+
+      OVERLOADED = {
+        "where" => {
+          "()" => "ActiveRecord::QueryMethods::WhereChain<$T::ActiveRecord_Relation>",
+          "(*args)" => "$T::ActiveRecord_Relation",
+        },
+        "select" => {
+          "()" => [
+            "@yieldparam [$T]",
+            "@return [Array<$T>]",
+          ],
+          "(*args)" => "$T::ActiveRecord_Relation",
+        },
+        "find" => {
+          "(id)" => [
+            "@param id [Integer, String]",
+            "@return [$T]"
+          ],
+          "(*ids)" => "Array<$T>",
+        },
+        "take" => {
+          "()" => "T, nil",
+          "(limit)" => "Array<$T>",
+        },
+        "first" => {
+          "()" => "$T, nil",
+          "(limit)" => "Array<$T>",
+        },
+        "last" => {
+          "()" => "$T, nil",
+          "(limit)" => "Array<$T>"
+        },
+      }
     end
   end
 end

lib/solargraph/rails/types.yml

@@ -1,17 +1,9 @@
 ActionController::Metal#params:
   return: ["ActionController::Parameters"]
-ActiveRecord::FinderMethods#find:
-  return: ["self", "Array<self>"]
 ActionController::Cookies#cookies:
   return: ["ActionDispatch::Cookies::CookieJar"]
 ActionDispatch::Flash::FlashHash#now:
   return: ["ActionDispatch::Flash::FlashNow"]
-ActiveRecord::QueryMethods#where:
-  return: ["self", "ActiveRecord::Relation", "ActiveRecord::QueryMethods::WhereChain"]
-ActiveRecord::QueryMethods#not:
-  return: ["ActiveRecord::QueryMethods::WhereChain"]
-ActiveRecord::FinderMethods#find_by:
-  return: ["self", "nil"]
 Rails.application:
   return: ["Rails::Application"]
 ActionDispatch::Routing::RouteSet#draw:

lib/solargraph/rails/util.rb

@@ -4,23 +4,25 @@ module Util
       def self.build_public_method(
         ns,
         name,
+        comments: +"",
+        parameters: [],
         types: nil,
         location: nil,
         attribute: false,
         scope: :instance
       )
         opts = {
           name: name,
+          parameters: parameters,
           location: location,
           closure: ns,
           scope: scope,
           attribute: attribute
         }
 
-        comments = []
-        comments << "@return [#{types.join(',')}]" if types
+        comments << "\n@return [#{types.join(',')}]" if types
 
-        opts[:comments] = comments.join("\n")
+        opts[:comments] ||= comments
 
         Solargraph::Pin::Method.new(**opts)
       end