charon_lib/
utils.rs

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

charon_lib/
utils.rs
