|
| 1 | +//! Bulk AST node insertion via rusqlite. |
| 2 | +//! |
| 3 | +//! Bypasses the JS iteration loop by opening the SQLite database directly |
| 4 | +//! from Rust and inserting all AST nodes in a single transaction. |
| 5 | +//! Parent node IDs are resolved by querying the `nodes` table. |
| 6 | +
|
| 7 | +use std::collections::HashMap; |
| 8 | + |
| 9 | +use napi_derive::napi; |
| 10 | +use rusqlite::{params, Connection, OpenFlags}; |
| 11 | +use serde::{Deserialize, Serialize}; |
| 12 | + |
| 13 | +/// A single AST node to insert (received from JS). |
| 14 | +#[napi(object)] |
| 15 | +#[derive(Debug, Clone, Serialize, Deserialize)] |
| 16 | +pub struct AstInsertNode { |
| 17 | + pub line: u32, |
| 18 | + pub kind: String, |
| 19 | + pub name: String, |
| 20 | + pub text: Option<String>, |
| 21 | + pub receiver: Option<String>, |
| 22 | +} |
| 23 | + |
| 24 | +/// A batch of AST nodes for a single file. |
| 25 | +#[napi(object)] |
| 26 | +#[derive(Debug, Clone, Serialize, Deserialize)] |
| 27 | +pub struct FileAstBatch { |
| 28 | + pub file: String, |
| 29 | + pub nodes: Vec<AstInsertNode>, |
| 30 | +} |
| 31 | + |
| 32 | +/// A definition row from the `nodes` table used for parent resolution. |
| 33 | +struct NodeDef { |
| 34 | + id: i64, |
| 35 | + line: u32, |
| 36 | + end_line: Option<u32>, |
| 37 | +} |
| 38 | + |
| 39 | +/// Find the narrowest enclosing definition for a given source line. |
| 40 | +/// Returns the node ID of the best match, or None if no definition encloses this line. |
| 41 | +/// |
| 42 | +/// Mirrors the JS `findParentDef` semantics: a definition with `end_line = NULL` |
| 43 | +/// is treated as always enclosing, with a negative sentinel span so it is preferred |
| 44 | +/// over definitions that have an explicit (wider) `end_line`. |
| 45 | +fn find_parent_id(defs: &[NodeDef], line: u32) -> Option<i64> { |
| 46 | + let mut best_id: Option<i64> = None; |
| 47 | + let mut best_span: i64 = i64::MAX; |
| 48 | + for d in defs { |
| 49 | + if d.line <= line { |
| 50 | + let span: i64 = match d.end_line { |
| 51 | + Some(el) if el >= line => (el - d.line) as i64, |
| 52 | + Some(_) => continue, |
| 53 | + // JS: (def.endLine ?? 0) - def.line → negative, always preferred |
| 54 | + None => -(d.line as i64), |
| 55 | + }; |
| 56 | + if span < best_span { |
| 57 | + best_id = Some(d.id); |
| 58 | + best_span = span; |
| 59 | + } |
| 60 | + } |
| 61 | + } |
| 62 | + best_id |
| 63 | +} |
| 64 | + |
| 65 | +/// Bulk-insert AST nodes into the database, resolving `parent_node_id` |
| 66 | +/// from the `nodes` table. Runs all inserts in a single SQLite transaction. |
| 67 | +/// |
| 68 | +/// Returns the number of rows inserted. Returns 0 on any error (DB open |
| 69 | +/// failure, missing table, transaction failure). |
| 70 | +#[napi] |
| 71 | +pub fn bulk_insert_ast_nodes(db_path: String, batches: Vec<FileAstBatch>) -> u32 { |
| 72 | + if batches.is_empty() { |
| 73 | + return 0; |
| 74 | + } |
| 75 | + |
| 76 | + let flags = OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX; |
| 77 | + let mut conn = match Connection::open_with_flags(&db_path, flags) { |
| 78 | + Ok(c) => c, |
| 79 | + Err(_) => return 0, |
| 80 | + }; |
| 81 | + |
| 82 | + // Match the JS-side performance pragmas (including busy_timeout for WAL contention) |
| 83 | + let _ = conn.execute_batch( |
| 84 | + "PRAGMA synchronous = NORMAL; PRAGMA busy_timeout = 5000", |
| 85 | + ); |
| 86 | + |
| 87 | + // Bail out if the ast_nodes table doesn't exist (schema too old) |
| 88 | + let has_table: bool = conn |
| 89 | + .prepare("SELECT 1 FROM sqlite_master WHERE type='table' AND name='ast_nodes'") |
| 90 | + .and_then(|mut s| s.query_row([], |_| Ok(true))) |
| 91 | + .unwrap_or(false); |
| 92 | + if !has_table { |
| 93 | + return 0; |
| 94 | + } |
| 95 | + |
| 96 | + // ── Phase 1: Pre-fetch node definitions for parent resolution ──────── |
| 97 | + let mut file_defs: HashMap<String, Vec<NodeDef>> = HashMap::new(); |
| 98 | + { |
| 99 | + let Ok(mut stmt) = |
| 100 | + conn.prepare("SELECT id, line, end_line FROM nodes WHERE file = ?1") |
| 101 | + else { |
| 102 | + return 0; |
| 103 | + }; |
| 104 | + |
| 105 | + for batch in &batches { |
| 106 | + if batch.nodes.is_empty() || file_defs.contains_key(&batch.file) { |
| 107 | + continue; |
| 108 | + } |
| 109 | + let defs: Vec<NodeDef> = stmt |
| 110 | + .query_map(params![&batch.file], |row| { |
| 111 | + Ok(NodeDef { |
| 112 | + id: row.get(0)?, |
| 113 | + line: row.get(1)?, |
| 114 | + end_line: row.get(2)?, |
| 115 | + }) |
| 116 | + }) |
| 117 | + .map(|rows| rows.filter_map(|r| r.ok()).collect()) |
| 118 | + .unwrap_or_default(); |
| 119 | + file_defs.insert(batch.file.clone(), defs); |
| 120 | + } |
| 121 | + } // `stmt` dropped — releases the immutable borrow on `conn` |
| 122 | + |
| 123 | + // ── Phase 2: Bulk insert in a single transaction ───────────────────── |
| 124 | + let Ok(tx) = conn.transaction() else { |
| 125 | + return 0; |
| 126 | + }; |
| 127 | + |
| 128 | + let mut total = 0u32; |
| 129 | + { |
| 130 | + let Ok(mut insert_stmt) = tx.prepare( |
| 131 | + "INSERT INTO ast_nodes (file, line, kind, name, text, receiver, parent_node_id) \ |
| 132 | + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)", |
| 133 | + ) else { |
| 134 | + return 0; |
| 135 | + }; |
| 136 | + |
| 137 | + for batch in &batches { |
| 138 | + let empty = Vec::new(); |
| 139 | + let defs = file_defs.get(&batch.file).unwrap_or(&empty); |
| 140 | + |
| 141 | + for node in &batch.nodes { |
| 142 | + let parent_id = find_parent_id(defs, node.line); |
| 143 | + |
| 144 | + match insert_stmt.execute(params![ |
| 145 | + &batch.file, |
| 146 | + node.line, |
| 147 | + &node.kind, |
| 148 | + &node.name, |
| 149 | + &node.text, |
| 150 | + &node.receiver, |
| 151 | + parent_id, |
| 152 | + ]) { |
| 153 | + Ok(_) => total += 1, |
| 154 | + Err(_) => return 0, // abort; tx rolls back on drop |
| 155 | + } |
| 156 | + } |
| 157 | + } |
| 158 | + } // `insert_stmt` dropped |
| 159 | + |
| 160 | + if tx.commit().is_err() { |
| 161 | + return 0; |
| 162 | + } |
| 163 | + |
| 164 | + total |
| 165 | +} |
0 commit comments