Performance

.rubocop.yml

@@ -55,6 +55,9 @@ Style/IdenticalConditionalBranches:
 Style/IfInsideElse:
   Enabled: false
 
+Style/IfWithBooleanLiteralBranches:
+  Enabled: false
+
 Style/KeywordParametersOrder:
   Enabled: false
 

Gemfile.lock

@@ -2,7 +2,7 @@ PATH
   remote: .
   specs:
     syntax_tree (3.6.3)
-      prettier_print
+      prettier_print (>= 1.0.0)
 
 GEM
   remote: https://rubygems.org/
@@ -14,10 +14,10 @@ GEM
     parallel (1.22.1)
     parser (3.1.2.1)
       ast (~> 2.4.1)
-    prettier_print (0.1.0)
+    prettier_print (1.0.0)
     rainbow (3.1.1)
     rake (13.0.6)
-    regexp_parser (2.5.0)
+    regexp_parser (2.6.0)
     rexml (3.2.5)
     rubocop (1.36.0)
       json (~> 2.3)
@@ -38,7 +38,7 @@ GEM
       simplecov_json_formatter (~> 0.1)
     simplecov-html (0.12.3)
     simplecov_json_formatter (0.1.4)
-    unicode-display_width (2.2.0)
+    unicode-display_width (2.3.0)
 
 PLATFORMS
   arm64-darwin-21

bin/profile

@@ -6,22 +6,21 @@ require "bundler/inline"
 gemfile do
   source "https://rubygems.org"
   gem "stackprof"
+  gem "prettier_print"
 end
 
 $:.unshift(File.expand_path("../lib", __dir__))
 require "syntax_tree"
 
-GC.disable
-
 StackProf.run(mode: :cpu, out: "tmp/profile.dump", raw: true) do
-  filepath = File.expand_path("../lib/syntax_tree/node.rb", __dir__)
-  SyntaxTree.format(File.read(filepath))
+  Dir[File.join(RbConfig::CONFIG["libdir"], "**/*.rb")].each do |filepath|
+    SyntaxTree.format(SyntaxTree.read(filepath))
+  end
 end
 
-GC.enable
-
 File.open("tmp/flamegraph.html", "w") do |file|
   report = Marshal.load(IO.binread("tmp/profile.dump"))
+  StackProf::Report.new(report).print_text
   StackProf::Report.new(report).print_d3_flamegraph(file)
 end
 

lib/syntax_tree.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "delegate"
 require "etc"
 require "json"
 require "pp"
@@ -10,7 +9,6 @@
 
 require_relative "syntax_tree/formatter"
 require_relative "syntax_tree/node"
-require_relative "syntax_tree/parser"
 require_relative "syntax_tree/version"
 
 require_relative "syntax_tree/basic_visitor"
@@ -20,6 +18,18 @@
 require_relative "syntax_tree/visitor/match_visitor"
 require_relative "syntax_tree/visitor/pretty_print_visitor"
 
+require_relative "syntax_tree/parser"
+
+# We rely on Symbol#name being available, which is only available in Ruby 3.0+.
+# In case we're running on an older Ruby version, we polyfill it here.
+unless :+.respond_to?(:name)
+  class Symbol # rubocop:disable Style/Documentation
+    def name
+      to_s.freeze
+    end
+  end
+end
+
 # Syntax Tree is a suite of tools built on top of the internal CRuby parser. It
 # provides the ability to generate a syntax tree from source, as well as the
 # tools necessary to inspect and manipulate that syntax tree. It can be used to

lib/syntax_tree/formatter.rb

@@ -62,21 +62,39 @@ def format(node, stackable: true)
       # If there are comments, then we're going to format them around the node
       # so that they get printed properly.
       if node.comments.any?
-        leading, trailing = node.comments.partition(&:leading?)
+        trailing = []
+        last_leading = nil
 
-        # Print all comments that were found before the node.
-        leading.each do |comment|
-          comment.format(self)
-          breakable(force: true)
+        # First, we're going to print all of the comments that were found before
+        # the node. We'll also gather up any trailing comments that we find.
+        node.comments.each do |comment|
+          if comment.leading?
+            comment.format(self)
+            breakable(force: true)
+            last_leading = comment
+          else
+            trailing << comment
+          end
         end
 
         # If the node has a stree-ignore comment right before it, then we're
         # going to just print out the node as it was seen in the source.
         doc =
-          if leading.last&.ignore?
+          if last_leading&.ignore?
             range = source[node.location.start_char...node.location.end_char]
-            separator = -> { breakable(indent: false, force: true) }
-            seplist(range.split(/\r?\n/, -1), separator) { |line| text(line) }
+            first = true
+
+            range.each_line(chomp: true) do |line|
+              if first
+                first = false
+              else
+                breakable_return
+              end
+
+              text(line)
+            end
+
+            breakable_return if range.end_with?("\n")
           else
             node.format(self)
           end
@@ -101,12 +119,53 @@ def format_each(nodes)
       nodes.each { |node| format(node) }
     end
 
+    def grandparent
+      stack[-3]
+    end
+
     def parent
       stack[-2]
     end
 
     def parents
       stack[0...-1].reverse_each
     end
+
+    # This is a simplified version of prettyprint's group. It doesn't provide
+    # any of the more advanced options because we don't need them and they take
+    # up expensive computation time.
+    def group
+      contents = []
+      doc = Group.new(0, contents: contents)
+
+      groups << doc
+      target << doc
+
+      with_target(contents) { yield }
+      groups.pop
+      doc
+    end
+
+    # A similar version to the super, except that it calls back into the
+    # separator proc with the instance of `self`.
+    def seplist(list, sep = nil, iter_method = :each)
+      first = true
+      list.__send__(iter_method) do |*v|
+        if first
+          first = false
+        elsif sep
+          sep.call(self)
+        else
+          comma_breakable
+        end
+        yield(*v)
+      end
+    end
+
+    # This is a much simplified version of prettyprint's text. It avoids
+    # calculating width by pushing the string directly onto the target.
+    def text(string)
+      target << string
+    end
   end
 end