@voidhash/mimic 0.0.1-alpha.6 → 0.0.1-alpha.7

package/src/primitives/Struct.ts

@@ -6,7 +6,7 @@ import * as ProxyEnvironment from "../ProxyEnvironment";
 import * as Transform from "../Transform";
 import type { Primitive, PrimitiveInternal, MaybeUndefined, AnyPrimitive, Validator, InferState, InferProxy, InferSnapshot } from "../Primitive";
 import { ValidationError } from "../Primitive";
-import { runValidators } from "./shared";
+import { runValidators, applyDefaults, StructSetInput } from "./shared";
 
 
 /**
@@ -25,6 +25,16 @@ export type InferStructSnapshot<TFields extends Record<string, AnyPrimitive>> =
   readonly [K in keyof TFields]: InferSnapshot<TFields[K]>;
 };
 
+/**
+ * Maps a schema definition to a partial update type.
+ * For nested structs, allows recursive partial updates.
+ */
+export type StructUpdateValue<TFields extends Record<string, AnyPrimitive>> = {
+  readonly [K in keyof TFields]?: TFields[K] extends StructPrimitive<infer F, any>
+    ? StructUpdateValue<F> | InferState<TFields[K]>
+    : InferState<TFields[K]>;
+};
+
 /**
  * Maps a schema definition to its proxy type.
  * Provides nested field access + get()/set()/toSnapshot() methods for the whole struct.
@@ -34,8 +44,10 @@ export type StructProxy<TFields extends Record<string, AnyPrimitive>, TDefined e
 } & {
   /** Gets the entire struct value */
   get(): MaybeUndefined<InferStructState<TFields>, TDefined>;
-  /** Sets the entire struct value */
-  set(value: InferStructState<TFields>): void;
+  /** Sets the entire struct value (only fields that are required without defaults must be provided) */
+  set(value: StructSetInput<TFields>): void;
+  /** Updates only the specified fields (partial update, handles nested structs recursively) */
+  update(value: StructUpdateValue<TFields>): void;
   /** Returns a readonly snapshot of the struct for rendering */
   toSnapshot(): MaybeUndefined<InferStructSnapshot<TFields>, TDefined>;
 };
@@ -47,12 +59,14 @@ interface StructPrimitiveSchema<TFields extends Record<string, AnyPrimitive>> {
   readonly validators: readonly Validator<InferStructState<TFields>>[];
 }
 
-export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefined extends boolean = false>
-  implements Primitive<InferStructState<TFields>, StructProxy<TFields, TDefined>>
+export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefined extends boolean = false, THasDefault extends boolean = false>
+  implements Primitive<InferStructState<TFields>, StructProxy<TFields, TDefined>, TDefined, THasDefault>
 {
   readonly _tag = "StructPrimitive" as const;
   readonly _State!: InferStructState<TFields>;
   readonly _Proxy!: StructProxy<TFields, TDefined>;
+  readonly _TDefined!: TDefined;
+  readonly _THasDefault!: THasDefault;
 
   private readonly _schema: StructPrimitiveSchema<TFields>;
 
@@ -70,7 +84,7 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
   }
 
   /** Mark this struct as required */
-  required(): StructPrimitive<TFields, true> {
+  required(): StructPrimitive<TFields, true, THasDefault> {
     return new StructPrimitive({
       ...this._schema,
       required: true,
@@ -78,10 +92,12 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
   }
 
   /** Set a default value for this struct */
-  default(defaultValue: InferStructState<TFields>): StructPrimitive<TFields, true> {
+  default(defaultValue: StructSetInput<TFields>): StructPrimitive<TFields, true, true> {
+    // Apply defaults to the provided value
+    const merged = applyDefaults(this as AnyPrimitive, defaultValue as Partial<InferStructState<TFields>>) as InferStructState<TFields>;
     return new StructPrimitive({
       ...this._schema,
-      defaultValue,
+      defaultValue: merged,
     });
   }
 
@@ -91,7 +107,7 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
   }
 
   /** Add a custom validation rule (useful for cross-field validation) */
-  refine(fn: (value: InferStructState<TFields>) => boolean, message: string): StructPrimitive<TFields, TDefined> {
+  refine(fn: (value: InferStructState<TFields>) => boolean, message: string): StructPrimitive<TFields, TDefined, THasDefault> {
     return new StructPrimitive({
       ...this._schema,
       validators: [...this._schema.validators, { validate: fn, message }],
@@ -130,17 +146,47 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
         return snapshot as InferStructSnapshot<TFields>;
       };
 
-      // Create the base object with get/set/toSnapshot methods
+      // Create the base object with get/set/update/toSnapshot methods
       const base = {
         get: (): MaybeUndefined<InferStructState<TFields>, TDefined> => {
           const state = env.getState(operationPath) as InferStructState<TFields> | undefined;
           return (state ?? defaultValue) as MaybeUndefined<InferStructState<TFields>, TDefined>;
         },
-        set: (value: InferStructState<TFields>) => {
+        set: (value: StructSetInput<TFields>) => {
+          // Apply defaults for missing fields
+          const merged = applyDefaults(this as AnyPrimitive, value as Partial<InferStructState<TFields>>) as InferStructState<TFields>;
           env.addOperation(
-            Operation.fromDefinition(operationPath, this._opDefinitions.set, value)
+            Operation.fromDefinition(operationPath, this._opDefinitions.set, merged)
           );
         },
+        update: (value: StructUpdateValue<TFields>) => {
+          for (const key in value) {
+            if (Object.prototype.hasOwnProperty.call(value, key)) {
+              const fieldValue = value[key as keyof TFields];
+              if (fieldValue === undefined) continue; // Skip undefined values
+
+              const fieldPrimitive = fields[key as keyof TFields];
+              if (!fieldPrimitive) continue; // Skip unknown fields
+
+              const fieldPath = operationPath.append(key);
+              const fieldProxy = fieldPrimitive._internal.createProxy(env, fieldPath);
+
+              // Check if this is a nested struct and value is a plain object (partial update)
+              if (
+                fieldPrimitive._tag === "StructPrimitive" &&
+                typeof fieldValue === "object" &&
+                fieldValue !== null &&
+                !Array.isArray(fieldValue)
+              ) {
+                // Recursively update nested struct
+                (fieldProxy as { update: (v: unknown) => void }).update(fieldValue);
+              } else {
+                // Set the field value directly
+                (fieldProxy as { set: (v: unknown) => void }).set(fieldValue);
+              }
+            }
+          }
+        },
         toSnapshot: (): MaybeUndefined<InferStructSnapshot<TFields>, TDefined> => {
           const snapshot = buildSnapshot();
           return snapshot as MaybeUndefined<InferStructSnapshot<TFields>, TDefined>;
@@ -150,13 +196,16 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
       // Use a JavaScript Proxy to intercept field access
       return new globalThis.Proxy(base as StructProxy<TFields, TDefined>, {
         get: (target, prop, _receiver) => {
-          // Return base methods (get, set, toSnapshot)
+          // Return base methods (get, set, update, toSnapshot)
           if (prop === "get") {
             return target.get;
           }
           if (prop === "set") {
             return target.set;
           }
+          if (prop === "update") {
+            return target.update;
+          }
           if (prop === "toSnapshot") {
             return target.toSnapshot;
           }
@@ -176,7 +225,7 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
           return undefined;
         },
         has: (_target, prop) => {
-          if (prop === "get" || prop === "set" || prop === "toSnapshot") return true;
+          if (prop === "get" || prop === "set" || prop === "update" || prop === "toSnapshot") return true;
           if (typeof prop === "string" && prop in fields) return true;
           return false;
         },
@@ -343,6 +392,6 @@ export class StructPrimitive<TFields extends Record<string, AnyPrimitive>, TDefi
 /** Creates a new StructPrimitive with the given fields */
 export const Struct = <TFields extends Record<string, AnyPrimitive>>(
   fields: TFields
-): StructPrimitive<TFields, false> =>
+): StructPrimitive<TFields, false, false> =>
   new StructPrimitive({ required: false, defaultValue: undefined, fields, validators: [] });
 

package/src/primitives/Tree.ts

@@ -5,11 +5,11 @@ import * as OperationPath from "../OperationPath";
 import * as ProxyEnvironment from "../ProxyEnvironment";
 import * as Transform from "../Transform";
 import * as FractionalIndex from "../FractionalIndex";
-import type { Primitive, PrimitiveInternal, Validator, InferProxy } from "./shared";
-import { ValidationError } from "./shared";
+import type { Primitive, PrimitiveInternal, Validator, InferProxy, AnyPrimitive, StructSetInput } from "./shared";
+import { ValidationError, applyDefaults } from "./shared";
 import { runValidators } from "./shared";
 import type { AnyTreeNodePrimitive, InferTreeNodeType, InferTreeNodeDataState, InferTreeNodeChildren } from "./TreeNode";
-import { InferStructState } from "./Struct";
+import { InferStructState, StructUpdateValue, StructPrimitive } from "./Struct";
 
 
 /**
@@ -105,6 +105,23 @@ export type TreeNodeSnapshot<TNode extends AnyTreeNodePrimitive> = {
 export type InferTreeSnapshot<T extends TreePrimitive<any>> =
   T extends TreePrimitive<infer TRoot> ? TreeNodeSnapshot<TRoot> : never;
 
+/**
+ * Helper type to infer the update value type from a TreeNode's data
+ */
+export type TreeNodeUpdateValue<TNode extends AnyTreeNodePrimitive> =
+  TNode["data"] extends StructPrimitive<infer TFields, any, any>
+    ? StructUpdateValue<TFields>
+    : never;
+
+/**
+ * Helper type to infer the input type for node data (with defaults applied).
+ * Uses StructSetInput for struct data, making fields with defaults optional.
+ */
+export type TreeNodeDataSetInput<TNode extends AnyTreeNodePrimitive> =
+  TNode["data"] extends StructPrimitive<infer TFields, any, any>
+    ? StructSetInput<TFields>
+    : InferTreeNodeDataState<TNode>;
+
 /**
  * Typed proxy for a specific node type - provides type-safe data access
  */
@@ -117,6 +134,8 @@ export interface TypedNodeProxy<TNode extends AnyTreeNodePrimitive> {
   readonly data: InferProxy<TNode["data"]>;
   /** Get the raw node state */
   get(): TypedTreeNodeState<TNode>;
+  /** Updates only the specified data fields (partial update, handles nested structs recursively) */
+  update(value: TreeNodeUpdateValue<TNode>): void;
 }
 
 /**
@@ -158,40 +177,40 @@ export interface TreeProxy<TRoot extends AnyTreeNodePrimitive> {
   /** Gets a node proxy by ID with type narrowing capabilities */
   node(id: string): TreeNodeProxyBase<TRoot> | undefined;
   
-  /** Insert a new node as the first child */
+  /** Insert a new node as the first child (applies defaults for node data) */
   insertFirst<TNode extends AnyTreeNodePrimitive>(
     parentId: string | null,
     nodeType: TNode,
-    data: InferTreeNodeDataState<TNode>
+    data: TreeNodeDataSetInput<TNode>
   ): string;
   
-  /** Insert a new node as the last child */
+  /** Insert a new node as the last child (applies defaults for node data) */
   insertLast<TNode extends AnyTreeNodePrimitive>(
     parentId: string | null,
     nodeType: TNode,
-    data: InferTreeNodeDataState<TNode>
+    data: TreeNodeDataSetInput<TNode>
   ): string;
   
-  /** Insert a new node at a specific index among siblings */
+  /** Insert a new node at a specific index among siblings (applies defaults for node data) */
   insertAt<TNode extends AnyTreeNodePrimitive>(
     parentId: string | null,
     index: number,
     nodeType: TNode,
-    data: InferTreeNodeDataState<TNode>
+    data: TreeNodeDataSetInput<TNode>
   ): string;
   
-  /** Insert a new node after a sibling */
+  /** Insert a new node after a sibling (applies defaults for node data) */
   insertAfter<TNode extends AnyTreeNodePrimitive>(
     siblingId: string,
     nodeType: TNode,
-    data: InferTreeNodeDataState<TNode>
+    data: TreeNodeDataSetInput<TNode>
   ): string;
   
-  /** Insert a new node before a sibling */
+  /** Insert a new node before a sibling (applies defaults for node data) */
   insertBefore<TNode extends AnyTreeNodePrimitive>(
     siblingId: string,
     nodeType: TNode,
-    data: InferTreeNodeDataState<TNode>
+    data: TreeNodeDataSetInput<TNode>
   ): string;
   
   /** Remove a node and all its descendants */
@@ -218,6 +237,13 @@ export interface TreeProxy<TRoot extends AnyTreeNodePrimitive> {
     nodeType: TNode
   ): InferProxy<TNode["data"]>;
   
+  /** Updates only the specified data fields of a node (partial update) */
+  updateAt<TNode extends AnyTreeNodePrimitive>(
+    id: string,
+    nodeType: TNode,
+    value: TreeNodeUpdateValue<TNode>
+  ): void;
+  
   /** Convert tree to a nested snapshot for UI rendering */
   toSnapshot(): TreeNodeSnapshot<TRoot> | undefined;
 }
@@ -229,12 +255,14 @@ interface TreePrimitiveSchema<TRoot extends AnyTreeNodePrimitive> {
   readonly validators: readonly Validator<TreeState<TRoot>>[];
 }
 
-export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
-  implements Primitive<TreeState<TRoot>, TreeProxy<TRoot>>
+export class TreePrimitive<TRoot extends AnyTreeNodePrimitive, TDefined extends boolean = false, THasDefault extends boolean = false>
+  implements Primitive<TreeState<TRoot>, TreeProxy<TRoot>, TDefined, THasDefault>
 {
   readonly _tag = "TreePrimitive" as const;
   readonly _State!: TreeState<TRoot>;
   readonly _Proxy!: TreeProxy<TRoot>;
+  readonly _TDefined!: TDefined;
+  readonly _THasDefault!: THasDefault;
 
   private readonly _schema: TreePrimitiveSchema<TRoot>;
   private _nodeTypeRegistry: Map<string, AnyTreeNodePrimitive> | undefined;
@@ -271,7 +299,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
   }
 
   /** Mark this tree as required */
-  required(): TreePrimitive<TRoot> {
+  required(): TreePrimitive<TRoot, true, THasDefault> {
     return new TreePrimitive({
       ...this._schema,
       required: true,
@@ -279,7 +307,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
   }
 
   /** Set a default value for this tree */
-  default(defaultValue: TreeState<TRoot>): TreePrimitive<TRoot> {
+  default(defaultValue: TreeState<TRoot>): TreePrimitive<TRoot, true, true> {
     return new TreePrimitive({
       ...this._schema,
       defaultValue,
@@ -292,7 +320,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
   }
 
   /** Add a custom validation rule */
-  refine(fn: (value: TreeState<TRoot>) => boolean, message: string): TreePrimitive<TRoot> {
+  refine(fn: (value: TreeState<TRoot>) => boolean, message: string): TreePrimitive<TRoot, TDefined, THasDefault> {
     return new TreePrimitive({
       ...this._schema,
       validators: [...this._schema.validators, { validate: fn, message }],
@@ -404,11 +432,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
               );
             }
             const nodePath = operationPath.append(nodeState.id);
+            const dataProxy = nodeType.data._internal.createProxy(env, nodePath) as InferProxy<TNode["data"]>;
             return {
               id: nodeState.id,
               type: nodeType.type as InferTreeNodeType<TNode>,
-              data: nodeType.data._internal.createProxy(env, nodePath) as InferProxy<TNode["data"]>,
+              data: dataProxy,
               get: () => nodeState as TypedTreeNodeState<TNode>,
+              update: (value: TreeNodeUpdateValue<TNode>) => {
+                // Delegate to the data proxy's update method
+                (dataProxy as { update: (v: unknown) => void }).update(value);
+              },
             };
           },
           
@@ -474,7 +507,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
         insertFirst: <TNode extends AnyTreeNodePrimitive>(
           parentId: string | null,
           nodeType: TNode,
-          data: InferTreeNodeDataState<TNode>
+          data: TreeNodeDataSetInput<TNode>
         ): string => {
           const state = getCurrentState();
           const siblings = getOrderedChildren(state, parentId);
@@ -496,13 +529,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
             throw new ValidationError("Tree already has a root node");
           }
 
+          // Apply defaults to node data
+          const mergedData = applyDefaults(nodeType.data as AnyPrimitive, data as Partial<InferTreeNodeDataState<TNode>>) as InferTreeNodeDataState<TNode>;
+
           env.addOperation(
             Operation.fromDefinition(operationPath, this._opDefinitions.insert, {
               id,
               type: nodeType.type,
               parentId,
               pos,
-              data,
+              data: mergedData,
             })
           );
 
@@ -512,7 +548,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
         insertLast: <TNode extends AnyTreeNodePrimitive>(
           parentId: string | null,
           nodeType: TNode,
-          data: InferTreeNodeDataState<TNode>
+          data: TreeNodeDataSetInput<TNode>
         ): string => {
           const state = getCurrentState();
           const siblings = getOrderedChildren(state, parentId);
@@ -534,13 +570,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
             throw new ValidationError("Tree already has a root node");
           }
 
+          // Apply defaults to node data
+          const mergedData = applyDefaults(nodeType.data as AnyPrimitive, data as Partial<InferTreeNodeDataState<TNode>>) as InferTreeNodeDataState<TNode>;
+
           env.addOperation(
             Operation.fromDefinition(operationPath, this._opDefinitions.insert, {
               id,
               type: nodeType.type,
               parentId,
               pos,
-              data,
+              data: mergedData,
             })
           );
 
@@ -551,7 +590,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
           parentId: string | null,
           index: number,
           nodeType: TNode,
-          data: InferTreeNodeDataState<TNode>
+          data: TreeNodeDataSetInput<TNode>
         ): string => {
           const state = getCurrentState();
           const siblings = getOrderedChildren(state, parentId);
@@ -575,13 +614,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
             throw new ValidationError("Tree already has a root node");
           }
 
+          // Apply defaults to node data
+          const mergedData = applyDefaults(nodeType.data as AnyPrimitive, data as Partial<InferTreeNodeDataState<TNode>>) as InferTreeNodeDataState<TNode>;
+
           env.addOperation(
             Operation.fromDefinition(operationPath, this._opDefinitions.insert, {
               id,
               type: nodeType.type,
               parentId,
               pos,
-              data,
+              data: mergedData,
             })
           );
 
@@ -591,7 +633,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
         insertAfter: <TNode extends AnyTreeNodePrimitive>(
           siblingId: string,
           nodeType: TNode,
-          data: InferTreeNodeDataState<TNode>
+          data: TreeNodeDataSetInput<TNode>
         ): string => {
           const state = getCurrentState();
           const sibling = state.find(n => n.id === siblingId);
@@ -610,13 +652,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
           const parentType = getParentType(parentId);
           this._validateChildType(parentType, nodeType.type);
 
+          // Apply defaults to node data
+          const mergedData = applyDefaults(nodeType.data as AnyPrimitive, data as Partial<InferTreeNodeDataState<TNode>>) as InferTreeNodeDataState<TNode>;
+
           env.addOperation(
             Operation.fromDefinition(operationPath, this._opDefinitions.insert, {
               id,
               type: nodeType.type,
               parentId,
               pos,
-              data,
+              data: mergedData,
             })
           );
 
@@ -626,7 +671,7 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
         insertBefore: <TNode extends AnyTreeNodePrimitive>(
           siblingId: string,
           nodeType: TNode,
-          data: InferTreeNodeDataState<TNode>
+          data: TreeNodeDataSetInput<TNode>
         ): string => {
           const state = getCurrentState();
           const sibling = state.find(n => n.id === siblingId);
@@ -645,13 +690,16 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
           const parentType = getParentType(parentId);
           this._validateChildType(parentType, nodeType.type);
 
+          // Apply defaults to node data
+          const mergedData = applyDefaults(nodeType.data as AnyPrimitive, data as Partial<InferTreeNodeDataState<TNode>>) as InferTreeNodeDataState<TNode>;
+
           env.addOperation(
             Operation.fromDefinition(operationPath, this._opDefinitions.insert, {
               id,
               type: nodeType.type,
               parentId,
               pos,
-              data,
+              data: mergedData,
             })
           );
 
@@ -890,6 +938,29 @@ export class TreePrimitive<TRoot extends AnyTreeNodePrimitive>
           return nodeType.data._internal.createProxy(env, nodePath) as InferProxy<TNode["data"]>;
         },
 
+        updateAt: <TNode extends AnyTreeNodePrimitive>(
+          id: string,
+          nodeType: TNode,
+          value: TreeNodeUpdateValue<TNode>
+        ): void => {
+          // Get the node to verify its type
+          const state = getCurrentState();
+          const node = state.find(n => n.id === id);
+          if (!node) {
+            throw new ValidationError(`Node not found: ${id}`);
+          }
+          if (node.type !== nodeType.type) {
+            throw new ValidationError(
+              `Node is of type "${node.type}", not "${nodeType.type}"`
+            );
+          }
+
+          const nodePath = operationPath.append(id);
+          const dataProxy = nodeType.data._internal.createProxy(env, nodePath);
+          // Delegate to the data proxy's update method
+          (dataProxy as { update: (v: unknown) => void }).update(value);
+        },
+
         toSnapshot: (): TreeNodeSnapshot<TRoot> | undefined => {
           const state = getCurrentState();
           const rootNode = state.find(n => n.parentId === null);
@@ -1111,7 +1182,7 @@ export interface TreeOptions<TRoot extends AnyTreeNodePrimitive> {
 /** Creates a new TreePrimitive with the given root node type */
 export const Tree = <TRoot extends AnyTreeNodePrimitive>(
   options: TreeOptions<TRoot>
-): TreePrimitive<TRoot> =>
+): TreePrimitive<TRoot, false, false> =>
   new TreePrimitive({
     required: false,
     defaultValue: undefined,

package/src/primitives/Union.ts

@@ -7,14 +7,14 @@ import * as Transform from "../Transform";
 import type { Primitive, PrimitiveInternal, MaybeUndefined, AnyPrimitive, InferState, InferProxy, InferSnapshot } from "../Primitive";
 import { ValidationError } from "../Primitive";
 import { LiteralPrimitive } from "./Literal";
-import { StructPrimitive } from "./Struct";
-import { runValidators } from "./shared";
+import { StructPrimitive, InferStructState } from "./Struct";
+import { runValidators, applyDefaults, StructSetInput } from "./shared";
 
 
 /**
  * Type constraint for union variants - must be struct primitives
  */
-export type UnionVariants = Record<string, StructPrimitive<any, any>>;
+export type UnionVariants = Record<string, StructPrimitive<any, any, any>>;
 
 /**
  * Infer the union state type from variants
@@ -30,6 +30,16 @@ export type InferUnionSnapshot<TVariants extends UnionVariants> = {
   [K in keyof TVariants]: InferSnapshot<TVariants[K]>;
 }[keyof TVariants];
 
+/**
+ * Compute the input type for union.set() operations.
+ * For each variant, uses StructSetInput to make fields with defaults optional.
+ */
+export type UnionSetInput<TVariants extends UnionVariants> = {
+  [K in keyof TVariants]: TVariants[K] extends StructPrimitive<infer TFields, any, any>
+    ? StructSetInput<TFields>
+    : InferState<TVariants[K]>;
+}[keyof TVariants];
+
 /**
  * Proxy for accessing union variants
  */
@@ -37,8 +47,8 @@ export interface UnionProxy<TVariants extends UnionVariants, _TDiscriminator ext
   /** Gets the current union value */
   get(): MaybeUndefined<InferUnionState<TVariants>, TDefined>;
   
-  /** Sets the entire union value */
-  set(value: InferUnionState<TVariants>): void;
+  /** Sets the entire union value (applies defaults for variant fields) */
+  set(value: UnionSetInput<TVariants>): void;
   
   /** Access a specific variant's proxy (assumes the variant is active) */
   as<K extends keyof TVariants>(variant: K): InferProxy<TVariants[K]>;
@@ -59,12 +69,14 @@ interface UnionPrimitiveSchema<TVariants extends UnionVariants, TDiscriminator e
   readonly variants: TVariants;
 }
 
-export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator extends string = "type", TDefined extends boolean = false>
-  implements Primitive<InferUnionState<TVariants>, UnionProxy<TVariants, TDiscriminator, TDefined>>
+export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator extends string = "type", TDefined extends boolean = false, THasDefault extends boolean = false>
+  implements Primitive<InferUnionState<TVariants>, UnionProxy<TVariants, TDiscriminator, TDefined>, TDefined, THasDefault>
 {
   readonly _tag = "UnionPrimitive" as const;
   readonly _State!: InferUnionState<TVariants>;
   readonly _Proxy!: UnionProxy<TVariants, TDiscriminator, TDefined>;
+  readonly _TDefined!: TDefined;
+  readonly _THasDefault!: THasDefault;
 
   private readonly _schema: UnionPrimitiveSchema<TVariants, TDiscriminator>;
 
@@ -82,7 +94,7 @@ export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator exte
   }
 
   /** Mark this union as required */
-  required(): UnionPrimitive<TVariants, TDiscriminator, true> {
+  required(): UnionPrimitive<TVariants, TDiscriminator, true, THasDefault> {
     return new UnionPrimitive({
       ...this._schema,
       required: true,
@@ -90,10 +102,12 @@ export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator exte
   }
 
   /** Set a default value for this union */
-  default(defaultValue: InferUnionState<TVariants>): UnionPrimitive<TVariants, TDiscriminator, true> {
+  default(defaultValue: UnionSetInput<TVariants>): UnionPrimitive<TVariants, TDiscriminator, true, true> {
+    // Apply defaults to the variant
+    const merged = this._applyVariantDefaults(defaultValue as Partial<InferUnionState<TVariants>>);
     return new UnionPrimitive({
       ...this._schema,
-      defaultValue,
+      defaultValue: merged,
     });
   }
 
@@ -119,7 +133,7 @@ export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator exte
       const variant = this._schema.variants[key]!;
       const discriminatorField = variant.fields[this._schema.discriminator];
       if (discriminatorField && discriminatorField._tag === "LiteralPrimitive") {
-        const literalPrimitive = discriminatorField as LiteralPrimitive<any, any>;
+        const literalPrimitive = discriminatorField as LiteralPrimitive<any, any, any>;
         if (literalPrimitive.literal === discriminatorValue) {
           return key;
         }
@@ -128,6 +142,17 @@ export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator exte
     return undefined;
   }
 
+  /** Apply defaults to a variant value based on the discriminator */
+  private _applyVariantDefaults(value: Partial<InferUnionState<TVariants>>): InferUnionState<TVariants> {
+    const variantKey = this._findVariantKey(value as InferUnionState<TVariants>);
+    if (!variantKey) {
+      return value as InferUnionState<TVariants>;
+    }
+    
+    const variantPrimitive = this._schema.variants[variantKey]!;
+    return applyDefaults(variantPrimitive as AnyPrimitive, value) as InferUnionState<TVariants>;
+  }
+
   readonly _internal: PrimitiveInternal<InferUnionState<TVariants>, UnionProxy<TVariants, TDiscriminator, TDefined>> = {
     createProxy: (
       env: ProxyEnvironment.ProxyEnvironment,
@@ -141,9 +166,11 @@ export class UnionPrimitive<TVariants extends UnionVariants, TDiscriminator exte
           const state = env.getState(operationPath) as InferUnionState<TVariants> | undefined;
           return (state ?? defaultValue) as MaybeUndefined<InferUnionState<TVariants>, TDefined>;
         },
-        set: (value: InferUnionState<TVariants>) => {
+        set: (value: UnionSetInput<TVariants>) => {
+          // Apply defaults for the variant
+          const merged = this._applyVariantDefaults(value as Partial<InferUnionState<TVariants>>);
           env.addOperation(
-            Operation.fromDefinition(operationPath, this._opDefinitions.set, value)
+            Operation.fromDefinition(operationPath, this._opDefinitions.set, merged)
           );
         },
         as: <K extends keyof TVariants>(variant: K): InferProxy<TVariants[K]> => {
@@ -311,13 +338,13 @@ export interface UnionOptions<TVariants extends UnionVariants, TDiscriminator ex
 /** Creates a new UnionPrimitive with the given variants */
 export function Union<TVariants extends UnionVariants>(
   options: UnionOptions<TVariants, "type">
-): UnionPrimitive<TVariants, "type", false>;
+): UnionPrimitive<TVariants, "type", false, false>;
 export function Union<TVariants extends UnionVariants, TDiscriminator extends string>(
   options: UnionOptions<TVariants, TDiscriminator>
-): UnionPrimitive<TVariants, TDiscriminator, false>;
+): UnionPrimitive<TVariants, TDiscriminator, false, false>;
 export function Union<TVariants extends UnionVariants, TDiscriminator extends string = "type">(
   options: UnionOptions<TVariants, TDiscriminator>
-): UnionPrimitive<TVariants, TDiscriminator, false> {
+): UnionPrimitive<TVariants, TDiscriminator, false, false> {
   const discriminator = (options.discriminator ?? "type") as TDiscriminator;
   return new UnionPrimitive({
     required: false,