wikioxide 0.1.0

//! Resource limits mirroring MediaWiki's production defaults.

use crate::error::BuildError;

const MIB: usize = 1024 * 1024;

/// Expansion and parsing resource limits.
///
/// Defaults mirror Wikimedia production: expansion depth 100, one million preprocessor
/// nodes, 2 MiB post-expand include size, 2 MiB per template argument, 500 expensive
/// parser-function calls, unstrip depth 20 and 5 MB unstrip size. Exceeding a limit never
/// fails a parse — it emits the matching MediaWiki parse warning and degrades exactly the
/// way MediaWiki does.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Limits {
    expansion_depth: u16,
    node_count: u64,
    post_expand_bytes: usize,
    template_arg_bytes: usize,
    expensive_parser_functions: u32,
    unstrip_depth: u16,
    unstrip_bytes: usize,
    max_source_bytes: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            expansion_depth: 100,
            node_count: 1_000_000,
            post_expand_bytes: 2 * MIB,
            template_arg_bytes: 2 * MIB,
            expensive_parser_functions: 500,
            unstrip_depth: 20,
            unstrip_bytes: 5_000_000,
            max_source_bytes: 4 * MIB,
        }
    }
}

impl Limits {
    /// Maximum template/parser-function nesting depth (`$wgMaxTemplateDepth`).
    #[must_use]
    pub const fn expansion_depth(&self) -> u16 {
        self.expansion_depth
    }

    /// Maximum preprocessor nodes visited per parse (`$wgMaxPPNodeCount`).
    #[must_use]
    pub const fn node_count(&self) -> u64 {
        self.node_count
    }

    /// Maximum total bytes produced by transclusions (post-expand include size).
    #[must_use]
    pub const fn post_expand_bytes(&self) -> usize {
        self.post_expand_bytes
    }

    /// Maximum expanded size of a single template argument.
    #[must_use]
    pub const fn template_arg_bytes(&self) -> usize {
        self.template_arg_bytes
    }

    /// Maximum number of expensive parser-function calls (`$wgExpensiveParserFunctionLimit`).
    #[must_use]
    pub const fn expensive_parser_functions(&self) -> u32 {
        self.expensive_parser_functions
    }

    /// Maximum recursion depth while unstripping extension-tag markers.
    #[must_use]
    pub const fn unstrip_depth(&self) -> u16 {
        self.unstrip_depth
    }

    /// Maximum total unstripped bytes.
    #[must_use]
    pub const fn unstrip_bytes(&self) -> usize {
        self.unstrip_bytes
    }

    /// Maximum accepted input size; longer sources are truncated with a warning.
    #[must_use]
    pub const fn max_source_bytes(&self) -> usize {
        self.max_source_bytes
    }

    /// Returns a copy with the expansion depth replaced.
    #[must_use]
    pub const fn with_expansion_depth(mut self, value: u16) -> Self {
        self.expansion_depth = value;
        self
    }

    /// Returns a copy with the preprocessor node budget replaced.
    #[must_use]
    pub const fn with_node_count(mut self, value: u64) -> Self {
        self.node_count = value;
        self
    }

    /// Returns a copy with the post-expand include size replaced.
    #[must_use]
    pub const fn with_post_expand_bytes(mut self, value: usize) -> Self {
        self.post_expand_bytes = value;
        self
    }

    /// Returns a copy with the per-argument expansion cap replaced.
    #[must_use]
    pub const fn with_template_arg_bytes(mut self, value: usize) -> Self {
        self.template_arg_bytes = value;
        self
    }

    /// Returns a copy with the expensive parser-function budget replaced.
    #[must_use]
    pub const fn with_expensive_parser_functions(mut self, value: u32) -> Self {
        self.expensive_parser_functions = value;
        self
    }

    /// Returns a copy with the unstrip depth replaced.
    #[must_use]
    pub const fn with_unstrip_depth(mut self, value: u16) -> Self {
        self.unstrip_depth = value;
        self
    }

    /// Returns a copy with the unstrip size cap replaced.
    #[must_use]
    pub const fn with_unstrip_bytes(mut self, value: usize) -> Self {
        self.unstrip_bytes = value;
        self
    }

    /// Returns a copy with the maximum input size replaced.
    #[must_use]
    pub const fn with_max_source_bytes(mut self, value: usize) -> Self {
        self.max_source_bytes = value;
        self
    }

    /// Checks that the limits can actually drive a parse (used by the parser builder).
    ///
    /// Zero budgets (except the expensive parser-function budget, where zero legitimately
    /// means "expensive functions disallowed") and source sizes beyond the 32-bit span
    /// space are rejected.
    pub fn validate(&self) -> Result<(), BuildError> {
        if self.expansion_depth == 0 {
            return Err(BuildError::InvalidLimits(
                "expansion depth must be non-zero",
            ));
        }
        if self.node_count == 0 {
            return Err(BuildError::InvalidLimits("node count must be non-zero"));
        }
        if self.post_expand_bytes == 0 {
            return Err(BuildError::InvalidLimits(
                "post-expand size must be non-zero",
            ));
        }
        if self.template_arg_bytes == 0 {
            return Err(BuildError::InvalidLimits(
                "template-argument size must be non-zero",
            ));
        }
        if self.unstrip_depth == 0 {
            return Err(BuildError::InvalidLimits("unstrip depth must be non-zero"));
        }
        if self.unstrip_bytes == 0 {
            return Err(BuildError::InvalidLimits("unstrip size must be non-zero"));
        }
        if self.max_source_bytes == 0 {
            return Err(BuildError::InvalidLimits(
                "max source size must be non-zero",
            ));
        }
        if u32::try_from(self.max_source_bytes).is_err() {
            return Err(BuildError::InvalidLimits(
                "max source size must fit in 32 bits (span space)",
            ));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn defaults_mirror_wikimedia_production() {
        let limits = Limits::default();
        assert_eq!(limits.expansion_depth(), 100);
        assert_eq!(limits.node_count(), 1_000_000);
        assert_eq!(limits.post_expand_bytes(), 2_097_152);
        assert_eq!(limits.template_arg_bytes(), 2_097_152);
        assert_eq!(limits.expensive_parser_functions(), 500);
        assert_eq!(limits.unstrip_depth(), 20);
        assert_eq!(limits.unstrip_bytes(), 5_000_000);
        assert_eq!(limits.max_source_bytes(), 4_194_304);
        assert!(limits.validate().is_ok());
    }

    #[test]
    fn builder_style_setters_replace_each_field() {
        let limits = Limits::default()
            .with_expansion_depth(40)
            .with_node_count(1_000)
            .with_post_expand_bytes(64)
            .with_template_arg_bytes(32)
            .with_expensive_parser_functions(0)
            .with_unstrip_depth(5)
            .with_unstrip_bytes(1_000)
            .with_max_source_bytes(512);
        assert_eq!(limits.expansion_depth(), 40);
        assert_eq!(limits.node_count(), 1_000);
        assert_eq!(limits.post_expand_bytes(), 64);
        assert_eq!(limits.template_arg_bytes(), 32);
        assert_eq!(limits.expensive_parser_functions(), 0);
        assert_eq!(limits.unstrip_depth(), 5);
        assert_eq!(limits.unstrip_bytes(), 1_000);
        assert_eq!(limits.max_source_bytes(), 512);
        assert!(
            limits.validate().is_ok(),
            "zero expensive-pf budget is legal"
        );
        assert_ne!(limits, Limits::default());
    }

    #[test]
    fn zero_budgets_fail_validation() {
        let cases = [
            Limits::default().with_expansion_depth(0),
            Limits::default().with_node_count(0),
            Limits::default().with_post_expand_bytes(0),
            Limits::default().with_template_arg_bytes(0),
            Limits::default().with_unstrip_depth(0),
            Limits::default().with_unstrip_bytes(0),
            Limits::default().with_max_source_bytes(0),
        ];
        for limits in cases {
            assert!(matches!(
                limits.validate(),
                Err(BuildError::InvalidLimits(_))
            ));
        }
    }

    #[cfg(target_pointer_width = "64")]
    #[test]
    fn source_cap_must_fit_span_space() {
        let too_big = Limits::default().with_max_source_bytes(u32::MAX as usize + 1);
        assert!(too_big.validate().is_err());
    }
}