graphix_compiler/
env.rs

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