Skip to main content

fory_core/meta/
meta_string.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::ensure;
19use crate::error::Error;
20use crate::util::is_latin;
21use std::sync::OnceLock;
22
23// equal to "std::i16::MAX"
24const SHORT_MAX_VALUE: usize = 32767;
25// const HEADER_MASK:i64 = 0xff;
26
27pub static NAMESPACE_ENCODER: MetaStringEncoder = MetaStringEncoder::new('.', '_');
28pub static TYPE_NAME_ENCODER: MetaStringEncoder = MetaStringEncoder::new('$', '_');
29pub static FIELD_NAME_ENCODER: MetaStringEncoder = MetaStringEncoder::new('$', '_');
30
31pub static NAMESPACE_DECODER: MetaStringDecoder = MetaStringDecoder::new('.', '_');
32pub static FIELD_NAME_DECODER: MetaStringDecoder = MetaStringDecoder::new('$', '_');
33pub static TYPE_NAME_DECODER: MetaStringDecoder = MetaStringDecoder::new('$', '_');
34
35#[derive(Debug, PartialEq, Hash, Eq, Clone, Copy, Default)]
36#[repr(i16)]
37pub enum Encoding {
38    #[default]
39    Utf8 = 0x00,
40    LowerSpecial = 0x01,
41    LowerUpperDigitSpecial = 0x02,
42    FirstToLowerSpecial = 0x03,
43    AllToLowerSpecial = 0x04,
44}
45
46#[derive(Debug, Clone, Default)]
47pub struct MetaString {
48    pub original: String,
49    pub encoding: Encoding,
50    pub bytes: Vec<u8>,
51    pub strip_last_char: bool,
52    pub special_char1: char,
53    pub special_char2: char,
54}
55
56// Encoding is part of the wire key because the same bytes can decode to different names.
57// Derived and decoder-specific fields are not part of the encoded identity.
58impl PartialEq for MetaString {
59    fn eq(&self, other: &Self) -> bool {
60        self.encoding == other.encoding && self.bytes == other.bytes
61    }
62}
63
64impl Eq for MetaString {}
65
66impl std::hash::Hash for MetaString {
67    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
68        self.encoding.hash(state);
69        self.bytes.hash(state);
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn rejects_invalid_utf8_meta_string() {
79        let err = TYPE_NAME_DECODER
80            .decode(&[0xff], Encoding::Utf8)
81            .unwrap_err();
82        assert!(
83            err.to_string().contains("invalid UTF-8 meta string"),
84            "unexpected error: {err}"
85        );
86    }
87}
88
89static EMPTY: OnceLock<MetaString> = OnceLock::new();
90
91impl MetaString {
92    pub fn new(
93        original: String,
94        encoding: Encoding,
95        bytes: Vec<u8>,
96        special_char1: char,
97        special_char2: char,
98    ) -> Result<Self, Error> {
99        let mut strip_last_char = false;
100        if encoding != Encoding::Utf8 {
101            if bytes.is_empty() {
102                return Err(Error::encode_error("Encoded data cannot be empty"));
103            }
104            strip_last_char = (bytes[0] & 0x80) != 0;
105        }
106        Ok(MetaString {
107            original,
108            encoding,
109            bytes,
110            strip_last_char,
111            special_char1,
112            special_char2,
113        })
114    }
115
116    pub fn write_to(&self, writer: &mut crate::buffer::Writer) {
117        writer.write_var_u32(self.bytes.len() as u32);
118        writer.write_bytes(&self.bytes);
119    }
120
121    pub fn get_empty() -> &'static MetaString {
122        EMPTY.get_or_init(|| MetaString {
123            original: "".to_string(),
124            encoding: Encoding::default(),
125            bytes: vec![],
126            strip_last_char: false,
127            special_char1: '\0',
128            special_char2: '\0',
129        })
130    }
131}
132
133#[derive(Clone)]
134pub struct MetaStringDecoder {
135    pub special_char1: char,
136    pub special_char2: char,
137}
138
139#[derive(Clone)]
140pub struct MetaStringEncoder {
141    pub special_char1: char,
142    pub special_char2: char,
143}
144
145#[derive(Debug)]
146struct StringStatistics {
147    digit_count: usize,
148    upper_count: usize,
149    can_lower_upper_digit_special_encoded: bool,
150    can_lower_special_encoded: bool,
151}
152
153impl MetaStringEncoder {
154    pub const fn new(special_char1: char, special_char2: char) -> Self {
155        Self {
156            special_char1,
157            special_char2,
158        }
159    }
160
161    fn is_latin(&self, s: &str) -> bool {
162        is_latin(s)
163    }
164
165    fn _encode(&self, input: &str) -> Result<Option<MetaString>, Error> {
166        if input.is_empty() {
167            return Ok(Some(MetaString::new(
168                input.to_string(),
169                Encoding::Utf8,
170                vec![],
171                self.special_char1,
172                self.special_char2,
173            )?));
174        }
175
176        ensure!(
177            input.len() < SHORT_MAX_VALUE,
178            Error::encode_error(format!(
179                "Meta string is too long, max:{SHORT_MAX_VALUE}, current:{}",
180                input.len()
181            ))
182        );
183
184        if !self.is_latin(input) {
185            return Ok(Some(MetaString::new(
186                input.to_string(),
187                Encoding::Utf8,
188                input.as_bytes().to_vec(),
189                self.special_char1,
190                self.special_char2,
191            )?));
192        }
193
194        Ok(None)
195    }
196
197    pub fn encode(&self, input: &str) -> Result<MetaString, Error> {
198        if let Some(ms) = self._encode(input)? {
199            return Ok(ms);
200        }
201        let encoding = self.compute_encoding(input, None);
202        self.encode_with_encoding(input, encoding)
203    }
204
205    pub fn encode_with_encodings(
206        &self,
207        input: &str,
208        encodings: &[Encoding],
209    ) -> Result<MetaString, Error> {
210        if let Some(ms) = self._encode(input)? {
211            return Ok(ms);
212        }
213        let encoding = self.compute_encoding(input, Some(encodings));
214        self.encode_with_encoding(input, encoding)
215    }
216
217    fn compute_encoding(&self, input: &str, encodings: Option<&[Encoding]>) -> Encoding {
218        let allow = |e: Encoding| encodings.map_or(true, |opts| opts.contains(&e));
219        let statistics = self.compute_statistics(input);
220        if statistics.can_lower_special_encoded && allow(Encoding::LowerSpecial) {
221            return Encoding::LowerSpecial;
222        }
223        if statistics.can_lower_upper_digit_special_encoded {
224            if statistics.digit_count != 0 && allow(Encoding::LowerUpperDigitSpecial) {
225                return Encoding::LowerUpperDigitSpecial;
226            }
227            let upper_count: usize = statistics.upper_count;
228            if upper_count == 1
229                && input.chars().next().unwrap().is_uppercase()
230                && allow(Encoding::FirstToLowerSpecial)
231            {
232                return Encoding::FirstToLowerSpecial;
233            }
234            if ((input.len() + upper_count) * 5) < (input.len() * 6)
235                && allow(Encoding::AllToLowerSpecial)
236            {
237                return Encoding::AllToLowerSpecial;
238            }
239            if allow(Encoding::LowerUpperDigitSpecial) {
240                return Encoding::LowerUpperDigitSpecial;
241            }
242        }
243        Encoding::Utf8
244    }
245
246    fn compute_statistics(&self, chars: &str) -> StringStatistics {
247        let mut can_lower_upper_digit_special_encoded = true;
248        let mut can_lower_special_encoded = true;
249        let mut digit_count = 0;
250        let mut upper_count = 0;
251        for c in chars.chars() {
252            if can_lower_upper_digit_special_encoded
253                && !(c.is_lowercase()
254                    || c.is_uppercase()
255                    || c.is_ascii_digit()
256                    || (c == self.special_char1 || c == self.special_char2))
257            {
258                can_lower_upper_digit_special_encoded = false;
259            }
260            if can_lower_special_encoded
261                && !(c.is_lowercase() || matches!(c, '.' | '_' | '$' | '|'))
262            {
263                can_lower_special_encoded = false;
264            }
265            if c.is_ascii_digit() {
266                digit_count += 1;
267            }
268            if c.is_uppercase() {
269                upper_count += 1;
270            }
271        }
272        StringStatistics {
273            digit_count,
274            upper_count,
275            can_lower_upper_digit_special_encoded,
276            can_lower_special_encoded,
277        }
278    }
279
280    pub fn encode_with_encoding(
281        &self,
282        input: &str,
283        encoding: Encoding,
284    ) -> Result<MetaString, Error> {
285        if input.is_empty() {
286            return MetaString::new(
287                input.to_string(),
288                Encoding::Utf8,
289                vec![],
290                self.special_char1,
291                self.special_char2,
292            );
293        }
294        ensure!(
295            input.len() < SHORT_MAX_VALUE,
296            Error::encode_error(format!(
297                "Meta string is too long, max:{SHORT_MAX_VALUE}, current:{}",
298                input.len()
299            ))
300        );
301        ensure!(
302            encoding == Encoding::Utf8 || self.is_latin(input),
303            Error::encode_error("Non-ASCII characters in meta string are not allowed")
304        );
305
306        if input.is_empty() {
307            return MetaString::new(
308                input.to_string(),
309                Encoding::Utf8,
310                vec![],
311                self.special_char1,
312                self.special_char2,
313            );
314        };
315
316        match encoding {
317            Encoding::LowerSpecial => {
318                let encoded_data = self.encode_lower_special(input)?;
319                MetaString::new(
320                    input.to_string(),
321                    encoding,
322                    encoded_data,
323                    self.special_char1,
324                    self.special_char2,
325                )
326            }
327            Encoding::LowerUpperDigitSpecial => {
328                let encoded_data = self.encode_lower_upper_digit_special(input)?;
329                MetaString::new(
330                    input.to_string(),
331                    encoding,
332                    encoded_data,
333                    self.special_char1,
334                    self.special_char2,
335                )
336            }
337            Encoding::FirstToLowerSpecial => {
338                let encoded_data = self.encode_first_to_lower_special(input)?;
339                MetaString::new(
340                    input.to_string(),
341                    encoding,
342                    encoded_data,
343                    self.special_char1,
344                    self.special_char2,
345                )
346            }
347            Encoding::AllToLowerSpecial => {
348                let upper_count = input.chars().filter(|c| c.is_uppercase()).count();
349                let encoded_data = self.encode_all_to_lower_special(input, upper_count)?;
350                MetaString::new(
351                    input.to_string(),
352                    encoding,
353                    encoded_data,
354                    self.special_char1,
355                    self.special_char2,
356                )
357            }
358            Encoding::Utf8 => {
359                let encoded_data = input.as_bytes().to_vec();
360                MetaString::new(
361                    input.to_string(),
362                    Encoding::Utf8,
363                    encoded_data,
364                    self.special_char1,
365                    self.special_char2,
366                )
367            }
368        }
369    }
370
371    fn encode_generic(&self, input: &str, bits_per_char: u8) -> Result<Vec<u8>, Error> {
372        let total_bits: usize = input.len() * bits_per_char as usize + 1;
373        let byte_length: usize = (total_bits + 7) / 8;
374        let mut bytes = vec![0; byte_length];
375        let mut current_bit = 1;
376        for c in input.chars() {
377            let value = self.char_to_value(c, bits_per_char)?;
378            for i in (0..bits_per_char).rev() {
379                if (value & (1 << i)) != 0 {
380                    let byte_pos: usize = current_bit / 8;
381                    let bit_pos: usize = current_bit % 8;
382                    bytes[byte_pos] |= 1 << (7 - bit_pos);
383                }
384                current_bit += 1;
385            }
386        }
387        if byte_length * 8 >= total_bits + bits_per_char as usize {
388            bytes[0] |= 0x80;
389        }
390        Ok(bytes)
391    }
392    pub fn encode_lower_special(&self, input: &str) -> Result<Vec<u8>, Error> {
393        self.encode_generic(input, 5)
394    }
395
396    pub fn encode_lower_upper_digit_special(&self, input: &str) -> Result<Vec<u8>, Error> {
397        self.encode_generic(input, 6)
398    }
399
400    pub fn encode_first_to_lower_special(&self, input: &str) -> Result<Vec<u8>, Error> {
401        if input.is_empty() {
402            return self.encode_generic("", 5);
403        }
404
405        let mut iter = input.char_indices();
406        let (first_idx, first_char) = iter.next().unwrap();
407
408        let lower = first_char.to_lowercase().to_string();
409
410        // Fast path: if lowercase has the same byte length and is ASCII,
411        // we can modify the first byte directly without rebuilding the string.
412        if lower.len() == first_char.len_utf8() && first_char.is_ascii() {
413            let mut bytes = input.as_bytes().to_owned();
414            bytes[first_idx] = lower.as_bytes()[0];
415            return self.encode_generic(std::str::from_utf8(&bytes).unwrap(), 5);
416        }
417
418        // rebuild only the necessary prefix + suffix (still efficient).
419        let (_, rest) = input.split_at(first_char.len_utf8());
420        let mut result = String::with_capacity(input.len() + lower.len() - first_char.len_utf8());
421        result.push_str(&lower);
422        result.push_str(rest);
423        self.encode_generic(&result, 5)
424    }
425
426    pub fn encode_all_to_lower_special(
427        &self,
428        input: &str,
429        upper_count: usize,
430    ) -> Result<Vec<u8>, Error> {
431        let mut new_chars = Vec::with_capacity(input.len() + upper_count);
432        for c in input.chars() {
433            if c.is_uppercase() {
434                new_chars.push('|');
435                new_chars.push(c.to_lowercase().next().unwrap());
436            } else {
437                new_chars.push(c);
438            }
439        }
440        self.encode_generic(&new_chars.iter().collect::<String>(), 5)
441    }
442
443    fn char_to_value(&self, c: char, bits_per_char: u8) -> Result<u8, Error> {
444        match bits_per_char {
445            5 => match c {
446                'a'..='z' => Ok(c as u8 - b'a'),
447                '.' => Ok(26),
448                '_' => Ok(27),
449                '$' => Ok(28),
450                '|' => Ok(29),
451                _ => Err(Error::encode_error(format!(
452                    "Unsupported character for LOWER_UPPER_DIGIT_SPECIAL encoding: {c}",
453                )))?,
454            },
455            6 => match c {
456                'a'..='z' => Ok(c as u8 - b'a'),
457                'A'..='Z' => Ok(c as u8 - b'A' + 26),
458                '0'..='9' => Ok(c as u8 - b'0' + 52),
459                _ => {
460                    if c == self.special_char1 {
461                        Ok(62)
462                    } else if c == self.special_char2 {
463                        Ok(63)
464                    } else {
465                        Err(Error::encode_error(format!(
466                            "Invalid character value for LOWER_SPECIAL decoding: {c:?}",
467                        )))?
468                    }
469                }
470            },
471            _ => unreachable!(),
472        }
473    }
474}
475
476impl MetaStringDecoder {
477    pub const fn new(special_char1: char, special_char2: char) -> Self {
478        MetaStringDecoder {
479            special_char1,
480            special_char2,
481        }
482    }
483
484    pub fn decode(&self, encoded_data: &[u8], encoding: Encoding) -> Result<MetaString, Error> {
485        let str = {
486            if encoded_data.is_empty() {
487                Ok("".to_string())
488            } else {
489                match encoding {
490                    Encoding::LowerSpecial => self.decode_lower_special(encoded_data),
491                    Encoding::LowerUpperDigitSpecial => {
492                        self.decode_lower_upper_digit_special(encoded_data)
493                    }
494                    Encoding::FirstToLowerSpecial => {
495                        self.decode_rep_first_lower_special(encoded_data)
496                    }
497                    Encoding::AllToLowerSpecial => {
498                        self.decode_rep_all_to_lower_special(encoded_data)
499                    }
500                    Encoding::Utf8 => std::str::from_utf8(encoded_data)
501                        .map(str::to_owned)
502                        .map_err(|_| Error::encoding_error("invalid UTF-8 meta string")),
503                }
504            }
505        }?;
506        MetaString::new(
507            str,
508            encoding,
509            Vec::from(encoded_data),
510            self.special_char1,
511            self.special_char2,
512        )
513    }
514
515    fn decode_lower_special(&self, data: &[u8]) -> Result<String, Error> {
516        let mut decoded = String::new();
517        let total_bits: usize = data.len() * 8;
518        let strip_last_char = (data[0] & 0x80) != 0;
519        let bit_mask: usize = 0b11111;
520        let mut bit_index = 1;
521        while bit_index + 5 <= total_bits && !(strip_last_char && (bit_index + 2 * 5 > total_bits))
522        {
523            let byte_index = bit_index / 8;
524            let intra_byte_index = bit_index % 8;
525            let char_value: usize = if intra_byte_index > 3 {
526                ((((data[byte_index] as usize) << 8)
527                    | if byte_index + 1 < data.len() {
528                        data.get(byte_index + 1).cloned().unwrap() as usize & 0xFF
529                    } else {
530                        0
531                    })
532                    >> (11 - intra_byte_index))
533                    & bit_mask
534            } else {
535                ((data[byte_index] as usize) >> (3 - intra_byte_index)) & bit_mask
536            };
537            bit_index += 5;
538            decoded.push(self.decode_lower_special_char(char_value as u8)?);
539        }
540        Ok(decoded)
541    }
542
543    fn decode_lower_upper_digit_special(&self, data: &[u8]) -> Result<String, Error> {
544        let mut decoded = String::new();
545        let num_bits = data.len() * 8;
546        let strip_last_char = (data[0] & 0x80) != 0;
547        let mut bit_index = 1;
548        let bit_mask: usize = 0b111111;
549        while bit_index + 6 <= num_bits && !(strip_last_char && (bit_index + 2 * 6 > num_bits)) {
550            let byte_index = bit_index / 8;
551            let intra_byte_index = bit_index % 8;
552            let char_value: usize = if intra_byte_index > 2 {
553                ((((data[byte_index] as usize) << 8)
554                    | if byte_index + 1 < data.len() {
555                        data.get(byte_index + 1).cloned().unwrap() as usize & 0xFF
556                    } else {
557                        0
558                    })
559                    >> (10 - intra_byte_index))
560                    & bit_mask
561            } else {
562                ((data[byte_index] as usize) >> (2 - intra_byte_index)) & bit_mask
563            };
564            bit_index += 6;
565            decoded.push(self.decode_lower_upper_digit_special_char(char_value as u8)?);
566        }
567        Ok(decoded)
568    }
569
570    fn decode_lower_special_char(&self, char_value: u8) -> Result<char, Error> {
571        match char_value {
572            0..=25 => Ok((b'a' + char_value) as char),
573            26 => Ok('.'),
574            27 => Ok('_'),
575            28 => Ok('$'),
576            29 => Ok('|'),
577            _ => Err(Error::encode_error(format!(
578                "Invalid character value for LOWER_SPECIAL decoding: {char_value}",
579            )))?,
580        }
581    }
582
583    fn decode_lower_upper_digit_special_char(&self, char_value: u8) -> Result<char, Error> {
584        match char_value {
585            0..=25 => Ok((b'a' + char_value) as char),
586            26..=51 => Ok((b'A' + char_value - 26) as char),
587            52..=61 => Ok((b'0' + char_value - 52) as char),
588            62 => Ok(self.special_char1),
589            63 => Ok(self.special_char2),
590            _ => Err(Error::encode_error(format!(
591                "Invalid character value for LOWER_UPPER_DIGIT_SPECIAL decoding: {char_value}",
592            )))?,
593        }
594    }
595
596    fn decode_rep_first_lower_special(&self, data: &[u8]) -> Result<String, Error> {
597        let decoded_str = self.decode_lower_special(data)?;
598        let mut chars = decoded_str.chars();
599        match chars.next() {
600            Some(first_char) => {
601                let mut result = first_char.to_ascii_uppercase().to_string();
602                result.extend(chars);
603                Ok(result)
604            }
605            None => Ok(decoded_str),
606        }
607    }
608    fn decode_rep_all_to_lower_special(&self, data: &[u8]) -> Result<String, Error> {
609        let decoded_str = self.decode_lower_special(data)?;
610        let mut result = String::with_capacity(decoded_str.len());
611        let mut chars = decoded_str.chars();
612        while let Some(char) = chars.next() {
613            if char == '|' {
614                if let Some(next_char) = chars.next() {
615                    result.push(next_char.to_ascii_uppercase());
616                }
617            } else {
618                result.push(char);
619            }
620        }
621        Ok(result)
622    }
623}