kcode_k1_audio_classification_fragment_format/
lib.rs

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

kcode_k1_audio_classification_fragment_format/
lib.rs
