From 41f4227b3f8fc4765368e2b2e8d2a3f5c7701e42 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Thu, 8 Aug 2024 15:08:28 -0400
Subject: [PATCH] Add associations enhancement

---
 .../ruby_lsp_rails/indexing_enhancement.rb    | 50 +++++++++++++++++++
 .../indexing_enhancement_test.rb              | 33 ++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb b/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
index cb16f59d..3fbd69e9 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
@@ -23,11 +23,61 @@ def on_call_node(index, owner, node, file_path)
         case name
         when :extend
           handle_concern_extend(index, owner, node)
+        when :has_one, :has_many, :belongs_to, :has_and_belongs_to_many
+          handle_association(index, owner, node, file_path)
         end
       end
 
       private
 
+      sig do
+        params(
+          index: RubyIndexer::Index,
+          owner: RubyIndexer::Entry::Namespace,
+          node: Prism::CallNode,
+          file_path: String,
+        ).void
+      end
+      def handle_association(index, owner, node, file_path)
+        arguments = node.arguments&.arguments
+        return unless arguments
+
+        name_arg = arguments.first
+
+        name = case name_arg
+        when Prism::StringNode
+          name_arg.content
+        when Prism::SymbolNode
+          name_arg.value
+        end
+
+        return unless name
+
+        # Reader
+        index.add(RubyIndexer::Entry::Method.new(
+          name,
+          file_path,
+          name_arg.location,
+          name_arg.location,
+          [],
+          [RubyIndexer::Entry::Signature.new([])],
+          RubyIndexer::Entry::Visibility::PUBLIC,
+          owner,
+        ))
+
+        # Writer
+        index.add(RubyIndexer::Entry::Method.new(
+          "#{name}=",
+          file_path,
+          name_arg.location,
+          name_arg.location,
+          [],
+          [RubyIndexer::Entry::Signature.new([RubyIndexer::Entry::RequiredParameter.new(name: name.to_sym)])],
+          RubyIndexer::Entry::Visibility::PUBLIC,
+          owner,
+        ))
+      end
+
       sig do
         params(
           index: RubyIndexer::Index,
diff --git a/test/ruby_lsp_rails/indexing_enhancement_test.rb b/test/ruby_lsp_rails/indexing_enhancement_test.rb
index 3eb2bb5b..a4f5baeb 100644
--- a/test/ruby_lsp_rails/indexing_enhancement_test.rb
+++ b/test/ruby_lsp_rails/indexing_enhancement_test.rb
@@ -34,6 +34,39 @@ class Post < ActiveRecord::Base
         assert_includes(ancestors, "ActiveRecord::Store::ClassMethods")
         assert_includes(ancestors, "ActiveRecord::AttributeMethods::ClassMethods")
       end
+
+      test "associations" do
+        @index.index_single(RubyIndexer::IndexablePath.new(nil, "/fake.rb"), <<~RUBY)
+          class Post < ActiveRecord::Base
+            has_one :content
+            belongs_to :author
+            has_many :comments
+            has_and_belongs_to_many :tags
+          end
+        RUBY
+
+        assert_declaration_on_line("content", "Post", 2)
+        assert_declaration_on_line("content=", "Post", 2)
+
+        assert_declaration_on_line("author", "Post", 3)
+        assert_declaration_on_line("author=", "Post", 3)
+
+        assert_declaration_on_line("comments", "Post", 4)
+        assert_declaration_on_line("comments=", "Post", 4)
+
+        assert_declaration_on_line("tags", "Post", 5)
+        assert_declaration_on_line("tags=", "Post", 5)
+      end
+
+      private
+
+      def assert_declaration_on_line(method_name, class_name, line)
+        association_entries = @index.resolve_method(method_name, class_name)
+        refute_nil(association_entries)
+
+        association = association_entries.first
+        assert_equal(line, association.location.start_line)
+      end
     end
   end
 end