overcast/src/vdf.rs

use binrw::{until, BinRead, BinWrite, NullString};
use std::{borrow::Cow, collections::HashMap, fmt, fs::File, path::Path};

#[derive(BinRead, BinWrite, Clone, PartialEq)]
#[brw(little)]
enum VdfBlock {
    #[brw(magic(0u8))]
    Dict(NullString, VdfDict),
    #[brw(magic(1u8))]
    Str(NullString, NullString),
    #[brw(magic(2u8))]
    Int(NullString, u32),
    #[brw(magic(8u8))]
    EndDict,
}

#[derive(BinRead, BinWrite, Clone, PartialEq)]
#[brw(little)]
struct VdfDict(
    #[br(parse_with(until(|block: &VdfBlock| block == &VdfBlock::EndDict)))] Vec<VdfBlock>,
);

#[derive(Clone)]
pub enum VdfValue {
    Dict(Dictionary),
    Str(String),
    Int(u32),
}

impl VdfValue {
    pub fn into_vec(self) -> Option<Vec<VdfValue>> {
        if let VdfValue::Dict(dict) = self {
            // lazily assuming this is an array of strings
            Some(dict.into_values().collect())
        } else {
            None
        }
    }
}

pub type Dictionary = HashMap<String, VdfValue>;

macro_rules! as_str {
    ($var:expr) => {{
        use std::borrow::Borrow;
        String::from_utf8_lossy($var.borrow()).to_string()
    }};
}

// === VdfValue Conversions ===

// Conversion error type for TryFrom
#[derive(Debug)]
pub enum VdfError {
    WrongType,
}

impl fmt::Display for VdfError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            VdfError::WrongType => write!(f, "Attempt to convert value to wrong type"),
        }
    }
}

impl std::error::Error for VdfError {}

// macros

macro_rules! impl_convert_to_vdfvalue {
    ($t:ty, $e:tt, from_only) => {
        impl From<$t> for VdfValue {
            fn from(value: $t) -> Self {
                VdfValue::$e(value.into())
            }
        }
    };

    ($t:ty, $e:tt) => {
        impl_convert_to_vdfvalue!($t, $e, from_only);

        impl TryFrom<VdfValue> for $t {
            type Error = VdfError;
            fn try_from(value: VdfValue) -> Result<Self, Self::Error> {
                if let VdfValue::$e(x) = value {
                    Ok(x)
                } else {
                    Err(VdfError::WrongType)
                }
            }
        }
    };
}

impl_convert_to_vdfvalue!(u32, Int);
impl_convert_to_vdfvalue!(String, Str);
impl_convert_to_vdfvalue!(&String, Str, from_only);
impl_convert_to_vdfvalue!(&str, Str, from_only);
impl_convert_to_vdfvalue!(bool, Int, from_only);
impl_convert_to_vdfvalue!(Dictionary, Dict);
impl_convert_to_vdfvalue!(Cow<'_, str>, Str, from_only);

// bool requires a special TryFrom<VdfValue>
impl TryFrom<VdfValue> for bool {
    type Error = VdfError;
    fn try_from(value: VdfValue) -> Result<Self, Self::Error> {
        if let VdfValue::Int(x) = value {
            Ok(!matches!(x, 0))
        } else {
            Err(VdfError::WrongType)
        }
    }
}

// Arrays also require special TryFrom
impl<T> TryFrom<VdfValue> for Vec<T>
where
    T: TryFrom<VdfValue>,
{
    type Error = VdfError;
    fn try_from(value: VdfValue) -> Result<Self, Self::Error> {
        if let VdfValue::Dict(dict) = value {
            // lazily assuming this is an array of strings
            let results: Result<Vec<T>, T::Error> = dict.into_values().map(T::try_from).collect();
            results.map_err(|_| VdfError::WrongType)
        } else {
            Err(VdfError::WrongType)
        }
    }
}

impl<T> From<Vec<T>> for VdfValue
where
    T: Into<VdfValue>,
{
    fn from(value: Vec<T>) -> Self {
        let dict: HashMap<String, VdfValue> = value
            .into_iter()
            .enumerate()
            .map(|(idx, val)| (idx.to_string(), val.into()))
            .collect();
        VdfValue::Dict(dict)
    }
}

/// === VdfDict (raw) to Dictionary Conversions ===

impl From<VdfDict> for Dictionary {
    fn from(value: VdfDict) -> Self {
        let map: Dictionary = value
            .0
            .iter()
            .filter_map(|block| match block.clone() {
                VdfBlock::Dict(k, v) => Some((as_str!(k), VdfValue::Dict(v.into()))),
                VdfBlock::Str(k, v) => Some((as_str!(k), VdfValue::Str(as_str!(v)))),
                VdfBlock::Int(k, v) => Some((as_str!(k), VdfValue::Int(v))),
                VdfBlock::EndDict => None,
            })
            .collect();
        map
    }
}

impl From<Dictionary> for VdfDict {
    fn from(value: Dictionary) -> Self {
        let mut dict: Vec<VdfBlock> = value
            .into_iter()
            .map(|ent| {
                let key = NullString::from(ent.0.as_str());
                match ent.1 {
                    VdfValue::Dict(v) => VdfBlock::Dict(key, v.into()),
                    VdfValue::Int(v) => VdfBlock::Int(key, v),
                    VdfValue::Str(v) => VdfBlock::Str(key, NullString::from(v.as_str())),
                }
            })
            .collect();
        dict.push(VdfBlock::EndDict);
        VdfDict(dict)
    }
}

// === Exported Functions ===

pub fn read_file<P: AsRef<Path>>(path: P) -> Result<Dictionary, Box<dyn std::error::Error>> {
    let mut fd = File::open(path)?;
    Ok(VdfDict::read(&mut fd)?.into())
}

pub fn write_file<P: AsRef<Path>>(
    dict: Dictionary,
    path: P,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut fd = File::create(path)?;
    VdfDict::from(dict).write(&mut fd)?;
    Ok(())
}