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
@@ -87,7 +90,6 @@ def requires_parentheses?(argument)
           end
           return true if argument.range_type? || argument.operator_keyword?
           return true if argument.if_type? || argument.assignment?
-          return true if argument.any_block_type?
 
           KEYWORD_EXPRESSION_TYPES.include?(argument.type)
         end

spec/rubocop/cop/type_toolkit/prefer_not_nil_spec.rb

@@ -77,6 +77,8 @@ class PreferNotNilSpec < ::Minitest::Spec
                          ^^^^^^^^^^^^^^^^^ #{MSG}
             block_value = T.must(foo { bar })
                           ^^^^^^^^^^^^^^^^^^^ #{MSG}
+            block_with_arguments = T.must(foo(arg) { |value| value })
+                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #{MSG}
             defined_value = T.must(defined?(foo))
                             ^^^^^^^^^^^^^^^^^^^^^ #{MSG}
             def example
@@ -98,7 +100,8 @@ def explicit_super
             logical = (foo || bar).not_nil!
             grouped = (foo || bar).not_nil!
             assignment = (foo = bar).not_nil!
-            block_value = (foo { bar }).not_nil!
+            block_value = foo { bar }.not_nil!
+            block_with_arguments = foo(arg) { |value| value }.not_nil!
             defined_value = (defined?(foo)).not_nil!
             def example
               (yield foo).not_nil!
@@ -123,6 +126,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 +167,36 @@ 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 parentheses around a multiline receiver" do
+          assert_offense(<<~RUBY)
+            T.must(
+            ^^^^^^^ #{MSG}
+              foo.
+                bar
+            )
+          RUBY
+
+          assert_correction(<<~RUBY)
+            (
+              foo.
+                bar
+            ).not_nil!
+          RUBY
+        end
+
         it "preserves comments in multiline calls" do
           assert_offense(<<~RUBY)
             value = T.must(