From 6b39e90ea4f4675d248668e8fce4c1c6f1f4a315 Mon Sep 17 00:00:00 2001
From: Shadowfacts <me@shadowfacts.net>
Date: Wed, 8 Apr 2020 23:16:46 -0400
Subject: [PATCH] JavaScript Highlighter: significantly improve performance

---
 .../JavaScriptHighlighter.swift               | 280 ++++++++++++------
 MongoView/Views/JavaScriptEditorView.swift    |  83 ++++--
 2 files changed, 252 insertions(+), 111 deletions(-)

diff --git a/MongoView/Synax Highlighting/JavaScriptHighlighter.swift b/MongoView/Synax Highlighting/JavaScriptHighlighter.swift
index 758163c..bc6fe75 100644
--- a/MongoView/Synax Highlighting/JavaScriptHighlighter.swift	
+++ b/MongoView/Synax Highlighting/JavaScriptHighlighter.swift	
@@ -7,6 +7,9 @@
 //
 
 import AppKit
+import OSLog
+
+fileprivate let log = OSLog(subsystem: "space.vaccor.MongoView.JavaScriptHighlighter", category: .pointsOfInterest)
 
 fileprivate let identifiers: CharacterSet = {
     var set = CharacterSet.alphanumerics
@@ -22,29 +25,32 @@ fileprivate let operators = CharacterSet(charactersIn: "+-*/<>=")
 fileprivate let expressionEnds = CharacterSet(charactersIn: ",]});")
 
 class JavaScriptHighlighter {
-    private let text: String
-    private var attributed: NSMutableAttributedString
+    private var text: String!
+    private var attributed: NSMutableAttributedString!
     private var currentIndex: String.Index!
-    private var indent = ""
+    private var _indent = ""
     private(set) var tokens = [(token: TokenType, range: NSRange)]()
     var debug = false
-    
-    init(text: String) {
-        self.text = text
-        self.attributed = NSMutableAttributedString(attributedString: NSAttributedString(string: text))
-    }
-    
-    init(mutableAttributed: NSMutableAttributedString) {
-        self.text = mutableAttributed.string
-        self.attributed = mutableAttributed
-    }
-    
+
     private func print(_ str: @autoclosure () -> String) {
+        #if DEBUG
         if debug {
-            Swift.print("\(indent)\(str())")
+            Swift.print("\(_indent)\(str())")
         }
+        #endif
     }
     
+    private func indent() {
+        #if DEBUG
+        _indent += "  "
+        #endif
+    }
+    private func outdent() {
+        #if DEBUG
+        _indent = String(_indent.dropLast(2))
+        #endif
+    }
+
     private func range(from: String.Index, to: String.Index) -> NSRange {
         return NSRange(from..<to, in: text)
     }
@@ -65,14 +71,28 @@ class JavaScriptHighlighter {
         return text[currentIndex..<text.index(currentIndex, offsetBy: realLength)]
     }
     
-    @discardableResult
-    private func consume() -> Unicode.Scalar? {
-        let c = peek()
+    private func consume() {
         currentIndex = text.index(after: currentIndex)
+    }
+    
+    private func getAndConsume() -> Unicode.Scalar? {
+        let c = peek()
+        consume()
         return c
     }
     
-    func highlight() {
+    func highlight(text: String) {
+        self.highlight(attributed: NSMutableAttributedString(attributedString: NSAttributedString(string: text)))
+    }
+    
+    func highlight(attributed: NSMutableAttributedString) {
+        self.attributed = attributed
+        self.text = attributed.string
+        self.tokens = []
+        
+        os_signpost(.begin, log: log, name: "highlight")
+        defer { os_signpost(.end, log: log, name: "highlight") }
+        
         attributed.beginEditing()
         let fullRange = NSRange(location: 0, length: attributed.length)
         attributed.setAttributes([
@@ -88,19 +108,95 @@ class JavaScriptHighlighter {
         attributed.endEditing()
     }
     
-    private func emit(token: TokenType, range: NSRange) {
-        let color: NSColor
-        switch token {
-        case .string:
-            color = .systemRed
-        case .number:
-            color = .systemBlue
-        case .punctuation:
-            color = .systemTeal
-        case .identifier:
-            return
+    func inserted(character: Unicode.Scalar, index: Int) -> Bool {
+        os_signpost(.begin, log: log, name: "inserted(character:index:)")
+        defer { os_signpost(.end, log: log, name: "inserted(character:index:)") }
+
+        if let (token, range) = tokenAndRange(fullyContaining: index) {
+            switch token {
+            case .identifier:
+                let set = range.location == index ? identifierStarts : identifiers
+                guard set.contains(character) else { return false }
+            case let .string(stringChar):
+                guard character != stringChar || (index > 0 && Unicode.Scalar((text as NSString).character(at: index - 1)) == "\\"),
+                    character != "\\" else { return false }
+            case .number:
+                if character == "." && (attributed.string as NSString).substring(with: range).contains(Character(".")) {
+                    return false
+                }
+            default:
+                return false
+            }
+            var attributes: [NSAttributedString.Key: Any] = [
+                .font: NSFont.monospacedSystemFont(ofSize: 13, weight: .regular)
+            ]
+            if let color = token.color {
+                attributes[.foregroundColor] = color
+            }
+            attributed.addAttributes(attributes, range: NSRange(location: index, length: 1))
+            tokens = tokens.map { (token, range) in
+                if range.contains(index) {
+                    return (token, NSRange(location: range.location, length: range.length + 1))
+                } else if range.location >= index {
+                    return (token, NSRange(location: range.location + 1, length: range.length))
+                } else {
+                    return (token, range)
+                }
+            }
+            return true
+        }
+        return false
+    }
+    
+    func removed(index: Int) -> Bool {
+        os_signpost(.begin, log: log, name: "removed(at:)")
+        defer { os_signpost(.end, log: log, name: "removed(at:)") }
+
+        if let (token, range) = tokenAndRange(fullyContaining: index) {
+            switch token {
+            case .string(_):
+                guard index > range.location && index < range.location + range.length - 1 else { return false }
+                if index > 0 && Unicode.Scalar((text as NSString).character(at: index - 1)) == "\\" {
+                    // removed character
+                    return false
+                }
+            case .number:
+                break
+            default:
+                return false
+            }
+            tokens = tokens.map { (token, range) in
+                if range.contains(index) {
+                    return (token, NSRange(location: range.location, length: range.length - 1))
+                } else if range.location >= index {
+                    return (token, NSRange(location: range.location - 1, length: range.length))
+                } else {
+                    return (token, range)
+                }
+            }
+            return true
+        }
+        return false
+    }
+    
+    func token(at index: Int) -> TokenType? {
+        for (token, range) in tokens where range.contains(index) {
+            return token
+        }
+        return nil
+    }
+    
+    private func tokenAndRange(fullyContaining index: Int) -> (TokenType, NSRange)? {
+        for (token, range) in tokens where index > range.location && index < range.location + range.length {
+            return (token, range)
+        }
+        return nil
+    }
+    
+    private func emit(token: TokenType, range: NSRange) {
+        if let color = token.color {
+            attributed.addAttribute(.foregroundColor, value: color, range: range)
         }
-        attributed.addAttribute(.foregroundColor, value: color, range: range)
         tokens.append((token, range))
     }
 
@@ -129,6 +225,8 @@ class JavaScriptHighlighter {
             consumeDotLookup()
         } else if char == "?" {
             consumeTernaryExpression()
+        } else if expressionEnds.contains(char) {
+            return
         } else {
             consume()
         }
@@ -168,15 +266,20 @@ class JavaScriptHighlighter {
     
     private func consumeString() {
         let stringStart = currentIndex!
-        let startChar = consume()
-        while currentIndex < text.endIndex && peek() != startChar && (currentIndex == text.startIndex || text[text.index(before: currentIndex)] != "\\") {
+        let quote = peek()!
+        consume() // opening quote
+        var prevChar: Unicode.Scalar?
+        var char = peek()
+        while currentIndex < text.endIndex {
             consume()
-        }
-        if currentIndex < text.endIndex {
-            consume() // string closing quote
+            if char == quote && prevChar != "\\" {
+                break
+            }
+            prevChar = char
+            char = peek()
         }
         print("String: \(text[stringStart..<currentIndex])")
-        emit(token: .string, range: range(from: stringStart, to: currentIndex))
+        emit(token: .string(quote), range: range(from: stringStart, to: currentIndex))
     }
     
     private func consumeTemplateString() {
@@ -185,14 +288,14 @@ class JavaScriptHighlighter {
         func emitTemplateStringFragment() {
             guard stringFragmentStart != currentIndex else { return }
             print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
-            emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
+            emit(token: .string("`"), range: range(from: stringFragmentStart!, to: currentIndex))
             stringFragmentStart = currentIndex
         }
         
         consume() // opening `
         
-        while currentIndex < text.endIndex {
-            if peek() == "$" {
+        while let char = peek() {
+            if char == "$" {
                 emitTemplateStringFragment()
                 consume() // $
                 if peek() == "{" {
@@ -205,7 +308,7 @@ class JavaScriptHighlighter {
                     }
 
                 }
-            } else if peek() == "`" {
+            } else if char == "`" {
                 stringFragmentStart = stringFragmentStart ?? currentIndex
                 consume() // `
                 emitTemplateStringFragment()
@@ -221,11 +324,11 @@ class JavaScriptHighlighter {
     }
     
     private func consumeTemplateStringExpression() {
-        indent += "  "
+        indent()
         while currentIndex < text.endIndex && peek() != "}" {
             consumeExpression()
         }
-        indent = String(indent.dropLast(2))
+        outdent()
     }
 
     private func consumeOperator() {
@@ -237,11 +340,11 @@ class JavaScriptHighlighter {
         consume() // (
         print("Opening (")
         emit(token: .punctuation, range: prevCharRange())
-        indent += "  "
+        indent()
         while currentIndex < text.endIndex && peek() != ")" {
             consumeExpression()
         }
-        indent = String(indent.dropLast(2))
+        outdent()
         if currentIndex < text.endIndex {
             consume() // )
             print("Closing )")
@@ -253,34 +356,29 @@ class JavaScriptHighlighter {
         consume() // {
         print("Opening {")
         emit(token: .punctuation, range: prevCharRange())
-        indent += "  "
-        object:
-        while currentIndex < text.endIndex && peek() != "}" {
-            consumeObjectKey()
-            if peek() == ":" {
+        indent()
+        while let char = peek() {
+            if char == "}" {
+                consume() // }
+                print("Closing }")
+                emit(token: .punctuation, range: prevCharRange())
+                break
+            } else if char == "'" || char == "\"" {
+                let keyStart = currentIndex!
+                consumeString()
+                print("Object key: '\(text[keyStart..<currentIndex])'")
+            } else if identifierStarts.contains(char) {
+                let keyStart = currentIndex!
+                consumeIdentifier()
+                print("Object key: '\(text[keyStart..<currentIndex])'")
+            } else if char == ":" || char == "," {
                 consume() // :
                 emit(token: .punctuation, range: prevCharRange())
-                consumeWhitespace()
-                while currentIndex < text.endIndex && peek() != "," && peek() != "}" {
-                    consumeExpression()
-                }
-            }
-            consumeWhitespace()
-
-            if peek() == "," {
-                consume() // ,
-                emit(token: .punctuation, range: prevCharRange())
-                continue
             } else {
-                break
+                consumeExpression()
             }
         }
-        indent = String(indent.dropLast(2))
-        if currentIndex < text.endIndex {
-            consume() // }
-            print("Closing }")
-            emit(token: .punctuation, range: prevCharRange())
-        }
+        outdent()
     }
     
     private func consumeObjectKey() {
@@ -317,35 +415,28 @@ class JavaScriptHighlighter {
         consume() // [
         print("Opening [")
         emit(token: .punctuation, range: prevCharRange())
-        indent += "  "
-        array:
-        while currentIndex < text.endIndex {
-            consumeWhitespace()
-            if peek() == "," {
+        indent()
+        while let char = peek() {
+            if char == "]" {
+                consume() // ]
+                print("Closing ]")
+                emit(token: .punctuation, range: prevCharRange())
+                break
+            } else if char == "," {
                 consume() // ,
                 print("Array separator")
                 emit(token: .punctuation, range: prevCharRange())
-            } else if peek() == "]" {
-                break array
             } else {
-                print("Array element")
-                while let char = peek(), char != ",", char != "]" {
-                    consumeExpression()
-                }
+                consumeExpression()
             }
         }
-        indent = String(indent.dropLast(2))
-        if currentIndex < text.endIndex {
-            consume() // ]
-            print("Closing ]")
-            emit(token: .punctuation, range: prevCharRange())
-        }
+        outdent()
     }
     
     func consumeTernaryExpression() {
         consume() // ?
         print("Ternary expression")
-        indent += "  "
+        indent()
         print("Ternary true result")
         while let char = peek(), char != ":" {
             consumeExpression() // true result
@@ -355,7 +446,7 @@ class JavaScriptHighlighter {
         while let char = peek(), !expressionEnds.contains(char) {
             consumeExpression()
         }
-        indent = String(indent.dropLast(2))
+        outdent()
     }
     
 }
@@ -365,6 +456,19 @@ extension JavaScriptHighlighter {
         case identifier
         case punctuation
         case number
-        case string
+        case string(Unicode.Scalar)
+        
+        var color: NSColor? {
+            switch self {
+            case .string(_):
+                return .systemRed
+            case .number:
+                return .systemBlue
+            case .punctuation:
+                return .systemTeal
+            case .identifier:
+                return nil
+            }
+        }
     }
 }
diff --git a/MongoView/Views/JavaScriptEditorView.swift b/MongoView/Views/JavaScriptEditorView.swift
index 3f48823..b662cbe 100644
--- a/MongoView/Views/JavaScriptEditorView.swift
+++ b/MongoView/Views/JavaScriptEditorView.swift
@@ -10,24 +10,37 @@ import AppKit
 
 class JavaScriptEditorView: NSTextView {
     
-    var highlighter: JavaScriptHighlighter?
+    var highlighter = JavaScriptHighlighter()
+    private var isRehighlighting = false
     
     override var string: String {
-        didSet {
+        get {
+            super.string
+        }
+        set {
+            isRehighlighting = true
+            super.string = newValue
             rehighlight()
         }
     }
 
     func rehighlight() {
-        highlighter = JavaScriptHighlighter(mutableAttributed: self.textStorage!)
-        highlighter!.highlight()
+        isRehighlighting = true
+        highlighter.highlight(attributed: self.textStorage!)
+        isRehighlighting = false
+    }
+    
+    override func awakeFromNib() {
+        super.awakeFromNib()
+        
+        textStorage!.delegate = self
     }
     
     override func shouldChangeText(in affectedCharRange: NSRange, replacementString: String?) -> Bool {
         guard super.shouldChangeText(in: affectedCharRange, replacementString: replacementString) else {
             return false
         }
-        
+
         if affectedCharRange.length == 0,
             let string = replacementString,
             string.count == 1,
@@ -37,17 +50,12 @@ class JavaScriptEditorView: NSTextView {
 
         return true
     }
-
-    override func didChangeText() {
-        rehighlight()
-        
-        super.didChangeText()
-    }
     
     func tryAutocompleteCharacter(for range: NSRange, string inserted: String) -> Bool {
         if let end = autocompleteResultFor(string: inserted, in: range) {
             textStorage!.insert(NSAttributedString(string: "\(inserted)\(end)"), at: range.location)
-            didChangeText()
+            rehighlight()
+            super.didChangeText()
             setSelectedRange(NSRange(location: range.location + 1, length: 0))
             return true
         } else {
@@ -58,28 +66,57 @@ class JavaScriptEditorView: NSTextView {
     private func autocompleteResultFor(string: String, in range: NSRange) -> String? {
         switch string {
         case "'", "\"", "`":
-            return token(at: range.location) != .string ? string : nil
+            if case .string(_) = highlighter.token(at: range.location) {
+                return nil
+            } else {
+                return string
+            }
         case "(":
-            return token(at: range.location) != .string ? ")" : nil
+            if case .string(_) = highlighter.token(at: range.location) {
+                return nil
+            } else {
+                return ")"
+            }
         case "[":
-            return token(at: range.location) != .string ? "]" : nil
+            if case .string(_) = highlighter.token(at: range.location) {
+                return nil
+            } else {
+                return "]"
+            }
         case "{":
-            var prevChar: Unicode.Scalar?
             if range.location > 0 {
                 let index = self.string.index(self.string.startIndex, offsetBy: range.location - 1)
-                prevChar = self.string.unicodeScalars[index]
+                let prevChar = self.string.unicodeScalars[index]
+                if prevChar == "$" {
+                    return "}"
+                }
+            }
+            if case .string(_) = highlighter.token(at: range.location) {
+                return nil
+            } else {
+                return "}"
             }
-            return token(at: range.location) != .string || (prevChar == "$") ? "}" : nil
         default:
             return nil
         }
     }
 
-    private func token(at index: Int) -> JavaScriptHighlighter.TokenType? {
-        guard let highlighter = highlighter else { return nil }
-        for (token, range) in highlighter.tokens where range.contains(index) {
-            return token
+}
+
+extension JavaScriptEditorView: NSTextStorageDelegate {
+    func textStorage(_ textStorage: NSTextStorage, didProcessEditing editedMask: NSTextStorageEditActions, range editedRange: NSRange, changeInLength delta: Int) {
+        guard editedMask.contains(.editedCharacters), !isRehighlighting else { return }
+        if delta == 1 {
+            if let char = Unicode.Scalar((textStorage.string as NSString).character(at: editedRange.location)),
+                highlighter.inserted(character: char, index: editedRange.location) {
+                return
+            }
+        } else if delta == -1 {
+            if highlighter.removed(index: editedRange.location) {
+                return
+            }
         }
-        return nil
+        
+        rehighlight()
     }
 }