Be able to patch includes for single line

lib/bulletmark_repairer/corrector.rb

@@ -1,16 +1,33 @@
 # frozen_string_literal: true
 
 class Corrector < Parser::TreeRewriter
-  def on_send(node)
-    if node.location.expression.line <= line_no && line_no <= node.location.expression.last_line
-      insert_after node.children[0].location.expression, ".includes(#{associations})"
-    end
+  def on_def(node)
+    method = node.children.first
+    patched_methods[method] = nil
+    node.children.each { |child_node| insert_include(node: child_node, method:) }
 
     super
   end
 
   private
 
+  def patched_methods
+    @patched_methods ||= {}
+  end
+
+  def insert_include(node:, method:)
+    return if patched_methods[method]
+    return unless node.respond_to?(:children) && node.children.size > 0
+    return unless node.location.expression.line <= line_no && line_no <= node.location.expression.last_line
+
+    if node.children.last.in?(%i[each map])
+      insert_after node.children[0].location.expression, ".includes(#{associations})"
+      patched_methods[method] = true
+    else
+      node.children.each { |child_node| insert_include(node: child_node, method:) }
+    end
+  end
+
   def line_no
     return @line_no if @line_no
 

spec/fake_app/app/controllers/single_lines_controller.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+class SingleLinesController < ActionController::Base
+  def show
+    @plays = Play.joins(:actors)
+    @plays.each { |play| play.actors.map { |actor| "name: #{actor.name}" } }
+  end
+end

spec/fake_app/config/routes.rb

@@ -2,5 +2,6 @@
 
 BulletmarkRepairerTestApp::Application.routes.draw do
   resources :multiple_lines, only: [:index]
+  resource :single_line, only: [:show]
   resources :previous_lines, only: [:index]
 end

spec/requests/single_lines_spec.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+RSpec.describe SingleLinesController do
+  let(:original_src) do
+    <<-SRC
+  def show
+    @plays = Play.joins(:actors)
+    @plays.each { |play| play.actors.map { |actor| "name: \#{actor.name}" } }
+  end
+    SRC
+  end
+  let(:patched_src) do
+    <<-SRC
+  def show
+    @plays = Play.joins(:actors)
+    @plays.includes([:actors]).each { |play| play.actors.map { |actor| "name: \#{actor.name}" } }
+  end
+    SRC
+  end
+
+  before { create(:play, :blast) }
+
+  subject { get single_line_path }
+
+  it_behaves_like 'correctly patched'
+end