Skip to main content

fidget_core/compiler/
reg_tape.rs

1//! Tape used for evaluation
2use crate::compiler::{RegOp, RegisterAllocator, SsaTape};
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6/// Low-level tape for use with the Fidget virtual machine (or to be lowered
7/// further into machine instructions).
8#[derive(Clone, Default, Serialize, Deserialize)]
9pub struct RegTape {
10    tape: Vec<RegOp>,
11
12    /// Total allocated slots
13    ///
14    /// This is a continuous space of registers (`0..N`) and memory (`N..`),
15    /// where `N` is the parameter in [`RegTape::new`].
16    pub(super) slot_count: u32,
17}
18
19impl RegTape {
20    /// Lowers the tape to assembly with a particular register limit
21    ///
22    /// Note that if you _also_ want to simplify the tape, it's more efficient
23    /// to use [`VmData::simplify`](crate::vm::VmData::simplify), which
24    /// simultaneously simplifies **and** performs register allocation in a
25    /// single pass.
26    pub fn new<const N: usize>(ssa: &SsaTape) -> Self {
27        let mut alloc = RegisterAllocator::<N>::new(ssa.len());
28        for &op in ssa.iter() {
29            alloc.op(op)
30        }
31        alloc.finalize()
32    }
33
34    /// Repacks registers by frequency (so that register 0 is the most frequent)
35    pub fn repack(&mut self) {
36        let map = self.repack_map();
37        for op in &mut self.tape {
38            op.visit_regs_mut(|reg| *reg = map[reg]);
39        }
40    }
41
42    /// Returns a map for register repacking
43    ///
44    /// The map repacks registers in the tape by frequency, so that register 0
45    /// is the most frequent.
46    pub fn repack_map(&self) -> HashMap<u8, u8> {
47        let mut reg_counts: HashMap<u8, usize> = HashMap::new();
48        for op in &self.tape {
49            op.visit_regs(|reg| *reg_counts.entry(reg).or_default() += 1);
50        }
51        let mut sorted = reg_counts
52            .into_iter()
53            .map(|(reg, count)| (std::cmp::Reverse(count), reg))
54            .collect::<Vec<_>>();
55        sorted.sort_unstable();
56        sorted
57            .into_iter()
58            .enumerate()
59            .map(|(i, (_count, reg))| (reg, u8::try_from(i).unwrap()))
60            .collect()
61    }
62
63    /// Builds a new empty tape
64    pub(crate) fn empty() -> Self {
65        Self {
66            tape: vec![],
67            slot_count: 0,
68        }
69    }
70
71    /// Resets this tape, retaining its allocations
72    pub fn reset(&mut self) {
73        self.tape.clear();
74        self.slot_count = 0;
75    }
76
77    /// Returns the number of unique register and memory locations that are used
78    /// by this tape.
79    #[inline]
80    pub fn slot_count(&self) -> usize {
81        self.slot_count as usize
82    }
83    /// Returns the number of elements in the tape
84    #[inline]
85    pub fn len(&self) -> usize {
86        self.tape.len()
87    }
88    /// Returns `true` if the tape contains no elements
89    #[inline]
90    pub fn is_empty(&self) -> bool {
91        self.tape.is_empty()
92    }
93    /// Returns a front-to-back iterator
94    ///
95    /// This is the opposite of evaluation order; it will visit the root of the
96    /// tree first, and end at the leaves.
97    #[inline]
98    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &RegOp> {
99        self.into_iter()
100    }
101    #[inline]
102    pub(crate) fn push(&mut self, op: RegOp) {
103        self.tape.push(op)
104    }
105}
106
107impl<'a> IntoIterator for &'a RegTape {
108    type Item = &'a RegOp;
109    type IntoIter = std::slice::Iter<'a, RegOp>;
110    fn into_iter(self) -> Self::IntoIter {
111        self.tape.iter()
112    }
113}