From 409cebcb00bffafdbe529009f877c76f8fb804ee Mon Sep 17 00:00:00 2001
From: Nick Lockwood <nick@charcoaldesign.co.uk>
Date: Sat, 23 Nov 2024 21:43:29 +0000
Subject: [PATCH] Fix parsing `\.?` sequence

---
 Sources/Tokenizer.swift              | 16 ++++++++++--
 Tests/Rules/RedundantSelfTests.swift | 14 ++++++++++
 Tests/TokenizerTests.swift           | 39 ++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/Sources/Tokenizer.swift b/Sources/Tokenizer.swift
index 1dfb8891..560b7bda 100644
--- a/Sources/Tokenizer.swift
+++ b/Sources/Tokenizer.swift
@@ -1484,7 +1484,19 @@ public func tokenize(_ source: String) -> [Token] {
             assertionFailure()
             return
         }
-        while let nextToken: Token = index + 1 < tokens.count ? tokens[index + 1] : nil,
+        if index > 0, case .operator("\\", _) = tokens[index - 1] {
+            switch string {
+            case ".?.":
+                tokens[index ... index] = [.operator(".", .prefix), .operator("?", .postfix), .operator(".", .infix)]
+                return
+            case ".?":
+                tokens[index ... index] = [.operator(".", .prefix), .operator("?", .postfix)]
+                return
+            default:
+                break
+            }
+        }
+        while let nextToken = index + 1 < tokens.count ? tokens[index + 1] : nil,
               case let .operator(nextString, _) = nextToken, !nextString.hasPrefix("\\"),
               string.hasPrefix(".") || !nextString.contains(".")
         {
@@ -1497,7 +1509,7 @@ public func tokenize(_ source: String) -> [Token] {
             scopeIndexStack = scopeIndexStack.map { $0 > index ? $0 - 1 : $0 }
         }
         var index = index
-        while let prevToken: Token = index > 0 ? tokens[index - 1] : nil,
+        while let prevToken = index > 0 ? tokens[index - 1] : nil,
               case let .operator(prevString, _) = prevToken, !isUnwrapOperator(at: index - 1),
               !string.hasPrefix("\\"), prevString.hasPrefix(".") || !string.contains(".")
         {
diff --git a/Tests/Rules/RedundantSelfTests.swift b/Tests/Rules/RedundantSelfTests.swift
index 48daedb2..96ae8d7d 100644
--- a/Tests/Rules/RedundantSelfTests.swift
+++ b/Tests/Rules/RedundantSelfTests.swift
@@ -2485,6 +2485,20 @@ class RedundantSelfTests: XCTestCase {
         testFormatting(for: input, rule: .redundantSelf, options: options)
     }
 
+    func testNoInsertSelfInKeyPath() {
+        let input = """
+        class UserScreenPresenter: ScreenPresenter {
+            func onAppear() {
+                self.sessionInteractor.stage.compactMap(\\.?.session).latestValues(on: .main)
+            }
+
+            private var session: Session?
+        }
+        """
+        let options = FormatOptions(explicitSelf: .insert)
+        testFormatting(for: input, rule: .redundantSelf, options: options)
+    }
+
     // explicitSelf = .initOnly
 
     func testPreserveSelfInsideClassInit() {
diff --git a/Tests/TokenizerTests.swift b/Tests/TokenizerTests.swift
index df18a6bc..50784fa8 100644
--- a/Tests/TokenizerTests.swift
+++ b/Tests/TokenizerTests.swift
@@ -4524,6 +4524,45 @@ class TokenizerTests: XCTestCase {
         XCTAssertEqual(tokenize(input), output)
     }
 
+    func testAnonymousOptionalKeyPath() {
+        let input = "let foo = \\.?.bar"
+        let output: [Token] = [
+            .keyword("let"),
+            .space(" "),
+            .identifier("foo"),
+            .space(" "),
+            .operator("=", .infix),
+            .space(" "),
+            .operator("\\", .prefix),
+            .operator(".", .prefix),
+            .operator("?", .postfix),
+            .operator(".", .infix),
+            .identifier("bar"),
+        ]
+        XCTAssertEqual(tokenize(input), output)
+    }
+
+    func testAnonymousOptionalSubscriptKeyPath() {
+        let input = "let foo = \\.?[0].bar"
+        let output: [Token] = [
+            .keyword("let"),
+            .space(" "),
+            .identifier("foo"),
+            .space(" "),
+            .operator("=", .infix),
+            .space(" "),
+            .operator("\\", .prefix),
+            .operator(".", .prefix),
+            .operator("?", .postfix),
+            .startOfScope("["),
+            .number("0", .integer),
+            .endOfScope("]"),
+            .operator(".", .infix),
+            .identifier("bar"),
+        ]
+        XCTAssertEqual(tokenize(input), output)
+    }
+
     func testAttributeInsideGenericArguments() {
         let input = "Foo<(@MainActor () -> Void)?>(nil)"
         let output: [Token] = [