alex/beta-minecraft-rust

Minecraft Beta 1.7.3 Rewrite
root / crates / common / src / nbt.rs

//! Named Binary Tag, the format of `NBTBase` and its subclasses, with Java's modified UTF-8 strings.

use std::io::{self, Read, Write};

#[derive(Debug, thiserror::Error)]
pub enum NbtError {
    #[error(transparent)]
    Io(#[from] io::Error),
    #[error("unknown tag type {0}")]
    UnknownTag(u8),
    #[error("root tag must be a named compound tag")]
    BadRoot,
    #[error("invalid modified UTF-8 string")]
    Utf,
    #[error("negative length {0}")]
    Length(i32),
}

#[derive(Clone, Debug, PartialEq)]
pub enum Tag {
    Byte(i8),
    Short(i16),
    Int(i32),
    Long(i64),
    Float(f32),
    Double(f64),
    ByteArray(Vec<u8>),
    String(String),
    List(Vec<Tag>),
    Compound(Compound),
}

impl Tag {
    fn id(&self) -> u8 {
        match self {
            Self::Byte(_) => 1,
            Self::Short(_) => 2,
            Self::Int(_) => 3,
            Self::Long(_) => 4,
            Self::Float(_) => 5,
            Self::Double(_) => 6,
            Self::ByteArray(_) => 7,
            Self::String(_) => 8,
            Self::List(_) => 9,
            Self::Compound(_) => 10,
        }
    }

    fn write_payload(&self, out: &mut impl Write) -> io::Result<()> {
        match self {
            Self::Byte(v) => out.write_all(&v.to_be_bytes()),
            Self::Short(v) => out.write_all(&v.to_be_bytes()),
            Self::Int(v) => out.write_all(&v.to_be_bytes()),
            Self::Long(v) => out.write_all(&v.to_be_bytes()),
            Self::Float(v) => out.write_all(&v.to_be_bytes()),
            Self::Double(v) => out.write_all(&v.to_be_bytes()),
            Self::ByteArray(v) => {
                out.write_all(&(v.len() as i32).to_be_bytes())?;
                out.write_all(v)
            }
            Self::String(v) => write_utf(out, v),
            Self::List(items) => {
                // `NBTTagList` declares an empty list as a list of bytes.
                out.write_all(&[items.first().map_or(1, Tag::id)])?;
                out.write_all(&(items.len() as i32).to_be_bytes())?;
                items.iter().try_for_each(|item| item.write_payload(out))
            }
            Self::Compound(compound) => compound.write_payload(out),
        }
    }

    fn read_payload(id: u8, input: &mut impl Read) -> Result<Self, NbtError> {
        Ok(match id {
            1 => Self::Byte(i8::from_be_bytes(read_array(input)?)),
            2 => Self::Short(i16::from_be_bytes(read_array(input)?)),
            3 => Self::Int(i32::from_be_bytes(read_array(input)?)),
            4 => Self::Long(i64::from_be_bytes(read_array(input)?)),
            5 => Self::Float(f32::from_be_bytes(read_array(input)?)),
            6 => Self::Double(f64::from_be_bytes(read_array(input)?)),
            7 => {
                let mut bytes = vec![0; read_length(input)?];
                input.read_exact(&mut bytes)?;
                Self::ByteArray(bytes)
            }
            8 => Self::String(read_utf(input)?),
            9 => {
                let [element] = read_array(input)?;
                let len = read_length(input)?;
                Self::List(
                    (0..len)
                        .map(|_| Self::read_payload(element, input))
                        .collect::<Result<_, _>>()?,
                )
            }
            10 => {
                let mut compound = Compound::default();
                loop {
                    let [id] = read_array(input)?;
                    if id == 0 {
                        break;
                    }
                    let name = read_utf(input)?;
                    compound.insert(name, Self::read_payload(id, input)?);
                }
                Self::Compound(compound)
            }
            id => return Err(NbtError::UnknownTag(id)),
        })
    }
}

/// `NBTTagCompound`: named tags, kept in insertion order.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Compound(pub Vec<(String, Tag)>);

macro_rules! accessors {
    ($($get:ident $set:ident $variant:ident $ty:ty;)*) => {
        $(
            /// The value, or the original's default when missing or of another type.
            pub fn $get(&self, name: &str) -> $ty {
                match self.get(name) {
                    Some(Tag::$variant(v)) => *v,
                    _ => Default::default(),
                }
            }

            pub fn $set(&mut self, name: &str, value: $ty) {
                self.insert(name, Tag::$variant(value));
            }
        )*
    };
}

impl Compound {
    pub fn get(&self, name: &str) -> Option<&Tag> {
        self.0.iter().find(|(n, _)| n == name).map(|(_, tag)| tag)
    }

    pub fn contains(&self, name: &str) -> bool {
        self.get(name).is_some()
    }

    /// Replaces an existing tag of the same name in place, like `HashMap.put`.
    pub fn insert(&mut self, name: impl AsRef<str> + Into<String>, tag: Tag) {
        match self.0.iter_mut().find(|(n, _)| n == name.as_ref()) {
            Some(entry) => entry.1 = tag,
            None => self.0.push((name.into(), tag)),
        }
    }

    pub fn remove(&mut self, name: &str) -> Option<Tag> {
        let index = self.0.iter().position(|(n, _)| n == name)?;
        Some(self.0.remove(index).1)
    }

    accessors! {
        byte set_byte Byte i8;
        short set_short Short i16;
        int set_int Int i32;
        long set_long Long i64;
        float set_float Float f32;
        double set_double Double f64;
    }

    pub fn bool(&self, name: &str) -> bool {
        self.byte(name) != 0
    }

    pub fn set_bool(&mut self, name: &str, value: bool) {
        self.set_byte(name, i8::from(value));
    }

    pub fn string(&self, name: &str) -> &str {
        match self.get(name) {
            Some(Tag::String(v)) => v,
            _ => "",
        }
    }

    pub fn set_string(&mut self, name: &str, value: impl Into<String>) {
        self.insert(name, Tag::String(value.into()));
    }

    pub fn byte_array(&self, name: &str) -> &[u8] {
        match self.get(name) {
            Some(Tag::ByteArray(v)) => v,
            _ => &[],
        }
    }

    pub fn set_byte_array(&mut self, name: &str, value: impl Into<Vec<u8>>) {
        self.insert(name, Tag::ByteArray(value.into()));
    }

    pub fn compound(&self, name: &str) -> Option<&Compound> {
        match self.get(name) {
            Some(Tag::Compound(v)) => Some(v),
            _ => None,
        }
    }

    pub fn set_compound(&mut self, name: &str, value: Compound) {
        self.insert(name, Tag::Compound(value));
    }

    pub fn list(&self, name: &str) -> &[Tag] {
        match self.get(name) {
            Some(Tag::List(v)) => v,
            _ => &[],
        }
    }

    /// The compounds of a list, skipping elements of other types.
    pub fn compounds(&self, name: &str) -> impl Iterator<Item = &Compound> {
        self.list(name).iter().filter_map(|tag| match tag {
            Tag::Compound(c) => Some(c),
            _ => None,
        })
    }

    pub fn set_list(&mut self, name: &str, value: Vec<Tag>) {
        self.insert(name, Tag::List(value));
    }

    /// `CompressedStreamTools.read`: a named root compound. The root name is discarded.
    pub fn read(input: &mut impl Read) -> Result<Self, NbtError> {
        let [id] = read_array(input)?;
        if id != 10 {
            return Err(NbtError::BadRoot);
        }
        read_utf(input)?;
        match Tag::read_payload(id, input)? {
            Tag::Compound(compound) => Ok(compound),
            _ => Err(NbtError::BadRoot),
        }
    }

    /// `CompressedStreamTools.write` with the original's empty root name.
    pub fn write(&self, out: &mut impl Write) -> io::Result<()> {
        out.write_all(&[10])?;
        write_utf(out, "")?;
        self.write_payload(out)
    }

    fn write_payload(&self, out: &mut impl Write) -> io::Result<()> {
        for (name, tag) in &self.0 {
            out.write_all(&[tag.id()])?;
            write_utf(out, name)?;
            tag.write_payload(out)?;
        }
        out.write_all(&[0])
    }
}

fn read_array<const N: usize>(input: &mut impl Read) -> io::Result<[u8; N]> {
    let mut bytes = [0; N];
    input.read_exact(&mut bytes)?;
    Ok(bytes)
}

fn read_length(input: &mut impl Read) -> Result<usize, NbtError> {
    let len = i32::from_be_bytes(read_array(input)?);
    usize::try_from(len).map_err(|_| NbtError::Length(len))
}

/// `DataOutput.writeUTF`: modified UTF-8 with a 16-bit length, NUL as two bytes and
/// supplementary characters as surrogate pairs.
pub fn write_utf(out: &mut impl Write, text: &str) -> io::Result<()> {
    let mut bytes = Vec::with_capacity(text.len());
    for unit in text.encode_utf16() {
        match unit {
            0x0001..=0x007F => bytes.push(unit as u8),
            0x0000 | 0x0080..=0x07FF => {
                bytes.extend([0xC0 | (unit >> 6) as u8, 0x80 | (unit & 0x3F) as u8]);
            }
            _ => bytes.extend([
                0xE0 | (unit >> 12) as u8,
                0x80 | (unit >> 6 & 0x3F) as u8,
                0x80 | (unit & 0x3F) as u8,
            ]),
        }
    }
    let len =
        u16::try_from(bytes.len()).map_err(|_| io::Error::other("string too long for writeUTF"))?;
    out.write_all(&len.to_be_bytes())?;
    out.write_all(&bytes)
}

/// `DataInput.readUTF`.
pub fn read_utf(input: &mut impl Read) -> Result<String, NbtError> {
    let len = u16::from_be_bytes(read_array(input)?);
    let mut bytes = vec![0; usize::from(len)];
    input.read_exact(&mut bytes)?;
    let mut units = Vec::with_capacity(bytes.len());
    let mut i = 0;
    while i < bytes.len() {
        let b = u16::from(bytes[i]);
        let continuation = |j: usize| match bytes.get(j) {
            Some(&c) if c & 0xC0 == 0x80 => Ok(u16::from(c & 0x3F)),
            _ => Err(NbtError::Utf),
        };
        let (unit, width) = match b >> 4 {
            0..=7 => (b, 1),
            12 | 13 => ((b & 0x1F) << 6 | continuation(i + 1)?, 2),
            14 => (
                (b & 0x0F) << 12 | continuation(i + 1)? << 6 | continuation(i + 2)?,
                3,
            ),
            _ => return Err(NbtError::Utf),
        };
        units.push(unit);
        i += width;
    }
    Ok(String::from_utf16_lossy(&units))
}