workshop_rs/analysis/

element_count.rs

use std::collections::HashMap;
use std::fmt;

use crate::catalog::{Catalog, Kind};
use crate::core::source::Span;
use crate::wir::{self, Action, ActionId, Program, Value, ValueId};

/// The Workshop node category represented in an element-count report.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElementNodeKind {
    Rule,
    Condition,
    Action,
    Value,
}

/// One node's contribution and its nested element-count analysis.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ElementCountNode {
    pub kind: ElementNodeKind,
    /// An opaque identity unique within this report. It is not a WIR or
    /// storage arena index and has no meaning across reports.
    pub id: usize,
    /// The canonical Workshop or analysis name for this node.
    pub name: String,
    /// The authored source span, when the program retained one.
    pub span: Option<Span>,
    /// The node-local contribution before child counts and adjustments.
    pub base_count: usize,
    /// The signed node-local adjustment, such as a direct-argument reduction
    /// or hero-pair surcharge.
    pub adjustment: isize,
    /// The node's recursive count: `base_count + adjustment + children`.
    pub count: usize,
    /// Nested values, conditions, and actions in canonical source order.
    pub children: Vec<ElementCountNode>,
}

/// A structured element-count report for one canonical Workshop program.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ElementCountReport {
    /// The sum of all rule counts.
    pub total: usize,
    /// Rule nodes in canonical source/WIR order.
    pub rules: Vec<ElementCountNode>,
}

impl ElementCountNode {
    /// The number of nodes on the longest downward path starting at this
    /// node: `1` for a leaf, otherwise `1 + max(child heights)`.
    ///
    /// This is a structural measure of the report tree, not a client
    /// limit; it answers "how deeply nested is this subtree" without
    /// claiming anything about importability.
    pub fn height(&self) -> usize {
        1 + self
            .children
            .iter()
            .map(Self::height)
            .max()
            .unwrap_or_default()
    }

    /// The number of nodes in this subtree, including this node.
    ///
    /// Unlike `count`, which is the weighted client element cost, this is
    /// the plain cardinality of the report subtree.
    pub fn node_count(&self) -> usize {
        1 + self.children.iter().map(Self::node_count).sum::<usize>()
    }

    /// The number of value nodes in this node's own statement value trees:
    /// `Value` nodes reachable from this node without entering a descendant
    /// `Action` subtree.
    ///
    /// On an `Action` this counts its own argument trees; actions nested in
    /// block bodies are separate statements and are excluded. On a
    /// `Condition` it is the condition's value tree, and on a `Rule` it is
    /// the conditions' value trees only — each action carries its own
    /// statement. On a `Value` node this equals `node_count()`.
    ///
    /// A comparison keeps its operator as a `Value` node inside an `If` or
    /// `While` argument, while a rule condition absorbs it into the
    /// `Condition` node — so one source expression measures one node fewer
    /// as a rule condition than as a block condition.
    pub fn statement_value_nodes(&self) -> usize {
        if self.kind == ElementNodeKind::Value {
            return self.node_count();
        }
        self.children
            .iter()
            .map(|child| match child.kind {
                ElementNodeKind::Value | ElementNodeKind::Condition => {
                    child.statement_value_nodes()
                }
                ElementNodeKind::Rule | ElementNodeKind::Action => 0,
            })
            .sum()
    }
}

impl ElementCountReport {
    /// Return the per-rule total in source/WIR order.
    pub fn rule_counts(&self) -> impl Iterator<Item = (&str, usize)> {
        self.rules
            .iter()
            .map(|rule| (rule.name.as_str(), rule.count))
    }
}

/// A construct for which an exact canonical element count cannot be produced.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum ElementCountError {
    /// The input cannot be materialized or structurally validated as a
    /// canonical Workshop program.
    InvalidProgram { message: String },
    /// The input contains a construct for which this analyzer has no exact
    /// canonical count.
    Unsupported {
        kind: ElementNodeKind,
        name: String,
        span: Option<Span>,
        reason: String,
    },
    /// The internal graph contains a recursive value or action reference.
    Cycle { kind: ElementNodeKind, id: usize },
}

impl fmt::Display for ElementCountError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidProgram { message } => write!(formatter, "invalid program: {message}"),
            Self::Unsupported {
                kind,
                name,
                span,
                reason,
            } => write!(
                formatter,
                "unsupported {kind:?} '{name}'{}: {reason}",
                span.map_or_else(String::new, |span| format!(" at {span:?}"))
            ),
            Self::Cycle { kind, id } => write!(formatter, "cyclic {kind:?} reference at {id}"),
        }
    }
}

impl std::error::Error for ElementCountError {}

impl Program {
    pub(crate) fn element_count(
        &self,
        catalog: &Catalog,
    ) -> Result<ElementCountReport, ElementCountError> {
        self.validate()
            .map_err(|error| ElementCountError::InvalidProgram {
                message: error.to_string(),
            })?;
        crate::rules::validate::validate_wir(self, catalog, false).map_err(|error| {
            ElementCountError::InvalidProgram {
                message: error.to_string(),
            }
        })?;

        let mut counter = Counter {
            program: self,
            catalog,
            values: HashMap::new(),
            actions: HashMap::new(),
            next_node_id: 0,
            variable_name_arg: false,
        };
        let mut rules = Vec::with_capacity(self.rules.len());
        for rule in self.rules.iter() {
            rules.push(counter.rule(rule)?);
        }
        let total = rules.iter().map(|rule| rule.count).sum();
        Ok(ElementCountReport { total, rules })
    }
}

impl crate::program::Program {
    /// Count the canonical Workshop target represented by this program.
    pub fn element_count(
        &self,
        catalog: &Catalog,
    ) -> Result<ElementCountReport, ElementCountError> {
        let storage = self
            .to_wir()
            .map_err(|error| ElementCountError::InvalidProgram {
                message: error.to_string(),
            })?;
        storage.element_count(catalog)
    }
}

struct Counted {
    node: ElementCountNode,
    heroes: usize,
}

impl Counted {
    #[allow(clippy::too_many_arguments)]
    fn finish(
        kind: ElementNodeKind,
        id: usize,
        name: impl Into<String>,
        span: Option<Span>,
        base_count: usize,
        adjustment: isize,
        children: Vec<ElementCountNode>,
        heroes: usize,
    ) -> Self {
        let children_count: usize = children.iter().map(|child| child.count).sum();
        let count = (base_count as isize + children_count as isize + adjustment).max(0) as usize;
        Self {
            node: ElementCountNode {
                kind,
                id,
                name: name.into(),
                span,
                base_count,
                adjustment,
                count,
                children,
            },
            heroes,
        }
    }
}

struct Counter<'a> {
    program: &'a Program,
    catalog: &'a Catalog,
    values: HashMap<usize, usize>,
    actions: HashMap<usize, usize>,
    next_node_id: usize,
    variable_name_arg: bool,
}

impl Counter<'_> {
    fn next_node_id(&mut self) -> usize {
        let id = self.next_node_id;
        self.next_node_id += 1;
        id
    }

    fn rule(&mut self, rule: &wir::Rule) -> Result<ElementCountNode, ElementCountError> {
        let node_id = self.next_node_id();
        let mut children = Vec::with_capacity(rule.conditions.len() + rule.actions.len());
        for condition in &rule.conditions {
            children.push(self.condition(condition.value)?.node);
        }
        for (index, action) in rule.actions.iter().enumerate() {
            let rule_final = index + 1 == rule.actions.len();
            children.push(self.action(*action, rule_final)?.node);
        }
        Ok(Counted::finish(
            ElementNodeKind::Rule,
            node_id,
            &rule.name,
            rule.span,
            1,
            0,
            children,
            0,
        )
        .node)
    }

    fn condition(&mut self, id: ValueId) -> Result<Counted, ElementCountError> {
        let node_id = self.next_node_id();
        let Some(value) = self.program.values.get(id) else {
            return Err(ElementCountError::InvalidProgram {
                message: format!("dangling condition value {}", id.index()),
            });
        };
        let (children, heroes) = match &value.value {
            Value::Call { name, args } if wir::is_comparison_operator(name) => {
                let mut children = Vec::with_capacity(args.len());
                let mut heroes = 0;
                for argument in args {
                    let counted = self.value(*argument, true)?;
                    heroes += counted.heroes / 2 * 2;
                    children.push(counted.node);
                }
                (children, heroes)
            }
            _ => {
                let counted = self.value(id, true)?;
                (vec![counted.node], counted.heroes)
            }
        };
        Ok(Counted::finish(
            ElementNodeKind::Condition,
            node_id,
            "condition",
            value.span,
            1,
            pair_surcharge(heroes),
            children,
            heroes,
        ))
    }

    fn action(&mut self, id: ActionId, rule_final: bool) -> Result<Counted, ElementCountError> {
        let node_id = self.next_node_id();
        if let Some(&active_id) = self.actions.get(&id.index()) {
            return Err(ElementCountError::Cycle {
                kind: ElementNodeKind::Action,
                id: active_id,
            });
        }
        self.actions.insert(id.index(), node_id);
        let mut action = self.program.actions.get(id);
        while let Some(Action::Disabled { action: inner, .. }) = action {
            action = self.program.actions.get(*inner);
        }
        let Some(action) = action else {
            return Err(ElementCountError::InvalidProgram {
                message: format!("dangling action {}", id.index()),
            });
        };
        let result = self.action_inner(action, node_id, rule_final);
        self.actions.remove(&id.index());
        result
    }

    fn action_inner(
        &mut self,
        action: &Action,
        node_id: usize,
        rule_final: bool,
    ) -> Result<Counted, ElementCountError> {
        let span = action.span();
        let mut children = Vec::new();
        let mut heroes = 0;
        let mut base = 1;
        let name;
        match action {
            Action::SetGlobalVariable { value, .. }
            | Action::ModifyGlobalVariable { value, .. } => {
                name = "variable action";
                self.push_action_value(&mut children, &mut heroes, *value)?;
            }
            Action::SetPlayerVariable { player, value, .. }
            | Action::ModifyPlayerVariable { player, value, .. } => {
                name = "player variable action";
                self.push_action_value(&mut children, &mut heroes, *player)?;
                self.push_action_value(&mut children, &mut heroes, *value)?;
            }
            Action::AssignMember { target, value, .. } => {
                name = "member assignment";
                self.push_action_value(&mut children, &mut heroes, *target)?;
                self.push_action_value(&mut children, &mut heroes, *value)?;
            }
            Action::CallSubroutine { .. } => {
                name = "call subroutine";
            }
            Action::If {
                branches,
                else_body,
                ..
            } => {
                name = "if";
                base += branches.len().saturating_sub(1)
                    + usize::from(else_body.is_some())
                    + usize::from(!rule_final);
                for branch in branches {
                    self.push_action_value(&mut children, &mut heroes, branch.condition)?;
                    for nested in &branch.body {
                        children.push(self.action(*nested, false)?.node);
                    }
                }
                if let Some(body) = else_body {
                    for nested in body {
                        children.push(self.action(*nested, false)?.node);
                    }
                }
            }
            Action::While {
                condition, body, ..
            } => {
                name = "while";
                base += 1;
                self.push_action_value(&mut children, &mut heroes, *condition)?;
                for nested in body {
                    children.push(self.action(*nested, false)?.node);
                }
            }
            Action::ForGlobalVariable {
                start,
                stop,
                step,
                body,
                ..
            } => {
                name = "for global variable";
                base += 1;
                for value in [start, stop, step] {
                    self.push_action_value(&mut children, &mut heroes, *value)?;
                }
                for nested in body {
                    children.push(self.action(*nested, false)?.node);
                }
            }
            Action::ForPlayerVariable {
                player,
                start,
                stop,
                step,
                body,
                ..
            } => {
                name = "for player variable";
                base += 1;
                for value in [player, start, stop, step] {
                    self.push_action_value(&mut children, &mut heroes, *value)?;
                }
                for nested in body {
                    children.push(self.action(*nested, false)?.node);
                }
            }
            Action::Disabled { .. } => {
                return Err(ElementCountError::InvalidProgram {
                    message: "a disabled action wraps no enabled action".to_string(),
                });
            }
            Action::Call {
                name: action_name,
                args,
                ..
            } => {
                if self.catalog.entry(Kind::Action, action_name).is_none() {
                    return Err(ElementCountError::Unsupported {
                        kind: ElementNodeKind::Action,
                        name: action_name.clone(),
                        span,
                        reason: "the action is not a catalog identity".to_string(),
                    });
                }
                name = action_name.as_str();
                for (index, argument) in args.iter().enumerate() {
                    self.variable_name_arg = variable_slot(
                        self.program,
                        self.catalog,
                        Kind::Action,
                        action_name,
                        args,
                        index,
                    ) || (index == 0
                        && takes_variable(self.catalog, Kind::Action, action_name));
                    self.push_action_value(&mut children, &mut heroes, *argument)?;
                }
                base = (base as isize
                    + omitted_default_cost(
                        self.program,
                        self.catalog,
                        Kind::Action,
                        action_name,
                        args,
                        true,
                    ))
                .max(0) as usize;
            }
        }
        Ok(Counted::finish(
            ElementNodeKind::Action,
            node_id,
            name,
            span,
            base,
            pair_surcharge(heroes),
            children,
            heroes,
        ))
    }

    fn push_action_value(
        &mut self,
        children: &mut Vec<ElementCountNode>,
        heroes: &mut usize,
        id: ValueId,
    ) -> Result<(), ElementCountError> {
        let counted = self.value(id, true)?;
        *heroes += counted.heroes / 2 * 2;
        children.push(counted.node);
        Ok(())
    }

    fn value(&mut self, id: ValueId, top_level: bool) -> Result<Counted, ElementCountError> {
        let variable_name_arg = std::mem::take(&mut self.variable_name_arg);
        let node_id = self.next_node_id();
        if let Some(&active_id) = self.values.get(&id.index()) {
            return Err(ElementCountError::Cycle {
                kind: ElementNodeKind::Value,
                id: active_id,
            });
        }
        self.values.insert(id.index(), node_id);
        let Some(value) = self.program.values.get(id) else {
            return Err(ElementCountError::InvalidProgram {
                message: format!("dangling value {}", id.index()),
            });
        };
        let span = value.span;
        let result = match &value.value {
            Value::Number { .. } => self.value_node(node_id, "number", span, 2, vec![], 0),
            Value::String(_) => self.value_node(node_id, "string", span, 1, vec![], 0),
            Value::LocalizedString(_) => {
                self.value_node(node_id, "localized string", span, 2, vec![], 0)
            }
            Value::Bool(_) => self.value_node(node_id, "boolean", span, 1, vec![], 0),
            Value::Null => self.value_node(node_id, "null", span, 1, vec![], 0),
            Value::Array(elements) => {
                self.value_children(node_id, "array", span, 2, elements, None)
            }
            Value::Vector { x, y, z } => {
                self.value_children(node_id, "vector", span, 1, &[*x, *y, *z], None)
            }
            Value::Enum { value_type, .. } => {
                let heroes = usize::from(value_type == "Hero");
                let base = literal_enum_cost(value_type);
                self.value_node(node_id, value_type, span, base, vec![], heroes)
            }
            Value::GlobalVariable(_) => {
                let base = if variable_name_arg { 1 } else { 2 };
                self.value_node(node_id, "global variable", span, base, vec![], 0)
            }
            Value::PlayerVariable { player, .. } => {
                let base = if variable_name_arg { 0 } else { 2 };
                self.value_children(node_id, "player variable", span, base, &[*player], None)
            }
            Value::Subroutine(_) => self.value_node(node_id, "subroutine", span, 1, vec![], 0),
            Value::EventPlayer => self.value_node(node_id, "event player", span, 1, vec![], 0),
            Value::Call { name, args } => {
                if name == crate::wir::AMBIGUOUS_ENUM_CALL
                    && crate::wir::ambiguous_enum_parts(self.program, id).is_some()
                {
                    self.value_node(node_id, "ambiguous enum", span, 1, vec![], 0)
                } else {
                    if name != "memberAccess"
                        && self.catalog.entry(Kind::Value, name).is_none()
                        && self.catalog.entry(Kind::Operator, name).is_none()
                        && !crate::wir::is_canonical_helper(name)
                    {
                        return Err(ElementCountError::Unsupported {
                            kind: ElementNodeKind::Value,
                            name: name.clone(),
                            span,
                            reason: "the value is not a catalog identity".to_string(),
                        });
                    }
                    let child_ids: Vec<ValueId> = if name == "memberAccess" {
                        args.first()
                            .copied()
                            .into_iter()
                            .chain(args.iter().copied().skip(2))
                            .collect()
                    } else {
                        args.clone()
                    };
                    let base = if wir::is_comparison_operator(name)
                        || name == "array"
                        || name == "evaluateOnce"
                    {
                        2
                    } else {
                        1
                    };
                    let omitted = omitted_default_cost(
                        self.program,
                        self.catalog,
                        Kind::Value,
                        name,
                        args,
                        false,
                    );
                    let mut counted = self.value_children(
                        node_id,
                        name,
                        span,
                        (base as isize + omitted) as usize,
                        &child_ids,
                        Some((Kind::Value, name)),
                    )?;
                    let setting = workshop_setting_adjustment(name);
                    counted.node.adjustment += setting;
                    counted.node.count = (counted.node.count as isize + setting).max(0) as usize;
                    Ok(counted)
                }
            }
        }?;
        self.values.remove(&id.index());
        let mut result = result;
        if top_level {
            result.node.adjustment -= 1;
            result.node.count = (result.node.count as isize - 1).max(0) as usize;
        }
        Ok(result)
    }

    fn value_node(
        &self,
        id: usize,
        name: impl Into<String>,
        span: Option<Span>,
        base: usize,
        children: Vec<ElementCountNode>,
        heroes: usize,
    ) -> Result<Counted, ElementCountError> {
        Ok(Counted::finish(
            ElementNodeKind::Value,
            id,
            name,
            span,
            base,
            0,
            children,
            heroes,
        ))
    }

    fn value_children(
        &mut self,
        id: usize,
        name: impl Into<String>,
        span: Option<Span>,
        base: usize,
        ids: &[ValueId],
        owner: Option<(Kind, &str)>,
    ) -> Result<Counted, ElementCountError> {
        let mut children = Vec::with_capacity(ids.len());
        let mut heroes = 0;
        for (index, child) in ids.iter().enumerate() {
            self.variable_name_arg = owner.is_some_and(|(kind, name)| {
                variable_slot(self.program, self.catalog, kind, name, ids, index)
            });
            let counted = self.value(*child, false)?;
            heroes += counted.heroes;
            children.push(counted.node);
        }
        self.value_node(id, name, span, base, children, heroes)
    }
}

fn pair_surcharge(heroes: usize) -> isize {
    (heroes / 2) as isize
}

fn omitted_default_cost(
    program: &wir::Program,
    catalog: &Catalog,
    kind: Kind,
    name: &str,
    args: &[ValueId],
    direct: bool,
) -> isize {
    let Some(entry) = catalog.entry(kind, name) else {
        return 0;
    };
    let given = args.len() + usize::from(entry.wir_args_folded(program, args));
    (given..entry.param_count())
        .filter_map(|index| entry.param_default(index))
        .map(|default| {
            let cost = if default.parse::<f64>().is_ok() {
                2
            } else {
                default.split('.').next().map_or(1, literal_enum_cost) as isize
            };
            cost - isize::from(direct)
        })
        .sum()
}

fn literal_enum_cost(value_type: &str) -> usize {
    if matches!(value_type, "Team" | "Hero" | "Color" | "Button" | "Map") {
        2
    } else {
        1
    }
}

fn variable_slot(
    program: &wir::Program,
    catalog: &Catalog,
    kind: Kind,
    name: &str,
    args: &[ValueId],
    index: usize,
) -> bool {
    catalog.entry(kind, name).is_some_and(|entry| {
        let index = entry.wir_arg_param_index(program, args, index);
        entry.param_type(index) == Some("Variable")
            || entry
                .params()
                .get(index)
                .is_some_and(|param| param == "Variable")
    })
}

fn takes_variable(catalog: &Catalog, kind: Kind, name: &str) -> bool {
    catalog.entry(kind, name).is_some_and(|entry| {
        (0..entry.param_count()).any(|index| {
            entry.param_type(index) == Some("Variable")
                || entry
                    .params()
                    .get(index)
                    .is_some_and(|param| param == "Variable")
        })
    })
}

fn workshop_setting_adjustment(name: &str) -> isize {
    match name {
        "workshopSettingInteger"
        | "workshopSettingFloat"
        | "createWorkshopSettingInt"
        | "createWorkshopSettingFloat" => -3,
        "workshopSettingCombo" | "createWorkshopSettingEnum" => -2,
        _ => 0,
    }
}