Fix subsequent synchronous atom updates are not properly rendered

packages/focal/src/react/react.ts

@@ -312,7 +312,9 @@ class FakeComponent<P> {
     public props: LiftWrapperProps<P>
   ) {}
 
-  setState(newState: LiftWrapperState) {
+  setState(state: (LiftWrapperState | ((state: LiftWrapperState) => LiftWrapperState))) {
+    const newState = typeof state === 'function' ? state(this.state) : state
+
     if ('subscription' in newState)
       this.state.subscription = newState.subscription
     if ('renderCache' in newState)
@@ -390,8 +392,9 @@ class RenderOne<P> implements Subscription {
     const { component, props } = liftedComponent.props
     const renderCache = render(component, props, [value])
 
-    if (!structEq(liftedComponent.state.renderCache, renderCache))
-      liftedComponent.setState({ renderCache })
+    liftedComponent.setState(state =>
+      !structEq(state.renderCache, renderCache) ? { renderCache } : {}
+    )
   }
 
   private _handleCompleted() {
@@ -486,8 +489,9 @@ class RenderMany<P> implements Subscription {
     const { component, props } = liftedComponent.props
     const renderCache = render(component, props, this._values)
 
-    if (!structEq(liftedComponent.state.renderCache, renderCache))
-      liftedComponent.setState({ renderCache })
+    liftedComponent.setState(state =>
+      !structEq(state.renderCache, renderCache) ? { renderCache } : {}
+    )
   }
 
   private _handleCompleted(idx: number) {

packages/test/src/app.tsx

@@ -11,7 +11,8 @@ declare var require: (path: string) => any
 
 const tests: { name: string, test: TestComponent<any> }[] = [
   'lifted-bug',
-  'render-bug'
+  'render-bug',
+  'setstate-bug'
 ].map(name => {
   return { name, test: require(`./${name}/index.tsx`).default as TestComponent<any> }
 })

packages/test/src/setstate-bug/index.tsx

@@ -0,0 +1,74 @@
+import * as React from 'react'
+import { Atom, F } from '@grammarly/focal'
+import { Observable, Subject } from 'rxjs'
+
+class Test extends React.Component<{ trigger: Observable<void> }> {
+  list = Atom.create(['a', 'b', 'c'])
+
+  componentWillMount() {
+    this.props.trigger.subscribe(list => {
+      this.list.set(['a', 'b'])
+      this.list.set(['a', 'b', 'c'])
+    })
+  }
+
+  render() {
+    return (
+      <>
+        <h4>Focal</h4>
+        <F.ul>
+          {this.list.view(list => {
+            return list.map(l => <li key={l}>{l}</li>)
+          })}
+        </F.ul>
+      </>
+    )
+  }
+}
+
+class Test2 extends React.Component<{ trigger: Observable<void> }, { list: string[] }> {
+  constructor(props: any) {
+    super(props)
+    this.state = {
+      list: ['a', 'b', 'c']
+    }
+  }
+
+  componentDidMount() {
+    this.props.trigger.subscribe(() => {
+      this.setState({ list: ['a', 'b'] })
+      this.setState({ list: ['a', 'b', 'c'] })
+    })
+  }
+
+  render() {
+    const { list } = this.state
+    return (
+      <>
+        <h4>React</h4>
+        <ul>
+          {list.map(l => <li key={l}>{l}</li>)}
+        </ul>
+      </>
+    )
+  }
+}
+
+const trigger = new Subject<void>()
+
+function onClick() {
+  trigger.next()
+}
+
+const App = () => (
+  <div>
+    <input type='button' onClick={onClick} value='Click me' />
+    <Test trigger={trigger} />
+    <Test2 trigger={trigger} />
+  </div>
+)
+
+export default {
+  Component: App,
+  defaultState: 0
+}