krz/orgstar

A native macOS editor for org-mode files. editor org-mode swift

Sources/OrgIndex/IndexStore.swift

778058aafc7ddacee4a613b1da87c60aa0c8e113
orgstar/Sources/OrgIndex/IndexStore.swift history · blame · raw

408 lines · 18829 bytes

20 symbols in this file
  1import Foundation
  2import GRDB
  3
  4/// What the index knows about a file without reading it.
  5public struct FileState: Sendable, Equatable {
  6    public let kind: FileKind
  7    public let size: Int
  8    public let mtime: Double
  9    public let hash: String
 10    public let settingsVersion: Int
 11}
 12
 13/// A heading as found by a query. `contentHash` is the hash of the text the offsets refer to,
 14/// so a caller can tell whether they still apply to an open buffer.
 15public struct IndexedHeading: Sendable, Hashable {
 16    public let location: HeadingLocation
 17    public let level: Int
 18    /// Titles from the top-level ancestor down to the heading itself.
 19    public let outlinePath: [String]
 20}
 21
 22public struct HeadingLocation: Sendable, Hashable {
 23    public let path: String
 24    /// `fileOrdinal` for the file itself, before its first heading.
 25    public let ordinal: Int
 26    public let title: String
 27    public let start: Int
 28    public let contentHash: String
 29
 30    public init(path: String, ordinal: Int, title: String, start: Int, contentHash: String) {
 31        self.path = path
 32        self.ordinal = ordinal
 33        self.title = title
 34        self.start = start
 35        self.contentHash = contentHash
 36    }
 37
 38    public static let fileOrdinal = -1
 39
 40    /// The file before its first heading, titled by its name.
 41    public static func file(_ path: String, contentHash: String) -> HeadingLocation {
 42        HeadingLocation(path: path, ordinal: fileOrdinal, title: (path as NSString).lastPathComponent, start: 0, contentHash: contentHash)
 43    }
 44}
 45
 46/// One reconciliation's worth of changes, applied in a single transaction.
 47public struct IndexChange: Sendable {
 48    public var records: [FileRecord] = []
 49    /// Unchanged content with a new modification time.
 50    public var touches: [(path: String, mtime: Double)] = []
 51    /// Renamed files whose content didn't change.
 52    public var moves: [(from: String, to: String, mtime: Double)] = []
 53    public var removals: [String] = []
 54
 55    public init() {}
 56
 57    public var isEmpty: Bool { records.isEmpty && touches.isEmpty && moves.isEmpty && removals.isEmpty }
 58}
 59
 60/// The SQLite index. A cache: deleting it loses nothing that the files don't hold.
 61public final class IndexStore: Sendable {
 62    let database: DatabaseQueue
 63
 64    /// `path` nil opens an in-memory index.
 65    public init(path: String? = nil) throws {
 66        database = try path.map { try DatabaseQueue(path: $0) } ?? DatabaseQueue()
 67        try Self.migrator.migrate(database)
 68    }
 69
 70    static var migrator: DatabaseMigrator {
 71        var migrator = DatabaseMigrator()
 72        migrator.registerMigration("v1") { db in
 73            try db.execute(sql: """
 74                CREATE TABLE files (
 75                    id INTEGER PRIMARY KEY,
 76                    path TEXT NOT NULL UNIQUE,
 77                    root TEXT NOT NULL,
 78                    kind TEXT NOT NULL,
 79                    size INTEGER NOT NULL,
 80                    mtime REAL NOT NULL,
 81                    hash TEXT NOT NULL,
 82                    settings_version INTEGER NOT NULL,
 83                    parsed_at REAL NOT NULL
 84                );
 85                CREATE INDEX files_root ON files(root);
 86                CREATE TABLE headings (
 87                    id INTEGER PRIMARY KEY,
 88                    file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
 89                    ordinal INTEGER NOT NULL,
 90                    parent_ordinal INTEGER,
 91                    start_offset INTEGER NOT NULL,
 92                    end_offset INTEGER NOT NULL,
 93                    level INTEGER NOT NULL,
 94                    todo TEXT,
 95                    is_done INTEGER NOT NULL,
 96                    priority TEXT,
 97                    title TEXT NOT NULL,
 98                    outline_path TEXT NOT NULL,
 99                    org_id TEXT,
100                    archived INTEGER NOT NULL
101                );
102                CREATE INDEX headings_file ON headings(file_id);
103                CREATE INDEX headings_org_id ON headings(org_id);
104                CREATE TABLE tags (
105                    heading_id INTEGER NOT NULL REFERENCES headings(id) ON DELETE CASCADE,
106                    tag TEXT NOT NULL,
107                    inherited INTEGER NOT NULL
108                );
109                CREATE INDEX tags_heading ON tags(heading_id);
110                CREATE TABLE properties (
111                    heading_id INTEGER NOT NULL REFERENCES headings(id) ON DELETE CASCADE,
112                    key TEXT NOT NULL,
113                    value TEXT NOT NULL,
114                    inherited INTEGER NOT NULL
115                );
116                CREATE INDEX properties_heading ON properties(heading_id);
117                CREATE TABLE timestamps (
118                    heading_id INTEGER NOT NULL REFERENCES headings(id) ON DELETE CASCADE,
119                    kind TEXT NOT NULL,
120                    start_at TEXT NOT NULL,
121                    end_at TEXT,
122                    repeater TEXT,
123                    warning TEXT
124                );
125                CREATE INDEX timestamps_heading ON timestamps(heading_id);
126                CREATE TABLE clocks (
127                    heading_id INTEGER NOT NULL REFERENCES headings(id) ON DELETE CASCADE,
128                    start_at TEXT NOT NULL,
129                    end_at TEXT,
130                    minutes INTEGER
131                );
132                CREATE INDEX clocks_heading ON clocks(heading_id);
133                CREATE TABLE links (
134                    heading_id INTEGER NOT NULL REFERENCES headings(id) ON DELETE CASCADE,
135                    type TEXT NOT NULL,
136                    target TEXT NOT NULL
137                );
138                CREATE INDEX links_heading ON links(heading_id);
139                CREATE VIRTUAL TABLE headings_fts USING fts5(title, body, tokenize = 'unicode61 remove_diacritics 2');
140                CREATE TRIGGER headings_fts_delete AFTER DELETE ON headings BEGIN
141                    DELETE FROM headings_fts WHERE rowid = old.id;
142                END;
143                """)
144        }
145        migrator.registerMigration("v2") { db in
146            // Links before a file's first heading; files are read again to fill it.
147            try db.execute(sql: """
148                CREATE TABLE file_links (
149                    file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
150                    type TEXT NOT NULL,
151                    target TEXT NOT NULL
152                );
153                CREATE INDEX file_links_file ON file_links(file_id);
154                UPDATE files SET settings_version = -1;
155                """)
156        }
157        return migrator
158    }
159
160    // MARK: - Writing
161
162    public func write(_ record: FileRecord) throws {
163        var change = IndexChange()
164        change.records = [record]
165        try apply(change)
166    }
167
168    public func apply(_ change: IndexChange) throws {
169        guard !change.isEmpty else { return }
170        try database.write { db in
171            for path in change.removals {
172                try db.execute(sql: "DELETE FROM files WHERE path = ?", arguments: [path])
173            }
174            for move in change.moves {
175                try db.execute(sql: "UPDATE files SET path = ?, mtime = ? WHERE path = ?", arguments: [move.to, move.mtime, move.from])
176            }
177            for touch in change.touches {
178                try db.execute(sql: "UPDATE files SET mtime = ? WHERE path = ?", arguments: [touch.mtime, touch.path])
179            }
180            for record in change.records {
181                try insert(record, db)
182            }
183        }
184    }
185
186    private func insert(_ record: FileRecord, _ db: Database) throws {
187        try db.execute(sql: "DELETE FROM files WHERE path = ?", arguments: [record.path])
188        try db.execute(
189            sql: """
190                INSERT INTO files (path, root, kind, size, mtime, hash, settings_version, parsed_at)
191                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
192                """,
193            arguments: [record.path, record.root, record.kind.rawValue, record.size, record.mtime, record.hash,
194                        record.settingsVersion, Date().timeIntervalSince1970]
195        )
196        let fileID = db.lastInsertedRowID
197        for link in record.links {
198            try db.execute(sql: "INSERT INTO file_links VALUES (?, ?, ?)", arguments: [fileID, link.type, link.target])
199        }
200        for heading in record.headings {
201            try db.execute(
202                sql: """
203                    INSERT INTO headings (file_id, ordinal, parent_ordinal, start_offset, end_offset, level, todo,
204                                          is_done, priority, title, outline_path, org_id, archived)
205                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
206                    """,
207                arguments: [fileID, heading.ordinal, heading.parent, heading.start, heading.end, heading.level, heading.todo,
208                            heading.isDone, heading.priority, heading.title, heading.outlinePath.joined(separator: "\u{1F}"),
209                            heading.orgID, heading.archived]
210            )
211            let id = db.lastInsertedRowID
212            try db.execute(sql: "INSERT INTO headings_fts (rowid, title, body) VALUES (?, ?, ?)", arguments: [id, heading.title, heading.body])
213            for tag in heading.tags {
214                try db.execute(sql: "INSERT INTO tags VALUES (?, ?, ?)", arguments: [id, tag.name, tag.inherited])
215            }
216            for property in heading.properties {
217                try db.execute(sql: "INSERT INTO properties VALUES (?, ?, ?, ?)", arguments: [id, property.key, property.value, property.inherited])
218            }
219            for stamp in heading.timestamps {
220                try db.execute(
221                    sql: "INSERT INTO timestamps VALUES (?, ?, ?, ?, ?, ?)",
222                    arguments: [id, stamp.kind.rawValue, stamp.start, stamp.end, stamp.repeater, stamp.warning]
223                )
224            }
225            for clock in heading.clocks {
226                try db.execute(sql: "INSERT INTO clocks VALUES (?, ?, ?, ?)", arguments: [id, clock.start, clock.end, clock.minutes])
227            }
228            for link in heading.links {
229                try db.execute(sql: "INSERT INTO links VALUES (?, ?, ?)", arguments: [id, link.type, link.target])
230            }
231        }
232    }
233
234    // MARK: - Reading
235
236    /// Indexed files under `root`, by path.
237    public func fileStates(root: String) throws -> [String: FileState] {
238        try database.read { db in
239            let rows = try Row.fetchAll(db, sql: "SELECT path, kind, size, mtime, hash, settings_version FROM files WHERE root = ?", arguments: [root])
240            var states: [String: FileState] = [:]
241            for row in rows {
242                states[row["path"]] = FileState(
243                    kind: FileKind(rawValue: row["kind"]) ?? .org,
244                    size: row["size"], mtime: row["mtime"], hash: row["hash"], settingsVersion: row["settings_version"]
245                )
246            }
247            return states
248        }
249    }
250
251    public func files() throws -> [(path: String, kind: FileKind)] {
252        try database.read { db in
253            try Row.fetchAll(db, sql: "SELECT path, kind FROM files ORDER BY path").map {
254                (path: $0["path"], kind: FileKind(rawValue: $0["kind"]) ?? .org)
255            }
256        }
257    }
258
259    // MARK: - Queries
260    //
261    // `overlay` holds records for open documents with unsaved edits, keyed by path. Their rows
262    // replace that file's indexed rows in every query.
263
264    /// Headings whose title or body contain every word of `query` as a word prefix.
265    public func search(_ query: String, overlay: [String: FileRecord] = [:], limit: Int = 50) throws -> [HeadingLocation] {
266        let terms = query.split(whereSeparator: \.isWhitespace).map(String.init)
267        guard !terms.isEmpty else { return [] }
268        let match = terms.map { "\"" + $0.replacingOccurrences(of: "\"", with: "\"\"") + "\"*" }.joined(separator: " ")
269        let indexed = try database.read { db in
270            try Row.fetchAll(
271                db,
272                sql: """
273                    SELECT f.path, f.hash, h.ordinal, h.title, h.start_offset
274                    FROM headings_fts
275                    JOIN headings h ON h.id = headings_fts.rowid
276                    JOIN files f ON f.id = h.file_id
277                    WHERE headings_fts MATCH ?
278                    ORDER BY bm25(headings_fts)
279                    LIMIT ?
280                    """,
281                arguments: [match, limit + overlay.count * 10]
282            ).map(location)
283        }
284        // Unsaved buffers get the same word-prefix matching as the full-text index.
285        let prefixes = terms.flatMap(Self.words)
286        let live = overlay.values.sorted { $0.path < $1.path }.flatMap { record in
287            record.headings.filter { heading in
288                let words = Self.words(heading.title + "\n" + heading.body)
289                return prefixes.allSatisfy { prefix in words.contains { $0.hasPrefix(prefix) } }
290            }.map { location(record, $0) }
291        }
292        return Array((live + indexed.filter { overlay[$0.path] == nil }).prefix(limit))
293    }
294
295    /// Headings down to `maxLevel` in org files (not archives or conflict copies), with their
296    /// levels and outline paths, in file order: refile targets.
297    public func outline(maxLevel: Int) throws -> [IndexedHeading] {
298        try database.read { db in
299            try Row.fetchAll(
300                db,
301                sql: """
302                    SELECT f.path, f.hash, h.ordinal, h.title, h.start_offset, h.level, h.outline_path
303                    FROM headings h JOIN files f ON f.id = h.file_id
304                    WHERE h.level <= ? AND f.kind = 'org' ORDER BY f.path, h.ordinal
305                    """,
306                arguments: [maxLevel]
307            ).map { row in
308                IndexedHeading(
309                    location: location(row), level: row["level"],
310                    outlinePath: (row["outline_path"] as String).split(separator: "\u{1F}", omittingEmptySubsequences: false).map(String.init)
311                )
312            }
313        }
314    }
315
316    /// Headings with `:ID: id`. More than one means the ID is duplicated.
317    /// Every tag a heading has, most used first.
318    public func tags() throws -> [String] {
319        try database.read { db in
320            try String.fetchAll(db, sql: "SELECT tag FROM tags WHERE inherited = 0 GROUP BY tag ORDER BY count(*) DESC, tag")
321        }
322    }
323
324    public func headings(withID id: String, overlay: [String: FileRecord] = [:]) throws -> [HeadingLocation] {
325        let indexed = try database.read { db in
326            try Row.fetchAll(
327                db,
328                sql: """
329                    SELECT f.path, f.hash, h.ordinal, h.title, h.start_offset
330                    FROM headings h JOIN files f ON f.id = h.file_id
331                    WHERE h.org_id = ? ORDER BY f.path, h.ordinal
332                    """,
333                arguments: [id]
334            ).map(location)
335        }
336        let live = overlay.values.sorted { $0.path < $1.path }.flatMap { record in
337            record.headings.filter { $0.orgID == id }.map { location(record, $0) }
338        }
339        return indexed.filter { overlay[$0.path] == nil } + live
340    }
341
342    /// Every link and the heading it's in, or the file for links before its first heading,
343    /// unsaved buffers' in place of their files'.
344    public func links(overlay: [String: FileRecord] = [:]) throws -> [(source: HeadingLocation, target: String)] {
345        let indexed = try database.read { db in
346            try Row.fetchAll(
347                db,
348                sql: """
349                    SELECT f.path, f.hash, -1 AS ordinal, '' AS title, 0 AS start_offset, l.target
350                    FROM file_links l JOIN files f ON f.id = l.file_id WHERE f.kind = 'org'
351                    UNION ALL
352                    SELECT f.path, f.hash, h.ordinal, h.title, h.start_offset, l.target
353                    FROM links l JOIN headings h ON h.id = l.heading_id JOIN files f ON f.id = h.file_id
354                    WHERE f.kind = 'org' ORDER BY 1, 3
355                    """
356            ).map { row in
357                let location = location(row)
358                return (source: location.ordinal == HeadingLocation.fileOrdinal ? .file(location.path, contentHash: location.contentHash) : location,
359                        target: row["target"] as String)
360            }
361        }
362        let live = overlay.values.sorted { $0.path < $1.path }.flatMap { record in
363            record.links.map { (source: HeadingLocation.file(record.path, contentHash: record.hash), target: $0.target) }
364                + record.headings.flatMap { heading in heading.links.map { (source: location(record, heading), target: $0.target) } }
365        }
366        return indexed.filter { overlay[$0.source.path] == nil } + live
367    }
368
369    /// IDs used by more than one heading.
370    public func duplicateIDs(overlay: [String: FileRecord] = [:]) throws -> [String: [HeadingLocation]] {
371        let indexed = try database.read { db in
372            try Row.fetchAll(
373                db,
374                sql: """
375                    SELECT h.org_id, f.path, f.hash, h.ordinal, h.title, h.start_offset
376                    FROM headings h JOIN files f ON f.id = h.file_id
377                    WHERE h.org_id IN (SELECT org_id FROM headings WHERE org_id IS NOT NULL GROUP BY org_id HAVING count(*) > 1)
378                       OR (h.org_id IS NOT NULL AND ? > 0)
379                    ORDER BY f.path, h.ordinal
380                    """,
381                arguments: [overlay.count]
382            ).map { (id: $0["org_id"] as String, location: location($0)) }
383        }
384        var byID: [String: [HeadingLocation]] = [:]
385        for row in indexed where overlay[row.location.path] == nil {
386            byID[row.id, default: []].append(row.location)
387        }
388        for record in overlay.values.sorted(by: { $0.path < $1.path }) {
389            for heading in record.headings {
390                if let id = heading.orgID { byID[id, default: []].append(location(record, heading)) }
391            }
392        }
393        return byID.filter { $0.value.count > 1 }
394    }
395
396    /// Lower-cased runs of letters and digits, as the unicode61 tokenizer splits them.
397    static func words(_ text: String) -> [String] {
398        text.lowercased().split { !$0.isLetter && !$0.isNumber }.map(String.init)
399    }
400
401    private func location(_ row: Row) -> HeadingLocation {
402        HeadingLocation(path: row["path"], ordinal: row["ordinal"], title: row["title"], start: row["start_offset"], contentHash: row["hash"])
403    }
404
405    private func location(_ record: FileRecord, _ heading: HeadingRecord) -> HeadingLocation {
406        HeadingLocation(path: record.path, ordinal: heading.ordinal, title: heading.title, start: heading.start, contentHash: record.hash)
407    }
408}