Peak memory use optimizations.

Reduce peak memory usage of swift-format by discarding collections once they are no longer needed.
- Remove `Token` arrays from `beforeMap` and `afterMap` while processing the associated token.
- Empty `preVisitedExprs` after visitation finishes for the root expression of the tree.

By the end of a file, these maps can accumulate a meaningful number of tokens such that they significantly impact the peak memory usage.

This admittedly replaces a O(1) lookup operation with a O(N) lookup-and-remove operation, by replacing `subscript` with `removeValue(forKey:)`. In my testing, I found that this had no impact on the runtime. My hypothesis is that removing from these dictionaries keeps them small enough that O(N) runtime isn't a problem.

This set previously contained `ExprSyntax` objects and was maintained for the entire file. Once the visitor has finished with all children of the node that initiated inserting the contextual breaks, it will never visit any of the children again. This set can be emptied at that point to keep it from growing until the visitor finishes with the whole file.

The runtime is unchanged, but this reduces peak memory usage especially for large files.

- Large file (30k LOC): 108.3 MB peak to 72.1 MB peak
- Medium file (10k LOC) : 35.9 MB peak to 24.4 MB peak
- swift-protobuf library (many small files): 77.6 MB peak to 75.6 MB peak

Author: dylansturg

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -43,7 +43,11 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
   /// Lists the expressions that have been visited, from the outermost expression, where contextual
   /// breaks and start/end contextual breaking tokens have been inserted.
-  private var preVisitedExprs = Set<ExprSyntax>()
+  private var preVisitedExprs = Set<SyntaxIdentifier>()
+
+  /// Tracks the "root" exprs where previsiting for contextual breaks started so that
+  /// `preVisitedExprs` can be emptied after exiting an expr tree.
+  private var rootExprs = Set<SyntaxIdentifier>()
 
   /// Lists the tokens that are the closing or final delimiter of a node that shouldn't be split
   /// from the preceding token. When breaks are inserted around compound expressions, the breaks are
@@ -850,6 +854,10 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visitPost(_ node: MemberAccessExprSyntax) {
+    clearContextualBreakState(node)
+  }
+
   override func visit(_ node: FunctionCallExprSyntax) -> SyntaxVisitorContinueKind {
     preVisitInsertingContextualBreaks(node)
 
@@ -881,6 +889,10 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visitPost(_ node: FunctionCallExprSyntax) {
+    clearContextualBreakState(node)
+  }
+
   /// Arrange the given argument list (or equivalently, tuple expression list) as a list of function
   /// arguments.
   ///
@@ -1076,6 +1088,10 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visitPost(_ node: SubscriptExprSyntax) {
+    clearContextualBreakState(node)
+  }
+
   override func visit(_ node: ExpressionSegmentSyntax) -> SyntaxVisitorContinueKind {
     // TODO: For now, just use the raw text of the node and don't try to format it deeper. In the
     // future, we should find a way to format the expression but without wrapping so that at least
@@ -2271,7 +2287,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
   /// Appends the before-tokens of the given syntax token to the token stream.
   private func appendBeforeTokens(_ token: TokenSyntax) {
-    if let before = beforeMap[token] {
+    if let before = beforeMap.removeValue(forKey: token) {
       before.forEach(appendToken)
     }
   }
@@ -2331,7 +2347,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   ///   not incorrectly inserted into the token stream *after* a break or newline.
   private func appendAfterTokensAndTrailingComments(_ token: TokenSyntax) {
     let (wasLineComment, trailingCommentTokens) = afterTokensForTrailingComment(token)
-    let afterGroups = afterMap[token] ?? []
+    let afterGroups = afterMap.removeValue(forKey: token) ?? []
     var hasAppendedTrailingComment = false
 
     if !wasLineComment {
@@ -3177,23 +3193,34 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     appendToken(.break(.same, size: 0))
   }
 
+  /// Cleans up state related to inserting contextual breaks throughout expressions during
+  /// `visitPost` for an expression that is the root of an expression tree.
+  private func clearContextualBreakState<T: ExprSyntaxProtocol>(_ expr: T) {
+    let exprID = expr.id
+    if rootExprs.remove(exprID) != nil {
+      preVisitedExprs.removeAll()
+    }
+  }
+
   /// Visits the given expression node and all of the nested expression nodes, inserting tokens
   /// necessary for contextual breaking throughout the expression. Records the nodes that were
   /// visited so that they can be skipped later.
   private func preVisitInsertingContextualBreaks<T: ExprSyntaxProtocol & Equatable>(_ expr: T) {
-    let exprSyntax = ExprSyntax(expr)
-    if !preVisitedExprs.contains(exprSyntax) {
-      let (visited, _, _) = insertContextualBreaks(exprSyntax, isTopLevel: true)
-      preVisitedExprs.formUnion(visited)
+    let exprID = expr.id
+    if !preVisitedExprs.contains(exprID) {
+      rootExprs.insert(exprID)
+      insertContextualBreaks(ExprSyntax(expr), isTopLevel: true)
     }
   }
 
   /// Recursively visits nested expressions from the given expression inserting contextual breaking
   /// tokens. When visiting an expression node, `preVisitInsertingContextualBreaks(_:)` should be
   /// called instead of this helper.
+  @discardableResult
   private func insertContextualBreaks(_ expr: ExprSyntax, isTopLevel: Bool) -> (
-    [ExprSyntax], hasCompoundExpression: Bool, hasMemberAccess: Bool
+    hasCompoundExpression: Bool, hasMemberAccess: Bool
   ) {
+    preVisitedExprs.insert(expr.id)
     if let memberAccessExpr = expr.as(MemberAccessExprSyntax.self) {
       // When the member access is part of a calling expression, the break before the dot is
       // inserted when visiting the parent node instead so that the break is inserted before any
@@ -3202,21 +3229,18 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
           expr.parent?.isProtocol(CallingExprSyntaxProtocol.self) != true {
         before(memberAccessExpr.dot, tokens: .break(.contextual, size: 0))
       }
-      let children: [ExprSyntax]
       var hasCompoundExpression = false
       if let base = memberAccessExpr.base {
-        (children, hasCompoundExpression, _) = insertContextualBreaks(base, isTopLevel: false)
-      } else {
-        children = []
+        (hasCompoundExpression, _) = insertContextualBreaks(base, isTopLevel: false)
       }
       if isTopLevel {
         before(expr.firstToken, tokens: .contextualBreakingStart)
         after(expr.lastToken, tokens: .contextualBreakingEnd)
       }
-      return ([expr] + children, hasCompoundExpression, true)
+      return (hasCompoundExpression, true)
     } else if let callingExpr = expr.asProtocol(CallingExprSyntaxProtocol.self) {
       let calledExpression = callingExpr.calledExpression
-      let (children, hasCompoundExpression, hasMemberAccess) =
+      let (hasCompoundExpression, hasMemberAccess) =
         insertContextualBreaks(calledExpression, isTopLevel: false)
 
       let shouldGroup =
@@ -3241,7 +3265,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
         before(expr.firstToken, tokens: beforeTokens)
         after(expr.lastToken, tokens: afterTokens)
       }
-      return ([expr] + children, true, hasMemberAccess)
+      return (true, hasMemberAccess)
     }
 
     // Otherwise, it's an expression that isn't calling another expression (e.g. array or
@@ -3250,7 +3274,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     before(expr.firstToken, tokens: .contextualBreakingStart)
     after(expr.lastToken, tokens: .contextualBreakingEnd)
     let hasCompoundExpression = !expr.is(IdentifierExprSyntax.self)
-    return ([expr], hasCompoundExpression, false)
+    return (hasCompoundExpression, false)
   }
 }