Merge pull request #2891 from MAJKFL/sanitize-dollar-identifier

Sanitize dollar identifiers with leading zeros in `Identifier`.

Author: DougGregor

Sources/SwiftLexicalLookup/LookupName.swift

@@ -86,7 +86,7 @@ import SwiftSyntax
   /// `self` and `Self` identifers override implicit `self` and `Self` introduced by
   /// the `Foo` class declaration.
   var identifier: Identifier {
-    Identifier(name)
+    Identifier(canonicalName: name)
   }
 
   /// Position of this implicit name.

Sources/SwiftLexicalLookup/Scopes/ScopeImplementations.swift

@@ -188,10 +188,10 @@ import SwiftSyntax
         checkIdentifier(identifier, refersTo: name, at: lookUpPosition)
       }
 
-      if let dollarIdentifierStr = identifier?.dollarIdentifierStr {
+      if let identifier, identifier.isDollarIdentifier {
         signatureResults = LookupResult.getResultArray(
           for: self,
-          withNames: filteredCaptureNames + [LookupName.dollarIdentifier(self, strRepresentation: dollarIdentifierStr)]
+          withNames: filteredCaptureNames + [LookupName.dollarIdentifier(self, strRepresentation: identifier.name)]
         )
       } else {
         signatureResults =

Sources/SwiftSyntax/Identifier.swift

@@ -17,7 +17,10 @@ public struct Identifier: Equatable, Hashable, Sendable {
     String(syntaxText: raw.name)
   }
 
-  public let dollarIdentifierStr: String?
+  /// `true` if the identifier is a dollar identifier.
+  public var isDollarIdentifier: Bool {
+    raw.original.hasPrefix(SyntaxText("$")) && Int(String(syntaxText: raw.original).dropFirst()) != nil
+  }
 
   @_spi(RawSyntax)
   public let raw: RawIdentifier
@@ -26,61 +29,86 @@ public struct Identifier: Equatable, Hashable, Sendable {
   public init?(_ token: TokenSyntax) {
     switch token.tokenKind {
     case .identifier, .keyword(.self), .keyword(.Self):
-      self.raw = RawIdentifier(token.tokenView)
+      self.raw = RawIdentifier(token.tokenView.rawText)
       self.arena = token.raw.arenaReference
-
-      self.dollarIdentifierStr = nil
     case .dollarIdentifier(let dollarIdentifierStr):
-      self.raw = RawIdentifier(token.tokenView)
       self.arena = token.raw.arenaReference
 
-      self.dollarIdentifierStr = dollarIdentifierStr
+      if Self.isPaddedDollarIdentifier(dollarIdentfierStr: dollarIdentifierStr),
+        let newDollarIdentifierNumber = Int(dollarIdentifierStr.dropFirst())
+      {
+        let newDollarIdentifierStr = "$\(newDollarIdentifierNumber)"
+        let sanitizedDollarIdentifierSyntaxText = token.raw.arenaReference.intern(newDollarIdentifierStr)
+
+        self.raw = RawIdentifier(sanitizedDollarIdentifierSyntaxText)
+      } else {
+        self.raw = RawIdentifier(token.tokenView.rawText)
+      }
     default:
       return nil
     }
   }
 
-  public init(_ staticString: StaticString) {
-    self.raw = RawIdentifier(staticString)
-    self.arena = nil
-
-    let name = String(syntaxText: raw.name)
+  /// Create a new `Identifier` from given `canonicalName`.
+  ///
+  /// - Precondition: `canonicalName` is a canonical identifier i.e. doesn't
+  ///   use backticks and is not a dollar identifier with leading zeros.
+  public init(canonicalName: StaticString) {
+    precondition(
+      Self.isCanonicalRepresentation(canonicalName),
+      "\(canonicalName) is not a canonical identifier."
+    )
 
-    if name.first == "$" && Int(name.dropFirst()) != nil {
-      self.dollarIdentifierStr = name
-    } else {
-      self.dollarIdentifierStr = nil
-    }
+    self.raw = RawIdentifier(SyntaxText(canonicalName))
+    self.arena = nil
   }
 
   public static func == (lhs: Self, rhs: Self) -> Bool {
     lhs.name == rhs.name
   }
 
-  private static func getDollarIdentifierNumber(str: String) -> Int? {
-    guard str.first == "$" else { return nil }
+  /// Returns `true` if `staticString` is a canonical identifier i.e. doesn't
+  /// use backticks and is not a dollar identifier with leading zeros.
+  private static func isCanonicalRepresentation(_ staticString: StaticString) -> Bool {
+    let text = SyntaxText(staticString)
+
+    guard !Self.hasBackticks(text) else { return false }
+
+    let str = String(syntaxText: text)
+    let isDollarIdentifier = str.first == "$" && Int(str.dropFirst()) != nil
 
-    return Int(str.dropFirst())
+    return !(isDollarIdentifier && Self.isPaddedDollarIdentifier(dollarIdentfierStr: str))
+  }
+
+  /// Returns `true` if `rawText` doesn't use backticks.
+  fileprivate static func hasBackticks(_ rawText: SyntaxText) -> Bool {
+    let backtick = SyntaxText("`")
+    return rawText.count > 2 && rawText.hasPrefix(backtick) && rawText.hasSuffix(backtick)
+  }
+
+  /// Returns `true` if `dollarIdentfierStr` is not a
+  /// dollar identifier with leading zeros.
+  fileprivate static func isPaddedDollarIdentifier(dollarIdentfierStr: String) -> Bool {
+    dollarIdentfierStr.count > 2 && dollarIdentfierStr.hasPrefix("$0")
   }
 }
 
 @_spi(RawSyntax)
 public struct RawIdentifier: Equatable, Hashable, Sendable {
+  fileprivate let original: SyntaxText
   public let name: SyntaxText
 
   @_spi(RawSyntax)
-  fileprivate init(_ raw: RawSyntaxTokenView) {
-    let backtick = SyntaxText("`")
-    if raw.rawText.count > 2 && raw.rawText.hasPrefix(backtick) && raw.rawText.hasSuffix(backtick) {
-      let startIndex = raw.rawText.index(after: raw.rawText.startIndex)
-      let endIndex = raw.rawText.index(before: raw.rawText.endIndex)
-      self.name = SyntaxText(rebasing: raw.rawText[startIndex..<endIndex])
-    } else {
-      self.name = raw.rawText
+  fileprivate init(_ rawText: SyntaxText) {
+    self.original = rawText
+
+    guard Identifier.hasBackticks(rawText) else {
+      self.name = rawText
+      return
     }
-  }
 
-  fileprivate init(_ staticString: StaticString) {
-    name = SyntaxText(staticString)
+    let startIndex = rawText.index(after: rawText.startIndex)
+    let endIndex = rawText.index(before: rawText.endIndex)
+    self.name = SyntaxText(rebasing: rawText[startIndex..<endIndex])
   }
 }

Sources/SwiftSyntax/SyntaxArena.swift

@@ -291,6 +291,12 @@ struct SyntaxArenaRef: Hashable, @unchecked Sendable {
     return RetainedSyntaxArena(value)
   }
 
+  /// Copies a UTF8 sequence of `String` to the memory the referenced arena manages, and
+  /// returns the copied string as a ``SyntaxText``
+  func intern(_ value: String) -> SyntaxText {
+    self.value.intern(value)
+  }
+
   #if DEBUG || SWIFTSYNTAX_ENABLE_ASSERTIONS
   /// Accessor for the underlying's `SyntaxArena.hasParent`
   var hasParent: Bool {

Tests/SwiftLexicalLookupTest/NameLookupTests.swift

@@ -142,7 +142,7 @@ final class testNameLookup: XCTestCase {
         func foo() {
           let 0️⃣a = 1
           let x = 5️⃣{
-            print(2️⃣a, 3️⃣$0, 4️⃣$123)
+            print(2️⃣a, 3️⃣$0, 4️⃣$123, 6️⃣$00000001)
           }
         }
         """,
@@ -157,6 +157,9 @@ final class testNameLookup: XCTestCase {
         "4️⃣": [
           .fromScope(ClosureExprSyntax.self, expectedNames: [NameExpectation.dollarIdentifier("5️⃣", "$123")])
         ],
+        "6️⃣": [
+          .fromScope(ClosureExprSyntax.self, expectedNames: [NameExpectation.dollarIdentifier("5️⃣", "$1")])
+        ],
       ]
     )
   }