Improve multiline T.must autocorrection

lib/rubocop/cop/type_toolkit/prefer_not_nil.rb

@@ -54,13 +54,16 @@ def extract_t_must_argument(node)
         #: (RuboCop::AST::Node) -> String
         def replacement_for(argument)
           source = argument.source
+
           source = "(#{source})" if requires_parentheses?(argument)
+
           "#{source}.not_nil!"
         end
 
         #: (RuboCop::AST::SendNode, RuboCop::AST::Node, String) -> String
         def correction_for(node, argument, replacement)
-          return replacement unless node.multiline?
+          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
 
           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

spec/rubocop/cop/type_toolkit/prefer_not_nil_spec.rb

@@ -123,6 +123,19 @@ def explicit_super
           RUBY
         end
 
+        it "autocorrects multiline command calls" do
+          assert_offense(<<~RUBY)
+            T.must foo
+            ^^^^^^^^^^ #{MSG}
+              .bar
+          RUBY
+
+          assert_correction(<<~RUBY)
+            foo
+              .bar.not_nil!
+          RUBY
+        end
+
         it "does not add unnecessary parentheses to parenthesized calls" do
           assert_offense(<<~RUBY)
             value = T.must(fetch(value))
@@ -151,6 +164,19 @@ def explicit_super
           RUBY
         end
 
+        it "does not group an argument that makes the call multiline" do
+          assert_offense(<<~RUBY)
+            T.must(foo
+            ^^^^^^^^^^ #{MSG}
+              .bar)
+          RUBY
+
+          assert_correction(<<~RUBY)
+            foo
+              .bar.not_nil!
+          RUBY
+        end
+
         it "preserves comments in multiline calls" do
           assert_offense(<<~RUBY)
             value = T.must(