1111//===----------------------------------------------------------------------===//
1212
1313import SwiftSyntax
14+ import Foundation
1415
1516/// PrettyPrinter takes a Syntax node and outputs a well-formatted, re-indented reproduction of the
1617/// code as a String.
@@ -66,6 +67,19 @@ public class PrettyPrinter {
6667 private var configuration : Configuration { return context. configuration }
6768 private let maxLineLength : Int
6869 private var tokens : [ Token ]
70+ private var source : String
71+
72+ /// keep track of where formatting was disabled in the original source
73+ ///
74+ /// To format a selection, we insert `enableFormatting`/`disableFormatting` tokens into the
75+ /// stream when entering/exiting a selection range. Those tokens include utf8 offsets into the
76+ /// original source. When enabling formatting, we copy the text between `disabledPosition` and the
77+ /// current position to `outputBuffer`. From then on, we continue to format until the next
78+ /// `disableFormatting` token.
79+ private var disabledPosition : AbsolutePosition ? = nil
80+ /// true if we're currently formatting
81+ private var writingIsEnabled : Bool { disabledPosition == nil }
82+
6983 private var outputBuffer : String = " "
7084
7185 /// The number of spaces remaining on the current line.
@@ -172,11 +186,14 @@ public class PrettyPrinter {
172186 /// - printTokenStream: Indicates whether debug information about the token stream should be
173187 /// printed to standard output.
174188 /// - whitespaceOnly: Whether only whitespace changes should be made.
175- public init ( context: Context , node: Syntax , printTokenStream: Bool , whitespaceOnly: Bool ) {
189+ public init ( context: Context , source : String , node: Syntax , printTokenStream: Bool , whitespaceOnly: Bool ) {
176190 self . context = context
191+ self . source = source
177192 let configuration = context. configuration
178193 self . tokens = node. makeTokenStream (
179- configuration: configuration, operatorTable: context. operatorTable)
194+ configuration: configuration,
195+ selection: context. selection,
196+ operatorTable: context. operatorTable)
180197 self . maxLineLength = configuration. lineLength
181198 self . spaceRemaining = self . maxLineLength
182199 self . printTokenStream = printTokenStream
@@ -216,7 +233,9 @@ public class PrettyPrinter {
216233 }
217234
218235 guard numberToPrint > 0 else { return }
219- writeRaw ( String ( repeating: " \n " , count: numberToPrint) )
236+ if writingIsEnabled {
237+ writeRaw ( String ( repeating: " \n " , count: numberToPrint) )
238+ }
220239 lineNumber += numberToPrint
221240 isAtStartOfLine = true
222241 consecutiveNewlineCount += numberToPrint
@@ -238,13 +257,17 @@ public class PrettyPrinter {
238257 /// leading spaces that are required before the text itself.
239258 private func write( _ text: String ) {
240259 if isAtStartOfLine {
241- writeRaw ( currentIndentation. indentation ( ) )
260+ if writingIsEnabled {
261+ writeRaw ( currentIndentation. indentation ( ) )
262+ }
242263 spaceRemaining = maxLineLength - currentIndentation. length ( in: configuration)
243264 isAtStartOfLine = false
244- } else if pendingSpaces > 0 {
265+ } else if pendingSpaces > 0 && writingIsEnabled {
245266 writeRaw ( String ( repeating: " " , count: pendingSpaces) )
246267 }
247- writeRaw ( text)
268+ if writingIsEnabled {
269+ writeRaw ( text)
270+ }
248271 consecutiveNewlineCount = 0
249272 pendingSpaces = 0
250273 }
@@ -523,7 +546,9 @@ public class PrettyPrinter {
523546 }
524547
525548 case . verbatim( let verbatim) :
526- writeRaw ( verbatim. print ( indent: currentIndentation) )
549+ if writingIsEnabled {
550+ writeRaw ( verbatim. print ( indent: currentIndentation) )
551+ }
527552 consecutiveNewlineCount = 0
528553 pendingSpaces = 0
529554 lastBreak = false
@@ -569,6 +594,40 @@ public class PrettyPrinter {
569594 write ( " , " )
570595 spaceRemaining -= 1
571596 }
597+
598+ case . enableFormatting( let enabledPosition) :
599+ // if we're not disabled, we ignore the token
600+ if let disabledPosition {
601+ let start = source. utf8. index ( source. utf8. startIndex, offsetBy: disabledPosition. utf8Offset)
602+ let end : String . Index
603+ if let enabledPosition {
604+ end = source. utf8. index ( source. utf8. startIndex, offsetBy: enabledPosition. utf8Offset)
605+ } else {
606+ end = source. endIndex
607+ }
608+ var text = String ( source [ start..< end] )
609+ // strip trailing whitespace so that the next formatting can add the right amount
610+ if let nonWhitespace = text. rangeOfCharacter (
611+ from: CharacterSet . whitespaces. inverted, options: . backwards) {
612+ text = String ( text [ ..< nonWhitespace. upperBound] )
613+ }
614+
615+ writeRaw ( text)
616+ if text. hasSuffix ( " \n " ) {
617+ isAtStartOfLine = true
618+ consecutiveNewlineCount = 1
619+ } else {
620+ isAtStartOfLine = false
621+ consecutiveNewlineCount = 0
622+ }
623+ self . disabledPosition = nil
624+ }
625+
626+ case . disableFormatting( let newPosition) :
627+ // a second disable is ignored
628+ if writingIsEnabled {
629+ disabledPosition = newPosition
630+ }
572631 }
573632 }
574633
@@ -673,6 +732,10 @@ public class PrettyPrinter {
673732 let length = isSingleElement ? 0 : 1
674733 total += length
675734 lengths. append ( length)
735+
736+ case . enableFormatting, . disableFormatting:
737+ // no effect on length calculations
738+ lengths. append ( 0 )
676739 }
677740 }
678741
@@ -775,6 +838,14 @@ public class PrettyPrinter {
775838 case . contextualBreakingEnd:
776839 printDebugIndent ( )
777840 print ( " [END BREAKING CONTEXT Idx: \( idx) ] " )
841+
842+ case . enableFormatting( let pos) :
843+ printDebugIndent ( )
844+ print ( " [ENABLE FORMATTING utf8 offset: \( String ( describing: pos) ) ] " )
845+
846+ case . disableFormatting( let pos) :
847+ printDebugIndent ( )
848+ print ( " [DISABLE FORMATTING utf8 offset: \( pos) ] " )
778849 }
779850 }
780851
0 commit comments