Skip to main content

micromegas_transit/
dyn_string.rs

1use crate::{
2    InProcSerialize, InProcSize, string_codec::StringCodec, try_advance_window,
3    try_read_consume_pod, write_any,
4};
5use anyhow::Result;
6use bumpalo::Bump;
7use std::borrow::Cow;
8
9#[derive(Debug)]
10pub struct LegacyDynString(pub String);
11
12impl InProcSerialize for LegacyDynString {
13    const IN_PROC_SIZE: InProcSize = InProcSize::Dynamic;
14
15    fn get_value_size(&self) -> Option<u32> {
16        Some(self.0.len() as u32)
17    }
18
19    fn write_value(&self, buffer: &mut Vec<u8>) {
20        buffer.extend_from_slice(self.0.as_bytes());
21    }
22
23    #[allow(unsafe_code)]
24    unsafe fn read_value(window: &[u8]) -> Self {
25        Self(String::from_utf8(window.to_vec()).unwrap())
26    }
27}
28
29#[derive(Debug)]
30pub struct DynString(pub String);
31
32impl InProcSerialize for DynString {
33    const IN_PROC_SIZE: InProcSize = InProcSize::Dynamic;
34
35    fn get_value_size(&self) -> Option<u32> {
36        let header_size = 1 + // codec
37			std::mem::size_of::<u32>() as u32 // size in bytes
38			;
39        let string_size = self.0.len() as u32;
40        Some(header_size + string_size)
41    }
42
43    fn write_value(&self, buffer: &mut Vec<u8>) {
44        let codec = StringCodec::Utf8 as u8;
45        write_any(buffer, &codec);
46        let len = self.0.len() as u32;
47        write_any(buffer, &len);
48        buffer.extend_from_slice(self.0.as_bytes());
49    }
50
51    #[allow(unsafe_code)]
52    unsafe fn read_value(mut window: &[u8]) -> Self {
53        let res = read_advance_string(&mut window).unwrap();
54        assert_eq!(window.len(), 0);
55        Self(res)
56    }
57}
58
59/// Parse string from buffer, move buffer pointer forward.
60pub fn read_advance_string(window: &mut &[u8]) -> Result<String> {
61    let codec = StringCodec::try_from(try_read_consume_pod::<u8>(window)?)?;
62    let string_len_bytes: u32 = try_read_consume_pod(window)?;
63    if string_len_bytes as usize > window.len() {
64        anyhow::bail!(
65            "truncated string: need {string_len_bytes} bytes, have {}",
66            window.len()
67        );
68    }
69    let string_buffer = &window[0..(string_len_bytes as usize)];
70    *window = try_advance_window(window, string_len_bytes as usize)?;
71    match codec {
72        StringCodec::Ansi => {
73            // this would be typically be windows 1252, an extension to ISO-8859-1/latin1
74            // random people on the interwebs tell me that latin1's codepoints are a subset of utf8
75            // so I guess it's ok to treat it as utf8
76            Ok(String::from_utf8_lossy(string_buffer).to_string())
77        }
78        StringCodec::Wide => {
79            if !string_len_bytes.is_multiple_of(2) {
80                anyhow::bail!("wrong utf-16 buffer size");
81            }
82            // Decode UTF-16 LE without assuming the source bytes are 2-byte aligned.
83            let units = string_buffer
84                .chunks_exact(2)
85                .map(|pair| u16::from_le_bytes([pair[0], pair[1]]));
86            let s: String = char::decode_utf16(units)
87                .map(|r| r.unwrap_or(char::REPLACEMENT_CHARACTER))
88                .collect();
89            Ok(s)
90        }
91        StringCodec::Utf8 => Ok(String::from_utf8_lossy(string_buffer).to_string()),
92    }
93}
94
95/// Parse a string from the buffer, moving the buffer pointer forward.
96///
97/// Borrows the source buffer (`'a`) where possible: a valid-UTF-8 string is
98/// returned as a zero-copy slice of the buffer. Only the transcoded cases
99/// (UTF-16 wide, or lossy replacement of invalid bytes) allocate into the arena.
100pub fn read_advance_string_in<'a>(bump: &'a Bump, window: &mut &'a [u8]) -> Result<&'a str> {
101    let codec = StringCodec::try_from(try_read_consume_pod::<u8>(window)?)?;
102    let string_len_bytes: u32 = try_read_consume_pod(window)?;
103    if string_len_bytes as usize > window.len() {
104        anyhow::bail!(
105            "truncated string: need {string_len_bytes} bytes, have {}",
106            window.len()
107        );
108    }
109    let string_buffer = &window[0..(string_len_bytes as usize)];
110    *window = try_advance_window(window, string_len_bytes as usize)?;
111    match codec {
112        // Treat ANSI (windows-1252/latin1) as utf8, matching read_advance_string.
113        StringCodec::Ansi | StringCodec::Utf8 => match String::from_utf8_lossy(string_buffer) {
114            // Valid UTF-8: borrow the source buffer directly (zero-copy).
115            Cow::Borrowed(s) => Ok(s),
116            // Invalid bytes were replaced: the transcoded result lives in the arena.
117            Cow::Owned(s) => Ok(bump.alloc_str(&s)),
118        },
119        StringCodec::Wide => {
120            if !string_len_bytes.is_multiple_of(2) {
121                anyhow::bail!("wrong utf-16 buffer size");
122            }
123            // Decode UTF-16 LE without assuming the source bytes are 2-byte aligned.
124            let units = string_buffer
125                .chunks_exact(2)
126                .map(|pair| u16::from_le_bytes([pair[0], pair[1]]));
127            let s: String = char::decode_utf16(units)
128                .map(|r| r.unwrap_or(char::REPLACEMENT_CHARACTER))
129                .collect();
130            Ok(bump.alloc_str(&s))
131        }
132    }
133}