Correctly highlight generic return types (#69)

This change makes Splash correctly highlight generics that are returned from a function.
JohnSundell · May 13, 2019 · ec13df1 · ec13df1
1 parent 4675ffe
commit ec13df1
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/Sources/Splash/Grammar/SwiftGrammar.swift b/Sources/Splash/Grammar/SwiftGrammar.swift
@@ -326,6 +326,11 @@ private extension SwiftGrammar {
                 // Since the declaration might be on another line, we have to walk
                 // backwards through all tokens until we've found enough information.
                 for token in segment.tokens.all.reversed() {
+                    // Highlight return type generics as normal
+                    if token == "->" {
+                        return true
+                    }
+
                     if !foundOpeningBracket && token == "<" {
                         foundOpeningBracket = true
                     }

diff --git a/Tests/SplashTests/Tests/DeclarationTests.swift b/Tests/SplashTests/Tests/DeclarationTests.swift
@@ -197,6 +197,33 @@ final class DeclarationTests: SyntaxHighlighterTestCase {
         ])
     }
 
+    func testFunctionDeclarationWithGenericReturnType() {
+        let components = highlighter.highlight("""
+        func array() -> Array<Element> { return [] }
+        """)
+
+        XCTAssertEqual(components, [
+            .token("func", .keyword),
+            .whitespace(" "),
+            .plainText("array()"),
+            .whitespace(" "),
+            .plainText("->"),
+            .whitespace(" "),
+            .token("Array", .type),
+            .plainText("<"),
+            .token("Element", .type),
+            .plainText(">"),
+            .whitespace(" "),
+            .plainText("{"),
+            .whitespace(" "),
+            .token("return", .keyword),
+            .whitespace(" "),
+            .plainText("[]"),
+            .whitespace(" "),
+            .plainText("}")
+        ])
+    }
+
     func testGenericStructDeclaration() {
         let components = highlighter.highlight("struct MyStruct<A: Hello, B> {}")
 
@@ -882,6 +909,7 @@ extension DeclarationTests {
             ("testGenericFunctionDeclarationWithSingleConstraint", testGenericFunctionDeclarationWithSingleConstraint),
             ("testGenericFunctionDeclarationWithMultipleConstraints", testGenericFunctionDeclarationWithMultipleConstraints),
             ("testGenericFunctionDeclarationWithGenericParameter", testGenericFunctionDeclarationWithGenericParameter),
+            ("testFunctionDeclarationWithGenericReturnType", testFunctionDeclarationWithGenericReturnType),
             ("testGenericStructDeclaration", testGenericStructDeclaration),
             ("testClassDeclaration", testClassDeclaration),
             ("testCompactClassDeclarationWithInitializer", testCompactClassDeclarationWithInitializer),