Skip to main content

graphix_compiler/
env.rs

1use crate::{
2    expr::{ModPath, Sandbox},
3    typ::{TVar, Type},
4    BindId, Scope,
5};
6use anyhow::{anyhow, bail, Result};
7use arcstr::ArcStr;
8use compact_str::CompactString;
9use fxhash::{FxHashMap, FxHashSet};
10use immutable_chunkmap::{map::MapS as Map, set::SetS as Set};
11use netidx::path::Path;
12use poolshark::local::LPooled;
13use std::{fmt, iter, mem, ops::Bound};
14use triomphe::Arc;
15
16pub struct Bind {
17    pub id: BindId,
18    pub export: bool,
19    pub typ: Type,
20    pub doc: Option<ArcStr>,
21    pub scope: ModPath,
22    pub name: CompactString,
23}
24
25impl fmt::Debug for Bind {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        write!(f, "Bind {{ id: {:?}, export: {} }}", self.id, self.export,)
28    }
29}
30
31impl Clone for Bind {
32    fn clone(&self) -> Self {
33        Self {
34            id: self.id,
35            scope: self.scope.clone(),
36            name: self.name.clone(),
37            doc: self.doc.clone(),
38            export: self.export,
39            typ: self.typ.clone(),
40        }
41    }
42}
43
44#[derive(Debug, Clone)]
45pub struct TypeDef {
46    pub params: Arc<[(TVar, Option<Type>)]>,
47    pub typ: Type,
48    pub doc: Option<ArcStr>,
49}
50
51#[derive(Clone, Debug, Default)]
52pub struct Env {
53    pub by_id: Map<BindId, Bind>,
54    pub byref_chain: Map<BindId, BindId>,
55    pub binds: Map<ModPath, Map<CompactString, BindId>>,
56    pub used: Map<ModPath, Arc<Vec<ModPath>>>,
57    pub modules: Set<ModPath>,
58    pub typedefs: Map<ModPath, Map<CompactString, TypeDef>>,
59    pub catch: Map<ModPath, BindId>,
60}
61
62impl Env {
63    pub(super) fn clear(&mut self) {
64        let Self { by_id, binds, byref_chain, used, modules, typedefs, catch } = self;
65        *by_id = Map::new();
66        *binds = Map::new();
67        *byref_chain = Map::new();
68        *used = Map::new();
69        *modules = Set::new();
70        *typedefs = Map::new();
71        *catch = Map::new();
72    }
73
74    // restore the lexical environment to the state it was in at the
75    // snapshot `other`, but leave the bind and type environment
76    // alone.
77    pub(super) fn restore_lexical_env(&self, other: Self) -> Self {
78        Self {
79            binds: other.binds,
80            used: other.used,
81            modules: other.modules,
82            typedefs: other.typedefs,
83            by_id: self.by_id.clone(),
84            catch: self.catch.clone(),
85            byref_chain: self.byref_chain.clone(),
86        }
87    }
88
89    pub(super) fn restore_lexical_env_mut(&self, other: &mut Self) -> Self {
90        Self {
91            binds: mem::take(&mut other.binds),
92            used: mem::take(&mut other.used),
93            modules: mem::take(&mut other.modules),
94            typedefs: mem::take(&mut other.typedefs),
95            by_id: self.by_id.clone(),
96            catch: self.catch.clone(),
97            byref_chain: self.byref_chain.clone(),
98        }
99    }
100
101    pub fn apply_sandbox(&self, spec: &Sandbox) -> Result<Self> {
102        fn get_bind_name(n: &ModPath) -> Result<(&str, &str)> {
103            let dir = Path::dirname(&**n).ok_or_else(|| anyhow!("unknown module {n}"))?;
104            let k = Path::basename(&**n).ok_or_else(|| anyhow!("unknown module {n}"))?;
105            Ok((dir, k))
106        }
107        match spec {
108            Sandbox::Unrestricted => Ok(self.clone()),
109            Sandbox::Blacklist(bl) => {
110                let mut t = self.clone();
111                for n in bl.iter() {
112                    if t.modules.remove_cow(n) {
113                        t.binds.remove_cow(n);
114                        t.typedefs.remove_cow(n);
115                    } else {
116                        let (dir, k) = get_bind_name(n)?;
117                        let vals = t.binds.get_mut_cow(dir).ok_or_else(|| {
118                            anyhow!("no value {k} in module {dir} and no module {n}")
119                        })?;
120                        if let None = vals.remove_cow(&CompactString::from(k)) {
121                            bail!("no value {k} in module {dir} and no module {n}")
122                        }
123                    }
124                }
125                Ok(t)
126            }
127            Sandbox::Whitelist(wl) => {
128                let mut t = self.clone();
129                let mut modules = FxHashSet::default();
130                let mut names: FxHashMap<_, FxHashSet<_>> = FxHashMap::default();
131                for w in wl.iter() {
132                    if t.modules.contains(w) {
133                        modules.insert(w.clone());
134                    } else {
135                        let (dir, n) = get_bind_name(w)?;
136                        let dir = ModPath(Path::from(ArcStr::from(dir)));
137                        let n = CompactString::from(n);
138                        t.binds.get(&dir).and_then(|v| v.get(&n)).ok_or_else(|| {
139                            anyhow!("no value {n} in module {dir} and no module {w}")
140                        })?;
141                        names.entry(dir).or_default().insert(n);
142                    }
143                }
144                t.typedefs = t.typedefs.update_many(
145                    t.typedefs.into_iter().map(|(k, v)| (k.clone(), v.clone())),
146                    |k, v, _| {
147                        if modules.contains(&k) || names.contains_key(&k) {
148                            Some((k, v))
149                        } else {
150                            None
151                        }
152                    },
153                );
154                t.modules =
155                    t.modules.update_many(t.modules.into_iter().cloned(), |k, _| {
156                        if modules.contains(&k) || names.contains_key(&k) {
157                            Some(k)
158                        } else {
159                            None
160                        }
161                    });
162                t.binds = t.binds.update_many(
163                    t.binds.into_iter().map(|(k, v)| (k.clone(), v.clone())),
164                    |k, v, _| {
165                        if modules.contains(&k) {
166                            Some((k, v))
167                        } else if let Some(names) = names.get(&k) {
168                            let v = v.update_many(
169                                v.into_iter().map(|(k, v)| (k.clone(), v.clone())),
170                                |kn, vn, _| {
171                                    if names.contains(&kn) {
172                                        Some((kn, vn))
173                                    } else {
174                                        None
175                                    }
176                                },
177                            );
178                            Some((k, v))
179                        } else {
180                            None
181                        }
182                    },
183                );
184                Ok(t)
185            }
186        }
187    }
188
189    pub fn find_visible<T, F: FnMut(&str, &str) -> Option<T>>(
190        &self,
191        scope: &ModPath,
192        name: &ModPath,
193        mut f: F,
194    ) -> Option<T> {
195        let mut buf = CompactString::from("");
196        let name_scope = Path::dirname(&**name);
197        let name = Path::basename(&**name).unwrap_or("");
198        for scope in Path::dirnames(&**scope).rev() {
199            let used = self.used.get(scope);
200            let used = iter::once(scope)
201                .chain(used.iter().flat_map(|s| s.iter().map(|p| &***p)));
202            for scope in used {
203                let scope = name_scope
204                    .map(|ns| {
205                        buf.clear();
206                        buf.push_str(scope);
207                        if let Some(Path::SEP) = buf.chars().next_back() {
208                            buf.pop();
209                        }
210                        buf.push_str(ns);
211                        buf.as_str()
212                    })
213                    .unwrap_or(scope);
214                if let Some(res) = f(scope, name) {
215                    return Some(res);
216                }
217            }
218        }
219        None
220    }
221
222    pub fn lookup_bind(
223        &self,
224        scope: &ModPath,
225        name: &ModPath,
226    ) -> Option<(&ModPath, &Bind)> {
227        self.find_visible(scope, name, |scope, name| {
228            self.binds.get_full(scope).and_then(|(scope, vars)| {
229                vars.get(name)
230                    .and_then(|bid| self.by_id.get(bid).map(|bind| (scope, bind)))
231            })
232        })
233    }
234
235    pub fn lookup_typedef(&self, scope: &ModPath, name: &ModPath) -> Option<&TypeDef> {
236        self.find_visible(scope, name, |scope, name| {
237            self.typedefs.get(scope).and_then(|m| m.get(name))
238        })
239    }
240
241    /// lookup the bind id of the nearest catch handler in this scope
242    pub fn lookup_catch(&self, scope: &ModPath) -> Result<BindId> {
243        match Path::dirnames(&scope.0).rev().find_map(|scope| self.catch.get(scope)) {
244            Some(id) => Ok(*id),
245            None => bail!("there is no catch visible in {scope}"),
246        }
247    }
248
249    /// lookup binds in scope that match the specified partial
250    /// name. This is intended to be used for IDEs and interactive
251    /// shells, and is not used by the compiler.
252    pub fn lookup_matching(
253        &self,
254        scope: &ModPath,
255        part: &ModPath,
256    ) -> Vec<(CompactString, BindId)> {
257        let mut res = vec![];
258        self.find_visible(scope, part, |scope, part| {
259            if let Some(vars) = self.binds.get(scope) {
260                let r = vars.range::<str, _>((Bound::Included(part), Bound::Unbounded));
261                for (name, bind) in r {
262                    if name.starts_with(part) {
263                        res.push((name.clone(), *bind));
264                    }
265                }
266            }
267            None::<()>
268        });
269        res
270    }
271
272    /// lookup modules in scope that match the specified partial
273    /// name. This is intended to be used for IDEs and interactive
274    /// shells, and is not used by the compiler.
275    pub fn lookup_matching_modules(
276        &self,
277        scope: &ModPath,
278        part: &ModPath,
279    ) -> Vec<ModPath> {
280        let mut res = vec![];
281        self.find_visible(scope, part, |scope, part| {
282            let p = ModPath(Path::from(ArcStr::from(scope)).append(part));
283            for m in self.modules.range((Bound::Included(p.clone()), Bound::Unbounded)) {
284                if m.0.starts_with(&*p.0) {
285                    if let Some(m) = m.strip_prefix(scope) {
286                        if !m.trim().is_empty() {
287                            res.push(ModPath(Path::from(ArcStr::from(m))));
288                        }
289                    }
290                }
291            }
292            None::<()>
293        });
294        res
295    }
296
297    pub fn canonical_modpath(&self, scope: &ModPath, name: &ModPath) -> Option<ModPath> {
298        self.find_visible(scope, name, |scope, name| {
299            let p = ModPath(Path::from(ArcStr::from(scope)).append(name));
300            if self.modules.contains(&p) {
301                Some(p)
302            } else {
303                None
304            }
305        })
306    }
307
308    pub fn deftype(
309        &mut self,
310        scope: &ModPath,
311        name: &str,
312        params: Arc<[(TVar, Option<Type>)]>,
313        typ: Type,
314        doc: Option<ArcStr>,
315    ) -> Result<()> {
316        let defs = self.typedefs.get_or_default_cow(scope.clone());
317        if defs.get(name).is_some() {
318            bail!("{name} is already defined in scope {scope}")
319        } else {
320            let mut known: LPooled<FxHashMap<ArcStr, TVar>> = LPooled::take();
321            let mut declared: LPooled<FxHashSet<ArcStr>> = LPooled::take();
322            for (tv, tc) in params.iter() {
323                Type::TVar(tv.clone()).alias_tvars(&mut known);
324                if let Some(tc) = tc {
325                    tc.alias_tvars(&mut known);
326                }
327            }
328            typ.alias_tvars(&mut known);
329            for (tv, _) in params.iter() {
330                if !declared.insert(tv.name.clone()) {
331                    bail!("duplicate type variable {tv} in definition of {name}");
332                }
333            }
334            typ.check_tvars_declared(&mut declared)?;
335            for (_, t) in params.iter() {
336                if let Some(t) = t {
337                    t.check_tvars_declared(&mut declared)?;
338                }
339            }
340            for dec in declared.iter() {
341                if !known.contains_key(dec) {
342                    bail!("unused type parameter {dec} in definition of {name}")
343                }
344            }
345            defs.insert_cow(name.into(), TypeDef { params, typ, doc });
346            Ok(())
347        }
348    }
349
350    pub fn undeftype(&mut self, scope: &ModPath, name: &str) {
351        if let Some(defs) = self.typedefs.get_mut_cow(scope) {
352            defs.remove_cow(&CompactString::from(name));
353            if defs.len() == 0 {
354                self.typedefs.remove_cow(scope);
355            }
356        }
357    }
358
359    /// create a new binding. If an existing bind exists in the same
360    /// scope shadow it.
361    pub fn bind_variable(&mut self, scope: &ModPath, name: &str, typ: Type) -> &mut Bind {
362        let binds = self.binds.get_or_default_cow(scope.clone());
363        let mut existing = true;
364        let id = binds.get_or_insert_cow(CompactString::from(name), || {
365            existing = false;
366            BindId::new()
367        });
368        if existing {
369            *id = BindId::new();
370        }
371        self.by_id.get_or_insert_cow(*id, || Bind {
372            export: true,
373            id: *id,
374            scope: scope.clone(),
375            doc: None,
376            name: CompactString::from(name),
377            typ,
378        })
379    }
380
381    /// make the specified name an alias for `id`
382    pub fn alias_variable(&mut self, scope: &ModPath, name: &str, id: BindId) {
383        let binds = self.binds.get_or_default_cow(scope.clone());
384        binds.insert_cow(CompactString::from(name), id);
385    }
386
387    pub fn unbind_variable(&mut self, id: BindId) {
388        if let Some(b) = self.by_id.remove_cow(&id) {
389            if let Some(binds) = self.binds.get_mut_cow(&b.scope) {
390                binds.remove_cow(&b.name);
391                if binds.len() == 0 {
392                    self.binds.remove_cow(&b.scope);
393                }
394            }
395        }
396    }
397
398    pub fn use_in_scope(&mut self, scope: &Scope, name: &ModPath) -> Result<()> {
399        match self.canonical_modpath(&scope.lexical, name) {
400            None => bail!("use {name}: no such module {name}"),
401            Some(_) => {
402                let used = self.used.get_or_default_cow(scope.lexical.clone());
403                Ok(Arc::make_mut(used).push(name.clone()))
404            }
405        }
406    }
407
408    pub fn stop_use_in_scope(&mut self, scope: &Scope, name: &ModPath) {
409        if let Some(used) = self.used.get_mut_cow(&scope.lexical) {
410            Arc::make_mut(used).retain(|n| n != name);
411            if used.is_empty() {
412                self.used.remove_cow(&scope.lexical);
413            }
414        }
415    }
416}