Add support for jump-to-def for associations (#373)

* WIP: has_many POC

* Add belongs to

* Add has_one

* Add has_one_and_belongs_to_many

* Move association constant to separate support module

* Rubocop fixes

* WIP try to track class name on class enter

* refactor: use nesting for class name instead of hardcoding

* fix: handle invalid association name in server

* refactor: add location builder support class

* refactor: remove unused node lifecycles

* Add a class_name example

* Use reflect_on_association

* PR feedback

* Apply suggestions from code review

Co-authored-by: Vinicius Stock <[email protected]>

* Rubocop

---------

Co-authored-by: Vinicius Stock <[email protected]>

Author: maxveldink

lib/ruby_lsp/ruby_lsp_rails/addon.rb

@@ -5,7 +5,9 @@
 
 require_relative "../../ruby_lsp_rails/version"
 require_relative "support/active_support_test_case_helper"
+require_relative "support/associations"
 require_relative "support/callbacks"
+require_relative "support/location_builder"
 require_relative "runner_client"
 require_relative "hover"
 require_relative "code_lens"

lib/ruby_lsp/ruby_lsp_rails/definition.rb

@@ -57,7 +57,9 @@ def on_call_node_enter(node)
 
         return unless message
 
-        if Support::Callbacks::ALL.include?(message)
+        if Support::Associations::ALL.include?(message)
+          handle_association(node)
+        elsif Support::Callbacks::ALL.include?(message)
           handle_callback(node)
         elsif message.end_with?("_path") || message.end_with?("_url")
           handle_route(node)
@@ -86,23 +88,28 @@ def handle_callback(node)
       end
 
       sig { params(node: Prism::CallNode).void }
-      def handle_route(node)
-        result = @client.route_location(T.must(node.message))
+      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_location(
+          model_name: @nesting.join("::"),
+          association_name: association_name,
+        )
+
         return unless result
 
-        *file_parts, line = result.fetch(:location).split(":")
+        @response_builder << Support::LocationBuilder.line_location_from_s(result.fetch(:location))
+      end
 
-        # On Windows, file paths will look something like `C:/path/to/file.rb:123`. Only the last colon is the line
-        # number and all other parts compose the file path
-        file_path = file_parts.join(":")
+      sig { params(node: Prism::CallNode).void }
+      def handle_route(node)
+        result = @client.route_location(T.must(node.message))
+        return unless result
 
-        @response_builder << Interface::Location.new(
-          uri: URI::Generic.from_path(path: file_path).to_s,
-          range: Interface::Range.new(
-            start: Interface::Position.new(line: Integer(line) - 1, character: 0),
-            end: Interface::Position.new(line: Integer(line) - 1, character: 0),
-          ),
-        )
+        @response_builder << Support::LocationBuilder.line_location_from_s(result.fetch(:location))
       end
 
       sig { params(name: String).void }

lib/ruby_lsp/ruby_lsp_rails/runner_client.rb

@@ -98,6 +98,22 @@ def model(name)
         nil
       end
 
+      sig do
+        params(
+          model_name: String,
+          association_name: String,
+        ).returns(T.nilable(T::Hash[Symbol, T.untyped]))
+      end
+      def association_target_location(model_name:, association_name:)
+        make_request(
+          "association_target_location",
+          model_name: model_name,
+          association_name: association_name,
+        )
+      rescue => e
+        $stderr.puts("Ruby LSP Rails failed with #{e.message}: #{@stderr.read}")
+      end
+
       sig { params(name: String).returns(T.nilable(T::Hash[Symbol, T.untyped])) }
       def route_location(name)
         make_request("route_location", name: name)

lib/ruby_lsp/ruby_lsp_rails/server.rb

@@ -47,6 +47,8 @@ def execute(request, params)
           VOID
         when "model"
           resolve_database_info_from_model(params.fetch(:name))
+        when "association_target_location"
+          resolve_association_target(params)
         when "reload"
           ::Rails.application.reloader.reload!
           VOID
@@ -114,6 +116,29 @@ def resolve_database_info_from_model(model_name)
         { error: e.full_message(highlight: false) }
       end
 
+      def resolve_association_target(params)
+        const = ActiveSupport::Inflector.safe_constantize(params[:model_name])
+        unless active_record_model?(const)
+          return {
+            result: nil,
+          }
+        end
+
+        association_klass = const.reflect_on_association(params[:association_name].intern).klass
+
+        source_location = Object.const_source_location(association_klass.to_s)
+
+        {
+          result: {
+            location: source_location.first + ":" + source_location.second.to_s,
+          },
+        }
+      rescue NameError
+        {
+          result: nil,
+        }
+      end
+
       def active_record_model?(const)
         !!(
           const &&

lib/ruby_lsp/ruby_lsp_rails/support/associations.rb

@@ -0,0 +1,20 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Rails
+    module Support
+      module Associations
+        ALL = T.let(
+          [
+            "belongs_to",
+            "has_many",
+            "has_one",
+            "has_and_belongs_to_many",
+          ].freeze,
+          T::Array[String],
+        )
+      end
+    end
+  end
+end

lib/ruby_lsp/ruby_lsp_rails/support/location_builder.rb

@@ -0,0 +1,33 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Rails
+    module Support
+      class LocationBuilder
+        class << self
+          extend T::Sig
+
+          sig { params(location_string: String).returns(Interface::Location) }
+          def line_location_from_s(location_string)
+            *file_parts, line = location_string.split(":")
+
+            raise ArgumentError, "Invalid location string given" unless file_parts
+
+            # On Windows, file paths will look something like `C:/path/to/file.rb:123`. Only the last colon is the line
+            # number and all other parts compose the file path
+            file_path = file_parts.join(":")
+
+            Interface::Location.new(
+              uri: URI::Generic.from_path(path: file_path).to_s,
+              range: Interface::Range.new(
+                start: Interface::Position.new(line: Integer(line) - 1, character: 0),
+                end: Interface::Position.new(line: Integer(line) - 1, character: 0),
+              ),
+            )
+          end
+        end
+      end
+    end
+  end
+end

test/dummy/app/models/country.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class Country < ApplicationRecord
+end

test/dummy/app/models/label.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Label < ApplicationRecord
+  has_and_belongs_to_many :profiles
+end

test/dummy/app/models/profile.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+class Profile < ApplicationRecord
+  belongs_to :user
+  has_and_belongs_to_many :labels
+end

test/dummy/app/models/user.rb

@@ -5,6 +5,7 @@ class User < ApplicationRecord
   validates :first_name, presence: true
   has_one :profile
   scope :adult, -> { where(age: 18..) }
+  has_one :location, class_name: "Country"
 
   attr_readonly :last_name
 

test/dummy/db/migrate/20240521183115_create_countries.rb

@@ -0,0 +1,9 @@
+class CreateCountries < ActiveRecord::Migration[7.1]
+  def change
+    create_table :countries do |t|
+      t.string :name
+
+      t.timestamps
+    end
+  end
+end

test/dummy/db/migrate/20240521183200_add_country_to_user.rb

@@ -0,0 +1,5 @@
+class AddCountryToUser < ActiveRecord::Migration[7.1]
+  def change
+    add_reference :users, :country, null: false, foreign_key: true
+  end
+end

test/dummy/db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2024_05_03_153021) do
+ActiveRecord::Schema[7.1].define(version: 2024_05_21_183200) do
   create_table "composite_primary_keys", primary_key: ["order_id", "product_id"], force: :cascade do |t|
     t.integer "order_id"
     t.integer "product_id"
@@ -19,6 +19,12 @@
     t.datetime "updated_at", null: false
   end
 
+  create_table "countries", force: :cascade do |t|
+    t.string "name"
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+  end
+
   create_table "memberships", force: :cascade do |t|
     t.integer "user_id", null: false
     t.integer "organization_id", null: false
@@ -47,8 +53,11 @@
     t.integer "age"
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.integer "country_id", null: false
+    t.index ["country_id"], name: "index_users_on_country_id"
   end
 
   add_foreign_key "memberships", "organizations"
   add_foreign_key "memberships", "users"
+  add_foreign_key "users", "countries"
 end

test/ruby_lsp_rails/definition_test.rb

@@ -33,6 +33,101 @@ def baz; end
         assert_equal(14, response[1].range.end.character)
       end
 
+      test "recognizes has_many model associations" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 4 })
+          # typed: false
+
+          class Organization < ActiveRecord::Base
+            has_many :memberships
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "membership.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
+      test "recognizes belongs_to model associations" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 4 })
+          # typed: false
+
+          class Membership < ActiveRecord::Base
+            belongs_to :organization
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "organization.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
+      test "recognizes has_one model associations" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 4 })
+          # typed: false
+
+          class User < ActiveRecord::Base
+            has_one :profile
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "profile.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
+      test "recognizes has_and_belongs_to_many model associations" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 4 })
+          # typed: false
+
+          class Profile < ActiveRecord::Base
+            has_and_belongs_to_many :labels
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "label.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        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: 4 })
+          # typed: false
+
+          class User < ActiveRecord::Base
+            has_one :location, class_name: "Country"
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "country.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
       test "recognizes controller callback with string argument" do
         response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 10 })
           # typed: false

test/ruby_lsp_rails/runner_client_test.rb

@@ -30,6 +30,7 @@ class RunnerClientTest < ActiveSupport::TestCase
           ["age", "integer"],
           ["created_at", "datetime"],
           ["updated_at", "datetime"],
+          ["country_id", "integer"],
         ]
         response = T.must(@client.model("User"))
         assert_equal(columns, response.fetch(:columns))