diff --git a/Polychromatic/Categories/DVTTextStorage+PLYHighlightingHook.m b/Polychromatic/Categories/DVTTextStorage+PLYHighlightingHook.m
index 27b1e98..bfa0f1f 100644
--- a/Polychromatic/Categories/DVTTextStorage+PLYHighlightingHook.m
+++ b/Polychromatic/Categories/DVTTextStorage+PLYHighlightingHook.m
@@ -53,7 +53,9 @@ - (NSColor *)ply_colorAtCharacterIndex:(unsigned long long)index effectiveRange:
     {
         long long nodeType = [self nodeTypeAtCharacterIndex:newRange.location effectiveRange:effectiveRange context:context];
 
-        if (nodeType == 9)
+        if (nodeType == [DVTSourceNodeTypes registerNodeTypeNamed:@"xcode.syntax.identifier.variable"] ||
+            nodeType == [DVTSourceNodeTypes registerNodeTypeNamed:@"xcode.syntax.identifier.constant"] ||
+            nodeType == [DVTSourceNodeTypes registerNodeTypeNamed:@"xcode.syntax.identifier"])
         {
             PLYMockSwift *fauxSwiftService = (PLYMockSwift *)self.languageService;
             NSRange funcDefinitionRange = [fauxSwiftService methodDefinitionRangeAtIndex:newRange.location];