Package collections: trie for search

Sources/PackageCollections/CMakeLists.txt

@@ -19,6 +19,9 @@ add_library(PackageCollections
   Providers/JSONPackageCollectionProvider.swift
   Providers/PackageCollectionProvider.swift
   Providers/PackageMetadataProvider.swift
+  Search/InMemoryPackageCollectionsSearch.swift
+  Search/PackageCollectionsSearch.swift
+  Search/Trie.swift
   Storage/FilePackageCollectionsSourcesStorage.swift
   Storage/PackageCollectionsSourcesStorage.swift
   Storage/PackageCollectionsStorage.swift

Sources/PackageCollections/Model/TargetListResult.swift

@@ -48,11 +48,15 @@ extension PackageCollectionsModel.TargetListResult {
 
 extension PackageCollectionsModel.TargetListResult {
     /// Represents a package version
-    public struct PackageVersion: Hashable {
+    public struct PackageVersion: Hashable, 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/PackageCollections.swift

@@ -105,7 +105,7 @@ public struct PackageCollections: PackageCollectionsProtocol {
                 sources.forEach { source in
                     self.refreshCollectionFromSource(source: source) { refreshResult in
                         lock.withLock { refreshResults.append(refreshResult) }
-                        if refreshResults.count == (lock.withLock { sources.count }) {
+                        if sources.count == (lock.withLock { refreshResults.count }) {
                             let errors = refreshResults.compactMap { $0.failure }
                             callback(errors.isEmpty ? .success(sources) : .failure(MultipleErrors(errors)))
                         }

Sources/PackageCollections/Search/PackageCollectionsSearch.swift

@@ -0,0 +1,61 @@
+/*
+ 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 PackageModel
+
+protocol PackageCollectionsSearch {
+    /// Adds the given `PackageCollectionsModel.Collection` to the search index.
+    ///
+    /// - Parameters:
+    ///   - collection: The `PackageCollectionsModel.Collection` to index
+    ///   - callback: The closure to invoke when result becomes available
+    func index(collection: Model.Collection,
+               callback: @escaping (Result<Void, Error>) -> Void)
+
+    /// Removes the `PackageCollectionsModel.Collection` from the search index.
+    ///
+    /// - Parameters:
+    ///   - identifier: The identifier of the `PackageCollectionsModel.Collection` to remove
+    ///   - callback: The closure to invoke when result becomes available
+    func remove(identifier: Model.CollectionIdentifier,
+                callback: @escaping (Result<Void, Error>) -> Void)
+
+    /// Returns `PackageSearchResult.Item` for the given package identity.
+    ///
+    /// - Parameters:
+    ///   - identifier: The package identifier
+    ///   - collectionIdentifiers: Optional. The identifiers of the `PackageCollectionsModel.Collection`s to search under.
+    ///   - callback: The closure to invoke when result becomes available
+    func findPackage(identifier: PackageIdentity,
+                     collectionIdentifiers: [Model.CollectionIdentifier]?,
+                     callback: @escaping (Result<Model.PackageSearchResult.Item, Error>) -> Void)
+
+    /// Returns `PackageSearchResult` for the given search criteria.
+    ///
+    /// - Parameters:
+    ///   - identifiers: Optional. The identifiers of the `PackageCollectionsModel.Collection`s to search under.
+    ///   - query: The search query expression
+    ///   - callback: The closure to invoke when result becomes available
+    func searchPackages(identifiers: [Model.CollectionIdentifier]?,
+                        query: String,
+                        callback: @escaping (Result<Model.PackageSearchResult, Error>) -> Void)
+
+    /// Returns `TargetSearchResult` for the given search criteria.
+    ///
+    /// - Parameters:
+    ///   - identifiers: Optional. The identifiers of the `PackageCollectionsModel.Collection`s  to search under.
+    ///   - query: The search query expression
+    ///   - type: The search type
+    ///   - callback: The closure to invoke when result becomes available
+    func searchTargets(identifiers: [Model.CollectionIdentifier]?,
+                       query: String,
+                       type: Model.TargetSearchType,
+                       callback: @escaping (Result<Model.TargetSearchResult, Error>) -> Void)
+}

Sources/PackageCollections/Search/Trie.swift

@@ -0,0 +1,220 @@
+/*
+ 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
+
+final class 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)
+    }
+
+    /// 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)
+        }
+    }
+}

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