Avoid unnecessary multiline parentheses

lib/rubocop/cop/type_toolkit/prefer_not_nil.rb

@@ -63,7 +63,7 @@ def replacement_for(argument)
         #: (RuboCop::AST::SendNode, RuboCop::AST::Node, String) -> String
         def correction_for(node, argument, replacement)
           return replacement unless node.multiline? && node.parenthesized_call?
-          return replacement if argument.first_line == node.loc.begin.line && argument.last_line == node.loc.end.line
+          return replacement unless comments_inside_parentheses?(node) || contains_heredoc?(argument)
 
           grouped_range = node.source_range.with(begin_pos: node.loc.begin.begin_pos, end_pos: node.loc.end.end_pos)
           grouped_source = grouped_range.source
@@ -72,6 +72,26 @@ def correction_for(node, argument, replacement)
           "#{grouped_source}.not_nil!"
         end
 
+        #: (RuboCop::AST::Node) -> bool
+        def contains_heredoc?(node)
+          return true if node.loc.is_a?(Parser::Source::Map::Heredoc)
+
+          node.each_descendant(:any_str).any? do |descendant|
+            descendant.loc.is_a?(Parser::Source::Map::Heredoc)
+          end
+        end
+
+        #: (RuboCop::AST::SendNode) -> bool
+        def comments_inside_parentheses?(node)
+          contents_begin = node.loc.begin.end_pos
+          contents_end = node.loc.end.begin_pos
+
+          processed_source.comments.any? do |comment|
+            comment_range = comment.loc.expression
+            contents_begin <= comment_range.begin_pos && comment_range.end_pos <= contents_end
+          end
+        end
+
         #: (RuboCop::AST::SendNode) -> bool
         def nested_t_must?(node)
           node.each_ancestor(:send).any? do |ancestor|

spec/rubocop/cop/type_toolkit/prefer_not_nil_spec.rb

@@ -180,7 +180,7 @@ def explicit_super
           RUBY
         end
 
-        it "preserves parentheses around a multiline receiver" do
+        it "preserves a multiline receiver without parentheses" do
           assert_offense(<<~RUBY)
             T.must(
             ^^^^^^^ #{MSG}
@@ -190,10 +190,8 @@ def explicit_super
           RUBY
 
           assert_correction(<<~RUBY)
-            (
-              foo.
-                bar
-            ).not_nil!
+            foo.
+                bar.not_nil!
           RUBY
         end
 
@@ -214,7 +212,7 @@ def explicit_super
           RUBY
         end
 
-        it "autocorrects multiline calls with whitespace before the method" do
+        it "drops unnecessary grouping from multiline calls without comments" do
           assert_offense(<<~RUBY)
             first = T .must(
                     ^^^^^^^^ #{MSG}
@@ -228,12 +226,21 @@ def explicit_super
           RUBY
 
           assert_correction(<<~RUBY)
-            first = (
-              foo
-            ).not_nil!
-            second = (
-                bar
-              ).not_nil!
+            first = foo.not_nil!
+            second = bar.not_nil!
+          RUBY
+        end
+
+        it "preserves a multiline call chain without parentheses" do
+          assert_offense(<<~RUBY)
+            variant = T.must(
+                      ^^^^^^^ #{MSG}
+              InventoryItemVariant.preload(:variant).where(inventory_item_id: @inventory_item_id).first!.variant
+            )
+          RUBY
+
+          assert_correction(<<~RUBY)
+            variant = InventoryItemVariant.preload(:variant).where(inventory_item_id: @inventory_item_id).first!.variant.not_nil!
           RUBY
         end
 
@@ -252,6 +259,67 @@ def explicit_super
           RUBY
         end
 
+        it "preserves grouping around heredocs in multiline calls" do
+          assert_offense(<<~RUBY)
+            value = T.must(
+                    ^^^^^^^ #{MSG}
+              <<~TEXT
+                hello
+              TEXT
+            )
+          RUBY
+
+          assert_correction(<<~RUBY)
+            value = (
+              <<~TEXT
+                hello
+              TEXT
+            ).not_nil!
+          RUBY
+        end
+
+        it "preserves grouping around nested heredocs in multiline calls" do
+          assert_offense(<<~RUBY)
+            value = T.must(
+                    ^^^^^^^ #{MSG}
+              foo(<<~TEXT)
+                hello
+              TEXT
+            )
+          RUBY
+
+          assert_correction(<<~RUBY)
+            value = (
+              foo(<<~TEXT)
+                hello
+              TEXT
+            ).not_nil!
+          RUBY
+        end
+
+        it "preserves grouping around a heredoc T.must nested in another call" do
+          assert_offense(<<~RUBY)
+            value = foo(
+              T.must(
+              ^^^^^^^ #{MSG}
+                <<~TEXT
+                  hello
+                TEXT
+              )
+            )
+          RUBY
+
+          assert_correction(<<~RUBY)
+            value = foo(
+              (
+                <<~TEXT
+                  hello
+                TEXT
+              ).not_nil!
+            )
+          RUBY
+        end
+
         it "autocorrects nested T.must calls" do
           assert_offense(<<~RUBY)
             value = T.must(T.must(foo).bar)