1 //! Jump table representation.
2 //!
3 //! Jump tables are declared in the preamble and assigned an `ir::entities::JumpTable` reference.
4 //! The actual table of destinations is stored in a `JumpTableData` struct defined in this module.
5 
6 use crate::ir::BlockCall;
7 use crate::ir::instructions::ValueListPool;
8 use alloc::vec::Vec;
9 use core::fmt::{self, Display, Formatter};
10 use core::slice::{Iter, IterMut};
11 
12 #[cfg(feature = "enable-serde")]
13 use serde_derive::{Deserialize, Serialize};
14 
15 /// Contents of a jump table.
16 ///
17 /// All jump tables use 0-based indexing and are densely populated.
18 ///
19 /// The default block for the jump table is stored as the first element of the underlying vector.
20 /// It can be accessed through the `default_block` and `default_block_mut` functions. All blocks
21 /// may be iterated using the `all_branches` and `all_branches_mut` functions, which will both
22 /// iterate over the default block first.
23 #[derive(Debug, Clone, PartialEq, Hash)]
24 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
25 pub struct JumpTableData {
26     // Table entries.
27     table: Vec<BlockCall>,
28 }
29 
30 impl JumpTableData {
31     /// Create a new jump table with the provided blocks.
new(def: BlockCall, table: &[BlockCall]) -> Self32     pub fn new(def: BlockCall, table: &[BlockCall]) -> Self {
33         Self {
34             table: core::iter::once(def).chain(table.iter().copied()).collect(),
35         }
36     }
37 
38     /// Fetch the default block for this jump table.
default_block(&self) -> BlockCall39     pub fn default_block(&self) -> BlockCall {
40         *self.table.first().unwrap()
41     }
42 
43     /// Mutable access to the default block of this jump table.
default_block_mut(&mut self) -> &mut BlockCall44     pub fn default_block_mut(&mut self) -> &mut BlockCall {
45         self.table.first_mut().unwrap()
46     }
47 
48     /// The jump table and default block as a single slice. The default block will always be first.
all_branches(&self) -> &[BlockCall]49     pub fn all_branches(&self) -> &[BlockCall] {
50         self.table.as_slice()
51     }
52 
53     /// The jump table and default block as a single mutable slice. The default block will always
54     /// be first.
all_branches_mut(&mut self) -> &mut [BlockCall]55     pub fn all_branches_mut(&mut self) -> &mut [BlockCall] {
56         self.table.as_mut_slice()
57     }
58 
59     /// Access the jump table as a slice. This excludes the default block.
as_slice(&self) -> &[BlockCall]60     pub fn as_slice(&self) -> &[BlockCall] {
61         &self.table.as_slice()[1..]
62     }
63 
64     /// Access the jump table as a mutable slice. This excludes the default block.
as_mut_slice(&mut self) -> &mut [BlockCall]65     pub fn as_mut_slice(&mut self) -> &mut [BlockCall] {
66         &mut self.table.as_mut_slice()[1..]
67     }
68 
69     /// Returns an iterator to the jump table, excluding the default block.
70     #[deprecated(since = "7.0.0", note = "please use `.as_slice()` instead")]
iter(&self) -> Iter<'_, BlockCall>71     pub fn iter(&self) -> Iter<'_, BlockCall> {
72         self.as_slice().iter()
73     }
74 
75     /// Returns an iterator that allows modifying each value, excluding the default block.
76     #[deprecated(since = "7.0.0", note = "please use `.as_mut_slice()` instead")]
iter_mut(&mut self) -> IterMut<'_, BlockCall>77     pub fn iter_mut(&mut self) -> IterMut<'_, BlockCall> {
78         self.as_mut_slice().iter_mut()
79     }
80 
81     /// Clears all entries in this jump table, except for the default block.
clear(&mut self)82     pub fn clear(&mut self) {
83         self.table.drain(1..);
84     }
85 
86     /// Return a value that can display the contents of this jump table.
display<'a>(&'a self, pool: &'a ValueListPool) -> DisplayJumpTable<'a>87     pub fn display<'a>(&'a self, pool: &'a ValueListPool) -> DisplayJumpTable<'a> {
88         DisplayJumpTable { jt: self, pool }
89     }
90 }
91 
92 /// A wrapper for the context required to display a [JumpTableData].
93 pub struct DisplayJumpTable<'a> {
94     jt: &'a JumpTableData,
95     pool: &'a ValueListPool,
96 }
97 
98 impl<'a> Display for DisplayJumpTable<'a> {
fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result99     fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
100         write!(fmt, "{}, [", self.jt.default_block().display(self.pool))?;
101         if let Some((first, rest)) = self.jt.as_slice().split_first() {
102             write!(fmt, "{}", first.display(self.pool))?;
103             for block in rest {
104                 write!(fmt, ", {}", block.display(self.pool))?;
105             }
106         }
107         write!(fmt, "]")
108     }
109 }
110 
111 #[cfg(test)]
112 mod tests {
113     use super::JumpTableData;
114     use crate::entity::EntityRef;
115     use crate::ir::instructions::ValueListPool;
116     use crate::ir::{Block, BlockArg, BlockCall, Value};
117     use alloc::string::ToString;
118     use alloc::vec::Vec;
119 
120     #[test]
empty()121     fn empty() {
122         let mut pool = ValueListPool::default();
123         let def = BlockCall::new(Block::new(0), core::iter::empty(), &mut pool);
124 
125         let jt = JumpTableData::new(def, &[]);
126 
127         assert_eq!(jt.all_branches().get(0), Some(&def));
128 
129         assert_eq!(jt.as_slice().get(0), None);
130         assert_eq!(jt.as_slice().get(10), None);
131 
132         assert_eq!(jt.display(&pool).to_string(), "block0, []");
133 
134         assert_eq!(jt.all_branches(), [def]);
135         assert_eq!(jt.as_slice(), []);
136     }
137 
138     #[test]
insert()139     fn insert() {
140         let mut pool = ValueListPool::default();
141 
142         let v0 = Value::new(0);
143         let v1 = Value::new(1);
144 
145         let e0 = Block::new(0);
146         let e1 = Block::new(1);
147         let e2 = Block::new(2);
148 
149         let def = BlockCall::new(e0, core::iter::empty(), &mut pool);
150         let b1 = BlockCall::new(e1, core::iter::once(v0.into()), &mut pool);
151         let b2 = BlockCall::new(e2, core::iter::empty(), &mut pool);
152         let b3 = BlockCall::new(e1, core::iter::once(v1.into()), &mut pool);
153 
154         let jt = JumpTableData::new(def, &[b1, b2, b3]);
155 
156         assert_eq!(jt.default_block(), def);
157         assert_eq!(
158             jt.display(&pool).to_string(),
159             "block0, [block1(v0), block2, block1(v1)]"
160         );
161 
162         assert_eq!(jt.all_branches(), [def, b1, b2, b3]);
163         assert_eq!(jt.as_slice(), [b1, b2, b3]);
164 
165         assert_eq!(
166             jt.as_slice()[0].args(&pool).collect::<Vec<_>>(),
167             [BlockArg::Value(v0)]
168         );
169         assert_eq!(jt.as_slice()[1].args(&pool).collect::<Vec<_>>(), []);
170         assert_eq!(
171             jt.as_slice()[2].args(&pool).collect::<Vec<_>>(),
172             [BlockArg::Value(v1)]
173         );
174     }
175 }
176