Proposal: Encoding store mutations type-level

[devanshj]

Middlewares in zustand can and do mutate the store, which makes it impossible to type them with without some workaround, this proposal tries to encode mutations type-level to make zustand almost 100% type-safe and still extensible.

I'll walk you through how this came to be, because I think it makes one understand how it works and why certain elements of it are necessary.

PR #662 already implements this proposal so you can try it out if you want. And here's the final playground to try things out.

Index

This is less of a "table of content" and more of an "index", meaning the whole thing is single-read and there are no independent "sections" per se, they are made as such to aid the reader to track their progress because this document is rather too long to be a single-read.

• Independent mutations
  • Forwarding mutations up the tree
  • Forwarding mutations down the tree
  • Ordered writes
  • Subtractive mutations
• Dependent mutations
  • Higher kinded mutations
  • Higher kinded mutators
• Zustand-adapted version
• To-dos
  • Caveats
  • Others
• Action point

Independent mutations

Consider this scenario (I'll be making the api simpler only for the explanation, there are no api changes required)...

const withA = (a, f) => store => {
  store.a = a,
  return f(store);
}

const withB = (b, f) => store => {
  store.b = b,
  return f(store);
}

let storeOut = create(
  withB("b",
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

Here both storeOut and storeIn should be of type Store<unknown> & { a: string } & { b: string } (we don't care about inferring the state { count: number } for now)

So how would we approach this? First let's take a simpler version

const withA = (a, f) => store => {
  store.a = a,
  return f(store);
}

let storeOut = create(
  withA("a",
    storeIn =>
      ({ count: 0 })
  ),
)

One could begin with something like...

// https://tsplay.dev/w6XJ6m

declare const create:
  <T, M>
    ( initializer:
        & ((store: Store<T>) => T)
        & { __mutation?: M }
    ) =>
    Store<T> & M

type Store<T> =
  { get: () => T
  , set: (t: T) => void
  }
  
declare const withA: 
  <A, T>
    ( a: A
    , initializer:
        (store: Store<T> & { a: A }) => T
    ) =>
      & ((store: Store<T>) => T)
      & { __mutation?: { a: A } }


let storeOut = create(
  withA("a",
    storeIn =>
      ({ count: 0 })
  ),
)

And it works. storeOut and storeIn are Store<unknown> & { a: string }. But why does it work? Intuitively you could think that had the store been immutable the middlewares would have returned the mutation isn't it? That's what we're doing here, the middlewares not only return the initializer but also the mutation to the store.

Forwarding mutations up the tree

Okay now let's see if it works for a more complex example...

// https://tsplay.dev/mbGe3W

declare const create:
  <T, M>
    ( initializer:
        & ((store: Store<T>) => T)
        & { __mutation?: M }
    ) =>
      Store<T> & M

type Store<T> =
  { get: () => T
  , set: (t: T) => void
  }
  
declare const withA: 
  <A, T>
    ( a: A
    , initializer:
        (store: Store<T> & { a: A }) => T
    ) =>
      & ((store: Store<T>) => T)
      & { __mutation?: { a: A } }

declare const withB: 
  <B, T>
    ( b: B
    , initializer:
        (store: Store<T> & { b: B }) => T
    ) =>
      & ((store: Store<T>) => T)
      & { __mutation?: { b: B } }


let storeOut = create(
  withB("b",
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

Hmm, it doesn't work. storeOut is Store<unknown> & { b: string } and storeIn is Store<unknown> & { a: string }. Why is that? Because were aren't accommodating the fact that the "child" of withB ie withA itself comes with its mutation so we need to carry forward mutations up the tree. Let's do that¹...

  // https://tsplay.dev/mpvdzw

  declare const withA: 
    < A
    , T
+   , M extends {}
    >
      ( a: A
      , initializer:
          & ((store: Store<T> & { a: A }) => T)
+         & { __mutation?: M }
      ) =>
        & ((store: Store<T>) => T)
        & { __mutation?:
+            & M
             & { a: A }
          }

So when the child itself has a mutation M (ie when the initializer has { mutation?: M }) we carry it forward as { __mutation?: M & { a: A } } instead of just { __mutation?: { a: A } }.

Also the extends {} constraint on M is so that it doesn't get fixated to undefined because of the ? in { __mutation?: M }.

Forwarding mutations down the tree

Okay so now the mutation is forwarded upstream and storeOut is Store<unknown> & { a: string } & { b: string }. But storeIn is still Store<unknown> & { a: string }. That's because we have only forwarded mutations upstream, we also need to forward the mutations downstream. Let's see how we can do that...

  // https://tsplay.dev/m3aLqw

  declare const withA: 
    < A
    , T
    , Mc extends {}
+   , Mp extends {}
    >
      ( a: A
      , initializer:
          & ((store:
              & Store<T>
+             & Mp
              & { a: A }
            ) => T)
          & { __mutation?: Mc }
      ) =>
        & ((store:
            & Store<T>
+           & Mp
          ) => T)
        & { __mutation?: Mc & { a: A } }

But it doesn't compile. Why so? Let me first show you a minimal equivalent of the above...

// https://tsplay.dev/wR9kXW

declare const create:
  <T>(initializer: (store: Store<T>) => T) => Store<T>

type Store<T> =
  { get: () => T
  , set: (t: T) => void
  }
  
declare const identityMiddleware: 
  <T, S extends Store<T>>
    (initializer: (store: S) => T) =>
      (store: S) => T

let storeOut = create(
  identityMiddleware(
    storeIn =>
      ({ count: 0 })
  )
)

Now what's happening here? Why does it not compile? The problem is TypeScript can't contextually infer S, it eagerly resolves it to Store<unknown>. And why does it do that? Well it'll take me rather too long to develop what I know intuitively so this comment is the best I can offer right now.

Long story short it doesn't work, we need a workaround. Which looks something like this...

  // https://tsplay.dev/wX2ALm

  declare const withA: 
    < A
    , T
    , Mc extends {}
    , Mp extends {}
    >
      ( a: A
      , initializer:
          & ( ( store:
                & Store<T>
                & Mp
                & { a: A }
+             , __mutation: UnionToIntersection<Mp | { a: A }>
              ) => T
            )
          & { __mutation?: Mc }
      ) =>
        & ( ( store:
              & Store<T>
-             & Mp
+           , __mutation: Mp
            ) => T
          )
        & { __mutation?: Mc & { a: A } }

+ type UnionToIntersection<U> =
+   (U extends unknown ? (u: U) => void : never) extends (i: infer I) => void
+     ? I
+     : never

So instead of passing the mutation from parent as (store: Store<T> & Pm) => we're doing (store: Store<T>, __mutation: Pm) => which in spirit is the same thing. And instead of Pm & { a: A } we trick the compiler by doing UnionToIntersection<Pm | { a: A }> which is effectively resolves to the same thing.

Why does all this works? Well, I kinda vaguely know it in my "feels" but it's going to take me a lot of experimentation to come up with a good enough theory, so let's just skip that xD

So now finally storeOut and storeIn both are Store<unknown> & { a: string } & { b: string }

Ordered writes

What if we have a scenario like this...

const withA = (a, f) => store => {
  store.a = a;
  return f(store)
}

let storeOut = create(
  withA(1,
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

Here with our existing approach storeOut and storeIn would be Store<unknown> & { a: number } & { a: string } when it should have been Store<unknown> & { a: string }. That means we need to "write" instead of just intersecting and order of mutations matters too. Let's see how we can do that...

  // https://tsplay.dev/NnXdvW

  declare const withA: 
    < A
    , T
-   , Mc extends {}
+   , Mcs extends {}[] = []
-   , Mp extends {}
+   , Mps extends {}[] = []
    >
      ( a: A
      , initializer:
-         & ( ( store: Store<T> & Mp & { a: A }
+         & ( ( store: Mutate<Store<T>, [...Mps, { a: A }]>
-             , __mutation: UnionToIntersection<Mp | { a: A }>
+             , __mutations: [...Mp, { a: A }]
              ) => T
            )
-         & { __mutation?: Mp }
+         & { __mutations?: Mps }
      ) =>
        & ( ( store: Store<T>
-           , __mutation: Mp
+           , __mutations: Mps
            ) => T
          )
-       & { __mutation?: { a: A } & Mps  }
+       & { __mutations?: [{ a: A }, ...Mcs] }

+ type Mutate<S, Ms> =
+   Ms extends [] ? S :
+   Ms extends [infer M, ...infer Mr] ? Mutate<Write<S, M>, Mr> :
+   never
+ 
+ type Write<T, U> =
+   Omit<T, keyof U> & U

Mps happen before because they are come from the parent, Mcs happen after because they come from the child.

Note that we are doing = [] for Mcs and Mps because previously when there's no { __mutation } found (like in case of () => ({ count: 0 })) the compiler would resolve them to their constraints ie {} which also happened to be the default. Now the default is [] so we need to make it explicit. So technically speaking previously too we should have had Mc extends {} = {} but that's almost same as writing Mc extends {}.

So now storeOut and storeIn are Mutate<Store<unknown>, [{ a: string }, { a: number }]> except that in case of storeOut, number gets narrowed to 1. We can fix that by blocking inference of A at the places it isn't supposed to be inferred from...

  // https://tsplay.dev/Wk5d2N

  declare const withA: 
    < A
    , T
    , Mcs extends {}[] = []
    , Mps extends {}[] = []
    >
      ( a: A
      , initializer:
-         & ( ( store: Mutate<Store<T>, [...Mps, { a: A }]>
+         & ( ( store: Mutate<Store<T>, [...Mps, { a: NoInfer<A> }]>
-             , __mutations: [...Mps, { a: A }]
+             , __mutations: [...Mps, { a: NoInfer<A> }]
              ) => T
            )
          & { __mutations?: Cm }
      ) =>
        & ((store: Store<T>, __mutations: Pm) => T)
-       & { __mutations?: [{ a: A }, ...Mcs] }
+       & { __mutations?: [{ a: NoInfer<A> }, ...Mcs] }

+ type NoInfer<T> =
+   [T][T extends unknown ? 0 : never]

So now we're inferring A only from the argument, as it should be. And now number no longer gets narrowed to 1.

Also remember that child mutations can also happen before the current one, example...

const withA = (a, f) => store => {
  let state = f(store)
  store.a = a;
  return state;
}

let storeOut = create(
  withA(1,
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

In this case we'd expect storeOut.a to be number instead of string unlike the previous case. Depending on the implementation, the author can type the middleware accordingly, in this case we're first doing the child mutations so the type becomes...

  // https://tsplay.dev/wg6kbW

- & { __mutations?: [{ a: NoInfer<A> }, ...Mcs] }
+ & { __mutations?: [...Mcs, { a: NoInfer<A> }] }

Subtractive mutations

Although rare but it's important to note that now mutations can make a store no longer extend Store<T>, in that case some cases that should compile won't compile, example...

// https://tsplay.dev/WYBpEw

const withReadonly = f => store => {
  let readonlyStore = { ...store };
  delete readonlyStore.set;
  return f(store);
}

let storeOut = create(
  withReadonly(
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

If you scroll at the end of the error message you'd see it says "'undefined' is not assignable to type '(t: unknown) => void'" that's because withA's type says it requires Store<T> when in fact it doesn't use set and only sets a property a. Previously, we moved from (store: Store<T> & Pm) => to (store: Store<T>, __mutation: Pm) => void because the former didn't work. But now we can go back to that and use Mutate<Store<T>, Mps> which would solve our problem and make it compile...

// https://tsplay.dev/WvYdRm

-       & ((store: Store<T>, __mutations: Mps) => T)
+       & ((store: Mutate<Store<T>, Mps>, __mutations: Mps) => T)

But what if withA actually requires set and want the store to be of type Store<T>? In that case we can specify our requirements and then we get the original error back which is now warrant...

// https://tsplay.dev/WyO7dN

-       & ((store: Mutate<Store<T>, Mps>, __mutations: Mps) => T)
+       & ((store: Mutate<Store<T>, Mps> & Store<T>, __mutations: Mps) => T)

It's important to only specify only what we require in the implementation, for example if withA only cared about get it should specify only that would then make our code compile again as it should...

// https://tsplay.dev/W4y1ew

-       & ((store: Mutate<Store<T>, Mps> & Store<T>, __mutations: Mps) => T)
+       & ((store: Mutate<Store<T>, Mps> & Pick<Store<T>, "get">, __mutations: Mps) => T)

Dependent mutations

At any point if our mutation has a T on a covariant position, it no longer compiles. Basically if the middleware (mutation) "gives" or "returns" the state it means T is on a covariant position (like in case of persist, onHydrate and onFinishHydration "give" the state), it won't work. Here's a simple example...

// https://tsplay.dev/mx6dKw

const withSerialized = f => store => {
  store.getSerialized = () => serialize(store.get())
  // Edit: I realized later this is an stupid example
  // but I'm too lazy to edit it. We'll just pretend
  // `serialize` is an identity function.
  // Perhaps a more realistic but yet simple example
  // would be `getNormalised` or something where the
  // it returns `Normalized</* covariant */ T>`.
  // The realest example is ofc `persist` as I
  // mentioned above.

  return f(store)
}

let storeOut = create(
  withSerialized(
    withA("a",
      storeIn =>
        ({ count: 0 })
    )
  )
)

If you scroll to the end of the huge error you'd find something like WithSerialized<unknown> is not assignable to WithSerialized<never>. The thing for some reason typescript (this could be a bug) resolves T to never instead of unknown assuming it to be contravariant.

I can make it compile by making all signatures bivariant but still storeIn would have WithSerialized<never> instead of WithSerialized<unknown> so that doesn't work. It'd work if T is already concrete either via user explicitly typing it or via a helper. These solutions could work but let's actually solve the problem.

Higher kinded mutations

One thing is clear that we can't have T directly on the mutations, that is to say we can't have { a: A, t: T } as a mutation so instead we'll use an emulated higher kinded type (T, A) => { a: A, t: T } and inject T (and an optional payload like A in this case)...

  // https://tsplay.dev/N9poJw

+ interface StoreMutations<T, A> {}
+ type StoreMutationIdentifier = keyof StoreMutations<unknown, unknown>

  declare const withSerialized:
    < T
-   , Mcs extends {}[] = []
+   , Mcs extends [StoreMutationIdentifier, unknown][] = []
-   , Mps extends {}[] = []
+   , Mps extends [StoreMutationIdentifier, unknown][] = []
    >
      ( initializer:
-         & ( ( store: Mutate<Store<T>, [...Mps, WithSerialized<T>]>
+         & ( ( store: Mutate<Store<T>, [...Mps, ["WithSerialized", never]]>
-             , __mutations: [...Mps, WithSerialized<T>]
+             , __mutations: [...Mps, ["WithSerialized", never]]
              ) => T
            )
          & { __mutations?: Mcs }
      ) =>
        & ((store: Mutate<Store<T>, Mps> __mutations: Mps) => T)
-       & { __mutations?: [WithSerialized<T>, ...Mcs] }
+       & { __mutations?: [["WithSerialized", never], ...Mcs] }
  
+ interface StoreMutations<T, A>
+   { WithSerialized: { getSerialized: () => T }
+   }
  
  declare const withA: 
    < A
    , T
-   , Mcs extends {}[] = []
+   , Mcs extends [StoreMutationIdentifier, unknown][] = []
-   , Mps extends {}[] = []
+   , Mps extends [StoreMutationIdentifier, unknown][] = []
    >
      ( a: A
      , initializer:
-         & ( ( store: Mutate<Store<T>, [...Mps, { a: NoInfer<A> }]>
+         & ( ( store: Mutate<Store<T>, [...Mps, ["WithA", NoInfer<A>]]>
-             , __mutations: [...Mps, { a: NoInfer<A> }]
+             , __mutations: [...Mps, ["WithA", NoInfer<A>]]
              ) => T
            )
          & { __mutations?: Mcs }
      ) =>
        & ((store: Mutate<Store<T>, Mps>, __mutations: Mps) => T)
-       & { __mutations?: [{ a: NoInfer<A> }, ...Mcs] }
+       & { __mutations?: [["WithA", NoInfer<A>], ...Mcs] }
      
+ interface StoreMutations<T, A>
+   { WithA: { a: A }
+   }
  
- type Mutate<S, Ms> =
+ type Mutate<S, Ms, T = S extends { get: () => infer T } ? T : never> =
    Ms extends [] ? S :
-   Ms extends [infer M, ...infer Mr]
+   Ms extends [[infer I, infer A], ...infer Mr]   
      ? Mutate<
          Overwrite<
            S,
-           M
+           StoreMutations<T, A>[I & StoreMutationIdentifier]
          >,
          Mr
        > :
    never

So instead of passing around mutations M we pass "functions" (T, A) => M that take T and some optional payload. So writing...

type M<A> = ["WithFoo", A]

interface StoreMutations<T, A>
  { WithFoo: { t: T, a: A }
  }

...is in spirit same as writing a higher kinded type...

type M<A> = T -> { t: T, a: A }

So StoreMutations is a collection of "mutation creators". And now as you can see in the diff M becomes StoreMutations<T, A>[I & StoreMutationIdentifier]². Writing StoreMutations<T, A>[I] is in spirit writing storeMutations[i](t, a) which would return m.

StoreMutations is also a "global" collection so it can be augmented anywhere and new mutations can be added. It works across files thanks to module augmentation and declaration merging, so in reality adding a mutation would look like this...

// @file node_modules/zustand/index.ts
export interface StoreMutations<T, A> {}

// @file node_modules/zustand-with-serialized/index.ts
declare module "zustand" {
  interface StoreMutations<T, A>
    { WithSerialized: { getSerialized: () => T }
    }
}

Higher kinded mutators

Now that we are using higher kinded types, we can simplify things a bit and instead of newStore = overwrite(oldStore, storeMutations[mi](tFromOldStore, a)) just simply do newStore = storeMutators[mi](oldStore, a).

  // https://tsplay.dev/NrGk2m

- interface StoreMutations<T, A>
-   { WithSerialized: { getSerialized: () => T }
-   }

+ interface StoreMutators<S, A>
+   { WithSerialized:
+       S extends { get: () => infer T }
+         ? Write<S, { getSerialized: () => T }>
+         : never
+   }

- interface StoreMutations<T, A>
-   { WithA: { a: A }
-   }

+ interface StoreMutators<S, A>
+   { WithA: Write<S, { a: A }>
+   }
  
- type Mutate<S, Ms, T = S extends { get: () => infer T } ? T : never> =
+ type Mutate<S, Ms> =
    Ms extends [] ? S :
    Ms extends [[infer Mi, infer A], ...infer Mr]   
      ? Mutate<
-         Write<
-           S,
-           StoreMutations<T, A>[Mi & StoreMutationIdentifier]
-         >,
+         StoreMutators<S, A>[Mi & StoreMutatorIdentifier],
          Mr
        > :
    never

Zustand-adapted version

Here's a version that is compatible with zustand's api (same as in #662), uses unique symbols for identifiers, a reusable StoreInitializer, etc and an example middleware.

To-dos

This is still sort of incomplete here are some things left to figure out...

Caveats

It doesn't compile if you don't write all middlewares in place because it relies on contextual inference, meaning this doesn't compile...

// https://tsplay.dev/w6X96m

let foo = withB("b", () => ({ count: 0 }))
let store = create(withA("a", foo))

... but this does ...

// https://tsplay.dev/WPxDZW

let store = create(withA("a", withB("b", () => ({ count: 0 }))))

I recently noticed this caveat so I haven't really worked on it. Maybe I could make it at least compile will a little less accuracy in types.

Other

• Migration plan? What changes do users need to make?
• createWithState? Probably a good idea.
• If we'll be having createWithState do we still need higher kinded types? Maybe, maybe not.
• Is this pipeable? Can we make a custom pipe that works?
• ...

Action point

I'd like to know if this seems like a good idea or not, if yes then I'll continue with the to-dos. If no then how far should we go wrt to types?

Thanks for reading all this! Hopefully I've not made much mistakes. Feel free to ask questions

Footnotes

1. Note that the diffs through out the document are just gists, they are not exact it excludes some obvious changes elsewhere (like in create) and also renames variables, changes formatting etc. The diff is minimal enough to get the crux of what changed. Go to the playground link to see the full change ↩

2. The & StoreMutationIdentifier is to just satisfy the compiler that I indeed is a key of StoreMutations (as X & Y will always be a subtype of Y). I & StoreMutationIdentifier is same as I because I already is a subtype of StoreMutationIdentifier, eg "WithSerialized" & ("WithSerialized" | "WithA") is "WithSerialized". ↩

[dai-shi]

wow, it's too long. I don't follow the later part.

I like the approach with the phantom mutation type. That's something I wanted to achieve with #601 originally.
tbh, I don't know how this new things work out.
as it sounds, this may introduce too much complexity and still may not solve the developer experience very well.

so, my counter proposal is to live with the "contextual aware function" limitation of TS behavior, and give up propagation mutation type top to bottom. In such cases, developers need to explicitly type store at the top. the question is if the TS error message is fairly readable. otherwise, we'd need a good documentation around this.

what do you think? do you have any other proposal to balance between the complexity and DX?

[devanshj]

I don't follow the later part.

Haha that's okay, I can help you if you tell me which part exactly.

I agree on the complexity part but it's not userland, with createWithState users have to write zero type-annotations (except for the state) and yet their store's type will be 100% correct...

// https://tsplay.dev/WJRdgW
let storeOut = createWithState<{ count: number }>()(
  withA("a",
    withB("b",
      storeIn => ({ count: 0 })
    )
  )
)

// `storeOut` and `storeIn` are `Write<Write<Store<{ count: number }>, { a: string }, { b: string }>>

And there are many problems with the existing phantom mutation type approach that leads to a worse DX compared to this proposal, some examples... (checkout to 46f3cb2 to try them out)

This proposal is very robust, hence these two compile...

And there's nothing to lose with this proposal, even the caveat that I mentioned in the proposal is gone, it works with extracting middlewares too.

// https://tsplay.dev/mpvY7w

let withBBounded = withB("b", storeIn1 => ({ count: 0 }))
let storeOut1 = create(withA("a", withBBounded))

let storeOut2 = create(
  withA("a",
    withB("b",
      storeIn2 => ({ count: 0 })
    )
  )
)

// `storeIn1` is `Write<Store<object>, { b: string }>
// `storeOut1`, `storeIn2`, `storeOut2` are `Write<Write<Store<object>, { a: string }>, { b: string }>>

Here all stores have correct types, the one you'd expect them to have.

So in my eyes there are no downsides to this proposal. The only downside is the added complexity and migration BUT only for the people who author middlewares themselves. I'm not a zustand user but I think there will be few. We can send PRs to popular third-party middlewares (if any). Btw This is how authoring a middleware looks, a bit complex but no rocket science as such.
Those who only use our middlewares or popular third-party ones will have to do nothing, no migration, no visible complexity. We can also provide a temporary compatibility layer to work with v3 middlewares, which we'll remove in v5.

This proposal with createWithState (create is not going anywhere) will enable top-notch DX for users at the cost of slight complexity for authoring middlewares.

If you're not confident about it already how about I go ahead with this, then you see the final piece—the tests, the changes for migration, etc—and if you don't like I'll revert it.

Wdyt?

[dai-shi]

If you're not confident about it already how about I go ahead with this, then you see the final piece—the tests, the changes for migration, etc—and if you don't like I'll revert it.

maybe I misunderstood your nuance. so, if you are comfortable with continuing your approach, you can go on. My concerns are a) if I can understand the final src, b) if you can make the current tests pass, and c) migration path (= breaking changes including subtle ones).

And there are many problems with the existing phantom mutation type approach

don't you call your approach phantom mutation type? or is it phantom mutation type + alpha?

[devanshj]

so, if you are comfortable with continuing your approach, you can go on.

Oh okay cool!

a) if I can understand the final src b) if you can make the current tests pass, and c) migration path (= breaking changes including subtle ones)

You might not understand it fully but I think you'll be able to understand 90% of it, I'll help you out. The tests might not pass as they are of course, but with migration they'll. I'll try to minimize breaking changes and make the migration as easy as possible.

don't you call your approach phantom mutation type? or is it phantom mutation type + alpha?

I don't call it by a certain name haha. But let's call it "higher kinded mutators".

Ahhh looks like when you said "I like phantom mutation type approach" you meant didn't mean the old approach (like in #660, hence I was comparing them) but the overall approach including this proposal, gotcha. That was the misunderstanding I guess haha.

[dai-shi]

but the overall approach including this proposal, gotcha.

yeah.