refactor(withState): simplify the implementation

src/__tests__/withState.test.js

@@ -19,7 +19,7 @@ describe('withState', () => {
     expect(dummy.prop('pass')).toBe('through')
   })
 
-  it('also accepts a non-function, which is passed directly to setState()', () => {
+  it('accepts a non-function, which is passed directly to setState()', () => {
     const Counter = withState('counter', 'updateCounter', 0)(Dummy)
     const dummy = mount(<Counter />).find(Dummy)
     const { updateCounter } = dummy.props()
@@ -28,7 +28,7 @@ describe('withState', () => {
     expect(dummy.prop('counter')).toBe(18)
   })
 
-  it('also accepts initialState as function of props', () => {
+  it('accepts initialState as function of props', () => {
     const Counter = withState(
       'counter',
       'updateCounter',
@@ -43,7 +43,7 @@ describe('withState', () => {
     expect(dummy.prop('counter')).toBe(3)
   })
 
-  it('should be merged with other hoc', () => {
+  it('is merged with other HOCs', () => {
     const Component = compose(
       withProps({ initialCounter: 1 }),
       withState(

src/withState.js

@@ -45,31 +45,26 @@ import updateProps from './utils/updateProps'
  */
 const withState = (stateName, stateUpdaterName, initialState) =>
   updateProps((next) => {
-    let previousProps
-    let previousStateValue
+    let props
+    let state
 
-    const updateState = (nextState) => {
-      update(previousProps, callOrUse(nextState, previousStateValue))
-    }
-
-    const update = (props, stateValue) => {
+    const stateUpdater = (nextState) => {
+      state = callOrUse(nextState, state)
       next({
         ...props,
-        [stateName]: stateValue,
-        [stateUpdaterName]: updateState,
+        [stateName]: state,
+        [stateUpdaterName]: stateUpdater,
       })
-
-      previousStateValue = stateValue
-      previousProps = props
     }
 
-    return (props) => {
-      update(
-        props,
-        !previousProps
-          ? callOrUse(initialState, props)
-          : previousStateValue,
-      )
+    return (nextProps) => {
+      if (!props) state = callOrUse(initialState, nextProps)
+      props = nextProps
+      next({
+        ...props,
+        [stateName]: state,
+        [stateUpdaterName]: stateUpdater,
+      })
     }
   })