erg_compiler/
hir.rs

1/// defines High-level Intermediate Representation
2use std::fmt;
3use std::hash::{Hash, Hasher};
4
5use erg_common::consts::ERG_MODE;
6use erg_common::dict::Dict as HashMap;
7use erg_common::error::Location;
8#[allow(unused_imports)]
9use erg_common::log;
10use erg_common::pathutil::NormalizedPathBuf;
11use erg_common::set::Set as HashSet;
12use erg_common::traits::{Locational, NestedDisplay, NoTypeDisplay, Stream};
13use erg_common::{
14    enum_unwrap, fmt_option, fmt_option_map, fmt_vec, impl_display_for_enum,
15    impl_display_from_nested, impl_locational, impl_locational_for_enum,
16    impl_nested_display_for_chunk_enum, impl_nested_display_for_enum,
17    impl_no_type_display_for_enum, impl_stream,
18};
19use erg_common::{impl_from_trait_for_enum, impl_try_from_trait_for_enum, Str};
20
21use erg_parser::ast::{self, AscriptionKind};
22use erg_parser::ast::{
23    fmt_lines, DefId, DefKind, OperationKind, TypeBoundSpecs, TypeSpec, VarName,
24};
25use erg_parser::token::{Token, TokenKind, DOT};
26
27use crate::context::ControlKind;
28use crate::ty::constructors::{dict_t, set_t, tuple_t};
29use crate::ty::typaram::TyParam;
30use crate::ty::value::{GenTypeObj, ValueObj};
31use crate::ty::{Field, HasType, Type, VisibilityModifier};
32
33use crate::context::eval::type_from_token_kind;
34use crate::error::readable_name;
35use crate::varinfo::VarInfo;
36use crate::{impl_t, impl_t_for_enum};
37
38#[derive(Debug, Clone, PartialEq, Eq, Hash)]
39pub struct Literal {
40    pub value: ValueObj,
41    pub token: Token, // for Locational
42    t: Type,
43}
44
45impl_t!(Literal);
46
47impl NestedDisplay for Literal {
48    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
49        write!(f, "{} (: {})", self.token.content, self.t)
50    }
51}
52
53impl NoTypeDisplay for Literal {
54    fn to_string_notype(&self) -> String {
55        format!("{}", self.token.content)
56    }
57}
58
59impl_display_from_nested!(Literal);
60
61impl Locational for Literal {
62    #[inline]
63    fn loc(&self) -> Location {
64        self.token.loc()
65    }
66}
67
68impl TryFrom<Token> for Literal {
69    type Error = ();
70    fn try_from(token: Token) -> Result<Self, ()> {
71        let data =
72            ValueObj::from_str(type_from_token_kind(token.kind), token.content.clone()).ok_or(())?;
73        Ok(Self {
74            t: data.t(),
75            value: data,
76            token,
77        })
78    }
79}
80
81impl Literal {
82    #[inline]
83    pub fn is(&self, kind: TokenKind) -> bool {
84        self.token.is(kind)
85    }
86
87    #[inline]
88    pub fn is_doc_comment(&self) -> bool {
89        self.token.is(TokenKind::DocComment)
90    }
91}
92
93impl Literal {
94    pub fn new(value: ValueObj, token: Token) -> Self {
95        Self {
96            t: value.t(),
97            value,
98            token,
99        }
100    }
101}
102
103#[derive(Debug, Clone, PartialEq, Eq, Hash)]
104pub struct PosArg {
105    pub expr: Expr,
106}
107
108impl NestedDisplay for PosArg {
109    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result {
110        self.expr.fmt_nest(f, level)
111    }
112}
113
114impl NoTypeDisplay for PosArg {
115    fn to_string_notype(&self) -> String {
116        self.expr.to_string_notype()
117    }
118}
119
120impl_display_from_nested!(PosArg);
121
122impl Locational for PosArg {
123    fn loc(&self) -> Location {
124        self.expr.loc()
125    }
126}
127
128impl PosArg {
129    pub const fn new(expr: Expr) -> Self {
130        Self { expr }
131    }
132}
133
134#[derive(Debug, Clone, PartialEq, Eq, Hash)]
135pub struct KwArg {
136    pub keyword: Token,
137    pub expr: Expr,
138}
139
140impl NestedDisplay for KwArg {
141    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result {
142        writeln!(f, "{} := ", self.keyword)?;
143        self.expr.fmt_nest(f, level + 1)
144    }
145}
146
147impl NoTypeDisplay for KwArg {
148    fn to_string_notype(&self) -> String {
149        format!(
150            "{} := {}",
151            self.keyword.content,
152            self.expr.to_string_notype()
153        )
154    }
155}
156
157impl_display_from_nested!(KwArg);
158
159impl Locational for KwArg {
160    fn loc(&self) -> Location {
161        Location::concat(&self.keyword, &self.expr)
162    }
163}
164
165impl KwArg {
166    pub const fn new(keyword: Token, expr: Expr) -> Self {
167        Self { keyword, expr }
168    }
169}
170
171#[derive(Debug, Clone)]
172pub struct Args {
173    pub pos_args: Vec<PosArg>,
174    pub var_args: Option<Box<PosArg>>,
175    pub kw_args: Vec<KwArg>,
176    pub kw_var: Option<Box<PosArg>>,
177    pub paren: Option<(Location, Location)>,
178}
179
180impl PartialEq for Args {
181    fn eq(&self, other: &Self) -> bool {
182        self.pos_args == other.pos_args
183            && self.var_args == other.var_args
184            && self.kw_args == other.kw_args
185            && self.kw_var == other.kw_var
186    }
187}
188
189impl Eq for Args {}
190
191impl Hash for Args {
192    fn hash<H: Hasher>(&self, state: &mut H) {
193        self.pos_args.hash(state);
194        self.var_args.hash(state);
195        self.kw_args.hash(state);
196        self.kw_var.hash(state);
197    }
198}
199
200impl NestedDisplay for Args {
201    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result {
202        if !self.pos_args.is_empty() {
203            fmt_lines(self.pos_args.iter(), f, level)?;
204        }
205        if let Some(var_args) = &self.var_args {
206            write!(f, "*")?;
207            var_args.fmt_nest(f, level)?;
208        }
209        if !self.kw_args.is_empty() {
210            fmt_lines(self.kw_args.iter(), f, level)?;
211        }
212        if let Some(kw_var) = &self.kw_var {
213            write!(f, "**")?;
214            kw_var.fmt_nest(f, level)?;
215        }
216        Ok(())
217    }
218}
219
220impl NoTypeDisplay for Args {
221    fn to_string_notype(&self) -> String {
222        let mut s = String::new();
223        if !self.pos_args.is_empty() {
224            s += &self
225                .pos_args
226                .iter()
227                .map(|x| x.to_string_notype())
228                .fold("".to_string(), |acc, s| acc + &s + ", ");
229        }
230        if let Some(var_args) = &self.var_args {
231            s += &format!(", *{}", var_args.to_string_notype());
232        }
233        if !self.kw_args.is_empty() {
234            s += &self
235                .kw_args
236                .iter()
237                .map(|x| x.to_string_notype())
238                .fold("".to_string(), |acc, s| acc + &s + ", ");
239        }
240        if let Some(kw_var) = &self.kw_var {
241            s += &format!(", **{}", kw_var.to_string_notype());
242        }
243        s
244    }
245}
246
247// do not implement From<Vec<Expr>> to Args, because it will miss paren info (use `values`)
248
249impl_display_from_nested!(Args);
250
251impl Locational for Args {
252    fn loc(&self) -> Location {
253        if let Some((l, r)) = &self.paren {
254            let loc = Location::concat(l, r);
255            if !loc.is_unknown() {
256                return loc;
257            }
258        }
259        match (
260            self.pos_args.first().zip(self.pos_args.last()),
261            self.var_args.as_ref(),
262            self.kw_args.first().zip(self.kw_args.last()),
263        ) {
264            (Some((l, _)), _, Some((_, r))) => Location::concat(l, r),
265            (Some((l, _)), Some(r), None) => Location::concat(l, r.as_ref()),
266            (Some((l, r)), None, None) => Location::concat(l, r),
267            (None, Some(l), Some((_, r))) => Location::concat(l.as_ref(), r),
268            (None, None, Some((l, r))) => Location::concat(l, r),
269            (None, Some(l), None) => l.loc(),
270            (None, None, None) => Location::Unknown,
271        }
272    }
273}
274
275impl Args {
276    pub fn new(
277        pos_args: Vec<PosArg>,
278        var_args: Option<PosArg>,
279        kw_args: Vec<KwArg>,
280        kw_var: Option<PosArg>,
281        paren: Option<(Location, Location)>,
282    ) -> Self {
283        Self {
284            pos_args,
285            var_args: var_args.map(Box::new),
286            kw_args,
287            kw_var: kw_var.map(Box::new),
288            paren,
289        }
290    }
291
292    pub fn values(exprs: Vec<Expr>, paren: Option<(Location, Location)>) -> Self {
293        Self::pos_only(exprs.into_iter().map(PosArg::new).collect(), paren)
294    }
295
296    pub fn single(pos_arg: PosArg) -> Self {
297        Self::pos_only(vec![pos_arg], None)
298    }
299
300    pub fn pos_only(pos_args: Vec<PosArg>, paren: Option<(Location, Location)>) -> Self {
301        Self::new(pos_args, None, vec![], None, paren)
302    }
303
304    pub fn var_args(var_args: PosArg, paren: Option<(Location, Location)>) -> Self {
305        Self::new(vec![], Some(var_args), vec![], None, paren)
306    }
307
308    pub fn empty() -> Self {
309        Self::new(vec![], None, vec![], None, None)
310    }
311
312    /// pos_args.len() + kw_args.len()
313    #[inline]
314    pub fn len(&self) -> usize {
315        self.pos_args.len() + self.kw_args.len()
316    }
317
318    #[inline]
319    pub fn is_empty(&self) -> bool {
320        self.pos_args.is_empty() && self.var_args.is_none() && self.kw_args.is_empty()
321    }
322
323    #[inline]
324    pub fn kw_len(&self) -> usize {
325        self.kw_args.len()
326    }
327
328    pub fn push_pos(&mut self, pos: PosArg) {
329        self.pos_args.push(pos);
330    }
331
332    pub fn push_kw(&mut self, kw: KwArg) {
333        self.kw_args.push(kw);
334    }
335
336    pub fn remove(&mut self, index: usize) -> Expr {
337        if self.pos_args.get(index).is_some() {
338            self.pos_args.remove(index).expr
339        } else {
340            self.kw_args.remove(index - self.pos_args.len()).expr
341        }
342    }
343
344    /// try_remove((1, 2, z: 3), 2) == Some(3)
345    pub fn try_remove(&mut self, index: usize) -> Option<Expr> {
346        if self.pos_args.get(index).is_some() {
347            Some(self.pos_args.remove(index).expr)
348        } else {
349            self.kw_args.get(index - self.pos_args.len())?;
350            Some(self.kw_args.remove(index - self.pos_args.len()).expr)
351        }
352    }
353
354    pub fn try_remove_pos(&mut self, index: usize) -> Option<PosArg> {
355        self.pos_args.get(index)?;
356        Some(self.pos_args.remove(index))
357    }
358
359    pub fn try_remove_kw(&mut self, index: usize) -> Option<KwArg> {
360        self.kw_args.get(index)?;
361        Some(self.kw_args.remove(index))
362    }
363
364    pub fn get(&self, index: usize) -> Option<&Expr> {
365        if self.pos_args.get(index).is_some() {
366            self.pos_args.get(index).map(|a| &a.expr)
367        } else {
368            self.kw_args
369                .get(index - self.pos_args.len())
370                .map(|a| &a.expr)
371        }
372    }
373
374    pub fn last(&self) -> Option<&Expr> {
375        if self.kw_args.is_empty() {
376            self.pos_args.last().map(|a| &a.expr)
377        } else {
378            self.kw_args.last().map(|a| &a.expr)
379        }
380    }
381
382    pub fn remove_left_or_key(&mut self, key: &str) -> Option<Expr> {
383        if !self.pos_args.is_empty() {
384            Some(self.pos_args.remove(0).expr)
385        } else if let Some(pos) = self
386            .kw_args
387            .iter()
388            .position(|arg| &arg.keyword.inspect()[..] == key)
389        {
390            Some(self.kw_args.remove(pos).expr)
391        } else {
392            None
393        }
394    }
395
396    pub fn get_left_or_key(&self, key: &str) -> Option<&Expr> {
397        if !self.pos_args.is_empty() {
398            Some(&self.pos_args.first()?.expr)
399        } else if let Some(pos) = self
400            .kw_args
401            .iter()
402            .position(|arg| &arg.keyword.inspect()[..] == key)
403        {
404            Some(&self.kw_args.get(pos)?.expr)
405        } else {
406            None
407        }
408    }
409
410    pub fn get_mut_left_or_key(&mut self, key: &str) -> Option<&mut Expr> {
411        if !self.pos_args.is_empty() {
412            Some(&mut self.pos_args.get_mut(0)?.expr)
413        } else if let Some(pos) = self
414            .kw_args
415            .iter()
416            .position(|arg| &arg.keyword.inspect()[..] == key)
417        {
418            Some(&mut self.kw_args.get_mut(pos)?.expr)
419        } else {
420            None
421        }
422    }
423
424    pub fn insert_pos(&mut self, idx: usize, pos: PosArg) {
425        self.pos_args.insert(idx, pos);
426    }
427
428    pub fn get_kw(&self, keyword: &str) -> Option<&Expr> {
429        self.kw_args
430            .iter()
431            .find(|kw| kw.keyword.inspect() == keyword)
432            .map(|kw| &kw.expr)
433    }
434
435    pub fn iter(&self) -> impl Iterator<Item = &Expr> {
436        self.pos_args
437            .iter()
438            .map(|pos| &pos.expr)
439            .chain(self.var_args.iter().map(|var| &var.expr))
440            .chain(self.kw_args.iter().map(|kw| &kw.expr))
441            .chain(self.kw_var.iter().map(|kw| &kw.expr))
442    }
443}
444
445#[derive(Debug, Clone, PartialEq, Eq, Hash)]
446pub struct Identifier {
447    pub raw: ast::Identifier,
448    pub qual_name: Option<Str>,
449    pub vi: VarInfo,
450}
451
452impl NestedDisplay for Identifier {
453    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
454        write!(f, "{}", self.raw)?;
455        if let Some(qn) = &self.qual_name {
456            write!(f, "(qual_name: {qn})")?;
457        }
458        if self.vi.t != Type::Uninited {
459            write!(f, "(: {})", self.vi.t)?;
460        }
461        Ok(())
462    }
463}
464
465impl NoTypeDisplay for Identifier {
466    fn to_string_notype(&self) -> String {
467        self.raw.to_string()
468    }
469}
470
471impl_display_from_nested!(Identifier);
472
473impl HasType for Identifier {
474    #[inline]
475    fn ref_t(&self) -> &Type {
476        &self.vi.t
477    }
478    #[inline]
479    fn ref_mut_t(&mut self) -> Option<&mut Type> {
480        Some(&mut self.vi.t)
481    }
482    #[inline]
483    fn signature_t(&self) -> Option<&Type> {
484        None
485    }
486    #[inline]
487    fn signature_mut_t(&mut self) -> Option<&mut Type> {
488        None
489    }
490}
491
492impl Locational for Identifier {
493    fn loc(&self) -> Location {
494        self.raw.loc()
495    }
496}
497
498impl From<&Identifier> for Field {
499    fn from(ident: &Identifier) -> Self {
500        Self::new(ident.vis().clone(), ident.inspect().clone())
501    }
502}
503
504impl From<Identifier> for Expr {
505    fn from(ident: Identifier) -> Self {
506        Expr::Accessor(Accessor::Ident(ident))
507    }
508}
509
510impl Identifier {
511    pub const fn new(raw: ast::Identifier, qual_name: Option<Str>, vi: VarInfo) -> Self {
512        Self { raw, qual_name, vi }
513    }
514
515    pub fn static_public(name: &'static str) -> Self {
516        let ident = ast::Identifier::public_from_token(
517            Token::from_str(TokenKind::Dot, "."),
518            Token::static_symbol(name),
519        );
520        Self::bare(ident)
521    }
522
523    pub fn public(name: &str) -> Self {
524        let ident = ast::Identifier::public_from_token(
525            Token::from_str(TokenKind::Dot, "."),
526            Token::symbol(name),
527        );
528        Self::bare(ident)
529    }
530
531    pub fn private(name: &'static str) -> Self {
532        let ident = ast::Identifier::private_from_token(Token::static_symbol(name));
533        Self::bare(ident)
534    }
535
536    pub fn private_with_line(name: Str, line: u32) -> Self {
537        let ident = ast::Identifier::private_from_token(Token::symbol_with_line(&name, line));
538        Self::bare(ident)
539    }
540
541    pub fn public_with_line(dot: Token, name: Str, line: u32) -> Self {
542        let ident = ast::Identifier::public_from_token(dot, Token::symbol_with_line(&name, line));
543        Self::bare(ident)
544    }
545
546    pub const fn bare(ident: ast::Identifier) -> Self {
547        if ident.vis.is_public() {
548            Self::new(ident, None, VarInfo::const_default_public())
549        } else {
550            Self::new(ident, None, VarInfo::const_default_private())
551        }
552    }
553
554    pub fn call(self, args: Args) -> Call {
555        Call::new(Expr::Accessor(Accessor::Ident(self)), None, args)
556    }
557
558    pub fn method_call(self, attr_name: Identifier, args: Args) -> Call {
559        Call::new(Expr::Accessor(Accessor::Ident(self)), Some(attr_name), args)
560    }
561
562    pub fn is_py_api(&self) -> bool {
563        self.vi.py_name.is_some()
564    }
565
566    pub fn is_const(&self) -> bool {
567        self.raw.is_const()
568    }
569
570    pub fn is_discarded(&self) -> bool {
571        self.raw.is_discarded()
572    }
573
574    pub fn vis(&self) -> &VisibilityModifier {
575        &self.vi.vis.modifier
576    }
577
578    pub const fn inspect(&self) -> &Str {
579        self.raw.inspect()
580    }
581
582    /// show dot + name (no qual_name & type)
583    pub fn to_string_notype(&self) -> String {
584        NoTypeDisplay::to_string_notype(self)
585    }
586
587    pub fn is_procedural(&self) -> bool {
588        self.raw.is_procedural()
589    }
590
591    pub fn downgrade(self) -> ast::Identifier {
592        self.raw
593    }
594}
595
596#[derive(Debug, Clone, PartialEq, Eq, Hash)]
597pub struct Attribute {
598    pub obj: Box<Expr>,
599    pub ident: Identifier,
600}
601
602impl NestedDisplay for Attribute {
603    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
604        write!(f, "({}){}", self.obj, self.ident)
605    }
606}
607
608impl NoTypeDisplay for Attribute {
609    fn to_string_notype(&self) -> String {
610        format!(
611            "({}){}",
612            self.obj.to_string_notype(),
613            self.ident.to_string_notype()
614        )
615    }
616}
617
618impl_display_from_nested!(Attribute);
619impl_locational!(Attribute, obj, ident);
620
621impl HasType for Attribute {
622    #[inline]
623    fn ref_t(&self) -> &Type {
624        self.ident.ref_t()
625    }
626    #[inline]
627    fn ref_mut_t(&mut self) -> Option<&mut Type> {
628        self.ident.ref_mut_t()
629    }
630    #[inline]
631    fn signature_t(&self) -> Option<&Type> {
632        self.ident.signature_t()
633    }
634    #[inline]
635    fn signature_mut_t(&mut self) -> Option<&mut Type> {
636        self.ident.signature_mut_t()
637    }
638}
639
640impl Attribute {
641    pub fn new(obj: Expr, ident: Identifier) -> Self {
642        Self {
643            obj: Box::new(obj),
644            ident,
645        }
646    }
647}
648
649#[derive(Debug, Clone, PartialEq, Eq, Hash)]
650pub enum Accessor {
651    Ident(Identifier),
652    Attr(Attribute),
653}
654
655impl_nested_display_for_enum!(Accessor; Ident, Attr);
656impl_no_type_display_for_enum!(Accessor; Ident, Attr);
657impl_display_from_nested!(Accessor);
658impl_locational_for_enum!(Accessor; Ident, Attr);
659impl_t_for_enum!(Accessor; Ident, Attr);
660
661impl<'x> TryFrom<&'x Accessor> for &'x Identifier {
662    type Error = ();
663    fn try_from(acc: &'x Accessor) -> Result<Self, ()> {
664        match acc {
665            Accessor::Ident(ident) => Ok(ident),
666            _ => Err(()),
667        }
668    }
669}
670impl<'x> TryFrom<&'x Accessor> for &'x Attribute {
671    type Error = ();
672    fn try_from(acc: &'x Accessor) -> Result<Self, ()> {
673        match acc {
674            Accessor::Attr(attr) => Ok(attr),
675            _ => Err(()),
676        }
677    }
678}
679
680impl Accessor {
681    pub fn private_with_line(name: Str, line: u32) -> Self {
682        Self::Ident(Identifier::private_with_line(name, line))
683    }
684
685    pub fn public_with_line(name: Str, line: u32) -> Self {
686        Self::Ident(Identifier::public_with_line(DOT, name, line))
687    }
688
689    pub fn private(name: Token, vi: VarInfo) -> Self {
690        Self::Ident(Identifier::new(
691            ast::Identifier::private_from_token(name),
692            None,
693            vi,
694        ))
695    }
696
697    pub fn public(name: Token, vi: VarInfo) -> Self {
698        Self::Ident(Identifier::new(
699            ast::Identifier::public_from_token(DOT, name),
700            None,
701            vi,
702        ))
703    }
704
705    pub fn attr(obj: Expr, ident: Identifier) -> Self {
706        Self::Attr(Attribute::new(obj, ident))
707    }
708
709    pub fn var_info(&self) -> &VarInfo {
710        match self {
711            Self::Ident(ident) => &ident.vi,
712            Self::Attr(attr) => &attr.ident.vi,
713        }
714    }
715
716    pub fn show(&self) -> String {
717        match self {
718            Self::Ident(ident) => readable_name(ident.inspect()).to_string(),
719            Self::Attr(attr) => {
720                attr.obj
721                    .show_acc()
722                    .unwrap_or_else(|| attr.obj.ref_t().to_string())
723                    + "." // TODO: visibility
724                    + readable_name(attr.ident.inspect())
725            }
726        }
727    }
728
729    pub fn is_py_api(&self) -> bool {
730        match self {
731            Self::Ident(ident) => ident.is_py_api(),
732            Self::Attr(attr) => attr.ident.is_py_api(),
733        }
734    }
735
736    /// Unique name of the object (class, module, etc.) to which the accessor points
737    pub fn qual_name(&self) -> Option<Str> {
738        match self {
739            Self::Attr(attr) => attr.obj.qual_name().map(|obj_qn| {
740                if let Some(id_qn) = attr.ident.qual_name.as_ref() {
741                    Str::from(format!("{obj_qn}.{id_qn}"))
742                } else {
743                    Str::from(format!("{obj_qn}.{}", attr.ident.inspect()))
744                }
745            }),
746            Self::Ident(ident) => ident.qual_name.as_ref().cloned(),
747        }
748    }
749
750    pub fn local_name(&self) -> Option<&str> {
751        match self {
752            Self::Ident(ident) => ident
753                .qual_name
754                .as_ref()
755                .map(|s| {
756                    let mut seps = s.split_with(&[".", "::"]);
757                    seps.remove(seps.len() - 1)
758                })
759                .or_else(|| {
760                    let mut raw_parts = ident.inspect().split_with(&["'"]);
761                    // "'aaa'".split_with(&["'"]) == ["", "aaa", ""]
762                    if raw_parts.len() == 3 || raw_parts.len() == 4 {
763                        Some(raw_parts.remove(1))
764                    } else {
765                        Some(ident.inspect())
766                    }
767                }),
768            _ => None,
769        }
770    }
771
772    pub fn last_name(&self) -> &VarName {
773        match self {
774            Self::Ident(ident) => &ident.raw.name,
775            Self::Attr(attr) => &attr.ident.raw.name,
776        }
777    }
778
779    pub fn obj(&self) -> Option<&Expr> {
780        match self {
781            Self::Attr(attr) => Some(&attr.obj),
782            _ => None,
783        }
784    }
785
786    pub fn root_obj(&self) -> Option<&Expr> {
787        match self.obj() {
788            Some(expr @ Expr::Accessor(Accessor::Ident(_))) => Some(expr),
789            Some(Expr::Accessor(acc)) => acc.root_obj(),
790            Some(obj) => Some(obj),
791            None => None,
792        }
793    }
794
795    pub fn as_ident(&self) -> Option<&Identifier> {
796        match self {
797            Self::Ident(ident) => Some(ident),
798            _ => None,
799        }
800    }
801
802    pub fn as_attr(&self) -> Option<&Attribute> {
803        match self {
804            Self::Attr(attr) => Some(attr),
805            _ => None,
806        }
807    }
808}
809
810#[derive(Debug, Clone, PartialEq, Eq, Hash)]
811pub struct ListWithLength {
812    pub l_sqbr: Token,
813    pub r_sqbr: Token,
814    pub t: Type,
815    pub elem: Box<Expr>,
816    pub len: Option<Box<Expr>>,
817}
818
819impl NestedDisplay for ListWithLength {
820    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
821        write!(
822            f,
823            "[{}; {}](: {})",
824            self.elem,
825            fmt_option!(self.len, else "_"),
826            self.t
827        )
828    }
829}
830
831impl NoTypeDisplay for ListWithLength {
832    fn to_string_notype(&self) -> String {
833        format!(
834            "[{}; {}]",
835            self.elem.to_string_notype(),
836            fmt_option_map!(self.len, else "_", |len: &Expr| len.to_string_notype())
837        )
838    }
839}
840
841impl_display_from_nested!(ListWithLength);
842impl_locational!(ListWithLength, l_sqbr, elem, r_sqbr);
843impl_t!(ListWithLength);
844
845impl ListWithLength {
846    pub fn new(l_sqbr: Token, r_sqbr: Token, t: Type, elem: Expr, len: Option<Expr>) -> Self {
847        Self {
848            l_sqbr,
849            r_sqbr,
850            t,
851            elem: Box::new(elem),
852            len: len.map(Box::new),
853        }
854    }
855
856    pub const fn is_unsized(&self) -> bool {
857        self.len.is_none()
858    }
859}
860
861// TODO: generators
862#[derive(Debug, Clone, PartialEq, Eq, Hash)]
863pub struct ListComprehension {
864    pub l_sqbr: Token,
865    pub r_sqbr: Token,
866    pub t: Type,
867    pub elem: Box<Expr>,
868    pub guard: Box<Expr>,
869}
870
871impl NestedDisplay for ListComprehension {
872    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
873        write!(f, "[{} | {}](: {})", self.elem, self.guard, self.t)
874    }
875}
876
877impl NoTypeDisplay for ListComprehension {
878    fn to_string_notype(&self) -> String {
879        format!(
880            "[{} | {}]",
881            self.elem.to_string_notype(),
882            self.guard.to_string_notype()
883        )
884    }
885}
886
887impl_display_from_nested!(ListComprehension);
888impl_locational!(ListComprehension, l_sqbr, elem, r_sqbr);
889impl_t!(ListComprehension);
890
891#[derive(Debug, Clone, PartialEq, Eq, Hash)]
892pub struct NormalList {
893    pub l_sqbr: Token,
894    pub r_sqbr: Token,
895    pub t: Type,
896    pub elems: Args,
897}
898
899impl NestedDisplay for NormalList {
900    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
901        writeln!(f, "[")?;
902        self.elems.fmt_nest(f, level + 1)?;
903        write!(f, "\n{}](: {})", "    ".repeat(level), self.t)
904    }
905}
906
907impl NoTypeDisplay for NormalList {
908    fn to_string_notype(&self) -> String {
909        format!(
910            "[{}]",
911            self.elems
912                .pos_args
913                .iter()
914                .map(|arg| arg.to_string_notype())
915                .collect::<Vec<_>>()
916                .join(", ")
917        )
918    }
919}
920
921impl_display_from_nested!(NormalList);
922impl_locational!(NormalList, l_sqbr, elems, r_sqbr);
923impl_t!(NormalList);
924
925impl NormalList {
926    pub fn new(l_sqbr: Token, r_sqbr: Token, t: Type, elems: Args) -> Self {
927        Self {
928            l_sqbr,
929            r_sqbr,
930            t,
931            elems,
932        }
933    }
934
935    pub fn push(&mut self, elem: Expr) {
936        self.elems.push_pos(PosArg::new(elem));
937    }
938}
939
940#[derive(Debug, Clone, PartialEq, Eq, Hash)]
941pub enum List {
942    Normal(NormalList),
943    Comprehension(ListComprehension),
944    WithLength(ListWithLength),
945}
946
947impl_nested_display_for_enum!(List; Normal, Comprehension, WithLength);
948impl_no_type_display_for_enum!(List; Normal, Comprehension, WithLength);
949impl_display_for_enum!(List; Normal, Comprehension, WithLength);
950impl_locational_for_enum!(List; Normal, Comprehension, WithLength);
951impl_t_for_enum!(List; Normal, Comprehension, WithLength);
952
953impl List {
954    pub const fn is_unsized(&self) -> bool {
955        matches!(self, Self::WithLength(lis) if lis.is_unsized())
956    }
957}
958
959#[derive(Debug, Clone, PartialEq, Eq, Hash)]
960pub struct NormalTuple {
961    pub elems: Args,
962    pub(crate) t: Type,
963}
964
965impl NestedDisplay for NormalTuple {
966    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
967        writeln!(f, "(")?;
968        self.elems.fmt_nest(f, level + 1)?;
969        write!(f, "\n{})(: {})", "    ".repeat(level), self.t)
970    }
971}
972
973impl NoTypeDisplay for NormalTuple {
974    fn to_string_notype(&self) -> String {
975        format!(
976            "({})",
977            self.elems
978                .pos_args
979                .iter()
980                .map(|arg| arg.to_string_notype())
981                .collect::<Vec<_>>()
982                .join(", ")
983        )
984    }
985}
986
987impl_display_from_nested!(NormalTuple);
988impl_locational!(NormalTuple, elems);
989impl_t!(NormalTuple);
990
991impl NormalTuple {
992    pub fn new(elems: Args) -> Self {
993        let t = tuple_t(elems.pos_args.iter().map(|a| a.expr.t()).collect());
994        Self { elems, t }
995    }
996}
997
998#[derive(Debug, Clone, PartialEq, Eq, Hash)]
999pub enum Tuple {
1000    Normal(NormalTuple),
1001    // Comprehension(TupleComprehension),
1002}
1003
1004impl_nested_display_for_enum!(Tuple; Normal);
1005impl_no_type_display_for_enum!(Tuple; Normal);
1006impl_display_for_enum!(Tuple; Normal);
1007impl_locational_for_enum!(Tuple; Normal);
1008impl_t_for_enum!(Tuple; Normal);
1009
1010#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1011pub struct KeyValue {
1012    pub key: Expr,
1013    pub value: Expr,
1014}
1015
1016impl NestedDisplay for KeyValue {
1017    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1018        write!(f, "{}: {}", self.key, self.value)
1019    }
1020}
1021
1022impl NoTypeDisplay for KeyValue {
1023    fn to_string_notype(&self) -> String {
1024        format!(
1025            "{}: {}",
1026            self.key.to_string_notype(),
1027            self.value.to_string_notype()
1028        )
1029    }
1030}
1031
1032impl_display_from_nested!(KeyValue);
1033impl_locational!(KeyValue, key, value);
1034
1035impl KeyValue {
1036    pub const fn new(key: Expr, value: Expr) -> Self {
1037        Self { key, value }
1038    }
1039}
1040
1041#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1042pub struct NormalDict {
1043    pub l_brace: Token,
1044    pub r_brace: Token,
1045    pub t: Type,
1046    pub kvs: Vec<KeyValue>,
1047}
1048
1049impl_t!(NormalDict);
1050
1051impl NestedDisplay for NormalDict {
1052    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1053        write!(f, "{{{}}}(: {})", fmt_vec(&self.kvs), self.t)
1054    }
1055}
1056
1057impl NoTypeDisplay for NormalDict {
1058    fn to_string_notype(&self) -> String {
1059        format!(
1060            "{{{}}}",
1061            self.kvs
1062                .iter()
1063                .map(|kv| kv.to_string_notype())
1064                .collect::<Vec<_>>()
1065                .join(", ")
1066        )
1067    }
1068}
1069
1070impl_display_from_nested!(NormalDict);
1071impl_locational!(NormalDict, l_brace, r_brace);
1072
1073impl NormalDict {
1074    pub fn new(
1075        l_brace: Token,
1076        r_brace: Token,
1077        kv_ts: HashMap<TyParam, TyParam>,
1078        kvs: Vec<KeyValue>,
1079    ) -> Self {
1080        Self {
1081            l_brace,
1082            r_brace,
1083            t: dict_t(TyParam::Dict(kv_ts)),
1084            kvs,
1085        }
1086    }
1087}
1088
1089#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1090pub struct DictComprehension {
1091    pub l_sqbr: Token,
1092    pub r_sqbr: Token,
1093    pub t: Type,
1094    pub key: Box<Expr>,
1095    pub value: Box<Expr>,
1096    pub guard: Box<Expr>,
1097}
1098
1099impl NestedDisplay for DictComprehension {
1100    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1101        write!(
1102            f,
1103            "[{}: {} | {}](: {})",
1104            self.key, self.value, self.guard, self.t
1105        )
1106    }
1107}
1108
1109impl NoTypeDisplay for DictComprehension {
1110    fn to_string_notype(&self) -> String {
1111        format!(
1112            "[{}: {} | {}]",
1113            self.key.to_string_notype(),
1114            self.value.to_string_notype(),
1115            self.guard.to_string_notype()
1116        )
1117    }
1118}
1119
1120impl_display_from_nested!(DictComprehension);
1121impl_locational!(DictComprehension, l_sqbr, key, r_sqbr);
1122impl_t!(DictComprehension);
1123
1124#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1125pub enum Dict {
1126    Normal(NormalDict),
1127    Comprehension(DictComprehension),
1128}
1129
1130impl_nested_display_for_enum!(Dict; Normal, Comprehension);
1131impl_no_type_display_for_enum!(Dict; Normal, Comprehension);
1132impl_display_for_enum!(Dict; Normal, Comprehension);
1133impl_locational_for_enum!(Dict; Normal, Comprehension);
1134impl_t_for_enum!(Dict; Normal, Comprehension);
1135
1136impl Dict {
1137    pub fn empty() -> Self {
1138        let l_brace = Token::from_str(TokenKind::LBrace, "{");
1139        let r_brace = Token::from_str(TokenKind::RBrace, "}");
1140        Self::Normal(NormalDict::new(l_brace, r_brace, HashMap::new(), vec![]))
1141    }
1142}
1143
1144#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1145pub struct NormalSet {
1146    pub l_brace: Token,
1147    pub r_brace: Token,
1148    pub t: Type,
1149    pub elems: Args,
1150}
1151
1152impl NestedDisplay for NormalSet {
1153    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1154        writeln!(f, "{{")?;
1155        self.elems.fmt_nest(f, level + 1)?;
1156        write!(f, "\n{}}}(: {})", "    ".repeat(level), self.t)
1157    }
1158}
1159
1160impl NoTypeDisplay for NormalSet {
1161    fn to_string_notype(&self) -> String {
1162        format!(
1163            "{{{}}}",
1164            self.elems
1165                .pos_args
1166                .iter()
1167                .map(|e| e.to_string_notype())
1168                .collect::<Vec<_>>()
1169                .join(", ")
1170        )
1171    }
1172}
1173
1174impl_display_from_nested!(NormalSet);
1175impl_locational!(NormalSet, l_brace, elems, r_brace);
1176impl_t!(NormalSet);
1177
1178impl NormalSet {
1179    pub fn new(l_brace: Token, r_brace: Token, elem_t: Type, elems: Args) -> Self {
1180        let t = set_t(elem_t, TyParam::value(elems.len()));
1181        Self {
1182            l_brace,
1183            r_brace,
1184            t,
1185            elems,
1186        }
1187    }
1188}
1189
1190#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1191pub struct SetWithLength {
1192    pub l_brace: Token,
1193    pub r_brace: Token,
1194    pub t: Type,
1195    pub elem: Box<Expr>,
1196    pub len: Box<Expr>,
1197}
1198
1199impl NestedDisplay for SetWithLength {
1200    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1201        write!(f, "{{{}; {}}}(: {})", self.elem, self.len, self.t)
1202    }
1203}
1204
1205impl NoTypeDisplay for SetWithLength {
1206    fn to_string_notype(&self) -> String {
1207        format!(
1208            "{{{}; {}}}",
1209            self.elem.to_string_notype(),
1210            self.len.to_string_notype()
1211        )
1212    }
1213}
1214
1215impl_display_from_nested!(SetWithLength);
1216impl_locational!(SetWithLength, l_brace, elem, r_brace);
1217impl_t!(SetWithLength);
1218
1219impl SetWithLength {
1220    pub fn new(l_brace: Token, r_brace: Token, t: Type, elem: Expr, len: Expr) -> Self {
1221        Self {
1222            l_brace,
1223            r_brace,
1224            t,
1225            elem: Box::new(elem),
1226            len: Box::new(len),
1227        }
1228    }
1229}
1230
1231#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1232pub enum Set {
1233    Normal(NormalSet),
1234    WithLength(SetWithLength),
1235}
1236
1237impl_nested_display_for_enum!(Set; Normal, WithLength);
1238impl_no_type_display_for_enum!(Set; Normal, WithLength);
1239impl_display_for_enum!(Set; Normal, WithLength);
1240impl_locational_for_enum!(Set; Normal, WithLength);
1241impl_t_for_enum!(Set; Normal, WithLength);
1242
1243impl Set {
1244    pub fn empty() -> Self {
1245        let l_brace = Token::from_str(TokenKind::LBrace, "{");
1246        let r_brace = Token::from_str(TokenKind::RBrace, "}");
1247        Self::Normal(NormalSet::new(
1248            l_brace,
1249            r_brace,
1250            Type::Uninited,
1251            Args::empty(),
1252        ))
1253    }
1254}
1255
1256#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1257pub struct RecordAttrs(Vec<Def>);
1258
1259impl NestedDisplay for RecordAttrs {
1260    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1261        fmt_lines(self.0.iter(), f, level)
1262    }
1263}
1264
1265impl NoTypeDisplay for RecordAttrs {
1266    fn to_string_notype(&self) -> String {
1267        self.0
1268            .iter()
1269            .map(|a| a.to_string_notype())
1270            .collect::<Vec<_>>()
1271            .join("\n")
1272    }
1273}
1274
1275impl_display_from_nested!(RecordAttrs);
1276impl_stream!(RecordAttrs, Def);
1277
1278impl Locational for RecordAttrs {
1279    fn loc(&self) -> Location {
1280        if let Some((l, r)) = self.0.first().zip(self.0.last()) {
1281            Location::concat(l, r)
1282        } else {
1283            Location::Unknown
1284        }
1285    }
1286}
1287
1288impl From<Vec<Def>> for RecordAttrs {
1289    fn from(attrs: Vec<Def>) -> Self {
1290        Self(attrs)
1291    }
1292}
1293
1294impl RecordAttrs {
1295    pub fn len(&self) -> usize {
1296        self.0.len()
1297    }
1298
1299    pub fn is_empty(&self) -> bool {
1300        self.0.is_empty()
1301    }
1302
1303    pub fn push(&mut self, attr: Def) {
1304        self.0.push(attr);
1305    }
1306
1307    pub fn extend(&mut self, attrs: RecordAttrs) {
1308        self.0.extend(attrs.0);
1309    }
1310
1311    pub fn get(&self, name: &str) -> Option<&Def> {
1312        self.0.iter().find(|def| def.sig.ident().inspect() == name)
1313    }
1314}
1315
1316#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1317pub struct Record {
1318    l_brace: Token,
1319    r_brace: Token,
1320    pub attrs: RecordAttrs,
1321    pub(crate) t: Type,
1322}
1323
1324impl NestedDisplay for Record {
1325    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1326        writeln!(f, "{{")?;
1327        self.attrs.fmt_nest(f, level + 1)?;
1328        writeln!(f, "\n{}}}", "    ".repeat(level))
1329    }
1330}
1331
1332impl NoTypeDisplay for Record {
1333    fn to_string_notype(&self) -> String {
1334        format!(
1335            "{{{}}}",
1336            self.attrs
1337                .iter()
1338                .map(|a| a.to_string_notype())
1339                .collect::<Vec<_>>()
1340                .join("; ")
1341        )
1342    }
1343}
1344
1345impl_display_from_nested!(Record);
1346impl_locational!(Record, l_brace, attrs, r_brace);
1347impl_t!(Record);
1348
1349impl Record {
1350    pub fn new(l_brace: Token, r_brace: Token, attrs: RecordAttrs) -> Self {
1351        let rec = attrs
1352            .iter()
1353            .map(|def| (Field::from(def.sig.ident()), def.body.block.t()))
1354            .collect();
1355        let t = Type::Record(rec);
1356        Self {
1357            l_brace,
1358            r_brace,
1359            attrs,
1360            t,
1361        }
1362    }
1363
1364    pub fn push(&mut self, attr: Def) {
1365        let t = enum_unwrap!(&mut self.t, Type::Record);
1366        t.insert(Field::from(attr.sig.ident()), attr.body.block.t());
1367        self.attrs.push(attr);
1368    }
1369
1370    pub fn get(&self, name: &str) -> Option<&Def> {
1371        self.attrs.get(name)
1372    }
1373}
1374
1375#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1376pub struct BinOp {
1377    pub op: Token,
1378    pub lhs: Box<Expr>,
1379    pub rhs: Box<Expr>,
1380    pub info: VarInfo, // e.g. (Int, Int) -> Int
1381}
1382
1383impl NestedDisplay for BinOp {
1384    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1385        write!(
1386            f,
1387            "`{}`(: {})({}, {})",
1388            self.op.content, self.info.t, self.lhs, self.rhs
1389        )
1390    }
1391}
1392
1393impl NoTypeDisplay for BinOp {
1394    fn to_string_notype(&self) -> String {
1395        format!(
1396            "`{}`({}, {})",
1397            self.op.content,
1398            self.lhs.to_string_notype(),
1399            self.rhs.to_string_notype()
1400        )
1401    }
1402}
1403
1404impl HasType for BinOp {
1405    #[inline]
1406    fn ref_t(&self) -> &Type {
1407        self.info.t.return_t().unwrap()
1408    }
1409    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1410        self.info.t.mut_return_t()
1411    }
1412    #[inline]
1413    fn lhs_t(&self) -> &Type {
1414        self.info.t.lhs_t()
1415    }
1416    #[inline]
1417    fn rhs_t(&self) -> &Type {
1418        self.info.t.rhs_t()
1419    }
1420    #[inline]
1421    fn signature_t(&self) -> Option<&Type> {
1422        Some(&self.info.t)
1423    }
1424    #[inline]
1425    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1426        Some(&mut self.info.t)
1427    }
1428}
1429
1430impl_display_from_nested!(BinOp);
1431impl_locational!(BinOp, lhs, rhs);
1432
1433impl BinOp {
1434    pub fn new(op: Token, lhs: Expr, rhs: Expr, info: VarInfo) -> Self {
1435        Self {
1436            op,
1437            lhs: Box::new(lhs),
1438            rhs: Box::new(rhs),
1439            info,
1440        }
1441    }
1442}
1443
1444#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1445pub struct UnaryOp {
1446    pub op: Token,
1447    pub expr: Box<Expr>,
1448    pub info: VarInfo, // e.g. Neg -> Nat
1449}
1450
1451impl HasType for UnaryOp {
1452    #[inline]
1453    fn ref_t(&self) -> &Type {
1454        self.info.t.return_t().unwrap()
1455    }
1456    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1457        self.info.t.mut_return_t()
1458    }
1459    #[inline]
1460    fn lhs_t(&self) -> &Type {
1461        self.expr.ref_t()
1462    }
1463    #[inline]
1464    fn rhs_t(&self) -> &Type {
1465        panic!("invalid operation")
1466    }
1467    #[inline]
1468    fn signature_t(&self) -> Option<&Type> {
1469        Some(&self.info.t)
1470    }
1471    #[inline]
1472    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1473        Some(&mut self.info.t)
1474    }
1475}
1476
1477impl NestedDisplay for UnaryOp {
1478    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1479        writeln!(f, "`{}`(: {}):", self.op, self.info.t)?;
1480        self.expr.fmt_nest(f, level + 1)
1481    }
1482}
1483
1484impl NoTypeDisplay for UnaryOp {
1485    fn to_string_notype(&self) -> String {
1486        format!("`{}`({})", self.op.content, self.expr.to_string_notype())
1487    }
1488}
1489
1490impl_display_from_nested!(UnaryOp);
1491impl_locational!(UnaryOp, op, expr);
1492
1493impl UnaryOp {
1494    pub fn new(op: Token, expr: Expr, info: VarInfo) -> Self {
1495        Self {
1496            op,
1497            expr: Box::new(expr),
1498            info,
1499        }
1500    }
1501}
1502
1503#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1504pub struct Call {
1505    pub obj: Box<Expr>,
1506    pub attr_name: Option<Identifier>,
1507    pub args: Args,
1508}
1509
1510impl NestedDisplay for Call {
1511    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result {
1512        writeln!(f, "({}){}", self.obj, fmt_option!(self.attr_name),)?;
1513        if self.args.is_empty() {
1514            write!(f, "()")
1515        } else {
1516            writeln!(f, ":")?;
1517            self.args.fmt_nest(f, level + 1)
1518        }
1519    }
1520}
1521
1522impl NoTypeDisplay for Call {
1523    fn to_string_notype(&self) -> String {
1524        format!(
1525            "({}){}({})",
1526            self.obj.to_string_notype(),
1527            fmt_option!(self.attr_name),
1528            self.args.to_string_notype()
1529        )
1530    }
1531}
1532
1533impl_display_from_nested!(Call);
1534
1535impl HasType for Call {
1536    #[inline]
1537    fn ref_t(&self) -> &Type {
1538        if let Some(attr) = self.attr_name.as_ref() {
1539            attr.ref_t().return_t().unwrap_or(Type::FAILURE)
1540        } else {
1541            self.obj.ref_t().return_t().unwrap_or(Type::FAILURE)
1542        }
1543    }
1544    #[inline]
1545    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1546        if let Some(attr) = self.attr_name.as_mut() {
1547            attr.ref_mut_t()?.mut_return_t()
1548        } else {
1549            self.obj.ref_mut_t()?.mut_return_t()
1550        }
1551    }
1552    #[inline]
1553    fn lhs_t(&self) -> &Type {
1554        if let Some(attr) = self.attr_name.as_ref() {
1555            attr.ref_t().lhs_t()
1556        } else {
1557            self.obj.lhs_t()
1558        }
1559    }
1560    #[inline]
1561    fn rhs_t(&self) -> &Type {
1562        if let Some(attr) = self.attr_name.as_ref() {
1563            attr.ref_t().rhs_t()
1564        } else {
1565            self.obj.rhs_t()
1566        }
1567    }
1568    #[inline]
1569    fn signature_t(&self) -> Option<&Type> {
1570        if let Some(attr) = self.attr_name.as_ref() {
1571            Some(attr.ref_t())
1572        } else {
1573            Some(self.obj.ref_t())
1574        }
1575    }
1576    #[inline]
1577    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1578        if let Some(attr) = self.attr_name.as_mut() {
1579            attr.ref_mut_t()
1580        } else {
1581            if let Expr::Call(call) = self.obj.as_ref() {
1582                call.return_t()?;
1583            }
1584            self.obj.ref_mut_t()
1585        }
1586    }
1587}
1588
1589impl_locational!(Call, obj, args);
1590
1591impl Call {
1592    pub fn new(obj: Expr, attr_name: Option<Identifier>, args: Args) -> Self {
1593        Self {
1594            obj: Box::new(obj),
1595            attr_name,
1596            args,
1597        }
1598    }
1599
1600    pub fn call_signature_t(&self) -> &Type {
1601        if let Some(attr) = self.attr_name.as_ref() {
1602            attr.ref_t()
1603        } else {
1604            self.obj.ref_t()
1605        }
1606    }
1607
1608    pub fn is_method_call(&self) -> bool {
1609        self.call_signature_t().self_t().is_some()
1610    }
1611
1612    pub fn control_kind(&self) -> Option<ControlKind> {
1613        self.obj
1614            .show_acc()
1615            .and_then(|acc| ControlKind::try_from(&acc[..]).ok())
1616    }
1617
1618    pub fn additional_operation(&self) -> Option<OperationKind> {
1619        self.obj.show_acc().and_then(|s| match &s[..] {
1620            "import" => Some(OperationKind::Import),
1621            "pyimport" | "py" | "__import__" => Some(OperationKind::PyImport),
1622            "rsimport" => Some(OperationKind::RsImport),
1623            "Del" => Some(OperationKind::Del),
1624            "assert" => Some(OperationKind::Assert),
1625            _ => {
1626                if self.obj.qual_name() == Some("typing".into())
1627                    && self
1628                        .attr_name
1629                        .as_ref()
1630                        .is_some_and(|ident| &ident.inspect()[..] == "cast")
1631                {
1632                    return Some(OperationKind::Cast);
1633                }
1634                if self.obj.ref_t().is_callable() {
1635                    match self.attr_name.as_ref().map(|i| &i.inspect()[..]) {
1636                        Some("return") => Some(OperationKind::Return),
1637                        Some("yield") => Some(OperationKind::Yield),
1638                        _ => None,
1639                    }
1640                } else {
1641                    None
1642                }
1643            }
1644        })
1645    }
1646
1647    pub fn return_t(&self) -> Option<&Type> {
1648        if let Some(attr) = self.attr_name.as_ref() {
1649            attr.ref_t().return_t()
1650        } else {
1651            self.obj.ref_t().return_t()
1652        }
1653    }
1654}
1655
1656#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1657pub struct Block(Vec<Expr>);
1658
1659impl HasType for Block {
1660    #[inline]
1661    fn ref_t(&self) -> &Type {
1662        self.last()
1663            .map(|last| last.ref_t())
1664            .unwrap_or(Type::FAILURE)
1665    }
1666    #[inline]
1667    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1668        self.last_mut()?.ref_mut_t()
1669    }
1670    #[inline]
1671    fn t(&self) -> Type {
1672        self.last().map(|last| last.t()).unwrap_or(Type::Failure)
1673    }
1674    #[inline]
1675    fn signature_t(&self) -> Option<&Type> {
1676        self.last().and_then(|last| last.signature_t())
1677    }
1678    #[inline]
1679    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1680        self.last_mut().unwrap().signature_mut_t()
1681    }
1682}
1683
1684impl NestedDisplay for Block {
1685    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1686        fmt_lines(self.0.iter(), f, level)
1687    }
1688}
1689
1690impl NoTypeDisplay for Block {
1691    fn to_string_notype(&self) -> String {
1692        self.0
1693            .iter()
1694            .map(|e| e.to_string_notype())
1695            .collect::<Vec<_>>()
1696            .join("; ")
1697    }
1698}
1699
1700impl_display_from_nested!(Block);
1701impl_stream!(Block, Expr);
1702
1703impl Locational for Block {
1704    fn loc(&self) -> Location {
1705        Location::stream(&self.0)
1706    }
1707}
1708
1709impl Block {
1710    pub fn remove_def(&mut self, name: &str) -> Option<Def> {
1711        let mut i = 0;
1712        while i < self.0.len() {
1713            if let Expr::Def(def) = &self.0[i] {
1714                if def.sig.ident().inspect() == name {
1715                    return Def::try_from(self.0.remove(i)).ok();
1716                }
1717            }
1718            i += 1;
1719        }
1720        None
1721    }
1722
1723    pub fn get_def(&self, name: &str) -> Option<&Def> {
1724        self.0.iter().find_map(|e| match e {
1725            Expr::Def(def) if def.sig.ident().inspect() == name => Some(def),
1726            _ => None,
1727        })
1728    }
1729}
1730
1731#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1732pub struct Dummy(Vec<Expr>);
1733
1734impl HasType for Dummy {
1735    #[inline]
1736    fn ref_t(&self) -> &Type {
1737        Type::FAILURE
1738    }
1739    #[inline]
1740    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1741        None
1742    }
1743    #[inline]
1744    fn t(&self) -> Type {
1745        Type::Failure
1746    }
1747    #[inline]
1748    fn signature_t(&self) -> Option<&Type> {
1749        Some(Type::FAILURE)
1750    }
1751    #[inline]
1752    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1753        unreachable!()
1754    }
1755}
1756
1757impl NestedDisplay for Dummy {
1758    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
1759        writeln!(f, "Dummy:")?;
1760        fmt_lines(self.0.iter(), f, level)
1761    }
1762}
1763
1764impl NoTypeDisplay for Dummy {
1765    fn to_string_notype(&self) -> String {
1766        self.0
1767            .iter()
1768            .map(|e| e.to_string_notype())
1769            .collect::<Vec<_>>()
1770            .join("; ")
1771    }
1772}
1773
1774impl_display_from_nested!(Dummy);
1775impl_stream!(Dummy, Expr);
1776
1777impl Locational for Dummy {
1778    fn loc(&self) -> Location {
1779        Location::stream(&self.0)
1780    }
1781}
1782
1783#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1784pub struct VarSignature {
1785    pub ident: Identifier,
1786    /// if this flag is `true`, the variable is stored as `global`
1787    pub global: bool,
1788    pub t_spec: Option<Box<TypeSpecWithOp>>,
1789}
1790
1791impl NestedDisplay for VarSignature {
1792    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1793        write!(f, "{}", self.ident)
1794    }
1795}
1796
1797impl_display_from_nested!(VarSignature);
1798impl_locational!(VarSignature, ident);
1799
1800impl HasType for VarSignature {
1801    #[inline]
1802    fn ref_t(&self) -> &Type {
1803        self.ident.ref_t()
1804    }
1805    #[inline]
1806    fn ref_mut_t(&mut self) -> Option<&mut Type> {
1807        self.ident.ref_mut_t()
1808    }
1809    #[inline]
1810    fn signature_t(&self) -> Option<&Type> {
1811        self.ident.signature_t()
1812    }
1813    #[inline]
1814    fn signature_mut_t(&mut self) -> Option<&mut Type> {
1815        self.ident.signature_mut_t()
1816    }
1817}
1818
1819impl VarSignature {
1820    pub fn new(ident: Identifier, t_spec: Option<TypeSpecWithOp>) -> Self {
1821        Self {
1822            ident,
1823            global: false,
1824            t_spec: t_spec.map(Box::new),
1825        }
1826    }
1827
1828    pub fn global(ident: Identifier, t_spec: Option<TypeSpecWithOp>) -> Self {
1829        Self {
1830            ident,
1831            global: true,
1832            t_spec: t_spec.map(Box::new),
1833        }
1834    }
1835
1836    pub fn inspect(&self) -> &Str {
1837        self.ident.inspect()
1838    }
1839
1840    pub fn vis(&self) -> &VisibilityModifier {
1841        self.ident.vis()
1842    }
1843
1844    pub const fn name(&self) -> &VarName {
1845        &self.ident.raw.name
1846    }
1847}
1848
1849#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1850pub struct GlobSignature {
1851    pub dummy_ident: Identifier,
1852    pub vis: VisibilityModifier,
1853    pub names: Vec<VarName>,
1854    /// RHS type
1855    pub t: Type,
1856}
1857
1858impl_locational!(GlobSignature, dummy_ident);
1859impl_t!(GlobSignature);
1860impl_display_from_nested!(GlobSignature);
1861
1862impl NestedDisplay for GlobSignature {
1863    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
1864        write!(f, "{}*", self.vis)
1865    }
1866}
1867
1868impl GlobSignature {
1869    pub fn new(token: Token, vis: VisibilityModifier, names: Vec<VarName>, t: Type) -> Self {
1870        let loc = token.loc();
1871        let raw = ast::Identifier::new(
1872            ast::VisModifierSpec::Auto,
1873            VarName::from_str_and_loc(token.content, loc),
1874        );
1875        Self {
1876            dummy_ident: Identifier::new(
1877                raw,
1878                None,
1879                VarInfo {
1880                    t: t.clone(),
1881                    ..Default::default()
1882                },
1883            ),
1884            vis,
1885            names,
1886            t,
1887        }
1888    }
1889}
1890
1891/// Once the default_value is set to Some, all subsequent values must be Some
1892#[derive(Clone, Debug, PartialEq, Eq, Hash)]
1893pub struct NonDefaultParamSignature {
1894    pub raw: ast::NonDefaultParamSignature,
1895    pub vi: VarInfo,
1896    pub t_spec_as_expr: Option<Expr>,
1897}
1898
1899impl NestedDisplay for NonDefaultParamSignature {
1900    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, _level: usize) -> std::fmt::Result {
1901        write!(f, "{}(: {})", self.raw, self.vi.t)
1902    }
1903}
1904
1905impl_display_from_nested!(NonDefaultParamSignature);
1906
1907impl Locational for NonDefaultParamSignature {
1908    fn loc(&self) -> Location {
1909        self.raw.loc()
1910    }
1911}
1912
1913impl NonDefaultParamSignature {
1914    pub const fn new(
1915        sig: ast::NonDefaultParamSignature,
1916        vi: VarInfo,
1917        t_spec_as_expr: Option<Expr>,
1918    ) -> Self {
1919        Self {
1920            raw: sig,
1921            vi,
1922            t_spec_as_expr,
1923        }
1924    }
1925
1926    pub const fn inspect(&self) -> Option<&Str> {
1927        self.raw.pat.inspect()
1928    }
1929
1930    pub const fn name(&self) -> Option<&VarName> {
1931        self.raw.pat.name()
1932    }
1933}
1934
1935/// Once the default_value is set to Some, all subsequent values must be Some
1936#[derive(Clone, Debug, PartialEq, Eq, Hash)]
1937pub struct DefaultParamSignature {
1938    pub sig: NonDefaultParamSignature,
1939    pub default_val: Expr,
1940}
1941
1942impl NestedDisplay for DefaultParamSignature {
1943    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, _level: usize) -> std::fmt::Result {
1944        write!(f, "{} := {}", self.sig, self.default_val,)
1945    }
1946}
1947
1948impl NoTypeDisplay for DefaultParamSignature {
1949    fn to_string_notype(&self) -> String {
1950        format!("{} := {}", self.sig, self.default_val.to_string_notype())
1951    }
1952}
1953
1954impl_display_from_nested!(DefaultParamSignature);
1955
1956impl Locational for DefaultParamSignature {
1957    fn loc(&self) -> Location {
1958        Location::concat(&self.sig, &self.default_val)
1959    }
1960}
1961
1962impl DefaultParamSignature {
1963    pub const fn new(sig: NonDefaultParamSignature, default_val: Expr) -> Self {
1964        Self { sig, default_val }
1965    }
1966
1967    pub const fn inspect(&self) -> Option<&Str> {
1968        self.sig.inspect()
1969    }
1970
1971    pub const fn name(&self) -> Option<&VarName> {
1972        self.sig.name()
1973    }
1974}
1975
1976#[derive(Clone, Debug, PartialEq, Eq, Hash)]
1977pub enum GuardClause {
1978    Condition(Expr),
1979    Bind(Def),
1980}
1981
1982impl NestedDisplay for GuardClause {
1983    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, _level: usize) -> std::fmt::Result {
1984        match self {
1985            GuardClause::Condition(cond) => write!(f, "{}", cond),
1986            GuardClause::Bind(bind) => write!(f, "{}", bind),
1987        }
1988    }
1989}
1990
1991impl NoTypeDisplay for GuardClause {
1992    fn to_string_notype(&self) -> String {
1993        match self {
1994            GuardClause::Condition(cond) => cond.to_string_notype(),
1995            GuardClause::Bind(bind) => bind.to_string_notype(),
1996        }
1997    }
1998}
1999
2000impl_display_from_nested!(GuardClause);
2001
2002#[derive(Clone, Debug)]
2003pub struct Params {
2004    pub non_defaults: Vec<NonDefaultParamSignature>,
2005    pub var_params: Option<Box<NonDefaultParamSignature>>,
2006    pub defaults: Vec<DefaultParamSignature>,
2007    pub kw_var_params: Option<Box<NonDefaultParamSignature>>,
2008    pub guards: Vec<GuardClause>,
2009    pub parens: Option<(Location, Location)>,
2010}
2011
2012impl PartialEq for Params {
2013    fn eq(&self, other: &Self) -> bool {
2014        self.non_defaults == other.non_defaults
2015            && self.var_params == other.var_params
2016            && self.defaults == other.defaults
2017            && self.kw_var_params == other.kw_var_params
2018            && self.guards == other.guards
2019    }
2020}
2021
2022impl Eq for Params {}
2023
2024impl Hash for Params {
2025    fn hash<H: Hasher>(&self, state: &mut H) {
2026        self.non_defaults.hash(state);
2027        self.var_params.hash(state);
2028        self.defaults.hash(state);
2029        self.kw_var_params.hash(state);
2030        self.guards.hash(state);
2031    }
2032}
2033
2034impl fmt::Display for Params {
2035    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2036        write!(f, "({}", fmt_vec(&self.non_defaults))?;
2037        if let Some(var_params) = &self.var_params {
2038            if !self.non_defaults.is_empty() {
2039                write!(f, ", ")?;
2040            }
2041            write!(f, "*{var_params}")?;
2042        }
2043        if !self.defaults.is_empty() {
2044            if !self.non_defaults.is_empty() || self.var_params.is_some() {
2045                write!(f, ", ")?;
2046            }
2047            write!(f, "{}", fmt_vec(&self.defaults))?;
2048        }
2049        if let Some(kw_var_params) = &self.kw_var_params {
2050            if !self.non_defaults.is_empty()
2051                || self.var_params.is_some()
2052                || !self.defaults.is_empty()
2053            {
2054                write!(f, ", ")?;
2055            }
2056            write!(f, "**{kw_var_params}")?;
2057        }
2058        if !self.guards.is_empty() {
2059            write!(f, " if ")?;
2060        }
2061        for (i, guard) in self.guards.iter().enumerate() {
2062            if i > 0 {
2063                write!(f, " and ")?;
2064            }
2065            write!(f, "{guard}")?;
2066        }
2067        write!(f, ")")
2068    }
2069}
2070
2071impl NoTypeDisplay for Params {
2072    fn to_string_notype(&self) -> String {
2073        format!(
2074            "({}, {}{}{})",
2075            fmt_vec(&self.non_defaults),
2076            fmt_option!("*", &self.var_params, ", "),
2077            self.defaults
2078                .iter()
2079                .map(|p| p.to_string_notype())
2080                .fold("".to_string(), |acc, e| acc + &e + ", "),
2081            fmt_option!(pre ", **", &self.kw_var_params),
2082        )
2083    }
2084}
2085
2086impl Locational for Params {
2087    fn loc(&self) -> Location {
2088        if let Some((l, r)) = &self.parens {
2089            let loc = Location::concat(l, r);
2090            if !loc.is_unknown() {
2091                return loc;
2092            }
2093        }
2094        match (
2095            self.non_defaults.first().zip(self.non_defaults.last()),
2096            self.var_params.as_ref(),
2097            self.defaults.first().zip(self.defaults.last()),
2098        ) {
2099            (Some((l, _)), _, Some((_, r))) => Location::concat(l, r),
2100            (Some((l, _)), Some(r), None) => Location::concat(l, r.as_ref()),
2101            (Some((l, r)), None, None) => Location::concat(l, r),
2102            (None, Some(l), Some((_, r))) => Location::concat(l.as_ref(), r),
2103            (None, None, Some((l, r))) => Location::concat(l, r),
2104            (None, Some(l), None) => l.loc(),
2105            (None, None, None) => Location::Unknown,
2106        }
2107    }
2108}
2109
2110type RawParams = (
2111    Vec<NonDefaultParamSignature>,
2112    Option<Box<NonDefaultParamSignature>>,
2113    Vec<DefaultParamSignature>,
2114    Option<Box<NonDefaultParamSignature>>,
2115    Option<(Location, Location)>,
2116);
2117
2118type RefRawParams<'a> = (
2119    &'a Vec<NonDefaultParamSignature>,
2120    &'a Option<Box<NonDefaultParamSignature>>,
2121    &'a Vec<DefaultParamSignature>,
2122    &'a Option<Box<NonDefaultParamSignature>>,
2123    &'a Option<(Location, Location)>,
2124);
2125
2126impl Params {
2127    pub const fn new(
2128        non_defaults: Vec<NonDefaultParamSignature>,
2129        var_params: Option<Box<NonDefaultParamSignature>>,
2130        defaults: Vec<DefaultParamSignature>,
2131        kw_var_params: Option<Box<NonDefaultParamSignature>>,
2132        guards: Vec<GuardClause>,
2133        parens: Option<(Location, Location)>,
2134    ) -> Self {
2135        Self {
2136            non_defaults,
2137            var_params,
2138            defaults,
2139            kw_var_params,
2140            guards,
2141            parens,
2142        }
2143    }
2144
2145    pub fn empty() -> Self {
2146        Self::new(vec![], None, vec![], None, vec![], None)
2147    }
2148
2149    pub fn single(sig: NonDefaultParamSignature) -> Self {
2150        Self::new(vec![sig], None, vec![], None, vec![], None)
2151    }
2152
2153    pub fn non_default_only(non_defaults: Vec<NonDefaultParamSignature>) -> Self {
2154        Self::new(non_defaults, None, vec![], None, vec![], None)
2155    }
2156
2157    pub fn var_params(sig: NonDefaultParamSignature) -> Self {
2158        Self::new(vec![], Some(Box::new(sig)), vec![], None, vec![], None)
2159    }
2160
2161    pub const fn ref_deconstruct(&self) -> RefRawParams<'_> {
2162        (
2163            &self.non_defaults,
2164            &self.var_params,
2165            &self.defaults,
2166            &self.kw_var_params,
2167            &self.parens,
2168        )
2169    }
2170
2171    pub fn sigs(&self) -> impl Iterator<Item = &NonDefaultParamSignature> {
2172        self.non_defaults
2173            .iter()
2174            .chain(self.var_params.as_deref())
2175            .chain(self.defaults.iter().map(|d| &d.sig))
2176    }
2177
2178    pub fn deconstruct(self) -> RawParams {
2179        (
2180            self.non_defaults,
2181            self.var_params,
2182            self.defaults,
2183            self.kw_var_params,
2184            self.parens,
2185        )
2186    }
2187
2188    #[inline]
2189    pub fn len(&self) -> usize {
2190        self.non_defaults.len() + self.defaults.len()
2191    }
2192
2193    #[inline]
2194    pub fn is_empty(&self) -> bool {
2195        self.len() == 0
2196    }
2197
2198    pub fn push_non_default(&mut self, sig: NonDefaultParamSignature) {
2199        self.non_defaults.push(sig);
2200    }
2201
2202    pub fn push_var(&mut self, sig: NonDefaultParamSignature) {
2203        self.var_params = Some(Box::new(sig));
2204    }
2205
2206    pub fn push_default(&mut self, sig: DefaultParamSignature) {
2207        self.defaults.push(sig);
2208    }
2209
2210    pub fn push_kw_var(&mut self, sig: NonDefaultParamSignature) {
2211        self.kw_var_params = Some(Box::new(sig));
2212    }
2213}
2214
2215pub type Decorator = Expr;
2216
2217#[allow(clippy::derived_hash_with_manual_eq)]
2218#[derive(Debug, Clone, Hash)]
2219pub struct SubrSignature {
2220    pub decorators: HashSet<Decorator>,
2221    pub ident: Identifier,
2222    pub bounds: TypeBoundSpecs,
2223    pub params: Params,
2224    pub return_t_spec: Option<Box<TypeSpecWithOp>>,
2225    pub captured_names: Vec<Identifier>,
2226}
2227
2228impl PartialEq for SubrSignature {
2229    fn eq(&self, other: &Self) -> bool {
2230        self.ident == other.ident
2231            && self.bounds == other.bounds
2232            && self.params == other.params
2233            && self.return_t_spec == other.return_t_spec
2234            && self.captured_names == other.captured_names
2235            && self.decorators.linear_eq(&other.decorators)
2236    }
2237}
2238
2239impl Eq for SubrSignature {}
2240
2241impl NestedDisplay for SubrSignature {
2242    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
2243        write!(
2244            f,
2245            "{}{} (: {})",
2246            self.ident.to_string_notype(),
2247            self.params,
2248            self.ident.t()
2249        )
2250    }
2251}
2252
2253impl NoTypeDisplay for SubrSignature {
2254    fn to_string_notype(&self) -> String {
2255        format!(
2256            "{}{}",
2257            self.ident.to_string_notype(),
2258            self.params.to_string_notype()
2259        )
2260    }
2261}
2262
2263impl_display_from_nested!(SubrSignature);
2264
2265impl Locational for SubrSignature {
2266    // TODO: decorators
2267    fn loc(&self) -> Location {
2268        if let Some(return_t_spec) = self.return_t_spec.as_deref() {
2269            Location::concat(&self.ident, return_t_spec)
2270        } else {
2271            let params = self.params.loc();
2272            if !params.is_unknown() {
2273                return Location::concat(&self.ident, &params);
2274            }
2275            self.ident.loc()
2276        }
2277    }
2278}
2279
2280impl HasType for SubrSignature {
2281    #[inline]
2282    fn ref_t(&self) -> &Type {
2283        self.ident.ref_t()
2284    }
2285    #[inline]
2286    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2287        self.ident.ref_mut_t()
2288    }
2289    #[inline]
2290    fn signature_t(&self) -> Option<&Type> {
2291        self.ident.signature_t()
2292    }
2293    #[inline]
2294    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2295        self.ident.signature_mut_t()
2296    }
2297}
2298
2299impl SubrSignature {
2300    pub fn new(
2301        decorators: HashSet<Decorator>,
2302        ident: Identifier,
2303        bounds: TypeBoundSpecs,
2304        params: Params,
2305        return_t_spec: Option<TypeSpecWithOp>,
2306        captured_names: Vec<Identifier>,
2307    ) -> Self {
2308        Self {
2309            decorators,
2310            ident,
2311            bounds,
2312            params,
2313            return_t_spec: return_t_spec.map(Box::new),
2314            captured_names,
2315        }
2316    }
2317
2318    pub fn is_procedural(&self) -> bool {
2319        self.ident.is_procedural()
2320    }
2321
2322    pub const fn name(&self) -> &VarName {
2323        &self.ident.raw.name
2324    }
2325}
2326
2327#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2328pub struct Lambda {
2329    pub params: Params,
2330    pub op: Token,
2331    pub return_t_spec: Option<TypeSpec>,
2332    pub captured_names: Vec<Identifier>,
2333    pub body: Block,
2334    pub id: usize,
2335    pub t: Type,
2336}
2337
2338impl NestedDisplay for Lambda {
2339    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2340        writeln!(f, "{} {} (: {})", self.params, self.op.content, self.t)?;
2341        self.body.fmt_nest(f, level + 1)
2342    }
2343}
2344
2345impl NoTypeDisplay for Lambda {
2346    fn to_string_notype(&self) -> String {
2347        format!(
2348            "{} {} {}",
2349            self.params.to_string_notype(),
2350            self.op.content,
2351            self.body.to_string_notype()
2352        )
2353    }
2354}
2355
2356impl_display_from_nested!(Lambda);
2357impl_locational!(Lambda, lossy params, body);
2358impl_t!(Lambda);
2359
2360impl Lambda {
2361    pub const fn new(
2362        id: usize,
2363        params: Params,
2364        op: Token,
2365        return_t_spec: Option<TypeSpec>,
2366        captured_names: Vec<Identifier>,
2367        body: Block,
2368        t: Type,
2369    ) -> Self {
2370        Self {
2371            id,
2372            params,
2373            op,
2374            return_t_spec,
2375            captured_names,
2376            body,
2377            t,
2378        }
2379    }
2380
2381    pub fn is_procedural(&self) -> bool {
2382        self.op.is(TokenKind::ProcArrow)
2383    }
2384
2385    pub fn name_to_string(&self) -> String {
2386        format!("::<lambda_{}>", self.id)
2387    }
2388}
2389
2390#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2391pub enum Signature {
2392    Var(VarSignature),
2393    Subr(SubrSignature),
2394    Glob(GlobSignature),
2395}
2396
2397impl_nested_display_for_chunk_enum!(Signature; Var, Subr, Glob);
2398impl_display_for_enum!(Signature; Var, Subr, Glob);
2399impl_t_for_enum!(Signature; Var, Subr, Glob);
2400impl_locational_for_enum!(Signature; Var, Subr, Glob);
2401
2402impl Signature {
2403    pub const fn is_subr(&self) -> bool {
2404        matches!(self, Self::Subr(_))
2405    }
2406
2407    pub const fn is_glob(&self) -> bool {
2408        matches!(self, Self::Glob(_))
2409    }
2410
2411    pub fn is_const(&self) -> bool {
2412        match self {
2413            Self::Var(v) => v.ident.is_const(),
2414            Self::Subr(s) => s.ident.is_const(),
2415            Self::Glob(_) => false,
2416        }
2417    }
2418
2419    pub fn is_procedural(&self) -> bool {
2420        match self {
2421            Self::Var(v) => v.ident.is_procedural(),
2422            Self::Subr(s) => s.ident.is_procedural(),
2423            Self::Glob(_) => false,
2424        }
2425    }
2426
2427    pub fn vis(&self) -> &VisibilityModifier {
2428        match self {
2429            Self::Var(v) => v.ident.vis(),
2430            Self::Subr(s) => s.ident.vis(),
2431            Self::Glob(g) => &g.vis,
2432        }
2433    }
2434
2435    pub const fn inspect(&self) -> &Str {
2436        match self {
2437            Self::Var(v) => v.ident.inspect(),
2438            Self::Subr(s) => s.ident.inspect(),
2439            Self::Glob(g) => g.dummy_ident.inspect(),
2440        }
2441    }
2442
2443    pub const fn ident(&self) -> &Identifier {
2444        match self {
2445            Self::Var(v) => &v.ident,
2446            Self::Subr(s) => &s.ident,
2447            Self::Glob(g) => &g.dummy_ident,
2448        }
2449    }
2450
2451    pub fn ident_mut(&mut self) -> &mut Identifier {
2452        match self {
2453            Self::Var(v) => &mut v.ident,
2454            Self::Subr(s) => &mut s.ident,
2455            Self::Glob(g) => &mut g.dummy_ident,
2456        }
2457    }
2458
2459    pub fn into_ident(self) -> Identifier {
2460        match self {
2461            Self::Var(v) => v.ident,
2462            Self::Subr(s) => s.ident,
2463            Self::Glob(g) => g.dummy_ident,
2464        }
2465    }
2466
2467    pub const fn name(&self) -> &VarName {
2468        match self {
2469            Self::Var(v) => v.name(),
2470            Self::Subr(s) => s.name(),
2471            Self::Glob(g) => &g.dummy_ident.raw.name,
2472        }
2473    }
2474
2475    pub fn t_spec(&self) -> Option<&TypeSpec> {
2476        match self {
2477            Self::Var(v) => v.t_spec.as_ref().map(|t| &t.raw.t_spec),
2478            Self::Subr(s) => s.return_t_spec.as_ref().map(|t| &t.raw.t_spec),
2479            Self::Glob(_) => None,
2480        }
2481    }
2482
2483    pub fn t_spec_with_op(&self) -> Option<&TypeSpecWithOp> {
2484        match self {
2485            Self::Var(v) => v.t_spec.as_deref(),
2486            Self::Subr(s) => s.return_t_spec.as_deref(),
2487            Self::Glob(_) => None,
2488        }
2489    }
2490
2491    pub fn params_mut(&mut self) -> Option<&mut Params> {
2492        match self {
2493            Self::Var(_) => None,
2494            Self::Subr(s) => Some(&mut s.params),
2495            Self::Glob(_) => None,
2496        }
2497    }
2498
2499    pub fn params(&self) -> Option<&Params> {
2500        match self {
2501            Self::Var(_) => None,
2502            Self::Subr(s) => Some(&s.params),
2503            Self::Glob(_) => None,
2504        }
2505    }
2506
2507    pub fn default_params(&self) -> Option<impl Iterator<Item = &DefaultParamSignature>> {
2508        match self {
2509            Self::Var(_) => None,
2510            Self::Subr(s) => Some(s.params.defaults.iter()),
2511            Self::Glob(_) => None,
2512        }
2513    }
2514
2515    pub fn captured_names(&self) -> &[Identifier] {
2516        match self {
2517            Self::Var(_) => &[],
2518            Self::Subr(s) => &s.captured_names,
2519            Self::Glob(_) => &[],
2520        }
2521    }
2522}
2523
2524#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2525pub struct DefBody {
2526    pub op: Token,
2527    pub block: Block,
2528    pub id: DefId,
2529}
2530
2531impl_locational!(DefBody, lossy op, block);
2532impl_t!(DefBody, delegate block);
2533
2534impl DefBody {
2535    pub const fn new(op: Token, block: Block, id: DefId) -> Self {
2536        Self { op, block, id }
2537    }
2538}
2539
2540#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2541pub struct Def {
2542    pub sig: Signature,
2543    pub body: DefBody,
2544}
2545
2546impl NestedDisplay for Def {
2547    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2548        writeln!(f, "{} {}", self.sig, self.body.op.content)?;
2549        self.body.block.fmt_nest(f, level + 1)
2550    }
2551}
2552
2553impl NoTypeDisplay for Def {
2554    fn to_string_notype(&self) -> String {
2555        format!(
2556            "{} {} {}",
2557            self.sig,
2558            self.body.op.content,
2559            self.body.block.to_string_notype()
2560        )
2561    }
2562}
2563
2564impl_display_from_nested!(Def);
2565impl_locational!(Def, sig, body);
2566
2567impl HasType for Def {
2568    #[inline]
2569    fn ref_t(&self) -> &Type {
2570        Type::NONE
2571    }
2572    #[inline]
2573    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2574        None
2575    }
2576    #[inline]
2577    fn signature_t(&self) -> Option<&Type> {
2578        None
2579    }
2580    #[inline]
2581    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2582        None
2583    }
2584}
2585
2586impl Def {
2587    pub const fn new(sig: Signature, body: DefBody) -> Self {
2588        Self { sig, body }
2589    }
2590
2591    pub fn def_kind(&self) -> DefKind {
2592        match self.body.block.first() {
2593            Some(Expr::Call(call)) => match call.obj.show_acc().as_ref().map(|n| &n[..]) {
2594                Some("Class") => DefKind::Class,
2595                Some("Inherit") => DefKind::Inherit,
2596                Some("Trait") => DefKind::Trait,
2597                Some("Subsume") => DefKind::Subsume,
2598                Some("Inheritable") => {
2599                    if let Some(Expr::Call(inner)) = call.args.get_left_or_key("Class") {
2600                        match inner.obj.show_acc().as_ref().map(|n| &n[..]) {
2601                            Some("Class") => DefKind::Class,
2602                            Some("Inherit") => DefKind::Inherit,
2603                            _ => DefKind::Other,
2604                        }
2605                    } else {
2606                        DefKind::Other
2607                    }
2608                }
2609                Some("Patch") => DefKind::Patch,
2610                Some("import") => DefKind::ErgImport,
2611                Some("pyimport") | Some("__import__") => DefKind::PyImport,
2612                Some("rsimport") => DefKind::RsImport,
2613                #[cfg(feature = "debug")]
2614                Some("py") => DefKind::PyImport,
2615                _ => DefKind::Other,
2616            },
2617            _ => DefKind::Other,
2618        }
2619    }
2620
2621    pub fn get_base(&self) -> Option<&Record> {
2622        match self.body.block.first()? {
2623            Expr::Call(call) => match call.obj.show_acc().as_ref().map(|n| &n[..]) {
2624                Some("Class") | Some("Trait") => {
2625                    if let Some(Expr::Record(rec)) = call.args.get_left_or_key("Base") {
2626                        Some(rec)
2627                    } else {
2628                        None
2629                    }
2630                }
2631                _ => None,
2632            },
2633            _ => None,
2634        }
2635    }
2636}
2637
2638#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2639pub struct Methods {
2640    pub class: Type,
2641    pub impl_trait: Option<Type>,
2642    pub defs: Block,
2643}
2644
2645impl NestedDisplay for Methods {
2646    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2647        writeln!(
2648            f,
2649            "{} {}",
2650            self.class,
2651            fmt_option!("|<: ", &self.impl_trait, "|"),
2652        )?;
2653        self.defs.fmt_nest(f, level + 1)
2654    }
2655}
2656
2657// TODO
2658impl NoTypeDisplay for Methods {
2659    fn to_string_notype(&self) -> String {
2660        format!(
2661            "{} {} {}",
2662            self.class,
2663            fmt_option!("|<: ", &self.impl_trait, "|"),
2664            self.defs.to_string_notype()
2665        )
2666    }
2667}
2668
2669impl_display_from_nested!(Methods);
2670impl_locational!(Methods, defs);
2671
2672impl HasType for Methods {
2673    #[inline]
2674    fn ref_t(&self) -> &Type {
2675        Type::NONE
2676    }
2677    #[inline]
2678    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2679        None
2680    }
2681    #[inline]
2682    fn signature_t(&self) -> Option<&Type> {
2683        None
2684    }
2685    #[inline]
2686    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2687        None
2688    }
2689}
2690
2691impl Methods {
2692    pub const fn new(class: Type, impl_trait: Option<Type>, defs: Block) -> Self {
2693        Self {
2694            class,
2695            impl_trait,
2696            defs,
2697        }
2698    }
2699}
2700
2701#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2702pub struct ClassDef {
2703    pub obj: Box<GenTypeObj>,
2704    pub sig: Signature,
2705    pub require_or_sup: Option<Box<Expr>>,
2706    /// The type of `new` that is automatically defined if not defined
2707    pub need_to_gen_new: bool,
2708    pub constructor: Type,
2709    pub methods_list: Vec<Methods>,
2710}
2711
2712impl NestedDisplay for ClassDef {
2713    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2714        self.sig.fmt_nest(f, level)?;
2715        writeln!(f, ":")?;
2716        fmt_lines(self.methods_list.iter(), f, level)
2717    }
2718}
2719
2720// TODO
2721impl NoTypeDisplay for ClassDef {
2722    fn to_string_notype(&self) -> String {
2723        let methods = self
2724            .methods_list
2725            .iter()
2726            .map(|m| m.to_string_notype())
2727            .collect::<Vec<_>>()
2728            .join("\n");
2729        format!("{}: {methods}", self.sig)
2730    }
2731}
2732
2733impl_display_from_nested!(ClassDef);
2734impl_locational!(ClassDef, sig, lossy methods_list);
2735
2736impl HasType for ClassDef {
2737    #[inline]
2738    fn ref_t(&self) -> &Type {
2739        Type::NONE
2740    }
2741    #[inline]
2742    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2743        None
2744    }
2745    #[inline]
2746    fn signature_t(&self) -> Option<&Type> {
2747        None
2748    }
2749    #[inline]
2750    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2751        None
2752    }
2753}
2754
2755impl ClassDef {
2756    pub fn new(
2757        obj: GenTypeObj,
2758        sig: Signature,
2759        require_or_sup: Option<Expr>,
2760        need_to_gen_new: bool,
2761        constructor: Type,
2762        methods_list: Vec<Methods>,
2763    ) -> Self {
2764        Self {
2765            obj: Box::new(obj),
2766            sig,
2767            require_or_sup: require_or_sup.map(Box::new),
2768            need_to_gen_new,
2769            constructor,
2770            methods_list,
2771        }
2772    }
2773
2774    pub fn all_methods(&self) -> impl Iterator<Item = &Expr> {
2775        self.methods_list.iter().flat_map(|m| m.defs.iter())
2776    }
2777
2778    pub fn all_methods_mut(&mut self) -> impl Iterator<Item = &mut Expr> {
2779        self.methods_list.iter_mut().flat_map(|m| m.defs.iter_mut())
2780    }
2781
2782    pub fn take_all_methods(methods_list: Vec<Methods>) -> Block {
2783        let mut joined = Block::empty();
2784        for methods in methods_list {
2785            joined.extend(methods.defs);
2786        }
2787        joined
2788    }
2789
2790    pub fn get_all_methods(methods_list: &[Methods]) -> Vec<&Expr> {
2791        let mut joined = vec![];
2792        for methods in methods_list {
2793            joined.extend(methods.defs.iter());
2794        }
2795        joined
2796    }
2797}
2798
2799#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2800pub struct PatchDef {
2801    pub sig: Signature,
2802    pub base: Box<Expr>,
2803    pub methods: Block,
2804}
2805
2806impl NestedDisplay for PatchDef {
2807    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2808        write!(f, "{} = Patch ", self.sig)?;
2809        self.base.fmt_nest(f, level)?;
2810        writeln!(f, ":")?;
2811        self.methods.fmt_nest(f, level + 1)
2812    }
2813}
2814
2815// TODO
2816impl NoTypeDisplay for PatchDef {
2817    fn to_string_notype(&self) -> String {
2818        format!(
2819            "{} = Patch {}: {}",
2820            self.sig,
2821            self.base.to_string_notype(),
2822            self.methods.to_string_notype()
2823        )
2824    }
2825}
2826
2827impl_display_from_nested!(PatchDef);
2828impl_locational!(PatchDef, sig, methods);
2829
2830impl HasType for PatchDef {
2831    #[inline]
2832    fn ref_t(&self) -> &Type {
2833        Type::NONE
2834    }
2835    #[inline]
2836    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2837        None
2838    }
2839    #[inline]
2840    fn signature_t(&self) -> Option<&Type> {
2841        None
2842    }
2843    #[inline]
2844    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2845        None
2846    }
2847}
2848
2849impl PatchDef {
2850    pub fn new(sig: Signature, base: Expr, methods: Block) -> Self {
2851        Self {
2852            sig,
2853            base: Box::new(base),
2854            methods,
2855        }
2856    }
2857}
2858
2859#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2860pub struct ReDef {
2861    pub attr: Accessor,
2862    pub block: Block,
2863}
2864
2865impl NestedDisplay for ReDef {
2866    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
2867        self.attr.fmt_nest(f, level)?;
2868        writeln!(f, " = ")?;
2869        self.block.fmt_nest(f, level + 1)
2870    }
2871}
2872
2873impl NoTypeDisplay for ReDef {
2874    fn to_string_notype(&self) -> String {
2875        format!(
2876            "{} = {}",
2877            self.attr.to_string_notype(),
2878            self.block.to_string_notype()
2879        )
2880    }
2881}
2882
2883impl_display_from_nested!(ReDef);
2884impl_locational!(ReDef, attr, block);
2885
2886impl HasType for ReDef {
2887    #[inline]
2888    fn ref_t(&self) -> &Type {
2889        Type::NONE
2890    }
2891    #[inline]
2892    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2893        None
2894    }
2895    #[inline]
2896    fn signature_t(&self) -> Option<&Type> {
2897        None
2898    }
2899    #[inline]
2900    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2901        None
2902    }
2903}
2904
2905impl ReDef {
2906    pub const fn new(attr: Accessor, block: Block) -> Self {
2907        Self { attr, block }
2908    }
2909}
2910
2911#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2912pub struct TypeSpecWithOp {
2913    pub raw: ast::TypeSpecWithOp,
2914    /// Required for dynamic type checking
2915    pub expr: Box<Expr>,
2916    pub spec_t: Type,
2917}
2918
2919impl NestedDisplay for TypeSpecWithOp {
2920    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
2921        write!(f, "{}", self.raw)
2922    }
2923}
2924
2925impl_display_from_nested!(TypeSpecWithOp);
2926impl_locational!(TypeSpecWithOp, raw);
2927
2928impl TypeSpecWithOp {
2929    pub fn new(raw: ast::TypeSpecWithOp, expr: Expr, spec_t: Type) -> Self {
2930        Self {
2931            raw,
2932            expr: Box::new(expr),
2933            spec_t,
2934        }
2935    }
2936
2937    pub fn kind(&self) -> AscriptionKind {
2938        match self.raw.op.kind {
2939            TokenKind::Colon => AscriptionKind::TypeOf,
2940            TokenKind::SubtypeOf => AscriptionKind::SubtypeOf,
2941            TokenKind::SupertypeOf => AscriptionKind::SupertypeOf,
2942            TokenKind::As => AscriptionKind::AsCast,
2943            _ => unreachable!(),
2944        }
2945    }
2946}
2947
2948#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2949pub struct TypeAscription {
2950    pub expr: Box<Expr>,
2951    pub spec: TypeSpecWithOp,
2952}
2953
2954impl NestedDisplay for TypeAscription {
2955    fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
2956        writeln!(f, "{} {}", self.expr, self.spec)
2957    }
2958}
2959
2960impl NoTypeDisplay for TypeAscription {
2961    fn to_string_notype(&self) -> String {
2962        format!("{}{}", self.expr.to_string_notype(), self.spec)
2963    }
2964}
2965
2966impl_display_from_nested!(TypeAscription);
2967impl_locational!(TypeAscription, expr, spec);
2968
2969impl HasType for TypeAscription {
2970    #[inline]
2971    fn ref_t(&self) -> &Type {
2972        if self.spec.kind().is_force_cast() {
2973            &self.spec.spec_t
2974        } else {
2975            self.expr.ref_t()
2976        }
2977    }
2978    #[inline]
2979    fn ref_mut_t(&mut self) -> Option<&mut Type> {
2980        if self.spec.kind().is_force_cast() {
2981            Some(&mut self.spec.spec_t)
2982        } else {
2983            self.expr.ref_mut_t()
2984        }
2985    }
2986    #[inline]
2987    fn signature_t(&self) -> Option<&Type> {
2988        self.expr.signature_t()
2989    }
2990    #[inline]
2991    fn signature_mut_t(&mut self) -> Option<&mut Type> {
2992        self.expr.signature_mut_t()
2993    }
2994}
2995
2996impl TypeAscription {
2997    pub fn new(expr: Expr, spec: TypeSpecWithOp) -> Self {
2998        Self {
2999            expr: Box::new(expr),
3000            spec,
3001        }
3002    }
3003}
3004
3005#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3006pub enum Expr {
3007    Literal(Literal),
3008    Accessor(Accessor),
3009    List(List),
3010    Tuple(Tuple),
3011    Set(Set),
3012    Dict(Dict),
3013    Record(Record),
3014    BinOp(BinOp),
3015    UnaryOp(UnaryOp),
3016    Call(Call),
3017    Lambda(Lambda),
3018    Def(Def),
3019    ClassDef(ClassDef),
3020    PatchDef(PatchDef),
3021    ReDef(ReDef),
3022    TypeAsc(TypeAscription),
3023    Code(Block),     // code object
3024    Compound(Block), // compound statement
3025    Import(Accessor),
3026    Dummy(Dummy), // for mapping to Python AST
3027}
3028
3029impl_nested_display_for_chunk_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Code, Compound, TypeAsc, Set, Import, Dummy);
3030impl_no_type_display_for_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Code, Compound, TypeAsc, Set, Import, Dummy);
3031impl_display_from_nested!(Expr);
3032impl_locational_for_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Code, Compound, TypeAsc, Set, Import, Dummy);
3033impl_t_for_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Code, Compound, TypeAsc, Set, Import, Dummy);
3034impl_from_trait_for_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Set, Dummy);
3035impl_try_from_trait_for_enum!(Expr; Literal, Accessor, List, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Def, ClassDef, PatchDef, ReDef, Set, Dummy);
3036
3037impl Default for Expr {
3038    fn default() -> Self {
3039        Self::Code(Block::default())
3040    }
3041}
3042
3043impl Expr {
3044    pub fn receiver_t(&self) -> Option<&Type> {
3045        match self {
3046            Self::Accessor(Accessor::Attr(attr)) => Some(attr.obj.ref_t()),
3047            Self::TypeAsc(t_asc) => t_asc.expr.receiver_t(),
3048            _other => None,
3049        }
3050    }
3051
3052    pub fn as_call(&self) -> Option<&Call> {
3053        <&Call>::try_from(self).ok()
3054    }
3055
3056    pub fn as_binop(&self) -> Option<&BinOp> {
3057        <&BinOp>::try_from(self).ok()
3058    }
3059
3060    pub fn as_unaryop(&self) -> Option<&UnaryOp> {
3061        <&UnaryOp>::try_from(self).ok()
3062    }
3063
3064    pub fn as_def(&self) -> Option<&Def> {
3065        <&Def>::try_from(self).ok()
3066    }
3067
3068    pub fn as_lambda(&self) -> Option<&Lambda> {
3069        <&Lambda>::try_from(self).ok()
3070    }
3071
3072    pub fn as_class_def(&self) -> Option<&ClassDef> {
3073        <&ClassDef>::try_from(self).ok()
3074    }
3075
3076    pub fn as_literal(&self) -> Option<&Literal> {
3077        <&Literal>::try_from(self).ok()
3078    }
3079
3080    pub fn as_accessor(&self) -> Option<&Accessor> {
3081        <&Accessor>::try_from(self).ok()
3082    }
3083
3084    pub fn as_ident(&self) -> Option<&Identifier> {
3085        self.as_accessor().and_then(|acc| acc.as_ident())
3086    }
3087
3088    pub fn as_attr(&self) -> Option<&Attribute> {
3089        self.as_accessor().and_then(|acc| acc.as_attr())
3090    }
3091
3092    pub fn show_acc(&self) -> Option<String> {
3093        match self {
3094            Expr::Accessor(acc) => Some(acc.show()),
3095            Expr::TypeAsc(t_asc) => t_asc.expr.show_acc(),
3096            _ => None,
3097        }
3098    }
3099
3100    pub fn var_info(&self) -> Option<&VarInfo> {
3101        match self {
3102            Expr::Accessor(acc) => Some(acc.var_info()),
3103            Expr::TypeAsc(t_asc) => t_asc.expr.var_info(),
3104            _ => None,
3105        }
3106    }
3107
3108    /// 参照するオブジェクト自体が持っている名前(e.g. Int.qual_name == Some("int"), Socket!.qual_name == Some("io.Socket!"))
3109    pub fn qual_name(&self) -> Option<Str> {
3110        match self {
3111            Expr::Accessor(acc) => acc.qual_name(),
3112            Expr::TypeAsc(tasc) => tasc.expr.qual_name(),
3113            _ => None,
3114        }
3115    }
3116
3117    /// e.g. Int.local_name == Some("int"), Socket!.local_name == Some("Socket!")
3118    pub fn local_name(&self) -> Option<&str> {
3119        match self {
3120            Expr::Accessor(acc) => acc.local_name(),
3121            Expr::TypeAsc(tasc) => tasc.expr.local_name(),
3122            _ => None,
3123        }
3124    }
3125
3126    pub fn is_py_api(&self) -> bool {
3127        match self {
3128            Expr::Accessor(acc) => acc.is_py_api(),
3129            Expr::TypeAsc(tasc) => tasc.expr.is_py_api(),
3130            _ => false,
3131        }
3132    }
3133
3134    pub fn is_acc(&self) -> bool {
3135        match self {
3136            Self::Accessor(_) => true,
3137            Self::TypeAsc(tasc) => tasc.expr.is_acc(),
3138            _ => false,
3139        }
3140    }
3141
3142    pub fn last_name(&self) -> Option<&VarName> {
3143        match self {
3144            Expr::Accessor(acc) => Some(acc.last_name()),
3145            Expr::TypeAsc(tasc) => tasc.expr.last_name(),
3146            _ => None,
3147        }
3148    }
3149
3150    pub fn is_type_asc(&self) -> bool {
3151        matches!(self, Expr::TypeAsc(_))
3152    }
3153
3154    pub fn is_doc_comment(&self) -> bool {
3155        match self {
3156            Expr::Literal(lit) => lit.is_doc_comment(),
3157            _ => false,
3158        }
3159    }
3160
3161    pub const fn name(&self) -> &'static str {
3162        match self {
3163            Self::Literal(_) => "literal",
3164            Self::Accessor(_) => "accessor",
3165            Self::List(_) => "list",
3166            Self::Tuple(_) => "tuple",
3167            Self::Dict(_) => "dict",
3168            Self::Set(_) => "set",
3169            Self::Record(_) => "record",
3170            Self::BinOp(_) => "binary operator call",
3171            Self::UnaryOp(_) => "unary operator call",
3172            Self::Call(_) => "call",
3173            Self::Lambda(_) => "lambda",
3174            Self::TypeAsc(_) => "type ascription",
3175            Self::Def(_) => "definition",
3176            Self::Code(_) => "code",
3177            Self::Compound(_) => "compound expression",
3178            Self::Import(_) => "import",
3179            Self::ClassDef(_) => "class definition",
3180            Self::PatchDef(_) => "patch definition",
3181            Self::ReDef(_) => "re-definition",
3182            Self::Dummy(_) => "dummy",
3183        }
3184    }
3185
3186    pub fn should_wrap(&self) -> bool {
3187        match self {
3188            Self::Literal(_)
3189            | Self::Accessor(_)
3190            | Self::Call(_)
3191            | Self::BinOp(_)
3192            | Self::UnaryOp(_) => true,
3193            Self::TypeAsc(t) => t.expr.should_wrap(),
3194            _ => false,
3195        }
3196    }
3197
3198    pub fn need_to_be_closed(&self) -> bool {
3199        match self {
3200            Self::BinOp(_) | Self::UnaryOp(_) | Self::Lambda(_) | Self::TypeAsc(_) => true,
3201            Self::Tuple(tup) => match tup {
3202                Tuple::Normal(tup) => tup.elems.paren.is_none(),
3203            },
3204            Self::Call(call) if ERG_MODE => call.args.paren.is_none(),
3205            _ => false,
3206        }
3207    }
3208
3209    pub fn call(self, args: Args) -> Call {
3210        match self {
3211            Self::Accessor(Accessor::Attr(attr)) => Call::new(*attr.obj, Some(attr.ident), args),
3212            other => Call::new(other, None, args),
3213        }
3214    }
3215
3216    pub fn call_expr(self, args: Args) -> Self {
3217        Self::Call(self.call(args))
3218    }
3219
3220    pub fn call1(self, expr: Expr) -> Self {
3221        self.call_expr(Args::single(PosArg::new(expr)))
3222    }
3223
3224    pub fn call2(self, expr1: Expr, expr2: Expr) -> Self {
3225        self.call_expr(Args::pos_only(
3226            vec![PosArg::new(expr1), PosArg::new(expr2)],
3227            None,
3228        ))
3229    }
3230
3231    pub fn method_call(self, attr_name: Identifier, args: Args) -> Call {
3232        Call::new(self, Some(attr_name), args)
3233    }
3234
3235    pub fn method_call_expr(self, attr_name: Identifier, args: Args) -> Self {
3236        Self::Call(self.method_call(attr_name, args))
3237    }
3238
3239    pub fn attr(self, ident: Identifier) -> Accessor {
3240        Accessor::attr(self, ident)
3241    }
3242
3243    pub fn attr_expr(self, ident: Identifier) -> Self {
3244        Self::Accessor(self.attr(ident))
3245    }
3246
3247    pub fn type_asc(self, t_spec: TypeSpecWithOp) -> TypeAscription {
3248        TypeAscription::new(self, t_spec)
3249    }
3250
3251    pub fn type_asc_expr(self, t_spec: TypeSpecWithOp) -> Self {
3252        Self::TypeAsc(self.type_asc(t_spec))
3253    }
3254
3255    /// Return the complexity of the expression in terms of type inference.
3256    /// For function calls, type inference is performed sequentially, starting with the least complex argument.
3257    pub fn complexity(&self) -> usize {
3258        match self {
3259            Self::Literal(_) | Self::TypeAsc(_) => 0,
3260            Self::Accessor(Accessor::Ident(_)) => 1,
3261            Self::Accessor(Accessor::Attr(attr)) => 1 + attr.obj.complexity(),
3262            Self::Tuple(Tuple::Normal(tup)) => {
3263                let mut sum = 0;
3264                for elem in tup.elems.pos_args.iter() {
3265                    sum += elem.expr.complexity();
3266                }
3267                sum
3268            }
3269            Self::List(List::Normal(lis)) => {
3270                let mut sum = 0;
3271                for elem in lis.elems.pos_args.iter() {
3272                    sum += elem.expr.complexity();
3273                }
3274                sum
3275            }
3276            Self::Dict(Dict::Normal(dic)) => {
3277                let mut sum = 0;
3278                for kv in dic.kvs.iter() {
3279                    sum += kv.key.complexity();
3280                    sum += kv.value.complexity();
3281                }
3282                sum
3283            }
3284            Self::Set(Set::Normal(set)) => {
3285                let mut sum = 0;
3286                for elem in set.elems.pos_args.iter() {
3287                    sum += elem.expr.complexity();
3288                }
3289                sum
3290            }
3291            Self::Record(rec) => {
3292                let mut sum = 0;
3293                for attr in rec.attrs.iter() {
3294                    for chunk in attr.body.block.iter() {
3295                        sum += chunk.complexity();
3296                    }
3297                }
3298                sum
3299            }
3300            Self::BinOp(bin) => 1 + bin.lhs.complexity() + bin.rhs.complexity(),
3301            Self::UnaryOp(unary) => 1 + unary.expr.complexity(),
3302            Self::Call(call) => {
3303                let mut sum = 1 + call.obj.complexity();
3304                for arg in call.args.pos_args.iter() {
3305                    sum += arg.expr.complexity();
3306                }
3307                if let Some(var_params) = call.args.var_args.as_ref() {
3308                    sum += var_params.expr.complexity();
3309                }
3310                for kw_arg in call.args.kw_args.iter() {
3311                    sum += kw_arg.expr.complexity();
3312                }
3313                sum
3314            }
3315            Self::Lambda(lambda) => {
3316                let mut sum = 1
3317                    + lambda.return_t_spec.is_none() as usize
3318                    + lambda
3319                        .params
3320                        .sigs()
3321                        .fold(0, |acc, sig| acc + sig.raw.t_spec.is_none() as usize);
3322                for chunk in lambda.body.iter() {
3323                    sum += chunk.complexity();
3324                }
3325                sum
3326            }
3327            _ => 5,
3328        }
3329    }
3330
3331    // TODO: structural types
3332    pub fn try_from_type(typ: Type) -> Result<Expr, Type> {
3333        Ok(Expr::Accessor(Accessor::Ident(
3334            Identifier::public_with_line(Token::DUMMY, typ.local_name(), 0),
3335        )))
3336    }
3337}
3338
3339/// Toplevel grammar unit
3340#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3341pub struct Module(Vec<Expr>);
3342
3343impl NestedDisplay for Module {
3344    fn fmt_nest(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result {
3345        fmt_lines(self.0.iter(), f, level)
3346    }
3347}
3348
3349impl_display_from_nested!(Module);
3350
3351impl NoTypeDisplay for Module {
3352    fn to_string_notype(&self) -> String {
3353        self.0
3354            .iter()
3355            .map(|e| e.to_string_notype())
3356            .collect::<Vec<_>>()
3357            .join("\n")
3358    }
3359}
3360
3361impl Locational for Module {
3362    fn loc(&self) -> Location {
3363        Location::stream(&self.0)
3364    }
3365}
3366
3367impl_stream!(Module, Expr);
3368
3369impl Module {
3370    pub fn get_attr(&self, name: &str) -> Option<&Def> {
3371        self.0.iter().find_map(|e| match e {
3372            Expr::Def(def) if def.sig.ident().inspect() == name => Some(def),
3373            _ => None,
3374        })
3375    }
3376}
3377
3378/// High-level Intermediate Representation
3379/// AST with type information added
3380#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3381pub struct HIR {
3382    pub name: Str,
3383    pub dependencies: HashSet<NormalizedPathBuf>,
3384    pub module: Module,
3385}
3386
3387impl std::fmt::Display for HIR {
3388    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3389        write!(f, "{}", self.module)
3390    }
3391}
3392
3393impl Default for HIR {
3394    fn default() -> Self {
3395        Self {
3396            name: Str::ever("<module>"),
3397            dependencies: HashSet::default(),
3398            module: Module(vec![]),
3399        }
3400    }
3401}
3402
3403impl HIR {
3404    pub fn new(name: Str, module: Module) -> Self {
3405        Self {
3406            name,
3407            dependencies: HashSet::default(),
3408            module,
3409        }
3410    }
3411
3412    pub fn with_dependencies(self, deps: HashSet<NormalizedPathBuf>) -> Self {
3413        Self {
3414            dependencies: self.dependencies.concat(deps),
3415            ..self
3416        }
3417    }
3418}
3419
3420#[cfg(test)]
3421mod tests {
3422    use super::*;
3423
3424    use std::mem::size_of;
3425
3426    #[test]
3427    fn check_structs_size() {
3428        println!("Expr: {}", size_of::<Expr>());
3429        println!("Literal: {}", size_of::<Literal>());
3430        println!("Accessor: {}", size_of::<Accessor>());
3431        println!("List: {}", size_of::<List>());
3432        println!("Tuple: {}", size_of::<Tuple>());
3433        println!("Dict: {}", size_of::<Dict>());
3434        println!("Record: {}", size_of::<Record>());
3435        println!("BinOp: {}", size_of::<BinOp>());
3436        println!("UnaryOp: {}", size_of::<UnaryOp>());
3437        println!("Call: {}", size_of::<Call>());
3438        println!("Lambda: {}", size_of::<Lambda>());
3439        println!("Def: {}", size_of::<Def>());
3440        println!("ClassDef: {}", size_of::<ClassDef>());
3441        println!("GenTypeObj: {}", size_of::<GenTypeObj>());
3442        println!("Signature: {}", size_of::<Signature>());
3443        println!("VarSignature: {}", size_of::<VarSignature>());
3444        println!("Identifier: {}", size_of::<Identifier>());
3445        println!("VarInfo: {}", size_of::<VarInfo>());
3446        println!("Visibility: {}", size_of::<crate::ty::Visibility>());
3447        println!("VarKind: {}", size_of::<crate::varinfo::VarKind>());
3448        println!("ContextKind: {}", size_of::<crate::context::ContextKind>());
3449        println!("TraitImpl: {}", size_of::<crate::context::TraitImpl>());
3450        println!("SubrSignature: {}", size_of::<SubrSignature>());
3451        println!("TypeBoundSpecs: {}", size_of::<TypeBoundSpecs>());
3452        println!("Params: {}", size_of::<Params>());
3453        println!("TypeSpecWithOp: {}", size_of::<TypeSpecWithOp>());
3454        println!("TypeSpec: {}", size_of::<TypeSpec>());
3455        println!("GlobSignature: {}", size_of::<GlobSignature>());
3456        println!("PatchDef: {}", size_of::<PatchDef>());
3457        println!("ReDef: {}", size_of::<ReDef>());
3458        println!("TypeAsc: {}", size_of::<TypeAscription>());
3459        println!("Code: {}", size_of::<Block>());
3460        println!("Compound: {}", size_of::<Block>());
3461        println!("Import: {}", size_of::<Accessor>());
3462        println!("Dummy: {}", size_of::<Dummy>());
3463        println!("Module: {}", size_of::<Module>());
3464        println!("HIR: {}", size_of::<HIR>());
3465        println!("Type: {}", size_of::<Type>());
3466    }
3467}