frontend: Support generic impls for index traits

dinfuehr · Nov 12, 2024 · 2569aa1 · 2569aa1
1 parent 67537a5
commit 2569aa1
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 7 deletions.
diff --git a/dora-frontend/src/typeck/call.rs b/dora-frontend/src/typeck/call.rs
@@ -211,16 +211,16 @@ fn check_expr_call_expr(
             .get_method_for_trait_method_id(trait_method_id)
             .expect("method not found");
 
-        let call_type = CallType::Method(expr_type.clone(), method_id, SourceTypeArray::empty());
+        let call_type = CallType::Method(expr_type.clone(), method_id, impl_match.bindings.clone());
         ck.analysis
             .map_calls
             .insert_or_replace(e.id, Arc::new(call_type));
 
         let method = ck.sa.fct(method_id);
 
-        check_args_compatible_fct(ck, method, arguments, &SourceTypeArray::empty(), None);
+        check_args_compatible_fct(ck, method, arguments, &impl_match.bindings, None);
 
-        let return_type = method.return_type();
+        let return_type = specialize_type(ck.sa, method.return_type(), &impl_match.bindings);
         ck.analysis.set_ty(e.id, return_type.clone());
 
         return_type

diff --git a/dora-frontend/src/typeck/expr.rs b/dora-frontend/src/typeck/expr.rs
@@ -15,6 +15,7 @@ use crate::sema::{
     LazyLambdaCreationData, LazyLambdaId, ModuleDefinitionId, NestedVarId, Param, Params, Sema,
     SourceFileId, TraitDefinitionId,
 };
+use crate::specialize_type;
 use crate::ty::TraitType;
 use crate::typeck::{
     arg_allows, check_args_compatible, check_expr_break_and_continue, check_expr_call,
@@ -316,17 +317,25 @@ fn check_expr_assign_call(ck: &mut TypeCheck, e: &ast::ExprBinType) {
             .expect("missing alias");
         let impl_item_type_alias = ck.sa.alias(impl_item_type_alias_id);
 
-        let call_type = CallType::Method(expr_type.clone(), method_id, SourceTypeArray::empty());
+        let call_type = CallType::Method(expr_type.clone(), method_id, impl_match.bindings.clone());
         ck.analysis
             .map_calls
             .insert_or_replace(e.id, Arc::new(call_type));
 
         let method = ck.sa.fct(method_id);
         let index_param = &method.params.params[1..2];
-        check_args_compatible(ck, index_param, None, args, &SourceTypeArray::empty(), None);
+        check_args_compatible(ck, index_param, None, args, &impl_match.bindings, None);
 
-        if !arg_allows(ck.sa, impl_item_type_alias.ty(), value_type.clone(), None) {
-            let exp = ck.ty_name(&impl_item_type_alias.ty());
+        let impl_item_type_alias_ty = impl_item_type_alias.ty();
+        let impl_item_type_alias_ty =
+            specialize_type(ck.sa, impl_item_type_alias_ty, &impl_match.bindings);
+        if !arg_allows(
+            ck.sa,
+            impl_item_type_alias_ty.clone(),
+            value_type.clone(),
+            None,
+        ) {
+            let exp = ck.ty_name(&impl_item_type_alias_ty);
             let got = ck.ty_name(&value_type);
             ck.sa.report(
                 ck.file_id,

diff --git a/dora-frontend/src/typeck/tests.rs b/dora-frontend/src/typeck/tests.rs
@@ -5300,6 +5300,50 @@ fn class_index_set() {
     ")
 }
 
+#[test]
+fn class_index_get_generic() {
+    ok("
+        class Foo[T] { a: T, b: T }
+        impl[T] std::traits::IndexGet for Foo[T] {
+            type Index = Int;
+            type Item = T;
+            fn get(index: Self::Index): Self::Item {
+                if index == 0 {
+                    self.a
+                } else {
+                    assert(index == 1);
+                    self.b
+                }
+            }
+        }
+        fn f(x: Foo[Float64]): Float64 {
+            x(0)
+        }
+    ")
+}
+
+#[test]
+fn class_index_set_generic() {
+    ok("
+        class Foo[T] { a: T, b: T }
+        impl[T] std::traits::IndexSet for Foo[T] {
+            type Index = Int;
+            type Item = T;
+            fn set(index: Self::Index, value: Self::Item) {
+                if index == 0 {
+                    self.a = value;
+                } else {
+                    assert(index == 1);
+                    self.b = value;
+                }
+            }
+        }
+        fn f(x: Foo[Float64], value: Float64) {
+            x(0) = value;
+        }
+    ")
+}
+
 #[test]
 fn class_index_set_wrong_type() {
     err(