Merge pull request swiftlang#199 from dylansturg/less_memory_makes_me_happy

Peak memory use optimizations.

Author: allevato

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)
  }
 }