1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use rnode::{Visit, VisitMut, VisitMutWith, VisitWith};
use rustc_hash::FxHashMap;

use crate::Type;

/// Replaces all types which matches `matcher` with `replacer`.
///
///  - `replacer` is called iff `matcher` returns `true`.
///
///  - `matcher` is optimization, and it **should not** be recursive because
/// `replace_type` invokes it recursively.
pub fn replace_type<M, R>(ty: &mut Type, matcher: M, replacer: R)
where
    M: Fn(&Type) -> bool,
    R: Fn(&mut Type) -> Option<Type>,
{
    let mut cache = FxHashMap::default();
    ty.visit_mut_with(&mut TypeReplacer {
        cache: &mut cache,

        matcher: &matcher,
        replacer: &replacer,
    });
}

type Cache = FxHashMap<*const (), bool>;

struct TypeReplacer<'a, M, R>
where
    M: Fn(&Type) -> bool,
    R: Fn(&mut Type) -> Option<Type>,
{
    cache: &'a mut Cache,

    matcher: &'a M,
    replacer: &'a R,
}

impl<M, R> VisitMut<Type> for TypeReplacer<'_, M, R>
where
    M: Fn(&Type) -> bool,
    R: Fn(&mut Type) -> Option<Type>,
{
    fn visit_mut(&mut self, ty: &mut Type) {
        {
            let mut finder = Finder {
                cache: self.cache,
                matcher: self.matcher,
                found: false,
            };
            ty.visit_with(&mut finder);
            if !finder.found {
                return;
            }
        }

        if (self.matcher)(ty) {
            if let Some(new_ty) = (self.replacer)(ty) {
                *ty = new_ty;
                return;
            }
        }

        ty.normalize_mut();
        ty.visit_mut_children_with(self);
    }
}

struct Finder<'a, M> {
    cache: &'a mut Cache,
    matcher: &'a M,
    found: bool,
}

impl<M> Visit<Type> for Finder<'_, M>
where
    M: Fn(&Type) -> bool,
{
    fn visit(&mut self, ty: &Type) {
        if (self.matcher)(ty) {
            self.found = true;
        }

        if self.found {
            return;
        }

        let key = ty as *const Type as *const ();

        if let Some(v) = self.cache.get(&key).copied() {
            self.found |= v;
            return;
        }

        ty.visit_children_with(self);

        self.cache.insert(key, self.found);
    }
}