Refactor a bit.

impl/memory/raw_datastore_data.go

Index: impl/memory/raw_datastore_data.go
diff --git a/memory/raw_datastore_data.go b/impl/memory/raw_datastore_data.go
similarity index 64%
rename from memory/raw_datastore_data.go
rename to impl/memory/raw_datastore_data.go
index e88e43a931ca46a4df4681a6ac8bc359f6085034..9ea3624bd9c9225c5f60011dfba8186d06b0a317 100644
--- a/memory/raw_datastore_data.go
+++ b/impl/memory/raw_datastore_data.go
@@ -11,10 +11,9 @@ import (
 	"sync"
 	"sync/atomic"
 
-	"golang.org/x/net/context"
-
 	"github.com/luci/gae"
-	"github.com/luci/gae/helper"
+	rds "github.com/luci/gae/service/rawdatastore"
+	"golang.org/x/net/context"
 )
 
 //////////////////////////////// dataStoreData /////////////////////////////////
@@ -49,19 +48,19 @@ func (d *dataStoreData) Unlock() {
 
 /////////////////////////// indicies(dataStoreData) ////////////////////////////
 
-func groupMetaKey(key gae.DSKey) []byte {
-	return keyBytes(helper.WithoutContext,
-		helper.NewDSKey("", "", "__entity_group__", "", 1, helper.DSKeyRoot(key)))
+func groupMetaKey(key rds.Key) []byte {
+	return keyBytes(rds.WithoutContext,
+		rds.NewKey("", "", "__entity_group__", "", 1, rds.KeyRoot(key)))
 }
 
-func groupIDsKey(key gae.DSKey) []byte {
-	return keyBytes(helper.WithoutContext,
-		helper.NewDSKey("", "", "__entity_group_ids__", "", 1, helper.DSKeyRoot(key)))
+func groupIDsKey(key rds.Key) []byte {
+	return keyBytes(rds.WithoutContext,
+		rds.NewKey("", "", "__entity_group_ids__", "", 1, rds.KeyRoot(key)))
 }
 
 func rootIDsKey(kind string) []byte {
-	return keyBytes(helper.WithoutContext,
-		helper.NewDSKey("", "", "__entity_root_ids__", kind, 0, nil))
+	return keyBytes(rds.WithoutContext,
+		rds.NewKey("", "", "__entity_root_ids__", kind, 0, nil))
 }
 
 func curVersion(ents *memCollection, key []byte) int64 {
@@ -71,7 +70,7 @@ func curVersion(ents *memCollection, key []byte) int64 {
 			panic(err) // memory corruption
 		}
 		pl, ok := pm["__version__"]
-		if ok && len(pl) > 0 && pl[0].Type() == gae.DSPTInt {
+		if ok && len(pl) > 0 && pl[0].Type() == rds.PTInt {
 			return pl[0].Value().(int64)
 		}
 		panic(fmt.Errorf("__version__ property missing or wrong: %v", pm))
@@ -82,20 +81,20 @@ func curVersion(ents *memCollection, key []byte) int64 {
 func incrementLocked(ents *memCollection, key []byte) int64 {
 	ret := curVersion(ents, key) + 1
 	buf := &bytes.Buffer{}
-	helper.WriteDSPropertyMap(
-		buf, gae.DSPropertyMap{"__version__": {gae.MkDSPropertyNI(ret)}}, helper.WithContext)
+	rds.WritePropertyMap(
+		buf, rds.PropertyMap{"__version__": {rds.MkPropertyNI(ret)}}, rds.WithContext)
 	ents.Set(key, buf.Bytes())
 	return ret
 }
 
-func (d *dataStoreData) entsKeyLocked(key gae.DSKey) (*memCollection, gae.DSKey) {
+func (d *dataStoreData) entsKeyLocked(key rds.Key) (*memCollection, rds.Key) {
 	coll := "ents:" + key.Namespace()
 	ents := d.store.GetCollection(coll)
 	if ents == nil {
 		ents = d.store.SetCollection(coll, nil)
 	}
 
-	if helper.DSKeyIncomplete(key) {
+	if rds.KeyIncomplete(key) {
 		idKey := []byte(nil)
 		if key.Parent() == nil {
 			idKey = rootIDsKey(key.Kind())
@@ -103,21 +102,21 @@ func (d *dataStoreData) entsKeyLocked(key gae.DSKey) (*memCollection, gae.DSKey)
 			idKey = groupIDsKey(key)
 		}
 		id := incrementLocked(ents, idKey)
-		key = helper.NewDSKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
+		key = rds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
 	}
 
 	return ents, key
 }
 
-func (d *dataStoreData) put(ns string, key gae.DSKey, pls gae.DSPropertyLoadSaver) (gae.DSKey, error) {
-	keys, errs := d.putMulti(ns, []gae.DSKey{key}, []gae.DSPropertyLoadSaver{pls})
+func (d *dataStoreData) put(ns string, key rds.Key, pls rds.PropertyLoadSaver) (rds.Key, error) {
+	keys, errs := d.putMulti(ns, []rds.Key{key}, []rds.PropertyLoadSaver{pls})
 	if errs == nil {
 		return keys[0], nil
 	}
 	return nil, gae.SingleError(errs)
 }
 
-func (d *dataStoreData) putMulti(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver) ([]gae.DSKey, error) {
+func (d *dataStoreData) putMulti(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver) ([]rds.Key, error) {
 	pmaps, err := putMultiPrelim(ns, keys, plss)
 	if err != nil {
 		return nil, err
@@ -125,12 +124,12 @@ func (d *dataStoreData) putMulti(ns string, keys []gae.DSKey, plss []gae.DSPrope
 	return d.putMultiInner(keys, pmaps)
 }
 
-func putMultiPrelim(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver) ([]gae.DSPropertyMap, error) {
+func putMultiPrelim(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver) ([]rds.PropertyMap, error) {
 	err := multiValid(keys, plss, ns, true, false)
 	if err != nil {
 		return nil, err
 	}
-	pmaps := make([]gae.DSPropertyMap, len(keys))
+	pmaps := make([]rds.PropertyMap, len(keys))
 	lme := gae.LazyMultiError{Size: len(keys)}
 	for i, pls := range plss {
 		pm, err := pls.Save(false)
@@ -140,30 +139,30 @@ func putMultiPrelim(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver)
 	return pmaps, lme.Get()
 }
 
-func (d *dataStoreData) putMultiInner(keys []gae.DSKey, data []gae.DSPropertyMap) ([]gae.DSKey, error) {
-	retKeys := make([]gae.DSKey, len(keys))
+func (d *dataStoreData) putMultiInner(keys []rds.Key, data []rds.PropertyMap) ([]rds.Key, error) {
+	retKeys := make([]rds.Key, len(keys))
 	lme := gae.LazyMultiError{Size: len(keys)}
 	for i, k := range keys {
 		buf := &bytes.Buffer{}
-		helper.WriteDSPropertyMap(buf, data[i], helper.WithoutContext)
+		rds.WritePropertyMap(buf, data[i], rds.WithoutContext)
 		dataBytes := buf.Bytes()
 
-		rKey, err := func() (ret gae.DSKey, err error) {
+		rKey, err := func() (ret rds.Key, err error) {
 			d.rwlock.Lock()
 			defer d.rwlock.Unlock()
 
 			ents, ret := d.entsKeyLocked(k)
 			incrementLocked(ents, groupMetaKey(ret))
 
-			old := ents.Get(keyBytes(helper.WithoutContext, ret))
-			oldPM := gae.DSPropertyMap(nil)
+			old := ents.Get(keyBytes(rds.WithoutContext, ret))
+			oldPM := rds.PropertyMap(nil)
 			if old != nil {
 				if oldPM, err = rpmWoCtx(old, ret.Namespace()); err != nil {
 					return
 				}
 			}
 			updateIndicies(d.store, ret, oldPM, data[i])
-			ents.Set(keyBytes(helper.WithoutContext, ret), dataBytes)
+			ents.Set(keyBytes(rds.WithoutContext, ret), dataBytes)
 			return
 		}()
 		lme.Assign(i, err)
@@ -172,7 +171,7 @@ func (d *dataStoreData) putMultiInner(keys []gae.DSKey, data []gae.DSPropertyMap
 	return retKeys, lme.Get()
 }
 
-func getMultiInner(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver, getColl func() (*memCollection, error)) error {
+func getMultiInner(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver, getColl func() (*memCollection, error)) error {
 	if err := multiValid(keys, plss, ns, false, true); err != nil {
 		return err
 	}
@@ -185,15 +184,15 @@ func getMultiInner(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver,
 	}
 	if ents == nil {
 		for i := range keys {
-			lme.Assign(i, gae.ErrDSNoSuchEntity)
+			lme.Assign(i, rds.ErrNoSuchEntity)
 		}
 		return lme.Get()
 	}
 
 	for i, k := range keys {
-		pdata := ents.Get(keyBytes(helper.WithoutContext, k))
+		pdata := ents.Get(keyBytes(rds.WithoutContext, k))
 		if pdata == nil {
-			lme.Assign(i, gae.ErrDSNoSuchEntity)
+			lme.Assign(i, rds.ErrNoSuchEntity)
 			continue
 		}
 
@@ -208,11 +207,11 @@ func getMultiInner(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver,
 	return lme.Get()
 }
 
-func (d *dataStoreData) get(ns string, key gae.DSKey, pls gae.DSPropertyLoadSaver) error {
-	return gae.SingleError(d.getMulti(ns, []gae.DSKey{key}, []gae.DSPropertyLoadSaver{pls}))
+func (d *dataStoreData) get(ns string, key rds.Key, pls rds.PropertyLoadSaver) error {
+	return gae.SingleError(d.getMulti(ns, []rds.Key{key}, []rds.PropertyLoadSaver{pls}))
 }
 
-func (d *dataStoreData) getMulti(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver) error {
+func (d *dataStoreData) getMulti(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver) error {
 	return getMultiInner(ns, keys, plss, func() (*memCollection, error) {
 		d.rwlock.RLock()
 		s := d.store.Snapshot()
@@ -222,19 +221,19 @@ func (d *dataStoreData) getMulti(ns string, keys []gae.DSKey, plss []gae.DSPrope
 	})
 }
 
-func (d *dataStoreData) del(ns string, key gae.DSKey) (err error) {
-	return gae.SingleError(d.delMulti(ns, []gae.DSKey{key}))
+func (d *dataStoreData) del(ns string, key rds.Key) (err error) {
+	return gae.SingleError(d.delMulti(ns, []rds.Key{key}))
 }
 
-func (d *dataStoreData) delMulti(ns string, keys []gae.DSKey) error {
+func (d *dataStoreData) delMulti(ns string, keys []rds.Key) error {
 	lme := gae.LazyMultiError{Size: len(keys)}
 	toDel := make([][]byte, 0, len(keys))
 	for i, k := range keys {
-		if !helper.DSKeyValid(k, ns, false) {
-			lme.Assign(i, gae.ErrDSInvalidKey)
+		if !rds.KeyValid(k, ns, false) {
+			lme.Assign(i, rds.ErrInvalidKey)
 			continue
 		}
-		toDel = append(toDel, keyBytes(helper.WithoutContext, k))
+		toDel = append(toDel, keyBytes(rds.WithoutContext, k))
 	}
 	err := lme.Get()
 	if err != nil {
@@ -253,7 +252,7 @@ func (d *dataStoreData) delMulti(ns string, keys []gae.DSKey) error {
 		incrementLocked(ents, groupMetaKey(k))
 		kb := toDel[i]
 		old := ents.Get(kb)
-		oldPM := gae.DSPropertyMap(nil)
+		oldPM := rds.PropertyMap(nil)
 		if old != nil {
 			if oldPM, err = rpmWoCtx(old, ns); err != nil {
 				lme.Assign(i, err)
@@ -274,7 +273,7 @@ func (d *dataStoreData) canApplyTxn(obj memContextObj) bool {
 		if len(muts) == 0 { // read-only
 			continue
 		}
-		k, err := helper.ReadDSKey(bytes.NewBufferString(rk), helper.WithContext, "", "")
+		k, err := rds.ReadKey(bytes.NewBufferString(rk), rds.WithContext, "", "")
 		if err != nil {
 			panic(err)
 		}
@@ -312,7 +311,7 @@ func (d *dataStoreData) applyTxn(c context.Context, obj memContextObj) {
 	}
 }
 
-func (d *dataStoreData) mkTxn(o *gae.DSTransactionOptions) memContextObj {
+func (d *dataStoreData) mkTxn(o *rds.TransactionOptions) memContextObj {
 	return &txnDataStoreData{
 		// alias to the main datastore's so that testing code can have primitive
 		// access to break features inside of transactions.
@@ -328,8 +327,8 @@ func (d *dataStoreData) endTxn() {}
 /////////////////////////////// txnDataStoreData ///////////////////////////////
 
 type txnMutation struct {
-	key  gae.DSKey
-	data gae.DSPropertyMap
+	key  rds.Key
+	data rds.PropertyMap
 }
 
 type txnDataStoreData struct {
@@ -363,7 +362,7 @@ func (td *txnDataStoreData) endTxn() {
 func (*txnDataStoreData) applyTxn(context.Context, memContextObj) {
 	panic("txnDataStoreData cannot apply transactions")
 }
-func (*txnDataStoreData) mkTxn(*gae.DSTransactionOptions) memContextObj {
+func (*txnDataStoreData) mkTxn(*rds.TransactionOptions) memContextObj {
 	panic("impossible")
 }
 
@@ -387,8 +386,8 @@ func (td *txnDataStoreData) run(f func() error) error {
 //
 // Returns an error if this key causes the transaction to cross too many entity
 // groups.
-func (td *txnDataStoreData) writeMutation(getOnly bool, key gae.DSKey, data gae.DSPropertyMap) error {
-	rk := string(keyBytes(helper.WithContext, helper.DSKeyRoot(key)))
+func (td *txnDataStoreData) writeMutation(getOnly bool, key rds.Key, data rds.PropertyMap) error {
+	rk := string(keyBytes(rds.WithContext, rds.KeyRoot(key)))
 
 	td.Lock()
 	defer td.Unlock()
@@ -414,21 +413,21 @@ func (td *txnDataStoreData) writeMutation(getOnly bool, key gae.DSKey, data gae.
 	return nil
 }
 
-func (td *txnDataStoreData) put(ns string, key gae.DSKey, pls gae.DSPropertyLoadSaver) (gae.DSKey, error) {
-	keys, errs := td.putMulti(ns, []gae.DSKey{key}, []gae.DSPropertyLoadSaver{pls})
+func (td *txnDataStoreData) put(ns string, key rds.Key, pls rds.PropertyLoadSaver) (rds.Key, error) {
+	keys, errs := td.putMulti(ns, []rds.Key{key}, []rds.PropertyLoadSaver{pls})
 	if errs == nil {
 		return keys[0], nil
 	}
 	return nil, gae.SingleError(errs)
 }
 
-func (td *txnDataStoreData) putMulti(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver) ([]gae.DSKey, error) {
+func (td *txnDataStoreData) putMulti(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver) ([]rds.Key, error) {
 	pmaps, err := putMultiPrelim(ns, keys, plss)
 	if err != nil {
 		return nil, err
 	}
 
-	retKeys := make([]gae.DSKey, len(keys))
+	retKeys := make([]rds.Key, len(keys))
 	lme := gae.LazyMultiError{Size: len(keys)}
 	for i, k := range keys {
 		func() {
@@ -443,11 +442,11 @@ func (td *txnDataStoreData) putMulti(ns string, keys []gae.DSKey, plss []gae.DSP
 	return retKeys, lme.Get()
 }
 
-func (td *txnDataStoreData) get(ns string, key gae.DSKey, pls gae.DSPropertyLoadSaver) error {
-	return gae.SingleError(td.getMulti(ns, []gae.DSKey{key}, []gae.DSPropertyLoadSaver{pls}))
+func (td *txnDataStoreData) get(ns string, key rds.Key, pls rds.PropertyLoadSaver) error {
+	return gae.SingleError(td.getMulti(ns, []rds.Key{key}, []rds.PropertyLoadSaver{pls}))
 }
 
-func (td *txnDataStoreData) getMulti(ns string, keys []gae.DSKey, plss []gae.DSPropertyLoadSaver) error {
+func (td *txnDataStoreData) getMulti(ns string, keys []rds.Key, plss []rds.PropertyLoadSaver) error {
 	return getMultiInner(ns, keys, plss, func() (*memCollection, error) {
 		lme := gae.LazyMultiError{Size: len(keys)}
 		for i, k := range keys {
@@ -457,15 +456,15 @@ func (td *txnDataStoreData) getMulti(ns string, keys []gae.DSKey, plss []gae.DSP
 	})
 }
 
-func (td *txnDataStoreData) del(ns string, key gae.DSKey) error {
-	return gae.SingleError(td.delMulti(ns, []gae.DSKey{key}))
+func (td *txnDataStoreData) del(ns string, key rds.Key) error {
+	return gae.SingleError(td.delMulti(ns, []rds.Key{key}))
 }
 
-func (td *txnDataStoreData) delMulti(ns string, keys []gae.DSKey) error {
+func (td *txnDataStoreData) delMulti(ns string, keys []rds.Key) error {
 	lme := gae.LazyMultiError{Size: len(keys)}
 	for i, k := range keys {
-		if !helper.DSKeyValid(k, ns, false) {
-			lme.Assign(i, gae.ErrDSInvalidKey)
+		if !rds.KeyValid(k, ns, false) {
+			lme.Assign(i, rds.ErrInvalidKey)
 		} else {
 			lme.Assign(i, td.writeMutation(false, k, nil))
 		}
@@ -473,31 +472,31 @@ func (td *txnDataStoreData) delMulti(ns string, keys []gae.DSKey) error {
 	return lme.Get()
 }
 
-func keyBytes(ctx helper.DSKeyContext, key gae.DSKey) []byte {
+func keyBytes(ctx rds.KeyContext, key rds.Key) []byte {
 	buf := &bytes.Buffer{}
-	helper.WriteDSKey(buf, ctx, key)
+	rds.WriteKey(buf, ctx, key)
 	return buf.Bytes()
 }
 
-func rpmWoCtx(data []byte, ns string) (gae.DSPropertyMap, error) {
-	return helper.ReadDSPropertyMap(bytes.NewBuffer(data), helper.WithoutContext, globalAppID, ns)
+func rpmWoCtx(data []byte, ns string) (rds.PropertyMap, error) {
+	return rds.ReadPropertyMap(bytes.NewBuffer(data), rds.WithoutContext, globalAppID, ns)
 }
 
-func rpm(data []byte) (gae.DSPropertyMap, error) {
-	return helper.ReadDSPropertyMap(bytes.NewBuffer(data), helper.WithContext, "", "")
+func rpm(data []byte) (rds.PropertyMap, error) {
+	return rds.ReadPropertyMap(bytes.NewBuffer(data), rds.WithContext, "", "")
 }
 
-func multiValid(keys []gae.DSKey, plss []gae.DSPropertyLoadSaver, ns string, potentialKey, allowSpecial bool) error {
-	vfn := func(k gae.DSKey) bool {
-		return !helper.DSKeyIncomplete(k) && helper.DSKeyValid(k, ns, allowSpecial)
+func multiValid(keys []rds.Key, plss []rds.PropertyLoadSaver, ns string, potentialKey, allowSpecial bool) error {
+	vfn := func(k rds.Key) bool {
+		return !rds.KeyIncomplete(k) && rds.KeyValid(k, ns, allowSpecial)
 	}
 	if potentialKey {
-		vfn = func(k gae.DSKey) bool {
+		vfn = func(k rds.Key) bool {
 			// adds an id to k if it's incomplete.
-			if helper.DSKeyIncomplete(k) {
-				k = helper.NewDSKey(k.AppID(), k.Namespace(), k.Kind(), "", 1, k.Parent())
+			if rds.KeyIncomplete(k) {
+				k = rds.NewKey(k.AppID(), k.Namespace(), k.Kind(), "", 1, k.Parent())
 			}
-			return helper.DSKeyValid(k, ns, allowSpecial)
+			return rds.KeyValid(k, ns, allowSpecial)
 		}
 	}
 
@@ -510,7 +509,7 @@ func multiValid(keys []gae.DSKey, plss []gae.DSPropertyLoadSaver, ns string, pot
 	lme := gae.LazyMultiError{Size: len(keys)}
 	for i, k := range keys {
 		if !vfn(k) {
-			lme.Assign(i, gae.ErrDSInvalidKey)
+			lme.Assign(i, rds.ErrInvalidKey)
 		}
 	}
 	return lme.Get()