1use uqa_analysis::TokenTerm;
10
11mod allocation;
12
13use crate::{StorageBackendError, StorageBackendResult};
14
15#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
19pub struct TokenTermKey(Vec<u8>);
20
21impl TokenTermKey {
22 pub fn from_text(text: &str) -> Self {
23 let mut bytes = Vec::with_capacity(text.len() + 1);
24 bytes.push(0);
25 bytes.extend_from_slice(text.as_bytes());
26 Self(bytes)
27 }
28
29 pub fn from_term(term: &TokenTerm) -> Self {
30 if let Some(text) = term.as_str() {
31 return Self::from_text(text);
32 }
33 let units = term.utf16();
34 let mut bytes = Vec::with_capacity(units.len() * 2 + 1);
35 bytes.push(1);
36 for unit in units.iter() {
37 bytes.extend_from_slice(&unit.to_be_bytes());
38 }
39 Self(bytes)
40 }
41
42 pub fn from_bytes(bytes: Vec<u8>) -> StorageBackendResult<Self> {
43 Self::validate(&bytes)?;
44 Ok(Self(bytes))
45 }
46
47 pub fn as_str(&self) -> Option<&str> {
49 (self.0[0] == 0).then(|| std::str::from_utf8(&self.0[1..]).expect("validated UTF-8 key"))
50 }
51
52 pub fn as_bytes(&self) -> &[u8] {
53 &self.0
54 }
55
56 pub fn into_bytes(self) -> Vec<u8> {
57 self.0
58 }
59
60 pub fn to_term(&self) -> TokenTerm {
61 match self.0[0] {
62 0 => TokenTerm::from(std::str::from_utf8(&self.0[1..]).expect("validated UTF-8 key")),
63 1 => TokenTerm::from_utf16(
64 self.0[1..]
65 .chunks_exact(2)
66 .map(|pair| u16::from_be_bytes([pair[0], pair[1]]))
67 .collect(),
68 ),
69 _ => unreachable!("validated token key tag"),
70 }
71 }
72
73 fn validate(bytes: &[u8]) -> StorageBackendResult<()> {
74 let invalid = || StorageBackendError::Other("invalid canonical token term key".into());
75 match bytes.split_first() {
76 Some((0, bytes)) => std::str::from_utf8(bytes)
77 .map(|_| ())
78 .map_err(|_| invalid()),
79 Some((1, bytes)) if bytes.len() % 2 == 0 => {
80 let units = bytes
81 .chunks_exact(2)
82 .map(|pair| u16::from_be_bytes([pair[0], pair[1]]));
83 if !char::decode_utf16(units).any(|unit| unit.is_err()) {
84 return Err(invalid());
85 }
86 Ok(())
87 }
88 _ => Err(invalid()),
89 }
90 }
91}
92
93impl From<String> for TokenTermKey {
94 fn from(text: String) -> Self {
95 Self::from_text(&text)
96 }
97}
98
99impl From<&str> for TokenTermKey {
100 fn from(text: &str) -> Self {
101 Self::from_text(text)
102 }
103}