onedrop_hlsl/

texture_plan.rs

//! Per-preset user-texture binding plan + the `tex2D` rewriter that
//! consumes it.
//!
//! Built by the renderer at preset load time:
//!
//! 1. [`scan_user_samplers`] extracts `sampler sampler_X;` declarations from
//!    the preset's HLSL.
//! 2. The renderer resolves each name to a [`TextureSlot`] via its texture
//!    pool. `sampler_rand0X` becomes a deterministic pick; literal names
//!    become a pool lookup.
//! 3. The slots are passed to [`crate::translate_shader_with_plan`] and the
//!    wrapper so the emitted WGSL points at the right binding.

use regex::Regex;
use std::sync::LazyLock;

/// Maximum simultaneously-bound user textures per preset. Sized to cover
/// the in-the-wild distribution: a typical preset survey shows ≤ 4
/// distinct `sampler sampler_X;` declarations per shader; 8 leaves
/// headroom without blowing past the comp pipeline's bind-group budget.
/// Presets that exceed this cap are translated with the first 8 slots
/// used and the rest falling through to the standard fallback path
/// (sampler_main).
pub const MAX_USER_TEXTURE_SLOTS: usize = 8;

/// WGSL binding name for user-texture slot `slot`. The codegen wrapper
/// declares `var sampler_user_<n>_texture: texture_2d<f32>` for each
/// slot `0..MAX_USER_TEXTURE_SLOTS`, and this is what the translator
/// emits in `textureSample(...)` calls for plan-routed samplers.
pub fn user_texture_binding_name(slot: usize) -> String {
    format!("sampler_user_{slot}_texture")
}

/// One slot in a [`TextureBindingPlan`]. Carries the resolved pool
/// texture name (or `None` when the renderer should fall back to the
/// 1×1 white fallback) and the texture's `vec4<f32>(w, h, 1/w, 1/h)`
/// for the `texsize_<NAME>` constant the wrapper emits.
#[derive(Debug, Clone, PartialEq)]
pub struct TextureSlot {
    /// Logical pool name (lowercase filename stem) that the renderer
    /// must bind into this slot. `None` → no texture available; bind
    /// the fallback. Stable across reloads as long as the pool's disk
    /// layout doesn't change.
    pub pool_name: Option<String>,
    /// MD2 `texsize_<NAME>` constant for the texture that's bound.
    /// Emitted inline by the wrapper so user shader code can read the
    /// dimensions without going through a uniform. `[1.0; 4]` for the
    /// fallback case.
    pub texsize: [f32; 4],
}

/// Per-preset mapping from HLSL sampler names to the comp pass's
/// user-texture slots.
///
/// `aliases` carries the surface-level names the translator should
/// treat as referencing each slot. A single slot may carry multiple
/// aliases — `sampler sampler_clouds;` and `sampler sampler_fw_clouds;`
/// both resolve to the same logical `clouds` texture but pick different
/// filter samplers when sampled. The alias map maps a full sampler name
/// (e.g. `"sampler_fw_clouds"`) to `(slot_index, sampler_kind)`.
#[derive(Debug, Clone, Default)]
pub struct TextureBindingPlan {
    slots: Vec<TextureSlot>,
    aliases: Vec<(String, usize, &'static str)>,
}

impl TextureBindingPlan {
    /// Empty plan — `translate_shader_with_plan` behaves identically to
    /// `translate_shader`. Used by callers that don't load disk
    /// textures (CLI batch translate, tests) and as the default for
    /// `translate_shader` itself.
    pub fn empty() -> Self {
        Self::default()
    }

    /// Number of filled slots. Equal to the comp pipeline's
    /// user-texture bind count for this preset (the renderer fills the
    /// remaining slots up to `MAX_USER_TEXTURE_SLOTS` with the fallback
    /// texture).
    pub fn slot_count(&self) -> usize {
        self.slots.len()
    }

    /// Iterate slots in binding order. Index → slot.
    pub fn slots(&self) -> &[TextureSlot] {
        &self.slots
    }

    /// Register one logical texture (`pool_name = Some("clouds")` or
    /// `None` for fallback) along with the set of HLSL surface aliases
    /// that reference it. Returns the slot index, or `None` when the
    /// cap was already hit. Subsequent calls with new aliases for an
    /// existing `pool_name` re-use the same slot.
    pub fn add_slot(
        &mut self,
        pool_name: Option<String>,
        texsize: [f32; 4],
        aliases: &[(String, &'static str)],
    ) -> Option<usize> {
        // Same `pool_name` shouldn't get a second slot — collapse onto
        // the existing one.
        let slot = if let Some(idx) = self
            .slots
            .iter()
            .position(|s| s.pool_name == pool_name && pool_name.is_some())
        {
            idx
        } else if self.slots.len() >= MAX_USER_TEXTURE_SLOTS {
            log::warn!(
                "user texture slot cap ({MAX_USER_TEXTURE_SLOTS}) reached; \
                 alias {:?} routes through fallback",
                aliases.iter().map(|(n, _)| n).collect::<Vec<_>>()
            );
            return None;
        } else {
            self.slots.push(TextureSlot { pool_name, texsize });
            self.slots.len() - 1
        };
        for (alias, sampler_kind) in aliases {
            // Skip duplicate aliases — keeps the first-write-wins
            // semantics for the WGSL sampler binding (so the translator
            // sees a stable routing if the renderer ever re-adds an
            // alias).
            if self.aliases.iter().any(|(a, _, _)| a == alias) {
                continue;
            }
            self.aliases.push((alias.clone(), slot, sampler_kind));
        }
        Some(slot)
    }

    /// Map an HLSL `tex2D` sampler name to a `(slot_index, sampler_kind)`.
    /// Returns `None` for built-in samplers (those are handled by
    /// [`resolve_tex2d_sampler`] before the plan is consulted) and for
    /// names the plan never registered an alias for.
    pub fn lookup_tex2d(&self, sampler_name: &str) -> Option<(usize, &'static str)> {
        self.aliases
            .iter()
            .find(|(a, _, _)| a == sampler_name)
            .map(|(_, slot, kind)| (*slot, *kind))
    }
}

/// Extract every `sampler sampler_X;` declaration from a MilkDrop comp
/// shader HLSL and return the **logical name** (the part after
/// `sampler_`) for each occurrence that isn't already a built-in.
///
/// Built-ins skipped:
/// - `main` (and its `fw/fc/pw/pc_main` variants)
/// - `noise_lq` / `noise_mq` / `noise_hq` (and the same set prefixed by
///   `fw/fc/pw/pc_`)
/// - `noisevol_lq` / `noisevol_hq` (same prefix set)
/// - `blur1` / `blur2` / `blur3`
///
/// The returned list preserves declaration order with duplicates
/// removed. The renderer uses this to know which textures it needs to
/// bind for this preset.
///
/// **Format note** — `sampler sampler_X` and `sampler2D sampler_X` both
/// match (MD2 uses both spellings; the type matters to HLSL but not to
/// us since we strip the decl entirely before translation).
pub fn scan_user_samplers(hlsl: &str) -> Vec<UserSamplerRef> {
    static SAMPLER_DECL: LazyLock<Regex> = LazyLock::new(|| {
        Regex::new(
            r"(?m)^\s*(?:sampler|texture|texture2D|sampler2D)\s+(sampler_[A-Za-z_][A-Za-z0-9_]*)\s*;",
        )
        .unwrap()
    });

    let mut out: Vec<UserSamplerRef> = Vec::new();
    for cap in SAMPLER_DECL.captures_iter(hlsl) {
        let full_name = cap.get(1).unwrap().as_str();
        if is_builtin_sampler_name(full_name) {
            continue;
        }
        if out.iter().any(|s| s.full_name == full_name) {
            continue;
        }
        let (logical, sampler_kind) = decompose_sampler_name(full_name);
        out.push(UserSamplerRef {
            full_name: full_name.to_string(),
            logical_name: logical,
            sampler_kind,
        });
    }
    out
}

/// One parsed `sampler sampler_X;` declaration. The renderer consumes
/// this to populate a [`TextureBindingPlan`].
#[derive(Debug, Clone, PartialEq)]
pub struct UserSamplerRef {
    /// Full HLSL sampler name as written: `sampler_clouds`,
    /// `sampler_fw_clouds`, `sampler_rand02_smalltiled`, …
    pub full_name: String,
    /// Logical texture name (filter prefix stripped):
    /// `sampler_fw_clouds` → `"clouds"`, `sampler_rand02` → `"rand02"`,
    /// `sampler_rand02_smalltiled` → `"rand02_smalltiled"`.
    pub logical_name: String,
    /// WGSL sampler binding the translator should pair with this
    /// texture. `sampler_fw_*` → `"sampler_fw"`, plain `sampler_*` →
    /// `"sampler_fw"` (MD2 default for user textures is linear+wrap).
    pub sampler_kind: &'static str,
}

/// Decompose an HLSL user-sampler name into `(logical_name,
/// sampler_kind)`. The logical name is what the renderer looks up in
/// the texture pool; the sampler kind is the WGSL sampler binding used
/// at the `textureSample` call site.
pub(crate) fn decompose_sampler_name(full: &str) -> (String, &'static str) {
    if let Some(rest) = full.strip_prefix("sampler_fw_") {
        return (rest.to_string(), "sampler_fw");
    }
    if let Some(rest) = full.strip_prefix("sampler_fc_") {
        return (rest.to_string(), "sampler_fc");
    }
    if let Some(rest) = full.strip_prefix("sampler_pw_") {
        return (rest.to_string(), "sampler_pw");
    }
    if let Some(rest) = full.strip_prefix("sampler_pc_") {
        return (rest.to_string(), "sampler_pc");
    }
    if let Some(rest) = full.strip_prefix("sampler_") {
        // MD2's default sampler state for user textures is linear+wrap,
        // matching what most authored shaders expect for tiled patterns
        // (`clouds`, `worms`, `lichen`, …).
        return (rest.to_string(), "sampler_fw");
    }
    (full.to_string(), "sampler_main")
}

/// `true` for sampler names the codegen wrapper already binds.
/// Filtering these out keeps the user-texture plan focused on
/// disk-loaded textures.
fn is_builtin_sampler_name(full: &str) -> bool {
    matches!(
        full,
        "sampler_main"
            | "sampler_fw_main"
            | "sampler_fc_main"
            | "sampler_pw_main"
            | "sampler_pc_main"
            | "sampler_blur1"
            | "sampler_blur2"
            | "sampler_blur3"
            | "sampler_noise_lq"
            | "sampler_noise_mq"
            | "sampler_noise_hq"
            | "sampler_noisevol_lq"
            | "sampler_noisevol_hq"
    ) || matches!(
        decompose_sampler_name(full).0.as_str(),
        // Prefixed noise — `sampler_fw_noise_hq` etc. are routed by
        // resolve_tex2d_sampler.
        "noise_lq"
            | "noise_mq"
            | "noise_hq"
            | "noisevol_lq"
            | "noisevol_hq"
            | "main"
            | "blur1"
            | "blur2"
            | "blur3"
    )
}

/// Resolve an MD2 sampler name to the `(texture_binding,
/// sampler_binding, recognised)` triple the WGSL wrapper actually
/// exposes. `recognised = true` suppresses the `/*was: …*/` debug
/// comment for bindings we routed deliberately (vs. the fallback case).
///
/// 1. **Noise pack textures** — `sampler_noise_lq` / `_mq` / `_hq`
///    route onto the dedicated noise textures using `sampler_pw`
///    (point+wrap), matching MD2's high-frequency sampling intent.
/// 2. **fw/fc/pw/pc variants of main** — `sampler_fw_main`,
///    `_fc_main`, `_pw_main`, `_pc_main` keep `sampler_main_texture`
///    but switch the sampler to the matching filter × address mode.
/// 3. **Anything else** — falls back to `(sampler_main_texture,
///    sampler_main)` so unknown / user-loaded sampler names still
///    produce valid WGSL. The original name is kept in a trailing
///    comment for debugging.
pub(crate) fn resolve_tex2d_sampler(name: &str) -> (&'static str, &'static str, bool) {
    match name {
        "sampler_noise_lq"
        | "sampler_pw_noise_lq"
        | "sampler_fw_noise_lq"
        | "sampler_pc_noise_lq"
        | "sampler_fc_noise_lq" => ("sampler_noise_lq_texture", noise_sampler_for(name), true),
        "sampler_noise_mq"
        | "sampler_pw_noise_mq"
        | "sampler_fw_noise_mq"
        | "sampler_pc_noise_mq"
        | "sampler_fc_noise_mq" => ("sampler_noise_mq_texture", noise_sampler_for(name), true),
        "sampler_noise_hq"
        | "sampler_pw_noise_hq"
        | "sampler_fw_noise_hq"
        | "sampler_pc_noise_hq"
        | "sampler_fc_noise_hq" => ("sampler_noise_hq_texture", noise_sampler_for(name), true),
        "sampler_fw_main" => ("sampler_main_texture", "sampler_fw", true),
        "sampler_fc_main" => ("sampler_main_texture", "sampler_fc", true),
        "sampler_pw_main" => ("sampler_main_texture", "sampler_pw", true),
        "sampler_pc_main" => ("sampler_main_texture", "sampler_pc", true),
        // Already-resolved or unknown — fall through to main+linear.
        _ => ("sampler_main_texture", "sampler_main", false),
    }
}

/// For a noise sampler name embedded in one of the 4 MD2 variants,
/// pick the sampler binding. Plain `sampler_noise_lq` (no prefix)
/// defaults to the point+wrap variant — that's how MD2 samples noise
/// pixel-by-pixel when the author wants the raw value.
pub(crate) fn noise_sampler_for(name: &str) -> &'static str {
    if name.starts_with("sampler_fw_") {
        "sampler_fw"
    } else if name.starts_with("sampler_fc_") {
        "sampler_fc"
    } else if name.starts_with("sampler_pc_") {
        "sampler_pc"
    } else {
        "sampler_pw"
    }
}

/// Rewrite `tex2D(<sampler>, <uv>)` to a WGSL `textureSample` call
/// against the `(texture, sampler)` pair resolved from the MD2 sampler
/// name. Paren-balanced because the uv expression often contains nested
/// calls (`tex2D(s, uv + offset(t))`).
pub(crate) fn replace_texture_sampling_with_plan(code: &str, plan: &TextureBindingPlan) -> String {
    let bytes = code.as_bytes();
    let mut out = String::with_capacity(code.len());
    let mut i = 0usize;

    while i < bytes.len() {
        if i + 6 <= bytes.len()
            && &bytes[i..i + 5] == b"tex2D"
            && bytes[i + 5] == b'('
            && (i == 0 || !(bytes[i - 1].is_ascii_alphanumeric() || bytes[i - 1] == b'_'))
        {
            let arg_start = i + 6;
            let mut j = arg_start;
            let mut depth = 1i32;
            let mut comma = None;
            while j < bytes.len() {
                match bytes[j] {
                    b'(' => depth += 1,
                    b')' => {
                        depth -= 1;
                        if depth == 0 {
                            break;
                        }
                    }
                    b',' if depth == 1 && comma.is_none() => comma = Some(j),
                    _ => {}
                }
                j += 1;
            }
            if let Some(c) = comma
                && j < bytes.len()
            {
                let sampler = code[arg_start..c].trim();
                let uv = code[c + 1..j].trim();
                // Built-in routing wins first (main / fw_main / noise_*).
                let (tex, smp, recognised) = resolve_tex2d_sampler(sampler);
                if recognised {
                    out.push_str(&format!("textureSample({tex}, {smp}, {uv})"));
                } else if let Some((slot, sampler_kind)) = plan.lookup_tex2d(sampler) {
                    // Plan-driven routing for user-loaded textures.
                    let tex_name = user_texture_binding_name(slot);
                    out.push_str(&format!("textureSample({tex_name}, {sampler_kind}, {uv})"));
                } else {
                    out.push_str(&format!(
                        "textureSample({tex}, {smp}, {uv}) /*was: {sampler}*/"
                    ));
                }
                i = j + 1;
                continue;
            }
        }
        out.push(bytes[i] as char);
        i += 1;
    }

    out
}