From bc20ac969c1b2a770bf315497e0a66d7367f7fec Mon Sep 17 00:00:00 2001
From: John Sundell <john@sundell.co>
Date: Mon, 4 Nov 2019 11:31:08 +0100
Subject: [PATCH] Support highlighting custom compiler directives (#88)

Add support for highlighting the `#warning` and `#error` directives.
---
 .../Extensions/Sequence/Sequence+AnyOf.swift  |  2 +-
 Sources/Splash/Grammar/SwiftGrammar.swift     | 11 +++++---
 .../SplashTests/Tests/PreprocessorTests.swift | 26 ++++++++++++++++++-
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/Sources/Splash/Extensions/Sequence/Sequence+AnyOf.swift b/Sources/Splash/Extensions/Sequence/Sequence+AnyOf.swift
index 0af8aae..232f5c6 100644
--- a/Sources/Splash/Extensions/Sequence/Sequence+AnyOf.swift
+++ b/Sources/Splash/Extensions/Sequence/Sequence+AnyOf.swift
@@ -11,7 +11,7 @@ internal extension Sequence where Element: Equatable {
         return contains(anyOf: candidates)
     }
 
-    func contains(anyOf candidates: [Element]) -> Bool {
+    func contains<S: Sequence>(anyOf candidates: S) -> Bool where S.Element == Element {
         for candidate in candidates {
             if contains(candidate) {
                 return true
diff --git a/Sources/Splash/Grammar/SwiftGrammar.swift b/Sources/Splash/Grammar/SwiftGrammar.swift
index 9596bfb..1dbf34d 100644
--- a/Sources/Splash/Grammar/SwiftGrammar.swift
+++ b/Sources/Splash/Grammar/SwiftGrammar.swift
@@ -80,14 +80,19 @@ private extension SwiftGrammar {
 
     struct PreprocessingRule: SyntaxRule {
         var tokenType: TokenType { return .preprocessing }
-        private let tokens = ["#if", "#endif", "#elseif", "#else"]
+        private let controlFlowTokens: Set<String> = ["#if", "#endif", "#elseif", "#else"]
+        private let directiveTokens: Set<String> = ["#warning", "#error"]
 
         func matches(_ segment: Segment) -> Bool {
-            if segment.tokens.current.isAny(of: tokens) {
+            if segment.tokens.current.isAny(of: controlFlowTokens) {
                 return true
             }
 
-            return segment.tokens.onSameLine.contains(anyOf: tokens)
+            if segment.tokens.current.isAny(of: directiveTokens) {
+                return true
+            }
+
+            return segment.tokens.onSameLine.contains(anyOf: controlFlowTokens)
         }
     }
 
diff --git a/Tests/SplashTests/Tests/PreprocessorTests.swift b/Tests/SplashTests/Tests/PreprocessorTests.swift
index 254e6d2..746f46d 100644
--- a/Tests/SplashTests/Tests/PreprocessorTests.swift
+++ b/Tests/SplashTests/Tests/PreprocessorTests.swift
@@ -80,6 +80,28 @@ final class PreprocessorTests: SyntaxHighlighterTestCase {
         ])
     }
 
+    func testWarningDirective() {
+        let components = highlighter.highlight(#"#warning("Hey!")"#)
+
+        XCTAssertEqual(components, [
+            .token("#warning", .preprocessing),
+            .plainText("("),
+            .token(#""Hey!""#, .string),
+            .plainText(")")
+        ])
+    }
+
+    func testErrorDirective() {
+        let components = highlighter.highlight(#"#error("No!")"#)
+
+        XCTAssertEqual(components, [
+            .token("#error", .preprocessing),
+            .plainText("("),
+            .token(#""No!""#, .string),
+            .plainText(")")
+        ])
+    }
+
     func testAllTestsRunOnLinux() {
         XCTAssertTrue(TestCaseVerifier.verifyLinuxTests((type(of: self)).allTests))
     }
@@ -91,7 +113,9 @@ extension PreprocessorTests {
             ("testPreprocessing", testPreprocessing),
             ("testSelector", testSelector),
             ("testFunctionAttribute", testFunctionAttribute),
-            ("testAvailabilityCheck", testAvailabilityCheck)
+            ("testAvailabilityCheck", testAvailabilityCheck),
+            ("testWarningDirective", testWarningDirective),
+            ("testErrorDirective", testErrorDirective)
         ]
     }
 }