From 9e89484e088475d5e2560f3cbce8b9e9fb6e90ff Mon Sep 17 00:00:00 2001
From: Manuel Wudka-Robles <mwudka@gmail.com>
Date: Thu, 22 Dec 2022 17:32:55 -0800
Subject: [PATCH] Do not change visibility when sorting

The SortNodes visitor also sorts Private and Protected nodes. This can cause
the visibility of methods to change. For example, a private method called
"a_thing" will always be moved before the private specifier, thus becoming
public. This change extends the SortNodes visitor to individually sort the
public, private, and protected nodes to preserve their visibility.
---
 lib/rbi/rewriters/sort_nodes.rb       | 29 ++++++++++++++++++---------
 test/rbi/rewriters/sort_nodes_test.rb | 25 +++++++++++++++++++++++
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/lib/rbi/rewriters/sort_nodes.rb b/lib/rbi/rewriters/sort_nodes.rb
index 2374883c..fe6e8cb7 100644
--- a/lib/rbi/rewriters/sort_nodes.rb
+++ b/lib/rbi/rewriters/sort_nodes.rb
@@ -13,18 +13,29 @@ def visit(node)
         return unless node.is_a?(Tree)
         visit_all(node.nodes)
         original_order = node.nodes.map.with_index.to_h
-        node.nodes.sort! do |a, b|
-          # First we try to compare the nodes by their node rank (based on the node type)
-          res = node_rank(a) <=> node_rank(b)
-          next res if res != 0 # we can sort the nodes by their rank, let's stop here
 
-          # Then, if the nodes ranks are the same (res == 0), we try to compare the nodes by their name
-          res = node_name(a) <=> node_name(b)
-          next res if res && res != 0 # we can sort the nodes by their name, let's stop here
+        # The child nodes could contain private/protected markers. If so, they should not be moved in the file.
+        # Otherwise, some methods could see their privacy change. To avoid that problem, divide the array of child
+        # nodes into chunks based on whether any Visibility nodes appear, and sort the chunks independently. This
+        # applies the ordering rules from the node_rank method as much as possible, while preserving visibility.
+        sorted_nodes = node.nodes.chunk do |n|
+          n.is_a?(Visibility)
+        end.flat_map do |_, nodes|
+          nodes.sort! do |a, b|
+            # First we try to compare the nodes by their node rank (based on the node type)
+            res = node_rank(a) <=> node_rank(b)
+            next res if res != 0 # we can sort the nodes by their rank, let's stop here
 
-          # Finally, if the two nodes have the same rank and the same name or at least one node is anonymous then,
-          T.must(original_order[a]) <=> T.must(original_order[b]) # we keep the original order
+            # Then, if the nodes ranks are the same (res == 0), we try to compare the nodes by their name
+            res = node_name(a) <=> node_name(b)
+            next res if res && res != 0 # we can sort the nodes by their name, let's stop here
+
+            # Finally, if the two nodes have the same rank and the same name or at least one node is anonymous then,
+            T.must(original_order[a]) <=> T.must(original_order[b]) # we keep the original order
+          end
         end
+
+        node.nodes.replace(sorted_nodes)
       end
 
       private
diff --git a/test/rbi/rewriters/sort_nodes_test.rb b/test/rbi/rewriters/sort_nodes_test.rb
index 50659e92..27a95eef 100644
--- a/test/rbi/rewriters/sort_nodes_test.rb
+++ b/test/rbi/rewriters/sort_nodes_test.rb
@@ -310,5 +310,30 @@ class D < ::T::Struct; end
         class E; end
       RBI
     end
+
+    def test_sort_doesnt_change_privacy
+      rbi = Tree.new
+      rbi << Public.new
+      rbi << Method.new("c") # 0
+      rbi << Private.new     # 1
+      rbi << Method.new("a") # 2
+      rbi << Protected.new   # 3
+      rbi << Method.new("b") # 4
+      rbi << Public.new
+      rbi << Method.new("aa")
+
+      rbi.sort_nodes!
+
+      assert_equal(<<~RBI, rbi.string)
+        public
+        def c; end
+        private
+        def a; end
+        protected
+        def b; end
+        public
+        def aa; end
+      RBI
+    end
   end
 end