Skip to main content

opt_lang/
budget.rs

1//! The work budget that bounds every pass.
2
3/// A budget of work steps that every pass draws from.
4///
5/// A step is one unit of analysis or rewrite work (roughly: one instruction,
6/// block, edge, or value visited). Each pass charges its cost before it changes
7/// anything; a pass that cannot be paid for does nothing and marks the budget
8/// exhausted, so the function is always left valid. Every pass is linear in the
9/// function except where noted on [`PassKind`](crate::PassKind), and those
10/// non-linear parts (nested loops in LICM, repeated re-examination in copy
11/// propagation, hash collisions in GVN) are charged step by step, so a hostile
12/// function cannot make optimization take more than `budget` steps.
13///
14/// # Examples
15///
16/// ```
17/// use opt_lang::Budget;
18///
19/// let mut b = Budget::new(10);
20/// assert_eq!(b.remaining(), 10);
21/// assert!(!b.is_exhausted());
22/// assert!(Budget::unlimited().remaining() > 1 << 60);
23/// # let _ = &mut b;
24/// ```
25#[derive(Clone, Copy, PartialEq, Eq, Debug)]
26pub struct Budget {
27    remaining: u64,
28    exhausted: bool,
29}
30
31impl Budget {
32    /// A budget of `steps` steps.
33    ///
34    /// # Examples
35    ///
36    /// ```
37    /// assert_eq!(opt_lang::Budget::new(5).remaining(), 5);
38    /// ```
39    #[must_use]
40    pub const fn new(steps: u64) -> Budget {
41        Budget {
42            remaining: steps,
43            exhausted: false,
44        }
45    }
46
47    /// A budget that never runs out in practice (`u64::MAX` steps).
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// assert!(!opt_lang::Budget::unlimited().is_exhausted());
53    /// ```
54    #[must_use]
55    pub const fn unlimited() -> Budget {
56        Budget::new(u64::MAX)
57    }
58
59    /// The steps left.
60    ///
61    /// # Examples
62    ///
63    /// ```
64    /// assert_eq!(opt_lang::Budget::new(0).remaining(), 0);
65    /// ```
66    #[must_use]
67    pub const fn remaining(&self) -> u64 {
68        self.remaining
69    }
70
71    /// Whether a pass was refused (or stopped early) because the budget ran out.
72    ///
73    /// # Examples
74    ///
75    /// ```
76    /// use ir_lang::{Linkage, Module, Signature};
77    /// use opt_lang::{Budget, PassKind};
78    ///
79    /// let mut m = Module::new("m");
80    /// let f = m.declare_function("f", &Signature::new(&[], &[]), Linkage::Export)?;
81    /// m.build(f)?.ret(&[])?;
82    /// let mut budget = Budget::new(0);
83    /// let changed = opt_lang::run_pass(&mut m.edit(f)?, PassKind::Dce, &mut budget)?;
84    /// assert!(!changed && budget.is_exhausted());
85    /// # Ok::<(), Box<dyn std::error::Error>>(())
86    /// ```
87    #[must_use]
88    pub const fn is_exhausted(&self) -> bool {
89        self.exhausted
90    }
91
92    /// Takes `steps` if they are available; otherwise marks the budget exhausted
93    /// and returns `false` (taking nothing).
94    pub(crate) fn charge(&mut self, steps: u64) -> bool {
95        if self.remaining >= steps {
96            self.remaining -= steps;
97            true
98        } else {
99            self.exhausted = true;
100            false
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn test_charge_takes_or_refuses() {
111        let mut b = Budget::new(10);
112        assert!(b.charge(4));
113        assert_eq!(b.remaining(), 6);
114        assert!(!b.charge(7));
115        assert!(b.is_exhausted());
116        assert_eq!(b.remaining(), 6);
117        assert!(b.charge(6));
118        assert_eq!(b.remaining(), 0);
119    }
120}