Speed up

Gemfile

@@ -7,3 +7,4 @@ gemspec
 gem 'benchmark-ips'
 gem 'parser'
 gem 'ruby_parser'
+gem 'stackprof'

Gemfile.lock

@@ -22,6 +22,7 @@ GEM
       simplecov_json_formatter (~> 0.1)
     simplecov-html (0.12.3)
     simplecov_json_formatter (0.1.3)
+    stackprof (0.2.17)
 
 PLATFORMS
   x86_64-darwin-19
@@ -35,6 +36,7 @@ DEPENDENCIES
   rake
   ruby_parser
   simplecov
+  stackprof
   syntax_tree!
 
 BUNDLED WITH

bin/profile

@@ -0,0 +1,16 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require 'bundler/setup'
+require 'stackprof'
+
+filepath = File.expand_path('../lib/syntax_tree', __dir__)
+require_relative filepath
+
+GC.disable
+
+StackProf.run(mode: :cpu, out: "tmp/profile.dump", raw: true) do
+  SyntaxTree.format(File.read("#{filepath}.rb"))
+end
+
+GC.enable

lib/syntax_tree/prettyprint.rb

@@ -540,37 +540,6 @@ def indent(part = IndentPart)
     end
   end
 
-  # This is a visitor that can be passed to PrettyPrint.visit that will
-  # propagate BreakParent nodes all of the way up the tree. When a BreakParent
-  # is encountered, it will break the surrounding group, and then that group
-  # will break its parent, and so on.
-  class PropagateBreaksVisitor
-    attr_reader :groups, :visited
-
-    def initialize
-      @groups = []
-      @visited = []
-    end
-
-    def on_enter(doc)
-      case doc
-      when BreakParent
-        groups.last&.break
-      when Group
-        groups << doc
-        return false if visited.include?(doc)
-
-        visited << doc
-      end
-
-      true
-    end
-
-    def on_exit(doc)
-      groups.last&.break if doc.is_a?(Group) && groups.pop.break?
-    end
-  end
-
   # When printing, you can optionally specify the value that should be used
   # whenever a group needs to be broken onto multiple lines. In this case the
   # default is \n.
@@ -751,7 +720,7 @@ def flush
     # First, ensure that we've propagated all of the necessary break-parent
     # nodes throughout the tree.
     doc = groups.first
-    PrettyPrint.visit(doc, PropagateBreaksVisitor.new)
+    propagate_breaks(doc)
 
     # This represents how far along the current line we are. It gets reset
     # back to 0 when we encounter a newline.
@@ -1132,6 +1101,37 @@ def fits?(next_command, rest_commands, remaining)
     false
   end
 
+  # This will walk the doc tree and propagate BreakParent nodes all of the way
+  # up the tree. When a BreakParent is encountered, it will break the
+  # surrounding group, and then that group will break its parent, and so on.
+  def propagate_breaks(doc)
+    marker = Object.new
+    stack = [doc]
+    groups = []
+
+    while doc = stack.pop
+      if doc == marker
+        groups.last&.break if stack.pop.is_a?(Group) && groups.pop.break?
+      else
+        stack += [doc, marker]
+        enter = false
+
+        case doc
+        when BreakParent
+          groups.last&.break
+        when Group
+          groups << doc
+          stack += doc.contents.reverse
+        when IfBreak
+          stack += doc.break_contents.reverse
+          stack += doc.flat_contents.reverse
+        when Align, Indent, LineSuffix
+          stack += doc.contents.reverse
+        end
+      end
+    end
+  end
+
   # Resets the group stack and target array so that this pretty printer object
   # can continue to be used before calling flush again if desired.
   def reset