use std::collections::BTreeMap;
use std::ops::ControlFlow;
use std::path::{Path, PathBuf};
use std::sync::{Condvar, Mutex};
use crate::ast::{api as ast_api, Language as AstLanguage, Symbol};
use super::file_table::fnv1a64;
use super::imports;
use super::trigram::{extract_trigrams, Trigram};
use super::walker::MAX_FILE_BYTES;
const DEFAULT_FACT_WORKERS: usize = 6;
const DEDICATED_BUILD_THREADS: usize = 2;
pub(super) fn default_workers() -> usize {
std::thread::available_parallelism()
.map_or(1, usize::from)
.saturating_sub(DEDICATED_BUILD_THREADS)
.clamp(1, DEFAULT_FACT_WORKERS)
}
const QUEUED_FACTS_PER_WORKER: usize = 8;
pub(super) struct FileJob {
pub abs: PathBuf,
pub rel: String,
pub language: harn_languages::Language,
pub size_bytes: u64,
pub mtime_ms: i64,
pub known_hash: Option<u64>,
}
pub(super) enum FileFacts {
Skipped,
Unchanged {
hash: u64,
},
Changed(Box<ParsedFile>),
}
pub(super) struct ParsedFile {
pub hash: u64,
pub language: String,
pub imports: Vec<String>,
pub line_count: u32,
pub trigrams: Vec<Trigram>,
pub words: Vec<(String, Vec<u32>)>,
pub graph: Option<GraphFacts>,
}
pub(crate) struct GraphFacts {
pub language: AstLanguage,
pub symbols: Vec<Symbol>,
pub call_sites: Vec<(String, u32)>,
}
impl GraphFacts {
pub(crate) fn extract(path: &str, language: AstLanguage, source: &str) -> Self {
let parse_phase = super::timing::Phase::start("graph.parse");
let parsed = ast_api::parse_with_symbols(source, language);
drop(parse_phase);
match parsed {
Ok((tree, symbols)) => {
let _phase = super::timing::Phase::start("facts.call_sites");
let call_sites = super::symbol_graph::extract_call_sites_from_tree(&tree, source);
Self {
language,
symbols,
call_sites,
}
}
Err(err) => {
tracing::debug!(
"code_index: tree-sitter parse failed for `{path}`: {err}; \
symbol graph slice will be Module-only"
);
Self {
language,
symbols: Vec::new(),
call_sites: Vec::new(),
}
}
}
}
}
pub(super) fn extract(job: &FileJob) -> FileFacts {
let read_phase = super::timing::Phase::start("facts.read");
if job.size_bytes > MAX_FILE_BYTES {
return FileFacts::Skipped;
}
let Ok(source) = std::fs::read_to_string(&job.abs) else {
return FileFacts::Skipped;
};
if source.len() > MAX_FILE_BYTES as usize {
return FileFacts::Skipped;
}
let hash = fnv1a64(source.as_bytes());
drop(read_phase);
if job.known_hash == Some(hash) {
return FileFacts::Unchanged { hash };
}
FileFacts::Changed(Box::new(parse_source(job, &source, hash)))
}
fn parse_source(job: &FileJob, source: &str, hash: u64) -> ParsedFile {
let language = job.language.name().to_string();
let imports = imports::extract_imports(source, &language);
let line_count = crate::text::count_lines(source.as_bytes()) as u32;
let trigrams = {
let _phase = super::timing::Phase::start("facts.trigrams");
extract_trigrams(source).into_iter().collect()
};
let words = {
let _phase = super::timing::Phase::start("facts.words");
super::words::occurrence_lines(source)
};
let graph = AstLanguage::detect(Path::new(&job.rel), None)
.map(|language| GraphFacts::extract(&job.rel, language, source));
ParsedFile {
hash,
language,
imports,
line_count,
trigrams,
words,
graph,
}
}
pub(super) fn publish_in_order(
jobs: &[FileJob],
workers: usize,
publish: impl FnMut(&FileJob, FileFacts) -> ControlFlow<()>,
) {
publish_extracted_in_order(jobs, workers, &extract, publish);
}
fn publish_extracted_in_order<E>(
jobs: &[FileJob],
workers: usize,
extract: &E,
mut publish: impl FnMut(&FileJob, FileFacts) -> ControlFlow<()>,
) where
E: Fn(&FileJob) -> FileFacts + Sync,
{
let workers = workers.min(jobs.len());
if workers <= 1 {
for job in jobs {
if publish(job, extract(job)).is_break() {
return;
}
}
return;
}
let window = Window::new(workers * QUEUED_FACTS_PER_WORKER);
super::parallel::scope(|scope| {
let _stop = StopOnDrop(&window);
for index in 0..workers {
let window = &window;
scope
.spawn_named(format!("harn-code-index-facts-{index}"), move || {
window.work(jobs, extract);
})
.expect("spawn a code-index fact worker");
}
for (index, job) in jobs.iter().enumerate() {
let facts = window.take(index);
if publish(job, facts).is_break() {
return;
}
}
});
}
struct Window {
state: Mutex<WindowState>,
changed: Condvar,
capacity: usize,
}
#[derive(Default)]
struct WindowState {
next_claim: usize,
taken: usize,
ready: BTreeMap<usize, FileFacts>,
stopped: bool,
lost: Option<usize>,
}
impl Window {
fn new(capacity: usize) -> Self {
Self {
state: Mutex::default(),
changed: Condvar::new(),
capacity,
}
}
fn work(&self, jobs: &[FileJob], extract: &impl Fn(&FileJob) -> FileFacts) {
let mut panic_flag = PanicFlag {
window: self,
claimed: None,
};
loop {
let index = {
let mut state = self.lock();
loop {
if state.stopped || state.next_claim >= jobs.len() {
return;
}
if state.next_claim < state.taken + self.capacity {
break;
}
state = self.wait(state);
}
state.next_claim += 1;
state.next_claim - 1
};
panic_flag.claimed = Some(index);
let facts = extract(&jobs[index]);
let mut state = self.lock();
state.ready.insert(index, facts);
panic_flag.claimed = None;
self.changed.notify_all();
}
}
fn take(&self, index: usize) -> FileFacts {
let _phase = super::timing::Phase::start("build.facts_wait");
let mut state = self.lock();
loop {
if let Some(facts) = state.ready.remove(&index) {
state.taken = index + 1;
self.changed.notify_all();
return facts;
}
assert!(
state.lost != Some(index) && !(state.stopped && index >= state.next_claim),
"a code-index fact worker panicked"
);
state = self.wait(state);
}
}
fn stop(&self) {
self.lock().stopped = true;
self.changed.notify_all();
}
fn lock(&self) -> std::sync::MutexGuard<'_, WindowState> {
self.state
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn wait<'a>(
&self,
state: std::sync::MutexGuard<'a, WindowState>,
) -> std::sync::MutexGuard<'a, WindowState> {
self.changed
.wait(state)
.unwrap_or_else(|poison| poison.into_inner())
}
}
struct StopOnDrop<'a>(&'a Window);
impl Drop for StopOnDrop<'_> {
fn drop(&mut self) {
self.0.stop();
}
}
struct PanicFlag<'a> {
window: &'a Window,
claimed: Option<usize>,
}
impl Drop for PanicFlag<'_> {
fn drop(&mut self) {
if std::thread::panicking() {
let mut state = self.window.lock();
state.lost = match (state.lost, self.claimed) {
(Some(lost), Some(claimed)) => Some(lost.min(claimed)),
(lost, claimed) => lost.or(claimed),
};
state.stopped = true;
self.window.changed.notify_all();
}
}
}
#[cfg(test)]
#[path = "file_facts_tests.rs"]
mod tests;