Updated a nested Set in a Map leaks a `DraftSet` in Immer 11

[markerikson]

Per this RTK issue:

• Immer 11 breaking optimistic updates reduxjs/redux-toolkit#5159

A user reported that they had a case where updating a nested Set wasn't working. They linked a repro at https://github.com/keybase/client/pull/28706/files , which has:

import {enableMapSet} from 'immer'
import {create} from 'zustand'
import {immer} from 'zustand/middleware/immer'
enableMapSet()

type Reaction = {
  timestamp: number
  username: string
}
type ReactionDesc = {
  users: Set<Reaction>
}
type Reactions = ReadonlyMap<string, ReactionDesc>

type MinimalMessage = {
  conversationIDKey: string
  ordinal: number
  reactions?: Reactions
  type: 'text'
}

type Store = {
  messageMap: Map<number, MinimalMessage>
}

type State = Store & {
  toggleLocalReaction: (p: {
    decorated: string
    emoji: string
    targetOrdinal: number
    username: string
  }) => void
}

const initialOrdinal = 1
const initialMessage: MinimalMessage = {
  conversationIDKey: 'test-convo',
  ordinal: initialOrdinal,
  type: 'text',
}

const initialMessageMap = new Map<number, MinimalMessage>()
initialMessageMap.set(initialOrdinal, initialMessage)

export const useStore = create<State>()(
  immer((set, get) => ({
    messageMap: initialMessageMap,
    toggleLocalReaction: (p: {decorated: string; emoji: string; targetOrdinal: number; username: string}) => {
      const {emoji, targetOrdinal, username} = p
      set(s => {
        const m = s.messageMap.get(targetOrdinal)
        if (m) {
          const rs = {
            users: m.reactions?.get(emoji)?.users ?? new Set(),
          }
          if (!m.reactions) {
            m.reactions = new Map()
          }
          m.reactions.set(emoji, rs)
          const existing = [...rs.users].find(r => r.username === username)
          if (existing) {
            rs.users.delete(existing)
          }
          rs.users.add({timestamp: Date.now(), username})
        }
      })
    },
  }))
)

After turning that repro into a test, I do see a DraftSet leaking:

// logging
Message:  { id: 2, reactions: Map(1) { '👍' => { users: [DraftSet [Set]] } } }

 FAIL  __tests__/rtk-5159-zustand.ts > RTK #5159 - Zustand Immer Middleware (DEVELOPMENT mode) > Pattern 1: Map/Set Operations (Keybase pattern) > add reaction to message with existing reactions
Error: [Immer] Cannot use a proxy that has been revoked. Did you pass an object from inside an immer function to an async process? {}
 ❯ Module.die src/utils/errors.ts:45:9
     43|   const e = errors[error]
     44|   const msg = isFunction(e) ? e.apply(null, args as any) : e
     45|   throw new Error(`[Immer] ${msg}`)
       |         ^
     46|  }
     47|  throw new Error(
 ❯ assertUnrevoked src/plugins/mapset.ts:365:23
 ❯ DraftSet.values src/plugins/mapset.ts:278:4
 ❯ DraftSet.[Symbol.iterator] src/plugins/mapset.ts:306:16
 ❯ __tests__/rtk-5159-zustand.ts:139:24

[markerikson]

Okay, got this one figured out:

The repro takes an existing DraftSet instance, puts it in a new plain object, and puts the plain object into a DraftMap. That effectively cuts the chain of assignments, and we never tried to fix up the DraftSet in the plain object. We already have a handleCrossReference method that does this, it just wasn't being called inside of DraftMap.set().