Merge "Fix SafeIteratorWithAddition in case of deletion of first element, while we are at it." into oc-support-26.0-dev
diff --git a/app-toolkit/common/src/main/java/android/arch/core/internal/SafeIterableMap.java b/app-toolkit/common/src/main/java/android/arch/core/internal/SafeIterableMap.java
index 16a7607..00e102f 100644
--- a/app-toolkit/common/src/main/java/android/arch/core/internal/SafeIterableMap.java
+++ b/app-toolkit/common/src/main/java/android/arch/core/internal/SafeIterableMap.java
@@ -300,18 +300,19 @@
 
     private class IteratorWithAdditions implements Iterator<Map.Entry<K, V>>, SupportRemove<K, V> {
         private Entry<K, V> mCurrent;
-        private boolean mFirstStep = true;
+        private boolean mBeforeStart = true;
 
         @Override
         public void supportRemove(@NonNull Entry<K, V> entry) {
             if (entry == mCurrent) {
                 mCurrent = mCurrent.mPrevious;
+                mBeforeStart = mCurrent == null;
             }
         }
 
         @Override
         public boolean hasNext() {
-            if (mFirstStep) {
+            if (mBeforeStart) {
                 return mStart != null;
             }
             return mCurrent != null && mCurrent.mNext != null;
@@ -319,8 +320,8 @@
 
         @Override
         public Map.Entry<K, V> next() {
-            if (mFirstStep) {
-                mFirstStep = false;
+            if (mBeforeStart) {
+                mBeforeStart = false;
                 mCurrent = mStart;
             } else {
                 mCurrent = mCurrent != null ? mCurrent.mNext : null;
diff --git a/app-toolkit/common/src/test/java/android/arch/core/internal/SafeIterableMapTest.java b/app-toolkit/common/src/test/java/android/arch/core/internal/SafeIterableMapTest.java
index 4b25e34..d879543 100644
--- a/app-toolkit/common/src/test/java/android/arch/core/internal/SafeIterableMapTest.java
+++ b/app-toolkit/common/src/test/java/android/arch/core/internal/SafeIterableMapTest.java
@@ -200,6 +200,20 @@
     }
 
     @Test
+    public void testRemoveDuringIteration4() {
+        SafeIterableMap<Integer, Boolean> map = mapOf(1, 2);
+        int[] expected = new int[]{1, 2};
+        int index = 0;
+        for (Entry<Integer, Boolean> entry : map) {
+            assertThat(entry.getKey(), is(expected[index++]));
+            if (index == 1) {
+                map.remove(1);
+            }
+        }
+        assertThat(index, is(2));
+    }
+
+    @Test
     public void testAdditionDuringIteration() {
         SafeIterableMap<Integer, Boolean> map = mapOf(1, 2, 3, 4);
         int[] expected = new int[]{1, 2, 3, 4};
@@ -317,6 +331,22 @@
     }
 
     @Test
+    public void testIteratorWithAddition5() {
+        SafeIterableMap<Integer, Boolean> map = mapOf(1, 2);
+        int[] expected = new int[]{1, 2};
+        int index = 0;
+        Iterator<Entry<Integer, Boolean>> iterator = map.iteratorWithAdditions();
+        while (iterator.hasNext()) {
+            Entry<Integer, Boolean> entry = iterator.next();
+            assertThat(entry.getKey(), is(expected[index++]));
+            if (index == 1) {
+                map.remove(1);
+            }
+        }
+        assertThat(index, is(2));
+    }
+
+    @Test
     public void testDescendingIteration() {
         SafeIterableMap<Integer, Boolean> map = mapOf(1, 2, 3, 4);
         int[] expected = new int[]{4, 3, 2, 1};
diff --git a/lifecycle/extensions/src/test/java/android/arch/lifecycle/LiveDataTest.java b/lifecycle/extensions/src/test/java/android/arch/lifecycle/LiveDataTest.java
index f401e1c..1802d94 100644
--- a/lifecycle/extensions/src/test/java/android/arch/lifecycle/LiveDataTest.java
+++ b/lifecycle/extensions/src/test/java/android/arch/lifecycle/LiveDataTest.java
@@ -368,6 +368,22 @@
     }
 
     @Test
+    public void testRemoveDuringSetValue() {
+        mRegistry.handleLifecycleEvent(ON_START);
+        final Observer observer1 = spy(new Observer<String>() {
+            @Override
+            public void onChanged(String o) {
+                mLiveData.removeObserver(this);
+            }
+        });
+        Observer<String> observer2 = (Observer<String>) mock(Observer.class);
+        mLiveData.observeForever(observer1);
+        mLiveData.observe(mOwner, observer2);
+        mLiveData.setValue("gt");
+        verify(observer2).onChanged("gt");
+    }
+
+    @Test
     public void testDataChangeDuringStateChange() {
         mRegistry.handleLifecycleEvent(ON_START);
         mRegistry.addObserver(new LifecycleObserver() {