Skip to main content

charon_lib/
utils.rs

1use itertools::Itertools;
2use macros::EnumAsGetters;
3
4pub mod hash_cons;
5pub use hash_cons::*;
6
7pub static TAB_INCR: &str = "    ";
8
9/// Custom function to pretty-print elements from an iterator
10/// The output format is:
11/// ```text
12/// [
13///   elem_0,
14///   ...
15///   elem_n
16/// ]
17/// ```
18pub fn pretty_display_list<T>(
19    t_to_string: impl Fn(T) -> String,
20    it: impl IntoIterator<Item = T>,
21) -> String {
22    let mut elems = it
23        .into_iter()
24        .map(t_to_string)
25        .map(|x| format!("  {},\n", x))
26        .peekable();
27    if elems.peek().is_none() {
28        "[]".to_owned()
29    } else {
30        format!("[\n{}]", elems.format(""))
31    }
32}
33
34/// Implement `From` and `TryFrom` to wrap/unwrap enum variants with a single payload.
35#[macro_export]
36macro_rules! impl_from_enum {
37    ($enum:ident::$variant:ident($ty:ty)) => {
38        impl From<$ty> for $enum {
39            fn from(x: $ty) -> Self {
40                $enum::$variant(x)
41            }
42        }
43        impl TryFrom<$enum> for $ty {
44            type Error = ();
45            fn try_from(e: $enum) -> Result<Self, Self::Error> {
46                match e {
47                    $enum::$variant(x) => Ok(x),
48                    _ => Err(()),
49                }
50            }
51        }
52    };
53}
54
55/// Yield `None` then infinitely many `Some(x)`.
56pub fn repeat_except_first<T: Clone>(x: T) -> impl Iterator<Item = Option<T>> {
57    [None].into_iter().chain(std::iter::repeat(Some(x)))
58}
59
60/// An enum to manage potentially-cyclic computations.
61#[derive(Debug, EnumAsGetters)]
62pub enum CycleDetector<T> {
63    /// We haven't analyzed this yet.
64    Unprocessed,
65    /// Sentinel value that we set when starting the computation on an item. If we ever encounter
66    /// this, we know we encountered a loop that we can't handle.
67    Processing,
68    /// Sentinel value we put when encountering a cycle, so we can know that happened.
69    Cyclic,
70    /// The final result of the computation.
71    Processed(T),
72}
73
74impl<T> CycleDetector<T> {
75    /// If this item hadn't been processed, return `true` and record it as `Processing`, otherwise
76    /// return `false`. If this item is already processing, record a cycle.
77    pub fn start_processing(&mut self) -> bool {
78        match self {
79            CycleDetector::Unprocessed => {
80                *self = CycleDetector::Processing;
81                true
82            }
83            CycleDetector::Processing => {
84                *self = CycleDetector::Cyclic;
85                false
86            }
87            CycleDetector::Cyclic | CycleDetector::Processed(_) => false,
88        }
89    }
90
91    pub fn done_processing(&mut self, x: T) {
92        *self = CycleDetector::Processed(x)
93    }
94}
95
96impl<T> Default for CycleDetector<T> {
97    fn default() -> Self {
98        CycleDetector::Unprocessed
99    }
100}
101
102pub mod type_map {
103    use rustc_hash::FxHashMap;
104    use std::{
105        any::{Any, TypeId},
106        marker::PhantomData,
107    };
108
109    pub trait Mappable: Any + Send + Sync {}
110    impl<T> Mappable for T where T: Any + Send + Sync {}
111
112    pub trait Mapper {
113        type Value<T: Mappable>: Mappable;
114    }
115
116    /// A map that maps types to values in a generic manner: we store for each type `T` a value of
117    /// type `M::Value<T>`.
118    pub struct TypeMap<M> {
119        data: FxHashMap<TypeId, Box<dyn Mappable>>,
120        phantom: PhantomData<M>,
121    }
122
123    impl<M: Mapper> TypeMap<M> {
124        pub fn get<T: Mappable>(&self) -> Option<&M::Value<T>> {
125            self.data
126                .get(&TypeId::of::<T>())
127                // We must be careful to not accidentally cast the box itself as `dyn Any`.
128                .map(|val: &Box<dyn Mappable>| &**val)
129                .and_then(|val: &dyn Mappable| (val as &dyn Any).downcast_ref())
130        }
131
132        pub fn get_mut<T: Mappable>(&mut self) -> Option<&mut M::Value<T>> {
133            self.data
134                .get_mut(&TypeId::of::<T>())
135                // We must be careful to not accidentally cast the box itself as `dyn Any`.
136                .map(|val: &mut Box<dyn Mappable>| &mut **val)
137                .and_then(|val: &mut dyn Mappable| (val as &mut dyn Any).downcast_mut())
138        }
139
140        pub fn insert<T: Mappable>(&mut self, val: M::Value<T>) -> Option<Box<M::Value<T>>> {
141            self.data
142                .insert(TypeId::of::<T>(), Box::new(val))
143                .and_then(|val: Box<dyn Mappable>| (val as Box<dyn Any>).downcast().ok())
144        }
145
146        pub fn or_insert_with<T: Mappable>(
147            &mut self,
148            f: impl FnOnce() -> M::Value<T>,
149        ) -> &mut M::Value<T> {
150            if self.get::<T>().is_none() {
151                self.insert(f());
152            }
153            self.get_mut::<T>().unwrap()
154        }
155        pub fn or_default<T: Mappable>(&mut self) -> &mut M::Value<T>
156        where
157            M::Value<T>: Default,
158        {
159            self.or_insert_with(Default::default)
160        }
161    }
162
163    impl<M> Default for TypeMap<M> {
164        fn default() -> Self {
165            Self {
166                data: Default::default(),
167                phantom: Default::default(),
168            }
169        }
170    }
171}
172
173pub mod hash_by_addr {
174    use serde::{Deserialize, Serialize};
175    use std::{
176        hash::{Hash, Hasher},
177        ops::Deref,
178    };
179
180    /// A wrapper around a smart pointer that hashes and compares the contents by the address of
181    /// the pointee.
182    #[derive(Debug, Clone, Serialize, Deserialize)]
183    pub struct HashByAddr<T>(pub T);
184
185    impl<T: Deref> HashByAddr<T> {
186        fn addr(&self) -> *const T::Target {
187            self.0.deref()
188        }
189    }
190
191    impl<T: Eq + Deref> Eq for HashByAddr<T> {}
192
193    impl<T: PartialEq + Deref> PartialEq for HashByAddr<T> {
194        fn eq(&self, other: &Self) -> bool {
195            std::ptr::addr_eq(self.addr(), other.addr())
196        }
197    }
198
199    impl<T: Hash + Deref> Hash for HashByAddr<T> {
200        fn hash<H: Hasher>(&self, state: &mut H) {
201            self.addr().hash(state);
202        }
203    }
204}
205
206pub mod serialize_map_to_array {
207    use core::{fmt, marker::PhantomData};
208    use std::{
209        collections::hash_map::RandomState,
210        hash::{BuildHasher, Hash},
211    };
212
213    use indexmap::IndexMap as SeqHashMap;
214    use serde::{
215        Deserialize, Deserializer, Serialize,
216        de::{SeqAccess, Visitor},
217        ser::Serializer,
218    };
219    use serde_state::{DeserializeState, SerializeState};
220
221    #[derive(Serialize, Deserialize, SerializeState, DeserializeState)]
222    struct KeyValue<K, V> {
223        key: K,
224        value: V,
225    }
226
227    /// A converter between an `SeqHashMap` and a sequence of named key-value pairs.
228    pub struct SeqHashMapToArray<K, V, U = RandomState>(PhantomData<(K, V, U)>);
229
230    impl<K, V, U> SeqHashMapToArray<K, V, U> {
231        /// Serializes the given `map` to an array of named key-values.
232        pub fn serialize<S>(map: &SeqHashMap<K, V, U>, serializer: S) -> Result<S::Ok, S::Error>
233        where
234            K: Serialize,
235            V: Serialize,
236            S: Serializer,
237        {
238            serializer.collect_seq(map.into_iter().map(|(key, value)| KeyValue { key, value }))
239        }
240        pub fn serialize_state<S, State: ?Sized>(
241            map: &SeqHashMap<K, V, U>,
242            state: &State,
243            serializer: S,
244        ) -> Result<S::Ok, S::Error>
245        where
246            K: SerializeState<State>,
247            V: SerializeState<State>,
248            S: Serializer,
249        {
250            serializer.collect_seq(
251                map.into_iter().map(|(key, value)| {
252                    serde_state::WithState::new(KeyValue { key, value }, state)
253                }),
254            )
255        }
256
257        /// Deserializes from an array of named key-values.
258        pub fn deserialize<'de, D>(deserializer: D) -> Result<SeqHashMap<K, V, U>, D::Error>
259        where
260            K: Deserialize<'de> + Eq + Hash,
261            V: Deserialize<'de>,
262            U: BuildHasher + Default,
263            D: Deserializer<'de>,
264        {
265            struct SeqHashMapToArrayVisitor<K, V, U>(PhantomData<(K, V, U)>);
266
267            impl<'de, K, V, U> Visitor<'de> for SeqHashMapToArrayVisitor<K, V, U>
268            where
269                K: Deserialize<'de> + Eq + Hash,
270                V: Deserialize<'de>,
271                U: BuildHasher + Default,
272            {
273                type Value = SeqHashMap<K, V, U>;
274
275                fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
276                    formatter.write_str("a list of key-value objects")
277                }
278
279                fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
280                    let mut map = SeqHashMap::<K, V, U>::default();
281                    while let Some(entry) = seq.next_element::<KeyValue<K, V>>()? {
282                        map.insert(entry.key, entry.value);
283                    }
284                    Ok(map)
285                }
286            }
287            let map =
288                deserializer.deserialize_seq(SeqHashMapToArrayVisitor::<K, V, U>(PhantomData))?;
289            Ok(map)
290        }
291        /// Deserializes from an array of named key-values.
292        pub fn deserialize_state<'de, D, State: ?Sized>(
293            state: &State,
294            deserializer: D,
295        ) -> Result<SeqHashMap<K, V, U>, D::Error>
296        where
297            K: DeserializeState<'de, State> + Eq + Hash,
298            V: DeserializeState<'de, State>,
299            U: BuildHasher + Default,
300            D: Deserializer<'de>,
301        {
302            struct SeqHashMapToArrayVisitor<'a, State: ?Sized, K, V, U>(
303                &'a State,
304                PhantomData<(K, V, U)>,
305            );
306
307            impl<'de, State: ?Sized, K, V, U> Visitor<'de> for SeqHashMapToArrayVisitor<'_, State, K, V, U>
308            where
309                K: DeserializeState<'de, State> + Eq + Hash,
310                V: DeserializeState<'de, State>,
311                U: BuildHasher + Default,
312            {
313                type Value = SeqHashMap<K, V, U>;
314
315                fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
316                    formatter.write_str("a list of key-value objects")
317                }
318
319                fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
320                    let mut map = SeqHashMap::default();
321                    let seed =
322                        serde_state::__private::wrap_deserialize_seed::<KeyValue<K, V>, _>(self.0);
323                    while let Some(entry) = seq.next_element_seed(seed)? {
324                        map.insert(entry.key, entry.value);
325                    }
326                    Ok(map)
327                }
328            }
329            let map = deserializer
330                .deserialize_seq(SeqHashMapToArrayVisitor::<_, K, V, U>(state, PhantomData))?;
331            Ok(map)
332        }
333    }
334}
335
336// This is the amount of bytes that need to be left on the stack before increasing the size. It
337// must be at least as large as the stack required by any code that does not call
338// `ensure_sufficient_stack`.
339const RED_ZONE: usize = 100 * 1024; // 100k
340
341// Only the first stack that is pushed, grows exponentially (2^n * STACK_PER_RECURSION) from then
342// on. Values taken from rustc.
343const STACK_PER_RECURSION: usize = 1024 * 1024; // 1MB
344
345/// Grows the stack on demand to prevent stack overflow. Call this in strategic locations to "break
346/// up" recursive calls. E.g. most statement visitors can benefit from this.
347#[inline]
348pub fn ensure_sufficient_stack<R>(f: impl FnOnce() -> R) -> R {
349    stacker::maybe_grow(RED_ZONE, STACK_PER_RECURSION, f)
350}
351
352/// Returns the values of the command-line options that match `find_arg`. The options are built-in
353/// to be of the form `--arg=value` or `--arg value`.
354pub fn arg_values<'a, T: AsRef<str>>(
355    args: &'a [T],
356    needle: &'a str,
357) -> impl Iterator<Item = &'a str> {
358    struct ArgFilter<'a, T> {
359        args: std::slice::Iter<'a, T>,
360        needle: &'a str,
361    }
362    impl<'a, T: AsRef<str>> Iterator for ArgFilter<'a, T> {
363        type Item = &'a str;
364        fn next(&mut self) -> Option<Self::Item> {
365            while let Some(arg) = self.args.next() {
366                let mut split_arg = arg.as_ref().splitn(2, '=');
367                if split_arg.next() == Some(self.needle) {
368                    return match split_arg.next() {
369                        // `--arg=value` form
370                        arg @ Some(_) => arg,
371                        // `--arg value` form
372                        None => self.args.next().map(|x| x.as_ref()),
373                    };
374                }
375            }
376            None
377        }
378    }
379    ArgFilter {
380        args: args.iter(),
381        needle,
382    }
383}
384
385pub fn arg_value<'a, T: AsRef<str>>(args: &'a [T], needle: &'a str) -> Option<&'a str> {
386    arg_values(args, needle).next()
387}