immer
Version:
Create your next immutable state by mutating the current one
305 lines (273 loc) • 6.84 kB
text/typescript
// types only!
import {
ImmerState,
AnyMap,
AnySet,
MapState,
SetState,
DRAFT_STATE,
getCurrentScope,
latest,
isDraftable,
createProxy,
loadPlugin,
markChanged,
die,
ArchType,
each
} from "../internal"
export function enableMapSet() {
class DraftMap extends Map {
[DRAFT_STATE]: MapState
constructor(target: AnyMap, parent?: ImmerState) {
super()
this[DRAFT_STATE] = {
type_: ArchType.Map,
parent_: parent,
scope_: parent ? parent.scope_ : getCurrentScope()!,
modified_: false,
finalized_: false,
copy_: undefined,
assigned_: undefined,
base_: target,
draft_: this as any,
isManual_: false,
revoked_: false
}
}
get size(): number {
return latest(this[DRAFT_STATE]).size
}
has(key: any): boolean {
return latest(this[DRAFT_STATE]).has(key)
}
set(key: any, value: any) {
const state: MapState = this[DRAFT_STATE]
assertUnrevoked(state)
if (!latest(state).has(key) || latest(state).get(key) !== value) {
prepareMapCopy(state)
markChanged(state)
state.assigned_!.set(key, true)
state.copy_!.set(key, value)
state.assigned_!.set(key, true)
}
return this
}
delete(key: any): boolean {
if (!this.has(key)) {
return false
}
const state: MapState = this[DRAFT_STATE]
assertUnrevoked(state)
prepareMapCopy(state)
markChanged(state)
if (state.base_.has(key)) {
state.assigned_!.set(key, false)
} else {
state.assigned_!.delete(key)
}
state.copy_!.delete(key)
return true
}
clear() {
const state: MapState = this[DRAFT_STATE]
assertUnrevoked(state)
if (latest(state).size) {
prepareMapCopy(state)
markChanged(state)
state.assigned_ = new Map()
each(state.base_, key => {
state.assigned_!.set(key, false)
})
state.copy_!.clear()
}
}
forEach(cb: (value: any, key: any, self: any) => void, thisArg?: any) {
const state: MapState = this[DRAFT_STATE]
latest(state).forEach((_value: any, key: any, _map: any) => {
cb.call(thisArg, this.get(key), key, this)
})
}
get(key: any): any {
const state: MapState = this[DRAFT_STATE]
assertUnrevoked(state)
const value = latest(state).get(key)
if (state.finalized_ || !isDraftable(value)) {
return value
}
if (value !== state.base_.get(key)) {
return value // either already drafted or reassigned
}
// despite what it looks, this creates a draft only once, see above condition
const draft = createProxy(value, state)
prepareMapCopy(state)
state.copy_!.set(key, draft)
return draft
}
keys(): IterableIterator<any> {
return latest(this[DRAFT_STATE]).keys()
}
values(): IterableIterator<any> {
const iterator = this.keys()
return {
[Symbol.iterator]: () => this.values(),
next: () => {
const r = iterator.next()
/* istanbul ignore next */
if (r.done) return r
const value = this.get(r.value)
return {
done: false,
value
}
}
} as any
}
entries(): IterableIterator<[any, any]> {
const iterator = this.keys()
return {
[Symbol.iterator]: () => this.entries(),
next: () => {
const r = iterator.next()
/* istanbul ignore next */
if (r.done) return r
const value = this.get(r.value)
return {
done: false,
value: [r.value, value]
}
}
} as any
}
[Symbol.iterator]() {
return this.entries()
}
}
function proxyMap_<T extends AnyMap>(target: T, parent?: ImmerState): T {
// @ts-ignore
return new DraftMap(target, parent)
}
function prepareMapCopy(state: MapState) {
if (!state.copy_) {
state.assigned_ = new Map()
state.copy_ = new Map(state.base_)
}
}
class DraftSet extends Set {
[DRAFT_STATE]: SetState
constructor(target: AnySet, parent?: ImmerState) {
super()
this[DRAFT_STATE] = {
type_: ArchType.Set,
parent_: parent,
scope_: parent ? parent.scope_ : getCurrentScope()!,
modified_: false,
finalized_: false,
copy_: undefined,
base_: target,
draft_: this,
drafts_: new Map(),
revoked_: false,
isManual_: false
}
}
get size(): number {
return latest(this[DRAFT_STATE]).size
}
has(value: any): boolean {
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
// bit of trickery here, to be able to recognize both the value, and the draft of its value
if (!state.copy_) {
return state.base_.has(value)
}
if (state.copy_.has(value)) return true
if (state.drafts_.has(value) && state.copy_.has(state.drafts_.get(value)))
return true
return false
}
add(value: any): any {
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
if (!this.has(value)) {
prepareSetCopy(state)
markChanged(state)
state.copy_!.add(value)
}
return this
}
delete(value: any): any {
if (!this.has(value)) {
return false
}
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
prepareSetCopy(state)
markChanged(state)
return (
state.copy_!.delete(value) ||
(state.drafts_.has(value)
? state.copy_!.delete(state.drafts_.get(value))
: /* istanbul ignore next */ false)
)
}
clear() {
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
if (latest(state).size) {
prepareSetCopy(state)
markChanged(state)
state.copy_!.clear()
}
}
values(): IterableIterator<any> {
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
prepareSetCopy(state)
return state.copy_!.values()
}
entries(): IterableIterator<[any, any]> {
const state: SetState = this[DRAFT_STATE]
assertUnrevoked(state)
prepareSetCopy(state)
return state.copy_!.entries()
}
keys(): IterableIterator<any> {
return this.values()
}
[Symbol.iterator]() {
return this.values()
}
forEach(cb: any, thisArg?: any) {
const iterator = this.values()
let result = iterator.next()
while (!result.done) {
cb.call(thisArg, result.value, result.value, this)
result = iterator.next()
}
}
}
function proxySet_<T extends AnySet>(target: T, parent?: ImmerState): T {
// @ts-ignore
return new DraftSet(target, parent)
}
function prepareSetCopy(state: SetState) {
if (!state.copy_) {
// create drafts for all entries to preserve insertion order
state.copy_ = new Set()
state.base_.forEach(value => {
if (isDraftable(value)) {
const draft = createProxy(value, state)
state.drafts_.set(value, draft)
state.copy_!.add(draft)
} else {
state.copy_!.add(value)
}
})
}
}
function assertUnrevoked(state: any /*ES5State | MapState | SetState*/) {
if (state.revoked_) die(3, JSON.stringify(latest(state)))
}
loadPlugin("MapSet", {proxyMap_, proxySet_})
}