twmap 0.15.0

Parse, edit and save Teeworlds and DDNet maps
Documentation
use az::{CheckedAs, UnwrappedAs};
use fixed::types::{I17F15, I22F10};
use image::RgbaImage;
use std::collections::{HashMap, HashSet};
use std::convert::TryInto;
use std::mem;

use crate::compression::compress;
use crate::datafile::RawDatafile;
use crate::map::parse;
use crate::{
    BezierCurve, CurveKind, EmbeddedImage, Env, EnvPoint, Envelope, Error, Image, Layer, LayerKind,
    Position, TwMap, Version, Volume,
};

const SANITIZE_OPTIONS: sanitize_filename::Options = sanitize_filename::Options {
    windows: true,
    truncate: true,
    replacement: "",
};
const OPT_U16_RANGE: std::ops::Range<i32> = -1..u16::MAX as i32;

impl TwMap {
    // loss_fix, and load unknown external mapres from fs.
    pub fn lossy_fix_with_fs(data: &[u8]) -> Result<Vec<u8>, Error> {
        fn load_mapres_from_fs(name: &str, version: Version) -> Option<RgbaImage> {
            let file = twstorage::open_file(format!("mapres/{name}.png"), version.into()).ok()?;
            let reader = image::ImageReader::with_format(
                std::io::BufReader::new(file),
                image::ImageFormat::Png,
            );
            Some(reader.decode().ok()?.into_rgba8())
        }
        Self::lossy_fix(data, load_mapres_from_fs)
    }

    /// Best effort attempt to fix a map.
    /// This does neither guarantee a fixed map, nor unintended side effects.
    pub fn lossy_fix<F: Fn(&str, Version) -> Option<RgbaImage>>(
        mut data: &[u8],
        load_nonstandard_external_mapres: F,
    ) -> Result<Vec<u8>, Error> {
        // Fix extra data at the end
        if data.len() > 36 && RawDatafile::parse(data).is_err() {
            let size = get_i32(data, 2).wrapping_add(16);
            if let Some(expected_size) = size.checked_as::<usize>() {
                if expected_size < data.len() {
                    data = &data[..expected_size];
                }
            }
        }
        let mut df = RawDatafile::parse(data)?.to_datafile();
        let ex_index =
            <parse::ExType as parse::ItemParseErrorTrait>::parse_all(&df, &HashMap::new())?;
        let ex_index: HashMap<[u8; 16], u16> = ex_index
            .into_iter()
            .map(|parse::ExType { uuid, type_id }| (uuid, type_id))
            .collect();
        let mut bezier_data_exists = ex_index
            .contains_key(&item_type_uuid(parse::ItemType::EnvPointsBezierData))
            || df
                .items
                .get(&item_type_id(parse::ItemType::Envelope))
                .is_some_and(|items| {
                    items
                        .iter()
                        .any(|i| i.item_data.first().is_some_and(|v| *v > 3))
                });

        let mut items = HashMap::new();
        mem::swap(&mut items, &mut df.items);
        for (id, items) in &mut items {
            let Some(item_type) = detect_item_type(*id, &ex_index) else {
                continue;
            };
            for item in items {
                match item_type {
                    parse::ItemType::Version => {
                        if let Some(v) = item.item_data.first_mut() {
                            *v = 1;
                        }
                    }
                    parse::ItemType::Image => {
                        let Some(name_data_index) = item.item_data.get_mut(4) else {
                            continue;
                        };
                        if *name_data_index == -1 {
                            *name_data_index = df.data_items.len().unwrapped_as();
                            let name = b"\0";
                            df.data_items.push((compress(name).into(), 1))
                        }
                    }
                    parse::ItemType::Envelope => {
                        if item.item_data.first().is_some_and(|v| *v >= 4) {
                            bezier_data_exists = true;
                        }
                        // Set envelopes to color to preserve all values
                        if let Some(env_type) = item.item_data.get_mut(1) {
                            *env_type = 4;
                        }
                    }
                    parse::ItemType::EnvPoints => {
                        // Set all curves to non-bezier if no bezier data given
                        if !bezier_data_exists {
                            item.item_data.chunks_mut(6).for_each(|d| {
                                if d[1] == 5 {
                                    d[1] = 4;
                                }
                            });
                        }
                    }
                    parse::ItemType::AutoMapperConfig => {
                        // Ensure automapper config index is in bounds
                        if let Some(config) = item.item_data.get_mut(3) {
                            if !OPT_U16_RANGE.contains(config) {
                                *config = -1;
                            }
                        }
                    }
                    parse::ItemType::Layer => {
                        if item.layer_kind() == LayerKind::Quads {
                            // Fix quads envelope indices
                            let Some(data_index) = item.item_data.get(5) else {
                                continue;
                            };
                            let Ok(mut data) = df.decompressed_data_item(*data_index) else {
                                continue;
                            };
                            data.chunks_mut(mem::size_of::<parse::BinaryQuad>())
                                .for_each(|q| {
                                    // Fix position envelope index
                                    if !OPT_U16_RANGE.contains(&get_i32(q, 34)) {
                                        set_i32(q, 34, -1);
                                    }
                                    // Fix color envelope index
                                    if !OPT_U16_RANGE.contains(&get_i32(q, 36)) {
                                        set_i32(q, 36, -1);
                                    }
                                });
                            let decompressed_len = data.len();
                            df.data_items[*data_index as usize] =
                                (compress(&data).into(), decompressed_len);
                        }
                    }
                    _ => {}
                }
            }
        }
        mem::swap(&mut items, &mut df.items);

        let mut map = TwMap::parse_datafile_unchecked(&df)?;

        for image in &mut map.images {
            // Ensure sanitized image filename
            *image.name_mut() =
                sanitize_filename::sanitize_with_options(image.name(), SANITIZE_OPTIONS);
            let size = image.size();
            let fallback_image = Image::Embedded(EmbeddedImage {
                name: image.name().clone(),
                image: RgbaImage::from_pixel(16, 16, [255; 4].into()).into(),
            });
            match image {
                Image::External(ex) => {
                    if crate::constants::is_external_name(&ex.name, map.version) {
                        continue;
                    }
                    // Here we fix external images with non-default mapres names
                    if let Some(mapres) = load_nonstandard_external_mapres(&ex.name, map.version) {
                        *image = Image::Embedded(EmbeddedImage {
                            name: ex.name.clone(),
                            image: mapres.into(),
                        });
                        continue;
                    }
                    // Otherwise just remove them
                    *image = fallback_image;
                }
                Image::Embedded(_) => {
                    // Remove too small
                    if size.w < 2 || size.h < 2 {
                        *image = fallback_image;
                    }
                }
            }
        }
        for group in &mut map.groups {
            group.layers.retain(|l| !matches!(l, Layer::Invalid(_)));
        }

        // Map envelopes to their correct kind
        let envelopes = map.envelopes.clone();
        let original_envelope_count = envelopes.len();
        // Remove envelope indices that are out-of-bounds
        map.edit_env_indices(|i| i.filter(|&i| (i as usize) < original_envelope_count));
        map.envelopes.extend(envelopes.iter().map(|e| {
            let Envelope::Color(c) = e else {
                // We have turned all envelopes to color beforehand.
                unreachable!()
            };
            Envelope::Position(conv_env(c, &color_to_position))
        }));
        map.envelopes.extend(envelopes.iter().map(|e| {
            let Envelope::Color(c) = e else {
                // We have turned all envelopes to color beforehand.
                unreachable!()
            };
            Envelope::Sound(conv_env(c, &color_to_volume))
        }));
        fn checked_offset(i: &mut Option<u16>, offset: usize, factor: usize) -> Option<()> {
            let offset = offset.checked_mul(factor)?;
            let offset: u16 = offset.try_into().ok()?;
            if let Some(i) = i.as_mut() {
                *i = i.checked_add(offset)?;
            }
            Some(())
        }
        for q in map
            .groups
            .iter_mut()
            .flat_map(|g| g.layers.iter_mut())
            .filter_map(|l| {
                if let Layer::Quads(q) = l {
                    Some(q)
                } else {
                    None
                }
            })
            .flat_map(|q| q.quads.iter_mut())
        {
            checked_offset(&mut q.position_env, original_envelope_count, 1).unwrap();
        }
        for s in map
            .groups
            .iter_mut()
            .flat_map(|g| g.layers.iter_mut())
            .filter_map(|l| {
                if let Layer::Sounds(s) = l {
                    Some(s)
                } else {
                    None
                }
            })
            .flat_map(|s| s.sources.iter_mut())
        {
            checked_offset(&mut s.position_env, original_envelope_count, 1).unwrap();
            checked_offset(&mut s.sound_env, original_envelope_count, 2).unwrap();
        }
        map.remove_unused_envelopes();

        map.load_unchecked()?;
        // Fix modified padding bytes / skip field
        map.edit_tiles::<crate::edit::ZeroUnusedParts>();

        let tilemap_image_indices: HashSet<u16> = map
            .groups
            .iter_mut()
            .flat_map(|g| g.layers.iter_mut())
            .filter_map(|l| {
                if let Layer::Tiles(t) = l {
                    Some(t)
                } else {
                    None
                }
            })
            .filter_map(|l| l.image)
            .collect();
        for index in tilemap_image_indices {
            let Some(image) = map.images.get_mut(index as usize) else {
                continue;
            };
            let Image::Embedded(emb) = image else {
                continue;
            };
            let rgba = emb.image.unwrap_mut();
            if rgba.width().is_multiple_of(16) && rgba.height().is_multiple_of(16) {
                continue;
            }
            *rgba = image::imageops::resize(
                rgba,
                rgba.width().next_multiple_of(16),
                rgba.height().next_multiple_of(16),
                image::imageops::FilterType::Triangle,
            );
        }

        let mut save_data = Vec::new();
        map.save(&mut save_data)?;
        Ok(save_data)
    }
}

fn get_i32(data: &[u8], idx: usize) -> i32 {
    i32::from_le_bytes(data[idx * 4..(idx + 1) * 4].try_into().unwrap())
}

fn set_i32(data: &mut [u8], idx: usize, num: i32) {
    data[idx * 4..(idx + 1) * 4].copy_from_slice(&num.to_le_bytes())
}

fn detect_item_type(type_id: u16, ex: &parse::ExTypeIndex) -> Option<parse::ItemType> {
    for item_type in [
        parse::ItemType::Version,
        parse::ItemType::Info,
        parse::ItemType::Image,
        parse::ItemType::Envelope,
        parse::ItemType::Group,
        parse::ItemType::Layer,
        parse::ItemType::EnvPoints,
        parse::ItemType::EnvPointsBezierData,
        parse::ItemType::Sound,
        parse::ItemType::ExType,
        parse::ItemType::AutoMapperConfig,
    ] {
        match item_type.identifier() {
            parse::Identifier::TypeId(id) => {
                if id == type_id {
                    return Some(item_type);
                }
            }
            parse::Identifier::Uuid(uuid) => {
                let Some(id) = ex.get(&uuid) else { continue };
                if *id == type_id {
                    return Some(item_type);
                }
            }
        }
    }
    None
}

fn item_type_uuid(item_type: parse::ItemType) -> [u8; 16] {
    if let parse::Identifier::Uuid(uuid) = item_type.identifier() {
        uuid
    } else {
        panic!()
    }
}

fn item_type_id(item_type: parse::ItemType) -> u16 {
    if let parse::Identifier::TypeId(id) = item_type.identifier() {
        id
    } else {
        panic!()
    }
}

fn color_to_position(c: vek::Rgba<I22F10>) -> Position {
    Position {
        offset: vek::Vec2::new(
            I17F15::from_bits(c.r.to_bits()),
            I17F15::from_bits(c.g.to_bits()),
        ),
        rotation: c.b,
    }
}

fn color_to_volume(c: vek::Rgba<I22F10>) -> Volume {
    Volume(c.r)
}

fn conv_env<T: Copy, F: Fn(vek::Rgba<I22F10>) -> T>(env: &Env<vek::Rgba<I22F10>>, f: &F) -> Env<T> {
    use CurveKind::*;
    Env {
        name: env.name.clone(),
        synchronized: env.synchronized,
        points: env
            .points
            .iter()
            .map(|p| EnvPoint {
                time: p.time,
                content: f(p.content),
                curve: match p.curve {
                    Step => Step,
                    Linear => Linear,
                    Slow => Slow,
                    Fast => Fast,
                    Smooth => Smooth,
                    Bezier(b) => Bezier(BezierCurve {
                        handle_l: b.handle_l.map(f),
                        handle_r: b.handle_r.map(f),
                    }),
                    Unknown(x) => Unknown(x),
                },
            })
            .collect(),
    }
}