Improve ValueType.simplified()

nicklockwood · nicklockwood · commit 6f58795737a9 · 2025-10-12T17:54:06.000+01:00
diff --git a/ShapeScript/Types.swift b/ShapeScript/Types.swift
@@ -125,7 +125,14 @@ extension ValueType {
     func simplified() -> ValueType {
         switch self {
         case let .union(types):
-            let types = types.sorted()
+            let types = types.flatMap {
+                switch $0.simplified() {
+                case let .union(types):
+                    return types
+                default:
+                    return [$0]
+                }
+            }.sorted()
             guard var result = types.first.map({ [$0] }) else {
                 return self
             }
@@ -137,9 +144,15 @@ extension ValueType {
                 result.append(type)
             }
             return result.count == 1 ? result[0] : .union(Set(result))
+        case let .list(type):
+            return .list(type.simplified())
+        case let .tuple(types):
+            return .tuple(types.map { $0.simplified() })
+        case let .object(fields):
+            return .object(fields.mapValues { $0.simplified() })
         case .any, .color, .texture, .material, .boolean, .font, .number, .radians,
              .halfturns, .vector, .size, .rotation, .string, .text, .path, .mesh,
-             .polygon, .point, .range, .partialRange, .bounds, .tuple, .list, .object:
+             .polygon, .point, .range, .partialRange, .bounds:
             return self
         }
     }
diff --git a/ShapeScriptTests/TypesystemTests.swift b/ShapeScriptTests/TypesystemTests.swift
@@ -506,7 +506,7 @@ final class TypesystemTests: XCTestCase {
             }
         }
         """)
-        XCTAssertEqual(type.parameterType, .tuple([.union([.numberOrVector, .list(.any)]), .number]))
+        XCTAssertEqual(type.parameterType, .tuple([.union([.number, .list(.any)]), .number]))
         XCTAssertEqual(type.returnType, .void)
     }
 
@@ -521,9 +521,7 @@ final class TypesystemTests: XCTestCase {
             }
         }
         """)
-        XCTAssertEqual(type.parameterType, .tuple([
-            .union([.number, .list(.number), .texture]),
-        ]))
+        XCTAssertEqual(type.parameterType, .tuple([.union([.numberOrVector, .texture]).simplified()]))
         XCTAssertEqual(type.returnType, .void)
     }
 
@@ -683,7 +681,7 @@ final class TypesystemTests: XCTestCase {
         }
         """
         let type = try blockType(in: input)
-        XCTAssertEqual(type.childTypes, .union([.mesh, .number, .list(.number)]))
+        XCTAssertEqual(type.childTypes, .union([.mesh, .numberOrVector]).simplified())
         XCTAssertEqual(type.returnType, .mesh)
 
         // No arguments
@@ -779,6 +777,31 @@ final class TypesystemTests: XCTestCase {
         XCTAssertEqual(type, .any)
     }
 
+    func testUnionSubtypeSimplification() {
+        XCTAssertEqual(ValueType.union([.number, .any]).simplified(), .any)
+        XCTAssertEqual(ValueType.union([.any, .boolean]).simplified(), .any)
+        XCTAssertEqual(ValueType.union([.list(.any), .list(.number)]).simplified(), .list(.any))
+    }
+
+    func testNestedUnionSimplification() {
+        XCTAssertEqual(
+            ValueType.union([.boolean, .union([.number, .string])]).simplified(),
+            .union([.boolean, .number, .string])
+        )
+        XCTAssertEqual(
+            ValueType.list(.union([.number, .any])).simplified(),
+            .list(.any)
+        )
+        XCTAssertEqual(
+            ValueType.tuple([.union([.number, .any]), .union([.any, .boolean])]).simplified(),
+            .tuple([.any, .any])
+        )
+        XCTAssertEqual(
+            ValueType.object(["foo": .union([.number, .any]), "bar": .union([.any, .boolean])]).simplified(),
+            .object(["foo": .any, "bar": .any])
+        )
+    }
+
     // MARK: Type conversion
 
     func testCastNumberToNumberTuple() {
