kcode_k1_audio_classification_fragment_format/
lib.rs

1use std::path::{Component, Path, PathBuf};
2
3pub use kcode_k1_audio_classification_format_error::FormatError;
4pub use kcode_k1_transaction_id::TxId;
5pub use kcode_speaker_v3_analysis::{FeatureVector24, LocalSpeakerLabel};
6
7const FRAGMENT_VERSION: u8 = 1;
8const STAGED_KIND: u8 = 1;
9const FINAL_KIND: u8 = 2;
10const ANALYSIS_SLOT_OFFSET: usize = 16;
11const CONFIRMATION_SLOT_OFFSET: usize = 32;
12const BODY_OFFSET: usize = 48;
13const BASE64_ALPHABET: &[u8; 64] =
14    b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
15
16#[derive(Debug, Clone, PartialEq)]
17pub struct StagedSpeakerV1 {
18    pub speaker: LocalSpeakerLabel,
19    pub language: String,
20    pub features: FeatureVector24,
21    pub usable_for_training: bool,
22}
23
24#[derive(Debug, Clone, PartialEq)]
25pub struct StagedFragmentV1 {
26    pub analysis_txid: TxId,
27    pub transcript: String,
28    pub speakers: Vec<StagedSpeakerV1>,
29}
30
31#[derive(Debug, Clone, PartialEq)]
32pub struct FinalSpeakerV1 {
33    pub speaker: LocalSpeakerLabel,
34    pub person_id: String,
35    pub language: String,
36    pub features: FeatureVector24,
37    pub usable_for_training: bool,
38}
39
40#[derive(Debug, Clone, PartialEq)]
41pub struct FinalFragmentV1 {
42    pub analysis_txid: TxId,
43    pub confirmation_txid: TxId,
44    pub transcript: String,
45    pub speakers: Vec<FinalSpeakerV1>,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub struct TxIdSlot {
50    txid: TxId,
51}
52
53impl TxIdSlot {
54    pub const LEN: usize = 16;
55    pub const PADDING_LEN: usize = 4;
56
57    pub const fn new(txid: TxId) -> Self {
58        Self { txid }
59    }
60
61    pub const fn txid(self) -> TxId {
62        self.txid
63    }
64
65    pub fn encode(self) -> [u8; Self::LEN] {
66        let mut output = [0; Self::LEN];
67        output[..12].copy_from_slice(self.txid.as_bytes());
68        output
69    }
70
71    pub fn decode(bytes: &[u8]) -> Result<Self, FormatError> {
72        if bytes.len() != Self::LEN {
73            return Err(FormatError::InvalidTxIdSlotLength(bytes.len()));
74        }
75        if bytes[12..].iter().any(|byte| *byte != 0) {
76            return Err(FormatError::NonZeroPadding);
77        }
78        let mut txid = [0; 12];
79        txid.copy_from_slice(&bytes[..12]);
80        Ok(Self::new(TxId::from_bytes(txid)))
81    }
82}
83
84pub fn encode_staged_fragment(value: &StagedFragmentV1) -> Result<Vec<u8>, FormatError> {
85    let mut output = encode_fragment_prefix(STAGED_KIND, value.analysis_txid, None);
86    append_string(&mut output, &value.transcript)?;
87    append_length(&mut output, value.speakers.len())?;
88    for speaker in &value.speakers {
89        append_string(&mut output, &speaker.speaker.to_string())?;
90        append_string(&mut output, &speaker.language)?;
91        append_features(&mut output, &speaker.features)?;
92        output.push(u8::from(speaker.usable_for_training));
93    }
94    Ok(output)
95}
96
97pub fn decode_staged_fragment(bytes: &[u8]) -> Result<StagedFragmentV1, FormatError> {
98    let (analysis_txid, _) = decode_fragment_header(bytes, STAGED_KIND, false)?;
99    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
100    let transcript = decoder.read_string()?;
101    let speaker_count = decoder.read_length()?;
102    let mut speakers = Vec::new();
103    for _ in 0..speaker_count {
104        speakers.push(StagedSpeakerV1 {
105            speaker: decoder.read_speaker_label()?,
106            language: decoder.read_string()?,
107            features: decoder.read_features()?,
108            usable_for_training: decoder.read_boolean()?,
109        });
110    }
111    decoder.finish()?;
112    Ok(StagedFragmentV1 {
113        analysis_txid,
114        transcript,
115        speakers,
116    })
117}
118
119pub fn encode_final_fragment(value: &FinalFragmentV1) -> Result<Vec<u8>, FormatError> {
120    let mut output = encode_fragment_prefix(
121        FINAL_KIND,
122        value.analysis_txid,
123        Some(value.confirmation_txid),
124    );
125    append_string(&mut output, &value.transcript)?;
126    append_length(&mut output, value.speakers.len())?;
127    for speaker in &value.speakers {
128        append_string(&mut output, &speaker.speaker.to_string())?;
129        append_string(&mut output, &speaker.person_id)?;
130        append_string(&mut output, &speaker.language)?;
131        append_features(&mut output, &speaker.features)?;
132        output.push(u8::from(speaker.usable_for_training));
133    }
134    Ok(output)
135}
136
137pub fn decode_final_fragment(bytes: &[u8]) -> Result<FinalFragmentV1, FormatError> {
138    let (analysis_txid, confirmation_txid) = decode_fragment_header(bytes, FINAL_KIND, true)?;
139    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
140    let transcript = decoder.read_string()?;
141    let speaker_count = decoder.read_length()?;
142    let mut speakers = Vec::new();
143    for _ in 0..speaker_count {
144        speakers.push(FinalSpeakerV1 {
145            speaker: decoder.read_speaker_label()?,
146            person_id: decoder.read_string()?,
147            language: decoder.read_string()?,
148            features: decoder.read_features()?,
149            usable_for_training: decoder.read_boolean()?,
150        });
151    }
152    decoder.finish()?;
153    Ok(FinalFragmentV1 {
154        analysis_txid,
155        confirmation_txid: confirmation_txid
156            .expect("final fragment header always contains a confirmation slot"),
157        transcript,
158        speakers,
159    })
160}
161
162pub fn txid_path(txid: TxId) -> PathBuf {
163    let encoded =
164        String::from_utf8(encode_txid_base64(txid).to_vec()).expect("base64 alphabet is ASCII");
165    PathBuf::from(&encoded[..1]).join(format!("{}.dat", &encoded[1..]))
166}
167
168pub fn txid_from_path(path: impl AsRef<Path>) -> Result<TxId, FormatError> {
169    let mut components = path.as_ref().components();
170    let shard = normal_utf8_component(components.next())?;
171    let filename = normal_utf8_component(components.next())?;
172    if components.next().is_some() || shard.len() != 1 {
173        return Err(FormatError::InvalidPath);
174    }
175    let name = filename
176        .strip_suffix(".dat")
177        .ok_or(FormatError::InvalidPath)?;
178    if name.len() != 15 {
179        return Err(FormatError::InvalidPath);
180    }
181    let mut encoded = [0; 16];
182    encoded[0] = shard.as_bytes()[0];
183    encoded[1..].copy_from_slice(name.as_bytes());
184    decode_txid_base64(encoded)
185}
186
187fn encode_fragment_prefix(
188    kind: u8,
189    analysis_txid: TxId,
190    confirmation_txid: Option<TxId>,
191) -> Vec<u8> {
192    let mut output = vec![0; BODY_OFFSET];
193    output[0] = FRAGMENT_VERSION;
194    output[1] = kind;
195    output[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET]
196        .copy_from_slice(&TxIdSlot::new(analysis_txid).encode());
197    if let Some(txid) = confirmation_txid {
198        output[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET]
199            .copy_from_slice(&TxIdSlot::new(txid).encode());
200    }
201    output
202}
203
204fn decode_fragment_header(
205    bytes: &[u8],
206    expected_kind: u8,
207    has_confirmation: bool,
208) -> Result<(TxId, Option<TxId>), FormatError> {
209    if bytes.len() < BODY_OFFSET {
210        return Err(FormatError::Truncated);
211    }
212    if bytes[0] != FRAGMENT_VERSION {
213        return Err(FormatError::UnsupportedVersion(bytes[0]));
214    }
215    if bytes[1] != expected_kind {
216        return Err(FormatError::InvalidFragmentKind(bytes[1]));
217    }
218    let reserved = &bytes[2..ANALYSIS_SLOT_OFFSET];
219    if reserved.iter().any(|byte| *byte != 0) {
220        return Err(FormatError::NonZeroReserved);
221    }
222    let analysis_txid =
223        TxIdSlot::decode(&bytes[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET])?.txid();
224    let confirmation_slot = &bytes[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET];
225    let confirmation_txid = if has_confirmation {
226        Some(TxIdSlot::decode(confirmation_slot)?.txid())
227    } else {
228        if confirmation_slot.iter().any(|byte| *byte != 0) {
229            return Err(FormatError::NonZeroStagedConfirmation);
230        }
231        None
232    };
233    Ok((analysis_txid, confirmation_txid))
234}
235
236fn append_length(output: &mut Vec<u8>, length: usize) -> Result<(), FormatError> {
237    let length = u64::try_from(length).map_err(|_| FormatError::LengthOverflow)?;
238    output.extend_from_slice(&length.to_le_bytes());
239    Ok(())
240}
241
242fn append_bytes(output: &mut Vec<u8>, bytes: &[u8]) -> Result<(), FormatError> {
243    append_length(output, bytes.len())?;
244    output.extend_from_slice(bytes);
245    Ok(())
246}
247
248fn append_string(output: &mut Vec<u8>, value: &str) -> Result<(), FormatError> {
249    append_bytes(output, value.as_bytes())
250}
251
252fn append_features(output: &mut Vec<u8>, value: &FeatureVector24) -> Result<(), FormatError> {
253    let bytes = postcard::to_allocvec(value).map_err(|_| FormatError::InvalidFeatureBody)?;
254    append_bytes(output, &bytes)
255}
256
257struct BodyDecoder<'a> {
258    bytes: &'a [u8],
259    position: usize,
260}
261
262impl<'a> BodyDecoder<'a> {
263    fn new(bytes: &'a [u8]) -> Self {
264        Self { bytes, position: 0 }
265    }
266
267    fn take(&mut self, length: usize) -> Result<&'a [u8], FormatError> {
268        let end = self
269            .position
270            .checked_add(length)
271            .ok_or(FormatError::LengthOverflow)?;
272        let value = self
273            .bytes
274            .get(self.position..end)
275            .ok_or(FormatError::Truncated)?;
276        self.position = end;
277        Ok(value)
278    }
279
280    fn read_length(&mut self) -> Result<usize, FormatError> {
281        let mut bytes = [0; 8];
282        bytes.copy_from_slice(self.take(8)?);
283        usize::try_from(u64::from_le_bytes(bytes)).map_err(|_| FormatError::LengthOverflow)
284    }
285
286    fn read_bytes(&mut self) -> Result<&'a [u8], FormatError> {
287        let length = self.read_length()?;
288        self.take(length)
289    }
290
291    fn read_string(&mut self) -> Result<String, FormatError> {
292        Ok(std::str::from_utf8(self.read_bytes()?)
293            .map_err(|_| FormatError::InvalidUtf8)?
294            .to_owned())
295    }
296
297    fn read_speaker_label(&mut self) -> Result<LocalSpeakerLabel, FormatError> {
298        self.read_string()?
299            .parse()
300            .map_err(|_| FormatError::InvalidSpeakerLabel)
301    }
302
303    fn read_features(&mut self) -> Result<FeatureVector24, FormatError> {
304        let (value, remaining) = postcard::take_from_bytes(self.read_bytes()?)
305            .map_err(|_| FormatError::InvalidFeatureBody)?;
306        if !remaining.is_empty() {
307            return Err(FormatError::InvalidFeatureBody);
308        }
309        Ok(value)
310    }
311
312    fn read_boolean(&mut self) -> Result<bool, FormatError> {
313        match self.take(1)?[0] {
314            0 => Ok(false),
315            1 => Ok(true),
316            value => Err(FormatError::InvalidBoolean(value)),
317        }
318    }
319
320    fn finish(self) -> Result<(), FormatError> {
321        if self.position == self.bytes.len() {
322            Ok(())
323        } else {
324            Err(FormatError::TrailingBytes)
325        }
326    }
327}
328
329fn normal_utf8_component(component: Option<Component<'_>>) -> Result<&str, FormatError> {
330    match component {
331        Some(Component::Normal(value)) => value.to_str().ok_or(FormatError::InvalidPath),
332        _ => Err(FormatError::InvalidPath),
333    }
334}
335
336fn encode_txid_base64(txid: TxId) -> [u8; 16] {
337    let bytes = txid.into_bytes();
338    let mut encoded = [0; 16];
339    for group in 0..4 {
340        let input = group * 3;
341        let output = group * 4;
342        encoded[output] = BASE64_ALPHABET[(bytes[input] >> 2) as usize];
343        encoded[output + 1] =
344            BASE64_ALPHABET[(((bytes[input] & 3) << 4) | (bytes[input + 1] >> 4)) as usize];
345        encoded[output + 2] =
346            BASE64_ALPHABET[(((bytes[input + 1] & 15) << 2) | (bytes[input + 2] >> 6)) as usize];
347        encoded[output + 3] = BASE64_ALPHABET[(bytes[input + 2] & 63) as usize];
348    }
349    encoded
350}
351
352fn decode_txid_base64(encoded: [u8; 16]) -> Result<TxId, FormatError> {
353    let mut bytes = [0; 12];
354    for group in 0..4 {
355        let input = group * 4;
356        let output = group * 3;
357        let first = decode_base64_character(encoded[input])?;
358        let second = decode_base64_character(encoded[input + 1])?;
359        let third = decode_base64_character(encoded[input + 2])?;
360        let fourth = decode_base64_character(encoded[input + 3])?;
361        bytes[output] = (first << 2) | (second >> 4);
362        bytes[output + 1] = (second << 4) | (third >> 2);
363        bytes[output + 2] = (third << 6) | fourth;
364    }
365    Ok(TxId::from_bytes(bytes))
366}
367
368fn decode_base64_character(value: u8) -> Result<u8, FormatError> {
369    BASE64_ALPHABET
370        .iter()
371        .position(|candidate| *candidate == value)
372        .and_then(|index| u8::try_from(index).ok())
373        .ok_or(FormatError::InvalidPath)
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    fn txid(seed: u8) -> TxId {
381        TxId::from_bytes([seed; 12])
382    }
383
384    fn staged_fragment() -> StagedFragmentV1 {
385        StagedFragmentV1 {
386            analysis_txid: txid(20),
387            transcript: "hello".into(),
388            speakers: vec![StagedSpeakerV1 {
389                speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
390                language: "English".into(),
391                features: FeatureVector24::default(),
392                usable_for_training: true,
393            }],
394        }
395    }
396
397    fn final_fragment() -> FinalFragmentV1 {
398        FinalFragmentV1 {
399            analysis_txid: txid(30),
400            confirmation_txid: txid(40),
401            transcript: "hello".into(),
402            speakers: vec![FinalSpeakerV1 {
403                speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
404                person_id: "person-1".into(),
405                language: "English".into(),
406                features: FeatureVector24::default(),
407                usable_for_training: false,
408            }],
409        }
410    }
411
412    #[test]
413    fn fragments_and_slots_roundtrip_with_exact_offsets() {
414        let staged = staged_fragment();
415        let staged_bytes = encode_staged_fragment(&staged).expect("encode staged");
416        assert_eq!(decode_staged_fragment(&staged_bytes), Ok(staged));
417
418        let final_value = final_fragment();
419        let final_bytes = encode_final_fragment(&final_value).expect("encode final");
420        assert_eq!(decode_final_fragment(&final_bytes), Ok(final_value));
421
422        assert_eq!(
423            TxIdSlot::decode(&[0; 15]),
424            Err(FormatError::InvalidTxIdSlotLength(15))
425        );
426    }
427
428    #[test]
429    fn malformed_headers_slots_and_bodies_are_rejected() {
430        let staged = staged_fragment();
431        let bytes = encode_staged_fragment(&staged).expect("encode staged");
432        let transcript = BODY_OFFSET + 8;
433        let label = transcript + staged.transcript.len() + 16;
434        let feature_length = label + staged.speakers[0].speaker.to_string().len() + 15;
435        let feature = feature_length + 8;
436        let boolean = feature
437            + postcard::to_allocvec(&staged.speakers[0].features)
438                .expect("features")
439                .len();
440        for (index, value, error) in [
441            (0, 2, FormatError::UnsupportedVersion(2)),
442            (1, 2, FormatError::InvalidFragmentKind(2)),
443            (2, 1, FormatError::NonZeroReserved),
444            (28, 1, FormatError::NonZeroPadding),
445            (32, 1, FormatError::NonZeroStagedConfirmation),
446            (transcript, 0xff, FormatError::InvalidUtf8),
447            (label, b'x', FormatError::InvalidSpeakerLabel),
448            (feature, 2, FormatError::InvalidFeatureBody),
449            (boolean, 2, FormatError::InvalidBoolean(2)),
450        ] {
451            let mut malformed = bytes.clone();
452            malformed[index] = value;
453            assert_eq!(decode_staged_fragment(&malformed), Err(error));
454        }
455        assert_eq!(
456            decode_staged_fragment(&bytes[..47]),
457            Err(FormatError::Truncated)
458        );
459        let mut trailing = bytes;
460        trailing.push(0);
461        assert_eq!(
462            decode_staged_fragment(&trailing),
463            Err(FormatError::TrailingBytes)
464        );
465    }
466
467    #[test]
468    fn every_path_shard_roundtrips_and_noncanonical_paths_fail() {
469        for (index, expected) in BASE64_ALPHABET.iter().enumerate() {
470            let mut bytes = [0; 12];
471            bytes[0] = (index as u8) << 2;
472            let txid = TxId::from_bytes(bytes);
473            let path = txid_path(txid);
474            assert_eq!(path.to_string_lossy().as_bytes()[0], *expected);
475            assert_eq!(txid_from_path(&path), Ok(txid));
476        }
477        for path in [
478            "/A/AAAAAAAAAAAAAAA.dat",
479            "AA/AAAAAAAAAAAAAAA.dat",
480            "A/AAAAAAAAAAAAAA!.dat",
481            "A/AAAAAAAAAAAAAAA.dat/extra",
482        ] {
483            assert_eq!(txid_from_path(path), Err(FormatError::InvalidPath));
484        }
485    }
486}
kcode_k1_audio_classification_fragment_format/lib.rs

kcode_k1_audio_classification_fragment_format/
lib.rs
