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};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElementNodeKind {
Rule,
Condition,
Action,
Value,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ElementCountNode {
pub kind: ElementNodeKind,
pub id: usize,
pub name: String,
pub span: Option<Span>,
pub base_count: usize,
pub adjustment: isize,
pub count: usize,
pub children: Vec<ElementCountNode>,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ElementCountReport {
pub total: usize,
pub rules: Vec<ElementCountNode>,
}
impl ElementCountNode {
pub fn height(&self) -> usize {
1 + self
.children
.iter()
.map(Self::height)
.max()
.unwrap_or_default()
}
pub fn node_count(&self) -> usize {
1 + self.children.iter().map(Self::node_count).sum::<usize>()
}
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 {
pub fn rule_counts(&self) -> impl Iterator<Item = (&str, usize)> {
self.rules
.iter()
.map(|rule| (rule.name.as_str(), rule.count))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum ElementCountError {
InvalidProgram { message: String },
Unsupported {
kind: ElementNodeKind,
name: String,
span: Option<Span>,
reason: String,
},
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,
name_slot: 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 {
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,
name_slot: 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.name_slot = variable_slot(self.catalog, Kind::Action, action_name, 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.catalog,
Kind::Action,
action_name,
args.len(),
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 name_slot = std::mem::take(&mut self.name_slot);
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 name_slot { 1 } else { 2 };
self.value_node(node_id, "global variable", span, base, vec![], 0)
}
Value::PlayerVariable { player, .. } => {
let base = if name_slot { 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) {
2
} else if name == "array" || name == "evaluateOnce" {
2
} else {
1
};
let omitted =
omitted_default_cost(self.catalog, Kind::Value, name, args.len(), 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.name_slot =
owner.is_some_and(|(kind, name)| variable_slot(self.catalog, kind, name, 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(
catalog: &Catalog,
kind: Kind,
name: &str,
given: usize,
direct: bool,
) -> isize {
let Some(entry) = catalog.entry(kind, name) else {
return 0;
};
(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(catalog: &Catalog, kind: Kind, name: &str, index: usize) -> bool {
catalog.entry(kind, name).is_some_and(|entry| {
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| variable_slot(catalog, kind, name, index))
})
}
fn workshop_setting_adjustment(name: &str) -> isize {
match name {
"workshopSettingInteger"
| "workshopSettingFloat"
| "createWorkshopSettingInt"
| "createWorkshopSettingFloat" => -3,
"workshopSettingCombo" | "createWorkshopSettingEnum" => -2,
_ => 0,
}
}