1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use crate::internal::encodings::varint::*;
use crate::prelude::*;
use rle::RLE;
use std::vec::IntoIter;

// TODO: Consider compressed unicode (SCSU?) for String in general,
// but in particular for schema strings. As schema strings need only
// be compared and usually not displayed we can do bit-for-bit comparisons
// (Make sure that's true for SCSU, which may allow multiple encodings!)

#[cfg(feature = "write")]
pub fn write_str<O: EncodeOptions>(value: &str, stream: &mut WriterStream<'_, O>) {
    write_usize(value.len(), stream);
    stream.bytes.extend_from_slice(value.as_bytes());
}

#[cfg(feature = "read")]
fn read_str_len<'a>(len: usize, bytes: &'a [u8], offset: &'_ mut usize) -> ReadResult<&'a str> {
    let utf8 = read_bytes(len, bytes, offset)?;
    Ok(std::str::from_utf8(utf8)?)
}

#[cfg(feature = "read")]
pub fn read_str<'a>(bytes: &'a [u8], offset: &'_ mut usize) -> ReadResult<&'a str> {
    let len = decode_prefix_varint(bytes, offset)? as usize;
    read_str_len(len, bytes, offset)
}

#[cfg(feature = "write")]
impl Writable for String {
    type WriterArray = Vec<&'static str>;
    fn write_root<O: EncodeOptions>(&self, stream: &mut WriterStream<'_, O>) -> RootTypeId {
        let value = self.as_str();
        match value.len() {
            0 => RootTypeId::Str0,
            1 => {
                stream.bytes.push(value.as_bytes()[0]);
                RootTypeId::Str1
            }
            2 => {
                stream.bytes.extend_from_slice(value.as_bytes());
                RootTypeId::Str2
            }
            3 => {
                stream.bytes.extend_from_slice(value.as_bytes());
                RootTypeId::Str3
            }
            _ => {
                let b = value.as_bytes();
                encode_prefix_varint(b.len() as u64, stream.bytes);
                stream.bytes.extend_from_slice(b);
                RootTypeId::Str
            }
        }
    }
}

#[cfg(feature = "write")]
impl WriterArray<String> for Vec<&'static str> {
    fn buffer<'a, 'b: 'a>(&'a mut self, value: &'b String) {
        // TODO: Working around lifetime issues for lack of GAT
        // A quick check makes this appear to be sound, since the signature
        // requires that the value outlive self.
        //
        // The big safety problem is that whe then give these references
        // away when flushing. We happen to know that nothing saves the references,
        // but when things like threading come into play it's hard to know.
        self.push(unsafe { std::mem::transmute(value.as_str()) });
    }
    fn flush<O: EncodeOptions>(self, stream: &mut WriterStream<'_, O>) -> ArrayTypeId {
        profile!("WriterArray::flush");

        let compressors = (
            Utf8Compressor,
            RLE::new((Utf8Compressor,)),
            Dictionary::new((Utf8Compressor,)),
        );

        compress(&self, stream, &compressors)
    }
}

#[cfg(feature = "read")]
impl Readable for String {
    // TODO: Use lifetimes to make this read lazy rather than IntoIter
    type ReaderArray = IntoIter<String>;
    fn read(sticks: DynRootBranch<'_>, _options: &impl DecodeOptions) -> ReadResult<Self> {
        profile!("Readable::read");
        match sticks {
            DynRootBranch::String(s) => Ok(s.to_owned()),
            _ => Err(ReadError::SchemaMismatch),
        }
    }
}

#[cfg(feature = "read")]
impl InfallibleReaderArray for IntoIter<String> {
    type Read = String;

    fn new_infallible(sticks: DynArrayBranch<'_>, options: &impl DecodeOptions) -> ReadResult<Self> {
        profile!("ReaderArray::new");

        match sticks {
            DynArrayBranch::String(bytes) => {
                #[cfg(feature = "profile")]
                let _g = flame::start_guard("String");

                let strs = read_all(&bytes, |b, o| read_str(b, o).and_then(|v| Ok(v.to_owned())))?;
                Ok(strs.into_iter())
            }
            DynArrayBranch::RLE { runs, values } => {
                let rle = RleIterator::new(runs, values, options, |values| Self::new_infallible(values, options))?;
                let all = rle.collect::<Vec<_>>();
                Ok(all.into_iter())
            }
            DynArrayBranch::Dictionary { indices, values } => {
                let dict = DictionaryIterator::new(indices, values, options, |values| Self::new_infallible(values, options))?;
                let all = dict.collect::<Vec<_>>();
                Ok(all.into_iter())
            }
            _ => Err(ReadError::SchemaMismatch),
        }
    }
    fn read_next_infallible(&mut self) -> Self::Read {
        self.next().unwrap_or_default()
    }
}

#[cfg(feature = "write")]
impl<'a> Compressor<&'a str> for Utf8Compressor {
    fn fast_size_for(&self, data: &[&'a str]) -> Option<usize> {
        profile!("Compressor::fast_size_for");
        let mut total = 0;
        for s in data {
            total += size_for_varint(s.len() as u64);
            total += s.as_bytes().len();
        }
        Some(total)
    }
    fn compress<O: EncodeOptions>(&self, data: &[&'a str], stream: &mut WriterStream<'_, O>) -> Result<ArrayTypeId, ()> {
        profile!("Compressor::compress");

        stream.write_with_len(|stream| {
            for value in data.iter() {
                encode_prefix_varint(value.len() as u64, stream.bytes);
                stream.bytes.extend_from_slice(value.as_bytes());
            }
        });

        Ok(ArrayTypeId::Utf8)
    }
}