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