From cbc820836a3e6c15fa2517153018cca1a14c56c7 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       | 39 ++++++++++++++++++++-------
 test/rbi/rewriters/sort_nodes_test.rb | 20 ++++++++++++++
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/lib/rbi/rewriters/sort_nodes.rb b/lib/rbi/rewriters/sort_nodes.rb
index 2374883c..37f5d21d 100644
--- a/lib/rbi/rewriters/sort_nodes.rb
+++ b/lib/rbi/rewriters/sort_nodes.rb
@@ -13,17 +13,38 @@ 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 the nodes into subranges bounded by
+        # array boundary or private/protected nodes, and sort them individually.
+        private_protected_indices = [-1]
+        private_protected_indices += node.nodes.each_with_index.select do |node, idx| 
+          node.is_a?(Private) || node.is_a?(Protected)
+        end.map(&:last)
+        private_protected_indices << node.nodes.length
 
-          # 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
+        private_protected_indices.each_cons(2) do |sort_start_range, sort_end_range|
+          # The start index will be the private/protected node which should not be moved (or -1 for the first element), 
+          # so add 1 to the index so that it doesn't get moved.
+          sort_start_range = T.must(sort_start_range) + 1
+          sort_end_range = T.must(sort_end_range)
+          # The end index will be the next private/protected node or the length array, so the range should be exclusive to avoid touching
+          # the private/protected marker or going off the end of the array
+          range_to_sort = (sort_start_range...sort_end_range)
+
+          # Sort just the current range in place.
+          node.nodes[range_to_sort] = T.must(node.nodes[range_to_sort]).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
+  
+            # 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
       end
 
diff --git a/test/rbi/rewriters/sort_nodes_test.rb b/test/rbi/rewriters/sort_nodes_test.rb
index 50659e92..26902617 100644
--- a/test/rbi/rewriters/sort_nodes_test.rb
+++ b/test/rbi/rewriters/sort_nodes_test.rb
@@ -310,5 +310,25 @@ class D < ::T::Struct; end
         class E; end
       RBI
     end
+
+    def test_sort_doesnt_change_privacy
+      rbi = Tree.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.sort_nodes!
+
+      assert_equal(<<~RBI, rbi.string)
+        def c; end
+        private
+        def a; end
+        protected
+        def b; end
+      RBI
+    end
+
   end
 end