use std::fmt;
use std::time::Duration;
use async_trait::async_trait;
use boundgraph_core::{Condensate, CondensatePayload, Digest, SchemaId, SlicerId};
use serde::{Deserialize, Serialize};
#[async_trait]
pub trait Store: Send + Sync + fmt::Debug {
async fn put(&self, bytes: &[u8]) -> Result<Digest, StoreError>;
async fn get(&self, digest: &Digest) -> Result<Vec<u8>, StoreError>;
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum StoreError {
#[error("nothing is stored at {0}")]
NotFound(Digest),
#[error("the content store failed: {0}")]
Unavailable(String),
}
pub trait TokenCount: Send + Sync + fmt::Debug {
fn count(&self, text: &str) -> usize;
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Estimate;
impl TokenCount for Estimate {
fn count(&self, text: &str) -> usize {
if text.is_empty() {
return 0;
}
let bytes = text.len();
let lines = text.lines().count();
bytes.div_ceil(3) + lines
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ByteCeiling;
impl TokenCount for ByteCeiling {
fn count(&self, text: &str) -> usize {
text.len()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommandOutcome {
pub command: String,
pub exit: Option<i32>,
pub duration: Duration,
}
impl CommandOutcome {
pub fn new(command: impl Into<String>, exit: i32, duration: Duration) -> Self {
Self {
command: command.into(),
exit: Some(exit),
duration,
}
}
pub fn succeeded(&self) -> bool {
self.exit == Some(0)
}
pub fn summary(&self) -> String {
let status = match self.exit {
Some(0) => "ok".to_string(),
Some(code) => format!("exit {code}"),
None => "killed".to_string(),
};
format!(
"$ {} — {status} in {}ms",
self.command,
self.duration.as_millis()
)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Raw {
pub bytes: Vec<u8>,
pub command: Option<CommandOutcome>,
pub hint: Option<String>,
}
impl Raw {
pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
Self {
bytes: bytes.into(),
command: None,
hint: None,
}
}
#[must_use]
pub fn from_command(mut self, outcome: CommandOutcome) -> Self {
self.command = Some(outcome);
self
}
#[must_use]
pub fn hinting(mut self, hint: impl Into<String>) -> Self {
self.hint = Some(hint.into());
self
}
pub fn text(&self) -> std::borrow::Cow<'_, str> {
String::from_utf8_lossy(&self.bytes)
}
pub fn len(&self) -> usize {
self.bytes.len()
}
pub fn is_empty(&self) -> bool {
self.bytes.is_empty()
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct Slice {
pub schema: SchemaId,
pub payload: serde_json::Value,
pub lossy: bool,
}
impl Slice {
pub fn complete(schema: SchemaId, payload: serde_json::Value) -> Self {
Self {
schema,
payload,
lossy: false,
}
}
pub fn lossy(schema: SchemaId, payload: serde_json::Value) -> Self {
Self {
schema,
payload,
lossy: true,
}
}
pub fn rendered(&self) -> String {
serde_json::to_string(&self.payload).unwrap_or_default()
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum SliceError {
#[error("slicer {slicer} produced {produced} tokens against a cap of {cap}")]
OverCap {
slicer: String,
produced: usize,
cap: usize,
},
#[error("a cap of {cap} tokens cannot hold even a summary (needs at least {minimum})")]
CapTooSmall {
cap: usize,
minimum: usize,
},
#[error("the raw output could not be stored: {0}")]
Store(#[from] StoreError),
#[error("{slicer} cannot slice this input: {detail}")]
Unsuitable {
slicer: String,
detail: String,
},
}
impl SliceError {
pub fn can_fall_back(&self) -> bool {
matches!(self, Self::Unsuitable { .. })
}
}
#[async_trait]
pub trait Slicer: Send + Sync + fmt::Debug {
fn id(&self) -> SlicerId;
fn counter(&self) -> &dyn TokenCount {
&Estimate
}
fn minimum_cap(&self) -> usize {
16
}
async fn slice(&self, raw: &Raw, cap: usize) -> Result<Slice, SliceError>;
async fn condense(
&self,
raw: &Raw,
cap: usize,
store: &dyn Store,
) -> Result<Condensate, SliceError> {
let minimum = self.minimum_cap();
if cap < minimum {
return Err(SliceError::CapTooSmall { cap, minimum });
}
let digest = store.put(&raw.bytes).await?;
let slice = self.slice(raw, cap).await?;
let produced = self.counter().count(&slice.rendered());
if produced > cap {
return Err(SliceError::OverCap {
slicer: self.id().name.clone(),
produced,
cap,
});
}
Ok(Condensate::new(
slice.schema,
CondensatePayload::inline(&slice.payload).map_err(|e| SliceError::Unsuitable {
slicer: self.id().name.clone(),
detail: format!("the summary could not be encoded: {e}"),
})?,
digest,
cap,
slice.lossy,
self.id(),
))
}
}
#[derive(Debug, Default)]
pub struct MemoryStore {
blobs: std::sync::Mutex<std::collections::BTreeMap<Digest, Vec<u8>>>,
}
impl MemoryStore {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.blobs.lock().expect("store lock").len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl Store for MemoryStore {
async fn put(&self, bytes: &[u8]) -> Result<Digest, StoreError> {
let digest = Digest::of(bytes);
self.blobs
.lock()
.expect("store lock")
.insert(digest, bytes.to_vec());
Ok(digest)
}
async fn get(&self, digest: &Digest) -> Result<Vec<u8>, StoreError> {
self.blobs
.lock()
.expect("store lock")
.get(digest)
.cloned()
.ok_or(StoreError::NotFound(*digest))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[derive(Debug)]
struct Fixed {
payload: serde_json::Value,
lossy: bool,
}
impl Fixed {
fn new(payload: serde_json::Value) -> Self {
Self {
payload,
lossy: false,
}
}
}
#[async_trait]
impl Slicer for Fixed {
fn id(&self) -> SlicerId {
SlicerId::new("fixed", 1)
}
async fn slice(&self, _raw: &Raw, _cap: usize) -> Result<Slice, SliceError> {
Ok(Slice {
schema: SchemaId::new("test.fixed.v1", 1),
payload: self.payload.clone(),
lossy: self.lossy,
})
}
}
fn raw() -> Raw {
Raw::new(b"the complete original output".to_vec())
}
#[tokio::test]
async fn condense_records_the_slicer_id_and_the_cap() {
let store = MemoryStore::new();
let condensate = Fixed::new(json!("small"))
.condense(&raw(), 100, &store)
.await
.unwrap();
assert_eq!(condensate.slicer, SlicerId::new("fixed", 1));
assert_eq!(condensate.schema, SchemaId::new("test.fixed.v1", 1));
assert_eq!(condensate.cap, 100);
}
#[tokio::test]
async fn the_original_is_preserved_and_addressed_by_the_condensate() {
let store = MemoryStore::new();
let raw = raw();
let condensate = Fixed::new(json!("small"))
.condense(&raw, 100, &store)
.await
.unwrap();
assert_eq!(condensate.raw, Digest::of(&raw.bytes));
assert_eq!(store.get(&condensate.raw).await.unwrap(), raw.bytes);
}
#[tokio::test]
async fn the_original_survives_a_slicer_that_fails() {
#[derive(Debug)]
struct Fails;
#[async_trait]
impl Slicer for Fails {
fn id(&self) -> SlicerId {
SlicerId::new("fails", 1)
}
async fn slice(&self, _raw: &Raw, _cap: usize) -> Result<Slice, SliceError> {
Err(SliceError::Unsuitable {
slicer: "fails".into(),
detail: "not my kind of input".into(),
})
}
}
let store = MemoryStore::new();
let raw = raw();
assert!(Fails.condense(&raw, 100, &store).await.is_err());
assert_eq!(
store.get(&Digest::of(&raw.bytes)).await.unwrap(),
raw.bytes,
"the original must outlive a failed slice"
);
}
#[tokio::test]
async fn a_slicer_that_overruns_its_cap_is_refused_not_truncated() {
let store = MemoryStore::new();
let long = json!("x".repeat(4_000));
let err = Fixed::new(long)
.condense(&raw(), 50, &store)
.await
.unwrap_err();
let SliceError::OverCap {
slicer,
produced,
cap,
} = err
else {
panic!("expected OverCap, got {err:?}");
};
assert_eq!(slicer, "fixed");
assert_eq!(cap, 50);
assert!(produced > 50);
assert!(!err_can_fall_back(produced, cap));
}
fn err_can_fall_back(produced: usize, cap: usize) -> bool {
SliceError::OverCap {
slicer: "fixed".into(),
produced,
cap,
}
.can_fall_back()
}
#[tokio::test]
async fn a_cap_too_small_for_any_summary_is_its_own_error() {
let store = MemoryStore::new();
let err = Fixed::new(json!("small"))
.condense(&raw(), 3, &store)
.await
.unwrap_err();
assert!(matches!(err, SliceError::CapTooSmall { cap: 3, .. }));
assert!(
!err.can_fall_back(),
"another slicer has the same window problem"
);
}
#[tokio::test]
async fn only_an_unsuitable_input_invites_a_fallback() {
let unsuitable = SliceError::Unsuitable {
slicer: "ast".into(),
detail: "does not parse".into(),
};
assert!(unsuitable.can_fall_back());
}
#[tokio::test]
async fn a_slice_exactly_at_the_cap_is_accepted() {
let store = MemoryStore::new();
let payload = json!("abc");
let exact = Estimate.count(&serde_json::to_string(&payload).unwrap());
assert!(
Fixed::new(payload)
.condense(&raw(), exact.max(16), &store)
.await
.is_ok()
);
}
#[tokio::test]
async fn lossy_is_carried_through_honestly() {
let store = MemoryStore::new();
let clean = Fixed::new(json!("all of it"))
.condense(&raw(), 100, &store)
.await
.unwrap();
assert!(!clean.lossy);
let dropped = Fixed {
payload: json!("some of it"),
lossy: true,
}
.condense(&raw(), 100, &store)
.await
.unwrap();
assert!(dropped.lossy);
assert!(dropped.needs_raw_for_completeness());
}
#[tokio::test]
async fn slicing_the_same_input_twice_gives_the_same_condensate() {
let store = MemoryStore::new();
let slicer = Fixed::new(json!({ "b": 2, "a": 1 }));
let first = slicer.condense(&raw(), 100, &store).await.unwrap();
for _ in 0..8 {
assert_eq!(slicer.condense(&raw(), 100, &store).await.unwrap(), first);
}
}
#[test]
fn the_default_estimate_is_pessimistic_rather_than_accurate() {
let prose = "the quick brown fox jumps over the lazy dog";
assert!(Estimate.count(prose) > prose.len() / 4);
assert_eq!(Estimate.count(""), 0);
}
#[test]
fn the_estimate_charges_for_lines_because_slicer_output_is_line_dense() {
let one_line = "aaaaaaaaaaa";
let three_lines = "aaa\naaa\naaa";
assert_eq!(
one_line.len(),
three_lines.len(),
"same bytes, different shape"
);
assert!(Estimate.count(three_lines) > Estimate.count(one_line));
}
#[test]
fn the_byte_ceiling_is_a_genuine_upper_bound() {
for text in ["", "a", "hello world", "{\"k\":[1,2,3]}"] {
assert_eq!(ByteCeiling.count(text), text.len());
assert!(
ByteCeiling.count(text) >= Estimate.count(text) || text.len() < 4,
"the bound should not sit below the estimate for real inputs"
);
}
}
#[tokio::test]
async fn storing_the_same_bytes_twice_stores_them_once() {
let store = MemoryStore::new();
let a = store.put(b"identical").await.unwrap();
let b = store.put(b"identical").await.unwrap();
assert_eq!(a, b);
assert_eq!(store.len(), 1);
}
#[tokio::test]
async fn an_address_nothing_was_stored_at_is_a_typed_miss() {
let store = MemoryStore::new();
let err = store.get(&Digest::of(b"never stored")).await.unwrap_err();
assert!(matches!(err, StoreError::NotFound(_)));
}
#[tokio::test]
async fn a_store_failure_stops_the_slice() {
#[derive(Debug)]
struct Broken;
#[async_trait]
impl Store for Broken {
async fn put(&self, _bytes: &[u8]) -> Result<Digest, StoreError> {
Err(StoreError::Unavailable("disk full".into()))
}
async fn get(&self, digest: &Digest) -> Result<Vec<u8>, StoreError> {
Err(StoreError::NotFound(*digest))
}
}
let err = Fixed::new(json!("small"))
.condense(&raw(), 100, &Broken)
.await
.unwrap_err();
assert!(matches!(err, SliceError::Store(_)));
}
#[test]
fn raw_text_survives_bytes_that_are_not_utf8() {
let raw = Raw::new(vec![b'o', b'k', 0xff, b'!']);
assert!(raw.text().contains("ok"));
assert_eq!(raw.len(), 4);
assert!(!raw.is_empty());
}
#[test]
fn a_command_outcome_summarises_status_and_timing() {
let ok = CommandOutcome::new("cargo build", 0, Duration::from_millis(1_200));
assert!(ok.succeeded());
assert!(ok.summary().contains("cargo build"));
assert!(ok.summary().contains("ok"));
assert!(ok.summary().contains("1200ms"));
let failed = CommandOutcome::new("cargo test", 101, Duration::from_millis(50));
assert!(!failed.succeeded());
assert!(failed.summary().contains("exit 101"));
let killed = CommandOutcome {
command: "sleep 100".into(),
exit: None,
duration: Duration::from_secs(9),
};
assert!(killed.summary().contains("killed"));
}
}