instanceof asks a union or an optional that may hold a class at run time

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -5140,6 +5140,89 @@ class MLIRGenImpl
             });
     }
 
+    // A union or an optional with a class among what it may hold.
+    bool mayHoldClass(mlir::Type type)
+    {
+        if (auto optionalType = dyn_cast<mlir_ts::OptionalType>(type))
+        {
+            return isa<mlir_ts::ClassType>(optionalType.getElementType()) || mayHoldClass(optionalType.getElementType());
+        }
+
+        if (auto unionType = dyn_cast<mlir_ts::UnionType>(type))
+        {
+            return llvm::any_of(unionType.getTypes(), [](mlir::Type member) { return isa<mlir_ts::ClassType>(member); });
+        }
+
+        return false;
+    }
+
+    // `u instanceof C` for a union or an optional that may hold a class. Whether it holds one is
+    // known only at run time, and which class only to the instance's vtable: every class member of
+    // a tagged union has the tag "class". Anything else it may hold - null, undefined, a number - is
+    // an instance of nothing.
+    ValueOrLogicalResult mlirGenInstanceOfMaybeClass(mlir::Location location, mlir::Value value, mlir::Value classRefVal,
+                                                     const GenContext &genContext)
+    {
+        auto notInstance = [&](mlir::Type) {
+            return ValueOrLogicalResult(builder.create<mlir_ts::ConstantOp>(location, getBooleanType(), builder.getBoolAttr(false)));
+        };
+
+        // the instance, asked through its own vtable
+        auto askInstance = [&](mlir::Value instance) -> ValueOrLogicalResult {
+            auto thisPtr = cast(location, getOpaqueType(), instance, genContext);
+            if (thisPtr.failed_or_no_value())
+            {
+                return mlir::failure();
+            }
+
+            return mlirGenInstanceOfOpaque(location, V(thisPtr), classRefVal, genContext);
+        };
+
+        MLIRCodeLogicHelper mclh(builder, location, compileOptions);
+        if (auto optionalType = dyn_cast<mlir_ts::OptionalType>(value.getType()))
+        {
+            auto hasValue = builder.create<mlir_ts::HasValueOp>(location, getBooleanType(), value);
+            return mclh.conditionalValue(
+                hasValue,
+                [&]() -> ValueOrLogicalResult {
+                    auto inner = builder.create<mlir_ts::ValueOp>(location, optionalType.getElementType(), value);
+                    return isa<mlir_ts::ClassType>(inner.getType())
+                        ? askInstance(inner.getResult())
+                        : mlirGenInstanceOfMaybeClass(location, inner, classRefVal, genContext);
+                },
+                notInstance);
+        }
+
+        auto unionType = mlir::cast<mlir_ts::UnionType>(value.getType());
+        auto classMember = *llvm::find_if(unionType.getTypes(), [](mlir::Type member) { return isa<mlir_ts::ClassType>(member); });
+
+        // `C | null` needs no tag: it is the class's pointer, null when it holds null
+        mlir::Type baseType;
+        if (!mth.isUnionTypeNeedsTag(location, unionType, baseType))
+        {
+            auto isSet = cast(location, getBooleanType(), value, genContext);
+            if (isSet.failed_or_no_value())
+            {
+                return mlir::failure();
+            }
+
+            return mclh.conditionalValue(
+                V(isSet),
+                [&]() { return askInstance(builder.create<mlir_ts::GetValueFromUnionOp>(location, classMember, value).getResult()); },
+                notInstance);
+        }
+
+        // a tagged union: a class is in it when its tag says so, and the payload is its pointer
+        auto typeOfValue = builder.create<mlir_ts::GetTypeInfoFromUnionOp>(location, getStringType(), value);
+        auto classStrConst = builder.create<mlir_ts::ConstantOp>(location, getStringType(), builder.getStringAttr("class"));
+        auto isClass = builder.create<mlir_ts::StringCompareOp>(location, getBooleanType(), typeOfValue, classStrConst,
+                                                                builder.getI32IntegerAttr((int)SyntaxKind::EqualsEqualsToken));
+        return mclh.conditionalValue(
+            isClass,
+            [&]() { return askInstance(builder.create<mlir_ts::GetValueFromUnionOp>(location, classMember, value).getResult()); },
+            notInstance);
+    }
+
     // Asks `thisPtrValue` whether it is a `classRefVal` through slot 0 of `vtablePtr`: the class's
     // `.instanceOf` for a class vtable, and the implementer's for an interface vtable.
     ValueOrLogicalResult mlirGenInstanceOfThroughVTable(mlir::Location location, mlir::Value vtablePtr, mlir::Value thisPtrValue,
@@ -5290,6 +5373,16 @@ class MLIRGenImpl
             {
                 return mlirGenInstanceOfInterface(location, resultLeftValue, resultRightValue, genContext);
             }
+
+            // a union or an optional that may hold a class: asked at run time, not by its type
+            auto refType = dyn_cast<mlir_ts::RefType>(resultLeftValue.getType());
+            if (mayHoldClass(refType ? refType.getElementType() : resultLeftValue.getType()))
+            {
+                auto leftValue = refType
+                    ? builder.create<mlir_ts::LoadOp>(location, refType.getElementType(), resultLeftValue).getResult()
+                    : resultLeftValue;
+                return mlirGenInstanceOfMaybeClass(location, leftValue, resultRightValue, genContext);
+            }
         }
 #endif
 

tslang/test/tester/CMakeLists.txt

@@ -468,6 +468,7 @@ tslang_add_test(NAME test-compile-00-in COMMAND test-runner "${PROJECT_SOURCE_DI
 tslang_add_test(NAME test-compile-00-in-method-names COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00in_method_names.ts")
 tslang_add_test(NAME test-compile-00-class-structural-extends COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_structural_extends.ts")
 tslang_add_test(NAME test-compile-00-instanceof COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
+tslang_add_test(NAME test-compile-00-instanceof-union COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-compile-00-class COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class.ts")
 tslang_add_test(NAME test-compile-00-class-new COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_new.ts")
 tslang_add_test(NAME test-compile-00-capture-in-new-arguments COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00capture_in_new_arguments.ts")
@@ -950,6 +951,7 @@ tslang_add_test(NAME test-jit-00-in COMMAND test-runner -jit "${PROJECT_SOURCE_D
 tslang_add_test(NAME test-jit-00-in-method-names COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00in_method_names.ts")
 tslang_add_test(NAME test-jit-00-class-structural-extends COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_structural_extends.ts")
 tslang_add_test(NAME test-jit-00-instanceof COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
+tslang_add_test(NAME test-jit-00-instanceof-union COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-jit-00-class COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class.ts")
 tslang_add_test(NAME test-jit-00-class-new COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_new.ts")
 tslang_add_test(NAME test-jit-00-capture-in-new-arguments COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00capture_in_new_arguments.ts")
@@ -1764,6 +1766,7 @@ set(TSLANG_CORPUS
     00infer.ts
     00instanceof.ts
     00instanceof_namespace_narrowing.ts
+    00instanceof_union.ts
     00interface_captures.ts
     00interface_conjunction.ts
     00interface_function_typed_field.ts

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

@@ -0,0 +1,58 @@
+// `x instanceof C` where `x` is a union or an optional that may hold a class: whether it holds one
+// is known only at run time, and which class only to the instance's vtable. It used to fold to the
+// compile-time question "is the union type C", which is false, so every one of these was false.
+class C {
+    constructor(public x: number) {}
+}
+
+class Sub extends C {
+    constructor(x: number) {
+        super(x);
+    }
+}
+
+class D {
+    y = 1;
+}
+
+function isC(v: C | D) {
+    return v instanceof C;
+}
+
+function main() {
+    // a nullable class: a pointer that may be null
+    let e: C | null = new C(1);
+    assert(e instanceof C, "C | null holding a C");
+    if (e) assert(e instanceof C, "narrowed C | null");
+    e = null;
+    assert(!(e instanceof C), "C | null holding null");
+
+    // a subclass held as its base, nullable
+    let s: C | null = new Sub(2);
+    assert(s instanceof C && s instanceof Sub, "C | null holding a Sub");
+
+    // a union of classes: tagged, every class member's tag is "class"
+    let b: C | D = new C(3);
+    assert(b instanceof C && !(b instanceof D), "C | D holding a C");
+    b = new D();
+    assert(b instanceof D && !(b instanceof C), "C | D holding a D");
+    assert(isC(new C(4)) && !isC(new D()), "C | D parameter");
+
+    // a class and a number
+    let f: C | number = new C(5);
+    assert(f instanceof C, "C | number holding a C");
+    f = 6;
+    assert(!(f instanceof C), "C | number holding a number");
+
+    // an optional class
+    let o: C | undefined = new C(7);
+    assert(o instanceof C, "C | undefined holding a C");
+    o = undefined;
+    assert(!(o instanceof C), "C | undefined holding undefined");
+
+    // a plain class still asks its own vtable
+    const c = new C(8);
+    assert(c instanceof C && !(c instanceof Sub), "plain C");
+
+    print("done.");
+}