valence/crates/valence_anvil/src/to_valence.rs

use std::borrow::Cow;

use num_integer::{div_ceil, Integer};
use thiserror::Error;
use valence_biome::BiomeId;
use valence_block::{BlockEntityKind, BlockKind, PropName, PropValue};
use valence_core::ident::Ident;
use valence_instance::{BlockEntity, Chunk};
use valence_nbt::{Compound, List, Value};

#[derive(Clone, Debug, Error)]
#[non_exhaustive]
pub enum ToValenceError {
    #[error("missing chunk sections")]
    MissingSections,
    #[error("missing chunk section Y")]
    MissingSectionY,
    #[error("missing block states")]
    MissingBlockStates,
    #[error("missing block palette")]
    MissingBlockPalette,
    #[error("invalid block palette length")]
    BadBlockPaletteLen,
    #[error("missing block name in palette")]
    MissingBlockName,
    #[error("unknown block name of \"{0}\"")]
    UnknownBlockName(String),
    #[error("unknown property name of \"{0}\"")]
    UnknownPropName(String),
    #[error("property value of block is not a string")]
    BadPropValueType,
    #[error("unknown property value of \"{0}\"")]
    UnknownPropValue(String),
    #[error("missing packed block state data in section")]
    MissingBlockStateData,
    #[error("unexpected number of longs in block state data")]
    BadBlockLongCount,
    #[error("invalid block palette index")]
    BadBlockPaletteIndex,
    #[error("missing biomes")]
    MissingBiomes,
    #[error("missing biome palette")]
    MissingBiomePalette,
    #[error("invalid biome palette length")]
    BadBiomePaletteLen,
    #[error("biome name is not a valid resource identifier")]
    BadBiomeName,
    #[error("missing biome name")]
    MissingBiomeName,
    #[error("missing packed biome data in section")]
    MissingBiomeData,
    #[error("unexpected number of longs in biome data")]
    BadBiomeLongCount,
    #[error("invalid biome palette index")]
    BadBiomePaletteIndex,
    #[error("missing block entities")]
    MissingBlockEntity,
    #[error("missing block entity ident")]
    MissingBlockEntityIdent,
    #[error("invalid block entity ident of \"{0}\"")]
    UnknownBlockEntityIdent(String),
    #[error("invalid block entity position")]
    InvalidBlockEntityPosition,
}

/// Takes an Anvil chunk in NBT form and writes its data to a Valence [`Chunk`].
/// An error is returned if the NBT data does not match the expected structure
/// for an Anvil chunk.
///
/// # Arguments
///
/// - `nbt`: The Anvil chunk to read from. This is usually the value returned by
///   [`AnvilWorld::read_chunk`].
/// - `chunk`: The Valence chunk to write to.
/// - `sect_offset`: A constant to add to all sector Y positions in `nbt`. After
///   applying the offset, only the sectors in the range
///   `0..chunk.sector_count()` are written.
/// - `map_biome`: A function to map biome resource identifiers in the NBT data
///   to Valence [`BiomeId`]s.
///
/// [`AnvilWorld::read_chunk`]: crate::AnvilWorld::read_chunk
pub fn to_valence<F, const LOADED: bool>(
    nbt: &Compound,
    chunk: &mut Chunk<LOADED>,
    sect_offset: i32,
    mut map_biome: F,
) -> Result<(), ToValenceError>
where
    F: FnMut(Ident<&str>) -> BiomeId,
{
    let Some(Value::List(List::Compound(sections))) = nbt.get("sections") else {
        return Err(ToValenceError::MissingSections)
    };

    let mut converted_block_palette = vec![];
    let mut converted_biome_palette = vec![];

    for section in sections {
        let Some(Value::Byte(sect_y)) = section.get("Y") else {
            return Err(ToValenceError::MissingSectionY)
        };

        let adjusted_sect_y = *sect_y as i32 + sect_offset;

        if adjusted_sect_y < 0 || adjusted_sect_y as usize >= chunk.section_count() {
            // Section is out of bounds. Skip it.
            continue;
        }

        let Some(Value::Compound(block_states)) = section.get("block_states") else {
            return Err(ToValenceError::MissingBlockStates)
        };

        let Some(Value::List(List::Compound(palette))) = block_states.get("palette") else {
            return Err(ToValenceError::MissingBlockPalette)
        };

        if !(1..BLOCKS_PER_SECTION).contains(&palette.len()) {
            return Err(ToValenceError::BadBlockPaletteLen);
        }

        converted_block_palette.clear();

        for block in palette {
            let Some(Value::String(name)) = block.get("Name") else {
                return Err(ToValenceError::MissingBlockName)
            };

            let Some(block_kind) = BlockKind::from_str(ident_path(name)) else {
                return Err(ToValenceError::UnknownBlockName(name.into()))
            };

            let mut state = block_kind.to_state();

            if let Some(Value::Compound(properties)) = block.get("Properties") {
                for (key, value) in properties {
                    let Value::String(value) = value else {
                        return Err(ToValenceError::BadPropValueType)
                    };

                    let Some(prop_name) = PropName::from_str(key) else {
                        return Err(ToValenceError::UnknownPropName(key.into()))
                    };

                    let Some(prop_value) = PropValue::from_str(value) else {
                        return Err(ToValenceError::UnknownPropValue(value.into()))
                    };

                    state = state.set(prop_name, prop_value);
                }
            }

            converted_block_palette.push(state);
        }

        if converted_block_palette.len() == 1 {
            chunk.fill_block_states(adjusted_sect_y as usize, converted_block_palette[0]);
        } else {
            debug_assert!(converted_block_palette.len() > 1);

            let Some(Value::LongArray(data)) = block_states.get("data") else {
                return Err(ToValenceError::MissingBlockStateData)
            };

            let bits_per_idx = bit_width(converted_block_palette.len() - 1).max(4);
            let idxs_per_long = 64 / bits_per_idx;
            let long_count = div_ceil(BLOCKS_PER_SECTION, idxs_per_long);
            let mask = 2_u64.pow(bits_per_idx as u32) - 1;

            if long_count != data.len() {
                return Err(ToValenceError::BadBlockLongCount);
            };

            let mut i = 0;
            for &long in data.iter() {
                let u64 = long as u64;

                for j in 0..idxs_per_long {
                    if i >= BLOCKS_PER_SECTION {
                        break;
                    }

                    let idx = (u64 >> (bits_per_idx * j)) & mask;

                    let Some(block) = converted_block_palette.get(idx as usize).cloned() else {
                        return Err(ToValenceError::BadBlockPaletteIndex)
                    };

                    let x = i % 16;
                    let z = i / 16 % 16;
                    let y = i / (16 * 16);

                    chunk.set_block_state(x, adjusted_sect_y as usize * 16 + y, z, block);

                    i += 1;
                }
            }
        }

        let Some(Value::Compound(biomes)) = section.get("biomes") else {
            return Err(ToValenceError::MissingBiomes)
        };

        let Some(Value::List(List::String(palette))) = biomes.get("palette") else {
            return Err(ToValenceError::MissingBiomePalette)
        };

        if !(1..BIOMES_PER_SECTION).contains(&palette.len()) {
            return Err(ToValenceError::BadBiomePaletteLen);
        }

        converted_biome_palette.clear();

        for biome_name in palette {
            let Ok(ident) = Ident::<Cow<str>>::new(biome_name) else {
                return Err(ToValenceError::BadBiomeName)
            };

            converted_biome_palette.push(map_biome(ident.as_str_ident()));
        }

        if converted_biome_palette.len() == 1 {
            chunk.fill_biomes(adjusted_sect_y as usize, converted_biome_palette[0]);
        } else {
            debug_assert!(converted_biome_palette.len() > 1);

            let Some(Value::LongArray(data)) = biomes.get("data") else {
                return Err(ToValenceError::MissingBiomeData)
            };

            let bits_per_idx = bit_width(converted_biome_palette.len() - 1);
            let idxs_per_long = 64 / bits_per_idx;
            let long_count = div_ceil(BIOMES_PER_SECTION, idxs_per_long);
            let mask = 2_u64.pow(bits_per_idx as u32) - 1;

            if long_count != data.len() {
                return Err(ToValenceError::BadBiomeLongCount);
            };

            let mut i = 0;
            for &long in data.iter() {
                let u64 = long as u64;

                for j in 0..idxs_per_long {
                    if i >= BIOMES_PER_SECTION {
                        break;
                    }

                    let idx = (u64 >> (bits_per_idx * j)) & mask;

                    let Some(biome) = converted_biome_palette.get(idx as usize).cloned() else {
                        return Err(ToValenceError::BadBiomePaletteIndex)
                    };

                    let x = i % 4;
                    let z = i / 4 % 4;
                    let y = i / (4 * 4);

                    chunk.set_biome(x, adjusted_sect_y as usize * 4 + y, z, biome);

                    i += 1;
                }
            }
        }
    }

    let Some(Value::List(block_entities)) = nbt.get("block_entities") else {
        return Err(ToValenceError::MissingBlockEntity);
    };

    if let List::Compound(block_entities) = block_entities {
        for comp in block_entities {
            let Some(Value::String(ident)) = comp.get("id") else {
                return Err(ToValenceError::MissingBlockEntityIdent);
            };
            let Ok(ident) = Ident::new(&ident[..]) else {
                return Err(ToValenceError::UnknownBlockEntityIdent(ident.clone()));
            };
            let Some(kind) = BlockEntityKind::from_ident(ident.as_str_ident()) else {
                return Err(ToValenceError::UnknownBlockEntityIdent(ident.as_str().to_string()));
            };
            let block_entity = BlockEntity {
                kind,
                nbt: comp.clone(),
            };
            let Some(Value::Int(x)) = comp.get("x") else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };
            let Ok(x) = usize::try_from(x.mod_floor(&16)) else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };
            let Some(Value::Int(y)) = comp.get("y") else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };
            let Ok(y) = usize::try_from(y + sect_offset * 16) else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };
            let Some(Value::Int(z)) = comp.get("z") else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };
            let Ok(z) = usize::try_from(z.mod_floor(&16)) else {
                return Err(ToValenceError::InvalidBlockEntityPosition);
            };

            chunk.set_block_entity(x, y, z, block_entity);
        }
    }

    Ok(())
}

const BLOCKS_PER_SECTION: usize = 16 * 16 * 16;
const BIOMES_PER_SECTION: usize = 4 * 4 * 4;

/// Gets the path part of a resource identifier.
fn ident_path(ident: &str) -> &str {
    match ident.rsplit_once(':') {
        Some((_, after)) => after,
        None => ident,
    }
}

/// Returns the minimum number of bits needed to represent the integer `n`.
const fn bit_width(n: usize) -> usize {
    (usize::BITS - n.leading_zeros()) as _
}