b3-rs 0.1.0

//! Item key encoding/decoding.

use crate::alloc_prelude::*;
use nano_leb128::ULEB128;

// Key type bits in the item control byte:
// (largely copied from b3/item_header.py)
//
// +------------+------------+
// | key type   | key type   |
// +------------+------------+
//   0     0  (0x0)  no key
//   0     1  (0x4)  UVarInt
//   1     0  (0x8)  UTF8 bytes
//   1     1  (0xC)  raw bytes

/// An item key.
#[derive(Debug, Clone, PartialEq)]
pub enum ItemKey {
    /// No key is set for this item.
    NoKey,

    /// The key is a `UVarInt` (unsigned LEB128).
    IntegerKey(u64),

    /// The key is a UTF-8 encoded string.
    ///
    /// The key is encoded as the length of the UTF-8 string in bytes as a
    /// `UVarInt` (unsigned LEB128), and then the bytes of the string.
    StringKey(String),

    /// The key is a byte array.
    ///
    /// The key is encoded as the length of the byte array as a `UVarInt`
    /// (unsigned LEB128), and then the bytes.
    BytesKey(Vec<u8>),
}

impl ItemKey {
    /// Return the bits that would be set in the item control byte for the
    /// key type.
    pub fn type_bits(&self) -> u8 {
        match self {
            Self::NoKey => 0b00000000,
            Self::IntegerKey(_) => 0b00010000,
            Self::StringKey(_) => 0b00100000,
            Self::BytesKey(_) => 0b00110000,
        }
    }

    /// Encode the key into it's byte representation.
    pub fn encode(&self) -> Result<Vec<u8>, crate::Error> {
        let mut out: Vec<u8> = Vec::new();
        let mut len_bytes = [0u8; 10];

        match self {
            // Empty vector if there's no key
            Self::NoKey => {}

            // Unsigned LEB128 for an integer key
            Self::IntegerKey(i) => {
                let count = ULEB128::from(*i).write_into(&mut len_bytes)?;
                out.extend(&len_bytes[0..count]);
            }

            // String key - unsigned LEB128 of byte length, then the bytes
            Self::StringKey(s) => {
                let b = s.as_bytes();
                let count = ULEB128::from(b.len() as u64).write_into(&mut len_bytes)?;
                out.extend(&len_bytes[0..count]);
                out.extend(b);
            }

            // Byte key - unsigned LEB128 of length, then the bytes
            Self::BytesKey(b) => {
                let count = ULEB128::from(b.len() as u64).write_into(&mut len_bytes)?;
                out.extend(&len_bytes[0..count]);
                out.extend(b);
            }
        }

        Ok(out)
    }

    pub fn decode(base: u8, ext: &[u8]) -> Result<(Self, usize), crate::Error> {
        match base & 0b00110000 {
            // No key
            0b00000000 => Ok((Self::NoKey, 0)),

            // Integer key
            0b00010000 => {
                let (val, len) = ULEB128::read_from(&ext)?;
                let val = u64::from(val);

                Ok((Self::IntegerKey(val), len))
            }

            // String key
            0b00100000 => {
                let (val, len) = ULEB128::read_from(&ext)?;
                let val = u64::from(val);
                let total = (val as usize) + len;

                let s = String::from_utf8(Vec::from(&ext[len..total]))?;
                Ok((Self::StringKey(s), total))
            }

            // Bytes key
            0b00110000 => {
                let (val, len) = ULEB128::read_from(&ext)?;
                let val = u64::from(val);
                let total = (val as usize) + len;

                let b = Vec::from(&ext[len..total]);
                Ok((Self::BytesKey(b), total))
            }

            _ => unreachable!(),
        }
    }
}

impl From<u64> for ItemKey {
    fn from(i: u64) -> ItemKey {
        ItemKey::IntegerKey(i)
    }
}

impl From<&str> for ItemKey {
    fn from(s: &str) -> ItemKey {
        ItemKey::StringKey(String::from(s))
    }
}

impl From<String> for ItemKey {
    fn from(s: String) -> ItemKey {
        ItemKey::StringKey(s)
    }
}

impl From<Vec<u8>> for ItemKey {
    fn from(v: Vec<u8>) -> ItemKey {
        ItemKey::BytesKey(v)
    }
}

impl From<&[u8]> for ItemKey {
    fn from(v: &[u8]) -> ItemKey {
        ItemKey::BytesKey(Vec::from(v))
    }
}

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

    #[test]
    fn encode_key_empty() {
        let key = ItemKey::NoKey;
        assert_eq!(Vec::<u8>::new(), key.encode().unwrap())
    }

    #[test]
    fn encode_key_integer() {
        let key = ItemKey::IntegerKey(624485);
        assert_eq!(vec![0xE5, 0x8E, 0x26], key.encode().unwrap())
    }

    #[test]
    fn encode_key_string() {
        let key = ItemKey::StringKey(String::from("A"));
        assert_eq!(vec![0x01, 0x41], key.encode().unwrap())
    }

    #[test]
    fn encode_key_bytes() {
        let key = ItemKey::BytesKey(vec![0xC0, 0xFF, 0xEE]);
        assert_eq!(vec![0x03, 0xC0, 0xFF, 0xEE], key.encode().unwrap())
    }

    #[test]
    fn decode_key_empty() {
        let (key, len) = ItemKey::decode(0b00000000, &[]).unwrap();
        assert_eq!(key, ItemKey::NoKey);
        assert_eq!(len, 0);
    }

    #[test]
    fn decode_key_integer() {
        let (key, len) = ItemKey::decode(0b00010000, &[0x01, 0x00]).unwrap();
        assert_eq!(key, ItemKey::IntegerKey(1));
        assert_eq!(len, 1);

        let (key, len) = ItemKey::decode(0b00010000, &[0xE5, 0x8E, 0x26, 0x00]).unwrap();
        assert_eq!(key, ItemKey::IntegerKey(624485));
        assert_eq!(len, 3);
    }

    #[test]
    fn decode_key_string() {
        let (key, len) = ItemKey::decode(0b00100000, &[0x03, 0x41, 0x41, 0x41]).unwrap();
        assert_eq!(key, ItemKey::StringKey(String::from("AAA")));
        assert_eq!(len, 4);
    }

    #[test]
    fn decode_key_bytes() {
        let (key, len) = ItemKey::decode(0b00110000, &[0x03, 0xC0, 0xFF, 0xEE]).unwrap();
        assert_eq!(key, ItemKey::BytesKey(vec![0xC0, 0xFF, 0xEE]));
        assert_eq!(len, 4);
    }
}