Xstate: context replace

[XaxisJosh]

Question:
Is there a way to completely replace context in an action (instead of just merging new values into it)?

It seems the only action that is provided to interact with context is assign. Similar to Object.assign it will take the current context and and merge updates into it. But what if I only want some keys on context in certain states?

In the below example as soon as CLICK is called once both oddClicks and evenClicks will exist on context.

Is there some there some way to completely replace the context?

I'm interested in this because I would like to use a type union for the context to prevent accessing of values that aren't on the context in a given state.

Or am I not approaching the problem in the correct way?

const buttonMachine = Machine({
  id: 'buttonMachine',
  initial: 'odd',
  states: {
    odd: {
      on: {
        CLICK: {
          target: 'even',
          actions: 'oddClicked',
        },
      },
    },
    even: {
      on: {
        CLICK: {
          target: 'odd',
          actions: 'evenClicked',
        },
      },
    },
  },
  context: {
    oddClicks: 1,
  },
}, {
  actions: {
    evenClicked: actions.assign({
      oddClicks: ctx => ctx.evenClicks + 1,
    }),
    oddClicked: actions.assign({
      evenClicks: ctx => ctx.oddClicks + 1,
    }),
  },
});

[davidkpiano]

If you do that, won't ctx.evenClicks at some point be undefined? Then you'd be adding undefined + 1, which is no good.

[XaxisJosh]

I think that the only time evenClicks is undefined would be at the initial state. In theory typescript/flow should prevent me from doing undefined + 1.

What i'm trying to avoid is that in some states I don't want data to be on an object that shouldn't be referenced (that may have been valid in earlier states). And there doesn't seem to be a way to do that via the current assign method.

In the above case i would have types like:

type EvenState = {
  kind: 'even',
  evenClicks: number,
};
type OddState = {
  kind: 'odd',
  oddClicks: number,
};
type MyStates = OddState | EvenState;

[davidkpiano]

I don't think it makes sense for properties in the context to just "not exist" when they've previously existed - would this suffice (or be the same)?

type EvenState = {
  kind: 'even',
  evenClicks: number,
  oddClicks: undefined
}

// etc.

We can probably make an enhancement to StateSchema that narrows down the TContext type:

interface StateSchema {
  states: {
    even: {
      context: { evenClicks: number; oddClicks: undefined; }
    },
    odd: {
      context: { evenClicks: undefined; oddClicks: number; }
    }
  }
}

[ksaldana1]

I was experimenting with how we could encode more invariants into our state machine and I got a bit stuck on implementing a few of my "ideals". The goal is to extend "making impossible states impossible" to compile time contracts.

1. Context is derived from the StateSchema

Something like @davidkpiano's example above could work for simple cases.

// probably not the best type code
type PluckContext<T> = T extends StateSchema
  ? T['context'] // assuming we add context property to StateSchema
  : never;

type ContextMapFromStateSchema<T extends StateSchema> = {
  [K in keyof T['states']]: PluckContext<T['states'][K]>
};

type ContextUnionFromContextMap<
  T extends ContextMapFromStateSchema<any>
> = T[keyof T];

export type ContextUnionFromStateSchema<
  T extends StateSchema
> = ContextUnionFromContextMap<ContextMapFromStateSchema<T>>;

// example 
interface StateSchema {
  states: {
    even: {
      context: { value: 'EVEN' };
    };
    odd: {
      context: { value: 'ODD' };
    };
  };
}

type Union = ContextUnionFromStateSchema<StateSchema> 
// { value: "EVEN" } | { value: "ODD" } 

type Map = ContextMapFromStateSchema<StateSchema>
/*
  {
    even: {
        value: "EVEN"
    };
    odd: {
        value: "ODD"
    }
 }
*/

2. Context is correctly narrowed inside StateNodeConfig

Using the above example, I want this behavior:

  states: {
      even: { /* Inside here Context type is { value: "EVEN" }  */},
      odd: { /* Inside here Context type is { value: "ODD" } */ }
  }

3. Context is correctly narrowed inside TransitionConfig (StateNode Context + Target Context)

Using the above example, I want this behavior:

 const states = {
  even: {
    ON: {
      CLICK: {
        target: 'odd',
        actions: assign((_ctx /* ctx here is { value: "EVEN" } */) => ({
          value: 'ODD',
        })),
      },
    },
  },
  odd: {
    ON: {
      CLICK: {
        target: 'even', // <- this would need to be generic
        actions: assign((_ctx /* ctx here is { value: "ODD" } */) => ({
          // EXPECTED: ERROR value "ODD" not assignable to value "EVEN"
          value: 'ODD', // Correct return value inferred from "target" property above
        })),
      },
    },
  },
};

Getting 1 to work isn't the hardest thing here--deriving from the StateSchema instead of an explicit Context interface is not too different and the changes to the typings could be pretty minimal (at least in theory). The recursive nature of all of this makes it less straightforward than I'm explaining here.

2 and 3 are a lot bigger undertakings (not even sure if fully possible). There would need to be more generics passed around... for example in 3 target: 'even' would need to be a generic (TTarget?) to correctly infer the behavior of the actions property on that TransitionConfig.

Not sure if this is the best place for these thoughts, or if this is something we'd even want to pursue. These are just some things I think would be really powerful if we could achieve. However these types of invariants would probably force a smaller API surface area to make it possible. Those changes would probably makes this lib less useful for JS users, so I'm not sure how valuable it really is in the long-term.

Regardless this is just me brain dumping after playing around with the types in this library.

[davidkpiano]

On a potentially related note, it would be great if we could instead infer the state schema of individual states, but I'm not sure how to go about doing that in a simple way: https://stackoverflow.com/questions/54664214/typescript-relate-to-own-property-in-type/54666233

[ksaldana1]

@davidkpiano Not sure if I'm fully capturing what your requirements are in that SO question, but here's a possible quick solution:

interface Config<T extends Record<string, Config<any>>> {
  initial?: keyof T;
  states?: T;
}

declare function createSomething<T extends Record<string, Config<any>>>(
  config: Config<T>
): void;

// valid
createSomething({
  states: {
    foo: { states: {} },
    bar: { states: {} },
    baz: { states: {} },
  },
  initial: 'foo',
});

// error Type '"invalid"' is not assignable to type '"foo" | "bar" | "baz"'
createSomething({
  states: {
    foo: { states: {} },
    bar: { states: {} },
    baz: { states: {} },
  },
  initial: 'invalid',
});

[ksaldana1]

Recursion is where this stuff starts getting tricky. The above code isn't going to narrow your child members correctly, so you'd have to wrap nested Configs yourself to get a proper type out of it.

interface Config<T extends Record<string, Config<any>>> {
  initial?: keyof T;
  states?: T;
}

declare function createConfig<T extends Record<string, Config<any>>>(
  config: Config<T>
): Config<T>;

const config = createConfig({
  states: {
    foo: createConfig({
      states: {
        blah: {},
        blam: {},
      },
      initial: 'blah',
    }),
    bar: createConfig({
      states: {
        counter: createConfig({
          states: {
            odd: {},
            even: {},
          },
          initial: 'oddz', // ERROR Type '"oddz"' is not assignable to type '"odd" | "even"'
        }),
      },
    }),
    baz: {},
  },
  initial: 'foo',
});

[ksaldana1]

Iterating on the above examples, here's a wrapper that enforces a proper initial + initialContext pairing.

interface StateSchema<T = any> {
  meta?: any;
  states?: Record<string, StateSchema>;
  context?: T;
}

type PluckContext<T> = T extends StateSchema<infer TContext> ? TContext : never;

type ContextMapFromStateSchema<T extends StateSchema> = {
  [K in keyof T['states']]: PluckContext<T['states'][K]>
};

type ContextUnionFromContextMap<T extends ContextMapFromStateSchema<any>> = T[keyof T];

type ContextUnionFromStateSchema<
  T extends StateSchema
> = ContextUnionFromContextMap<ContextMapFromStateSchema<T>>;

interface SchemaConfig<T extends StateSchema, K extends keyof T['states']>
  extends StateSchema {
  initial: K;
  initialContext?: ContextMapFromStateSchema<T>[K];
}

declare function createSchema<T extends StateSchema, K extends keyof T['states']>(
  schema: SchemaConfig<T, K>
): SchemaConfig<T, K>;

interface ExampleSchema {
  states: {
    odd: {
      context: {
        value: 'ODD';
      };
    };
    even: {
      context: {
        value: 'EVEN';
      };
    };
  };
}

/*
  Unfortunate that I have to pass in the second generic argument here
  Hopefully something that https://github.com/Microsoft/TypeScript/pull/26349
  can fix in the future. Despite the slight annoyance, this gives us good safety
*/ 
const validSchemaConfig = createSchema<ExampleSchema, 'even'>({
  states: {
    odd: {},
    even: {},
  },
  initial: 'even',
  initialContext: {
    value: 'EVEN', 
  },
});

const invalidSchemaConfig = createSchema<ExampleSchema, 'even'>({
  states: {
    odd: {},
    even: {},
  },
  initial: 'odd', // ERROR: Type '"odd"' is not assignable to type '"even"'
  initialContext: {
    value: 'ODD', // ERROR: Type '"ODD"' is not assignable to type '"EVEN"'
  },
});

Sorry for polluting this issue thread, I can move the discussion somewhere else if there's a more appropriate place. This issue just seemed most closely related to this stuff.

[ksaldana1]

I spent a bit more time tonight looking into my invariant wishes and seeing if what I wanted was even conceptually possible in TS's type system. I was having trouble getting what I wanted with the current typings in the library—because they are incredibly useful and serve a wide variety of use cases that extend beyond my narrow niche needs, but that also makes encoding invariants difficult.

I took a step back and worked with my own types to try to isolate my problem a bit. One of my favorite parts of xstate over Redux is how much it forces me to think deeply about my transitions. By having a mechanism to deal with unexpected transitions, we solve a large category of problems that usually pop up in UI development. But I still think there's similar value in designing those expected transitions and I'd like to infer even more from that upfront design.

We had already considered inferring context from our state schema—but why stop there? What if we made transitions a first class type in the schema?

interface StateSchema<TEvents extends EventObject, TAllStates extends string> {
  states?: { [K in TAllStates]: StateNode<TEvents, TAllStates> };
}
interface StateNode<
  TEvents extends EventObject,
  TAllStates extends string,
  TContext = any
> extends StateSchema<TEvents, TAllStates> {
  context: TContext;
  transitions: Array<Transition<TEvents, TAllStates>>;
}
interface Transition<TEvents extends EventObject, TAllStates extends string> {
  to: TAllStates;
  event: TEvents;
}

// Example Schema
enum States {
  EVEN = 'EVEN',
  ODD = 'ODD',
}

enum EventTypes {
  CLICK = 'CLICK',
  DOUBLE_CLICK = 'DOUBLE_CLICK',
}

interface ClickEvent {
  type: EventTypes.CLICK;
}

interface DoubleClickEvent {
  type: EventTypes.DOUBLE_CLICK;
}

interface ExampleSchema {
  states: {
    [States.EVEN]: {
      context: {
        value: 'EVEN';
      };
      transitions: [
        { to: States.ODD; event: ClickEvent },
        { to: States.EVEN; event: DoubleClickEvent }
      ];
    };
    [States.ODD]: {
      context: {
        value: 'ODD';
      };
      transitions: [
        { to: States.EVEN; event: ClickEvent },
        { to: States.ODD; event: DoubleClickEvent }
      ];
    };
  };
}

Bear with me for a second... I recognize this has a whole host of issues and limitations in design (like really important runtime information missing to function) and is super boilerplate-y non-sense, but the idea is that by declaring this stuff up front we can provide a really awesome developer experience downstream in the implementation. This satisfies all 3 things (context derived from schema, context scoped to state nodes, context parameter + return typesafe in actions) I mentioned in my earlier comment, as well as providing a few other nice things (exhaustiveness of transition implementations, easy dead code removal).

const config: TransitionConfig<ExampleSchema> = {
  [States.ODD]: {
    on: {
      CLICK: (_ctx /* ctx is { value: 'ODD' } */ , _event) => {
        return { value: "EVE" }; // ERROR Type '"EVE"' is not assignable to type '"EVEN"'
        // because we declared that CLICK transitions to States.EVEN, Context must equal { value: "EVEN" }
      },
      DOUBLE_CLICK: (_ctx, _event) => {
        return { value: "ODD" };
      }
    }
  },
  [States.EVEN]: {
    on: { // ERROR Property 'DOUBLE_CLICK' is missing in type '{ CLICK: (_ctx: { value: "EVEN"; }, _event: ClickEvent) => { value: "ODD"; }; }
      // we forgot to implement our DOUBLE_CLICK transition according to our schema
      CLICK: (_ctx /* ctx is { value: "EVEN" } */, _event) => {
        return { value: "ODD" };
      },
    }
  }
};

The idea is that your TS "schema" is really the source of truth. There is no event union type or context interface—there is only the schema. Adding or removing states / transitions / events to this interface should propagate changes to the config implementations.

This stuff veered off into space pretty hard and I know most of this is complete non-sense to think about for a library of this popularity. I mostly just wanted to document some of my learnings using xstate and fusing it with my fondness for static type systems. Maybe there's some nugget that someone can pull away that's useful from these ideas.

I saved the TransitionConfig code for down here... you are warned right now. This could ugh.... probably be better... but it's late and I was trying to get something to work. When you have code that looks this awful I probably need some better type abstractions....

type TransitionConfig<T extends StateSchema<any, any>> = {
  [K in keyof T['states']]: {
    // oh god oh no
    on: {
      [E in T['states'][K]['transitions'][number]['event']['type']]: ActionFunction<
        ContextMapFromStateSchema<T>[K],
        Extract<T['states'][K]['transitions'][number]['event'], { type: E }>,
        ContextMapFromStateSchema<T>[Extract<
          T['states'][K]['transitions'][number],
          { event: { type: E } }
        >['to']]
      >
    };
  }
};

Full code snippets here.

[ksaldana1]

Been reading more about FSM implementations in other languages to see what type of language and type guarantees they are striving for. Since I'm already well into the clouds here, I wanted to take a look at what the Haskell community was doing when I came across this cool blog post. Part 2 is pretty relevant to the stuff I am talking about here and I'm going to steal their good vocabulary going forward. What I was calling the "Schema" (states + context + transitions) I think is better named "Protocol" in this case.

I got some of the recursive trickiness to work yesterday.

import { Lookup, SchemaConfig, StateProtocol } from './types';

// Pedestrian Protocol
enum PedestrianStates {
  WALK = 'WALK',
  WAIT = 'WAIT',
  STOP = 'STOP',
}

enum PedestrianEventTypes {
  PED_TIMER = 'PED_TIMER',
}

interface PedestrianTimerEvent {
  type: PedestrianEventTypes.PED_TIMER;
}

interface PedestrianProtocol {
  states: {
    [PedestrianStates.WALK]: {
      context: { value: 'RED.WALK' };
      transitions: [{ to: PedestrianStates.WAIT; event: PedestrianTimerEvent }];
    };
    [PedestrianStates.WAIT]: {
      context: { value: 'RED.WAIT' };
      transitions: [{ to: PedestrianStates.STOP; event: PedestrianTimerEvent }];
    };
    [PedestrianStates.STOP]: {
      context: { value: 'RED.STOP' };
      transitions: [];
    };
  };
}

// Light Protocol
enum LightStates {
  GREEN = 'GREEN',
  YELLOW = 'YELLOW',
  RED = 'RED',
}

enum LightEventTypes {
  TIMER = 'TIMER',
}

interface TimerEvent {
  type: LightEventTypes.TIMER;
}

interface LightProtocol {
  states: {
    [LightStates.GREEN]: {
      context: { value: 'GREEN' };
      transitions: [{ to: LightStates.YELLOW; event: TimerEvent }];
    };
    [LightStates.YELLOW]: {
      context: { value: 'YELLOW' };
      transitions: [{ to: LightStates.RED; event: TimerEvent }];
    };
    [LightStates.RED]: {
      context: { value: 'RED.WALK' };
      transitions: [{ to: LightStates.GREEN; event: TimerEvent }];
      states: PedestrianProtocol;
    };
  };
}
````

```typescript
// Light Configuration
const config: SchemaConfig<LightProtocol, LightStates.RED> = {
  initial: LightStates.RED,
  context: { value: 'RED.WALK' },
  states: {
    GREEN: {
      on: {
        TIMER: (_ctx, _event) => {
          return { value: 'YELLOW' };
        },
      },
    },
    RED: {
      on: {
        TIMER: (_ctx, _event) => {
          return { value: 'GREEN' };
        },
      },
      states: {
        STOP: {
          on: {},
        },
        WAIT: {
          on: {
            PED_TIMER: (_ctx, _event) => {
              return { value: 'RED.STOP' };
            },
          },
        },
        WALK: {
          on: {
            PED_TIMER: (_ctx, _event) => {
              return { value: 'RED.WAIT' };
            },
          },
        },
      },
    },
    YELLOW: {
      on: {
        TIMER: (_ctx, _event) => {
          return { value: 'RED.WALK' };
        },
      },
    },
  },
};

This has all the same type guarantees as the earlier stuff. Something else I worked on was the idea of a type-safe "match" function. Excuse the terrible type code here:

function matchFactory<T extends StateProtocol<any, any>>() {
  function match<K extends keyof T['states']>(k: K): boolean;
  function match<
    K extends keyof T['states'],
    K1 extends keyof Lookup<Lookup<T['states'][K], 'states'>, 'states'>
  >(k: K, k1: K1): boolean;
  function match<
    K extends keyof T['states'],
    K1 extends keyof Lookup<Lookup<T['states'][K], 'states'>, 'states'>,
    K2 extends keyof Lookup<Lookup<Lookup<T['states'][K], 'states'>, K1>, 'states'>
  >(k: K, k1: K1, k2: K2): boolean;
  function match<
    K extends keyof T['states'],
    K1 extends keyof Lookup<Lookup<T['states'][K], 'states'>, 'states'>,
    K2 extends keyof Lookup<Lookup<Lookup<T['states'][K], 'states'>, K1>, 'states'>,
    K3 extends keyof Lookup<
      Lookup<Lookup<Lookup<T['states'][K], 'states'>, K1>, K2>,
      'states'
    >
  >(k: K, k1: K1, k2: K2, k3: K3): boolean;
  function match(...args: any[]) {
    // mock impl
    return true;
  }
  return match;
}

This code again looks terrible, but the API it provides for state.matches could be pretty dope.

I don't think the current typings in the library are limiting us from doing something like this (except reasonable aversion to putting horribly complicated type code in the codebase).

[davidkpiano]

This is really awesome stuff @ksaldana1 - I'll look at it deeper soon. I'm imagining that a state schema can be _generated_ from a machine definition 🤔

[ksaldana1]

@davidkpiano I have put the code up in a GH repo if you want to reference it later. I have also started a thread in the state charts Spectrum channel. I've gone well off-topic here, so probably a more appropriate place to chat about this stuff.