midenc_hir/
constants.rs

1use alloc::{collections::BTreeMap, sync::Arc, vec, vec::Vec};
2use core::{fmt, str::FromStr};
3
4use crate::define_attr_type;
5
6pub trait IntoBytes {
7    fn into_bytes(self) -> Vec<u8>;
8}
9impl IntoBytes for Vec<u8> {
10    #[inline(always)]
11    fn into_bytes(self) -> Vec<u8> {
12        self
13    }
14}
15impl IntoBytes for i8 {
16    #[inline]
17    fn into_bytes(self) -> Vec<u8> {
18        vec![self as u8]
19    }
20}
21impl IntoBytes for i16 {
22    #[inline]
23    fn into_bytes(self) -> Vec<u8> {
24        self.to_le_bytes().to_vec()
25    }
26}
27
28/// A handle to a constant
29#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
30pub struct ConstantId(u32);
31impl ConstantId {
32    pub fn new(id: usize) -> Self {
33        assert!(id <= u32::MAX as usize);
34        Self(id as u32)
35    }
36
37    #[inline]
38    pub const fn as_usize(&self) -> usize {
39        self.0 as usize
40    }
41}
42impl fmt::Debug for ConstantId {
43    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
44        fmt::Display::fmt(self, f)
45    }
46}
47impl fmt::Display for ConstantId {
48    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
49        write!(f, "const{}", &self.0)
50    }
51}
52
53define_attr_type!(ConstantId);
54
55impl crate::AttrPrinter for ConstantId {
56    fn print(
57        &self,
58        _flags: &crate::OpPrintingFlags,
59        context: &crate::Context,
60    ) -> crate::formatter::Document {
61        let data = context.get_constant(*self);
62        crate::formatter::display(data)
63    }
64}
65
66/// This type represents the raw data of a constant.
67///
68/// The data is expected to be in little-endian order.
69#[derive(Debug, Clone, PartialEq, Eq, Default, PartialOrd, Ord, Hash)]
70pub struct ConstantData(Vec<u8>);
71impl ConstantData {
72    /// Return the number of bytes in the constant.
73    pub fn len(&self) -> usize {
74        self.0.len()
75    }
76
77    /// Check if the constant contains any bytes.
78    pub fn is_empty(&self) -> bool {
79        self.0.is_empty()
80    }
81
82    /// Return the data as a slice.
83    pub fn as_slice(&self) -> &[u8] {
84        self.0.as_slice()
85    }
86
87    /// Append bytes to this constant
88    pub fn append(mut self, bytes: impl IntoBytes) -> Self {
89        let mut bytes = bytes.into_bytes();
90        self.0.append(&mut bytes);
91        self
92    }
93
94    /// Grow the size of the constant data in bytes to `expected_size`, zero-extending
95    /// the data by writing zeroes to the newly-added high-order bytes.
96    pub fn zext(mut self, expected_size: usize) -> Self {
97        assert!(
98            self.len() <= expected_size,
99            "the constant is already larger than {expected_size} bytes"
100        );
101        self.0.resize(expected_size, 0);
102        self
103    }
104
105    /// Attempt to convert this constant data to a `u32` value
106    pub fn as_u32(&self) -> Option<u32> {
107        let bytes = self.as_slice();
108        if bytes.len() != 4 {
109            return None;
110        }
111        let bytes = bytes.as_ptr() as *const [u8; 4];
112        Some(u32::from_le_bytes(unsafe { bytes.read() }))
113    }
114}
115impl From<ConstantData> for Vec<u8> {
116    #[inline(always)]
117    fn from(data: ConstantData) -> Self {
118        data.0
119    }
120}
121impl FromIterator<u8> for ConstantData {
122    fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
123        Self(iter.into_iter().collect())
124    }
125}
126impl From<Vec<u8>> for ConstantData {
127    fn from(v: Vec<u8>) -> Self {
128        Self(v)
129    }
130}
131impl<const N: usize> From<[u8; N]> for ConstantData {
132    fn from(v: [u8; N]) -> Self {
133        Self(v.to_vec())
134    }
135}
136impl From<&[u8]> for ConstantData {
137    fn from(v: &[u8]) -> Self {
138        Self(v.to_vec())
139    }
140}
141impl fmt::Display for ConstantData {
142    /// Print the constant data in hexadecimal format, e.g. 0x000102030405060708090a0b0c0d0e0f.
143    ///
144    /// The printed form of the constant renders the bytes in big-endian order, for readability.
145    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
146        fmt::LowerHex::fmt(self, f)
147    }
148}
149impl fmt::LowerHex for ConstantData {
150    /// Print the constant data in hexadecimal format, e.g. 0x000102030405060708090a0b0c0d0e0f.
151    ///
152    /// The printed form of the constant renders the bytes in the same order as the data.
153    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
154        if !self.is_empty() {
155            if f.alternate() {
156                f.write_str("0x")?;
157            }
158            for byte in self.0.iter().rev() {
159                write!(f, "{byte:02x}")?;
160            }
161        }
162        Ok(())
163    }
164}
165impl FromStr for ConstantData {
166    type Err = ();
167
168    #[inline]
169    fn from_str(s: &str) -> Result<Self, Self::Err> {
170        Self::from_str_be(s).map_err(|_| ())
171    }
172}
173impl ConstantData {
174    pub fn from_str_be(s: &str) -> Result<Self, &'static str> {
175        const NOT_EVEN: &str = "invalid hex-encoded data: expected an even number of hex digits";
176        const NOT_HEX: &str = "invalid hex-encoded data: contains invalid hex digits";
177
178        let s = s.strip_prefix("0x").unwrap_or(s);
179        let len = s.len();
180        if !len.is_multiple_of(2) {
181            return Err(NOT_EVEN);
182        }
183        // Parse big-endian
184        let pairs = len / 2;
185        let mut data = Vec::with_capacity(pairs);
186        let mut chars = s.chars();
187        while let Some(a) = chars.next() {
188            let a = a.to_digit(16).ok_or(NOT_HEX)?;
189            let b = chars.next().unwrap().to_digit(16).ok_or(NOT_HEX)?;
190            data.push(((a << 4) + b) as u8);
191        }
192
193        Ok(Self(data))
194    }
195
196    pub fn from_str_le(s: &str) -> Result<Self, &'static str> {
197        let mut data = Self::from_str_be(s)?;
198        // Make little-endian
199        data.0.reverse();
200        Ok(data)
201    }
202}
203
204/// This maintains the storage for constants used within a function
205#[derive(Default)]
206pub(crate) struct ConstantPool {
207    /// This mapping maintains the insertion order as long as Constants are created with
208    /// sequentially increasing integers.
209    ///
210    /// It is important that, by construction, no entry in that list gets removed. If that ever
211    /// need to happen, don't forget to update the `Constant` generation scheme.
212    constants: BTreeMap<ConstantId, Arc<ConstantData>>,
213
214    /// Mapping of hashed `ConstantData` to the index into the other hashmap.
215    ///
216    /// This allows for deduplication of entries into the `handles_to_values` mapping.
217    cache: BTreeMap<Arc<ConstantData>, ConstantId>,
218}
219impl ConstantPool {
220    /// Returns true if the pool is empty
221    #[allow(unused)]
222    pub fn is_empty(&self) -> bool {
223        self.constants.is_empty()
224    }
225
226    /// Returns the number of constants in this pool
227    pub fn len(&self) -> usize {
228        self.constants.len()
229    }
230
231    /// Retrieve the constant data as a reference-counted pointer, given a handle.
232    pub fn get(&self, id: ConstantId) -> Arc<ConstantData> {
233        Arc::clone(&self.constants[&id])
234    }
235
236    /// Retrieve the constant data by reference given a handle.
237    pub fn get_by_ref(&self, id: ConstantId) -> &ConstantData {
238        self.constants[&id].as_ref()
239    }
240
241    /// Returns true if this pool contains the given constant data
242    #[allow(unused)]
243    pub fn contains(&self, data: &ConstantData) -> bool {
244        self.cache.contains_key(data)
245    }
246
247    /// Insert constant data into the pool, returning a handle for later referencing; when constant
248    /// data is inserted that is a duplicate of previous constant data, the existing handle will be
249    /// returned.
250    pub fn insert(&mut self, data: ConstantData) -> ConstantId {
251        if let Some(cst) = self.cache.get(&data) {
252            return *cst;
253        }
254
255        let data = Arc::new(data);
256        let id = ConstantId::new(self.len());
257        self.constants.insert(id, Arc::clone(&data));
258        self.cache.insert(data, id);
259        id
260    }
261
262    /// Same as [ConstantPool::insert], but for data already allocated in an [Arc].
263    #[allow(unused)]
264    pub fn insert_arc(&mut self, data: Arc<ConstantData>) -> ConstantId {
265        if let Some(cst) = self.cache.get(data.as_ref()) {
266            return *cst;
267        }
268
269        let id = ConstantId::new(self.len());
270        self.constants.insert(id, Arc::clone(&data));
271        self.cache.insert(data, id);
272        id
273    }
274
275    /// Traverse the contents of the pool
276    #[inline]
277    #[allow(unused)]
278    pub fn iter(&self) -> impl Iterator<Item = (ConstantId, Arc<ConstantData>)> + '_ {
279        self.constants.iter().map(|(k, v)| (*k, Arc::clone(v)))
280    }
281}
midenc_hir/constants.rs

midenc_hir/
constants.rs
