netidx_bscript/
env.rs

1use crate::{
2    expr::{Arg, ModPath},
3    typ::{FnType, TVar, Type},
4    BindId, Ctx, InitFn, LambdaId, UserEvent,
5};
6use 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<C: Ctx, E: UserEvent> {
16    pub id: LambdaId,
17    pub env: Env<C, E>,
18    pub scope: ModPath,
19    pub argspec: Arc<[Arg]>,
20    pub typ: Arc<FnType>,
21    pub init: InitFn<C, E>,
22}
23
24impl<C: Ctx, E: UserEvent> fmt::Debug for LambdaDef<C, 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    scope: ModPath,
36    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(Clone)]
59pub struct TypeDef {
60    pub params: Arc<[(TVar, Option<Type>)]>,
61    pub typ: Type,
62}
63
64pub struct Env<C: Ctx, E: UserEvent> {
65    pub by_id: Map<BindId, Bind>,
66    pub lambdas: Map<LambdaId, Weak<LambdaDef<C, E>>>,
67    pub binds: Map<ModPath, Map<CompactString, BindId>>,
68    pub used: Map<ModPath, Arc<Vec<ModPath>>>,
69    pub modules: Set<ModPath>,
70    pub typedefs: Map<ModPath, Map<CompactString, TypeDef>>,
71}
72
73impl<C: Ctx, E: UserEvent> Clone for Env<C, E> {
74    fn clone(&self) -> Self {
75        Self {
76            by_id: self.by_id.clone(),
77            binds: self.binds.clone(),
78            used: self.used.clone(),
79            modules: self.modules.clone(),
80            typedefs: self.typedefs.clone(),
81            lambdas: self.lambdas.clone(),
82        }
83    }
84}
85
86impl<C: Ctx, E: UserEvent> Env<C, E> {
87    pub(super) fn new() -> Self {
88        Self {
89            by_id: Map::new(),
90            binds: Map::new(),
91            used: Map::new(),
92            modules: Set::new(),
93            typedefs: Map::new(),
94            lambdas: Map::new(),
95        }
96    }
97
98    pub(super) fn clear(&mut self) {
99        let Self { by_id, binds, used, modules, typedefs, lambdas } = self;
100        *by_id = Map::new();
101        *binds = Map::new();
102        *used = Map::new();
103        *modules = Set::new();
104        *typedefs = Map::new();
105        *lambdas = Map::new();
106    }
107
108    // restore the lexical environment to the state it was in at the
109    // snapshot `other`, but leave the bind and type environment
110    // alone.
111    pub(super) fn restore_lexical_env(&self, other: Self) -> Self {
112        Self {
113            binds: other.binds,
114            used: other.used,
115            modules: other.modules,
116            typedefs: other.typedefs,
117            by_id: self.by_id.clone(),
118            lambdas: self.lambdas.clone(),
119        }
120    }
121
122    pub fn find_visible<R, F: FnMut(&str, &str) -> Option<R>>(
123        &self,
124        scope: &ModPath,
125        name: &ModPath,
126        mut f: F,
127    ) -> Option<R> {
128        let mut buf = CompactString::from("");
129        let name_scope = Path::dirname(&**name);
130        let name = Path::basename(&**name).unwrap_or("");
131        for scope in Path::dirnames(&**scope).rev() {
132            let used = self.used.get(scope);
133            let used = iter::once(scope)
134                .chain(used.iter().flat_map(|s| s.iter().map(|p| &***p)));
135            for scope in used {
136                let scope = name_scope
137                    .map(|ns| {
138                        buf.clear();
139                        buf.push_str(scope);
140                        if let Some(Path::SEP) = buf.chars().next_back() {
141                            buf.pop();
142                        }
143                        buf.push_str(ns);
144                        buf.as_str()
145                    })
146                    .unwrap_or(scope);
147                if let Some(res) = f(scope, name) {
148                    return Some(res);
149                }
150            }
151        }
152        None
153    }
154
155    pub fn lookup_bind(
156        &self,
157        scope: &ModPath,
158        name: &ModPath,
159    ) -> Option<(&ModPath, &Bind)> {
160        self.find_visible(scope, name, |scope, name| {
161            self.binds.get_full(scope).and_then(|(scope, vars)| {
162                vars.get(name)
163                    .and_then(|bid| self.by_id.get(bid).map(|bind| (scope, bind)))
164            })
165        })
166    }
167
168    pub fn lookup_typedef(&self, scope: &ModPath, name: &ModPath) -> Option<&TypeDef> {
169        self.find_visible(scope, name, |scope, name| {
170            self.typedefs.get(scope).and_then(|m| m.get(name))
171        })
172    }
173
174    /// lookup binds in scope that match the specified partial
175    /// name. This is intended to be used for IDEs and interactive
176    /// shells, and is not used by the compiler.
177    pub fn lookup_matching(
178        &self,
179        scope: &ModPath,
180        part: &ModPath,
181    ) -> Vec<(CompactString, BindId)> {
182        let mut res = vec![];
183        self.find_visible(scope, part, |scope, part| {
184            if let Some(vars) = self.binds.get(scope) {
185                let r = vars.range::<str, _>((Bound::Included(part), Bound::Unbounded));
186                for (name, bind) in r {
187                    if name.starts_with(part) {
188                        res.push((name.clone(), *bind));
189                    }
190                }
191            }
192            None::<()>
193        });
194        res
195    }
196
197    /// lookup modules in scope that match the specified partial
198    /// name. This is intended to be used for IDEs and interactive
199    /// shells, and is not used by the compiler.
200    pub fn lookup_matching_modules(
201        &self,
202        scope: &ModPath,
203        part: &ModPath,
204    ) -> Vec<ModPath> {
205        let mut res = vec![];
206        self.find_visible(scope, part, |scope, part| {
207            let p = ModPath(Path::from(ArcStr::from(scope)).append(part));
208            for m in self.modules.range((Bound::Included(p.clone()), Bound::Unbounded)) {
209                if m.0.starts_with(&*p.0) {
210                    if let Some(m) = m.strip_prefix(scope) {
211                        if !m.trim().is_empty() {
212                            res.push(ModPath(Path::from(ArcStr::from(m))));
213                        }
214                    }
215                }
216            }
217            None::<()>
218        });
219        res
220    }
221
222    pub fn canonical_modpath(&self, scope: &ModPath, name: &ModPath) -> Option<ModPath> {
223        self.find_visible(scope, name, |scope, name| {
224            let p = ModPath(Path::from(ArcStr::from(scope)).append(name));
225            if self.modules.contains(&p) {
226                Some(p)
227            } else {
228                None
229            }
230        })
231    }
232
233    pub fn deftype(
234        &mut self,
235        scope: &ModPath,
236        name: &str,
237        params: Arc<[(TVar, Option<Type>)]>,
238        typ: Type,
239    ) -> Result<()> {
240        let defs = self.typedefs.get_or_default_cow(scope.clone());
241        if defs.get(name).is_some() {
242            bail!("{name} is already defined in scope {scope}")
243        } else {
244            thread_local! {
245                static KNOWN: RefCell<FxHashMap<ArcStr, TVar>> = RefCell::new(FxHashMap::default());
246                static DECLARED: RefCell<FxHashSet<ArcStr>> = RefCell::new(FxHashSet::default());
247            }
248            KNOWN.with_borrow_mut(|known| {
249                known.clear();
250                for (tv, tc) in params.iter() {
251                    Type::TVar(tv.clone()).alias_tvars(known);
252                    if let Some(tc) = tc {
253                        tc.alias_tvars(known);
254                    }
255                }
256                typ.alias_tvars(known);
257            });
258            DECLARED.with_borrow_mut(|declared| {
259                declared.clear();
260                for (tv, _) in params.iter() {
261                    if !declared.insert(tv.name.clone()) {
262                        bail!("duplicate type variable {tv} in definition of {name}");
263                    }
264                }
265                typ.check_tvars_declared(declared)?;
266                for (_, t) in params.iter() {
267                    if let Some(t) = t {
268                        t.check_tvars_declared(declared)?;
269                    }
270                }
271                Ok::<_, anyhow::Error>(())
272            })?;
273            KNOWN.with_borrow(|known| {
274                DECLARED.with_borrow(|declared| {
275                    for dec in declared {
276                        if !known.contains_key(dec) {
277                            bail!("unused type parameter {dec} in definition of {name}")
278                        }
279                    }
280                    Ok(())
281                })
282            })?;
283            defs.insert_cow(name.into(), TypeDef { params, typ });
284            Ok(())
285        }
286    }
287
288    pub fn undeftype(&mut self, scope: &ModPath, name: &str) {
289        if let Some(defs) = self.typedefs.get_mut_cow(scope) {
290            defs.remove_cow(&CompactString::from(name));
291            if defs.len() == 0 {
292                self.typedefs.remove_cow(scope);
293            }
294        }
295    }
296
297    // create a new binding. If an existing bind exists in the same
298    // scope shadow it.
299    pub fn bind_variable(&mut self, scope: &ModPath, name: &str, typ: Type) -> &mut Bind {
300        let binds = self.binds.get_or_default_cow(scope.clone());
301        let mut existing = true;
302        let id = binds.get_or_insert_cow(CompactString::from(name), || {
303            existing = false;
304            BindId::new()
305        });
306        if existing {
307            *id = BindId::new();
308        }
309        self.by_id.get_or_insert_cow(*id, || Bind {
310            export: true,
311            id: *id,
312            scope: scope.clone(),
313            doc: None,
314            name: CompactString::from(name),
315            typ,
316        })
317    }
318
319    pub fn unbind_variable(&mut self, id: BindId) {
320        if let Some(b) = self.by_id.remove_cow(&id) {
321            if let Some(binds) = self.binds.get_mut_cow(&b.scope) {
322                binds.remove_cow(&b.name);
323                if binds.len() == 0 {
324                    self.binds.remove_cow(&b.scope);
325                }
326            }
327        }
328    }
329}