A union of literals takes part in arithmetic as its literals' type

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -6407,6 +6407,25 @@ class MLIRGenImpl
         return mlir::success();
     }
 
+    // The type of a union whose members are all literals of one type (`0.5 | 1` is a number,
+    // `"a" | "b"` a string), or null for any other type.
+    mlir::Type getLiteralUnionBaseType(mlir::Location location, mlir::Type type)
+    {
+        auto unionType = dyn_cast<mlir_ts::UnionType>(type);
+        if (!unionType || !llvm::all_of(unionType.getTypes(), [](mlir::Type member) { return isa<mlir_ts::LiteralType>(member); }))
+        {
+            return mlir::Type();
+        }
+
+        mlir::Type baseType;
+        if (mth.isUnionTypeNeedsTag(location, unionType, baseType) || !baseType || baseType == type)
+        {
+            return mlir::Type();
+        }
+
+        return baseType;
+    }
+
     mlir::LogicalResult unwrapForBinaryOp(mlir::Location location, SyntaxKind opCode, mlir::Value &leftExpressionValue,
                                           mlir::Value &rightExpressionValue, const GenContext &genContext)
     {
@@ -6426,6 +6445,19 @@ class MLIRGenImpl
         {
             CAST(rightExpressionValue, location, rightType.getElementType(), rightExpressionValue, genContext);
         }
+
+        // A union of literals, as `c ? 0.5 : 1.0` is (`0.5 | 1`), is its literals' type too. Left a
+        // union, it was cast to the other operand's type: `int * (c ? 0.5 : 1.0)` multiplied s32s,
+        // 0.5 became 0, and the product was 0.
+        if (auto leftType = getLiteralUnionBaseType(location, leftExpressionValue.getType()))
+        {
+            CAST(leftExpressionValue, location, leftType, leftExpressionValue, genContext);
+        }
+
+        if (auto rightType = getLiteralUnionBaseType(location, rightExpressionValue.getType()))
+        {
+            CAST(rightExpressionValue, location, rightType, rightExpressionValue, genContext);
+        }
         // end of hack
 
         if (leftExpressionValue.getType() != rightExpressionValue.getType())

tslang/test/tester/CMakeLists.txt

@@ -272,6 +272,7 @@ tslang_add_test(NAME test-compile-00-numbers COMMAND test-runner "${PROJECT_SOUR
 tslang_add_test(NAME test-compile-00-equals COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00equals.ts")
 tslang_add_test(NAME test-compile-00-bool-arith-ops COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_arith_ops.ts")
 tslang_add_test(NAME test-compile-00-int64-shift-bitwise COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00int64_shift_bitwise.ts")
+tslang_add_test(NAME test-compile-00-literal-union-arith COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00literal_union_arith.ts")
 tslang_add_test(NAME test-compile-00-bool-const-array COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_const_array.ts")
 tslang_add_test(NAME test-compile-00-mixed-type-ops COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00mixed_type_ops.ts")
 tslang_add_test(NAME test-compile-00-generator-manual-next COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next.ts")
@@ -805,6 +806,7 @@ tslang_add_test(NAME test-jit-00-numbers COMMAND test-runner -jit "${PROJECT_SOU
 tslang_add_test(NAME test-jit-00-equals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00equals.ts")
 tslang_add_test(NAME test-jit-00-bool-arith-ops COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_arith_ops.ts")
 tslang_add_test(NAME test-jit-00-int64-shift-bitwise COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00int64_shift_bitwise.ts")
+tslang_add_test(NAME test-jit-00-literal-union-arith COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00literal_union_arith.ts")
 tslang_add_test(NAME test-jit-00-bool-const-array COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_const_array.ts")
 tslang_add_test(NAME test-jit-00-mixed-type-ops COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00mixed_type_ops.ts")
 tslang_add_test(NAME test-jit-00-generator-manual-next COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next.ts")
@@ -1794,6 +1796,7 @@ set(TSLANG_CORPUS
     00async_result_types.ts
     00bool_arith_ops.ts
     00int64_shift_bitwise.ts
+    00literal_union_arith.ts
     00bool_const_array.ts
     00break_continue_scope_exit.ts
     00break_continue.ts

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

@@ -0,0 +1,49 @@
+// regression test: a union of literals, as a ternary of two literals is, takes part in arithmetic
+// as its literals' type.
+//
+// bug found and fixed: `c ? 0.5 : 1.0` is the union `0.5 | 1`, and a binary operator cast that
+// union to the other operand's type: `int * (c ? 0.5 : 1.0)` multiplied two s32s, 0.5 became 0,
+// and the product was 0.
+
+function scale(split: boolean, width: int): number {
+    return width * (split ? 0.5 : 1.0);
+}
+
+function main() {
+    let split = true;
+    let held = false;
+    const width = 100;
+
+    assert(width * (split ? 0.5 : 1.0) == 50);
+    assert((split ? 0.5 : 0.25) * width == 50);
+    assert(width * (held ? 0.5 : 0.25) == 25);
+    assert(100 * (split ? 0.5 : 1.0) == 50);
+    assert(width + (split ? 0.5 : 1.0) == 100.5);
+    assert(width - (split ? 0.5 : 1.0) == 99.5);
+    assert(scale(true, 100) == 50);
+    assert(scale(false, 100) == 100);
+
+    const w: number = 100;
+    assert(w * (split ? 0.5 : 1.0) == 50);
+
+    // through a variable
+    const half = split ? 0.5 : 1.0;
+    assert(width * half == 50);
+
+    const left = split ? 448.0 : 96.5;
+    assert(left + 10.5 == 458.5);
+    let m = left * 2;
+    m = m / 3;
+    assert(m > 298.6 && m < 298.7);
+
+    assert((held ? 1.0 : 0.0) - (split ? 1.0 : 0.0) == -1);
+
+    // comparisons and strings still work
+    assert(half == 0.5);
+    assert(half != 1.0);
+    const s = held ? "a" : "b";
+    assert(s + "!" == "b!");
+    assert(s == "b");
+
+    print("done.");
+}