T3 interblag real-estate

Typescript: immutable state updates using Lenses

March 13, 2022

React developers will be painfully aware of how annoying it is to change some value deep within a nested object, while not modifying the original (redux anyone?). Though even outside of react, immutability is often desired.

Of course, the common pattern is to shove it all into a one-liner that does way too much. Here's how it tends to go:

const book = {
  title: 'Salt',
  author: 'Mark Kurlansky',
  chapters: [{
    index: 1,
    title: 'A Mandate of Salt',
    pages: [...]
  }]
};

const titleOfChapter1Changed = {
  ...book,
  chapters: book.chapters.map((chapter, i) => {
    if (i !== 0)
      return chapter;

    return {
      ...chapter,
      title: 'Fish, Fowl, and Pharoahs'
    };
  })
};

I'm not saying it couldn't be written better, but I've seen plenty of code like this in the wild. In my opinion it's a definite code smell, and if a lot state updates are happening this way in your codebase it usually points to some design issues within the wider context of the application.

Anyway, there exist programming languages such as Haskell where everything is immutable, so inevitably the community must have encountered this problem before. And of course, the Haskell folk being the Haskell folk, they have thought about everything a million years before everyone else, and they solved it.

With immutable state updates, two people in particular—Twan van Laarhoven and Edward Kmett—completely knocked it out of the park, with their lens concepts!

Lenses are more or less a tuple consisting of a getter and a setter, which come with interesting mathematical properties—but this being javascript, we'll throw all the beautiful theory out the window and bastardize the original idea. Sorry for egging you on.

For this blog post in particular I'd like to focus on the setting part, because that's where all the action is.

Setters

First, let's see if we can't generalize the code snippet above. We can roughly chop the code into three pieces:

This sounds like there are two distinct types of state updates happening here: Updating an object property, and changing an array value at a certain index.

Let's implement these functions.

First, the function that updates an object property. Note that I'm creating a function that returns a function.

const setProp = (prop) => (obj, update) => ({
  ...obj,
  [prop]: update(obj[prop])
});

Example:

const book = {
  title: 'Salt',
  author: 'Mark Kurlansky',
  chapters: [{
    title: 'A Mandate of Salt',
    pages: [...]
  }]
};

const authorSetter = setProp('author');
const updatedBook = authorSetter(book, x => 'Steven Spielberg');

console.log(updatedBook); // {title: 'Salt', author: 'Steven Spielberg', ...};
console.log(book); // {title: 'Salt', author: 'Mark Kurlansky', ...};

Secondly, the function that updates an array value.

const setArrayAtIndex = (index) => (arr, update) => arr.map((x, i) =>
  i === index ? update(x) : x
);

Example:

const numbers = [10, 20, 30, 40];

const secondElementSetter = setArrayAtIndex(1);
const updatedNumbers = secondElementSetter(numbers, x => 300);

console.log(updatedNumbers); // [10, 300, 30, 40]
console.log(numbers); // [10, 20, 30, 40]

Putting setters together

It may not be immediately obvious, but we can chain setters. Since a setter returns the object that went in, more or less, and since it also takes a callback, we can very easily nest them inside each other.

To update the title of the first chapter in our book:

const book = {
  title: 'Salt',
  author: 'Mark Kurlansky',
  chapters: [{
    title: 'A Mandate of Salt',
    pages: [...]
  }]
};

const chaptersSetter = setProp('chapters');
const firstChapterSetter = setArrayAtIndex(0);
const titleSetter = setProp('title');

const updatedBook = chaptersSetter(book, chapters =>
  firstChapterSetter(chapters, chapter =>
    titleSetter(chapter, title =>
      'Fish, Fowl, and Pharoahs'
    )
  )
);

console.log(updatedBook);

This looks kind of neat: We've abstracted two very different functions--one for updating objects, one for updating arrays--into a similar looking interface, namely that of a Setter. A Setter being a function which takes an object and an update function, and returns the updated object. In an immutable fashion!

Though it also can't be denied that this code is reminiscent of 2010's callback hell, which is not really what anyone wants. But, if promises have taught us anything, it's that we can turn callback hell back into imperative greatness.

We can define a function that applies setters, one after the other, using recursion.

const compose = (...setters) => (obj, update) => {
  if (setters.length === 0)
    return update(obj);

  const [setter, ...rest] = setters;  
  return setter(obj, value => compose(...rest)(value, update));
};

This function might be a bit on the abstract side, but lets see it in use:

const setTitleOfFirstChapter = compose(
  setProp('chapters'),
  setArrayAtIndex(0),
  setProp('title')
);

const updatedBook = setTitleOfFirstChapter(book, x => 'Fish, Fowl, and Pharoahs');
console.log(updatedBook);
console.log(book); // Still no changes!

Isn't that cool? Compose takes a list of setters and returns a new setter.

Ramda, and how we can do better

Actually, I'm not the first one to port these ideas from Haskell to Javascript. In particular, Ramda has been around since forever. I've never used it, but as far as I can tell, it takes the lens/setter concept and tries to knead it into a lodash-like library.

In particular it exports a lensPath function that looks extremely similar to what we are doing here:

const firstChapterTitleLens = R.lensPath(['chapters', 0, 'title']);

const updatedBook = R.set(firstChapterTitleLens, 'Fish, Fowl and Pharoahs', book);

Which looks even better than what I came up with, in my opinion. Our code above tries very hard to be in a very functional style, while Ramda's R.lensPath takes what I would call a very "JavaScript-appropriate" approach: Just chuck the properties you would like to look up in succession into an array and call it a day.

But as far as developer experience goes, I claim we can do even better! Using TypeScript's types, and JavaScript proxies, we can actually get some type-safety up in this mother.

Typescript + Proxy :)

Since we are moving to typescript, let me formalize our Book datatype first.

export interface Book {
  title: string;
  author: string;
  chapters: Chapter[];
}

export interface Chapter {
  index: number;
  title: string;
  pages: Page[];
}

export interface Page {
  pagenumber: number;
  paragraphs: Paragraph[];
}

export interface Paragraph {
  words: Word[];
}

export interface Word {
  word: string;
}

Okay.

You know how JavaScript proxies let us capture object property access? If we have a proxy and access a property, like so: myProxy.foo, then the Proxie's get-trap is called with "foo" as an argument.

We can use that to "record" the path which was taken down an object hierarchy into an array, like one which R.lensPath takes as an argument.

export type TLensIndex = string | number | symbol;
export const LensProxy = <T extends object>(path: TLensIndex[] = []): T =>
  new Proxy<T>({} as any as T, {
    get(target, prop, receiver) {
      return LensProxy<any>([...path, prop]);
    }
  });

const bookLens = LensProxy<Book>()
  .chapters[0]
  .title;

The best part of it: Types, and Intellisense support! Since we are fooling typescript into thinking our proxy is of type Book, it will happily autocomplete our lens for us.

Though we need a way to get the recorded path out of our proxy. I chose to do it with a custom Symbol:

export const lensProps = Symbol('lensProps');

export type TLensIndex = string | number | symbol;
export const LensProxy = <T extends object>(path: TLensIndex[] = []): T =>
  new Proxy<T>({} as any as T, {
    get(target, prop, receiver) {
      if (prop === lensProps)
        return path; 

      return LensProxy<any>([...path, prop]);
    }
  });

And now we can inspect the array that was generated by our lens.

const bookLens = LensProxy<Book>()
  .chapters[0]
  .pages[10]
  .pagenumber;

const path = (bookLens as any)[lensProps];

console.log(path); // ["chapters", "0", "pages", "10", "pagenumber"]

we can use this mechanism to neatly interop with Ramda:

const book = {
  title: "Salt",
  author: "Mark Kurlansky",
  chapters: [{
    title: "A Mandate of Salt",
    pages: []
  }]
};

const titlePath = LensProxy<Book>()
  .chapters[0]
  .title;

const lens = R.lensPath((titlePath as any)[lensProps]);
const updatedBook = R.set(lens, "Fish, Fowl and Pharoahs", book);

console.log(updatedBook);

Noice.

We've turned mutable state updates into immutable state updates, and since we are doing it via proxy it doesn't even feel like much has changed at all.

Let me perhaps give a really cool usecase for our LensProxy: An alternative to react reducers.

Reducer? Hardly knew 'er

Since we want to be cool, like Ramda, we'll rename our LensProxy to simply L. I'll also get rid of the ramda dependency and reimplement R.set because bringing in a huge library for a single function is kind of overkill. Then, I'll add a simple react hook... and voila!

const BookView = () => {
  const [book, setBook] = useLens(generateBook());
  const chapterOneTitleLens = L<Book>().chapters[0].title;

  return (
    <>
      <textarea readOnly value={bookToString(book)} />
      <button
        onClick={() => setBook(chapterOneTitleLens, "Fish, Fowl and Pharoahs")}>
        Change Title!
      </button>
    </>
  );

};

render(<BookView />, document.getElementById("root"));

For illustration purposes, I've hacked together a demo over at codesandbox where you can inspect the working code in action.

As you can see this approach might render most reducers unneccessary, since complicated state updates can now happen succinctly inside the component itself.

Of course, some actions in a reducer might require multiple state updates, or would translate to code that's not easy to model with a single lens. So it also wouldn't be a bad approach to use lenses in conjunction with reducers, instead of instead of them.

The elephant in the room: immerjs

Yeeeeeeeeah... I've got nothing. immerjs took the same concepts, but they did it 10 times better.

const nextBook = produce(book, draft => {
  draft.chapters[0].title = "Fish, Fowl and Pharoahs";
});

I mean look at that. It makes immutable state updates look like plain javascript. It's beautiful, really.

So yeah whatever, just go ahead and use immerjs. Sorry for wasting your time, thanks for reading!

-Flo