An inferred return type keeps undefined and null

tslang/include/TypeScript/MLIRLogic/MLIRGenContext.h

@@ -52,6 +52,10 @@ struct PassResult
 
     mlir::Type functionReturnType;
     bool functionReturnTypeShouldBeProvided;
+    // a `return;` without a value: with value returns as well, the function returns T | undefined
+    bool hasBareReturn = false;
+    // a generator's next(): its {value, done} is Iterator<T>'s, so a finished one's undefined value is not in T
+    bool iteratorResult = false;
     llvm::StringMap<ts::VariableDeclarationDOM::TypePtr> outerVariables;
     SmallVector<mlir_ts::FieldInfo> extraFieldsInThisContext;
 };

tslang/include/TypeScript/MLIRLogic/MLIRTypeHelper.h

@@ -3628,6 +3628,57 @@ class MLIRTypeHelper
         return mlir_ts::TupleType::get(context, resultFields);
     }
 
+    // A function's inferred return type is the union of what it returns: `return "a"` with `return undefined` is
+    // string | undefined, as in TypeScript. mergeType absorbs undefined and null into the other type, as an
+    // initializer would; here they are kept, and so per field of a returned tuple ({value, done} of a generator).
+    mlir::Type mergeReturnType(mlir::Location location, mlir::Type existType, mlir::Type currentType, bool& merged)
+    {
+        merged = false;
+        if (existType == currentType)
+        {
+            merged = true;
+            return existType;
+        }
+
+        auto isUndefinedOrNull = [](mlir::Type type) {
+            return isa<mlir_ts::UndefinedType>(type) || isa<mlir_ts::NullType>(type);
+        };
+
+        if (isUndefinedOrNull(existType) != isUndefinedOrNull(currentType))
+        {
+            merged = true;
+            return getUnionType(location, existType, currentType);
+        }
+
+        auto existTuple = dyn_cast<mlir_ts::TupleType>(existType);
+        auto currentTuple = dyn_cast<mlir_ts::TupleType>(currentType);
+        if (existTuple && currentTuple && existTuple.size() == currentTuple.size())
+        {
+            llvm::SmallVector<mlir_ts::FieldInfo> resultFields;
+            for (auto [existField, currentField] : llvm::zip(existTuple.getFields(), currentTuple.getFields()))
+            {
+                if (existField.id != currentField.id)
+                {
+                    return mergeType(location, existType, currentType, merged);
+                }
+
+                auto fieldMerged = false;
+                auto fieldType = mergeReturnType(location, existField.type, currentField.type, fieldMerged);
+                if (!fieldType)
+                {
+                    return mergeType(location, existType, currentType, merged);
+                }
+
+                resultFields.push_back({existField.id, fieldType, existField.isConditional, existField.accessLevel});
+            }
+
+            merged = true;
+            return mlir_ts::TupleType::get(context, resultFields);
+        }
+
+        return mergeType(location, existType, currentType, merged);
+    }
+
     mlir::Type mergeType(mlir::Location location, mlir::Type existType, mlir::Type currentType, bool& merged)
     {
         merged = false;

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -453,6 +453,8 @@ namespace mlirgen
             SmallVector<mlir::Block *> cleanUpsList;
             SmallVector<mlir::Operation *> cleanUpOpsList;
             PassResult passResultData;
+            passResultData.iteratorResult =
+                (functionLikeDeclarationBaseAST->internalFlags & InternalFlags::GeneratorNext) == InternalFlags::GeneratorNext;
             int discoverState = 1;
 
             GenContext genContextWithPassResult{};
@@ -525,6 +527,13 @@ namespace mlirgen
 
                 funcProto->setDiscovered(true);
                 auto discoveredType = passResult->functionReturnType;
+
+                // a bare `return;` beside value returns: T | undefined
+                if (passResult->hasBareReturn && !hasExplicitReturnType && discoveredType && !mth.isNoneType(discoveredType)
+                    && !isa<mlir_ts::VoidType>(discoveredType))
+                {
+                    discoveredType = getUnionType(loc(functionLikeDeclarationBaseAST), discoveredType, getUndefinedType());
+                }
                 // a "none" discoveredType means inference didn't converge (see above) -
                 // never let that silently overwrite an already-known (explicit or
                 // previously-cached) return type.
@@ -723,7 +732,7 @@ namespace mlirgen
         auto nextMethodDecl =
             nf.createMethodDeclaration(undefined, undefined, nf.createIdentifier(S(ITERATOR_NEXT)), undefined,
                                        undefined, undefined, undefined, nextBody);
-        nextMethodDecl->internalFlags |= InternalFlags::VarsInObjectContext;
+        nextMethodDecl->internalFlags |= InternalFlags::VarsInObjectContext | InternalFlags::GeneratorNext;
 
         // copy location info, to fix issue with names of anonymous functions
         nextMethodDecl->pos = functionLikeDeclarationBaseAST->pos;

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -3133,23 +3133,19 @@ class MLIRGenImpl
                 return mlir::success();
             }
 
-            // TODO: undefined & null should be processed as union type
-            auto undefType = getUndefinedType();
-            auto nullType = getNullType();
-
-            // filter out types, such as: undefined, objects with undefined values etc
-            if (type == undefType || type == nullType)
+            // a generator's next() returns {value: T, done} of Iterator<T>: the undefined value of a finished
+            // generator is not part of T, it reads as T's zero (00generator_done_value_zero.ts)
+            auto iteratorResult = genContext.passResult->iteratorResult;
+            if (iteratorResult && (isa<mlir_ts::UndefinedType>(type) || isa<mlir_ts::NullType>(type)))
             {
                 return mlir::failure();
             }
 
-            // if (mth.hasUndefines(type))
-            // {
-            //     return mlir::failure();
-            // }
-
+            // otherwise `return undefined` / `return null` are part of the union too (mergeReturnType)
             auto merged = false;
-            auto resultReturnType = mth.mergeType(location, genContext.passResult->functionReturnType, type, merged);            
+            auto resultReturnType = iteratorResult
+                ? mth.mergeType(location, genContext.passResult->functionReturnType, type, merged)
+                : mth.mergeReturnType(location, genContext.passResult->functionReturnType, type, merged);
 
             LLVM_DEBUG(dbgs() << "\n!! return type: " << resultReturnType << "");
 

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -512,6 +512,25 @@ namespace mlirgen
             return mlirGenReturnValue(location, expressionValue, false, genContext);
         }
 
+        // `return;` in a function that returns a value returns undefined: its type is T | undefined
+        // (discovery adds undefined when it saw a bare return beside value ones)
+        if (genContext.passResult)
+        {
+            genContext.passResult->hasBareReturn = true;
+        }
+
+        if (auto returnType = getExplicitReturnTypeOfCurrentFunction(genContext))
+        {
+            if (!isa<mlir_ts::VoidType>(returnType))
+            {
+                NodeFactory nf(NodeFactoryFlags::None);
+                auto undefinedReturn = nf.createReturnStatement(nf.createIdentifier(S(UNDEFINED_NAME)));
+                undefinedReturn->pos = returnStatementAST->pos;
+                undefinedReturn->_end = returnStatementAST->_end;
+                return mlirGen(undefinedReturn, genContext);
+            }
+        }
+
         EXIT_IF_FAILED(mlirGenScopeExit(location, DisposeDepth::FullStack, {}, &genContext));
 
         builder.create<mlir_ts::ReturnOp>(location);

tslang/test/tester/CMakeLists.txt

@@ -395,6 +395,7 @@ tslang_add_test(NAME test-compile-00-length-with-negative-merge COMMAND test-run
 tslang_add_test(NAME test-compile-00-length-signed-arithmetic COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_signed_arithmetic.ts")
 tslang_add_test(NAME test-compile-00-optional-fields COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00optional_fields.ts")
 tslang_add_test(NAME test-compile-00-generator-return COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_return.ts")
+tslang_add_test(NAME test-compile-00-return-inference-undefined COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00return_inference_undefined.ts")
 tslang_add_test(NAME test-compile-00-intersection-type COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type.ts")
 tslang_add_test(NAME test-compile-00-intersection-type-generic COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type_generic.ts")
 tslang_add_test(NAME test-compile-00-length COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00length.ts")
@@ -873,6 +874,7 @@ tslang_add_test(NAME test-jit-00-length-with-negative-merge COMMAND test-runner
 tslang_add_test(NAME test-jit-00-length-signed-arithmetic COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_signed_arithmetic.ts")
 tslang_add_test(NAME test-jit-00-optional-fields COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00optional_fields.ts")
 tslang_add_test(NAME test-jit-00-generator-return COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_return.ts")
+tslang_add_test(NAME test-jit-00-return-inference-undefined COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00return_inference_undefined.ts")
 tslang_add_test(NAME test-jit-00-union-global-init COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00union_global_init.ts")
 tslang_add_test(NAME test-jit-00-intersection-type COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type.ts")
 tslang_add_test(NAME test-jit-00-intersection-type-generic COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type_generic.ts")

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

@@ -0,0 +1,76 @@
+// An inferred return type is the union of what the function returns: `undefined` or `null` beside a value
+// gives `T | undefined` / `T | null` in either order, and a bare `return;` beside value returns is undefined
+function laterUndefined(c: boolean) {
+    if (c) return "a";
+    return undefined;
+}
+
+function firstUndefined(c: boolean) {
+    if (c) return undefined;
+    return "a";
+}
+
+function bareReturn(c: boolean) {
+    if (c) return "a";
+    return;
+}
+
+function integer(c: boolean) {
+    if (c) return 1;
+    return undefined;
+}
+
+function object(c: boolean) {
+    if (c) return { v: "x" };
+    return undefined;
+}
+
+function orNull(c: boolean) {
+    if (c) return "a";
+    return null;
+}
+
+function onlyBare(c: boolean) {
+    if (c) return;
+    sideEffect = 1;
+}
+
+let sideEffect = 0;
+
+function declared(c: boolean): string | undefined {
+    if (c) return;
+    return "x";
+}
+
+class WithBareReturn {
+    v = 1;
+    constructor(skip: boolean) {
+        if (skip) return;
+        this.v = 2;
+    }
+}
+
+function main() {
+    assert(laterUndefined(false) === undefined, "later undefined");
+    assert(laterUndefined(true) == "a", "later undefined: value");
+    assert(firstUndefined(true) === undefined, "first undefined");
+    assert(firstUndefined(false) == "a", "first undefined: value");
+    assert(bareReturn(false) === undefined, "bare return");
+    assert(bareReturn(true) == "a", "bare return: value");
+    assert(integer(false) === undefined, "integer");
+    assert(integer(true) == 1, "integer: value");
+    assert(object(false) === undefined, "object");
+    assert(object(true)?.v == "x", "object: value");
+    assert(orNull(false) === null, "null");
+    assert(orNull(true) == "a", "null: value");
+
+    onlyBare(false);
+    assert(sideEffect == 1, "void function with a bare return");
+
+    assert(declared(true) === undefined, "declared, bare return");
+    assert(declared(false) == "x", "declared, value");
+
+    assert(new WithBareReturn(true).v == 1 && new WithBareReturn(false).v == 2, "constructor bare return");
+
+    print("done.");
+}

tslang/ts-new-parser/enums.h

@@ -562,7 +562,9 @@ enum class InternalFlags : number
     // whose mutable `step` state must be shared across aliases)
     BoxAsObject = 1 << 11,
     // a `return` of a generator's own body: it finishes the generator (mlirGen(ReturnStatement))
-    GeneratorReturn = 1 << 12
+    GeneratorReturn = 1 << 12,
+    // the next() a generator is rewritten into: it returns Iterator<T>'s {value: T, done} (processReturnType)
+    GeneratorNext = 1 << 13
 };
 
 ENUM_OPS(InternalFlags)