new an Immer

Preface

Immer combines Copy-on-write mechanism with ES6 Proxy features, providing an exceptionally concise immutable data manipulation method:

const myStructure = {
  a: [1, 2, 3],
  b: 0
};
const copy = produce(myStructure, () => {
  // nothings to do
});
const modified = produce(myStructure, myStructure => {
  myStructure.a.push(4);
  myStructure.b++;
});

copy === myStructure  // true
modified !== myStructure  // true
JSON.stringify(modified) === JSON.stringify({ a: [1, 2, 3, 4], b: 1 })  // true
JSON.stringify(myStructure) === JSON.stringify({ a: [1, 2, 3], b: 0 })  // true

How exactly is this achieved?

I. Goal

Immer has only one core API:

produce(currentState, producer: (draftState) => void): nextState

So, as long as we manually implement an equivalent produce function, we can figure out Immer's secret

II. Ideas

Carefully observing produce's usage, not hard to discover 5 characteristics (see comments):

const myStructure = {
  a: [1, 2, 3],
  b: 0
};
const copy = produce(myStructure, () => {});
const modified = produce(myStructure, myStructure => {
  // 1.Access draftState in producer function, just like accessing original value currentState
  myStructure.a.push(4);
  myStructure.b++;
});

// 2.If not modifying draftState in producer, reference remains unchanged, both point to original value
copy === myStructure  // true
// 3.If draftState is modified, reference changes, produce() returns new value
modified !== myStructure  // true
// 4.Operations on draftState in producer function all apply to new value
JSON.stringify(modified) === JSON.stringify({ a: [1, 2, 3, 4], b: 1 })  // true
// 5.Operations on draftState in producer function don't affect original value
JSON.stringify(myStructure) === JSON.stringify({ a: [1, 2, 3], b: 0 })  // true

That is:

Copy only on write (see comment 2, comment 3)
Read operations are proxied to original value (see comment 1)
Write operations are proxied to copied value (see comment 4, comment 5)

So, simple skeleton has emerged from water:

function produce(currentState, producer) {
  const copy = null;
  const draftState = new Proxy(currentState, {
    get(target, key, receiver) {
      // todo proxy read operations to original value
    },
    set() {
      if (!mutated) {
        mutated = true;
        // todo create copied value
      }
      // todo proxy write operations to copied value
    }
  });
  producer(draftState);
  return copy || currentState;
}

Additionally, since Proxy can only listen to current level property access, proxy relationships also need to be created on demand:

Root node pre-creates a Proxy, all intermediate nodes accessed on object tree (or root nodes of newly added subtrees) need to create corresponding Proxies

And each Proxy only creates copied value when listening to write operations (direct assignment, native data manipulation APIs, etc.) (so-called Copy-on-write), and proxies all subsequent write operations to copied value

Finally, integrate these copied values with original values, get data manipulation result

Therefore, Immer = Copy-on-write + Proxy

III. Specific Implementation

According to above analysis, implementation mainly divided into 3 parts:

Proxy: Create on demand, proxy read/write operations
Copy: Copy on demand (Copy-on-write)
Integrate: Establish association between copied values and original values, deep merge original values with copied values

Proxy

After getting original value, first create Proxy for root node, get draftState for producer to operate:

function produce(original, producer) {
  const draft = proxy(original);
  //...
}

Most critical of course is proxying get, set operations on original value:

function proxy(original, onWrite) {
  // Store proxy relationships and copied values
  let draftState = {
    originalValue: original,
    draftValue: Array.isArray(original) ? [] : Object.create(Object.getPrototypeOf(original)),
    mutated: false,
    onWrite
  };

  // Create root node proxy
  const draft = new Proxy(original, {
    // Read operations (proxy property access)
    get(target, key, receiver) {
      if (typeof original[key] === 'object' && original[key] !== null) {
        // Existing properties not of basic value type, create next level proxy
        return proxyProp(original[key], key, draftState, onWrite);
      }
      else {
        // If modified, get latest state directly from draft
        if (draftState.mutated) {
          return draftValue[key];
        }

        // Non-existent, or existing properties with basic values, proxy to original value
        return Reflect.get(target, key, receiver);
      }
    },

    // Write operations (proxy data modification)
    set(target, key, value) {
      // If new value is not basic value type, create next level proxy
      if (typeof value === 'object') {
        proxyProp(value, key, draftState, onWrite);
      }
      // Copy on first write
      copyOnWrite(draftState);
      // After copying, write directly
      draftValue[key] = value;
      return true;
    }
  });

  return draft;
}

P.S. Additionally, many other read/write methods also need proxying, such as has, ownKeys, deleteProperty, etc., processing methods are similar, won't elaborate here

Copy

This is the copyOnWrite function appeared above:

function copyOnWrite(draftState) {
  const { originalValue, draftValue, mutated, onWrite } = draftState;
  if (!mutated) {
    draftState.mutated = true;
    // Only hang on parent draftValue when next level has modifications
    if (onWrite) {
      onWrite(draftValue);
    }
    // Copy on first write
    copyProps(draftValue, originalValue);
  }
}

Only on first write (!mutated) copy remaining properties on original value to draftValue

Specially, shallow copy needs to note property descriptors, [Symbol](/articles/symbol-es6 笔记 7/) properties and other details:

// Skip properties already on target
function copyProps(target, source) {
  if (Array.isArray(target)) {
    for (let i = 0; i < source.length; i++) {
      // Skip properties already modified at deeper level
      if (!(i in target)) {
        target[i] = source[i];
      }
    }
  }
  else {
    Reflect.ownKeys(source).forEach(key => {
      const desc = Object.getOwnPropertyDescriptor(source, key);
      // Skip existing properties
      if (!(key in target)) {
        Object.defineProperty(target, key, desc);
      }
    });
  }
}

P.S. Reflect.ownKeys can return all property names of object (including Symbol property names and string property names)

Integrate

To integrate copied values with original values, first need to establish two relationships:

Association between proxy and original value, copied value: Root node's proxy needs to bring out result
Association between lower level copied values and ancestor copied values: Copied values need to easily correspond to result tree

For first problem, just need to expose draftState corresponding to proxy object:

const INTERNAL_STATE_KEY = Symbol('state');
function proxy(original, onWrite) {
  let draftState = {
    originalValue: original,
    draftValue,
    mutated: false,
    onWrite
  };
  const draft = new Proxy(original, {
    get(target, key, receiver) {
      // Establish association from proxy to draft value
      if (key === INTERNAL_STATE_KEY) {
        return draftState;
      }
      //...
    }
  }
}

As for second problem, can establish association between lower level copied values and ancestor copied values through onWrite hook:

// Create next level proxy
function proxyProp(propValue, propKey, hostDraftState) {
  const { originalValue, draftValue, onWrite } = hostDraftState;
  // When next level property has write operation
  const onPropWrite = (value) => {
    // Create parent copied value on demand
    if (!draftValue.mutated) {
      hostDraftState.mutated = true;
      // Copy all host properties
      copyProps(draftValue, originalValue);
    }
    // Hang child copied value on it (establish parent-child relationship of copied values)
    draftValue[propKey] = value;
    // Notify ancestors, build complete copied value tree upward
    if (onWrite) {
      onWrite(draftValue);
    }
  };
  return proxy(propValue, onPropWrite);
}

That is to say, when deep properties have first write operation, copy on demand upward, construct copied value tree

At this point, mission accomplished:

function produce(original, producer) {
  const draft = proxy(original);
  // Modify draft
  producer(draft);
  // Take out draft internal state
  const { originalValue, draftValue, mutated } = draft[INTERNAL_STATE_KEY];
  // Patch modified new value on
  const next = mutated ? draftValue : originalValue;
  return next;
}

IV. Online Demo

Given hand-crafted version is slightly more concise than original version, might as well drop an m, call it imer:

Git repo: ayqy/imer
npm package: imer

V. Compare with Immer

Compared with genuine version, there are two differences in implementation approach:

Different way of creating proxy: imer uses new Proxy, immer uses Proxy.revocable()
Different integration approach: imer reverse-builds copied value tree, immer forward-traverses proxy object tree

Proxy created through Proxy.revocable() can revoke proxy relationship, safer

And Immer's forward traversal of proxy object tree is also a quite clever approach:

When the producer finally ends, it will just walk through the proxy tree, and, if a proxy is modified, take the copy; or, if not modified, simply return the original node. This process results in a tree that is structurally shared with the previous state. And that is basically all there is to it.

Much more intuitive than onWrite reverse-building copied value tree, worth learning from

P.S. Additionally, Immer doesn't support Object.defineProperty(), Object.setPrototypeOf() operations, while hand-crafted imer supports all proxy operations