pub(crate) mod backend;
pub(crate) mod batch_topology;
mod boundary;
pub(crate) use boundary::derive_pattern_boundary_context;
#[cfg(feature = "gpu")]
pub(crate) use boundary::regex_match_byte_upper_bound;
#[cfg(test)]
pub(crate) use boundary::scan_chunk_boundaries as scan_chunk_boundaries_for_test;
#[cfg(test)]
pub(crate) use boundary::MAX_BOUNDARY_SEAM_BYTES;
mod csr;
pub(crate) use csr::CsrU32;
mod extract;
pub(crate) use crate::gpu_matcher_cache as gpu_cache;
#[cfg(all(test, feature = "gpu"))]
pub(crate) use gpu_cache::gpu_matcher_cache_dir_from_base;
mod gpu_forced;
#[cfg(any(feature = "gpu", test))]
mod gpu_forced_helpers;
mod gpu_lazy;
mod gpu_lazy_helpers;
mod gpu_literal_scratch;
#[cfg(feature = "gpu")]
pub(crate) mod gpu_region_batch;
#[cfg(feature = "gpu")]
mod gpu_region_dispatch;
#[cfg(feature = "gpu")]
mod gpu_region_dispatch_helpers;
#[cfg(feature = "gpu")]
pub(crate) use crate::gpu::GpuResidentLiteralSlot;
mod gpu_stack;
mod hot_patterns;
pub(crate) mod phase2;
pub(crate) mod phase2_anchor;
#[cfg(test)]
pub(crate) use phase2_anchor::required_prefix_literals as phase2_required_prefix_literals_for_test;
pub(crate) use phase2_anchor::Phase2AnchorIndex;
pub(crate) use phase2_anchor::{
required_prefix_literals_with_cap, CONFIRMED_MAX_LITERALS_PER_PATTERN,
};
pub(crate) mod phase1_admission;
mod phase2_anchor_scan;
mod phase2_compiled;
mod phase2_compiled_anchored;
pub(crate) mod phase2_entropy;
#[path = "phase2/first_bigram.rs"]
mod phase2_first_bigram;
pub(crate) mod phase2_generic;
#[cfg(feature = "gpu")]
mod phase2_gpu_dfa;
#[cfg(feature = "gpu")]
pub(crate) use phase2_gpu_dfa::{compile_phase2_gpu_catalog_artifact, Phase2GpuDfaCatalogCache};
#[cfg(feature = "simd")]
mod phase2_hs;
#[cfg(feature = "gpu")]
pub(crate) use crate::gpu_input_budget;
#[cfg(feature = "simd")]
pub(crate) use phase2_hs::compile_phase2_scope_program;
#[cfg(all(test, feature = "simd"))]
pub(crate) use phase2_hs::hs_prefilter_requires_host_regex as hs_prefilter_requires_host_regex_for_test;
#[cfg(all(test, feature = "simd"))]
pub(crate) use phase2_hs::Phase2HsEngine;
mod phase2_prefilter;
pub(crate) use crate::phase2_truncate;
#[cfg_attr(not(feature = "simd"), allow(unused_imports))]
pub(crate) use phase2_prefilter::canonical_phase2_scope_indices;
mod process;
pub(crate) use crate::scan_profile as profile;
mod recovery;
pub use recovery::{BackendRecoveryReceipt, CoalescedScanOutcome, RecoveredInputRange};
mod scan;
mod vocab_absence;
pub(crate) use scan::{vocab_path_class, vocab_previously_clean};
mod scan_coalesced;
#[cfg(feature = "simd")]
pub(crate) use scan_coalesced::ReusableSimdTriggerCache;
pub(crate) mod scan_filters;
pub(crate) mod scan_postprocess;
pub(crate) use scan_postprocess::{
build_confirmed_suffix_gate_with_hints, confirmed_anchor::ConfirmedAnchorIndex,
};
#[path = "scan_postprocess/confirmed_extract.rs"]
mod scan_postprocess_confirmed_extract;
pub(crate) use scan_postprocess_confirmed_extract::{
exercise_confirmed_offsets_scratch_for_test, HOT_DIRECT_OFFSETS_CEILING,
};
#[path = "scan_postprocess/fragments.rs"]
mod scan_postprocess_fragments;
#[cfg(feature = "ml")]
#[path = "scan_postprocess/ml.rs"]
mod scan_postprocess_ml;
#[cfg(all(test, feature = "ml"))]
pub(crate) use scan_postprocess_ml::finalize_pending_match_for_test;
#[path = "scan_postprocess/companion_gate.rs"]
mod scan_postprocess_companion_gate;
#[path = "scan_postprocess/profile.rs"]
mod scan_postprocess_profile;
#[path = "scan_postprocess/suffix_gate.rs"]
mod scan_postprocess_suffix_gate;
pub(crate) mod trigger_bitmap;
mod windowed;
mod windowed_support;
pub(crate) use backend::PreparedChunk;
#[cfg(feature = "simd")]
pub(crate) use backend::{
build_packed_simd_compile_plan, build_simd_compile_plan, SimdPhase1CompilePlan,
SimdPhase1Prefilter,
};
#[cfg(test)]
pub(crate) use boundary::scan_chunk_boundaries;
#[cfg(test)]
pub(crate) use gpu_forced_helpers::gpu_forced_unavailable_message;
#[cfg(test)]
pub(crate) use phase2::{phase2_gate_stats_dump, take_mark_stats};
pub(crate) use scan_postprocess_companion_gate::{
companion_arms, companions_allow, companions_deny_absent,
};
pub(crate) use scan_postprocess_suffix_gate::suffix_gate_literals;
pub(crate) use windowed::{reject_oversized_window_chunk, MAX_WINDOW_CHUNK_BYTES};
pub(crate) use windowed_support::{absolute_line, absolute_offset, ceil_char_boundary};
pub use windowed_support::{
floor_char_boundary, line_number_for_offset, next_window_offset, record_window_match,
window_chunk, window_end_offset, window_ranges,
};
use crate::compiled_scanner::{GpuBackendAcquisitionFailure, GpuBackendPeers, SelectedGpuPeer};
use crate::types::*;
use aho_corasick::AhoCorasick;
use keyhog_core::{Chunk, RawMatch};
use std::sync::Arc;
use std::sync::OnceLock;
pub(crate) const MAX_INNER_LOOP_ITERS: usize = 1_000_000;
pub(crate) const BIGRAM_BLOOM_MIN_CHUNK_BYTES: usize = 64;
pub(crate) const MAX_RETAINED_WORKER_SCRATCH_BYTES: usize = crate::types::MAX_SCAN_CHUNK_BYTES;
#[derive(Default)]
pub(crate) struct CandidateScratch {
pub(crate) candidates: Vec<(u32, u32)>,
pub(crate) active_eligible: Vec<usize>,
pub(crate) literal_ids: Vec<u32>,
}
pub(crate) fn release_candidate_scratch(values: &mut CandidateScratch) {
values.candidates.clear();
values.active_eligible.clear();
values.literal_ids.clear();
let retained_bytes = values
.candidates
.capacity()
.saturating_mul(std::mem::size_of::<(u32, u32)>())
.saturating_add(
values
.active_eligible
.capacity()
.saturating_mul(std::mem::size_of::<usize>()),
)
.saturating_add(
values
.literal_ids
.capacity()
.saturating_mul(std::mem::size_of::<u32>()),
);
if retained_bytes > MAX_RETAINED_WORKER_SCRATCH_BYTES {
*values = CandidateScratch::default();
}
}
const MAX_IDLE_CANDIDATE_SCRATCH_BUFFERS: usize = 4;
static CANDIDATE_SCRATCH_POOL: std::sync::Mutex<Vec<CandidateScratch>> =
std::sync::Mutex::new(Vec::new());
fn release_idle_candidate_scratch() {
CANDIDATE_SCRATCH_POOL
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clear();
}
pub(crate) fn with_candidate_scratch<R>(f: impl FnOnce(&mut CandidateScratch) -> R) -> R {
let mut values = CANDIDATE_SCRATCH_POOL
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.pop()
.unwrap_or_default();
let result = f(&mut values);
release_candidate_scratch(&mut values);
if values.candidates.capacity() != 0
|| values.active_eligible.capacity() != 0
|| values.literal_ids.capacity() != 0
{
let mut pool = CANDIDATE_SCRATCH_POOL
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if pool.len() < MAX_IDLE_CANDIDATE_SCRATCH_BUFFERS {
pool.push(values);
}
}
result
}
#[cfg(test)]
pub(crate) fn candidate_scratch_idle_count_for_test() -> usize {
CANDIDATE_SCRATCH_POOL
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.len()
}
#[cfg(test)]
#[path = "../../tests/unit/scratch_pool_retention.rs"]
mod scratch_pool_retention;
#[cfg(test)]
#[path = "../../tests/unit/worker_scratch_bounds.rs"]
mod worker_scratch_bounds;
pub(crate) use phase1_admission::{Phase1Admission, Phase1AdmissionPlanIdentityError};
pub use phase1_admission::{
Phase1AdmissionPlan, Phase1AdmissionSummary, Phase2KeywordTriggerSummary,
};
pub(crate) enum ScannerBackendState {
Census {
peers: GpuBackendPeers,
failures: Vec<GpuBackendAcquisitionFailure>,
#[cfg(feature = "gpu")]
resident_literal_cuda: std::sync::Mutex<GpuResidentLiteralSlot>,
#[cfg(feature = "gpu")]
resident_literal_metal: std::sync::Mutex<GpuResidentLiteralSlot>,
#[cfg(feature = "gpu")]
resident_literal_wgpu: std::sync::Mutex<GpuResidentLiteralSlot>,
},
SelectedHost(crate::hw_probe::ScanBackend),
SelectedGpu {
peer: SelectedGpuPeer,
#[cfg(feature = "gpu")]
resident_literal: std::sync::Mutex<GpuResidentLiteralSlot>,
},
Disabled,
}
impl ScannerBackendState {
pub(crate) fn selected_backend(&self) -> Option<crate::hw_probe::ScanBackend> {
match self {
Self::SelectedHost(backend) => Some(*backend),
Self::SelectedGpu { peer, .. } => Some(peer.backend()),
Self::Census { .. } | Self::Disabled => None,
}
}
pub(crate) fn gpu_backend(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<&Arc<dyn vyre::VyreBackend>> {
match self {
Self::Census { peers, .. } => peers.get(backend),
Self::SelectedGpu { peer, .. } => peer.get(backend),
Self::SelectedHost(_) | Self::Disabled => None,
}
}
pub(crate) fn gpu_backend_available(&self, backend: crate::hw_probe::ScanBackend) -> bool {
match self {
Self::Census { peers, .. } => match backend {
crate::hw_probe::ScanBackend::GpuCuda => peers.cuda_available,
crate::hw_probe::ScanBackend::GpuMetal => peers.metal_available,
crate::hw_probe::ScanBackend::GpuWgpu => peers.wgpu_available,
_ => false,
},
Self::SelectedGpu { peer, .. } => peer.backend() == backend && peer.available,
Self::SelectedHost(_) | Self::Disabled => false,
}
}
pub(crate) fn gpu_backend_acquired(&self, backend: crate::hw_probe::ScanBackend) -> bool {
match self {
Self::Census { peers, .. } => peers.initialized(backend).is_some(),
Self::SelectedGpu { peer, .. } => peer.initialized(backend).is_some(),
Self::SelectedHost(_) | Self::Disabled => false,
}
}
pub(crate) fn gpu_backend_device_identity(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<String> {
match self {
Self::Census { peers, .. } => peers
.initialized(backend)
.and_then(|peer| peer.device_identity.clone())
.or_else(|| match backend {
crate::hw_probe::ScanBackend::GpuCuda => peers.cuda_device_identity.clone(),
crate::hw_probe::ScanBackend::GpuMetal => peers.metal_device_identity.clone(),
crate::hw_probe::ScanBackend::GpuWgpu => peers.wgpu_device_identity.clone(),
_ => None,
}),
Self::SelectedGpu { peer, .. } if peer.backend() == backend => peer
.initialized(backend)
.and_then(|acquired| acquired.device_identity.clone())
.or_else(|| peer.device_identity.clone()),
Self::SelectedGpu { .. } | Self::SelectedHost(_) | Self::Disabled => None,
}
}
pub(crate) fn gpu_backend_runtime_identity(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<String> {
match self {
Self::Census { peers, .. } => match backend {
crate::hw_probe::ScanBackend::GpuCuda => peers.cuda_runtime_identity.clone(),
crate::hw_probe::ScanBackend::GpuMetal => peers.metal_runtime_identity.clone(),
crate::hw_probe::ScanBackend::GpuWgpu => peers.wgpu_runtime_identity.clone(),
_ => None,
},
Self::SelectedGpu { peer, .. } if peer.backend() == backend => {
peer.runtime_identity.clone()
}
Self::SelectedGpu { .. } | Self::SelectedHost(_) | Self::Disabled => None,
}
}
pub(crate) fn gpu_backend_is_software(&self, backend: crate::hw_probe::ScanBackend) -> bool {
match self {
Self::Census { peers, .. } => peers.initialized(backend).map_or_else(
|| backend == crate::hw_probe::ScanBackend::GpuWgpu && peers.wgpu_is_software,
|peer| peer.is_software,
),
Self::SelectedGpu { peer, .. } if peer.backend() == backend => peer
.initialized(backend)
.map_or(peer.is_software, |acquired| acquired.is_software),
Self::SelectedGpu { .. } | Self::SelectedHost(_) | Self::Disabled => true,
}
}
pub(crate) fn gpu_backend_initialization_error(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<&str> {
match self {
Self::Census {
peers, failures, ..
} => peers.initialization_error(backend).or_else(|| {
failures
.iter()
.find(|failure| {
failure.backend
== match backend {
crate::hw_probe::ScanBackend::GpuCuda => "cuda",
crate::hw_probe::ScanBackend::GpuMetal => "metal",
crate::hw_probe::ScanBackend::GpuWgpu => "wgpu",
_ => return false,
}
})
.map(|failure| failure.diagnostic.as_str())
}),
Self::SelectedGpu { peer, .. } => peer.initialization_error(backend),
Self::SelectedHost(_) | Self::Disabled => None,
}
}
pub(crate) fn gpu_availability(&self) -> crate::gpu::GpuBackendAvailability {
match self {
Self::Census { peers, .. } => peers.availability(),
Self::SelectedGpu { peer, .. } => {
let available = peer.available;
crate::gpu::GpuBackendAvailability {
cuda: peer.backend() == crate::hw_probe::ScanBackend::GpuCuda && available,
metal: peer.backend() == crate::hw_probe::ScanBackend::GpuMetal && available,
wgpu: peer.backend() == crate::hw_probe::ScanBackend::GpuWgpu && available,
}
}
Self::SelectedHost(_) | Self::Disabled => crate::gpu::GpuBackendAvailability::default(),
}
}
pub(crate) fn gpu_candidate_backends(
&self,
) -> impl Iterator<Item = crate::hw_probe::ScanBackend> {
let backends = match self {
Self::Census { .. } => [
Some(crate::hw_probe::ScanBackend::GpuCuda),
Some(crate::hw_probe::ScanBackend::GpuMetal),
Some(crate::hw_probe::ScanBackend::GpuWgpu),
],
Self::SelectedGpu { peer, .. } => [Some(peer.backend()), None, None],
Self::SelectedHost(_) | Self::Disabled => [None, None, None],
};
backends.into_iter().flatten()
}
#[cfg(feature = "gpu")]
pub(crate) fn gpu_backend_adapter_identity(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<(u32, u32, bool, Option<&str>)> {
let peer = match self {
Self::Census { peers, .. } => peers.initialized(backend),
Self::SelectedGpu { peer, .. } => peer.initialized(backend),
Self::SelectedHost(_) | Self::Disabled => None,
}?;
Some((
peer.adapter_vendor,
peer.adapter_device,
peer.is_software,
peer.device_identity.as_deref(),
))
}
#[cfg(feature = "gpu")]
pub(crate) fn gpu_resident_timed_dispatch_supported(
&self,
backend: crate::hw_probe::ScanBackend,
) -> bool {
match self {
Self::Census { peers, .. } => peers.resident_timed_dispatch_supported(backend),
Self::SelectedGpu { peer, .. } => peer.resident_timed_dispatch_supported(backend),
Self::SelectedHost(_) | Self::Disabled => false,
}
}
}
impl CompiledScanner {
pub(crate) fn selected_backend(&self) -> Option<crate::hw_probe::ScanBackend> {
self.backend_state.selected_backend()
}
pub(crate) fn gpu_backend(
&self,
backend: crate::hw_probe::ScanBackend,
) -> Option<&Arc<dyn vyre::VyreBackend>> {
self.backend_state.gpu_backend(backend)
}
pub fn finish_partition(&self) {
self.fragment_cache.clear();
self.reusable_phase1_evidence.lock().clear();
release_idle_candidate_scratch();
}
}
pub struct CompiledScanner {
pub(crate) detector_digest: u64,
pub(crate) vocab_stage_absence_cache: dashmap::DashMap<
crate::engine::scan::VocabAbsenceKey,
crate::engine::scan::VocabStageAbsence,
ahash::RandomState,
>,
pub(crate) entropy_config_digest_cache: parking_lot::Mutex<Option<[u8; 32]>>,
pub(crate) compiled_plan_digest: [u8; 32],
pub(crate) fragment_cache: crate::fragment_cache::FragmentCache,
pub(crate) reusable_phase1_evidence:
parking_lot::Mutex<phase1_admission::ReusablePhase1EvidenceCache>,
pub(crate) ac: Option<AhoCorasick>,
pub(crate) backend_state: ScannerBackendState,
#[cfg(feature = "gpu")]
pub(crate) direct_gpu_resident_dispatch: std::sync::Mutex<()>,
pub(crate) quantized_confidence_authenticated: bool,
pub(crate) gpu_literals: Option<Arc<Vec<Vec<u8>>>>,
#[cfg(feature = "gpu")]
pub(crate) gpu_max_literal_len: usize,
pub(crate) gpu_matcher: OnceLock<Option<vyre::scan::GpuLiteralSet>>,
pub(crate) gpu_last_degrade_reason: std::sync::Mutex<Option<String>>,
pub(crate) gpu_degrade_count: std::sync::atomic::AtomicU64,
pub(crate) autoroute_gpu_shared_cold_ns: std::sync::atomic::AtomicU64,
pub(crate) static_intern: Arc<crate::static_intern::StaticInterner>,
pub(crate) detector_plans: crate::detector_plan::CompiledDetectorPlans,
pub(crate) assignment_keyword_matcher:
std::sync::Mutex<crate::assignment_keyword_matcher::AssignmentKeywordMatcherCache>,
#[cfg(feature = "gpu")]
pub(crate) ac_match_upper_bounds: Option<Vec<Option<usize>>>,
pub(crate) ac_map: Vec<CompiledPattern>,
pub(crate) structural_confirmed_patterns: CsrU32,
pub(crate) pattern_boundary_context: boundary::BoundaryContextBytes,
pub(crate) suffix_gate_ac: Option<scan_postprocess_suffix_gate::LazyConfirmedSuffixGate>,
pub(crate) ac_suffix_gate: CsrU32,
pub(crate) hot_confirmed_by_pattern: Vec<bool>,
pub(crate) confirmed_anchor_index:
Option<scan_postprocess::confirmed_anchor::ConfirmedAnchorIndex>,
pub(crate) prefix_propagation: CsrU32,
pub(crate) phase2_patterns: Vec<(CompiledPattern, Vec<String>)>,
pub(crate) structural_phase2_patterns: CsrU32,
pub(crate) same_prefix_patterns: CsrU32,
pub(crate) phase2_keyword_to_patterns: CsrU32,
pub(crate) phase2_keyword_count: usize,
pub(crate) phase2_always_anchor_literal_count: usize,
#[cfg(feature = "gpu")]
pub(crate) confirmed_anchor_literal_count: usize,
#[cfg(feature = "gpu")]
pub(crate) generic_keyword_literal_count: usize,
pub(crate) phase2_always_active_indices: Vec<usize>,
pub(crate) phase2_always_active_prefilter: Option<phase2::Phase2AlwaysActivePrefilter>,
pub(crate) phase2_anchor_index: Option<phase2_anchor::Phase2AnchorIndex>,
#[cfg(feature = "gpu")]
pub(crate) phase2_gpu_dfa: phase2_gpu_dfa::Phase2GpuDfaCatalogCache,
pub(crate) tuning: phase2::ScannerTuning,
#[cfg(feature = "simd")]
pub(crate) simd_candidate_available: bool,
#[cfg(feature = "simd")]
pub(crate) simd_compile_plan: std::sync::Mutex<Option<SimdPhase1CompilePlan>>,
#[cfg(feature = "simd")]
pub(crate) simd_prefilter: OnceLock<std::result::Result<SimdPhase1Prefilter, String>>,
#[cfg(feature = "simd")]
pub(crate) simd_initialization_ns: std::sync::atomic::AtomicU64,
#[cfg(feature = "simdsieve")]
pub(crate) hot_pattern_slots: Vec<crate::simdsieve_prefilter::HotPatternSlot>,
pub config: ScannerConfig,
pub(crate) route_classification: Arc<phase1_admission::RouteClassificationPlan>,
#[cfg(debug_assertions)]
pub(crate) phase2_keyword_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) generic_keyword_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) phase2_prefilter_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) phase1_trigger_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) normalization_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) confirmed_pattern_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) entropy_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) multiline_admission_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) line_index_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) decoder_admission_scanned_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) direct_scan_absence_skipped_bytes: std::sync::atomic::AtomicU64,
#[cfg(debug_assertions)]
pub(crate) direct_scan_absence_batches: std::sync::atomic::AtomicU64,
#[cfg(feature = "simd")]
pub(crate) reusable_simd_triggers: parking_lot::Mutex<scan_coalesced::ReusableSimdTriggerCache>,
#[cfg(debug_assertions)]
pub(crate) simd_phase2_tail_absence_skipped_bytes: std::sync::atomic::AtomicU64,
}
impl CompiledScanner {
pub fn detector_signature_sources(&self) -> std::collections::HashSet<Arc<str>> {
self.ac_map
.iter()
.chain(self.phase2_patterns.iter().map(|(pattern, _)| pattern))
.filter(|pattern| !pattern.homoglyph_variant)
.map(|pattern| pattern.regex.cloned_source())
.chain(self.detector_plans.companion_signature_sources())
.collect()
}
pub fn declared_detector_min_confidence(&self) -> impl Iterator<Item = (&str, f64)> + '_ {
self.detector_plans.declared_min_confidence()
}
}
const _: () = {
const fn assert_send_sync<T: Send + Sync>() {}
let _ = assert_send_sync::<CompiledScanner>; };