use core::{
num::NonZeroU64,
sync::atomic::{AtomicU64, Ordering},
};
use smol_str::{SmolStr, format_smolstr};
use crate::{
core::{AlignmentUnit, OovDecision, OovDetection, OovEvent, OovResolution},
runner::alignment_pool::AlignWorkItem,
types::{AlignmentError, AlignmentFailure, ChunkId, WorkFailure},
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct JobId(NonZeroU64);
impl JobId {
pub(crate) fn next() -> Self {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let raw = COUNTER.fetch_add(1, Ordering::Relaxed);
Self(NonZeroU64::new(raw).expect("JobId counter overflowed u64"))
}
}
#[derive(Debug)]
#[must_use = "a detection does nothing until it is decided"]
pub struct JobDetection {
job: JobId,
chunk_id: ChunkId,
registry: NonZeroU64,
units: Vec<OovDetection>,
}
impl JobDetection {
pub(crate) fn new(job: &AlignWorkItem, registry: NonZeroU64, units: Vec<OovDetection>) -> Self {
Self {
job: job.id(),
chunk_id: job.chunk_id(),
registry,
units,
}
}
#[must_use]
pub const fn chunk_id(&self) -> ChunkId {
self.chunk_id
}
pub fn units(&self) -> &[OovDetection] {
&self.units
}
pub fn decide(self, mut policy: impl FnMut(&OovEvent) -> OovDecision) -> JobResolution {
JobResolution {
job: self.job,
chunk_id: self.chunk_id,
registry: self.registry,
units: self
.units
.into_iter()
.map(|unit| unit.decide(&mut policy))
.collect(),
}
}
}
#[derive(Debug)]
#[must_use = "a resolution does nothing until run_one_alignment applies it"]
pub struct JobResolution {
job: JobId,
chunk_id: ChunkId,
registry: NonZeroU64,
units: Vec<OovResolution>,
}
impl JobResolution {
#[must_use]
pub const fn chunk_id(&self) -> ChunkId {
self.chunk_id
}
pub fn units(&self) -> &[OovResolution] {
&self.units
}
pub(crate) fn into_units_for(
self,
job: &AlignWorkItem,
registry: NonZeroU64,
) -> Result<Vec<OovResolution>, WorkFailure> {
let refuse = |message: SmolStr| {
WorkFailure::Alignment(AlignmentError::Tokenization(AlignmentFailure::new(
message,
job.language().clone(),
)))
};
if self.job != job.id() || self.chunk_id != job.chunk_id() {
return Err(refuse(format_smolstr!(
"this JobResolution was detected for another job (chunk {}), not for this \
AlignWorkItem (chunk {}). A resolution applies only to the work item its detection \
read; detect this job with `AlignmentSet::detect_oov(&job)` and decide that.",
self.chunk_id,
job.chunk_id(),
)));
}
if self.registry != registry {
return Err(refuse(SmolStr::new_static(
"this JobResolution was detected through another AlignmentSet. A resolution \
applies only through the set that read it; detect this job on the set that aligns it.",
)));
}
let expected = job_units(job);
if !self
.units
.iter()
.map(OovResolution::unit)
.eq(expected.iter().copied())
{
return Err(refuse(format_smolstr!(
"this JobResolution has units {:?}, but the job's units are {expected:?}",
self
.units
.iter()
.map(OovResolution::unit)
.collect::<Vec<_>>(),
)));
}
Ok(self.units)
}
}
pub(crate) fn job_units(job: &AlignWorkItem) -> Vec<AlignmentUnit> {
if job.runs().is_empty() {
vec![AlignmentUnit::Whole]
} else {
(0..job.runs().len()).map(AlignmentUnit::Run).collect()
}
}