Use correct type for this in methods for generic structs

1 files changed, 43 insertions, 0 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
index 09df039c..c1954ea5 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
@@ -397,4 +397,47 @@ public class TypeCheckerStructDefinitionTests {
             ))
         );
     }
+
+    @Test
+    public void canTypeCheckMethodOnGenericStructReturningThis() {
+        var untyped = UntypedStructDefinitionNode.builder()
+            .withName("ExampleStruct")
+            .withTypeLevelParams(List.of(Untyped.typeParam("T")))
+            .addMethod(
+                UntypedFunctionDefinitionNode.builder()
+                    .addPositionalParam("this", Optional.empty())
+                    .withName("a")
+                    .withReturnType(Untyped.constructedType(
+                        Untyped.typeLevelReference("ExampleStruct"),
+                        List.of(Untyped.typeLevelReference("T"))
+                    ))
+                    .addStatement(Untyped.functionReturn(Untyped.reference("this")))
+                    .build()
+            )
+            .build();
+        var context = TypeCheckerContext.EMPTY;
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        var structTypeInfo = typed.context().toTypeInfos()
+            .structTypeInfo(Types.struct(NAMESPACE_NAME, "ExampleStruct"))
+            .orElseThrow();
+        assertThat(
+            structTypeInfo,
+            has("methods", x -> x.methods(), isSequence(
+                equalTo(new MethodInfo(
+                    "a",
+                    FunctionType.builder()
+                        .withReturnType(Types.constructedType(
+                            Types.typeConstructor(
+                                List.of(TypeParam.forStructType(Types.struct(NAMESPACE_NAME, "ExampleStruct"), "T")),
+                                Types.struct(NAMESPACE_NAME, "ExampleStruct")
+                            ),
+                            List.of(TypeParam.forStructType(Types.struct(NAMESPACE_NAME, "ExampleStruct"), "T"))
+                        ))
+                        .build()
+                ))
+            ))
+        );
+    }
 }