use crate::disk_cache::{
DiskDiagnosticsCacheSlotV0, disk_diagnostics_cache_slot_for_resolve,
is_disk_diagnostics_cache_kill_switch_value,
};
use crate::style_diagnostics::finish_style_diagnostics_value_with_shared_reachability;
use crate::style_diagnostics_snapshot::{
attach_workspace_snapshot_id_to_diagnostics, current_style_workspace_snapshot_id,
};
use crate::{
LspQueryReadView, LspShellState, LspStyleDiagnosticsRenderInputsV0,
protocol::is_style_document_uri, query_adapter::query_style_hover_candidate_from_lsp,
resolution_inputs_for_workspace_uri, source_documents_from_open_documents,
state::LspResolverIdentityIndexMemo, style_hover_candidates_for_document,
style_sources_from_open_documents,
};
use omena_query::{
OmenaQuerySourceDocumentInputV0, OmenaQueryStyleHoverCandidateV0,
OmenaQueryStyleMemoDatabaseV0, OmenaQueryStyleMemoHostV0, OmenaQueryStyleResolutionInputsV0,
OmenaQueryStyleSourceInputV0, OmenaResolverStyleModuleConfirmationIdentityIndexV0,
build_omena_resolver_style_module_confirmation_identity_index,
prepare_committed_workspace_wave_substrate,
resolve_committed_workspace_style_diagnostics_from_view_with_identity_index_and_wave_substrate,
};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
pub(crate) const PARALLEL_STYLE_WAVE_MIN_PARALLEL_TARGETS: usize = 2;
const PARALLEL_STYLE_WAVE_MAX_THREADS: usize = 4;
pub(crate) const PARALLEL_STYLE_DIAGNOSTICS_ENV_KILL_SWITCH: &str =
"OMENA_LSP_PARALLEL_DIAGNOSTICS";
fn parallel_style_diagnostics_kill_switch_engaged() -> bool {
std::env::var(PARALLEL_STYLE_DIAGNOSTICS_ENV_KILL_SWITCH)
.is_ok_and(|value| is_disk_diagnostics_cache_kill_switch_value(value.as_str()))
}
pub(crate) struct ResolvedParallelStyleTargetV0 {
pub(crate) diagnostics: Value,
pub(crate) disk_cache_slot: Option<DiskDiagnosticsCacheSlotV0>,
}
#[derive(PartialEq)]
struct ParallelStyleWaveSurfaceV0 {
style_sources: Vec<OmenaQueryStyleSourceInputV0>,
source_documents: Vec<OmenaQuerySourceDocumentInputV0>,
resolution_inputs: OmenaQueryStyleResolutionInputsV0,
}
struct ParallelStyleWaveTargetPlanV0 {
index: usize,
target_uri: String,
document_text: String,
query_candidates: Vec<OmenaQueryStyleHoverCandidateV0>,
disk_cache_slot: Option<DiskDiagnosticsCacheSlotV0>,
oracle_cached_diagnostics: Option<Value>,
}
struct ParallelStyleWaveCachedTargetPlanV0 {
index: usize,
diagnostics: Value,
snapshot_id: Option<omena_query::OmenaWorkspaceSnapshotIdV0>,
}
fn resolver_identity_index_for_parallel_style_wave(
state: &dyn LspQueryReadView,
surface: &ParallelStyleWaveSurfaceV0,
) -> Arc<OmenaResolverStyleModuleConfirmationIdentityIndexV0> {
let available_style_paths = surface
.style_sources
.iter()
.map(|source| source.style_path.clone())
.collect::<Vec<_>>();
let disk_style_path_identities = surface.resolution_inputs.disk_style_path_identities.clone();
{
let memo = state.resolver_identity_index_memo_lock();
if let Some(memo) = memo.as_ref()
&& memo.available_style_paths == available_style_paths
&& memo.disk_style_path_identities == disk_style_path_identities
{
return Arc::clone(&memo.index);
}
}
let available_style_path_refs = available_style_paths
.iter()
.map(String::as_str)
.collect::<BTreeSet<_>>();
let index = Arc::new(
build_omena_resolver_style_module_confirmation_identity_index(
&available_style_path_refs,
disk_style_path_identities.as_slice(),
),
);
*state.resolver_identity_index_memo_lock() = Some(LspResolverIdentityIndexMemo {
available_style_paths,
disk_style_path_identities,
index: Arc::clone(&index),
});
index
}
fn snapshot_id_for_parallel_style_surface(
state: &dyn LspQueryReadView,
_surface: &ParallelStyleWaveSurfaceV0,
_package_manifests: &[omena_query::OmenaQueryStylePackageManifestV0],
_external_sifs: &[omena_query::OmenaQueryExternalSifInputV0],
) -> Option<omena_query::OmenaWorkspaceSnapshotIdV0> {
current_style_workspace_snapshot_id(state)
}
pub(crate) fn resolved_parallel_style_wave_targets(
state: &LspShellState,
document_uris: &[String],
min_parallel_targets: usize,
) -> BTreeMap<usize, ResolvedParallelStyleTargetV0> {
resolved_parallel_style_wave_targets_with_abort(
state,
document_uris,
min_parallel_targets,
None,
)
}
pub(crate) type ParallelStyleWaveItemSinkV0<'sink> =
&'sink (dyn Fn(usize, Value, Option<DiskDiagnosticsCacheSlotV0>) + Sync);
pub(crate) fn resolved_parallel_style_wave_targets_with_abort(
state: &LspShellState,
document_uris: &[String],
min_parallel_targets: usize,
abort: Option<(&std::sync::atomic::AtomicU64, u64)>,
) -> BTreeMap<usize, ResolvedParallelStyleTargetV0> {
resolved_parallel_style_wave_targets_with_abort_and_sink(
state,
document_uris,
min_parallel_targets,
abort,
None,
)
}
pub(crate) fn resolved_parallel_style_wave_targets_with_abort_and_sink(
state: &LspShellState,
document_uris: &[String],
min_parallel_targets: usize,
abort: Option<(&std::sync::atomic::AtomicU64, u64)>,
on_item: Option<ParallelStyleWaveItemSinkV0<'_>>,
) -> BTreeMap<usize, ResolvedParallelStyleTargetV0> {
let mut host_slot = state.style_memo_host.borrow_mut();
let host = host_slot.get_or_insert_with(OmenaQueryStyleMemoHostV0::new);
resolved_parallel_style_wave_targets_with_host(
state,
document_uris,
min_parallel_targets,
abort,
on_item,
host,
Some(state),
)
}
pub(crate) fn resolved_parallel_style_wave_targets_from_read_view_with_abort_and_sink(
state: &dyn LspQueryReadView,
document_uris: &[String],
min_parallel_targets: usize,
abort: Option<(&std::sync::atomic::AtomicU64, u64)>,
on_item: Option<ParallelStyleWaveItemSinkV0<'_>>,
) -> BTreeMap<usize, ResolvedParallelStyleTargetV0> {
let mut host = OmenaQueryStyleMemoHostV0::new();
resolved_parallel_style_wave_targets_with_host(
state,
document_uris,
min_parallel_targets,
abort,
on_item,
&mut host,
None,
)
}
fn resolved_parallel_style_wave_targets_with_host(
state: &dyn LspQueryReadView,
document_uris: &[String],
min_parallel_targets: usize,
abort: Option<(&std::sync::atomic::AtomicU64, u64)>,
on_item: Option<ParallelStyleWaveItemSinkV0<'_>>,
memo_host: &mut OmenaQueryStyleMemoHostV0,
reverse_dependency_state: Option<&LspShellState>,
) -> BTreeMap<usize, ResolvedParallelStyleTargetV0> {
let mut resolved = BTreeMap::new();
let min_parallel_targets = min_parallel_targets.max(PARALLEL_STYLE_WAVE_MIN_PARALLEL_TARGETS);
if parallel_style_diagnostics_kill_switch_engaged() {
return resolved;
}
let candidate_indices = document_uris
.iter()
.enumerate()
.filter(|(_, uri)| is_style_document_uri(uri.as_str()))
.filter(|(_, uri)| {
state
.document(uri.as_str())
.is_some_and(|document| document.style_summary.is_some())
})
.map(|(index, _)| index)
.collect::<Vec<_>>();
if candidate_indices.len() < min_parallel_targets {
return resolved;
}
let package_manifests = state.query_resolution().package_manifests.as_slice();
let external_sifs = state.query_resolution().external_sifs.as_slice();
let configured_severity = state.query_diagnostics().severity;
let deep_analysis = state.query_diagnostics().deep_analysis;
let oracle_engaged = crate::disk_cache::disk_diagnostics_cache_oracle_engaged();
let mut shared_surface: Option<Arc<ParallelStyleWaveSurfaceV0>> = None;
let mut wave_cache_plan: Option<crate::disk_cache::DiskDiagnosticsCacheWavePlanV0> = None;
let mut cached_hits: Vec<ParallelStyleWaveCachedTargetPlanV0> = Vec::new();
let mut group: Vec<ParallelStyleWaveTargetPlanV0> = Vec::new();
for index in candidate_indices {
let uri = document_uris[index].as_str();
let Some(document) = state.document(uri) else {
continue;
};
let Some((_, candidates)) = style_hover_candidates_for_document(document) else {
continue;
};
let surface = ParallelStyleWaveSurfaceV0 {
style_sources: style_sources_from_open_documents(
state,
document.workspace_folder_uri.as_deref(),
Some(document.uri.as_str()),
),
source_documents: source_documents_from_open_documents(
state,
document.workspace_folder_uri.as_deref(),
),
resolution_inputs: resolution_inputs_for_workspace_uri(
state,
document.workspace_folder_uri.as_deref(),
),
};
match shared_surface.as_ref() {
None => {
let surface = Arc::new(surface);
shared_surface = Some(Arc::clone(&surface));
wave_cache_plan = crate::disk_cache::disk_diagnostics_cache_wave_plan_v1(
&crate::disk_cache::DiskDiagnosticsCacheEnvironmentComponentsV1 {
style_sources: surface.style_sources.as_slice(),
source_documents: surface.source_documents.as_slice(),
package_manifests,
external_sifs,
resolution_inputs: &surface.resolution_inputs,
severity: configured_severity,
deep_analysis,
},
);
}
Some(shared) if **shared == surface => {}
Some(_) => continue,
}
let disk_cache_slot = wave_cache_plan.as_ref().and_then(|plan| {
disk_diagnostics_cache_slot_for_resolve(
state,
document.workspace_folder_uri.as_deref(),
document.uri.as_str(),
plan,
)
});
let mut oracle_cached_diagnostics = None;
if let Some(slot) = disk_cache_slot.as_ref()
&& let Some(cached_diagnostics) = slot.load()
{
let cached_snapshot_id = slot.load_workspace_snapshot_id();
if oracle_engaged {
oracle_cached_diagnostics = Some(cached_diagnostics);
} else {
cached_hits.push(ParallelStyleWaveCachedTargetPlanV0 {
index,
diagnostics: cached_diagnostics,
snapshot_id: cached_snapshot_id,
});
continue;
}
}
group.push(ParallelStyleWaveTargetPlanV0 {
index,
target_uri: document.uri.clone(),
document_text: document.text.clone(),
query_candidates: candidates
.iter()
.map(query_style_hover_candidate_from_lsp)
.collect(),
disk_cache_slot,
oracle_cached_diagnostics,
});
}
let Some(shared_surface) = shared_surface else {
return resolved;
};
if group.len() < min_parallel_targets {
if !cached_hits.is_empty() {
let snapshot_id = snapshot_id_for_parallel_style_surface(
state,
shared_surface.as_ref(),
package_manifests,
external_sifs,
);
for cached_hit in cached_hits {
let diagnostics = attach_workspace_snapshot_id_to_diagnostics(
cached_hit.diagnostics,
cached_hit.snapshot_id.or(snapshot_id),
);
if let Some(sink) = on_item {
sink(cached_hit.index, diagnostics.clone(), None);
}
resolved.insert(
cached_hit.index,
ResolvedParallelStyleTargetV0 {
diagnostics,
disk_cache_slot: None,
},
);
}
}
return resolved;
}
let sync = memo_host.sync_workspace_for_parallel_resolve(
shared_surface.style_sources.as_slice(),
shared_surface.source_documents.as_slice(),
package_manifests,
external_sifs,
&shared_surface.resolution_inputs,
);
let Some(sync) = sync else {
return resolved;
};
let snapshot_id = Some(omena_query::OmenaWorkspaceSnapshotIdV0::from_revision(
sync.revision,
));
if let Some(reverse_dependency_state) = reverse_dependency_state {
crate::diagnostics_scheduler::refresh_reverse_dependency_index_memo(
reverse_dependency_state,
sync.revision.value,
&sync.committed_graph.cross_file_summary,
state.query_tide_ledger().epoch(),
);
}
for cached_hit in cached_hits {
let diagnostics = attach_workspace_snapshot_id_to_diagnostics(
cached_hit.diagnostics,
cached_hit.snapshot_id.or(snapshot_id),
);
if let Some(sink) = on_item {
sink(cached_hit.index, diagnostics.clone(), None);
}
resolved.insert(
cached_hit.index,
ResolvedParallelStyleTargetV0 {
diagnostics,
disk_cache_slot: None,
},
);
}
let files_by_path = sync
.files
.iter()
.map(|(style_path, file)| (style_path.as_str(), *file))
.collect::<BTreeMap<_, _>>();
let pool_items = group
.into_iter()
.filter_map(|plan| {
files_by_path
.get(plan.target_uri.as_str())
.copied()
.map(|file| (plan, file))
})
.collect::<Vec<_>>();
if pool_items.len() < min_parallel_targets {
return resolved;
}
let thread_count = std::thread::available_parallelism()
.map(std::num::NonZero::get)
.unwrap_or(1)
.min(PARALLEL_STYLE_WAVE_MAX_THREADS)
.min(pool_items.len());
let Ok(pool) = rayon::ThreadPoolBuilder::new()
.num_threads(thread_count)
.build()
else {
return resolved;
};
let workspace = sync.workspace;
let committed_graph = std::sync::Arc::new(sync.committed_graph.clone());
let resolver_identity_index =
resolver_identity_index_for_parallel_style_wave(state, shared_surface.as_ref());
let wave_substrate = {
let db = OmenaQueryStyleMemoDatabaseV0::from_handle(sync.handle.clone());
std::sync::Arc::new(prepare_committed_workspace_wave_substrate(
&db,
workspace,
committed_graph.as_ref(),
Some(resolver_identity_index.as_ref()),
))
};
let shared_reachability = std::sync::Arc::new(
crate::streaming_ifds_diagnostics::shared_streaming_reachability_for_lsp(
&committed_graph.cross_file_summary,
),
);
let read_set_index = std::sync::Arc::new(omena_query::reverse_dependency_index_from_edges_v0(
committed_graph.cross_file_summary.edges.as_slice(),
));
let computed: Vec<Option<(Value, Option<DiskDiagnosticsCacheSlotV0>)>> = pool.install(|| {
pool_items
.par_iter()
.map_with(
(
sync.handle.clone(),
committed_graph.clone(),
resolver_identity_index.clone(),
wave_substrate.clone(),
shared_reachability.clone(),
read_set_index.clone(),
snapshot_id,
),
|worker_state, (plan, file)| {
if let Some((watch, generation)) = abort
&& watch.load(std::sync::atomic::Ordering::Relaxed) != generation
{
return None;
}
let handle = worker_state.0.clone();
let committed_graph = worker_state.1.clone();
let resolver_identity_index = worker_state.2.clone();
let wave_substrate = worker_state.3.clone();
let shared_reachability = worker_state.4.clone();
let read_set_index = worker_state.5.clone();
let snapshot_id = worker_state.6;
let value = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let db = OmenaQueryStyleMemoDatabaseV0::from_handle(handle);
let summary =
resolve_committed_workspace_style_diagnostics_from_view_with_identity_index_and_wave_substrate(
&db,
workspace,
*file,
wave_substrate.as_ref(),
resolver_identity_index.as_ref(),
);
finish_style_diagnostics_value_with_shared_reachability(
&LspStyleDiagnosticsRenderInputsV0 {
document_uri: plan.target_uri.as_str(),
document_text: plan.document_text.as_str(),
query_candidates: plan.query_candidates.as_slice(),
snapshot_id,
deep_analysis,
configured_severity,
},
summary,
Some(&committed_graph.cross_file_summary),
Some(shared_reachability.as_ref()),
)
}))
.ok();
let slot_with_read_set = value.as_ref().and_then(|_| {
plan.disk_cache_slot.clone().map(|mut slot| {
slot.set_read_set_paths(
omena_query::diagnostics_read_set_for_target_v0(
read_set_index.as_ref(),
plan.target_uri.as_str(),
),
);
slot
})
});
if let (Some(sink), Some(value)) = (on_item, value.as_ref()) {
sink(plan.index, value.clone(), slot_with_read_set.clone());
}
value.map(|value| (value, slot_with_read_set))
},
)
.collect()
});
drop(pool);
drop(sync);
for ((plan, _), computed_target) in pool_items.into_iter().zip(computed) {
if let Some((diagnostics, disk_cache_slot)) = computed_target {
if let Some(cached) = plan.oracle_cached_diagnostics.as_ref() {
let cached =
attach_workspace_snapshot_id_to_diagnostics(cached.clone(), snapshot_id);
crate::disk_cache::record_disk_diagnostics_cache_oracle_outcome(
plan.target_uri.as_str(),
cached == diagnostics,
);
}
resolved.insert(
plan.index,
ResolvedParallelStyleTargetV0 {
diagnostics,
disk_cache_slot,
},
);
}
}
resolved
}