charon_lib/
utils.rs

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

charon_lib/
utils.rs
