Implement the first livecycle compoent: LiveData

This class allows observing an object for the lifetime of a
lifecycle provider.

Test: LiveDataTest
Bug: 32342385

Change-Id: I392e1f6da5b9363317f16d1313d1167d53dc18a3
diff --git a/lifecycle/extensions/build.gradle b/lifecycle/extensions/build.gradle
index 7b32435..1bfa9e4 100644
--- a/lifecycle/extensions/build.gradle
+++ b/lifecycle/extensions/build.gradle
@@ -19,9 +19,13 @@
             proguardFiles getDefaultProguardFile('proguard-android.txt'), 'proguard-rules.pro'
         }
     }
-}
 
+    testOptions {
+        unitTests.returnDefaultValues = true
+    }
+}
 dependencies {
     compile project(":runtime")
     testCompile 'junit:junit:4.12'
-}
+    testCompile "org.mockito:mockito-core:1.9.5"
+}
\ No newline at end of file
diff --git a/lifecycle/extensions/src/main/java/com/android/support/lifecycle/LiveData.java b/lifecycle/extensions/src/main/java/com/android/support/lifecycle/LiveData.java
new file mode 100644
index 0000000..d95796f
--- /dev/null
+++ b/lifecycle/extensions/src/main/java/com/android/support/lifecycle/LiveData.java
@@ -0,0 +1,304 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.support.lifecycle;
+
+import static com.android.support.lifecycle.Lifecycle.DESTROYED;
+import static com.android.support.lifecycle.Lifecycle.STARTED;
+
+import android.support.annotation.MainThread;
+import android.support.annotation.Nullable;
+import android.support.annotation.VisibleForTesting;
+
+/**
+ * LiveData is a data reference that can be observed withing a given lifecycle.
+ * <p>
+ * The Observers of LiveData must specify their LifecycleProvider, which allows LiveData to observe
+ * the provider's state changes and unsubscribe the observer when necessary.
+ *
+ * @param <T> The type of data hold by this instance
+ */
+@SuppressWarnings({"WeakerAccess", "unused"})
+public class LiveData<T> {
+    private static final int START_VERSION = -1;
+    private static final Object NOT_SET = new Object();
+
+    @VisibleForTesting
+    ObserverSet<LifecycleBoundObserver> mObservers =
+            new ObserverSet<LifecycleBoundObserver>() {
+                @Override
+                protected boolean checkEquality(LifecycleBoundObserver existing,
+                        LifecycleBoundObserver added) {
+                    if (existing.observer == added.observer) {
+                        if (existing.provider != added.provider) {
+                            throw new IllegalArgumentException("Cannot add the same observer twice"
+                                    + " to the LiveData");
+                        }
+                        return true;
+                    }
+                    return false;
+                }
+
+                @Override
+                protected void onAdded(LifecycleBoundObserver observer) {
+                    observer.onAdded();
+                    mObserverCount++;
+                    if (mObserverCount == 1) {
+                        onHasObserversChanged(true);
+                    }
+                    if (observer.active) {
+                        mActiveCount++;
+                        if (mActiveCount == 1) {
+                            onHasActiveObserversChanged(true);
+                        }
+                    }
+                    if (mData != NOT_SET) {
+                        //noinspection unchecked
+                        observer.considerNotify((T) mData, mVersion);
+                    }
+                }
+
+                @Override
+                protected void onRemoved(LifecycleBoundObserver observer) {
+                    if (observer.active) {
+                        mActiveCount--;
+                        if (mActiveCount == 0) {
+                            onHasActiveObserversChanged(false);
+                        }
+                    }
+                    mObserverCount--;
+                    if (mObserverCount == 0) {
+                        onHasObserversChanged(false);
+                    }
+                    observer.onRemoved();
+                }
+            };
+
+    // how many observers are in active state
+    private int mActiveCount = 0;
+    // how many observers do we have
+    private int mObserverCount = 0;
+
+    private Object mData = NOT_SET;
+    private int mVersion = START_VERSION;
+
+    private ObserverSet.Callback<LifecycleBoundObserver> mDispatchCallback =
+            new ObserverSet.Callback<LifecycleBoundObserver>() {
+                @Override
+                public void run(LifecycleBoundObserver observer) {
+                    //noinspection unchecked
+                    observer.considerNotify((T) mData, mVersion);
+                }
+            };
+
+    /**
+     * Adds the given observer to the observers list within the lifespan of the given provider. The
+     * events are dispatched on the main thread. If LiveData already has data set, it is instantly
+     * delivered to the observer before this call returns.
+     * <p>
+     * The observer will only receive events if the provider is in {@link Lifecycle#STARTED} or
+     * {@link Lifecycle#RESUMED} state (active).
+     * <p>
+     * If the provider moves to the {@link Lifecycle#DESTROYED} state, the observer will
+     * automatically be removed.
+     * <p>
+     * When data changes while the {@code provider} is not active, it will not receive any updates.
+     * If it becomes active again, it will receive the last available data automatically.
+     * <p>
+     * LiveData keeps a strong reference to the observer and the provider as long as the given
+     * LifecycleProvider is not destroyed. When it is destroyed, LiveData removes references to
+     * the observer & the provider.
+     * <p>
+     * If the given provider is already in {@link Lifecycle#DESTROYED} state, LiveData ignores the
+     * call.
+     * <p>
+     * If the given provider, observer tuple is already in the list, the call is ignored.
+     * If the observer is already in the list with another provider, LiveData throws an
+     * {@link IllegalArgumentException}.
+     *
+     * @param provider The LifecycleProvider which controls the observer
+     * @param observer The observer that will receive the events
+     */
+    @MainThread
+    public void observe(LifecycleProvider provider, Observer<T> observer) {
+        if (provider.getLifecycle().getCurrentState() == DESTROYED) {
+            // ignore
+            return;
+        }
+        final LifecycleBoundObserver wrapper = new LifecycleBoundObserver(provider, observer);
+        mObservers.add(wrapper);
+    }
+
+    /**
+     * Removes the given observer from the observers list.
+     *
+     * @param observer The Observer to receive events.
+     */
+    @MainThread
+    public void removeObserver(final Observer<T> observer) {
+        // TODO make it efficient
+        mObservers.forEach(new ObserverSet.Callback<LifecycleBoundObserver>() {
+            @Override
+            public void run(LifecycleBoundObserver key) {
+                if (key.observer == observer) {
+                    mObservers.remove(key);
+                }
+            }
+        });
+    }
+
+    /**
+     * Removes all observers that are tied to the given LifecycleProvider.
+     *
+     * @param provider The provider scope for the observers to be removed.
+     */
+    @MainThread
+    public void removeObservers(final LifecycleProvider provider) {
+        // TODO make it efficient
+        mObservers.forEach(new ObserverSet.Callback<LifecycleBoundObserver>() {
+            @Override
+            public void run(LifecycleBoundObserver key) {
+                if (key.provider == provider) {
+                    mObservers.remove(key);
+                }
+            }
+        });
+    }
+
+    /**
+     * Sets the value. If there are active observers, the value will be dispatched to them.
+     *
+     * @param value The new value
+     */
+    @MainThread
+    public void setValue(T value) {
+        mVersion++;
+        mData = value;
+        mObservers.forEach(mDispatchCallback);
+    }
+
+    /**
+     * Returns the current value.
+     * Note that calling this method on a background thread does not guarantee that the latest
+     * value set will be received.
+     *
+     * @return the current value
+     */
+    @Nullable
+    public T getValue() {
+        Object data = mData;
+        if (mData != NOT_SET) {
+            //noinspection unchecked
+            return (T) mData;
+        }
+        return null;
+    }
+
+    /**
+     * Called when the number of active observers change between 0 and 1.
+     * <p>
+     * If your LiveData observes another resource (e.g. {@link android.hardware.SensorManager},
+     * this should be the place where you enable / disable that observability.
+     *
+     * @param hasActiveObservers True if there are active observers, false otherwise.
+     */
+    protected void onHasActiveObserversChanged(boolean hasActiveObservers) {
+    }
+
+    /**
+     * Called when the number of observers change between 0 and 1.
+     * <p>
+     * Since there are no observers on this LiveData, this might be a good place to clear it from
+     * its owner object.
+     *
+     * @param hasObservers True if there are 1 or more observers, false otherwise.
+     */
+    protected void onHasObserversChanged(boolean hasObservers) {
+    }
+
+
+    class LifecycleBoundObserver implements LifecycleObserver {
+        public final LifecycleProvider provider;
+        public final Observer<T> observer;
+        public boolean active;
+        public int lastVersion = START_VERSION;
+
+        LifecycleBoundObserver(LifecycleProvider provider, Observer<T> observer) {
+            this.provider = provider;
+            this.observer = observer;
+        }
+
+        private void onAdded() {
+            active = isActiveState(provider.getLifecycle().getCurrentState());
+            provider.getLifecycle().addObserver(this);
+        }
+
+        public void onRemoved() {
+            provider.getLifecycle().removeObserver(this);
+        }
+
+        void considerNotify(T data, int version) {
+            if (!active) {
+                return;
+            }
+            if (lastVersion >= version) {
+                return;
+            }
+            lastVersion = version;
+            observer.onChanged(data);
+        }
+
+        @SuppressWarnings("unused")
+        @OnLifecycleEvent(Lifecycle.ANY)
+        void onStateChange() {
+            if (provider.getLifecycle().getCurrentState() == DESTROYED) {
+                removeInternal(this);
+                return;
+            }
+            boolean activeNow = isActiveState(provider.getLifecycle().getCurrentState());
+            if (active != activeNow) {
+                active = activeNow;
+                handleActiveStateChange(this);
+            }
+        }
+    }
+
+    private void removeInternal(LifecycleBoundObserver observer) {
+        mObservers.remove(observer);
+    }
+
+    @SuppressWarnings("unchecked")
+    private void handleActiveStateChange(LifecycleBoundObserver observer) {
+        if (observer.active) {
+            mActiveCount++;
+            if (mActiveCount == 1) {
+                onHasActiveObserversChanged(true);
+            }
+            if (mData != NOT_SET) {
+                observer.considerNotify((T) mData, mVersion);
+            }
+        } else {
+            mActiveCount--;
+            if (mActiveCount == 0) {
+                onHasActiveObserversChanged(false);
+            }
+        }
+    }
+
+    static boolean isActiveState(@Lifecycle.State int state) {
+        return state >= STARTED;
+    }
+}
diff --git a/lifecycle/extensions/src/main/java/com/android/support/lifecycle/Observer.java b/lifecycle/extensions/src/main/java/com/android/support/lifecycle/Observer.java
new file mode 100644
index 0000000..76d8e48
--- /dev/null
+++ b/lifecycle/extensions/src/main/java/com/android/support/lifecycle/Observer.java
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.support.lifecycle;
+
+import android.support.annotation.Nullable;
+
+/**
+ * A simple callback that can receive updates from an Observable.
+ *
+ * @param <T> The type of the parameter
+ */
+public interface Observer<T> {
+    /**
+     * Called when the data is changed.
+     * @param t  The new data
+     */
+    void onChanged(@Nullable T t);
+}
diff --git a/lifecycle/extensions/src/test/java/com/android/support/lifecycle/LiveDataTest.java b/lifecycle/extensions/src/test/java/com/android/support/lifecycle/LiveDataTest.java
new file mode 100644
index 0000000..0ef8a97
--- /dev/null
+++ b/lifecycle/extensions/src/test/java/com/android/support/lifecycle/LiveDataTest.java
@@ -0,0 +1,255 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.support.lifecycle;
+
+import static com.android.support.lifecycle.Lifecycle.ON_CREATE;
+import static com.android.support.lifecycle.Lifecycle.ON_DESTROY;
+import static com.android.support.lifecycle.Lifecycle.ON_START;
+import static com.android.support.lifecycle.Lifecycle.ON_STOP;
+
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import org.junit.Before;
+import org.junit.Test;
+
+@SuppressWarnings({"unchecked"})
+public class LiveDataTest {
+    private PublicLiveData<String> mLiveData;
+    private LifecycleProvider mProvider;
+    private LifecycleRegistry mRegistry;
+    private MethodExec mActiveObserversChanged;
+    private MethodExec mObserversChanged;
+
+    @Before
+    public void init() {
+        mLiveData = new PublicLiveData<>();
+        mProvider = mock(LifecycleProvider.class);
+        mRegistry = new LifecycleRegistry(mProvider);
+        when(mProvider.getLifecycle()).thenReturn(mRegistry);
+        mActiveObserversChanged = mock(MethodExec.class);
+        mObserversChanged = mock(MethodExec.class);
+        mLiveData.activeObserversChanged = mActiveObserversChanged;
+        mLiveData.observersChanged = mObserversChanged;
+    }
+
+    @Test
+    public void testObserverToggle() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mLiveData.observe(mProvider, observer);
+
+        verify(mActiveObserversChanged, never()).onCall(anyBoolean());
+        verify(mObserversChanged).onCall(true);
+        reset(mObserversChanged, mActiveObserversChanged);
+
+        mLiveData.removeObserver(observer);
+        verify(mActiveObserversChanged, never()).onCall(anyBoolean());
+        verify(mObserversChanged).onCall(false);
+        verifyNoMoreInteractions(mActiveObserversChanged, mObserversChanged);
+    }
+
+    @Test
+    public void testActiveObserverToggle() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mLiveData.observe(mProvider, observer);
+
+        verify(mActiveObserversChanged, never()).onCall(anyBoolean());
+        verify(mObserversChanged).onCall(true);
+
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(mActiveObserversChanged).onCall(true);
+        reset(mObserversChanged, mActiveObserversChanged);
+
+        mRegistry.handleLifecycleEvent(ON_STOP);
+        verify(mActiveObserversChanged).onCall(false);
+        verify(mObserversChanged, never()).onCall(anyBoolean());
+
+        reset(mObserversChanged, mActiveObserversChanged);
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(mActiveObserversChanged).onCall(true);
+        verify(mObserversChanged, never()).onCall(anyBoolean());
+
+        reset(mObserversChanged, mActiveObserversChanged);
+        mLiveData.removeObserver(observer);
+        verify(mActiveObserversChanged).onCall(false);
+        verify(mObserversChanged).onCall(false);
+
+        verifyNoMoreInteractions(mActiveObserversChanged, mObserversChanged);
+    }
+
+    @Test
+    public void testReAddSameObserverTuple() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mLiveData.observe(mProvider, observer);
+        mLiveData.observe(mProvider, observer);
+        assertThat(mLiveData.mObservers.size(), is(1));
+    }
+
+    @Test
+    public void testAdd2ObserversWithSameProviderAndRemove() {
+        Observer<String> o1 = (Observer<String>) mock(Observer.class);
+        Observer<String> o2 = (Observer<String>) mock(Observer.class);
+        mLiveData.observe(mProvider, o1);
+        mLiveData.observe(mProvider, o2);
+        assertThat(mLiveData.mObservers.size(), is(2));
+        verify(mActiveObserversChanged, never()).onCall(anyBoolean());
+        verify(mObserversChanged).onCall(true);
+
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(mActiveObserversChanged).onCall(true);
+        mLiveData.setValue("a");
+        verify(o1).onChanged("a");
+        verify(o2).onChanged("a");
+
+        reset(mObserversChanged);
+        mLiveData.removeObservers(mProvider);
+
+        verify(mObserversChanged).onCall(false);
+        assertThat(mLiveData.mObservers.size(), is(0));
+        assertThat(mRegistry.size(), is(0));
+    }
+
+    @Test
+    public void testAddSameObserverIn2LifecycleProviders() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        LifecycleProvider provider2 = mock(LifecycleProvider.class);
+        LifecycleRegistry registry2 = new LifecycleRegistry(provider2);
+        when(provider2.getLifecycle()).thenReturn(registry2);
+
+        mLiveData.observe(mProvider, observer);
+        Throwable throwable = null;
+        try {
+            mLiveData.observe(provider2, observer);
+        } catch (Throwable t) {
+            throwable = t;
+        }
+        assertThat(throwable, instanceOf(IllegalArgumentException.class));
+        assertThat(throwable.getMessage(),
+                is("Cannot add the same observer twice to the LiveData"));
+    }
+
+    @Test
+    public void testRemoveDestroyedObserver() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mLiveData.observe(mProvider, observer);
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(mObserversChanged).onCall(true);
+        verify(mActiveObserversChanged).onCall(true);
+        assertThat(mLiveData.mObservers.size(), is(1));
+        reset(mObserversChanged, mActiveObserversChanged);
+
+        mRegistry.handleLifecycleEvent(ON_DESTROY);
+        assertThat(mLiveData.mObservers.size(), is(0));
+        verify(mObserversChanged).onCall(false);
+        verify(mActiveObserversChanged).onCall(false);
+    }
+
+    @Test
+    public void testInactiveRegistry() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mRegistry.handleLifecycleEvent(ON_DESTROY);
+        mLiveData.observe(mProvider, observer);
+        assertThat(mLiveData.mObservers.size(), is(0));
+    }
+
+    @Test
+    public void testNotifyActiveInactive() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mRegistry.handleLifecycleEvent(ON_CREATE);
+        mLiveData.observe(mProvider, observer);
+        mLiveData.setValue("a");
+        verify(observer, never()).onChanged(anyString());
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(observer).onChanged("a");
+
+        mLiveData.setValue("b");
+        verify(observer).onChanged("b");
+
+        mRegistry.handleLifecycleEvent(ON_STOP);
+        mLiveData.setValue("c");
+        verify(observer, never()).onChanged("c");
+
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(observer).onChanged("c");
+
+        reset(observer);
+        mRegistry.handleLifecycleEvent(ON_STOP);
+        mRegistry.handleLifecycleEvent(ON_START);
+        verify(observer, never()).onChanged(anyString());
+    }
+
+    @Test
+    public void testStopObservingProvider_onDestroy() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mRegistry.handleLifecycleEvent(ON_CREATE);
+        mLiveData.observe(mProvider, observer);
+        assertThat(mRegistry.size(), is(1));
+        mRegistry.handleLifecycleEvent(ON_DESTROY);
+        assertThat(mRegistry.size(), is(0));
+    }
+
+    @Test
+    public void testStopObservingProvider_onStopObserving() {
+        Observer<String> observer = (Observer<String>) mock(Observer.class);
+        mRegistry.handleLifecycleEvent(ON_CREATE);
+        mLiveData.observe(mProvider, observer);
+        assertThat(mRegistry.size(), is(1));
+
+        mLiveData.removeObserver(observer);
+        assertThat(mRegistry.size(), is(0));
+    }
+
+    @SuppressWarnings("WeakerAccess")
+    static class PublicLiveData<T> extends LiveData<T> {
+        // cannot spy due to internal calls
+        public MethodExec activeObserversChanged;
+        public MethodExec observersChanged;
+        @Override
+        public void setValue(T value) {
+            super.setValue(value);
+        }
+
+        @Override
+        public void onHasActiveObserversChanged(boolean hasActiveObservers) {
+            super.onHasActiveObserversChanged(hasActiveObservers);
+            if (activeObserversChanged != null) {
+                activeObserversChanged.onCall(hasActiveObservers);
+            }
+        }
+
+        @Override
+        public void onHasObserversChanged(boolean hasObservers) {
+            super.onHasObserversChanged(hasObservers);
+            if (observersChanged != null) {
+                observersChanged.onCall(hasObservers);
+            }
+        }
+    }
+
+    interface MethodExec {
+        void onCall(boolean value);
+    }
+}
diff --git a/lifecycle/runtime/src/main/java/com/android/support/lifecycle/LifecycleRegistry.java b/lifecycle/runtime/src/main/java/com/android/support/lifecycle/LifecycleRegistry.java
index 9360b7e..7ceb1a9 100644
--- a/lifecycle/runtime/src/main/java/com/android/support/lifecycle/LifecycleRegistry.java
+++ b/lifecycle/runtime/src/main/java/com/android/support/lifecycle/LifecycleRegistry.java
@@ -17,6 +17,7 @@
 package com.android.support.lifecycle;
 
 import android.support.annotation.NonNull;
+import android.support.v4.util.Pair;
 
 /**
  * An implementation of {@link Lifecycle} that can handle multiple observers.
@@ -30,7 +31,18 @@
     /**
      * Custom list that keeps observers and can handle removals / additions during traversal.
      */
-    private ObserverList mObserverList = new ObserverList();
+    private ObserverSet<Pair<LifecycleObserver, GenericLifecycleObserver>> mObserverSet =
+            new ObserverSet<Pair<LifecycleObserver, GenericLifecycleObserver>>() {
+                @Override
+                protected boolean checkEquality(Pair<LifecycleObserver, GenericLifecycleObserver>
+                        existing, Pair<LifecycleObserver, GenericLifecycleObserver> added) {
+                    return existing.first == added.first;
+                }
+            };
+
+    /**
+     * Current state
+     */
     @State
     private int mState;
     /**
@@ -44,12 +56,15 @@
      */
     private final LifecycleProvider mLifecycleProvider;
 
-    private final ObserverList.Callback mDispatchCallback = new ObserverList.Callback() {
-        @Override
-        public void run(GenericLifecycleObserver observer) {
-            observer.onStateChanged(mLifecycleProvider, mLastEvent);
-        }
-    };
+
+    private final ObserverSet.Callback<Pair<LifecycleObserver, GenericLifecycleObserver>>
+            mDispatchCallback =
+            new ObserverSet.Callback<Pair<LifecycleObserver, GenericLifecycleObserver>>() {
+                @Override
+                public void run(Pair<LifecycleObserver, GenericLifecycleObserver> pair) {
+                    pair.second.onStateChanged(mLifecycleProvider, mLastEvent);
+                }
+            };
 
     /**
      * Creates a new LifecycleRegistry for the given provider.
@@ -57,7 +72,7 @@
      * You should usually create this inside your LifecycleProvider class's constructor and hold
      * onto the same instance.
      *
-     * @param provider     The owner LifecycleProvider
+     * @param provider The owner LifecycleProvider
      */
     public LifecycleRegistry(@NonNull LifecycleProvider provider) {
         mLifecycleProvider = provider;
@@ -79,17 +94,17 @@
         mLastEvent = event;
         // TODO fake intermediate events
         mState = getStateAfter(event);
-        mObserverList.forEach(mDispatchCallback);
+        mObserverSet.forEach(mDispatchCallback);
     }
 
     @Override
     public void addObserver(LifecycleObserver observer) {
-        mObserverList.add(observer);
+        mObserverSet.add(new Pair<>(observer, Lifecycling.getCallback(observer)));
     }
 
     @Override
     public void removeObserver(LifecycleObserver observer) {
-        mObserverList.remove(observer);
+        mObserverSet.remove(new Pair<LifecycleObserver, GenericLifecycleObserver>(observer, null));
     }
 
     /**
@@ -98,7 +113,7 @@
      * @return The number of observers.
      */
     public int size() {
-        return mObserverList.size();
+        return mObserverSet.size();
     }
 
     @Override
diff --git a/lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverList.java b/lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverSet.java
similarity index 67%
rename from lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverList.java
rename to lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverSet.java
index e1e8d5a..f2d3d56 100644
--- a/lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverList.java
+++ b/lifecycle/runtime/src/main/java/com/android/support/lifecycle/ObserverSet.java
@@ -23,40 +23,46 @@
 import java.util.Iterator;
 import java.util.List;
 
-class ObserverList {
+/**
+ * A set class that can keep a list of items and allows modifications while traversing.
+ * It is NOT thread safe.
+ * @param <T> Item type
+ */
+abstract class ObserverSet<T> {
     private boolean mLocked = false;
     private static final Object TO_BE_ADDED = new Object();
     private static final Object TO_BE_REMOVED = new Object();
-    private List<Pair<LifecycleObserver, Object>> mPendingModifications = new ArrayList<>(0);
-    private List<Pair<LifecycleObserver, GenericLifecycleObserver>> mData = new ArrayList<>(5);
+    private List<Pair<T, Object>> mPendingModifications = new ArrayList<>(0);
+    private List<T> mData = new ArrayList<>(5);
 
     private static final String ERR_RE_ADD = "Trying to re-add an already existing observer. "
             + "Ignoring the call for ";
     private static final String WARN_RE_REMOVE = "Trying to remove a non-existing observer. ";
 
-    private boolean exists(LifecycleObserver observer) {
+    private boolean exists(T observer) {
         final int size = mData.size();
         for (int i = 0; i < size; i++) {
-            Pair<LifecycleObserver, GenericLifecycleObserver> pair = mData.get(i);
-            if (pair.first == observer) {
+            if (checkEquality(mData.get(i), observer)) {
                 return true;
             }
         }
         return false;
     }
 
-    private boolean hasPendingRemoval(LifecycleObserver observer) {
+    protected abstract boolean checkEquality(T existing, T added);
+
+    private boolean hasPendingRemoval(T observer) {
         final int size = mPendingModifications.size();
         for (int i = 0; i < size; i++) {
-            Pair<LifecycleObserver, Object> pair = mPendingModifications.get(i);
-            if (pair.first == observer && pair.second == TO_BE_REMOVED) {
+            Pair<T, Object> pair = mPendingModifications.get(i);
+            if (checkEquality(pair.first, observer) && pair.second == TO_BE_REMOVED) {
                 return true;
             }
         }
         return false;
     }
 
-    void add(LifecycleObserver observer) {
+    void add(T observer) {
         if (!mLocked) {
             addInternal(observer);
             return;
@@ -64,16 +70,20 @@
         mPendingModifications.add(new Pair<>(observer, TO_BE_ADDED));
     }
 
-    private void addInternal(LifecycleObserver observer) {
+    private void addInternal(T observer) {
         if (exists(observer)) {
             Log.e(LifecycleRegistry.TAG, ERR_RE_ADD + observer);
             return;
         }
-        GenericLifecycleObserver genericObserver = Lifecycling.getCallback(observer);
-        mData.add(new Pair<>(observer, genericObserver));
+        mData.add(observer);
+        onAdded(observer);
     }
 
-    void remove(LifecycleObserver observer) {
+    protected void onAdded(T observer) {
+
+    }
+
+    void remove(T observer) {
         if (!mLocked) {
             removeInternal(observer);
             return;
@@ -81,27 +91,31 @@
         mPendingModifications.add(new Pair<>(observer, TO_BE_REMOVED));
     }
 
-    private void removeInternal(LifecycleObserver observer) {
-        Iterator<Pair<LifecycleObserver, GenericLifecycleObserver>> iterator = mData.iterator();
+    private void removeInternal(T observer) {
+        Iterator<T> iterator = mData.iterator();
         while (iterator.hasNext()) {
-            if (iterator.next().first == observer) {
+            T next = iterator.next();
+            if (checkEquality(next, observer)) {
                 iterator.remove();
+                onRemoved(next);
                 return;
             }
         }
         Log.w(LifecycleRegistry.TAG, WARN_RE_REMOVE + observer);
     }
 
-    void forEach(Callback func) {
+    protected void onRemoved(T item) {}
+
+    void forEach(Callback<T> func) {
         mLocked = true;
         try {
             final int size = mData.size();
             for (int i = 0; i < size; i++) {
-                Pair<LifecycleObserver, GenericLifecycleObserver> pair = mData.get(i);
-                if (hasPendingRemoval(pair.first)) {
+                T item = mData.get(i);
+                if (hasPendingRemoval(item)) {
                     continue;
                 }
-                func.run(pair.second);
+                func.run(item);
             }
         } finally {
             mLocked = false;
@@ -116,7 +130,7 @@
         final int size = mPendingModifications.size();
         //noinspection StatementWithEmptyBody
         for (int i = 0; i < size; i++) {
-            Pair<LifecycleObserver, Object> pair = mPendingModifications.get(i);
+            Pair<T, Object> pair = mPendingModifications.get(i);
             if (pair.second == TO_BE_REMOVED) {
                 removeInternal(pair.first);
             } else { // TO_BE_ADDED
@@ -130,8 +144,7 @@
         return mData.size();
     }
 
-    interface Callback {
-        void run(GenericLifecycleObserver observer);
+    interface Callback<T> {
+        void run(T key);
     }
-
 }
diff --git a/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverListTest.java b/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverListTest.java
deleted file mode 100644
index 835e938..0000000
--- a/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverListTest.java
+++ /dev/null
@@ -1,222 +0,0 @@
-/*
- * Copyright (C) 2016 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.support.lifecycle;
-
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.fail;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-@RunWith(JUnit4.class)
-public class ObserverListTest {
-    private ObserverList mObserverList;
-
-    @Before
-    public void init() {
-        mObserverList = new ObserverList();
-    }
-
-    @Test
-    public void add() {
-        GenericLifecycleObserver observer = mock(GenericLifecycleObserver.class);
-        mObserverList.add(observer);
-        assertThat(mObserverList.size(), is(1));
-        assertThat(collect().get(0), is(observer));
-        mObserverList.remove(observer);
-        assertThat(mObserverList.size(), is(0));
-    }
-
-    @Test
-    public void addTwice() {
-        GenericLifecycleObserver observer = mock(GenericLifecycleObserver.class);
-        mObserverList.add(observer);
-        assertThat(mObserverList.size(), is(1));
-        mObserverList.add(observer);
-        assertThat(mObserverList.size(), is(1));
-    }
-
-    @Test
-    public void addTwiceWhileIterating() {
-        final ObserveAll observer = spy(new ObserveAll() {
-            @Override
-            public void onAny() {
-                mObserverList.add(this);
-                mObserverList.add(this);
-                mObserverList.add(this);
-            }
-        });
-        mObserverList.add(observer);
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver ignored) {
-                observer.onAny();
-            }
-        });
-        verify(observer, times(1)).onAny();
-        Assert.assertThat(mObserverList.size(), is(1));
-    }
-
-    @Test
-    public void addRemoveAdd() {
-        final ObserveAll observer2 = mock(ObserveAll.class);
-        final ObserveAll observer1 = spy(new ObserveAll() {
-            @Override
-            public void onAny() {
-                mObserverList.remove(this);
-                mObserverList.add(this);
-            }
-        });
-        mObserverList.add(observer1);
-        mObserverList.add(observer2);
-        assertThat(collectObservers(),
-                equalTo(Arrays.asList((LifecycleObserver) observer1, observer2)));
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver glo) {
-                ((ObserveAll) glo.getReceiver()).onAny();
-            }
-        });
-        verify(observer1, times(1)).onAny();
-        verify(observer2, times(1)).onAny();
-        // because 1 has been removed and re-added, it should get the next event after o1
-        assertThat(collectObservers(),
-                equalTo(Arrays.asList((LifecycleObserver) observer2, observer1)));
-    }
-
-    @Test
-    public void remove() {
-        GenericLifecycleObserver observer = mock(GenericLifecycleObserver.class);
-        LifecycleObserver obj = mock(LifecycleObserver.class);
-        when(observer.getReceiver()).thenReturn(obj);
-        mObserverList.add(observer);
-        assertThat(mObserverList.size(), is(1));
-        assertThat(collect().get(0), is(observer));
-        mObserverList.remove(observer);
-        assertThat(mObserverList.size(), is(0));
-    }
-
-    @Test
-    public void removeWhileTraversing() {
-        GenericLifecycleObserver observer1 = mock(GenericLifecycleObserver.class);
-        final GenericLifecycleObserver observer2 = mock(GenericLifecycleObserver.class);
-        mObserverList.add(observer1);
-        mObserverList.add(observer2);
-        final AtomicBoolean first = new AtomicBoolean(true);
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver observer) {
-                if (first.getAndSet(false)) {
-                    mObserverList.remove(observer2);
-                } else {
-                    fail("should never receive this call");
-                }
-            }
-        });
-    }
-
-    @Test
-    public void removeObjectWithMultipleCallbacksWhenTraversing() {
-        // if the removed object has multiple callbacks and already received one, should receive
-        // all.
-        final AtomicBoolean first = new AtomicBoolean(true);
-        final StartedObserverWith2Methods observer = spy(new StartedObserverWith2Methods() {
-            @Override
-            public void onStarted1() {
-                handle();
-            }
-
-            @Override
-            public void onStarted2() {
-                handle();
-            }
-
-            private void handle() {
-                if (first.getAndSet(false)) {
-                    mObserverList.remove(this);
-                }
-            }
-        });
-        mObserverList.add(observer);
-
-        final LifecycleProvider lifecycleProvider = mock(LifecycleProvider.class);
-        Lifecycle lifecycle = mock(Lifecycle.class);
-        when(lifecycleProvider.getLifecycle()).thenReturn(lifecycle);
-        when(lifecycle.getCurrentState()).thenReturn(Lifecycle.STARTED);
-
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver observer) {
-                observer.onStateChanged(lifecycleProvider, Lifecycle.ON_START);
-            }
-        });
-
-        verify(observer).onStarted1();
-        verify(observer).onStarted2();
-    }
-
-    private List<GenericLifecycleObserver> collect() {
-        final ArrayList<GenericLifecycleObserver> items = new ArrayList<>();
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver observer) {
-                items.add(observer);
-            }
-        });
-        return items;
-    }
-
-    private List<LifecycleObserver> collectObservers() {
-        final ArrayList<LifecycleObserver> items = new ArrayList<>();
-        mObserverList.forEach(new ObserverList.Callback() {
-            @Override
-            public void run(GenericLifecycleObserver observer) {
-                items.add((LifecycleObserver) observer.getReceiver());
-            }
-        });
-        return items;
-    }
-
-    @SuppressWarnings("unused")
-    private interface StartedObserverWith2Methods extends LifecycleObserver {
-        @OnLifecycleEvent(Lifecycle.ON_START)
-        void onStarted1();
-
-        @OnLifecycleEvent(Lifecycle.ON_START)
-        void onStarted2();
-    }
-
-    private interface ObserveAll extends LifecycleObserver {
-        @OnLifecycleEvent(Lifecycle.ANY)
-        void onAny();
-    }
-}
diff --git a/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverSetTest.java b/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverSetTest.java
new file mode 100644
index 0000000..100edcb
--- /dev/null
+++ b/lifecycle/runtime/src/test/java/com/android/support/lifecycle/ObserverSetTest.java
@@ -0,0 +1,180 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.support.lifecycle;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+@RunWith(JUnit4.class)
+public class ObserverSetTest {
+    private ObserverSet<String> mObserverSet;
+
+    @Before
+    public void init() {
+        mObserverSet = new ObserverSet<String>() {
+            @Override
+            protected boolean checkEquality(String existing, String added) {
+                return existing.equals(added);
+            }
+        };
+    }
+
+    private void add(String value) {
+        mObserverSet.add(value);
+    }
+
+    private void remove(String value) {
+        mObserverSet.remove(value);
+    }
+
+    @Test
+    public void add() {
+        add("a");
+        assertThat(mObserverSet.size(), is(1));
+        assertThat(collect().get(0), is("a"));
+        remove("a");
+        assertThat(mObserverSet.size(), is(0));
+    }
+
+    @Test
+    public void addTwice() {
+        add("a");
+        assertThat(mObserverSet.size(), is(1));
+        add("a");
+        assertThat(mObserverSet.size(), is(1));
+    }
+
+    @Test
+    public void addTwiceWhileIterating() {
+        add("a");
+        final AtomicInteger cnt = new AtomicInteger(0);
+        mObserverSet.forEach(new ObserverSet.Callback<String>() {
+            @Override
+            public void run(String key) {
+                add("a");
+                cnt.incrementAndGet();
+            }
+        });
+        assertThat(cnt.get(), is(1));
+        assertThat(mObserverSet.size(), is(1));
+    }
+
+    @Test
+    public void addRemoveAdd() {
+        final ObserverSet<ObserveAll> list = new ObserverSet<ObserveAll>() {
+            @Override
+            protected boolean checkEquality(ObserveAll existing,
+                    ObserveAll added) {
+                return existing == added;
+            }
+        };
+        final ObserveAll observer2 = mock(ObserveAll.class);
+        final ObserveAll observer1 = spy(new ObserveAll() {
+            @Override
+            public void onAny() {
+                list.remove(this);
+                list.add(this);
+            }
+        });
+        list.add(observer1);
+        list.add(observer2);
+        assertThat(collect(list), equalTo(Arrays.asList(observer1, observer2)));
+        list.forEach(new ObserverSet.Callback<ObserveAll>() {
+            @Override
+            public void run(ObserveAll key) {
+                key.onAny();
+            }
+        });
+
+        verify(observer1, times(1)).onAny();
+        verify(observer2, times(1)).onAny();
+        // because 1 has been removed and re-added, it should get the next event after o1
+        assertThat(collect(list), equalTo(Arrays.asList(observer2, observer1)));
+    }
+
+    @Test
+    public void remove() {
+        add("a");
+        assertThat(mObserverSet.size(), is(1));
+        assertThat(collect().get(0), is("a"));
+        remove("a");
+        assertThat(mObserverSet.size(), is(0));
+    }
+
+    @Test
+    public void removeWhileTraversing() {
+        add("a");
+        add("b");
+        final AtomicBoolean first = new AtomicBoolean(true);
+        mObserverSet.forEach(new ObserverSet.Callback<String>() {
+            @Override
+            public void run(String key) {
+                if (first.getAndSet(false)) {
+                    remove("b");
+                } else {
+                    fail("should never receive this call");
+                }
+            }
+        });
+    }
+
+    private List<String> collect() {
+        return collect(mObserverSet);
+    }
+
+    private <T> List<T> collect(ObserverSet<T> list) {
+        final ArrayList<T> items = new ArrayList<>();
+        list.forEach(new ObserverSet.Callback<T>() {
+            @Override
+            public void run(T value) {
+                items.add(value);
+            }
+        });
+        return items;
+    }
+
+    @SuppressWarnings("unused")
+    private interface StartedObserverWith2Methods extends LifecycleObserver {
+        @OnLifecycleEvent(Lifecycle.ON_START)
+        void onStarted1();
+
+        @OnLifecycleEvent(Lifecycle.ON_START)
+        void onStarted2();
+    }
+
+    private interface ObserveAll extends LifecycleObserver {
+        @OnLifecycleEvent(Lifecycle.ANY)
+        void onAny();
+    }
+}