Adding support for ArrayMap return type in @MapInfo.

Adding ArrayMap to the list of supported map return type for Dao query functions. This type is very similar to the Map type handling, hence functionality has been added to MapQueryResultAdapter.kt to handle ArrayMap to avoid duplicate code.

Test: MultimapQueryTest.java
Bug: 138910317
Change-Id: Ic7985d3dbeb242158b9ae33989279131032a67c8
diff --git a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/MusicDao.java b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/MusicDao.java
index 76e5981..713b8c9 100644
--- a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/MusicDao.java
+++ b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/dao/MusicDao.java
@@ -16,6 +16,7 @@
 
 package androidx.room.integration.testapp.dao;
 
+import androidx.collection.ArrayMap;
 import androidx.lifecycle.LiveData;
 import androidx.room.Dao;
 import androidx.room.Insert;
@@ -250,10 +251,21 @@
     @Query("SELECT * FROM Artist JOIN Image ON Artist.mArtistName = Image.mArtistInImage")
     ImmutableMap<Artist, Long> getAllArtistsWithAlbumCoverYear();
 
+    @RewriteQueriesToDropUnusedColumns
+    @MapInfo(valueColumn = "mImageYear")
+    @Query("SELECT * FROM Artist JOIN Image ON Artist.mArtistName = Image.mArtistInImage")
+    ArrayMap<Artist, Long> getAllArtistsWithAlbumCoverYearArrayMap();
+
     @MapInfo(keyColumn = "mImageYear")
     @RawQuery
     ImmutableMap<Long, Artist> getAllAlbumCoverYearToArtistsWithRawQuery(SupportSQLiteQuery query);
 
+    @MapInfo(keyColumn = "mImageYear")
+    @RawQuery
+    ArrayMap<Long, Artist> getAllAlbumCoverYearToArtistsWithRawQueryArrayMap(
+            SupportSQLiteQuery query
+    );
+
     @MapInfo(keyColumn = "mAlbumCover", valueColumn = "mIsActive")
     @Query("SELECT * FROM Image JOIN Artist ON Artist.mArtistName = Image.mArtistInImage")
     ImmutableMap<ByteBuffer, Boolean> getAlbumCoversWithBandActivity();
diff --git a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/MultimapQueryTest.java b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/MultimapQueryTest.java
index 090db86..ab71b1c 100644
--- a/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/MultimapQueryTest.java
+++ b/room/integration-tests/testapp/src/androidTest/java/androidx/room/integration/testapp/test/MultimapQueryTest.java
@@ -24,6 +24,7 @@
 import android.content.Context;
 
 import androidx.arch.core.executor.testing.CountingTaskExecutorRule;
+import androidx.collection.ArrayMap;
 import androidx.lifecycle.Lifecycle;
 import androidx.lifecycle.LiveData;
 import androidx.lifecycle.testing.TestLifecycleOwner;
@@ -935,6 +936,37 @@
     }
 
     @Test
+    public void testImageYearToArtistRawQueryArrayMap() {
+        mMusicDao.addArtists(mRhcp, mAcDc, mTheClash, mPinkFloyd);
+        mMusicDao.addImages(mPinkFloydAlbumCover, mRhcpAlbumCover);
+
+        ArrayMap<Long, Artist> imageToArtistsMap =
+                mMusicDao.getAllAlbumCoverYearToArtistsWithRawQueryArrayMap(
+                        new SimpleSQLiteQuery(
+                                "SELECT * FROM Image JOIN Artist ON Artist.mArtistName = Image"
+                                        + ".mArtistInImage"
+                        )
+                );
+
+        assertThat(imageToArtistsMap.get(2006L)).isEqualTo(mRhcp);
+        assertThat(
+                imageToArtistsMap.keySet()).containsExactlyElementsIn(Arrays.asList(2006L, 1973L)
+        );
+    }
+
+
+    @Test
+    public void testArtistToImageYearArrayMap() {
+        mMusicDao.addArtists(mRhcp, mAcDc, mTheClash, mPinkFloyd);
+        mMusicDao.addImages(mPinkFloydAlbumCover, mRhcpAlbumCover);
+
+        ArrayMap<Artist, Long> artistNameToImagesMap =
+                mMusicDao.getAllArtistsWithAlbumCoverYearArrayMap();
+
+        assertThat(artistNameToImagesMap.get(mRhcp)).isEqualTo(2006L);
+    }
+
+    @Test
     public void testAlbumCoversWithBandActivity() {
         mMusicDao.addArtists(mRhcp, mAcDc, mTheClash, mPinkFloyd);
         mMusicDao.addImages(mPinkFloydAlbumCover, mRhcpAlbumCover);
diff --git a/room/room-compiler/src/main/kotlin/androidx/room/solver/TypeAdapterStore.kt b/room/room-compiler/src/main/kotlin/androidx/room/solver/TypeAdapterStore.kt
index 3a55bb7..6216452 100644
--- a/room/room-compiler/src/main/kotlin/androidx/room/solver/TypeAdapterStore.kt
+++ b/room/room-compiler/src/main/kotlin/androidx/room/solver/TypeAdapterStore.kt
@@ -19,6 +19,7 @@
 import androidx.room.compiler.processing.XType
 import androidx.room.compiler.processing.isArray
 import androidx.room.compiler.processing.isEnum
+import androidx.room.ext.CollectionTypeNames.ARRAY_MAP
 import androidx.room.ext.CommonTypeNames
 import androidx.room.ext.GuavaBaseTypeNames
 import androidx.room.ext.isEntityElement
@@ -572,7 +573,9 @@
                 valueRowAdapter = valueRowAdapter,
                 immutableClassName = immutableClassName
             )
-        } else if (typeMirror.isTypeOf(java.util.Map::class)) {
+        } else if (typeMirror.isTypeOf(java.util.Map::class) ||
+            typeMirror.rawType.typeName == ARRAY_MAP
+        ) {
             val keyTypeArg = typeMirror.typeArguments[0].extendsBoundOrSelf()
             val mapValueTypeArg = typeMirror.typeArguments[1].extendsBoundOrSelf()
 
@@ -623,7 +626,8 @@
                         valueTypeArg = valueTypeArg,
                         keyRowAdapter = keyRowAdapter,
                         valueRowAdapter = valueRowAdapter,
-                        valueCollectionType = mapValueTypeArg
+                        valueCollectionType = mapValueTypeArg,
+                        isArrayMap = typeMirror.rawType.typeName == ARRAY_MAP
                     )
                 } else {
                     context.logger.e(
@@ -655,7 +659,8 @@
                     valueTypeArg = mapValueTypeArg,
                     keyRowAdapter = keyRowAdapter,
                     valueRowAdapter = valueRowAdapter,
-                    valueCollectionType = null
+                    valueCollectionType = null,
+                    isArrayMap = typeMirror.rawType.typeName == ARRAY_MAP
                 )
             }
         }
diff --git a/room/room-compiler/src/main/kotlin/androidx/room/solver/query/result/MapQueryResultAdapter.kt b/room/room-compiler/src/main/kotlin/androidx/room/solver/query/result/MapQueryResultAdapter.kt
index 0328383..d1f4465 100644
--- a/room/room-compiler/src/main/kotlin/androidx/room/solver/query/result/MapQueryResultAdapter.kt
+++ b/room/room-compiler/src/main/kotlin/androidx/room/solver/query/result/MapQueryResultAdapter.kt
@@ -17,6 +17,7 @@
 package androidx.room.solver.query.result
 
 import androidx.room.compiler.processing.XType
+import androidx.room.ext.CollectionTypeNames.ARRAY_MAP
 import androidx.room.ext.L
 import androidx.room.ext.T
 import androidx.room.solver.CodeGenScope
@@ -28,7 +29,8 @@
     override val valueTypeArg: XType,
     private val keyRowAdapter: RowAdapter,
     private val valueRowAdapter: RowAdapter,
-    private val valueCollectionType: XType?
+    private val valueCollectionType: XType?,
+    isArrayMap: Boolean = false
 ) : QueryResultAdapter(listOf(keyRowAdapter, valueRowAdapter)), MultimapQueryResultAdapter {
     private val declaredToConcreteCollection = mapOf<ClassName, ClassName>(
         ClassName.get(List::class.java) to ClassName.get(ArrayList::class.java),
@@ -54,14 +56,14 @@
     }
 
     private val mapType = ParameterizedTypeName.get(
-        ClassName.get(Map::class.java),
+        if (isArrayMap) ARRAY_MAP else ClassName.get(Map::class.java),
         keyTypeArg.typeName,
         declaredValueType
     )
 
     // LinkedHashMap is used as impl to preserve key ordering for ordered query results.
     private val mapImplType = ParameterizedTypeName.get(
-        ClassName.get(LinkedHashMap::class.java),
+        if (isArrayMap) ARRAY_MAP else ClassName.get(LinkedHashMap::class.java),
         keyTypeArg.typeName,
         declaredValueType
     )