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 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 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 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 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 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 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}