TSLANG_ISSUES 13: generic inference through Ref<T> and union members by type

tslang/lib/TypeScript/MLIRGenGenerics.cpp

@@ -158,6 +158,22 @@ namespace mlirgen
         return false;
     }
 
+    // `Ref<T>` from a `Ref<Item>`. Without it a parameter `Ref<T>`, or a tuple alias holding one
+    // (`Vec<T> = [first: Ref<T> | null, ...]`), inferred nothing, and T had to be given explicitly.
+    bool MLIRGenImpl::tryInferRef(mlir::Location location, mlir::Type templateType, mlir::Type concreteType,
+                     StringMap<mlir::Type> &results, const GenContext &genContext)
+    {
+        auto tempRef = dyn_cast<mlir_ts::RefType>(templateType);
+        auto typeRef = dyn_cast<mlir_ts::RefType>(concreteType);
+        if (!tempRef || !typeRef)
+        {
+            return false;
+        }
+
+        inferType(location, tempRef.getElementType(), typeRef.getElementType(), results, genContext);
+        return true;
+    }
+
     bool MLIRGenImpl::tryInferTuple(mlir::Location location, mlir::Type templateType, mlir::Type concreteType,
                        StringMap<mlir::Type> &results, const GenContext &genContext)
     {
@@ -260,6 +276,36 @@ namespace mlirgen
                 return true;
             }
 
+            // a member both unions have (`null` of `T | null` and `Item | null`) binds nothing: pair
+            // up the rest. By position alone, the members' order decided, and `T` could take `null`.
+            SmallVector<mlir::Type> tempRest;
+            SmallVector<mlir::Type> typeRest;
+            for (auto tempSubType : tempUnionType.getTypes())
+            {
+                if (!llvm::is_contained(types, tempSubType))
+                {
+                    tempRest.push_back(tempSubType);
+                }
+            }
+
+            for (auto subType : types)
+            {
+                if (!llvm::is_contained(tempUnionType.getTypes(), subType))
+                {
+                    typeRest.push_back(subType);
+                }
+            }
+
+            if (tempRest.size() == typeRest.size())
+            {
+                for (auto [index, tempSubType] : enumerate(tempRest))
+                {
+                    inferType(location, tempSubType, typeRest[index], results, genContext);
+                }
+
+                return true;
+            }
+
             for (auto [index, tempSubType] : enumerate(tempUnionType.getTypes()))
             {
                 inferType(location, tempSubType, types[index], results, genContext);
@@ -323,6 +369,12 @@ namespace mlirgen
             return;
         }
 
+        // Ref<T> -> Ref<T>
+        if (tryInferRef(location, templateType, concreteType, results, genContext))
+        {
+            return;
+        }
+
         // tuple -> tuple
         if (tryInferTuple(location, templateType, concreteType, results, genContext))
         {

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -1241,6 +1241,9 @@ class MLIRGenImpl
     bool tryInferArray(mlir::Location location, mlir::Type templateType, mlir::Type concreteType,
                        StringMap<mlir::Type> &results, const GenContext &genContext);
 
+    bool tryInferRef(mlir::Location location, mlir::Type templateType, mlir::Type concreteType,
+                     StringMap<mlir::Type> &results, const GenContext &genContext);
+
     // TODO: finish it
     template <typename T>
     bool tryInferTupleFields(mlir::Location location, mlir_ts::TupleType tempTuple, T typeTuple,

tslang/test/tester/CMakeLists.txt

@@ -488,6 +488,7 @@ tslang_add_test(NAME test-compile-00-rest-args-pack-owned COMMAND test-runner "$
 tslang_add_test(NAME test-compile-00-switch-char-string-cases COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_char_string_cases.ts")
 tslang_add_test(NAME test-compile-00-switch-default-fallthrough COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_default_fallthrough.ts")
 tslang_add_test(NAME test-compile-00-generic-class-constructor-inference COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_class_constructor_inference.ts")
+tslang_add_test(NAME test-compile-00-generic-ref-inference COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_ref_inference.ts")
 tslang_add_test(NAME test-compile-00-generator-done-value-zero COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_done_value_zero.ts")
 tslang_add_test(NAME test-compile-00-any-equality-undefined COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00any_equality_undefined.ts")
 tslang_add_test(NAME test-compile-00-owned-generators COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00owned_generators.ts")
@@ -1019,6 +1020,7 @@ tslang_add_test(NAME test-jit-00-rest-args-pack-owned COMMAND test-runner -jit "
 tslang_add_test(NAME test-jit-00-switch-char-string-cases COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_char_string_cases.ts")
 tslang_add_test(NAME test-jit-00-switch-default-fallthrough COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_default_fallthrough.ts")
 tslang_add_test(NAME test-jit-00-generic-class-constructor-inference COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_class_constructor_inference.ts")
+tslang_add_test(NAME test-jit-00-generic-ref-inference COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_ref_inference.ts")
 tslang_add_test(NAME test-jit-00-generator-done-value-zero COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_done_value_zero.ts")
 tslang_add_test(NAME test-jit-00-any-equality-undefined COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00any_equality_undefined.ts")
 tslang_add_test(NAME test-jit-00-owned-generators COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00owned_generators.ts")
@@ -1905,6 +1907,7 @@ set(TSLANG_CORPUS
     00generator_done_value_zero.ts
     00generic_arguments_name_conflict.ts
     00generic_class_constructor_inference.ts
+    00generic_ref_inference.ts
     00global_const_object_method.ts
     00globals.ts
     00globals2.ts

tslang/test/tester/tests/00generic_ref_inference.ts

@@ -0,0 +1,39 @@
+// regression test: a generic function's T is inferred through Ref<T>, also inside a generic tuple
+// alias, and through `T | null`.
+//
+// bug found and fixed: inference had no case for Ref<T>, so `count(v)` with `v: Vec<Item>`
+// (`Vec<T> = [first: Ref<T> | null, last: Ref<T> | null]`) failed with "not all types could be
+// inferred". And a union was matched member by member in order, so `T | null` against
+// `Item | null` could bind T to null.
+
+type Vec<T> = [first: Ref<T> | null, last: Ref<T> | null];
+type Item = [id: s32, size: s32];
+type Nullable<T> = [value: T | null];
+
+function count<T>(v: Vec<T>): int {
+    return <int>((<index><Opaque>v.last - <index><Opaque>v.first) / sizeof<T>());
+}
+
+function sizeOfRef<T>(r: Ref<T>): int {
+    return sizeof<T>();
+}
+
+function sizeOfNullable<T>(n: Nullable<T>): int {
+    return sizeof<T>();
+}
+
+function main() {
+    let items: Item[] = [[1, 2], [3, 4]];
+    const first = Ref(items[0]);
+    const v: Vec<Item> = [first, <Ref<Item>><Opaque>(<index><Opaque>first + 2 * sizeof<Item>())];
+
+    assert(count<Item>(v) == 2);
+    assert(count(v) == 2);
+
+    assert(sizeOfRef(first) == sizeof<Item>());
+
+    const n: Nullable<Item> = [items[1]];
+    assert(sizeOfNullable(n) == sizeof<Item>());
+
+    print("done.");
+}