1use std::collections::HashMap;
12use std::fmt;
13
14use crate::catalog::{Catalog, Kind};
15use crate::core::source::Span;
16use crate::wir::{self, Action, ActionId, Program, Value, ValueId};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ElementNodeKind {
21 Rule,
22 Condition,
23 Action,
24 Value,
25}
26
27#[derive(Debug, Clone)]
29pub struct ElementCountNode {
30 pub kind: ElementNodeKind,
31 pub id: usize,
34 pub name: String,
36 pub span: Option<Span>,
38 pub base_count: usize,
40 pub adjustment: isize,
43 pub count: usize,
45 pub children: Vec<ElementCountNode>,
47}
48
49#[derive(Debug, Clone)]
51pub struct ElementCountReport {
52 pub total: usize,
54 pub rules: Vec<ElementCountNode>,
56}
57
58impl ElementCountReport {
59 pub fn rule_counts(&self) -> impl Iterator<Item = (&str, usize)> {
61 self.rules
62 .iter()
63 .map(|rule| (rule.name.as_str(), rule.count))
64 }
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum ElementCountError {
70 InvalidProgram { message: String },
73 Unsupported {
76 kind: ElementNodeKind,
77 name: String,
78 span: Option<Span>,
79 reason: String,
80 },
81 Cycle {
83 kind: ElementNodeKind,
84 id: usize,
86 },
87}
88
89impl fmt::Display for ElementCountError {
90 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
91 match self {
92 Self::InvalidProgram { message } => write!(formatter, "invalid program: {message}"),
93 Self::Unsupported {
94 kind,
95 name,
96 span,
97 reason,
98 } => write!(
99 formatter,
100 "unsupported {kind:?} '{name}'{}: {reason}",
101 span.map_or_else(String::new, |span| format!(" at {span:?}"))
102 ),
103 Self::Cycle { kind, id } => write!(formatter, "cyclic {kind:?} reference at {id}"),
104 }
105 }
106}
107
108impl std::error::Error for ElementCountError {}
109
110impl Program {
111 pub(crate) fn element_count(
117 &self,
118 catalog: &Catalog,
119 ) -> Result<ElementCountReport, ElementCountError> {
120 self.validate()
121 .map_err(|error| ElementCountError::InvalidProgram {
122 message: error.to_string(),
123 })?;
124 crate::rules::validate::validate_wir(self, catalog).map_err(|error| {
125 ElementCountError::InvalidProgram {
126 message: error.to_string(),
127 }
128 })?;
129
130 let mut counter = Counter {
131 program: self,
132 catalog,
133 values: HashMap::new(),
134 actions: HashMap::new(),
135 next_node_id: 0,
136 };
137 let mut rules = Vec::with_capacity(self.rules.len());
138 for rule in self.rules.iter() {
139 rules.push(counter.rule(rule)?);
140 }
141 let total = rules.iter().map(|rule| rule.count).sum();
142 Ok(ElementCountReport { total, rules })
143 }
144}
145
146impl crate::program::Program {
147 pub fn element_count(
149 &self,
150 catalog: &Catalog,
151 ) -> Result<ElementCountReport, ElementCountError> {
152 let storage = self
153 .to_wir()
154 .map_err(|error| ElementCountError::InvalidProgram {
155 message: error.to_string(),
156 })?;
157 storage.element_count(catalog)
158 }
159}
160
161struct Counted {
162 node: ElementCountNode,
163 heroes: usize,
164}
165
166impl Counted {
167 #[allow(clippy::too_many_arguments)]
168 fn finish(
169 kind: ElementNodeKind,
170 id: usize,
171 name: impl Into<String>,
172 span: Option<Span>,
173 base_count: usize,
174 adjustment: isize,
175 children: Vec<ElementCountNode>,
176 heroes: usize,
177 ) -> Self {
178 let children_count: usize = children.iter().map(|child| child.count).sum();
179 let count = (base_count as isize + children_count as isize + adjustment).max(0) as usize;
180 Self {
181 node: ElementCountNode {
182 kind,
183 id,
184 name: name.into(),
185 span,
186 base_count,
187 adjustment,
188 count,
189 children,
190 },
191 heroes,
192 }
193 }
194}
195
196struct Counter<'a> {
197 program: &'a Program,
198 catalog: &'a Catalog,
199 values: HashMap<usize, usize>,
200 actions: HashMap<usize, usize>,
201 next_node_id: usize,
202}
203
204impl Counter<'_> {
205 fn next_node_id(&mut self) -> usize {
206 let id = self.next_node_id;
207 self.next_node_id += 1;
208 id
209 }
210
211 fn rule(&mut self, rule: &wir::Rule) -> Result<ElementCountNode, ElementCountError> {
212 let node_id = self.next_node_id();
213 let mut children = Vec::with_capacity(rule.conditions.len() + rule.actions.len());
214 for condition in &rule.conditions {
215 if condition.disabled {
216 return Err(ElementCountError::Unsupported {
217 kind: ElementNodeKind::Condition,
218 name: "disabled condition".to_string(),
219 span: self
220 .program
221 .values
222 .get(condition.value)
223 .and_then(|v| v.span),
224 reason: "the element cost of a disabled condition is not established"
225 .to_string(),
226 });
227 }
228 children.push(self.condition(condition.value)?.node);
229 }
230 for action in &rule.actions {
231 children.push(self.action(*action)?.node);
232 }
233 Ok(Counted::finish(
234 ElementNodeKind::Rule,
235 node_id,
236 &rule.name,
237 rule.span,
238 1,
239 0,
240 children,
241 0,
242 )
243 .node)
244 }
245
246 fn condition(&mut self, id: ValueId) -> Result<Counted, ElementCountError> {
247 let node_id = self.next_node_id();
248 let Some(value) = self.program.values.get(id) else {
249 return Err(ElementCountError::InvalidProgram {
250 message: format!("dangling condition value {}", id.index()),
251 });
252 };
253 let (children, heroes) = match &value.value {
254 Value::Call { name, args } if is_comparison(name) => {
255 let mut children = Vec::with_capacity(args.len());
256 let mut heroes = 0;
257 for argument in args {
258 let counted = self.value(*argument, true)?;
259 heroes += counted.heroes;
260 children.push(counted.node);
261 }
262 (children, heroes)
263 }
264 _ => {
265 let counted = self.value(id, true)?;
266 (vec![counted.node], counted.heroes)
267 }
268 };
269 Ok(Counted::finish(
270 ElementNodeKind::Condition,
271 node_id,
272 "condition",
273 value.span,
274 1,
275 pair_surcharge(heroes),
276 children,
277 heroes,
278 ))
279 }
280
281 fn action(&mut self, id: ActionId) -> Result<Counted, ElementCountError> {
282 let node_id = self.next_node_id();
283 if let Some(&active_id) = self.actions.get(&id.index()) {
284 return Err(ElementCountError::Cycle {
285 kind: ElementNodeKind::Action,
286 id: active_id,
287 });
288 }
289 self.actions.insert(id.index(), node_id);
290 let Some(action) = self.program.actions.get(id) else {
291 return Err(ElementCountError::InvalidProgram {
292 message: format!("dangling action {}", id.index()),
293 });
294 };
295 let result = self.action_inner(action, node_id);
296 self.actions.remove(&id.index());
297 result
298 }
299
300 fn action_inner(
301 &mut self,
302 action: &Action,
303 node_id: usize,
304 ) -> Result<Counted, ElementCountError> {
305 let span = action.span();
306 let mut children = Vec::new();
307 let mut heroes = 0;
308 let name;
309 match action {
310 Action::SetGlobalVariable { value, .. }
311 | Action::ModifyGlobalVariable { value, .. } => {
312 name = "variable action";
313 self.push_action_value(&mut children, &mut heroes, *value)?;
314 }
315 Action::SetPlayerVariable { player, value, .. }
316 | Action::ModifyPlayerVariable { player, value, .. } => {
317 name = "player variable action";
318 self.push_action_value(&mut children, &mut heroes, *player)?;
319 self.push_action_value(&mut children, &mut heroes, *value)?;
320 }
321 Action::AssignMember { target, value, .. } => {
322 name = "member assignment";
323 self.push_action_value(&mut children, &mut heroes, *target)?;
324 self.push_action_value(&mut children, &mut heroes, *value)?;
325 }
326 Action::CallSubroutine { .. } => {
327 name = "call subroutine";
328 }
329 Action::If {
330 branches,
331 else_body,
332 ..
333 } => {
334 name = "if";
335 for branch in branches {
336 self.push_action_value(&mut children, &mut heroes, branch.condition)?;
337 for nested in &branch.body {
338 children.push(self.action(*nested)?.node);
339 }
340 }
341 if let Some(body) = else_body {
342 for nested in body {
343 children.push(self.action(*nested)?.node);
344 }
345 }
346 }
347 Action::While {
348 condition, body, ..
349 } => {
350 name = "while";
351 self.push_action_value(&mut children, &mut heroes, *condition)?;
352 for nested in body {
353 children.push(self.action(*nested)?.node);
354 }
355 }
356 Action::ForGlobalVariable {
357 start,
358 stop,
359 step,
360 body,
361 ..
362 } => {
363 name = "for global variable";
364 for value in [start, stop, step] {
365 self.push_action_value(&mut children, &mut heroes, *value)?;
366 }
367 for nested in body {
368 children.push(self.action(*nested)?.node);
369 }
370 }
371 Action::ForPlayerVariable {
372 player,
373 start,
374 stop,
375 step,
376 body,
377 ..
378 } => {
379 name = "for player variable";
380 for value in [player, start, stop, step] {
381 self.push_action_value(&mut children, &mut heroes, *value)?;
382 }
383 for nested in body {
384 children.push(self.action(*nested)?.node);
385 }
386 }
387 Action::Disabled { .. } => {
388 return Err(ElementCountError::Unsupported {
389 kind: ElementNodeKind::Action,
390 name: "disabled action".to_string(),
391 span,
392 reason: "the element cost of a disabled action is not established".to_string(),
393 });
394 }
395 Action::Call {
396 name: action_name,
397 args,
398 ..
399 } => {
400 if self.catalog.entry(Kind::Action, action_name).is_none() {
401 return Err(ElementCountError::Unsupported {
402 kind: ElementNodeKind::Action,
403 name: action_name.clone(),
404 span,
405 reason: "the action is not a catalog identity".to_string(),
406 });
407 }
408 name = action_name.as_str();
409 for argument in args {
410 self.push_action_value(&mut children, &mut heroes, *argument)?;
411 }
412 }
413 }
414 Ok(Counted::finish(
415 ElementNodeKind::Action,
416 node_id,
417 name,
418 span,
419 1,
420 pair_surcharge(heroes),
421 children,
422 heroes,
423 ))
424 }
425
426 fn push_action_value(
427 &mut self,
428 children: &mut Vec<ElementCountNode>,
429 heroes: &mut usize,
430 id: ValueId,
431 ) -> Result<(), ElementCountError> {
432 let counted = self.value(id, true)?;
433 *heroes += counted.heroes;
434 children.push(counted.node);
435 Ok(())
436 }
437
438 fn value(&mut self, id: ValueId, top_level: bool) -> Result<Counted, ElementCountError> {
439 let node_id = self.next_node_id();
440 if let Some(&active_id) = self.values.get(&id.index()) {
441 return Err(ElementCountError::Cycle {
442 kind: ElementNodeKind::Value,
443 id: active_id,
444 });
445 }
446 self.values.insert(id.index(), node_id);
447 let Some(value) = self.program.values.get(id) else {
448 return Err(ElementCountError::InvalidProgram {
449 message: format!("dangling value {}", id.index()),
450 });
451 };
452 let span = value.span;
453 let result = match &value.value {
454 Value::Number { .. } => self.value_node(node_id, "number", span, 1, vec![], 0),
455 Value::String(_) => self.value_node(node_id, "string", span, 1, vec![], 0),
456 Value::LocalizedString(_) => {
457 self.value_node(node_id, "localized string", span, 2, vec![], 0)
458 }
459 Value::Bool(_) => self.value_node(node_id, "boolean", span, 1, vec![], 0),
460 Value::Null => self.value_node(node_id, "null", span, 1, vec![], 0),
461 Value::Array(elements) => self.value_children(node_id, "array", span, 2, elements),
462 Value::Vector { x, y, z } => {
463 self.value_children(node_id, "vector", span, 1, &[*x, *y, *z])
464 }
465 Value::Enum { value_type, .. } => {
466 let heroes = usize::from(value_type == "Hero");
467 self.value_node(node_id, value_type, span, 1, vec![], heroes)
468 }
469 Value::GlobalVariable(_) => {
470 self.value_node(node_id, "global variable", span, 1, vec![], 0)
471 }
472 Value::PlayerVariable { player, .. } => {
473 self.value_children(node_id, "player variable", span, 1, &[*player])
474 }
475 Value::Subroutine(_) => self.value_node(node_id, "subroutine", span, 1, vec![], 0),
476 Value::EventPlayer => self.value_node(node_id, "event player", span, 1, vec![], 0),
477 Value::Call { name, args } => {
478 if name == crate::wir::AMBIGUOUS_ENUM_CALL
479 && crate::wir::ambiguous_enum_parts(self.program, id).is_some()
480 {
481 self.value_node(node_id, "ambiguous enum", span, 1, vec![], 0)
482 } else {
483 if name != "memberAccess"
484 && self.catalog.entry(Kind::Value, name).is_none()
485 && self.catalog.entry(Kind::Operator, name).is_none()
486 && !is_canonical_helper(name)
487 {
488 return Err(ElementCountError::Unsupported {
489 kind: ElementNodeKind::Value,
490 name: name.clone(),
491 span,
492 reason: "the value is not a catalog identity".to_string(),
493 });
494 }
495 let child_ids: Vec<ValueId> = if name == "memberAccess" {
496 args.first()
497 .copied()
498 .into_iter()
499 .chain(args.iter().copied().skip(2))
500 .collect()
501 } else {
502 args.clone()
503 };
504 let base = if name == "array"
505 || name == "evalOnce"
506 || name.starts_with("workshopSetting")
507 || name.starts_with("createWorkshopSetting")
508 {
509 2
510 } else {
511 1
512 };
513 self.value_children(node_id, name, span, base, &child_ids)
514 }
515 }
516 }?;
517 self.values.remove(&id.index());
518 let mut result = result;
519 if top_level {
520 result.node.adjustment -= 1;
521 result.node.count = (result.node.count as isize - 1).max(0) as usize;
522 }
523 Ok(result)
524 }
525
526 fn value_node(
527 &self,
528 id: usize,
529 name: impl Into<String>,
530 span: Option<Span>,
531 base: usize,
532 children: Vec<ElementCountNode>,
533 heroes: usize,
534 ) -> Result<Counted, ElementCountError> {
535 Ok(Counted::finish(
536 ElementNodeKind::Value,
537 id,
538 name,
539 span,
540 base,
541 0,
542 children,
543 heroes,
544 ))
545 }
546
547 fn value_children(
548 &mut self,
549 id: usize,
550 name: impl Into<String>,
551 span: Option<Span>,
552 base: usize,
553 ids: &[ValueId],
554 ) -> Result<Counted, ElementCountError> {
555 let mut children = Vec::with_capacity(ids.len());
556 let mut heroes = 0;
557 for child in ids {
558 let counted = self.value(*child, false)?;
559 heroes += counted.heroes;
560 children.push(counted.node);
561 }
562 self.value_node(id, name, span, base, children, heroes)
563 }
564}
565
566fn pair_surcharge(heroes: usize) -> isize {
567 (heroes / 2) as isize
568}
569
570fn is_comparison(name: &str) -> bool {
571 matches!(name, "==" | "!=" | "<" | "<=" | ">" | ">=")
572}
573
574fn is_canonical_helper(name: &str) -> bool {
575 matches!(
576 name,
577 "memberAccess"
578 | "+"
579 | "-"
580 | "*"
581 | "/"
582 | "%"
583 | "add"
584 | "subtract"
585 | "multiply"
586 | "divide"
587 | "modulo"
588 | "min"
589 | "max"
590 | "raiseToPower"
591 | "appendToArray"
592 | "removeFromArray"
593 | "removeFromArrayByValue"
594 | "removeFromArrayByIndex"
595 )
596}