2 files changed, 25 insertions, 10 deletions

diff --git a/java/com/android/dialer/speeddial/database/SpeedDialEntryDao.java b/java/com/android/dialer/speeddial/database/SpeedDialEntryDao.java
index ce771c3c8..4d6ac2d7f 100644
--- a/java/com/android/dialer/speeddial/database/SpeedDialEntryDao.java
+++ b/java/com/android/dialer/speeddial/database/SpeedDialEntryDao.java
@@ -17,6 +17,7 @@
 package com.android.dialer.speeddial.database;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 
 /**
  * Interface that databases support speed dial entries should implement.
@@ -32,8 +33,10 @@ public interface SpeedDialEntryDao {
    * Insert new entries.
    *
    * <p>{@link SpeedDialEntry#id() ids} must be null.
+   *
+   * @return a map of the inserted entries to their new ids.
    */
-  void insert(ImmutableList<SpeedDialEntry> entries);
+  ImmutableMap<SpeedDialEntry, Long> insert(ImmutableList<SpeedDialEntry> entries);
 
   /**
    * Insert a new entry.
@@ -59,11 +62,12 @@ public interface SpeedDialEntryDao {
   /**
    * Inserts, updates and deletes rows all in on transaction.
    *
+   * @return a map of the inserted entries to their new ids.
    * @see #insert(ImmutableList)
    * @see #update(ImmutableList)
    * @see #delete(ImmutableList)
    */
-  void insertUpdateAndDelete(
+  ImmutableMap<SpeedDialEntry, Long> insertUpdateAndDelete(
       ImmutableList<SpeedDialEntry> entriesToInsert,
       ImmutableList<SpeedDialEntry> entriesToUpdate,
       ImmutableList<Long> entriesToDelete);
diff --git a/java/com/android/dialer/speeddial/database/SpeedDialEntryDatabaseHelper.java b/java/com/android/dialer/speeddial/database/SpeedDialEntryDatabaseHelper.java
index 137933fbe..544bb3613 100644
--- a/java/com/android/dialer/speeddial/database/SpeedDialEntryDatabaseHelper.java
+++ b/java/com/android/dialer/speeddial/database/SpeedDialEntryDatabaseHelper.java
@@ -26,6 +26,7 @@ import com.android.dialer.common.Assert;
 import com.android.dialer.common.database.Selection;
 import com.android.dialer.speeddial.database.SpeedDialEntry.Channel;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -132,30 +133,39 @@ public final class SpeedDialEntryDatabaseHelper extends SQLiteOpenHelper
   }
 
   @Override
-  public void insert(ImmutableList<SpeedDialEntry> entries) {
+  public ImmutableMap<SpeedDialEntry, Long> insert(ImmutableList<SpeedDialEntry> entries) {
     if (entries.isEmpty()) {
-      return;
+      return ImmutableMap.of();
     }
 
     SQLiteDatabase db = getWritableDatabase();
     db.beginTransaction();
     try {
-      insert(db, entries);
+      ImmutableMap<SpeedDialEntry, Long> insertedEntriesToIdsMap = insert(db, entries);
       db.setTransactionSuccessful();
+      return insertedEntriesToIdsMap;
     } finally {
       db.endTransaction();
       db.close();
     }
   }
 
-  private void insert(SQLiteDatabase writeableDatabase, ImmutableList<SpeedDialEntry> entries) {
+  private ImmutableMap<SpeedDialEntry, Long> insert(
+      SQLiteDatabase writeableDatabase, ImmutableList<SpeedDialEntry> entries) {
+    ImmutableMap.Builder<SpeedDialEntry, Long> insertedEntriesToIdsMap = ImmutableMap.builder();
     for (SpeedDialEntry entry : entries) {
       Assert.checkArgument(entry.id() == null);
-      if (writeableDatabase.insert(TABLE_NAME, null, buildContentValuesWithoutId(entry)) == -1L) {
+      long id = writeableDatabase.insert(TABLE_NAME, null, buildContentValuesWithoutId(entry));
+      if (id == -1L) {
         throw Assert.createUnsupportedOperationFailException(
             "Attempted to insert a row that already exists.");
       }
+      // It's impossible to insert two identical entries but this is an important assumption we need
+      // to verify because there's an assumption that each entry will correspond to exactly one id.
+      // ImmutableMap#put verifies this check for us.
+      insertedEntriesToIdsMap.put(entry, id);
     }
+    return insertedEntriesToIdsMap.build();
   }
 
   @Override
@@ -255,20 +265,21 @@ public final class SpeedDialEntryDatabaseHelper extends SQLiteOpenHelper
   }
 
   @Override
-  public void insertUpdateAndDelete(
+  public ImmutableMap<SpeedDialEntry, Long> insertUpdateAndDelete(
       ImmutableList<SpeedDialEntry> entriesToInsert,
       ImmutableList<SpeedDialEntry> entriesToUpdate,
       ImmutableList<Long> entriesToDelete) {
     if (entriesToInsert.isEmpty() && entriesToUpdate.isEmpty() && entriesToDelete.isEmpty()) {
-      return;
+      return ImmutableMap.of();
     }
     SQLiteDatabase db = getWritableDatabase();
     db.beginTransaction();
     try {
-      insert(db, entriesToInsert);
+      ImmutableMap<SpeedDialEntry, Long> insertedEntriesToIdsMap = insert(db, entriesToInsert);
       update(db, entriesToUpdate);
       delete(db, entriesToDelete);
       db.setTransactionSuccessful();
+      return insertedEntriesToIdsMap;
     } finally {
       db.endTransaction();
       db.close();