Use filter documents intead of shelling out to mongo every time

MongoView.xcodeproj/project.pbxproj

@@ -11,6 +11,8 @@
 		D60C863A23CA2DD100C9DB8E /* ServerConnectWindowController.xib in Resources */ = {isa = PBXBuildFile; fileRef = D60C863823CA2DD100C9DB8E /* ServerConnectWindowController.xib */; };
 		D60C863F23CA2E2100C9DB8E /* ServerConnectViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = D60C863D23CA2E2100C9DB8E /* ServerConnectViewController.swift */; };
 		D60C864023CA2E2100C9DB8E /* ServerConnectViewController.xib in Resources */ = {isa = PBXBuildFile; fileRef = D60C863E23CA2E2100C9DB8E /* ServerConnectViewController.xib */; };
+		D62408C12438CF550020E09F /* JavaScriptEditorView.swift in Sources */ = {isa = PBXBuildFile; fileRef = D62408C02438CF550020E09F /* JavaScriptEditorView.swift */; };
+		D624090F243903E90020E09F /* ExtendedJSON.swift in Sources */ = {isa = PBXBuildFile; fileRef = D624090E243903E90020E09F /* ExtendedJSON.swift */; };
 		D63CDEBE23C837DC0012D658 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = D63CDEBD23C837DC0012D658 /* AppDelegate.swift */; };
 		D63CDEC023C837DD0012D658 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = D63CDEBF23C837DD0012D658 /* Assets.xcassets */; };
 		D63CDEC323C837DD0012D658 /* MainMenu.xib in Resources */ = {isa = PBXBuildFile; fileRef = D63CDEC123C837DD0012D658 /* MainMenu.xib */; };
@@ -84,6 +86,8 @@
 		D60C863823CA2DD100C9DB8E /* ServerConnectWindowController.xib */ = {isa = PBXFileReference; lastKnownFileType = file.xib; path = ServerConnectWindowController.xib; sourceTree = "<group>"; };
 		D60C863D23CA2E2100C9DB8E /* ServerConnectViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ServerConnectViewController.swift; sourceTree = "<group>"; };
 		D60C863E23CA2E2100C9DB8E /* ServerConnectViewController.xib */ = {isa = PBXFileReference; lastKnownFileType = file.xib; path = ServerConnectViewController.xib; sourceTree = "<group>"; };
+		D62408C02438CF550020E09F /* JavaScriptEditorView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = JavaScriptEditorView.swift; sourceTree = "<group>"; };
+		D624090E243903E90020E09F /* ExtendedJSON.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ExtendedJSON.swift; sourceTree = "<group>"; };
 		D63CDEBA23C837DC0012D658 /* MongoView.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = MongoView.app; sourceTree = BUILT_PRODUCTS_DIR; };
 		D63CDEBD23C837DC0012D658 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
 		D63CDEBF23C837DD0012D658 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
@@ -149,8 +153,6 @@
 				D60C863823CA2DD100C9DB8E /* ServerConnectWindowController.xib */,
 				D63CDF3A23C838470012D658 /* DatabaseWindowController.swift */,
 				D63CDF3B23C838470012D658 /* DatabaseWindowController.xib */,
-				D6A7D095243541A400B46857 /* WindowStatusView.swift */,
-				D6A7D099243546B500B46857 /* WindowStatusView.xib */,
 			);
 			path = Windows;
 			sourceTree = "<group>";
@@ -168,6 +170,16 @@
 			path = "View Controllers";
 			sourceTree = "<group>";
 		};
+		D62408BF2438CF3C0020E09F /* Views */ = {
+			isa = PBXGroup;
+			children = (
+				D6A7D095243541A400B46857 /* WindowStatusView.swift */,
+				D6A7D099243546B500B46857 /* WindowStatusView.xib */,
+				D62408C02438CF550020E09F /* JavaScriptEditorView.swift */,
+			);
+			path = Views;
+			sourceTree = "<group>";
+		};
 		D63CDEB123C837DC0012D658 = {
 			isa = PBXGroup;
 			children = (
@@ -194,9 +206,11 @@
 				D63CDF3423C838190012D658 /* MongoController.swift */,
 				D63CDF3323C838190012D658 /* Node.swift */,
 				D6D4665223CB730C00F13B1B /* MongoEvaluator.swift */,
+				D624090E243903E90020E09F /* ExtendedJSON.swift */,
 				D6A7D0A22435880700B46857 /* Synax Highlighting */,
 				D60C863B23CA2DD600C9DB8E /* Windows */,
 				D60C863C23CA2DDD00C9DB8E /* View Controllers */,
+				D62408BF2438CF3C0020E09F /* Views */,
 				D63CDEBF23C837DD0012D658 /* Assets.xcassets */,
 				D63CDEC123C837DD0012D658 /* MainMenu.xib */,
 				D63CDEC423C837DD0012D658 /* Info.plist */,
@@ -337,6 +351,7 @@
 			files = (
 				D63CDEBE23C837DC0012D658 /* AppDelegate.swift in Sources */,
 				D6D4665323CB730C00F13B1B /* MongoEvaluator.swift in Sources */,
+				D624090F243903E90020E09F /* ExtendedJSON.swift in Sources */,
 				D63CDF3823C8381A0012D658 /* MongoController.swift in Sources */,
 				D60C863923CA2DD100C9DB8E /* ServerConnectWindowController.swift in Sources */,
 				D63CDF3723C8381A0012D658 /* Node.swift in Sources */,
@@ -345,6 +360,7 @@
 				D63CDF3C23C838470012D658 /* DatabaseWindowController.swift in Sources */,
 				D6A7D096243541A400B46857 /* WindowStatusView.swift in Sources */,
 				D6A7D0A42435885B00B46857 /* JavaScriptHighlighter.swift in Sources */,
+				D62408C12438CF550020E09F /* JavaScriptEditorView.swift in Sources */,
 				D63CDF4023C839010012D658 /* QueryViewController.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;

MongoView/Base.lproj/MainMenu.xib

@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.Cocoa.XIB" version="3.0" toolsVersion="16085" targetRuntime="MacOSX.Cocoa" propertyAccessControl="none" useAutolayout="YES" customObjectInstantitationMethod="direct">
+<document type="com.apple.InterfaceBuilder3.Cocoa.XIB" version="3.0" toolsVersion="16096" targetRuntime="MacOSX.Cocoa" propertyAccessControl="none" useAutolayout="YES" customObjectInstantitationMethod="direct">
     <dependencies>
-        <plugIn identifier="com.apple.InterfaceBuilder.CocoaPlugin" version="16085"/>
+        <plugIn identifier="com.apple.InterfaceBuilder.CocoaPlugin" version="16096"/>
     </dependencies>
     <objects>
         <customObject id="-2" userLabel="File's Owner" customClass="NSApplication">
@@ -368,14 +368,6 @@
                                     <action selector="refresh:" target="-1" id="S5L-sX-Z2j"/>
                                 </connections>
                             </menuItem>
-                            <menuItem title="Run Query" id="RSP-nH-a1I">
-                                <string key="keyEquivalent" base64-UTF8="YES">
-DQ
-</string>
-                                <connections>
-                                    <action selector="runQuery:" target="-1" id="jLo-qu-WjN"/>
-                                </connections>
-                            </menuItem>
                         </items>
                     </menu>
                 </menuItem>

MongoView/ExtendedJSON.swift

@@ -0,0 +1,28 @@
+//
+//  ExtendedJSON.swift
+//  MongoView
+//
+//  Created by Shadowfacts on 4/4/20.
+//  Copyright © 2020 Shadowfacts. All rights reserved.
+//
+
+import Foundation
+import JavaScriptCore
+
+struct ExtendedJSON {
+    private init() {}
+    
+    private static let context: JSContext = {
+        let context = JSContext()!
+        let objectId: @convention(block) (String) -> [String: String] = { (id) in
+            return ["$oid": id]
+        }
+        context.setObject(objectId, forKeyedSubscript: "ObjectId" as NSString)
+        return context
+    }()
+    
+    static func normalize(_ string: String) -> String? {
+        return context.evaluateScript("JSON.stringify(\(string))")?.toString()
+    }
+    
+}

MongoView/Synax Highlighting/JavaScriptHighlighter.swift

@@ -19,15 +19,23 @@ fileprivate let identifierStarts: CharacterSet = {
     return set
 }()
 fileprivate let operators = CharacterSet(charactersIn: "+-*/<>=")
+fileprivate let expressionEnds = CharacterSet(charactersIn: ",]});")
 
 class JavaScriptHighlighter {
     private let text: String
-    private var attributed: NSMutableAttributedString!
+    private var attributed: NSMutableAttributedString
     private var currentIndex: String.Index!
     private var indent = ""
+    private(set) var tokens = [(token: TokenType, range: NSRange)]()
     
     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: String) {
@@ -61,9 +69,7 @@ class JavaScriptHighlighter {
         return c
     }
     
-    func highlight(mutableAttributed: NSMutableAttributedString? = nil) -> NSAttributedString {
+    func highlight() {
-        attributed = mutableAttributed ?? NSMutableAttributedString(attributedString: NSAttributedString(string: text))
-        
         let fullRange = NSRange(location: 0, length: attributed.length)
         attributed.setAttributes([
             .foregroundColor: NSColor.textColor,
@@ -71,11 +77,9 @@ class JavaScriptHighlighter {
         ], range: fullRange)
 
         currentIndex = text.startIndex
-        while currentIndex < text.endIndex {
+        while let char = peek(), !expressionEnds.contains(char) {
             consumeExpression()
         }
-
-        return attributed
     }
     
     private func emit(token: TokenType, range: NSRange) {
@@ -91,6 +95,7 @@ class JavaScriptHighlighter {
             return
         }
         attributed.addAttribute(.foregroundColor, value: color, range: range)
+        tokens.append((token, range))
     }
 
     private func consumeExpression() {
@@ -116,6 +121,8 @@ class JavaScriptHighlighter {
             consumeArray()
         } else if char == "." {
             consumeDotLookup()
+        } else if char == "?" {
+            consumeTernaryExpression()
         } else {
             consume()
         }
@@ -168,32 +175,42 @@ class JavaScriptHighlighter {
     
     private func consumeTemplateString() {
         var stringFragmentStart: String.Index? = currentIndex
+        
+        func emitTemplateStringFragment() {
+            guard stringFragmentStart != currentIndex else { return }
+            print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
+            emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
+            stringFragmentStart = currentIndex
+        }
+        
         consume() // opening `
+        
         while currentIndex < text.endIndex {
-            if peek(length: 2) == "${" {
+            if peek() == "$" {
+                emitTemplateStringFragment()
                 consume() // $
-                consume() // {
+                if peek() == "{" {
-                print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
+                    consume() // {
-                emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
+                    emitTemplateStringFragment()
-                consumeTemplateStringExpression()
+                    consumeTemplateStringExpression()
-                stringFragmentStart = currentIndex
+                    stringFragmentStart = currentIndex
-                if currentIndex < text.endIndex && peek() == "}" {
+                    if currentIndex < text.endIndex && peek() == "}" {
-                    consume()
+                        consume() // }
+                    }
+
                 }
             } else if peek() == "`" {
                 stringFragmentStart = stringFragmentStart ?? currentIndex
                 consume() // `
-                print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
+                emitTemplateStringFragment()
-                emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
                 stringFragmentStart = nil
                 break
             } else {
                 consume()
             }
         }
-        if let start = stringFragmentStart {
+        if stringFragmentStart != nil {
-            print("Template string fragment: '\(text[start..<currentIndex])'")
+            emitTemplateStringFragment()
-            emit(token: .string, range: range(from: start, to: currentIndex))
         }
     }
     
@@ -286,7 +303,7 @@ class JavaScriptHighlighter {
         indent += "  "
         array:
         while currentIndex < text.endIndex && peek() != "]" {
-            print("Array element")
+            consumeWhitespace()
             while currentIndex < text.endIndex {
                 if peek() == "," {
                     consume() // ,
@@ -297,6 +314,7 @@ class JavaScriptHighlighter {
                     break array
                 } else {
                     indent += "  "
+                    print("Array element")
                     consumeExpression()
                     indent = String(indent.dropLast(2))
                 }
@@ -310,6 +328,22 @@ class JavaScriptHighlighter {
         }
     }
     
+    func consumeTernaryExpression() {
+        consume() // ?
+        print("Ternary expression")
+        indent += "  "
+        print("Ternary true result")
+        while let char = peek(), char != ":" {
+            consumeExpression() // true result
+        }
+        consume() // :
+        print("Ternary false result")
+        while let char = peek(), !expressionEnds.contains(char) {
+            consumeExpression()
+        }
+        indent = String(indent.dropLast(2))
+    }
+    
 }
 
 extension JavaScriptHighlighter {

MongoView/View Controllers/DatabaseViewController.swift

@@ -140,7 +140,7 @@ class DatabaseViewController: NSViewController {
             }
         } else if let collection = item as? DatabaseCollection {
             // only open a new window/tab if our own query has changed from the default, otherwise replace our query controller
-            if let queryViewController = queryViewController, queryViewController.hasQueryChanged {
+            if let queryViewController = queryViewController, queryViewController.hasFilterChanged {
                 (NSApplication.shared.delegate as! AppDelegate).newWindow(mongoController: mongoController, collection: collection)
             } else {
                 self.selectedCollection = collection
@@ -149,10 +149,6 @@ class DatabaseViewController: NSViewController {
         }
     }
 
-    @IBAction func runQuery(_ sender: Any) {
-        queryViewController?.runQuery()
-    }
-    
     @IBAction func refresh(_ sender: Any) {
         queryViewController?.refresh()
     }
@@ -160,7 +156,7 @@ class DatabaseViewController: NSViewController {
 
 extension DatabaseViewController: NSMenuItemValidation {
     func validateMenuItem(_ menuItem: NSMenuItem) -> Bool {
-        if menuItem.action == #selector(runQuery(_:)) || menuItem.action == #selector(refresh(_:)) {
+        if menuItem.action == #selector(refresh(_:)) {
             return queryViewController != nil
         }
         return true

MongoView/View Controllers/QueryViewController.swift

@@ -12,19 +12,19 @@ import MongoSwift
 class QueryViewController: NSViewController {
 
     @IBOutlet weak var verticalSplitView: NSSplitView!
-    @IBOutlet var queryTextView: NSTextView!
+    @IBOutlet var filterTextView: JavaScriptEditorView!
     @IBOutlet weak var outlineView: NSOutlineView!
     @IBOutlet weak var documentCountLabel: NSTextField!
     
     let mongoController: MongoController
     let collection: DatabaseCollection
     
-    var defaultQuery: String {
+    var defaultFilter: String {
-        "db.getCollection('\(collection.name)').find({}).toArray()"
+        "{}"
     }
     
-    var hasQueryChanged: Bool {
+    var hasFilterChanged: Bool {
-        return queryTextView.string != defaultQuery
+        return filterTextView.string != defaultFilter
     }
     
     var mostRecentQuery: String? = nil
@@ -50,12 +50,10 @@ class QueryViewController: NSViewController {
         verticalSplitView.delegate = self
         verticalSplitView.setHoldingPriority(.defaultHigh, forSubviewAt: 0)
         
-        queryTextView.font = .monospacedSystemFont(ofSize: 13, weight: .regular)
+        filterTextView.font = .monospacedSystemFont(ofSize: 13, weight: .regular)
-        queryTextView.isAutomaticQuoteSubstitutionEnabled = false
+        filterTextView.isAutomaticQuoteSubstitutionEnabled = false
-        queryTextView.string = defaultQuery
+        filterTextView.string = defaultFilter
-        queryTextView.delegate = self
+
-        highlightQuery()
-        
         outlineView.dataSource = self
         outlineView.delegate = self
         
@@ -77,38 +75,29 @@ class QueryViewController: NSViewController {
         
         view.window!.makeFirstResponder(outlineView)
     }
-    
-    func highlightQuery() {
-        _ = JavaScriptHighlighter(text: queryTextView.string).highlight(mutableAttributed: queryTextView.textStorage!)
-    }
-    
-    func refresh(reload: Bool = true) {
-        if let query = mostRecentQuery {
-            let connStr = "\(mongoController.connectionString)/\(collection.database)"
-            
-            rootNodes = MongoEvaluator.eval(command: query, connectingTo: connStr).map {
-                Node(value: $0)
-            }
-            
-            title = query
-            documentCountLabel.stringValue = "\(rootNodes.count) result\(rootNodes.count == 1 ? "" : "s")"
-        } else {
-            let documents = try! mongoController.collection(collection).find().all()
-            rootNodes = documents.map { Node(document: $0) }
 
-            title = "\(self.collection.database).\(self.collection.name)"
+    func refresh(reload: Bool = true) {
-            documentCountLabel.stringValue = "\(documents.count) document\(documents.count == 1 ? "" : "s")"
+        let filterText = filterTextView.string.trimmingCharacters(in: .whitespacesAndNewlines)
+        let filter: Document
+        if !filterText.isEmpty,
+            let normalized = ExtendedJSON.normalize(filterText),
+            let doc = try? Document(fromJSON: normalized) {
+            filter = doc
+        } else {
+            filter = [:]
         }
+        
+        let documents = try! mongoController.collection(collection).find(filter).all()
+        rootNodes = documents.map { Node(document: $0) }
+
+        title = "\(self.collection.database).\(self.collection.name)"
+        documentCountLabel.stringValue = "\(documents.count) document\(documents.count == 1 ? "" : "s")"
+        
         if reload {
             outlineView.reloadData()
         }
     }
     
-    func runQuery() {
-        mostRecentQuery = queryTextView.string
-        refresh()
-    }
-    
     func deleteRootNode(_ node: Node) {
         guard case let .document(doc) = node.value else { return }
         let alert = NSAlert()
@@ -176,16 +165,14 @@ class QueryViewController: NSViewController {
     
 }
 
-extension QueryViewController: NSTextViewDelegate {
-    func textDidChange(_ notification: Notification) {
-        highlightQuery()
-    }
-}
-
 extension QueryViewController: NSMenuItemValidation {
     func validateMenuItem(_ menuItem: NSMenuItem) -> Bool {
         if menuItem.action == #selector(deleteNode(_:)) {
-            return outlineView.clickedRow != -1
+            if outlineView.clickedRow != -1, let node = outlineView.item(atRow: outlineView.clickedRow) as? Node, node.parent == nil {
+                return true
+            } else {
+                return false
+            }
         }
         return true
     }

MongoView/View Controllers/QueryViewController.xib

@@ -1,15 +1,15 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.Cocoa.XIB" version="3.0" toolsVersion="16085" targetRuntime="MacOSX.Cocoa" propertyAccessControl="none" useAutolayout="YES" customObjectInstantitationMethod="direct">
+<document type="com.apple.InterfaceBuilder3.Cocoa.XIB" version="3.0" toolsVersion="16096" targetRuntime="MacOSX.Cocoa" propertyAccessControl="none" useAutolayout="YES" customObjectInstantitationMethod="direct">
     <dependencies>
-        <plugIn identifier="com.apple.InterfaceBuilder.CocoaPlugin" version="16085"/>
+        <plugIn identifier="com.apple.InterfaceBuilder.CocoaPlugin" version="16096"/>
         <capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
     </dependencies>
     <objects>
         <customObject id="-2" userLabel="File's Owner" customClass="QueryViewController" customModule="MongoView" customModuleProvider="target">
             <connections>
                 <outlet property="documentCountLabel" destination="bwd-po-tUv" id="5Ip-AZ-nVc"/>
+                <outlet property="filterTextView" destination="f8D-lV-IMK" id="dq6-lR-LeP"/>
                 <outlet property="outlineView" destination="uu9-9q-MWr" id="BEH-jX-h0S"/>
-                <outlet property="queryTextView" destination="f8D-lV-IMK" id="dq6-lR-LeP"/>
                 <outlet property="verticalSplitView" destination="E3D-CM-hk3" id="3Ks-1t-OFW"/>
                 <outlet property="view" destination="Hz6-mo-xeY" id="0bl-1N-x8E"/>
             </connections>
@@ -38,7 +38,7 @@
                                 <rect key="frame" x="0.0" y="0.0" width="726" height="86"/>
                                 <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
                                 <subviews>
-                                    <textView importsGraphics="NO" richText="NO" verticallyResizable="YES" allowsCharacterPickerTouchBarItem="NO" textCompletion="NO" id="f8D-lV-IMK">
+                                    <textView importsGraphics="NO" richText="NO" verticallyResizable="YES" allowsCharacterPickerTouchBarItem="NO" textCompletion="NO" id="f8D-lV-IMK" customClass="JavaScriptEditorView" customModule="MongoView" customModuleProvider="target">
                                         <rect key="frame" x="0.0" y="0.0" width="726" height="86"/>
                                         <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
                                         <color key="textColor" name="textColor" catalog="System" colorSpace="catalog"/>

MongoView/Views/JavaScriptEditorView.swift

@@ -0,0 +1,85 @@
+//
+//  JavaScriptEditorView.swift
+//  MongoView
+//
+//  Created by Shadowfacts on 4/4/20.
+//  Copyright © 2020 Shadowfacts. All rights reserved.
+//
+
+import AppKit
+
+class JavaScriptEditorView: NSTextView {
+    
+    var highlighter: JavaScriptHighlighter?
+    
+    override var string: String {
+        didSet {
+            rehighlight()
+        }
+    }
+
+    func rehighlight() {
+        highlighter = JavaScriptHighlighter(mutableAttributed: self.textStorage!)
+        highlighter!.highlight()
+    }
+    
+    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,
+            tryAutocompleteCharacter(for: affectedCharRange, string: string) {
+            return false
+        }
+
+        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()
+            setSelectedRange(NSRange(location: range.location + 1, length: 0))
+            return true
+        } else {
+            return false
+        }
+    }
+    
+    private func autocompleteResultFor(string: String, in range: NSRange) -> String? {
+        switch string {
+        case "'", "\"", "`":
+            return token(at: range.location) != .string ? string : nil
+        case "(":
+            return token(at: range.location) != .string ? ")" : nil
+        case "[":
+            return token(at: range.location) != .string ? "]" : nil
+        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]
+            }
+            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
+        }
+        return nil
+    }
+}

jstest/main.swift

@@ -8,6 +8,6 @@
 
 import Foundation
 
-let source = "{a: 1234 , 'b': 'blah', foo}"
+let source = "`$foo ${blah} bar`"
 _ = JavaScriptHighlighter(text: source).highlight()