Improve discriminated union implementation (#4556)

* WIP

* Simplify discriminated union logic

* Simplify disc union internals, improve perf

* Improve docs

* Typo

Author: colinhacks

packages/bench/discriminated-union.ts

@@ -1,4 +1,4 @@
-import { makeData, makeSchema, randomPick, randomString } from "./benchUtil.js";
+import { makeData, randomPick, randomString } from "./benchUtil.js";
 import { metabench } from "./metabench.js";
 
 
@@ -54,55 +54,54 @@ const z3DiscUnion = z3.discriminatedUnion("type", z3Union._def.options);
 
 function makeSchema(z: typeof z4){
    const z4fields = {
-  data1: z4.string(),
-  data2: z4.string(),
-  data3: z4.string(),
-  // data4: z4.string(),
-  // data5: z4.string(),
-  // data6: z4.string(),
-  // data7: z4.string(),
-  // data8: z4.string(),
-  // data9: z4.string(),
-  // data10: z4.string(),
-}
-const z4Union = z4.union([
-  z4.object({
-    type: z4.literal("a"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("b"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("c"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("d"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("e"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("f"),
-    ...z4fields
-  }),
-  z4.object({
-    type: z4.literal("g"),
-    ...z4fields
-  }),
-]);
-return z4Union;
+    data1: z.string(),
+    data2: z.string(),
+    data3: z.string(),
+    // data4: z.string(),
+    // data5: z.string(),
+    // data6: z.string(),
+    // data7: z.string(),
+    // data8: z.string(),
+    // data9: z.string(),
+    // data10: z.string(),
+  }
+  const z4Union = z.union([
+    z.object({
+      type: z.literal("a"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("b"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("c"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("d"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("e"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("f"),
+      ...z4fields
+    }),
+    z.object({
+      type: z.literal("g"),
+      ...z4fields
+    }),
+  ]);
+  return z4Union;
 
 }
 
 const z4Union = makeSchema(z4);
 const z4LibUnion  = makeSchema(z4lib as any);
-const z4LibDiscUnion  = z4lib.discriminatedUnion( z4LibUnion._def.options);
-
+const z4LibDiscUnion  = z4lib.discriminatedUnion("type", z4LibUnion._def.options as any);
 const z4DiscUnion = z4.discriminatedUnion("type", z4Union.def.options);
 
 const DATA = makeData(100, () => ({ 
@@ -137,7 +136,7 @@ console.dir(z4LibDiscUnion.parse(DATA[0]), {depth: null});
 
 
 const args=  {jitless: true}
-const bench = metabench("z.disriminatedUnion().parse", {
+const bench = metabench("z.discriminatedUnion().parse", {
   // z3() {
   //   for (const item of DATA) {
   //     z3Union.parse(item);

packages/bench/index.ts

@@ -6,7 +6,7 @@ async function run() {
   const files = process.argv[2].split(",").map((file) => import.meta.resolve(`./${file}`).replace("file://", ""));
 
   for (const file of files) {
-    await $`pnpm tsx --conditions @zod/source ${file}`;
+    await $`pnpm tsx ${file}`;
   }
 }
 

packages/bench/jit-union.ts

@@ -74,7 +74,7 @@ const DATA = makeData(100, () => ({
 const args=  {jitless: true}
 console.dir(z4Union.parse(DATA[0]), {depth: null});
 console.dir(z4Union.parse(DATA[0], args), {depth: null});
-const bench = metabench("z.disriminatedUnion().parse", {
+const bench = metabench("z.discriminatedUnion().parse", {
 
   v4_jit() {
     for (const item of DATA) {

packages/bench/package.json

@@ -6,7 +6,7 @@
     "arktype": "^2.1.19",
     "valibot": "^1.0.0",
     "zod": "workspace:*",
-    "zodnext": "npm:zod@next",
+    "zodnext": "npm:zod@^3.25.0",
     "zod3": "npm:zod@^3.23.7"
   },
   "scripts": {

packages/docs/content/v4/index.mdx

@@ -827,7 +827,7 @@ Replace `errorMap` with `error` (function syntax):
 
 ## Upgraded `z.discriminatedUnion()`
 
-Discriminated unions now support a number of schema types not previously supported, including unions, pipes, and nested objects:
+Discriminated unions now support a number of schema types not previously supported, including unions and pipes:
 
 ```ts
 const MyResult = z.discriminatedUnion("status", [
@@ -836,23 +836,22 @@ const MyResult = z.discriminatedUnion("status", [
   // union discriminator
   z.object({ status: z.union([z.literal("bbb"), z.literal("ccc")]) }),
   // pipe discriminator
-  z.object({ status: z.object({ value: z.literal("fail") }) }),
+  z.object({ status: z.literal("fail").transform(val => val.toUpperCase()) }),
 ]);
 ```
 
 Perhaps most importantly, discriminated unions now *compose*—you can use one discriminated union as a member of another.
 
 ```ts
 const BaseError = z.object({ status: z.literal("failed"), message: z.string() });
-const MyErrors = z.discriminatedUnion("code", [
-  BaseError.extend({ code: z.literal(400) }),
-  BaseError.extend({ code: z.literal(401) }),
-  BaseError.extend({ code: z.literal(500) })
-]);
 
 const MyResult = z.discriminatedUnion("status", [
   z.object({ status: z.literal("success"), data: z.string() }),
-  MyErrors
+  z.discriminatedUnion("code", [
+    BaseError.extend({ code: z.literal(400) }),
+    BaseError.extend({ code: z.literal(401) }),
+    BaseError.extend({ code: z.literal(500) })
+  ])
 ]);
 ```