datastore: Update AllocateIDs to take keys.

impl/memory/datastore_data.go

Index: impl/memory/datastore_data.go
diff --git a/impl/memory/datastore_data.go b/impl/memory/datastore_data.go
index 94a868878e6db035cb3de4e9b4e6758e8e514e2f..f530ff98aee7fab51abf3dac230e58d5f5d112cf 100644
--- a/impl/memory/datastore_data.go
+++ b/impl/memory/datastore_data.go
@@ -215,12 +215,31 @@ func incrementLocked(ents memCollection, key []byte, amt int) int64 {
 	return ret
 }
 
-func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) (int64, error) {
-	d.Lock()
-	defer d.Unlock()
+func (d *dataStoreData) allocateIDs(keys []*ds.Key) error {
+	start, err := func() (int64, error) {
+		d.Lock()
+		defer d.Unlock()
+
+		// Allocate a set of IDs in our key namespace. We know that "keys" has at
+		// least one entry and that all keys share the same namespace and entty
+		// type because checkfilter has already asserted this.
+		baseKey := keys[0]
+		ents := d.head.GetOrCreateCollection("ents:" + baseKey.Namespace())
+		start, err := d.allocateIDsLocked(ents, baseKey, len(keys))
+		if err != nil {
+			return -1, err
+		}
+		return start, nil
+	}()
+	if err != nil {
+		return err
+	}
 
-	ents := d.head.GetOrCreateCollection("ents:" + incomplete.Namespace())
-	return d.allocateIDsLocked(ents, incomplete, n)
+	// Update the keys in our "keys" slice with their new IDs.
+	for i, k := range keys {
+		keys[i] = ds.NewKey(k.AppID(), k.Namespace(), k.Kind(), "", start+int64(i), k.Parent())
+	}
+	return nil
 }
 
 func (d *dataStoreData) allocateIDsLocked(ents memCollection, incomplete *ds.Key, n int) (int64, error) {