Package collections: improve search performance

Sources/Basics/ConcurrencyHelpers.swift

@@ -126,7 +126,7 @@ public final class ThreadSafeArrayStore<Value> {
     }
 }
 
-/// Thread-safe value boxing  structure
+/// Thread-safe value boxing structure
 public final class ThreadSafeBox<Value> {
     private var underlying: Value?
     private let lock = Lock()

Sources/PackageCollections/CMakeLists.txt

@@ -23,6 +23,7 @@ add_library(PackageCollections
   Storage/PackageCollectionsSourcesStorage.swift
   Storage/PackageCollectionsStorage.swift
   Storage/SQLitePackageCollectionsStorage.swift
+  Storage/Trie.swift
   API.swift
   PackageCollections.swift
   PackageCollections+Configuration.swift

Sources/PackageCollections/Model/TargetListResult.swift

@@ -48,11 +48,15 @@ extension PackageCollectionsModel.TargetListResult {
 
 extension PackageCollectionsModel.TargetListResult {
     /// Represents a package version
-    public struct PackageVersion: Hashable, Encodable {
+    public struct PackageVersion: Hashable, Encodable, Comparable {
         /// The version
         public let version: TSCUtility.Version
 
         /// Package name
         public let packageName: String
+
+        public static func < (lhs: PackageVersion, rhs: PackageVersion) -> Bool {
+            lhs.version < rhs.version
+        }
     }
 }

Sources/PackageCollections/Storage/PackageCollectionsStorage.swift

@@ -53,7 +53,7 @@ public protocol PackageCollectionsStorage {
                         query: String,
                         callback: @escaping (Result<PackageCollectionsModel.PackageSearchResult, Error>) -> Void)
 
-    /// Returns optional `PackageSearchResult.Item` for the given package identity.
+    /// Returns `PackageSearchResult.Item` for the given package identity.
     ///
     /// - Parameters:
     ///   - identifier: The package identifier

Sources/PackageCollections/Storage/Trie.swift

@@ -0,0 +1,253 @@
+/*
+ This source file is part of the Swift.org open source project
+ Copyright (c) 2020 Apple Inc. and the Swift project authors
+ Licensed under Apache License v2.0 with Runtime Library Exception
+ See http://swift.org/LICENSE.txt for license information
+ See http://swift.org/CONTRIBUTORS.txt for Swift project authors
+ */
+
+import TSCBasic
+
+import PackageModel
+
+struct Trie<Document: Hashable> {
+    private typealias Node = TrieNode<Character, Document>
+
+    private let root: Node
+
+    init() {
+        self.root = Node()
+    }
+
+    /// Inserts a word and its document to the trie.
+    func insert(word: String, foundIn document: Document) {
+        guard !word.isEmpty else { return }
+
+        var currentNode = self.root
+        // Check if word already exists otherwise creates the node path
+        for character in word.lowercased() {
+            if let child = currentNode.children[character] {
+                currentNode = child
+            } else {
+                currentNode = currentNode.add(value: character)
+            }
+        }
+
+        currentNode.add(document: document)
+    }
+
+    /// Removes word occurrences found in the given document.
+    func remove(document: Document) {
+        func removeInSubTrie(root: Node, document: Document) {
+            if root.isTerminating {
+                root.remove(document: document)
+            }
+
+            // Clean up sub-tries
+            root.children.values.forEach {
+                removeInSubTrie(root: $0, document: document)
+            }
+
+            root.children.forEach { value, node in
+                // If a child node doesn't have children (i.e., there are no words under it),
+                // and itself is not a word, delete it since its path has become a deadend.
+                if node.isLeaf, !node.isTerminating {
+                    root.remove(value: value)
+                }
+            }
+        }
+
+        removeInSubTrie(root: self.root, document: document)
+    }
+
+    /// Removes word occurrences found in matching document(s).
+    func remove(where predicate: @escaping (Document) -> Bool) {
+        func removeInSubTrie(root: Node, where predicate: @escaping (Document) -> Bool) {
+            if root.isTerminating {
+                root.remove(where: predicate)
+            }
+
+            // Clean up sub-tries
+            root.children.values.forEach {
+                removeInSubTrie(root: $0, where: predicate)
+            }
+
+            root.children.forEach { value, node in
+                // If a child node doesn't have children (i.e., there are no words under it),
+                // and itself is not a word, delete it since its path has become a deadend.
+                if node.isLeaf, !node.isTerminating {
+                    root.remove(value: value)
+                }
+            }
+        }
+
+        removeInSubTrie(root: self.root, where: predicate)
+    }
+
+    /// Checks if the trie contains the exact word or words with matching prefix.
+    func contains(word: String, prefixMatch: Bool = false) -> Bool {
+        guard let node = self.findLastNodeOf(word: word) else {
+            return false
+        }
+        return prefixMatch || node.isTerminating
+    }
+
+    /// Finds the word in this trie and returns its documents.
+    func find(word: String) throws -> Set<Document> {
+        guard let node = self.findLastNodeOf(word: word), node.isTerminating else {
+            throw NotFoundError(word)
+        }
+        return node.documents
+    }
+
+    /// Finds words with matching prefix in this trie and returns their documents.
+    func findWithPrefix(_ prefix: String) throws -> [String: Set<Document>] {
+        guard let node = self.findLastNodeOf(word: prefix) else {
+            throw NotFoundError(prefix)
+        }
+
+        func wordsInSubTrie(root: Node, prefix: String) -> [String: Set<Document>] {
+            precondition(root.value != nil, "Sub-trie root's value should not be nil")
+
+            var subTrieWords = [String: Set<Document>]()
+
+            // Construct the new prefix by adding the sub-trie root's character
+            var previousCharacters = prefix
+            previousCharacters.append(root.value!.lowercased()) // !-safe; see precondition
+
+            // The root actually forms a word
+            if root.isTerminating {
+                subTrieWords[previousCharacters] = root.documents
+            }
+
+            // Collect all words under this sub-trie
+            root.children.values.forEach {
+                let childWords = wordsInSubTrie(root: $0, prefix: previousCharacters)
+                subTrieWords.merge(childWords, uniquingKeysWith: { _, child in child })
+            }
+
+            return subTrieWords
+        }
+
+        var words = [String: Set<Document>]()
+
+        let prefix = prefix.lowercased()
+        // The prefix is actually a word
+        if node.isTerminating {
+            words[prefix] = node.documents
+        }
+
+        node.children.values.forEach {
+            let childWords = wordsInSubTrie(root: $0, prefix: prefix)
+            words.merge(childWords, uniquingKeysWith: { _, child in child })
+        }
+
+        return words
+    }
+
+    /// Finds the last node in the path of the given word if it exists in this trie.
+    private func findLastNodeOf(word: String) -> Node? {
+        guard !word.isEmpty else { return nil }
+
+        var currentNode = self.root
+        // Traverse down the trie as far as we can
+        for character in word.lowercased() {
+            guard let child = currentNode.children[character] else {
+                return nil
+            }
+            currentNode = child
+        }
+        return currentNode
+    }
+}
+
+private final class TrieNode<T: Hashable, Document: Hashable> {
+    /// The value (i.e., character) that this node stores. `nil` if root.
+    let value: T?
+
+    /// The parent of this node. `nil` if root.
+    private weak var parent: TrieNode<T, Document>?
+
+    /// The children of this node identified by their corresponding value.
+    private var _children = [T: TrieNode<T, Document>]()
+    private let childrenLock = Lock()
+
+    /// If the path to this node forms a valid word, these are the documents where the word can be found.
+    private var _documents = Set<Document>()
+    private let documentsLock = Lock()
+
+    var isLeaf: Bool {
+        self.childrenLock.withLock {
+            self._children.isEmpty
+        }
+    }
+
+    /// `true` indicates the path to this node forms a valid word.
+    var isTerminating: Bool {
+        self.documentsLock.withLock {
+            !self._documents.isEmpty
+        }
+    }
+
+    var children: [T: TrieNode<T, Document>] {
+        self.childrenLock.withLock {
+            self._children
+        }
+    }
+
+    var documents: Set<Document> {
+        self.documentsLock.withLock {
+            self._documents
+        }
+    }
+
+    init(value: T? = nil, parent: TrieNode<T, Document>? = nil) {
+        self.value = value
+        self.parent = parent
+    }
+
+    /// Adds a subpath under this node.
+    func add(value: T) -> TrieNode<T, Document> {
+        self.childrenLock.withLock {
+            if let existing = self._children[value] {
+                return existing
+            }
+
+            let child = TrieNode<T, Document>(value: value, parent: self)
+            self._children[value] = child
+            return child
+        }
+    }
+
+    /// Removes a subpath from this node.
+    func remove(value: T) {
+        _ = self.childrenLock.withLock {
+            self._children.removeValue(forKey: value)
+        }
+    }
+
+    /// Adds a document in which the word formed by path leading to this node can be found.
+    func add(document: Document) {
+        _ = self.documentsLock.withLock {
+            self._documents.insert(document)
+        }
+    }
+
+    /// Removes a referenced document.
+    func remove(document: Document) {
+        _ = self.documentsLock.withLock {
+            self._documents.remove(document)
+        }
+    }
+
+    /// Removes documents that satisfy the given predicate.
+    func remove(where predicate: @escaping (Document) -> Bool) {
+        self.documentsLock.withLock {
+            for document in self._documents {
+                if predicate(document) {
+                    self._documents.remove(document)
+                }
+            }
+        }
+    }
+}

Tests/PackageCollectionsTests/PackageCollectionsStorageTests.swift

@@ -196,4 +196,37 @@ class PackageCollectionsStorageTests: XCTestCase {
         _ = try tsc_await { callback in storage.put(collection: mockCollections.last!, callback: callback) }
         XCTAssertEqual(list.count, mockCollections.count)
     }
+
+    func testPopulateTargetTrie() throws {
+        try testWithTemporaryDirectory { tmpPath in
+            let path = tmpPath.appending(component: "test.db")
+            let storage = SQLitePackageCollectionsStorage(path: path)
+            defer { XCTAssertNoThrow(try storage.close()) }
+
+            let mockCollections = makeMockCollections(count: 3)
+            try mockCollections.forEach { collection in
+                _ = try tsc_await { callback in storage.put(collection: collection, callback: callback) }
+            }
+
+            let targetName = mockCollections.last!.packages.last!.versions.last!.targets.last!.name
+
+            do {
+                let searchResult = try tsc_await { callback in storage.searchTargets(query: targetName, type: .exactMatch, callback: callback) }
+                XCTAssert(searchResult.items.count > 0, "should get results")
+            }
+
+            // Create another instance, which should read existing data and populate target trie with it.
+            // Since we are not calling `storage2.put`, there is no other way for target trie to get populated.
+            let storage2 = SQLitePackageCollectionsStorage(path: path)
+            defer { XCTAssertNoThrow(try storage2.close()) }
+
+            // populateTargetTrie is called in `.init`; call it again explicitly so we know when it's finished
+            try tsc_await { callback in storage2.populateTargetTrie(callback: callback) }
+
+            do {
+                let searchResult = try tsc_await { callback in storage2.searchTargets(query: targetName, type: .exactMatch, callback: callback) }
+                XCTAssert(searchResult.items.count > 0, "should get results")
+            }
+        }
+    }
 }