From 606114d1a031017b67f831d05bf9209ebff2a403 Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Tue, 6 Aug 2024 12:17:14 -0400
Subject: [PATCH] Introduce NestTopLevelMembers rewriter

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 lib/rbi.rb                                    |  1 +
 lib/rbi/rewriters/nest_top_level_members.rb   | 68 ++++++++++++++
 .../rewriters/nest_top_level_members_test.rb  | 89 +++++++++++++++++++
 3 files changed, 158 insertions(+)
 create mode 100644 lib/rbi/rewriters/nest_top_level_members.rb
 create mode 100644 test/rbi/rewriters/nest_top_level_members_test.rb

diff --git a/lib/rbi.rb b/lib/rbi.rb
index 69b53166..449a51b8 100644
--- a/lib/rbi.rb
+++ b/lib/rbi.rb
@@ -22,6 +22,7 @@ class Error < StandardError
 require "rbi/rewriters/merge_trees"
 require "rbi/rewriters/nest_singleton_methods"
 require "rbi/rewriters/nest_non_public_methods"
+require "rbi/rewriters/nest_top_level_members"
 require "rbi/rewriters/group_nodes"
 require "rbi/rewriters/remove_known_definitions"
 require "rbi/rewriters/attr_to_methods"
diff --git a/lib/rbi/rewriters/nest_top_level_members.rb b/lib/rbi/rewriters/nest_top_level_members.rb
new file mode 100644
index 00000000..709da2ae
--- /dev/null
+++ b/lib/rbi/rewriters/nest_top_level_members.rb
@@ -0,0 +1,68 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RBI
+  module Rewriters
+    # This rewriter moves top-level members into a top-level Object class
+    #
+    # Example:
+    # ~~~rb
+    # def foo; end
+    # attr_reader :bar
+    # ~~~
+    #
+    # will be rewritten to:
+    #
+    # ~~~rb
+    # class Object
+    #  def foo; end
+    #  attr_reader :bar
+    # end
+    # ~~~
+    class NestTopLevelMembers < Visitor
+      extend T::Sig
+
+      sig { void }
+      def initialize
+        super
+
+        @top_level_object_class = T.let(nil, T.nilable(Class))
+      end
+
+      sig { override.params(node: T.nilable(Node)).void }
+      def visit(node)
+        return unless node
+
+        case node
+        when Tree
+          visit_all(node.nodes.dup)
+        else
+          scope = node.parent_scope
+          unless scope
+            parent = node.parent_tree
+            raise unless parent
+
+            node.detach
+
+            unless @top_level_object_class
+              @top_level_object_class = Class.new("Object")
+              parent.nodes << @top_level_object_class
+            end
+
+            @top_level_object_class << node
+          end
+        end
+      end
+    end
+  end
+
+  class Tree
+    extend T::Sig
+
+    sig { void }
+    def nest_top_level_members!
+      visitor = Rewriters::NestTopLevelMembers.new
+      visitor.visit(self)
+    end
+  end
+end
diff --git a/test/rbi/rewriters/nest_top_level_members_test.rb b/test/rbi/rewriters/nest_top_level_members_test.rb
new file mode 100644
index 00000000..74f44d2d
--- /dev/null
+++ b/test/rbi/rewriters/nest_top_level_members_test.rb
@@ -0,0 +1,89 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RBI
+  class NestTopLevelMembersTest < Minitest::Test
+    include TestHelper
+
+    def test_nest_top_level_members
+      tree = parse_rbi(<<~RBI)
+        module Foo
+          def m1; end
+        end
+
+        def m2; end
+        def self.m3; end
+        attr_reader :foo
+        send!
+
+        class << self; end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        module Foo
+          def m1; end
+        end
+
+        class << self; end
+
+        class Object
+          def m2; end
+          def self.m3; end
+          attr_reader :foo
+          send!
+        end
+      RBI
+    end
+
+    def test_nest_top_level_members_reuse_the_same_class
+      tree = parse_rbi(<<~RBI)
+        def foo; end
+        class Foo; end
+        def bar; end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        class Foo; end
+
+        class Object
+          def foo; end
+          def bar; end
+        end
+      RBI
+    end
+
+    def test_nest_top_level_members_duplicate_existing_object_class
+      tree = parse_rbi(<<~RBI)
+        class Object
+          def foo; end
+        end
+        def bar; end
+        class Object
+          def baz; end
+        end
+      RBI
+
+      tree.nest_top_level_members!
+
+      assert_equal(<<~RBI, tree.string)
+        class Object
+          def foo; end
+        end
+
+        class Object
+          def baz; end
+        end
+
+        class Object
+          def bar; end
+        end
+      RBI
+    end
+  end
+end