From d01f7a6b2f3f9549f26251312c1307a13f1eff22 Mon Sep 17 00:00:00 2001
From: John Sundell <john.sundell@me.com>
Date: Sun, 14 Apr 2019 22:17:47 +0200
Subject: [PATCH] Correctly highlight closure argument shorthands (#66)

This patch makes Splash correctly highlight closure argument
shorthands (`$0`, `$1`, etc.), while also making the logic for
detecting a function call a bit simpler.
---
 Sources/Splash/Grammar/SwiftGrammar.swift  |  8 ++++--
 Tests/SplashTests/Tests/ClosureTests.swift | 32 +++++++++++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/Sources/Splash/Grammar/SwiftGrammar.swift b/Sources/Splash/Grammar/SwiftGrammar.swift
index 16b2690..5c63b32 100644
--- a/Sources/Splash/Grammar/SwiftGrammar.swift
+++ b/Sources/Splash/Grammar/SwiftGrammar.swift
@@ -144,8 +144,10 @@ private extension SwiftGrammar {
 
         func matches(_ segment: Segment) -> Bool {
             // Don't match against index-based closure arguments
-            guard segment.tokens.previous != "$" else {
-                return false
+            if let previous = segment.tokens.previous {
+                guard !previous.hasSuffix("$") else {
+                    return false
+                }
             }
 
             // Integers can be separated using "_", so handle that
@@ -231,7 +233,7 @@ private extension SwiftGrammar {
                 }
             }
 
-            return segment.tokens.next.isAny(of: "(", "()", "())", "(.", "({", "().", #"(\."#)
+            return segment.tokens.next?.starts(with: "(") ?? false
         }
     }
 
diff --git a/Tests/SplashTests/Tests/ClosureTests.swift b/Tests/SplashTests/Tests/ClosureTests.swift
index 4aceb2f..45e67fa 100644
--- a/Tests/SplashTests/Tests/ClosureTests.swift
+++ b/Tests/SplashTests/Tests/ClosureTests.swift
@@ -151,6 +151,35 @@ final class ClosureTests: SyntaxHighlighterTestCase {
         ])
     }
 
+    func testClosureArgumentShorthands() {
+        let components = highlighter.highlight("""
+        call {
+            print($0)
+            _ = $1
+            $2()
+        }
+        """)
+
+        XCTAssertEqual(components, [
+            .token("call", .call),
+            .whitespace(" "),
+            .plainText("{"),
+            .whitespace("\n    "),
+            .token("print", .call),
+            .plainText("($0)"),
+            .whitespace("\n    "),
+            .token("_", .keyword),
+            .whitespace(" "),
+            .plainText("="),
+            .whitespace(" "),
+            .plainText("$1"),
+            .whitespace("\n    "),
+            .plainText("$2()"),
+            .whitespace("\n"),
+            .plainText("}")
+        ])
+    }
+
     func testAllTestsRunOnLinux() {
         XCTAssertTrue(TestCaseVerifier.verifyLinuxTests((type(of: self)).allTests))
     }
@@ -166,7 +195,8 @@ extension ClosureTests {
             ("testClosureArgumentWithMultipleArguments", testClosureArgumentWithMultipleArguments),
             ("testEscapingClosureArgument", testEscapingClosureArgument),
             ("testPassingClosureAsArgument", testPassingClosureAsArgument),
-            ("testNestedEscapingClosure", testNestedEscapingClosure)
+            ("testNestedEscapingClosure", testNestedEscapingClosure),
+            ("testClosureArgumentShorthands", testClosureArgumentShorthands)
         ]
     }
 }