Remove sigs attached to methods defined with modifier calls

lib/spoom/deadcode/remover.rb

@@ -48,8 +48,10 @@ def apply_edit
 
           node = @node_context.node
           case node
-          when Prism::ClassNode, Prism::ModuleNode, Prism::DefNode
+          when Prism::ClassNode, Prism::ModuleNode
             delete_node_and_comments_and_sigs(@node_context)
+          when Prism::DefNode
+            delete_node_and_comments_and_sigs(modifier_call_context(node) || @node_context)
           when Prism::ConstantWriteNode, Prism::ConstantOperatorWriteNode,
                 Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
                 Prism::ConstantPathWriteNode, Prism::ConstantPathOperatorWriteNode,
@@ -65,6 +67,46 @@ def apply_edit
 
         private
 
+        # When a method is defined as the argument of a modifier call (e.g. `private def foo; end` or
+        # `private_class_method def self.foo; end`), the `def` node's parent is the call rather than the
+        # enclosing class or module body, so its sigs and comments are siblings of the call. Return a
+        # context targeting the outermost such call so they are removed together with the method.
+        #
+        # Modifiers are matched structurally rather than from a fixed list, so user-defined ones are
+        # handled too. A call only counts as a modifier when the `def` is its sole argument and it takes
+        # no block, so we never remove a call that does more than wrap the method (e.g.
+        # `register(:thing, def foo; end)`).
+        #: (Prism::DefNode def_node) -> NodeContext?
+        def modifier_call_context(def_node)
+          wrapped = def_node #: Prism::Node
+          nesting = @node_context.nesting.dup
+          outer_call = nil #: Prism::CallNode?
+          outer_nesting = nil #: Array[Prism::Node]?
+
+          while (ancestor = nesting.pop)
+            case ancestor
+            when Prism::ArgumentsNode
+              # An arguments node sits between a call and its arguments, keep climbing
+              next
+            when Prism::CallNode
+              # Stop unless the call wraps the node as its sole argument and takes no block, otherwise
+              # it does more than wrap the method and we must not remove it
+              args = ancestor.arguments&.arguments
+              break if ancestor.block
+              break unless args && args.size == 1 && args.first.equal?(wrapped)
+
+              wrapped = ancestor
+              outer_call = ancestor
+              outer_nesting = nesting.dup
+            else
+              break
+            end
+          end
+          return unless outer_call && outer_nesting
+
+          NodeContext.new(@old_source, @node_context.comments, outer_call, outer_nesting)
+        end
+
         #: (NodeContext context) -> void
         def delete_constant_assignment(context)
           case context.node

rbi/spoom.rbi

@@ -1519,6 +1519,9 @@ class Spoom::Deadcode::Remover::NodeRemover
   end
   def insert_accessor(node, send_context, was_removed:); end
 
+  sig { returns(T.nilable(::Spoom::Deadcode::Remover::NodeContext)) }
+  def modifier_call_context; end
+
   sig { params(start_char: ::Integer, end_char: ::Integer, replacement: ::String).void }
   def replace_chars(start_char, end_char, replacement); end
 

test/spoom/deadcode/remover_test.rb

@@ -940,6 +940,186 @@ def baz; end
         RB
       end
 
+      def test_removes_node_sig_with_visibility_modifier
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+
+            sig { void }
+            private_class_method def self.bar
+              something
+            end
+
+            def baz; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+
+            def baz; end
+          end
+        RB
+      end
+
+      def test_removes_node_sig_with_custom_modifier
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+
+            # A comment
+            sig { void }
+            some_custom_modifier def bar
+              something
+            end
+
+            def baz; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+
+            def baz; end
+          end
+        RB
+      end
+
+      def test_removes_node_sig_with_nested_modifiers
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+
+            sig { void }
+            private abstract def bar
+              something
+            end
+
+            def baz; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+
+            def baz; end
+          end
+        RB
+      end
+
+      def test_removes_node_sig_with_modifier_as_last_statement
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+
+            sig { void }
+            private def bar
+              something
+            end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+          end
+        RB
+      end
+
+      def test_removes_single_line_def_with_modifier
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+            private def bar; end
+            def baz; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+            def baz; end
+          end
+        RB
+      end
+
+      def test_removes_def_with_modifier_without_sig
+        res = remove(<<~RB, "bar")
+          class Foo
+            def foo; end
+
+            private_class_method def self.bar
+              something
+            end
+
+            def baz; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            def foo; end
+
+            def baz; end
+          end
+        RB
+      end
+
+      def test_keeps_call_when_def_is_not_the_sole_argument
+        # `bar` is a positional argument to `register`, not wrapped by a modifier, so only the `def`
+        # is removed and the call (along with its other arguments) is preserved.
+        res = remove(<<~RB, "bar")
+          class Foo
+            register(
+              :thing,
+              def bar
+                x
+              end,
+            )
+
+            def keep; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            register(
+              :thing,
+            )
+
+            def keep; end
+          end
+        RB
+      end
+
+      def test_keeps_sibling_arguments_when_def_is_not_the_sole_argument
+        res = remove(<<~RB, "bar")
+          class Foo
+            memoize(
+              def bar
+                x
+              end,
+              ttl: 60,
+            )
+
+            def keep; end
+          end
+        RB
+
+        assert_equal(<<~RB, res)
+          class Foo
+            memoize(
+              ttl: 60,
+            )
+
+            def keep; end
+          end
+        RB
+      end
+
       def test_removes_node_sig_and_comments
         res = remove(<<~RB, "bar")
           class Foo