Merge pull request #452 from erg-lang/proj-type-member

Support trait associated types
erg-lang · Aug 26, 2023 · 1c81a02 · 1c81a02
2 parents e51e7a2 + 0cf2522
commit 1c81a02
Show file tree

Hide file tree

Showing 13 changed files with 103 additions and 33 deletions.
diff --git a/crates/erg_compiler/context/eval.rs b/crates/erg_compiler/context/eval.rs
@@ -20,17 +20,20 @@ use erg_parser::desugar::Desugarer;
 use erg_parser::token::{Token, TokenKind};
 
 use crate::ty::constructors::{
-    array_t, dict_t, mono, named_free_var, poly, proj, proj_call, ref_, ref_mut, refinement, set_t,
-    subr_t, tp_enum, tuple_t, v_enum,
+    array_t, dict_t, mono, mono_q, named_free_var, poly, proj, proj_call, ref_, ref_mut,
+    refinement, set_t, subr_t, subtypeof, tp_enum, tuple_t, v_enum,
 };
 use crate::ty::free::{Constraint, HasLevel};
 use crate::ty::typaram::{OpKind, TyParam};
 use crate::ty::value::{GenTypeObj, TypeObj, ValueObj};
-use crate::ty::{ConstSubr, HasType, Predicate, SubrKind, Type, UserConstSubr, ValueArgs};
+use crate::ty::{
+    ConstSubr, HasType, Predicate, SubrKind, Type, UserConstSubr, ValueArgs, Visibility,
+};
 
 use crate::context::instantiate_spec::ParamKind;
 use crate::context::{ClassDefType, Context, ContextKind, RegistrationMode};
 use crate::error::{EvalError, EvalErrors, EvalResult, SingleEvalResult};
+use crate::varinfo::VarInfo;
 
 use super::instantiate::TyVarCache;
 use Type::{Failure, Never, Subr};
@@ -794,7 +797,7 @@ impl Context {
             let elem = record_ctx.eval_const_block(&attr.body.block)?;
             let ident = match &attr.sig {
                 Signature::Var(var) => match &var.pat {
-                    VarPattern::Ident(ident) => self.instantiate_field(ident)?,
+                    VarPattern::Ident(ident) => record_ctx.instantiate_field(ident)?,
                     other => {
                         return feature_error!(self, other.loc(), &format!("record field: {other}"))
                     }
@@ -803,6 +806,21 @@ impl Context {
                     return feature_error!(self, other.loc(), &format!("record field: {other}"))
                 }
             };
+            let name = VarName::from_str(ident.symbol.clone());
+            // T = Trait { .Output = Type; ... }
+            // -> .Output = Self(<: T).Output
+            if self.kind.is_trait() && self.convert_value_into_type(elem.clone()).is_ok() {
+                let slf = mono_q("Self", subtypeof(mono(self.name.clone())));
+                let t = ValueObj::builtin_type(slf.proj(ident.symbol.clone()));
+                record_ctx.consts.insert(name.clone(), t);
+            } else {
+                record_ctx.consts.insert(name.clone(), elem.clone());
+            }
+            let t = v_enum(set! { elem.clone() });
+            let vis = record_ctx.instantiate_vis_modifier(attr.sig.vis())?;
+            let vis = Visibility::new(vis, record_ctx.name.clone());
+            let vi = VarInfo::record_field(t, record_ctx.absolutize(attr.sig.loc()), vis);
+            record_ctx.locals.insert(name, vi);
             attrs.push((ident, elem));
         }
         Ok(ValueObj::Record(attrs.into_iter().collect()))
@@ -1795,10 +1813,7 @@ impl Context {
     pub(crate) fn convert_value_into_type(&self, val: ValueObj) -> Result<Type, ValueObj> {
         match val {
             ValueObj::Ellipsis => Ok(Type::Ellipsis),
-            ValueObj::Type(t) => match t {
-                TypeObj::Builtin { t, .. } => Ok(t),
-                TypeObj::Generated(gen) => Ok(gen.into_typ()),
-            },
+            ValueObj::Type(t) => Ok(t.into_typ()),
             ValueObj::Record(rec) => {
                 let mut fields = dict! {};
                 for (name, val) in rec.into_iter() {
@@ -1998,7 +2013,7 @@ impl Context {
             // obj: [T; N]|<: Add([T; M])|.Output == ValueObj::Type(<type [T; M+N]>)
             if let ValueObj::Type(quant_projected_t) = obj {
                 let projected_t = quant_projected_t.into_typ();
-                let (quant_sub, _) = self.get_type(&sub.qual_name()).unwrap();
+                let (quant_sub, _) = self.get_type_and_ctx(&sub.qual_name()).unwrap();
                 let _sup_subs = if let Some((sup, quant_sup)) = opt_sup.zip(methods.impl_of()) {
                     // T -> Int, M -> 2
                     match Substituter::substitute_typarams(self, &quant_sup, sup) {

diff --git a/crates/erg_compiler/context/generalize.rs b/crates/erg_compiler/context/generalize.rs
@@ -592,7 +592,9 @@ impl<'c, 'q, 'l, L: Locational> Dereferencer<'c, 'q, 'l, L> {
                             Ok(ty)
                         }
                         Err(errs) => {
-                            Type::FreeVar(fv).destructive_link(&Never);
+                            if !fv.is_generalized() {
+                                Type::FreeVar(fv).destructive_link(&Never);
+                            }
                             Err(errs)
                         }
                     }

diff --git a/crates/erg_compiler/context/inquire.rs b/crates/erg_compiler/context/inquire.rs
@@ -213,7 +213,7 @@ impl Context {
         self.get_mod(ident.inspect())
             .map(|ctx| vec![ctx])
             .or_else(|| {
-                let (typ, _) = self.get_type(ident.inspect())?;
+                let (typ, _) = self.get_type_and_ctx(ident.inspect())?;
                 self.get_nominal_super_type_ctxs(typ)
             })
             .or_else(|| self.rec_get_patch(ident.inspect()).map(|ctx| vec![ctx]))
@@ -2829,7 +2829,7 @@ impl Context {
         }
     }
 
-    pub(crate) fn get_type(&self, name: &str) -> Option<(&Type, &Context)> {
+    pub(crate) fn get_type_and_ctx(&self, name: &str) -> Option<(&Type, &Context)> {
         if let Some((t, ctx)) = self.rec_local_get_type(name) {
             return Some((t, ctx));
         }
@@ -2848,7 +2848,8 @@ impl Context {
     }
 
     pub fn get_type_info_by_str(&self, name: &str) -> Option<(&VarName, &VarInfo)> {
-        self.get_type(name).and_then(|(t, _)| self.get_type_info(t))
+        self.get_type_and_ctx(name)
+            .and_then(|(t, _)| self.get_type_info(t))
     }
 
     /// you should use `get_type` instead of this

diff --git a/crates/erg_compiler/context/instantiate_spec.rs b/crates/erg_compiler/context/instantiate_spec.rs
@@ -550,8 +550,14 @@ impl Context {
                         return Ok(t);
                     }
                 }
-                if let Some((typ, _)) = self.get_type(ident.inspect()) {
+                if let Some((typ, _)) = self.get_type_and_ctx(ident.inspect()) {
                     Ok(typ.clone())
+                } else if let Some(typ) = self
+                    .consts
+                    .get(ident.inspect())
+                    .and_then(|v| self.convert_value_into_type(v.clone()).ok())
+                {
+                    Ok(typ)
                 } else if not_found_is_qvar {
                     let tyvar = named_free_var(Str::rc(other), self.level, Constraint::Uninited);
                     tmp_tv_cache.push_or_init_tyvar(&ident.name, &tyvar, self);
@@ -715,7 +721,7 @@ impl Context {
                 Ok(Type::NamedTuple(ts))
             }
             other => {
-                let Some((typ, ctx)) = self.get_type(&Str::rc(other)) else {
+                let Some((typ, ctx)) = self.get_type_and_ctx(&Str::rc(other)) else {
                     return Err(TyCheckErrors::from(TyCheckError::no_type_error(
                         self.cfg.input.clone(),
                         line!() as usize,

diff --git a/crates/erg_compiler/context/register.rs b/crates/erg_compiler/context/register.rs
@@ -907,6 +907,7 @@ impl Context {
         Ok(())
     }
 
+    /// Registers type definitions of types and constants; unlike `register_const`, this does not evaluate terms.
     pub(crate) fn preregister_const(&mut self, block: &ast::Block) -> TyCheckResult<()> {
         let mut total_errs = TyCheckErrors::empty();
         for expr in block.iter() {

diff --git a/crates/erg_compiler/declare.rs b/crates/erg_compiler/declare.rs
@@ -528,7 +528,7 @@ impl ASTLowerer {
 
     fn get_tv_ctx(&self, ident: &ast::Identifier, args: &ast::Args) -> TyVarCache {
         let mut tv_ctx = TyVarCache::new(self.module.context.level, &self.module.context);
-        if let Some((t, _)) = self.module.context.get_type(ident.inspect()) {
+        if let Some((t, _)) = self.module.context.get_type_and_ctx(ident.inspect()) {
             for (tp, arg) in t.typarams().iter().zip(args.pos_args()) {
                 if let ast::Expr::Accessor(ast::Accessor::Ident(ident)) = &arg.expr {
                     tv_ctx.push_or_init_typaram(&ident.name, tp, &self.module.context);

diff --git a/crates/erg_compiler/lower.rs b/crates/erg_compiler/lower.rs
@@ -389,6 +389,17 @@ impl ASTLowerer {
         self.module
             .context
             .grow("<record>", ContextKind::Dummy, Private, None);
+        for attr in record.attrs.iter() {
+            if attr.sig.is_const() {
+                self.module
+                    .context
+                    .register_const_def(attr)
+                    .map_err(|errs| {
+                        self.pop_append_errs();
+                        errs
+                    })?;
+            }
+        }
         for attr in record.attrs.into_iter() {
             let attr = self.lower_def(attr).map_err(|errs| {
                 self.pop_append_errs();

diff --git a/crates/erg_compiler/ty/const_subr.rs b/crates/erg_compiler/ty/const_subr.rs
@@ -269,6 +269,11 @@ impl ConstSubr {
                             subr.default_params.clone(),
                             return_t,
                         );
+                        let subr_t = if subr_t.has_qvar() {
+                            subr_t.quantify()
+                        } else {
+                            subr_t
+                        };
                         return Some(subr_t);
                     }
                 }

diff --git a/crates/erg_compiler/ty/value.rs b/crates/erg_compiler/ty/value.rs
@@ -1454,23 +1454,10 @@ impl ValueObj {
     pub fn as_type(&self, ctx: &Context) -> Option<TypeObj> {
         match self {
             Self::Type(t) => Some(t.clone()),
-            Self::Record(rec) => {
-                let mut attr_ts = dict! {};
-                for (k, v) in rec.iter() {
-                    attr_ts.insert(k.clone(), v.as_type(ctx)?.typ().clone());
-                }
-                Some(TypeObj::builtin_type(Type::Record(attr_ts)))
-            }
-            Self::Subr(subr) => subr.as_type(ctx).map(TypeObj::builtin_type),
-            Self::Array(elems) | Self::Tuple(elems) => {
-                log!(err "as_type({})", erg_common::fmt_vec(elems));
-                None
-            }
-            Self::Dict(elems) => {
-                log!(err "as_type({elems})");
-                None
-            }
-            _other => None,
+            other => ctx
+                .convert_value_into_type(other.clone())
+                .ok()
+                .map(TypeObj::builtin_type),
         }
     }
 }

diff --git a/crates/erg_compiler/ty/vis.rs b/crates/erg_compiler/ty/vis.rs
@@ -115,6 +115,13 @@ impl Visibility {
         }
     }
 
+    pub fn public(namespace: Str) -> Self {
+        Self {
+            modifier: VisibilityModifier::Public,
+            def_namespace: namespace,
+        }
+    }
+
     pub const fn is_public(&self) -> bool {
         self.modifier.is_public()
     }

diff --git a/crates/erg_compiler/varinfo.rs b/crates/erg_compiler/varinfo.rs
@@ -361,6 +361,19 @@ impl VarInfo {
         )
     }
 
+    pub fn record_field(t: Type, def_loc: AbsLocation, vis: Visibility) -> Self {
+        Self::new(
+            t,
+            Immutable,
+            vis,
+            VarKind::Declared,
+            None,
+            None,
+            None,
+            def_loc,
+        )
+    }
+
     pub fn is_untyped_parameter(&self) -> bool {
         self.kind.is_parameter() && self.t.is_unbound_var()
     }

diff --git a/tests/should_ok/associated_types.er b/tests/should_ok/associated_types.er
@@ -0,0 +1,17 @@
+IO = Trait {
+    .Inp = Type
+    .Output = Type
+    .func = (self, x: .Inp) -> .Output
+}
+
+C = Class()
+C|<: IO|.
+    Inp = Int
+    Output = Bool
+    func self, i =
+        _ = self
+        i >= 0
+
+io x, y = x.func y
+
+assert io C.new(), 1
diff --git a/tests/test.rs b/tests/test.rs
@@ -27,6 +27,11 @@ fn exec_assert_cast_ok() -> Result<(), ()> {
     expect_success("tests/should_ok/assert_cast.er", 0)
 }
 
+#[test]
+fn exec_associated_types() -> Result<(), ()> {
+    expect_success("tests/should_ok/associated_types.er", 0)
+}
+
 #[test]
 fn exec_class() -> Result<(), ()> {
     expect_success("examples/class.er", 0)