wowlab_loadout/

codec.rs

use serde::Serialize;

wowlab_engine_macros::define_error! {
/// Errors produced while decoding a trait loadout string.
#[derive(Debug)]
pub struct TraitError {
    #[source]
    kind: TraitErrorKind,
}

#[derive(Debug, thiserror::Error)]
enum TraitErrorKind {
    #[error("Invalid characters in trait string")]
    InvalidCharacters,

    #[error("Trait string too short")]
    TooShort,

    #[error("Invalid bit read: not enough data")]
    NotEnoughData {
        #[source]
        source: Option<std::num::TryFromIntError>,
    },
}
}

impl TraitError {
    const fn invalid_characters() -> Self {
        Self {
            kind: TraitErrorKind::InvalidCharacters,
        }
    }

    const fn too_short() -> Self {
        Self {
            kind: TraitErrorKind::TooShort,
        }
    }

    const fn not_enough_data() -> Self {
        Self {
            kind: TraitErrorKind::NotEnoughData { source: None },
        }
    }

    /// Returns whether the input contains an unsupported character.
    #[must_use]
    pub const fn is_invalid_characters(&self) -> bool {
        matches!(self.kind, TraitErrorKind::InvalidCharacters)
    }

    /// Returns whether the input was empty.
    #[must_use]
    pub const fn is_too_short(&self) -> bool {
        matches!(self.kind, TraitErrorKind::TooShort)
    }

    /// Returns whether the encoded input ended in the middle of its header or a node.
    #[must_use]
    pub const fn is_not_enough_data(&self) -> bool {
        matches!(self.kind, TraitErrorKind::NotEnoughData { .. })
    }
}

impl From<std::num::TryFromIntError> for TraitError {
    fn from(source: std::num::TryFromIntError) -> Self {
        Self {
            kind: TraitErrorKind::NotEnoughData {
                source: Some(source),
            },
        }
    }
}

const BITS_PER_BASE64_CHAR: usize = 6;
const VERSION_BITS: usize = 8;
const SPEC_ID_BITS: usize = 16;
const TREE_HASH_LEN: usize = 16;
const RANK_BITS: usize = 6;
const CHOICE_INDEX_BITS: usize = 2;
const MIN_NODE_BITS: usize = 2;
const INVALID_BASE64_CHAR: u8 = 255;

/// A decoded trait-loadout header and its position-ordered node selections.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct DecodedTraitLoadout {
    pub version: u8,
    pub spec_id: u16,
    pub tree_hash: [u8; TREE_HASH_LEN],
    pub nodes: Vec<DecodedTraitNode>,
}

/// Selection state for one positional talent node.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[expect(
    clippy::struct_excessive_bools,
    reason = "flags map one-to-one to the stable client bitstream and serialized API shape"
)]
pub struct DecodedTraitNode {
    pub selected: bool,
    pub purchased: bool,
    pub partially_ranked: bool,
    pub ranks_purchased: Option<u8>,
    pub choice_node: bool,
    pub choice_index: Option<u8>,
}

const BASE64_URL_SAFE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

const BASE64_LOOKUP_SIZE: usize = 256;

const BASE64_CHAR_MAP: [u8; BASE64_LOOKUP_SIZE] = build_char_map();

const fn build_char_map() -> [u8; BASE64_LOOKUP_SIZE] {
    let standard = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let url_safe = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
    let mut map = [INVALID_BASE64_CHAR; BASE64_LOOKUP_SIZE];
    let mut i = 0u8;

    while (i as usize) < standard.len() {
        // #t(block: rust_unchecked_indexing) base64 index 0-63, fits in u8
        map[standard[i as usize] as usize] = i;
        i += 1;
    }

    i = 0;

    while (i as usize) < url_safe.len() {
        // #t(block: rust_unchecked_indexing) base64 index 0-63, fits in u8
        map[url_safe[i as usize] as usize] = i;
        i += 1;
    }

    map
}

struct BitReader<'a> {
    data: &'a str,
    position: usize,
    total_bits: usize,
}

impl<'a> BitReader<'a> {
    fn new(data: &'a str) -> Self {
        Self {
            data,
            position: 0,
            total_bits: data.len() * BITS_PER_BASE64_CHAR,
        }
    }

    fn read(&mut self, bit_count: usize) -> Result<u32, TraitError> {
        if self.position + bit_count > self.total_bits {
            return Err(TraitError::not_enough_data());
        }

        let mut value = 0u32;

        for i in 0..bit_count {
            let char_idx = self.position / BITS_PER_BASE64_CHAR;
            let bit_offset = self.position % BITS_PER_BASE64_CHAR;

            let byte = self.data.as_bytes().get(char_idx).copied().unwrap_or(0);
            // #t(rust_unchecked_indexing) byte is u8 and table has 256 entries, always in range
            let char_value = BASE64_CHAR_MAP[byte as usize];

            if char_value == INVALID_BASE64_CHAR {
                return Err(TraitError::invalid_characters());
            }

            let bit = (char_value >> bit_offset) & 1;

            value |= u32::from(bit) << i;
            self.position += 1;
        }

        Ok(value)
    }

    fn has_remaining(&self, bit_count: usize) -> bool {
        self.position + bit_count <= self.total_bits
    }
}

struct BitWriter {
    output: String,
    bit_position: usize,
    current_char: u8,
}

impl BitWriter {
    fn new() -> Self {
        Self {
            output: String::new(),
            bit_position: 0,
            current_char: 0,
        }
    }

    fn write(&mut self, value: u32, bit_count: usize) {
        for i in 0..bit_count {
            // #t(rust_lossy_cast) masked to single bit, value is 0 or 1
            let bit = ((value >> i) & 1) as u8;

            self.current_char |= bit << self.bit_position;
            self.bit_position += 1;

            if self.bit_position == BITS_PER_BASE64_CHAR {
                self.flush_char();
            }
        }
    }

    fn flush_char(&mut self) {
        // #t(block: rust_unchecked_indexing) current_char is max 63 from 6-bit accumulation, table has 64 entries
        self.output
            .push(BASE64_URL_SAFE[self.current_char as usize] as char);
        self.current_char = 0;
        self.bit_position = 0;
    }

    fn finish(mut self) -> String {
        if self.bit_position > 0 {
            self.flush_char();
        }

        self.output
    }
}

/// Decode a trait loadout string into structured data.
///
/// # Errors
///
/// Returns an error for invalid base64 characters, incomplete headers, or truncated node data.
pub fn decode_trait_loadout(talent_string: &str) -> Result<DecodedTraitLoadout, TraitError> {
    if talent_string.is_empty() {
        return Err(TraitError::too_short());
    }

    validate_base64_chars(talent_string)?;

    let mut reader = BitReader::new(talent_string);

    let version = read_u8(&mut reader, VERSION_BITS)?;
    let spec_id = read_u16(&mut reader, SPEC_ID_BITS)?;

    let mut tree_hash = [0u8; TREE_HASH_LEN];

    for byte in &mut tree_hash {
        *byte = read_u8(&mut reader, VERSION_BITS)?;
    }

    let remaining_bits = reader.total_bits.saturating_sub(reader.position);
    let mut nodes = Vec::with_capacity(remaining_bits / MIN_NODE_BITS);

    while reader.has_remaining(1) {
        nodes.push(decode_node(&mut reader)?);
    }

    Ok(DecodedTraitLoadout {
        version,
        spec_id,
        tree_hash,
        nodes,
    })
}

fn read_u8(reader: &mut BitReader<'_>, bit_count: usize) -> Result<u8, TraitError> {
    Ok(u8::try_from(reader.read(bit_count)?)?)
}

fn read_u16(reader: &mut BitReader<'_>, bit_count: usize) -> Result<u16, TraitError> {
    Ok(u16::try_from(reader.read(bit_count)?)?)
}

fn validate_base64_chars(s: &str) -> Result<(), TraitError> {
    for c in s.chars() {
        let valid = c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '-' | '_');

        if !valid {
            return Err(TraitError::invalid_characters());
        }
    }

    Ok(())
}

fn decode_node(reader: &mut BitReader) -> Result<DecodedTraitNode, TraitError> {
    let selected = reader.read(1)? == 1;

    if !selected {
        return Ok(DecodedTraitNode {
            selected: false,
            purchased: false,
            partially_ranked: false,
            ranks_purchased: None,
            choice_node: false,
            choice_index: None,
        });
    }

    let purchased = reader.read(1)? == 1;

    if !purchased {
        return Ok(DecodedTraitNode {
            selected: true,
            purchased: false,
            partially_ranked: false,
            ranks_purchased: None,
            choice_node: false,
            choice_index: None,
        });
    }

    let partially_ranked = reader.read(1)? == 1;
    let ranks_purchased = if partially_ranked {
        Some(read_u8(reader, RANK_BITS)?)
    } else {
        None
    };

    let choice_node = reader.read(1)? == 1;
    let choice_index = if choice_node {
        Some(read_u8(reader, CHOICE_INDEX_BITS)?)
    } else {
        None
    };

    Ok(DecodedTraitNode {
        selected: true,
        purchased: true,
        partially_ranked,
        ranks_purchased,
        choice_node,
        choice_index,
    })
}

/// Creates the version-1 loadout header for a spec with a zero tree hash and no explicit nodes.
#[must_use]
pub fn encode_minimal_loadout(spec_id: u16) -> String {
    let loadout = DecodedTraitLoadout {
        version: 1,
        spec_id,
        tree_hash: [0u8; TREE_HASH_LEN],
        nodes: Vec::new(),
    };

    encode_trait_loadout(&loadout)
}

/// Encodes a trait loadout with the URL-safe Base64 alphabet.
#[must_use]
pub fn encode_trait_loadout(loadout: &DecodedTraitLoadout) -> String {
    let mut writer = BitWriter::new();

    writer.write(u32::from(loadout.version), VERSION_BITS);
    writer.write(u32::from(loadout.spec_id), SPEC_ID_BITS);

    for &byte in &loadout.tree_hash {
        writer.write(u32::from(byte), VERSION_BITS);
    }

    for node in &loadout.nodes {
        encode_node(&mut writer, node);
    }

    writer.finish()
}

fn encode_node(writer: &mut BitWriter, node: &DecodedTraitNode) {
    writer.write(u32::from(node.selected), 1);

    if !node.selected {
        return;
    }

    writer.write(u32::from(node.purchased), 1);

    if !node.purchased {
        return;
    }

    writer.write(u32::from(node.partially_ranked), 1);

    if node.partially_ranked {
        writer.write(u32::from(node.ranks_purchased.unwrap_or(0)), RANK_BITS);
    }

    writer.write(u32::from(node.choice_node), 1);

    if node.choice_node {
        writer.write(u32::from(node.choice_index.unwrap_or(0)), CHOICE_INDEX_BITS);
    }
}

#[cfg(test)]
#[path = "codec_tests.rs"]
mod tests;