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

const MSB: u8 = 0b10000000;
const DROP_MSB: u8 = 0b01111111;


#[inline]
fn required_encoded_space_unsigned(mut v: u64) -> usize {
    if v == 0 {
        return 1;
    }
    let mut c = 0;
    while v > 0 {
        c += 1;
        v >>= 7;
    }
    c
}

/// Varint trait
pub trait EncodeVar: Sized + Copy {
    fn required_size(self) -> usize;
    fn real_encode(self, dst: &mut [u8]);
    fn encode_varint(self) -> Vec<u8> {
        let mut v = Vec::new();
        v.resize(self.required_size(), 0);
        self.real_encode(&mut v);
        v
    }
    fn decode_varint(src: &[u8]) -> (Self, usize);
}

impl EncodeVar for u64 {
    fn required_size(self) -> usize {
        required_encoded_space_unsigned(self)
    }
    fn real_encode(self, v: &mut [u8]) {
        assert!(v.len() >= self.required_size());
        let mut i = 0;
        let mut t = self;
        while t >= MSB as u64 {
            v[i] = t as u8 & (MSB - 1) | MSB;
            i += 1;
            t >>= 7;
        }
        v[i] = t as u8;
    }
    fn decode_varint(src: &[u8]) -> (Self, usize) {
        let mut result: u64 = 0;
        let mut shift = 0;

        for b in src.iter() {
            let msb_dropped = b & DROP_MSB;
            result |= (msb_dropped as u64) << shift;
            shift += 7;
            if b & MSB == 0 || shift > (10 * 7) {
                break;
            }
        }
        (result, (shift / 7) as usize)
    }
}
/// todo signed
macro_rules! impl_varint {
    ($t:ty, unsigned) => {
        impl EncodeVar for $t {
            fn required_size(self) -> usize {
                required_encoded_space_unsigned(self as u64)
            }
            fn real_encode(self, v: &mut [u8]) {
                (self as u64).real_encode(v);
            }
            fn encode_varint(self) -> Vec<u8> {
                (self as u64).encode_varint()
            }
            fn decode_varint(src: &[u8]) -> (Self, usize) {
                let (n, s) = u64::decode_varint(src);
                (n as Self, s)
            }
        }
    };
}

impl_varint!(usize, unsigned);
impl_varint!(u32, unsigned);
impl_varint!(u16, unsigned);
impl_varint!(u8, unsigned);

#[cfg(test)]
mod test {
    use super::EncodeVar;

    #[test]
    fn test_var_u32() {
        let mut value = 1u32;
        while value < std::u32::MAX / 10 {
            value *= 10;
//            println!("{}", value);
            let v = value.encode_varint();
            let (result, _) = u32::decode_varint(&v);
            assert_eq!(result, value);
        }
    }

    #[test]
    fn test_var_u64() {
        let mut value = std::u32::MAX as u64;
        while value < std::u64::MAX / 10 {
            value *= 10;
//            println!("{}", value);
            let v = value.encode_varint();
            let (result, _) = u64::decode_varint(&v);
            assert_eq!(result, value);
        }
    }
}