use super::engine::{TypeKind, UnsafeInferenceEngine};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InferencePolicy {
pub min_score: u16,
pub min_margin: u16,
}
impl InferencePolicy {
pub const PRELIMINARY: Self = Self {
min_score: 30,
min_margin: 10,
};
#[must_use]
pub const fn new(min_score: u16, min_margin: u16) -> Self {
Self {
min_score,
min_margin,
}
}
}
impl Default for InferencePolicy {
fn default() -> Self {
Self::PRELIMINARY
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InferenceDecision {
pub predicted: Option<TypeKind>,
pub best_candidate: TypeKind,
pub raw_score: u16,
pub runner_up_score: u16,
pub abstain_reason: Option<String>,
}
impl InferenceDecision {
#[must_use]
pub fn is_accepted(&self) -> bool {
self.predicted.is_some()
}
#[must_use]
pub fn accepted_score(&self) -> u16 {
if self.is_accepted() {
self.raw_score
} else {
0
}
}
}
pub struct CalibrationSample {
pub bytes: Vec<u8>,
pub size: usize,
pub expected: TypeKind,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct CalibrationReport {
pub min_score: u16,
pub min_margin: u16,
pub accepted: usize,
pub correct: usize,
pub abstained: usize,
pub precision: f64,
pub recall: f64,
}
impl CalibrationReport {
#[must_use]
pub fn f1(&self) -> f64 {
if self.precision + self.recall == 0.0 {
return 0.0;
}
2.0 * self.precision * self.recall / (self.precision + self.recall)
}
}
#[must_use]
pub fn evaluate(policy: InferencePolicy, samples: &[CalibrationSample]) -> CalibrationReport {
let mut report = CalibrationReport {
min_score: policy.min_score,
min_margin: policy.min_margin,
..CalibrationReport::default()
};
if samples.is_empty() {
return report;
}
for sample in samples {
let decision =
UnsafeInferenceEngine::decide_with_policy(&sample.bytes, sample.size, &policy);
match decision.predicted {
Some(kind) => {
report.accepted += 1;
if kind == sample.expected {
report.correct += 1;
}
}
None => report.abstained += 1,
}
}
report.precision = if report.accepted == 0 {
1.0
} else {
report.correct as f64 / report.accepted as f64
};
report.recall = report.correct as f64 / samples.len() as f64;
report
}
#[must_use]
pub fn calibrate(samples: &[CalibrationSample], min_precision: f64) -> CalibrationReport {
let mut best: Option<CalibrationReport> = None;
for min_score in (0..=80u16).step_by(10) {
for min_margin in (0..=30u16).step_by(5) {
let policy = InferencePolicy::new(min_score, min_margin);
let report = evaluate(policy, samples);
if report.precision < min_precision || report.accepted == 0 {
continue;
}
let better = match &best {
None => true,
Some(current) => match report.f1().partial_cmp(¤t.f1()) {
Some(std::cmp::Ordering::Greater) => true,
Some(std::cmp::Ordering::Equal) => report.min_margin > current.min_margin,
_ => false,
},
};
if better {
best = Some(report);
}
}
}
best.unwrap_or_else(|| evaluate(InferencePolicy::PRELIMINARY, samples))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_policy_is_preliminary() {
assert_eq!(
InferencePolicy::default(),
InferencePolicy::PRELIMINARY,
"the default policy must be the documented preliminary thresholds"
);
assert!(
InferencePolicy::default().min_score > 0,
"a zero minimum score would accept every guess"
);
assert!(
InferencePolicy::default().min_margin > 0,
"a zero margin would accept ties"
);
}
#[test]
fn abstained_decision_reports_zero_score() {
let decision = InferenceDecision {
predicted: None,
best_candidate: TypeKind::Vec,
raw_score: 25,
runner_up_score: 20,
abstain_reason: Some("below min_score".to_string()),
};
assert!(!decision.is_accepted(), "no type was accepted");
assert_eq!(
decision.accepted_score(),
0,
"an abstained decision must not report a confidence score"
);
}
#[test]
fn accepted_decision_keeps_raw_score() {
let decision = InferenceDecision {
predicted: Some(TypeKind::String),
best_candidate: TypeKind::String,
raw_score: 90,
runner_up_score: 15,
abstain_reason: None,
};
assert!(decision.is_accepted(), "the type must be accepted");
assert_eq!(decision.accepted_score(), 90, "the raw score must survive");
}
fn calibration_fixture() -> Vec<CalibrationSample> {
let mut samples = Vec::new();
let text = b"the quick brown fox jumps over the lazy dog".to_vec();
samples.push(CalibrationSample {
size: text.len(),
bytes: text,
expected: TypeKind::String,
});
let mut cstring = b"hello world".to_vec();
cstring.push(0);
samples.push(CalibrationSample {
size: cstring.len(),
bytes: cstring,
expected: TypeKind::CString,
});
let buffer: Vec<u8> = (0..=255u8).collect();
samples.push(CalibrationSample {
size: buffer.len(),
bytes: buffer,
expected: TypeKind::Buffer,
});
let control = vec![0xC2u8, 0x80, 0xC2, 0x81, 0xC2, 0x82];
samples.push(CalibrationSample {
size: control.len(),
bytes: control,
expected: TypeKind::Buffer,
});
samples
}
#[test]
fn preliminary_thresholds_are_backed_by_measurement() {
let samples = calibration_fixture();
let shipped = evaluate(InferencePolicy::PRELIMINARY, &samples);
let calibrated = calibrate(&samples, 0.85);
assert!(
shipped.accepted > 0,
"the shipped policy must not abstain on the whole fixture"
);
assert!(
shipped.precision >= 0.85,
"shipped thresholds must keep precision above 0.85, measured {} (accepted {}, abstained {})",
shipped.precision,
shipped.accepted,
shipped.abstained
);
assert!(
shipped.recall > 0.0,
"the shipped policy must recover at least one sample, got recall {}",
shipped.recall
);
assert!(
calibrated.precision >= shipped.precision - 0.2,
"calibration must not degrade precision below the shipped policy (calibrated {}, shipped {})",
calibrated.precision,
shipped.precision
);
assert!(
calibrated.f1() >= shipped.f1() - 1e-9,
"the calibrated policy must not be worse than the shipped one (calibrated {}, shipped {})",
calibrated.f1(),
shipped.f1()
);
assert!(
shipped.accepted >= 2 && shipped.accepted + shipped.abstained == samples.len(),
"measured acceptance on the calibration fixture changed: {shipped:?}"
);
assert_eq!(
shipped.correct, shipped.accepted,
"every accepted prediction must match its label: {shipped:?}"
);
assert!(
(shipped.precision - 1.0).abs() < 1e-9,
"measured precision must be 1.0, got {}",
shipped.precision
);
}
#[test]
fn evaluation_exposes_weak_thresholds() {
let samples = calibration_fixture();
let permissive = evaluate(InferencePolicy::new(0, 0), &samples);
assert_eq!(
permissive.abstained, 0,
"a zero-threshold policy never abstains"
);
assert_eq!(
permissive.accepted,
samples.len(),
"every sample must be accepted by the permissive policy"
);
assert!(
permissive.precision <= 1.0,
"precision must stay a fraction, got {}",
permissive.precision
);
}
}