Handle @Relation annotation with non-list type of more than 999 rows.

When fetching one-to-one relationships where there are more than 999
parent items we need to batch query the children. We already had the
mechanism to do this for one-to-many and many-to-many but for one-to-one
the same strategy failed because the batching code relied in the value
of the relating map to be a collection, a reference that can be passed
around and filled separately. Meanwhile for one-to-one we have to
copy the batch map into the original relationship map.

Bug: 143105450
Test: PojoWithRelationTest, BareRelationDatabaseTest
Change-Id: I00fac5076d4243e973d0e9e2801fef659d518e05
diff --git a/room/compiler/src/main/kotlin/androidx/room/ext/javapoet_ext.kt b/room/compiler/src/main/kotlin/androidx/room/ext/javapoet_ext.kt
index 0267637..de835eb8 100644
--- a/room/compiler/src/main/kotlin/androidx/room/ext/javapoet_ext.kt
+++ b/room/compiler/src/main/kotlin/androidx/room/ext/javapoet_ext.kt
@@ -121,6 +121,7 @@
 
 object CommonTypeNames {
     val LIST = ClassName.get("java.util", "List")
+    val MAP = ClassName.get("java.util", "Map")
     val SET = ClassName.get("java.util", "Set")
     val STRING = ClassName.get("java.lang", "String")
     val INTEGER = ClassName.get("java.lang", "Integer")
diff --git a/room/compiler/src/main/kotlin/androidx/room/writer/RelationCollectorMethodWriter.kt b/room/compiler/src/main/kotlin/androidx/room/writer/RelationCollectorMethodWriter.kt
index b2d3410..1ca5dba 100644
--- a/room/compiler/src/main/kotlin/androidx/room/writer/RelationCollectorMethodWriter.kt
+++ b/room/compiler/src/main/kotlin/androidx/room/writer/RelationCollectorMethodWriter.kt
@@ -18,6 +18,7 @@
 
 import androidx.room.ext.AndroidTypeNames
 import androidx.room.ext.CollectionTypeNames
+import androidx.room.ext.CommonTypeNames
 import androidx.room.ext.L
 import androidx.room.ext.N
 import androidx.room.ext.RoomTypeNames
@@ -27,6 +28,7 @@
 import androidx.room.solver.query.result.PojoRowAdapter
 import androidx.room.vo.RelationCollector
 import com.squareup.javapoet.ClassName
+import com.squareup.javapoet.CodeBlock
 import com.squareup.javapoet.MethodSpec
 import com.squareup.javapoet.ParameterSpec
 import com.squareup.javapoet.ParameterizedTypeName
@@ -72,6 +74,20 @@
                     collector.mapTypeName.rawType == CollectionTypeNames.LONG_SPARSE_ARRAY
             val usingArrayMap =
                     collector.mapTypeName.rawType == CollectionTypeNames.ARRAY_MAP
+            fun CodeBlock.Builder.addBatchPutAllStatement(tmpMapVar: String) {
+                if (usingArrayMap) {
+                    // When using ArrayMap there is ambiguity in the putAll() method, clear the
+                    // confusion by casting the temporary map.
+                    val disambiguityTypeName =
+                        ParameterizedTypeName.get(CommonTypeNames.MAP,
+                            collector.mapTypeName.typeArguments[0],
+                            collector.mapTypeName.typeArguments[1])
+                    addStatement("$N.putAll(($T) $L)",
+                        param, disambiguityTypeName, tmpMapVar)
+                } else {
+                    addStatement("$N.putAll($L)", param, tmpMapVar)
+                }
+            }
             if (usingLongSparseArray) {
                 beginControlFlow("if ($N.isEmpty())", param)
             } else {
@@ -99,16 +115,25 @@
                     addStatement("$T $L = 0", TypeName.INT, mapIndexVar)
                     addStatement("final $T $L = $N.size()", TypeName.INT, limitVar, param)
                     beginControlFlow("while($L < $L)", mapIndexVar, limitVar).apply {
-                        addStatement("$L.put($N.keyAt($L), $N.valueAt($L))",
-                            tmpMapVar, param, mapIndexVar, param, mapIndexVar)
+                        if (collector.relationTypeIsCollection) {
+                            addStatement("$L.put($N.keyAt($L), $N.valueAt($L))",
+                                tmpMapVar, param, mapIndexVar, param, mapIndexVar)
+                        } else {
+                            addStatement("$L.put($N.keyAt($L), null)",
+                                tmpMapVar, param, mapIndexVar)
+                        }
                         addStatement("$L++", mapIndexVar)
                     }
                 } else {
                     val mapKeyVar = scope.getTmpVar("_mapKey")
                     beginControlFlow("for($T $L : $L)",
                         collector.keyTypeName, mapKeyVar, KEY_SET_VARIABLE).apply {
-                        addStatement("$L.put($L, $N.get($L))",
-                            tmpMapVar, mapKeyVar, param, mapKeyVar)
+                        if (collector.relationTypeIsCollection) {
+                            addStatement("$L.put($L, $N.get($L))",
+                                tmpMapVar, mapKeyVar, param, mapKeyVar)
+                        } else {
+                            addStatement("$L.put($L, null)", tmpMapVar, mapKeyVar)
+                        }
                     }
                 }.apply {
                     addStatement("$L++", tmpIndexVar)
@@ -116,6 +141,11 @@
                         tmpIndexVar, RoomTypeNames.ROOM_DB).apply {
                         // recursively load that batch
                         addStatement("$L($L)", methodName, tmpMapVar)
+                        // for non collection relation, put the loaded batch in the original map,
+                        // not needed when dealing with collections since references are passed
+                        if (!collector.relationTypeIsCollection) {
+                            addBatchPutAllStatement(tmpMapVar)
+                        }
                         // clear nukes the backing data hence we create a new one
                         addStatement("$L = new $T($T.MAX_BIND_PARAMETER_CNT)",
                             tmpMapVar, collector.mapTypeName, RoomTypeNames.ROOM_DB)
@@ -125,6 +155,10 @@
                 beginControlFlow("if($L > 0)", tmpIndexVar).apply {
                     // load the last batch
                     addStatement("$L($L)", methodName, tmpMapVar)
+                    // for non collection relation, put the last batch in the original map
+                    if (!collector.relationTypeIsCollection) {
+                        addBatchPutAllStatement(tmpMapVar)
+                    }
                 }.endControlFlow()
                 addStatement("return")
             }.endControlFlow()
diff --git a/room/integration-tests/noappcompattestapp/src/androidTest/java/androidx/room/integration/noappcompat/BareRelationDatabaseTest.java b/room/integration-tests/noappcompattestapp/src/androidTest/java/androidx/room/integration/noappcompat/BareRelationDatabaseTest.java
index ff1230b..64ceb77 100644
--- a/room/integration-tests/noappcompattestapp/src/androidTest/java/androidx/room/integration/noappcompat/BareRelationDatabaseTest.java
+++ b/room/integration-tests/noappcompattestapp/src/androidTest/java/androidx/room/integration/noappcompat/BareRelationDatabaseTest.java
@@ -31,7 +31,7 @@
 import androidx.room.RoomDatabase;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
+import androidx.test.filters.LargeTest;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,7 +40,7 @@
 
 // More than a simple read & write, this test that we generate correct relationship collector
 // code that doesn't use androidx.collection
-@SmallTest
+@LargeTest
 @RunWith(AndroidJUnit4.class)
 @SuppressWarnings("WeakerAccess") // to avoid naming field with m
 public class BareRelationDatabaseTest {
@@ -62,6 +62,28 @@
         assertThat(result.pets.get(1).petId, is(2L));
     }
 
+    @Test
+    public void large_nonCollectionRelation() {
+        RelationDatabase db = Room.inMemoryDatabaseBuilder(
+                ApplicationProvider.getApplicationContext(), RelationDatabase.class)
+                .build();
+        UserPetDao dao = db.getDao();
+
+        int count = 2000;
+        db.runInTransaction(() -> {
+            for (int i = 1; i <= count; i++) {
+                dao.insertUser(new User(i));
+                dao.insertPet(new Pet(i, i));
+            }
+        });
+
+        List<UserAndPet> ownerAndPet = dao.getUsersWithPet();
+        assertThat(ownerAndPet.size(), is(count));
+        for (int i = 0; i < count; i++) {
+            assertThat(ownerAndPet.get(i).pet.petId, is(i + 1L));
+        }
+    }
+
     @Database(entities = {User.class, Pet.class}, version = 1, exportSchema = false)
     abstract static class RelationDatabase extends RoomDatabase {
         abstract UserPetDao getDao();
@@ -72,6 +94,9 @@
         @Query("SELECT * FROM User WHERE userId = :id")
         UserAndPets getUserWithPets(long id);
 
+        @Query("SELECT * FROM User")
+        List<UserAndPet> getUsersWithPet();
+
         @Insert
         void insertUser(User user);
 
@@ -107,4 +132,11 @@
         @Relation(parentColumn = "userId", entityColumn = "ownerId")
         public List<Pet> pets;
     }
+
+    static class UserAndPet {
+        @Embedded
+        public User user;
+        @Relation(parentColumn = "userId", entityColumn = "ownerId")
+        public Pet pet;
+    }
 }
diff --git a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/RobotsDao.java b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/RobotsDao.java
index ee4acd8..3cd476d 100644
--- a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/RobotsDao.java
+++ b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/RobotsDao.java
@@ -22,6 +22,7 @@
 import androidx.room.integration.testapp.vo.Cluster;
 import androidx.room.integration.testapp.vo.Hivemind;
 import androidx.room.integration.testapp.vo.Robot;
+import androidx.room.integration.testapp.vo.RobotAndHivemind;
 
 import java.util.List;
 import java.util.UUID;
@@ -40,4 +41,7 @@
 
     @Query("SELECT * FROM Robot WHERE mHiveId = :hiveId")
     List<Robot> getHiveRobots(UUID hiveId);
+
+    @Query("SELECT * FROM Robot")
+    List<RobotAndHivemind> getRobotsWithHivemind();
 }
diff --git a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/PojoWithRelationTest.java b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/PojoWithRelationTest.java
index 1cced11..3cb7571 100644
--- a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/PojoWithRelationTest.java
+++ b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/PojoWithRelationTest.java
@@ -28,6 +28,7 @@
 import androidx.room.integration.testapp.vo.PetAndOwner;
 import androidx.room.integration.testapp.vo.PetWithToyIds;
 import androidx.room.integration.testapp.vo.Robot;
+import androidx.room.integration.testapp.vo.RobotAndHivemind;
 import androidx.room.integration.testapp.vo.Toy;
 import androidx.room.integration.testapp.vo.User;
 import androidx.room.integration.testapp.vo.UserAndAllPets;
@@ -301,6 +302,44 @@
     }
 
     @Test
+    public void large_nonCollectionRelation() {
+        int count = 2000;
+        mDatabase.runInTransaction(() -> {
+            for (int i = 1; i <= count; i++) {
+                mUserDao.insert(TestUtil.createUser(i));
+                Pet pet = TestUtil.createPet(i);
+                pet.setUserId(i);
+                mPetDao.insertOrReplace(pet);
+            }
+        });
+
+        List<PetAndOwner> petAndOwners = mPetDao.allPetsWithOwners();
+        assertThat(petAndOwners.size(), is(count));
+        for (int i = 0; i < count; i++) {
+            assertThat(petAndOwners.get(i).getUser().getId(), is(i + 1));
+        }
+    }
+
+    @Test
+    public void large_nonCollectionRelation_withComplexKey() {
+        int count = 2000;
+        mDatabase.runInTransaction(() -> {
+            for (int i = 1; i <= count; i++) {
+                Hivemind hivemind = new Hivemind(UUID.randomUUID());
+                mRobotsDao.putHivemind(hivemind);
+                Robot robot = new Robot(UUID.randomUUID(), hivemind.mId);
+                mRobotsDao.putRobot(robot);
+            }
+        });
+
+        List<RobotAndHivemind> robotsWithHivemind = mRobotsDao.getRobotsWithHivemind();
+        assertThat(robotsWithHivemind.size(), is(count));
+        for (int i = 0; i < count; i++) {
+            assertThat(robotsWithHivemind.get(i).getHivemind(), is(notNullValue()));
+        }
+    }
+
+    @Test
     public void relationWithBlobKey() {
         UUID hiveId1 = UUID.randomUUID();
         UUID hiveId2 = UUID.randomUUID();
diff --git a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/vo/RobotAndHivemind.java b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/vo/RobotAndHivemind.java
new file mode 100644
index 0000000..965011f
--- /dev/null
+++ b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/vo/RobotAndHivemind.java
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2019 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 androidx.room.integration.testapp.vo;
+
+import androidx.room.Embedded;
+import androidx.room.Relation;
+
+public class RobotAndHivemind {
+
+    @Embedded
+    public final Robot mRobot;
+
+    @Relation(parentColumn = "mHiveId", entityColumn = "mId")
+    public final Hivemind mHivemind;
+
+    public RobotAndHivemind(Robot robot, Hivemind hivemind) {
+        mRobot = robot;
+        mHivemind = hivemind;
+    }
+
+    public Robot getRobot() {
+        return mRobot;
+    }
+
+    public Hivemind getHivemind() {
+        return mHivemind;
+    }
+}