chore: union method type

mtshiba · mtshiba · commit 8eb8cd7c4a70 · 2024-08-19T12:30:23.000+09:00
diff --git a/crates/els/completion.rs b/crates/els/completion.rs
@@ -23,7 +23,7 @@ use erg_compiler::erg_parser::token::TokenKind;
 use erg_compiler::hir::Expr;
 use erg_compiler::module::SharedCompilerResource;
 use erg_compiler::ty::{HasType, ParamTy, Type};
-use erg_compiler::varinfo::{AbsLocation, VarInfo};
+use erg_compiler::varinfo::{AbsLocation, Mutability, VarInfo, VarKind};
 use TokenKind::*;
 
 use lsp_types::{
@@ -35,17 +35,40 @@ use crate::_log;
 use crate::server::{ELSResult, Flags, RedirectableStdout, Server};
 use crate::util::{self, loc_to_pos, NormalizedUrl};
 
-fn comp_item_kind(vi: &VarInfo) -> CompletionItemKind {
-    match &vi.t {
+fn comp_item_kind(t: &Type, muty: Mutability) -> CompletionItemKind {
+    match t {
         Type::Subr(subr) if subr.self_t().is_some() => CompletionItemKind::METHOD,
         Type::Quantified(quant) if quant.self_t().is_some() => CompletionItemKind::METHOD,
         Type::Subr(_) | Type::Quantified(_) => CompletionItemKind::FUNCTION,
         Type::ClassType => CompletionItemKind::CLASS,
         Type::TraitType => CompletionItemKind::INTERFACE,
+        Type::Or(l, r) => {
+            let l = comp_item_kind(l, muty);
+            let r = comp_item_kind(r, muty);
+            if l == r {
+                l
+            } else if muty.is_const() {
+                CompletionItemKind::CONSTANT
+            } else {
+                CompletionItemKind::VARIABLE
+            }
+        }
+        Type::And(l, r) => {
+            let l = comp_item_kind(l, muty);
+            let r = comp_item_kind(r, muty);
+            if l == CompletionItemKind::VARIABLE {
+                r
+            } else {
+                l
+            }
+        }
+        Type::Refinement(r) => comp_item_kind(&r.t, muty),
+        Type::Bounded { sub, .. } => comp_item_kind(sub, muty),
         t if matches!(&t.qual_name()[..], "Module" | "PyModule" | "GenericModule") => {
             CompletionItemKind::MODULE
         }
-        _ if vi.muty.is_const() => CompletionItemKind::CONSTANT,
+        Type::Type => CompletionItemKind::CONSTANT,
+        _ if muty.is_const() => CompletionItemKind::CONSTANT,
         _ => CompletionItemKind::VARIABLE,
     }
 }
@@ -112,21 +135,24 @@ impl CompletionOrder {
 }
 
 pub struct CompletionOrderSetter<'b> {
-    vi: &'b VarInfo,
+    t: &'b Type,
+    kind: &'b VarKind,
     arg_pt: Option<&'b ParamTy>,
     mod_ctx: &'b Context, // for subtype judgement, not for variable lookup
     label: String,
 }
 
 impl<'b> CompletionOrderSetter<'b> {
     pub fn new(
-        vi: &'b VarInfo,
+        t: &'b Type,
+        kind: &'b VarKind,
         arg_pt: Option<&'b ParamTy>,
         mod_ctx: &'b Context,
         label: String,
     ) -> Self {
         Self {
-            vi,
+            t,
+            kind,
             arg_pt,
             mod_ctx,
             label,
@@ -140,7 +166,7 @@ impl<'b> CompletionOrderSetter<'b> {
         } else if self.label.starts_with('_') {
             orders.push(CompletionOrder::Escaped);
         }
-        if self.vi.kind.is_builtin() {
+        if self.kind.is_builtin() {
             orders.push(CompletionOrder::Builtin);
         }
         if self
@@ -152,11 +178,11 @@ impl<'b> CompletionOrderSetter<'b> {
         #[allow(clippy::blocks_in_conditions)]
         if self
             .arg_pt
-            .map_or(false, |pt| self.mod_ctx.subtype_of(&self.vi.t, pt.typ()))
+            .map_or(false, |pt| self.mod_ctx.subtype_of(self.t, pt.typ()))
         {
             orders.push(CompletionOrder::TypeMatched);
         } else if self.arg_pt.map_or(false, |pt| {
-            let Some(return_t) = self.vi.t.return_t() else {
+            let Some(return_t) = self.t.return_t() else {
                 return false;
             };
             if return_t.has_qvar() {
@@ -196,7 +222,7 @@ fn external_item(name: &str, vi: &VarInfo, mod_name: &str) -> CompletionItem {
     let mut item =
         CompletionItem::new_simple(format!("{name} (import from {mod_name})"), vi.t.to_string());
     item.sort_text = Some(format!("{}_{}", CompletionOrder::STD_ITEM, item.label));
-    item.kind = Some(comp_item_kind(vi));
+    item.kind = Some(comp_item_kind(&vi.t, vi.muty));
     let import = if PYTHON_MODE {
         format!("from {mod_name} import {name}\n")
     } else {
@@ -467,10 +493,16 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
                     continue;
                 }
                 let mut item = CompletionItem::new_simple(label, vi.t.to_string());
-                CompletionOrderSetter::new(vi, arg_pt.as_ref(), mod_ctx, item.label.clone())
-                    .set(&mut item);
+                CompletionOrderSetter::new(
+                    &vi.t,
+                    &vi.kind,
+                    arg_pt.as_ref(),
+                    mod_ctx,
+                    item.label.clone(),
+                )
+                .set(&mut item);
                 // item.sort_text = Some(format!("{}_{}", CompletionOrder::OtherNamespace, item.label));
-                item.kind = Some(comp_item_kind(vi));
+                item.kind = Some(comp_item_kind(&vi.t, vi.muty));
                 item.data = Some(Value::String(vi.def_loc.to_string()));
                 let import = if PYTHON_MODE {
                     format!("from {path} import {name}\n")
@@ -590,6 +622,25 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
             _log!(self, "module context not found: {uri}");
             return Ok(Some(CompletionResponse::Array(result)));
         };
+        if PYTHON_MODE {
+            if let Some(receiver_t) = &receiver_t {
+                for (field, ty) in mod_ctx.context.fields(receiver_t) {
+                    let mut item =
+                        CompletionItem::new_simple(field.symbol.to_string(), ty.to_string());
+                    CompletionOrderSetter::new(
+                        &ty,
+                        &VarKind::Builtin,
+                        arg_pt.as_ref(),
+                        &mod_ctx.context,
+                        item.label.clone(),
+                    )
+                    .set(&mut item);
+                    item.kind = Some(comp_item_kind(&ty, Mutability::Immutable));
+                    already_appeared.insert(item.label.clone());
+                    result.push(item);
+                }
+            }
+        }
         for (name, vi) in contexts.into_iter().flat_map(|ctx| ctx.local_dir()) {
             if comp_kind.should_be_method() && vi.vis.is_private() {
                 continue;
@@ -621,9 +672,15 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
             }
             let readable_t = mod_ctx.context.readable_type(vi.t.clone());
             let mut item = CompletionItem::new_simple(label, readable_t.to_string());
-            CompletionOrderSetter::new(vi, arg_pt.as_ref(), &mod_ctx.context, item.label.clone())
-                .set(&mut item);
-            item.kind = Some(comp_item_kind(vi));
+            CompletionOrderSetter::new(
+                &vi.t,
+                &vi.kind,
+                arg_pt.as_ref(),
+                &mod_ctx.context,
+                item.label.clone(),
+            )
+            .set(&mut item);
+            item.kind = Some(comp_item_kind(&vi.t, vi.muty));
             item.data = Some(Value::String(vi.def_loc.to_string()));
             already_appeared.insert(item.label.clone());
             result.push(item);
diff --git a/crates/erg_compiler/context/compare.rs b/crates/erg_compiler/context/compare.rs
@@ -215,7 +215,11 @@ impl Context {
                 Some((Type::Never, Type::Obj)) => (Absolutely, true),
                 _ => (Maybe, false),
             },
-            (Mono(n), Subr(_) | Quantified(_)) if &n[..] == "Subroutine" => (Absolutely, true),
+            (Mono(n), Subr(_) | Quantified(_))
+                if &n[..] == "Subroutine" || &n[..] == "GenericCallable" =>
+            {
+                (Absolutely, true)
+            }
             (lhs, rhs) if lhs.is_mono_value_class() && rhs.is_mono_value_class() => {
                 (Absolutely, false)
             }
diff --git a/crates/erg_compiler/context/initialize/classes.rs b/crates/erg_compiler/context/initialize/classes.rs
@@ -2673,6 +2673,9 @@ impl Context {
             Visibility::BUILTIN_PUBLIC,
         );
         bytes.register_trait_methods(mono(BYTES), bytes_seq);
+        bytes
+            .register_trait(self, poly(SEQUENCE, vec![ty_tp(Int)]))
+            .unwrap();
         let mut bytes_eq = Self::builtin_methods(Some(mono(EQ)), 2);
         bytes_eq.register_builtin_erg_impl(
             OP_EQ,
@@ -3581,6 +3584,9 @@ impl Context {
             Visibility::BUILTIN_PUBLIC,
         );
         bytearray_mut.register_trait_methods(mono(MUT_BYTEARRAY), bytearray_seq);
+        bytearray_mut
+            .register_trait(self, poly(SEQUENCE, vec![ty_tp(Int)]))
+            .unwrap();
         let t_append = pr_met(
             ref_mut(bytearray_mut_t.clone(), None),
             vec![kw(KW_ELEM, int_interval(IntervalOp::Closed, 0, 255))],
diff --git a/crates/erg_compiler/context/inquire.rs b/crates/erg_compiler/context/inquire.rs
@@ -1379,11 +1379,27 @@ impl Context {
                 }
                 return Ok(method.method_info.clone());
             }
-            Triple::Err(err) => {
+            Triple::Err(err) if ERG_MODE => {
                 return Err(err);
             }
             _ => {}
         }
+        if PYTHON_MODE {
+            if let Some(subr_t) = self.get_union_attr_type_by_name(attr_name) {
+                let muty = Mutability::from(&attr_name.inspect()[..]);
+                let vi = VarInfo::new(
+                    subr_t,
+                    muty,
+                    Visibility::DUMMY_PUBLIC,
+                    VarKind::Builtin,
+                    None,
+                    ContextKind::Dummy,
+                    None,
+                    AbsLocation::unknown(),
+                );
+                return Ok(vi);
+            }
+        }
         for patch in self.find_patches_of(obj.ref_t()) {
             if let Some(vi) = patch.get_current_scope_non_param(&attr_name.name) {
                 self.validate_visibility(attr_name, vi, input, namespace)?;
@@ -1508,7 +1524,7 @@ impl Context {
                 }
                 return Ok(method.method_info.clone());
             }
-            Triple::Err(err) => {
+            Triple::Err(err) if ERG_MODE => {
                 return Err(err);
             }
             _ => {}
@@ -3551,6 +3567,7 @@ impl Context {
             // if all methods have the same return type, the minimum type (has biggest param types) is selected
             // e.g. [Float -> Bool, Int -> Bool] => Float -> Bool
             // REVIEW: should [Int -> Bool, Str -> Bool] => (Str or Int) -> Bool?
+            // -> get_union_method_type
             if let Some(min) = self.min_type(candidates.iter().map(|mp| &mp.method_info.t)) {
                 let min_pair = candidates
                     .iter()
@@ -3574,6 +3591,55 @@ impl Context {
         ))
     }
 
+    // (Int -> Bool, Float -> Bool) => Int or Float -> Bool
+    fn get_union_method_type(&self, candidates: &[MethodPair]) -> Option<Type> {
+        let fst = candidates.first()?;
+        let mut kind = fst.method_info.t.subr_kind()?;
+        let mut union_nds = fst.method_info.t.non_default_params()?.clone();
+        let mut union_var = fst.method_info.t.var_params().cloned();
+        let mut union_ds = fst.method_info.t.default_params()?.clone();
+        let mut union_kw_var = fst.method_info.t.kw_var_params().cloned();
+        let mut union_return = fst.method_info.t.return_t()?.clone();
+        for cand in candidates.iter().skip(1) {
+            kind = kind | cand.method_info.t.subr_kind()?;
+            for (union, r) in union_nds
+                .iter_mut()
+                .zip(cand.method_info.t.non_default_params()?)
+            {
+                *union.typ_mut() = self.union(union.typ(), r.typ());
+            }
+            if let Some((union, r)) = union_var.as_mut().zip(cand.method_info.t.var_params()) {
+                *union.typ_mut() = self.union(union.typ(), r.typ());
+            }
+            for (union, r) in union_ds
+                .iter_mut()
+                .zip(cand.method_info.t.default_params()?)
+            {
+                *union.typ_mut() = self.union(union.typ(), r.typ());
+            }
+            if let Some((union, r)) = union_kw_var
+                .as_mut()
+                .zip(cand.method_info.t.kw_var_params())
+            {
+                *union.typ_mut() = self.union(union.typ(), r.typ());
+            }
+            union_return = self.union(&union_return, cand.method_info.t.return_t()?);
+        }
+        let subr = Type::Subr(SubrType::new(
+            kind,
+            union_nds,
+            union_var,
+            union_ds,
+            union_kw_var,
+            union_return,
+        ));
+        if subr.has_qvar() {
+            Some(subr.quantify())
+        } else {
+            Some(subr)
+        }
+    }
+
     /// Infer the receiver type from the attribute name.
     /// Returns an error if multiple candidates are found. If nothing is found, returns None.
     fn get_attr_type_by_name(
@@ -3595,6 +3661,20 @@ impl Context {
         }
     }
 
+    fn get_union_attr_type_by_name(&self, attr: &Identifier) -> Option<Type> {
+        if let Some(candidates) = self.method_to_traits.get(attr.inspect()) {
+            return self.get_union_method_type(candidates);
+        }
+        if let Some(candidates) = self.method_to_classes.get(attr.inspect()) {
+            return self.get_union_method_type(candidates);
+        }
+        if let Some(outer) = self.get_outer_scope_or_builtins() {
+            outer.get_union_attr_type_by_name(attr)
+        } else {
+            None
+        }
+    }
+
     fn _get_gen_t_require_attr_t<'a>(
         &'a self,
         gen: &'a GenTypeObj,
diff --git a/crates/erg_compiler/ty/mod.rs b/crates/erg_compiler/ty/mod.rs
@@ -973,6 +973,16 @@ impl From<TokenKind> for SubrKind {
     }
 }
 
+impl BitOr for SubrKind {
+    type Output = Self;
+    fn bitor(self, rhs: Self) -> Self {
+        match (self, rhs) {
+            (Self::Func, Self::Func) => Self::Func,
+            _ => Self::Proc,
+        }
+    }
+}
+
 impl SubrKind {
     pub const fn arrow(&self) -> Str {
         match self {
@@ -2488,6 +2498,17 @@ impl Type {
         }
     }
 
+    pub fn subr_kind(&self) -> Option<SubrKind> {
+        match self {
+            Self::FreeVar(fv) if fv.is_linked() => fv.crack().subr_kind(),
+            Self::Subr(subr) => Some(subr.kind),
+            Self::Refinement(refine) => refine.t.subr_kind(),
+            Self::Quantified(quant) => quant.subr_kind(),
+            Self::And(l, r) => l.subr_kind().and_then(|k| r.subr_kind().map(|k2| k | k2)),
+            _ => None,
+        }
+    }
+
     pub fn is_quantified_subr(&self) -> bool {
         match self {
             Self::FreeVar(fv) if fv.is_linked() => fv.crack().is_quantified_subr(),
