From b31cf6fe94238db0388edf74c3482b5b6a0c16bb Mon Sep 17 00:00:00 2001
From: Matt <85322+mattmassicotte@users.noreply.github.com>
Date: Tue, 10 Sep 2024 16:56:14 -0400
Subject: [PATCH] Throw if query is unsupported

---
 .../SwiftTreeSitterLayer/LanguageLayer.swift  |  2 +-
 Sources/SwiftTreeSitterLayer/Snapshots.swift  | 20 +++++++-------
 .../LanguageLayerTests.swift                  | 26 +++++++++++++++++++
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/Sources/SwiftTreeSitterLayer/LanguageLayer.swift b/Sources/SwiftTreeSitterLayer/LanguageLayer.swift
index e694ebf..b26b34e 100644
--- a/Sources/SwiftTreeSitterLayer/LanguageLayer.swift
+++ b/Sources/SwiftTreeSitterLayer/LanguageLayer.swift
@@ -2,7 +2,7 @@ import Foundation
 
 import SwiftTreeSitter
 
-public enum LanguageLayerError: Error {
+public enum LanguageLayerError: Error, Hashable {
 	case noRootNode
 	case queryUnavailable(String, Query.Definition)
 }
diff --git a/Sources/SwiftTreeSitterLayer/Snapshots.swift b/Sources/SwiftTreeSitterLayer/Snapshots.swift
index 39f50a3..7dccf82 100644
--- a/Sources/SwiftTreeSitterLayer/Snapshots.swift
+++ b/Sources/SwiftTreeSitterLayer/Snapshots.swift
@@ -35,9 +35,9 @@ public struct LanguageLayerSnapshot: Sendable {
 		return set
 	}
 
-	func queryTarget(for queryDef: Query.Definition) -> LanguageTreeQueryCursor.Target? {
+	func queryTarget(for queryDef: Query.Definition) throws -> LanguageTreeQueryCursor.Target {
 		guard let query = data.queries[queryDef] else {
-			return nil
+			throw LanguageLayerError.queryUnavailable(data.name, queryDef)
 		}
 
 		return (tree, query, depth, data.name)
@@ -47,9 +47,7 @@ public struct LanguageLayerSnapshot: Sendable {
 extension LanguageLayerSnapshot: Queryable {
 	/// Run a query against the snapshot.
 	public func executeQuery(_ queryDef: Query.Definition, in set: IndexSet) throws -> LanguageLayerQueryCursor {
-		guard let target = queryTarget(for: queryDef) else {
-			throw LanguageLayerError.queryUnavailable(data.name, queryDef)
-		}
+		let target = try queryTarget(for: queryDef)
 
 		return LanguageLayerQueryCursor(target: target, set: set)
 	}
@@ -77,13 +75,13 @@ public struct LanguageLayerTreeSnapshot: Sendable {
 		}
 	}
 
-	private func queryTargets(in set: IndexSet, for queryDef: Query.Definition) -> [LanguageTreeQueryCursor.Target] {
+	private func queryTargets(in set: IndexSet, for queryDef: Query.Definition) throws -> [LanguageTreeQueryCursor.Target] {
 		var targets = [LanguageTreeQueryCursor.Target]()
 
-		enumerateSnapshots(in: set) { snapshot in
-			if let target = snapshot.queryTarget(for: queryDef) {
-				targets.append(target)
-			}
+		try enumerateSnapshots(in: set) { snapshot in
+			let target = try snapshot.queryTarget(for: queryDef)
+
+			targets.append(target)
 		}
 
 		return targets
@@ -92,7 +90,7 @@ public struct LanguageLayerTreeSnapshot: Sendable {
 
 extension LanguageLayerTreeSnapshot: Queryable {
 	public func executeQuery(_ queryDef: Query.Definition, in set: IndexSet) throws -> LanguageTreeQueryCursor {
-		let targets = queryTargets(in: set, for: queryDef)
+		let targets = try queryTargets(in: set, for: queryDef)
 
 		return LanguageTreeQueryCursor(set: set, targets: targets)
 	}
diff --git a/Tests/SwiftTreeSitterLayerTests/LanguageLayerTests.swift b/Tests/SwiftTreeSitterLayerTests/LanguageLayerTests.swift
index a714cd3..0d164d0 100644
--- a/Tests/SwiftTreeSitterLayerTests/LanguageLayerTests.swift
+++ b/Tests/SwiftTreeSitterLayerTests/LanguageLayerTests.swift
@@ -258,6 +258,32 @@ let c = "var c = 1"
 
 		XCTAssertEqual(highlights3, expected3)
 	}
+
+	func testThrowsOnNoHighlightQuery() throws {
+		let language = Language(language: tree_sitter_swift())
+		let swiftConfig = LanguageConfiguration(
+			language,
+			name: "Swift",
+			queries: [:]
+		)
+
+		let config = LanguageLayer.Configuration(languageProvider: { name in
+			precondition(name == "swift")
+
+			return swiftConfig
+		})
+
+		let tree = try LanguageLayer(languageConfig: swiftConfig, configuration: config)
+
+		let text = """
+let a = "var a = 1"
+"""
+		tree.replaceContent(with: text)
+
+		XCTAssertThrowsError(try tree.highlights(in: NSRange(0..<10), provider: { _, _ in nil })) { error in
+			XCTAssertEqual(error as? LanguageLayerError, LanguageLayerError.queryUnavailable("Swift", .highlights))
+		}
+	}
 }
 
 #endif