Add TypeAscription

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)?;

compiler/erg_compiler/hir.rs

@@ -1485,6 +1485,49 @@ impl AttrDef {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct TypeAscription {
+    pub expr: Box<Expr>,
+    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_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);
+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> {

compiler/erg_compiler/lower.rs

@@ -968,6 +968,25 @@ impl ASTLowerer {
         }
     }
 
+    fn lower_type_asc(&mut self, tasc: ast::TypeAscription) -> LowerResult<hir::TypeAscription> {
+        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<hir::Expr> {
@@ -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}"),
         }
     }