Add go to definition support for through: and class_name: association options

lib/ruby_lsp/ruby_lsp_rails/addon.rb

@@ -129,7 +129,7 @@ def create_document_symbol_listener(response_builder, dispatcher)
       def create_definition_listener(response_builder, uri, node_context, dispatcher)
         return unless @global_state
 
-        Definition.new(@rails_runner_client, response_builder, node_context, @global_state.index, dispatcher)
+        Definition.new(@rails_runner_client, response_builder, node_context, @global_state.index, dispatcher, uri)
       end
 
       # @override

lib/ruby_lsp/ruby_lsp_rails/definition.rb

@@ -30,13 +30,14 @@ module Rails
     class Definition
       include Requests::Support::Common
 
-      #: (RunnerClient client, RubyLsp::ResponseBuilders::CollectionResponseBuilder[(Interface::Location | Interface::LocationLink)] response_builder, NodeContext node_context, RubyIndexer::Index index, Prism::Dispatcher dispatcher) -> void
-      def initialize(client, response_builder, node_context, index, dispatcher)
+      #: (RunnerClient client, RubyLsp::ResponseBuilders::CollectionResponseBuilder[(Interface::Location | Interface::LocationLink)] response_builder, NodeContext node_context, RubyIndexer::Index index, Prism::Dispatcher dispatcher, URI::Generic uri) -> void
+      def initialize(client, response_builder, node_context, index, dispatcher, uri)
         @client = client
         @response_builder = response_builder
         @node_context = node_context
         @nesting = node_context.nesting #: Array[String]
         @index = index
+        @uri = uri
 
         dispatcher.register(self, :on_call_node_enter, :on_symbol_node_enter, :on_string_node_enter)
       end
@@ -80,7 +81,7 @@ def handle_possible_dsl(node)
         return unless arguments
 
         if Support::Associations::ALL.include?(message)
-          handle_association(call_node)
+          handle_association(node, arguments)
         elsif Support::Callbacks::ALL.include?(message)
           handle_callback(node, call_node, arguments)
           handle_if_unless_conditional(node, call_node, arguments)
@@ -90,6 +91,107 @@ def handle_possible_dsl(node)
         end
       end
 
+      #: ((Prism::SymbolNode | Prism::StringNode) node, Array[Prism::Node] arguments) -> void
+      def handle_association(node, arguments)
+        association_name_node = arguments.first
+        through_node = extract_option_value(arguments, "through")
+        class_name_node = extract_option_value(arguments, "class_name")
+
+        case node
+        when association_name_node
+          handle_association_name(association_name_node, class_name_node)
+        when through_node
+          handle_through_option(node)
+        when class_name_node
+          goto_class(node.content)
+        end
+      end
+
+      #: (Array[Prism::Node] arguments, String option_name) -> Prism::Node?
+      def extract_option_value(arguments, option_name)
+        keyword_hash = arguments.find { |arg| arg.is_a?(Prism::KeywordHashNode) } #: as Prism::KeywordHashNode?
+        return unless keyword_hash
+
+        assoc = keyword_hash.elements.find do |element|
+          element.is_a?(Prism::AssocNode) &&
+            element.key.is_a?(Prism::SymbolNode) &&
+            element.key.value == option_name
+        end #: as Prism::AssocNode?
+
+        assoc&.value
+      end
+
+      #: (Prism::SymbolNode node, Prism::Node? class_name_node) -> void
+      def handle_association_name(node, class_name_node)
+        # If class_name is specified, use it directly from the index
+        if class_name_node.is_a?(Prism::StringNode)
+          goto_class(class_name_node.content)
+          return
+        end
+
+        # Otherwise, ask Rails for the associated model
+        result = @client.association_target(
+          model_name: @nesting.join("::"),
+          association_name: node.unescaped,
+        )
+
+        return unless result
+
+        @response_builder << Support::LocationBuilder.line_location_from_s(result.fetch(:location))
+      end
+
+      #: (String class_name) -> void
+      def goto_class(class_name)
+        entries = @index[class_name]
+        return unless entries
+
+        entries.each do |entry|
+          @response_builder << Interface::Location.new(
+            uri: entry.uri.to_s,
+            range: range_from_location(entry.location),
+            )
+        end
+      end
+
+      #: ((Prism::SymbolNode | Prism::StringNode) node) -> void
+      def handle_through_option(node)
+        return unless node.is_a?(Prism::SymbolNode)
+
+        association_call = find_association_in_nesting_nodes(node.unescaped)
+        return unless association_call
+
+        @response_builder << Interface::Location.new(
+          uri: @uri.to_s,
+          range: range_from_location(association_call.location),
+        )
+      end
+
+      #: (String association_name) -> Prism::CallNode?
+      def find_association_in_nesting_nodes(association_name)
+        nesting_nodes = @node_context.instance_variable_get(:@nesting_nodes) #: as Array[Prism::Node]
+
+        nesting_nodes.each do |nesting_node|
+          body = case nesting_node
+          when Prism::ClassNode, Prism::ModuleNode
+            nesting_node.body
+          end
+
+          next unless body.is_a?(Prism::StatementsNode)
+
+          match = body.body.find do |statement|
+            next unless statement.is_a?(Prism::CallNode)
+            next unless Support::Associations::ALL.include?(statement.message)
+
+            first_arg = statement.arguments&.arguments&.first
+            first_arg.is_a?(Prism::SymbolNode) && first_arg.unescaped == association_name
+          end #: as Prism::CallNode?
+
+          return match if match
+        end
+
+        nil
+      end
+
       #: ((Prism::SymbolNode | Prism::StringNode) node, Prism::CallNode call_node, Array[Prism::Node] arguments) -> void
       def handle_callback(node, call_node, arguments)
         focus_argument = arguments.find { |argument| argument == node }
@@ -125,23 +227,6 @@ def handle_validation(node, call_node, arguments)
         collect_definitions(name)
       end
 
-      #: (Prism::CallNode node) -> void
-      def handle_association(node)
-        first_argument = node.arguments&.arguments&.first
-        return unless first_argument.is_a?(Prism::SymbolNode)
-
-        association_name = first_argument.unescaped
-
-        result = @client.association_target(
-          model_name: @nesting.join("::"),
-          association_name: association_name,
-        )
-
-        return unless result
-
-        @response_builder << Support::LocationBuilder.line_location_from_s(result.fetch(:location))
-      end
-
       #: (Prism::CallNode node) -> void
       def handle_route(node)
         result = @client.route_location(

test/ruby_lsp_rails/definition_test.rb

@@ -110,19 +110,71 @@ class Profile < ActiveRecord::Base
         assert_equal(2, response[0].range.end.line)
       end
 
-      test "handles class_name argument for associations" do
-        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 11 })
+      test "handles clicking on class_name value for associations" do
+        source = <<~RUBY
+          # typed: false
+
+          class User < ActiveRecord::Base
+            has_one :location, class_name: "Country"
+          end
+        RUBY
+
+        country_model = File.join(dummy_root, "app", "models", "country.rb")
+        response = generate_definitions_for_source(source, { line: 3, character: 35 }, index_files: [country_model])
+
+        assert_equal(1, response.size)
+        assert_equal(URI::Generic.from_path(path: country_model).to_s, response[0].uri)
+      end
+
+      test "handles clicking on association name when class_name is specified" do
+        source = <<~RUBY
           # typed: false
 
           class User < ActiveRecord::Base
             has_one :location, class_name: "Country"
           end
         RUBY
 
+        country_model = File.join(dummy_root, "app", "models", "country.rb")
+        response = generate_definitions_for_source(source, { line: 3, character: 12 }, index_files: [country_model])
+
+        assert_equal(1, response.size)
+        assert_equal(URI::Generic.from_path(path: country_model).to_s, response[0].uri)
+      end
+
+      test "handles clicking on through option value" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 30 })
+          # typed: false
+
+          class Organization < ActiveRecord::Base
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal("file:///fake.rb", response[0].uri)
+        assert_equal(3, response[0].range.start.line)
+        assert_equal(2, response[0].range.start.character)
+        assert_equal(3, response[0].range.end.line)
+        assert_equal(23, response[0].range.end.character)
+      end
+
+      test "handles clicking on association name with through option" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 14 })
+          # typed: false
+
+          class Organization < ActiveRecord::Base
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+        RUBY
+
         assert_equal(1, response.size)
 
         assert_equal(
-          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "country.rb")).to_s,
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "user.rb")).to_s,
           response[0].uri,
         )
         assert_equal(2, response[0].range.start.line)
@@ -469,10 +521,15 @@ def name; end
 
       private
 
-      def generate_definitions_for_source(source, position)
+      def generate_definitions_for_source(source, position, index_files: [])
         with_server(source) do |server, uri|
           sleep(0.1) while RubyLsp::Addon.addons.first.instance_variable_get(:@rails_runner_client).is_a?(NullClient)
 
+          index_files.each do |file_path|
+            file_uri = URI::Generic.from_path(path: file_path)
+            server.global_state.index.index_single(file_uri, File.read(file_path))
+          end
+
           server.process_message(
             id: 1,
             method: "textDocument/definition",