fix: Collection iterations should yield observable values

.vscode/settings.json

@@ -8,7 +8,6 @@
     "css",
     "go",
     "handlebars",
-    "html",
     "jade",
     "javascriptreact",
     "json",

packages/observer/__tests__/collections/Map.spec.ts

@@ -236,14 +236,70 @@ describe('observer/collections', () => {
       expect(dummy).toBe(2)
     })
 
-    it('should observe iterated values (forEach)', () => {})
+    it('should observe nested values in iterations (forEach)', () => {
+      const map = observable(new Map([[1, { foo: 1 }]]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        map.forEach(value => {
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        })
+      })
+      expect(dummy).toBe(1)
+      ;(map.get(1) as any).foo++
+      expect(dummy).toBe(2)
+    })
 
-    it('should observe iterated values (values)', () => {})
+    it('should observe nested values in iterations (values)', () => {
+      const map = observable(new Map([[1, { foo: 1 }]]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const value of map.values()) {
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      ;(map.get(1) as any).foo++
+      expect(dummy).toBe(2)
+    })
 
-    it('should observe iterated values (keys)', () => {})
+    it('should observe nested values in iterations (entries)', () => {
+      const key = {}
+      const map = observable(new Map([[key, { foo: 1 }]]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const [key, value] of map.entries()) {
+          key
+          expect(isObservable(key)).toBe(true)
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      ;(map.get(key) as any).foo++
+      expect(dummy).toBe(2)
+    })
 
-    it('should observe iterated values (entries)', () => {})
-
-    it('should observe iterated values (for...of)', () => {})
+    it('should observe nested values in iterations (for...of)', () => {
+      const key = {}
+      const map = observable(new Map([[key, { foo: 1 }]]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const [key, value] of map) {
+          key
+          expect(isObservable(key)).toBe(true)
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      ;(map.get(key) as any).foo++
+      expect(dummy).toBe(2)
+    })
   })
 })

packages/observer/__tests__/collections/Set.spec.ts

@@ -286,5 +286,74 @@ describe('observer/collections', () => {
       expect(observed.has(value)).toBe(true)
       expect(set.has(value)).toBe(false)
     })
+
+    it('should observe nested values in iterations (forEach)', () => {
+      const set = observable(new Set([{ foo: 1 }]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        set.forEach(value => {
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        })
+      })
+      expect(dummy).toBe(1)
+      set.forEach(value => {
+        value.foo++
+      })
+      expect(dummy).toBe(2)
+    })
+
+    it('should observe nested values in iterations (values)', () => {
+      const set = observable(new Set([{ foo: 1 }]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const value of set.values()) {
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      set.forEach(value => {
+        value.foo++
+      })
+      expect(dummy).toBe(2)
+    })
+
+    it('should observe nested values in iterations (entries)', () => {
+      const set = observable(new Set([{ foo: 1 }]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const [key, value] of set.entries()) {
+          expect(isObservable(key)).toBe(true)
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      set.forEach(value => {
+        value.foo++
+      })
+      expect(dummy).toBe(2)
+    })
+
+    it('should observe nested values in iterations (for...of)', () => {
+      const set = observable(new Set([{ foo: 1 }]))
+      let dummy: any
+      autorun(() => {
+        dummy = 0
+        for (const value of set) {
+          expect(isObservable(value)).toBe(true)
+          dummy += value.foo
+        }
+      })
+      expect(dummy).toBe(1)
+      set.forEach(value => {
+        value.foo++
+      })
+      expect(dummy).toBe(2)
+    })
   })
 })

packages/observer/__tests__/immutable.spec.ts

@@ -250,6 +250,25 @@ describe('observer/immutable', () => {
         expect(map.get(key)).toBe(1)
         expect(warn).not.toHaveBeenCalled()
       })
+
+      if (Collection === Map) {
+        test('should retrive immutable values on iteration', () => {
+          const key1 = {}
+          const key2 = {}
+          const original = new Collection([[key1, {}], [key2, {}]])
+          const observed = immutable(original)
+          for (const [key, value] of observed) {
+            expect(isImmutable(key)).toBe(true)
+            expect(isImmutable(value)).toBe(true)
+          }
+          observed.forEach((value: any) => {
+            expect(isImmutable(value)).toBe(true)
+          })
+          for (const value of observed.values()) {
+            expect(isImmutable(value)).toBe(true)
+          }
+        })
+      }
     })
   })
 
@@ -299,6 +318,26 @@ describe('observer/immutable', () => {
         expect(set.has(key)).toBe(true)
         expect(warn).not.toHaveBeenCalled()
       })
+
+      if (Collection === Set) {
+        test('should retrive immutable values on iteration', () => {
+          const original = new Collection([{}, {}])
+          const observed = immutable(original)
+          for (const value of observed) {
+            expect(isImmutable(value)).toBe(true)
+          }
+          observed.forEach((value: any) => {
+            expect(isImmutable(value)).toBe(true)
+          })
+          for (const value of observed.values()) {
+            expect(isImmutable(value)).toBe(true)
+          }
+          for (const [v1, v2] of observed.entries()) {
+            expect(isImmutable(v1)).toBe(true)
+            expect(isImmutable(v2)).toBe(true)
+          }
+        })
+      }
     })
   })
 

packages/observer/src/baseHandlers.ts

@@ -11,7 +11,7 @@ const builtInSymbols = new Set(
     .filter(value => typeof value === 'symbol')
 )
 
-function makeGetter(isImmutable: boolean) {
+function createGetter(isImmutable: boolean) {
   return function get(target: any, key: string | symbol, receiver: any) {
     const res = Reflect.get(target, key, receiver)
     if (typeof key === 'symbol' && builtInSymbols.has(key)) {
@@ -86,7 +86,7 @@ function ownKeys(target: any): (string | number | symbol)[] {
 }
 
 export const mutableHandlers: ProxyHandler<any> = {
-  get: makeGetter(false),
+  get: createGetter(false),
   set,
   deleteProperty,
   has,
@@ -94,7 +94,7 @@ export const mutableHandlers: ProxyHandler<any> = {
 }
 
 export const immutableHandlers: ProxyHandler<any> = {
-  get: makeGetter(true),
+  get: createGetter(true),
 
   set(target: any, key: string | symbol, value: any, receiver: any): boolean {
     if (LOCKED) {

packages/observer/src/collectionHandlers.ts

@@ -3,13 +3,18 @@ import { track, trigger } from './autorun'
 import { OperationTypes } from './operations'
 import { LOCKED } from './lock'
 
-function get(target: any, key: any, toObservable: (t: any) => any): any {
+const isObject = (value: any) => value !== null && typeof value === 'object'
+const toObservable = (value: any) =>
+  isObject(value) ? observable(value) : value
+const toImmutable = (value: any) => (isObject(value) ? immutable(value) : value)
+
+function get(target: any, key: any, wrap: (t: any) => any): any {
   target = unwrap(target)
   key = unwrap(key)
   const proto: any = Reflect.getPrototypeOf(target)
   track(target, OperationTypes.GET, key)
   const res = proto.get.call(target, key)
-  return res !== null && typeof res === 'object' ? toObservable(res) : res
+  return wrap(res)
 }
 
 function has(key: any): boolean {
@@ -107,7 +112,58 @@ function clear() {
   return result
 }
 
-function makeImmutableMethod(method: Function, type: OperationTypes): Function {
+function createForEach(isImmutable: boolean) {
+  return function forEach(callback: Function, thisArg?: any) {
+    const observed = this
+    const target = unwrap(observed)
+    const proto: any = Reflect.getPrototypeOf(target)
+    const wrap = isImmutable ? toImmutable : toObservable
+    track(target, OperationTypes.ITERATE)
+    // important: create sure the callback is
+    // 1. invoked with the observable map as `this` and 3rd arg
+    // 2. the value received should be a corresponding observable/immutable.
+    function wrappedCallback(value: any, key: any) {
+      return callback.call(observed, wrap(value), wrap(key), observed)
+    }
+    return proto.forEach.call(target, wrappedCallback, thisArg)
+  }
+}
+
+function createIterableMethod(method: string | symbol, isImmutable: boolean) {
+  return function(...args: any[]) {
+    const target = unwrap(this)
+    const proto: any = Reflect.getPrototypeOf(target)
+    const isPair =
+      method === 'entries' ||
+      (method === Symbol.iterator && target instanceof Map)
+    const innerIterator = proto[method].apply(target, args)
+    const wrap = isImmutable ? toImmutable : toObservable
+    track(target, OperationTypes.ITERATE)
+    // return a wrapped iterator which returns observed versions of the
+    // values emitted from the real iterator
+    return {
+      // iterator protocol
+      next() {
+        const { value, done } = innerIterator.next()
+        return done
+          ? { value, done }
+          : {
+              value: isPair ? [wrap(value[0]), wrap(value[1])] : wrap(value),
+              done
+            }
+      },
+      // iterable protocol
+      [Symbol.iterator]() {
+        return this
+      }
+    }
+  }
+}
+
+function createImmutableMethod(
+  method: Function,
+  type: OperationTypes
+): Function {
   return function(...args: any[]) {
     if (LOCKED) {
       if (__DEV__) {
@@ -126,7 +182,7 @@ function makeImmutableMethod(method: Function, type: OperationTypes): Function {
 
 const mutableInstrumentations: any = {
   get(key: any) {
-    return get(this, key, observable)
+    return get(this, key, toObservable)
   },
   get size() {
     return size(this)
@@ -135,49 +191,49 @@ const mutableInstrumentations: any = {
   add,
   set,
   delete: deleteEntry,
-  clear
+  clear,
+  forEach: createForEach(false)
 }
 
 const immutableInstrumentations: any = {
   get(key: any) {
-    return get(this, key, immutable)
+    return get(this, key, toImmutable)
   },
   get size() {
     return size(this)
   },
   has,
-  add: makeImmutableMethod(add, OperationTypes.ADD),
-  set: makeImmutableMethod(set, OperationTypes.SET),
-  delete: makeImmutableMethod(deleteEntry, OperationTypes.DELETE),
-  clear: makeImmutableMethod(clear, OperationTypes.CLEAR)
+  add: createImmutableMethod(add, OperationTypes.ADD),
+  set: createImmutableMethod(set, OperationTypes.SET),
+  delete: createImmutableMethod(deleteEntry, OperationTypes.DELETE),
+  clear: createImmutableMethod(clear, OperationTypes.CLEAR),
+  forEach: createForEach(true)
 }
-;['forEach', 'keys', 'values', 'entries', Symbol.iterator].forEach(method => {
-  mutableInstrumentations[method] = immutableInstrumentations[
-    method
-  ] = function(...args: any[]) {
-    const target = unwrap(this)
-    const proto: any = Reflect.getPrototypeOf(target)
-    track(target, OperationTypes.ITERATE)
-    // TODO values retrived from iterations should also be observables
-    return proto[method].apply(target, args)
-  }
+
+const iteratorMethods = ['keys', 'values', 'entries', Symbol.iterator]
+iteratorMethods.forEach(method => {
+  mutableInstrumentations[method] = createIterableMethod(method, false)
+  immutableInstrumentations[method] = createIterableMethod(method, true)
 })
 
-function makeInstrumentationGetter(instrumentations: any) {
+function createInstrumentationGetter(instrumentations: any) {
   return function getInstrumented(
     target: any,
     key: string | symbol,
     receiver: any
   ) {
-    target = instrumentations.hasOwnProperty(key) ? instrumentations : target
+    target =
+      instrumentations.hasOwnProperty(key) && key in target
+        ? instrumentations
+        : target
     return Reflect.get(target, key, receiver)
   }
 }
 
 export const mutableCollectionHandlers: ProxyHandler<any> = {
-  get: makeInstrumentationGetter(mutableInstrumentations)
+  get: createInstrumentationGetter(mutableInstrumentations)
 }
 
 export const immutableCollectionHandlers: ProxyHandler<any> = {
-  get: makeInstrumentationGetter(immutableInstrumentations)
+  get: createInstrumentationGetter(immutableInstrumentations)
 }