diff --git a/compiler/erg_compiler/context/tyvar.rs b/compiler/erg_compiler/context/tyvar.rs index 6721fa22..2c708388 100644 --- a/compiler/erg_compiler/context/tyvar.rs +++ b/compiler/erg_compiler/context/tyvar.rs @@ -489,6 +489,7 @@ impl Context { } Ok(()) } + hir::Expr::TypeAsc(tasc) => self.resolve_expr_t(&mut tasc.expr), hir::Expr::Code(chunks) | hir::Expr::Compound(chunks) => { for chunk in chunks.iter_mut() { self.resolve_expr_t(chunk)?; diff --git a/compiler/erg_compiler/hir.rs b/compiler/erg_compiler/hir.rs index 32b08fee..29497f77 100644 --- a/compiler/erg_compiler/hir.rs +++ b/compiler/erg_compiler/hir.rs @@ -1485,6 +1485,49 @@ impl AttrDef { } } +#[derive(Debug, Clone)] +pub struct TypeAscription { + pub expr: Box, + pub spec: TypeSpec, +} + +impl NestedDisplay for TypeAscription { + fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result { + writeln!(f, "{}: {}", self.expr, self.spec) + } +} + +impl_display_from_nested!(TypeAscription); +impl_locational!(TypeAscription, expr, spec); + +impl HasType for TypeAscription { + #[inline] + fn ref_t(&self) -> &Type { + self.expr.ref_t() + } + #[inline] + fn ref_mut_t(&mut self) -> &mut Type { + self.expr.ref_mut_t() + } + #[inline] + fn signature_t(&self) -> Option<&Type> { + self.expr.signature_t() + } + #[inline] + fn signature_mut_t(&mut self) -> Option<&mut Type> { + self.expr.signature_mut_t() + } +} + +impl TypeAscription { + pub fn new(expr: Expr, spec: TypeSpec) -> Self { + Self { + expr: Box::new(expr), + spec, + } + } +} + #[derive(Debug, Clone)] pub enum Expr { Lit(Literal), @@ -1502,14 +1545,15 @@ pub enum Expr { Def(Def), ClassDef(ClassDef), AttrDef(AttrDef), + TypeAsc(TypeAscription), Code(Block), // code object Compound(Block), // compound statement } -impl_nested_display_for_chunk_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound); +impl_nested_display_for_chunk_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound, TypeAsc); impl_display_from_nested!(Expr); -impl_locational_for_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound); -impl_t_for_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound); +impl_locational_for_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound, TypeAsc); +impl_t_for_enum!(Expr; Lit, Accessor, Array, Tuple, Dict, Record, BinOp, UnaryOp, Call, Lambda, Decl, Def, ClassDef, AttrDef, Code, Compound, TypeAsc); impl Expr { pub fn receiver_t(&self) -> Option<&Type> { diff --git a/compiler/erg_compiler/lower.rs b/compiler/erg_compiler/lower.rs index 6e22fcc4..26297ab3 100644 --- a/compiler/erg_compiler/lower.rs +++ b/compiler/erg_compiler/lower.rs @@ -968,6 +968,25 @@ impl ASTLowerer { } } + fn lower_type_asc(&mut self, tasc: ast::TypeAscription) -> LowerResult { + log!(info "entered {}({tasc})", fn_name!()); + let t = self.ctx.instantiate_typespec( + &tasc.t_spec, + None, + &mut None, + RegistrationMode::Normal, + )?; + let expr = self.lower_expr(*tasc.expr)?; + self.ctx.sub_unify( + expr.ref_t(), + &t, + Some(expr.loc()), + None, + Some(&Str::from(expr.to_string())), + )?; + Ok(hir::TypeAscription::new(expr, tasc.t_spec)) + } + // Call.obj == Accessor cannot be type inferred by itself (it can only be inferred with arguments) // so turn off type checking (check=false) fn lower_expr(&mut self, expr: ast::Expr) -> LowerResult { @@ -985,6 +1004,7 @@ impl ASTLowerer { ast::Expr::Lambda(lambda) => Ok(hir::Expr::Lambda(self.lower_lambda(lambda)?)), ast::Expr::Def(def) => Ok(hir::Expr::Def(self.lower_def(def)?)), ast::Expr::ClassDef(defs) => Ok(hir::Expr::ClassDef(self.lower_class_def(defs)?)), + ast::Expr::TypeAsc(tasc) => Ok(hir::Expr::TypeAsc(self.lower_type_asc(tasc)?)), other => todo!("{other}"), } }