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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use std::borrow::Borrow;
use std::fmt;
use std::hash::Hash;

use erg_common::dict::Dict;
use erg_common::set::Set;
use erg_common::shared::{
    MappedRwLockReadGuard, MappedRwLockWriteGuard, RwLockReadGuard, RwLockWriteGuard, Shared,
};
use erg_common::Str;

use crate::context::TraitImpl;

/// Caches checked modules.
/// In addition to being queried here when re-imported, it is also used when linking
/// (Erg links all scripts defined in erg and outputs them to a single pyc file).
#[derive(Debug, Default)]
pub struct TraitImpls {
    /// * key: trait qualified name
    /// * value: set of trait impls
    cache: Dict<Str, Set<TraitImpl>>,
}

impl fmt::Display for TraitImpls {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "TraitImpls {{")?;
        for (name, impls) in self.cache.iter() {
            writeln!(f, "{name}: {impls}, ")?;
        }
        write!(f, "}}")
    }
}

impl TraitImpls {
    pub fn new() -> Self {
        Self { cache: Dict::new() }
    }

    pub fn get<P: Eq + Hash + ?Sized>(&self, path: &P) -> Option<&Set<TraitImpl>>
    where
        Str: Borrow<P>,
    {
        self.cache.get(path)
    }

    pub fn get_mut<Q: Eq + Hash + ?Sized>(&mut self, path: &Q) -> Option<&mut Set<TraitImpl>>
    where
        Str: Borrow<Q>,
    {
        self.cache.get_mut(path)
    }

    pub fn register(&mut self, name: Str, impls: Set<TraitImpl>) {
        self.cache.insert(name, impls);
    }

    pub fn remove<Q: Eq + Hash + ?Sized>(&mut self, path: &Q) -> Option<Set<TraitImpl>>
    where
        Str: Borrow<Q>,
    {
        self.cache.remove(path)
    }

    pub fn rename<Q: Eq + Hash + ?Sized>(&mut self, old: &Q, new: Str)
    where
        Str: Borrow<Q>,
    {
        if let Some(impls) = self.remove(old) {
            self.register(new, impls);
        }
    }

    pub fn initialize(&mut self) {
        self.cache.clear();
    }
}

#[derive(Debug, Clone, Default)]
pub struct SharedTraitImpls(Shared<TraitImpls>);

impl fmt::Display for SharedTraitImpls {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Shared{}", self.0)
    }
}

impl SharedTraitImpls {
    pub fn new() -> Self {
        Self(Shared::new(TraitImpls::new()))
    }

    pub fn get<Q: Eq + Hash + ?Sized>(
        &self,
        path: &Q,
    ) -> Option<MappedRwLockReadGuard<Set<TraitImpl>>>
    where
        Str: Borrow<Q>,
    {
        if self.0.borrow().get(path).is_some() {
            Some(RwLockReadGuard::map(self.0.borrow(), |tis| {
                tis.get(path).unwrap()
            }))
        } else {
            None
        }
    }

    pub fn get_mut<Q: Eq + Hash + ?Sized>(
        &self,
        path: &Q,
    ) -> Option<MappedRwLockWriteGuard<Set<TraitImpl>>>
    where
        Str: Borrow<Q>,
    {
        if self.0.borrow().get(path).is_some() {
            Some(RwLockWriteGuard::map(self.0.borrow_mut(), |tis| {
                tis.get_mut(path).unwrap()
            }))
        } else {
            None
        }
    }

    pub fn register(&self, name: Str, impls: Set<TraitImpl>) {
        self.0.borrow_mut().register(name, impls);
    }

    pub fn remove<Q: Eq + Hash + ?Sized>(&self, path: &Q) -> Option<Set<TraitImpl>>
    where
        Str: Borrow<Q>,
    {
        self.0.borrow_mut().remove(path)
    }

    pub fn rename<Q: Eq + Hash + ?Sized>(&self, old: &Q, new: Str)
    where
        Str: Borrow<Q>,
    {
        self.0.borrow_mut().rename(old, new);
    }

    pub fn ref_inner(&self) -> MappedRwLockReadGuard<Dict<Str, Set<TraitImpl>>> {
        RwLockReadGuard::map(self.0.borrow(), |tis| &tis.cache)
    }

    pub fn initialize(&self) {
        self.0.borrow_mut().initialize();
    }
}