From 370f685cb1c6259858eafd9424ba17e26ef6d3cc Mon Sep 17 00:00:00 2001
From: John Sundell <john@sundell.co>
Date: Wed, 17 Jul 2019 10:57:44 +0200
Subject: [PATCH] Support types and functions named with a leading underscore
 (#74)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This patch fixes syntax highlighting for types and functions which names
begin with an underscore. Previously these entities would be treated as
plain text, but now they’re highlighted correctly as either types or
function calls.
---
 Sources/Splash/Grammar/SwiftGrammar.swift    | 14 ++++++++++---
 Tests/SplashTests/Tests/StatementTests.swift | 22 +++++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/Sources/Splash/Grammar/SwiftGrammar.swift b/Sources/Splash/Grammar/SwiftGrammar.swift
index 1441837..3217478 100644
--- a/Sources/Splash/Grammar/SwiftGrammar.swift
+++ b/Sources/Splash/Grammar/SwiftGrammar.swift
@@ -190,7 +190,11 @@ private extension SwiftGrammar {
         }
 
         func matches(_ segment: Segment) -> Bool {
-            guard segment.tokens.current.startsWithLetter else {
+            let token = segment.tokens.current.trimmingCharacters(
+                in: CharacterSet(charactersIn: "_")
+            )
+
+            guard token.startsWithLetter else {
                 return false
             }
 
@@ -305,7 +309,11 @@ private extension SwiftGrammar {
                 }
             }
 
-            guard segment.tokens.current.isCapitalized else {
+            let token = segment.tokens.current.trimmingCharacters(
+                in: CharacterSet(charactersIn: "_")
+            )
+
+            guard token.isCapitalized else {
                 return false
             }
 
@@ -316,7 +324,7 @@ private extension SwiftGrammar {
             // The XCTAssert family of functions is a bit of an edge case,
             // since they start with capital letters. Since they are so
             // commonly used, we'll add a special case for them here:
-            guard !segment.tokens.current.starts(with: "XCTAssert") else {
+            guard !token.starts(with: "XCTAssert") else {
                 return false
             }
 
diff --git a/Tests/SplashTests/Tests/StatementTests.swift b/Tests/SplashTests/Tests/StatementTests.swift
index c06ddc1..70af5bd 100644
--- a/Tests/SplashTests/Tests/StatementTests.swift
+++ b/Tests/SplashTests/Tests/StatementTests.swift
@@ -330,6 +330,24 @@ final class StatementTests: SyntaxHighlighterTestCase {
         ])
     }
 
+    func testInitializingTypeWithLeadingUnderscore() {
+        let components = highlighter.highlight("_MyType()")
+
+        XCTAssertEqual(components, [
+            .token("_MyType", .type),
+            .plainText("()")
+        ])
+    }
+
+    func testCallingFunctionWithLeadingUnderscore() {
+        let components = highlighter.highlight("_myFunction()")
+
+        XCTAssertEqual(components, [
+            .token("_myFunction", .call),
+            .plainText("()")
+        ])
+    }
+
     func testAllTestsRunOnLinux() {
         XCTAssertTrue(TestCaseVerifier.verifyLinuxTests((type(of: self)).allTests))
     }
@@ -348,7 +366,9 @@ extension StatementTests {
             ("testSwitchStatementWithOptional", testSwitchStatementWithOptional),
             ("testForStatementWithStaticProperty", testForStatementWithStaticProperty),
             ("testForStatementWithContinue", testForStatementWithContinue),
-            ("testRepeatWhileStatement", testRepeatWhileStatement)
+            ("testRepeatWhileStatement", testRepeatWhileStatement),
+            ("testInitializingTypeWithLeadingUnderscore", testInitializingTypeWithLeadingUnderscore),
+            ("testCallingFunctionWithLeadingUnderscore", testCallingFunctionWithLeadingUnderscore)
         ]
     }
 }