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