1use anyhow::{Context, Result, anyhow, bail};
15use arena_terms::{Arena, Term, View};
16use indexmap::IndexMap;
17use smartstring::alias::String;
18use std::collections::HashSet;
19use std::fmt;
20use std::str::FromStr;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23#[repr(u8)]
24pub enum Fixity {
25 Fun = 0,
26 Prefix = 1,
27 Infix = 2,
28 Postfix = 3,
29}
30
31impl Fixity {
32 pub const COUNT: usize = 4;
33 pub const STRS: &[&str] = &["fun", "prefix", "infix", "postfix"];
34}
35
36impl From<Fixity> for String {
37 fn from(f: Fixity) -> Self {
38 Fixity::STRS[Into::<usize>::into(f)].into()
39 }
40}
41
42impl From<Fixity> for usize {
43 fn from(f: Fixity) -> Self {
44 f as usize
45 }
46}
47
48impl fmt::Display for Fixity {
49 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50 f.write_str(String::from(*self).as_str())
51 }
52}
53
54#[derive(Debug, Clone)]
55pub struct ParseFixityError(String);
56
57impl fmt::Display for ParseFixityError {
58 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59 write!(f, "invalid fixity: {}", self.0)
60 }
61}
62impl std::error::Error for ParseFixityError {}
63
64impl FromStr for Fixity {
65 type Err = ParseFixityError;
66 fn from_str(s: &str) -> Result<Self, Self::Err> {
67 match s {
68 "fun" => Ok(Fixity::Fun),
69 "prefix" => Ok(Fixity::Prefix),
70 "infix" => Ok(Fixity::Infix),
71 "postfix" => Ok(Fixity::Postfix),
72 other => Err(ParseFixityError(String::from(other))),
73 }
74 }
75}
76
77impl TryFrom<&str> for Fixity {
78 type Error = ParseFixityError;
79 fn try_from(s: &str) -> Result<Self, Self::Error> {
80 s.parse()
81 }
82}
83
84impl TryFrom<String> for Fixity {
85 type Error = ParseFixityError;
86 fn try_from(s: String) -> Result<Self, Self::Error> {
87 s.as_str().parse()
88 }
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92#[repr(u8)]
93pub enum Assoc {
94 None = 0,
95 Left = 1,
96 Right = 2,
97}
98
99impl Assoc {
100 pub const COUNT: usize = 3;
101 pub const STRS: &[&str] = &["none", "left", "right"];
102}
103
104impl From<Assoc> for String {
105 fn from(a: Assoc) -> Self {
106 Assoc::STRS[Into::<usize>::into(a)].into()
107 }
108}
109
110impl From<Assoc> for usize {
111 fn from(a: Assoc) -> Self {
112 a as usize
113 }
114}
115
116impl fmt::Display for Assoc {
117 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118 f.write_str(String::from(*self).as_str())
119 }
120}
121
122#[derive(Debug, Clone)]
123pub struct ParseAssocError(String);
124
125impl fmt::Display for ParseAssocError {
126 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
127 write!(f, "invalid associativity: {}", self.0)
128 }
129}
130impl std::error::Error for ParseAssocError {}
131
132impl FromStr for Assoc {
133 type Err = ParseAssocError;
134
135 fn from_str(s: &str) -> Result<Self, Self::Err> {
136 match s {
137 "none" => Ok(Assoc::None),
138 "left" => Ok(Assoc::Left),
139 "right" => Ok(Assoc::Right),
140 other => Err(ParseAssocError(String::from(other))),
141 }
142 }
143}
144
145impl TryFrom<&str> for Assoc {
146 type Error = ParseAssocError;
147 fn try_from(s: &str) -> Result<Self, Self::Error> {
148 s.parse()
149 }
150}
151
152impl TryFrom<String> for Assoc {
153 type Error = ParseAssocError;
154 fn try_from(s: String) -> Result<Self, Self::Error> {
155 s.as_str().parse()
156 }
157}
158
159#[derive(Debug, Clone)]
160pub struct OperArg {
161 pub name: String, pub default: Option<Term>,
163}
164
165pub const NON_OPER_PREC: usize = 0;
166pub const MIN_OPER_PREC: usize = 0;
167pub const MAX_OPER_PREC: usize = 1200;
168
169#[derive(Debug, Clone)]
170pub struct OperDef {
171 pub fixity: Fixity,
172 pub prec: usize, pub assoc: Assoc, pub args: Vec<OperArg>, pub rename_to: Option<Term>, pub embed_fixity: bool,
177}
178
179#[derive(Debug, Clone)]
180pub struct OperDefTab {
181 tab: [Option<OperDef>; Fixity::COUNT],
182}
183
184#[derive(Debug, Clone, Default)]
185pub struct OperDefs {
186 map: IndexMap<String, OperDefTab>,
187}
188
189static EMPTY_OPER_DEF_TAB: OperDefTab = OperDefTab::new();
190
191impl OperDef {
192 pub fn required_arity(fixity: Fixity) -> usize {
193 match fixity {
194 Fixity::Fun => 0,
195 Fixity::Prefix => 1,
196 Fixity::Infix => 2,
197 Fixity::Postfix => 1,
198 }
199 }
200}
201
202impl OperDefTab {
203 pub const fn new() -> Self {
204 Self {
205 tab: [const { None }; Fixity::COUNT],
206 }
207 }
208
209 pub fn is_fun(&self) -> bool {
210 self.tab[0].is_some()
211 }
212
213 pub fn is_oper(&self) -> bool {
214 self.tab[1..].iter().any(|x| x.is_some())
215 }
216
217 pub fn get_op_def(&self, fixity: Fixity) -> Option<&OperDef> {
218 self.tab[usize::from(fixity)].as_ref()
219 }
220}
221
222impl std::ops::Index<Fixity> for OperDefTab {
223 type Output = Option<OperDef>;
224 fn index(&self, i: Fixity) -> &Self::Output {
225 let i: usize = i.into();
226 &self.tab[i]
227 }
228}
229
230impl std::ops::IndexMut<Fixity> for OperDefTab {
231 fn index_mut(&mut self, i: Fixity) -> &mut Self::Output {
232 let i: usize = i.into();
233 &mut self.tab[i]
234 }
235}
236
237impl OperDefs {
238 pub fn new() -> Self {
239 Self {
240 map: IndexMap::new(),
241 }
242 }
243
244 pub fn try_from_ops(arena: &Arena, ops: Term) -> Result<Self> {
245 let mut oper_defs = Self::new();
246 oper_defs.define_opers(arena, ops)?;
247 Ok(oper_defs)
248 }
249
250 pub fn len(&self) -> usize {
251 self.map.len()
252 }
253
254 pub fn lookup(&self, name: &str) -> Option<usize> {
255 self.map.get_index_of(name)
256 }
257
258 pub fn get(&self, index: Option<usize>) -> &OperDefTab {
259 match index {
260 Some(index) => match self.map.get_index(index) {
261 Some((_, tab)) => tab,
262 None => &EMPTY_OPER_DEF_TAB,
263 },
264 None => &EMPTY_OPER_DEF_TAB,
265 }
266 }
267
268 pub fn define_oper(&mut self, arena: &Arena, op: Term) -> Result<()> {
278 const BOOLS: &[&str] = &["false", "true"];
279
280 let (_, [oper, fixity, prec, assoc, rename_to, embed_fixity]) =
281 op.unpack_func(arena, &["op"])?;
282
283 let (functor, args) = oper.unpack_func_any(arena, &[])?;
284 let name = functor.atom_name(arena)?;
285
286 let fixity = Fixity::try_from(fixity.unpack_atom(arena, Fixity::STRS)?)?;
287 let prec = prec.unpack_int(arena)?.try_into()?;
288 let assoc = Assoc::try_from(assoc.unpack_atom(arena, Assoc::STRS)?)?;
289 let embed_fixity = embed_fixity.unpack_atom(arena, BOOLS)? == "true";
290
291 let args = args
292 .into_iter()
293 .map(|arg| {
294 Ok(match arg.view(arena)? {
295 View::Atom(name) => OperArg {
296 name: String::from(name),
297 default: None,
298 },
299 View::Func(ar, _, _) => {
300 let (_, [name, term]) = arg.unpack_func(ar, &["="])?;
301 OperArg {
302 name: String::from(name.atom_name(ar)?),
303 default: Some(term),
304 }
305 }
306 _ => bail!("oper arg must be an atom or =(atom, term) in {:?}", name),
307 })
308 })
309 .collect::<Result<Vec<_>>>()?;
310
311 let required_arity = OperDef::required_arity(fixity);
312 if args.len() < required_arity {
313 bail!(
314 "operator {:?} requires at least {} argument(s)",
315 name,
316 required_arity
317 );
318 }
319
320 if args[..required_arity].iter().any(|x| x.default.is_some()) {
321 bail!("defaults are not allowed for required operator arguments");
322 }
323
324 let unique_arg_names: HashSet<_> = args.iter().map(|x| &x.name).cloned().collect();
325 if unique_arg_names.len() != args.len() {
326 bail!("duplicate arguments in {:?}", name);
327 }
328
329 let rename_to = match rename_to.view(arena)? {
330 View::Atom("none") => None,
331 View::Func(ar, _, _) => {
332 let (_, [rename_to]) = rename_to.unpack_func(ar, &["some"])?;
333 Some(rename_to)
334 }
335 _ => bail!("rename_to must be 'none' | some(atom)"),
336 };
337
338 if matches!(fixity, Fixity::Fun) && prec != NON_OPER_PREC {
339 bail!("{:?} must be assigned precedence 0", name);
340 }
341 if !matches!(fixity, Fixity::Fun) && (prec < MIN_OPER_PREC || prec > MAX_OPER_PREC) {
342 bail!(
343 "precedence {} is out of range for operator {:?} with type {:?} (expected {}–{})",
344 prec,
345 name,
346 fixity,
347 MIN_OPER_PREC,
348 MAX_OPER_PREC,
349 );
350 }
351 if matches!((fixity, assoc), (Fixity::Prefix, Assoc::Left))
352 || matches!((fixity, assoc), (Fixity::Postfix, Assoc::Right))
353 {
354 bail!(
355 "operator {:?} with type {:?} cannot have associativity {:?}",
356 name,
357 fixity,
358 assoc
359 );
360 }
361
362 #[cfg(false)]
365 if matches!((fixity, assoc), (Fixity::Fun, Assoc::Left | Assoc::Right)) {
366 bail!(
367 "{:?} with type {:?} cannot have associativity {:?}",
368 name,
369 fixity,
370 assoc
371 );
372 }
373
374 let tab = self
375 .map
376 .entry(String::from(name))
377 .or_insert_with(OperDefTab::new);
378
379 if matches!(fixity, Fixity::Fun) && tab.is_oper() {
380 bail!(
381 "cannot define {:?} with type {:?}; it is already defined as an operator with a different type",
382 name,
383 fixity,
384 );
385 }
386
387 if matches!(fixity, Fixity::Prefix | Fixity::Infix | Fixity::Postfix)
388 && tab.tab[Into::<usize>::into(Fixity::Fun)].is_some()
389 {
390 bail!(
391 "cannot define {:?} as an operator with type {:?}; it is already defined with type Fun",
392 name,
393 fixity,
394 );
395 }
396
397 if tab[fixity].is_some() {
398 bail!("cannot re-define {:?}", name);
399 }
400
401 tab[fixity] = Some(OperDef {
402 fixity,
403 prec,
404 assoc,
405 rename_to,
406 embed_fixity,
407 args,
408 });
409
410 Ok(())
411 }
412
413 pub fn define_opers(&mut self, arena: &Arena, term: Term) -> Result<()> {
414 match term.view(arena)? {
415 View::List(arena, ts, _) => {
416 for t in ts {
417 self.define_oper(arena, *t)?;
418 }
419 }
420 _ => {
421 self.define_oper(arena, term)?;
422 }
423 }
424 Ok(())
425 }
426}