datastore: variadic Get, Put, Exists, Delete.

service/datastore/multiarg.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package datastore
 
 import (
         "fmt"
         "reflect"
 
         "github.com/luci/luci-go/common/errors"
 )
 
 type multiArgType struct {
         getKey    func(aid, ns string, slot reflect.Value) (*Key, error)
         getPM     func(slot reflect.Value) (PropertyMap, error)
         getMetaPM func(slot reflect.Value) PropertyMap
         setPM     func(slot reflect.Value, pm PropertyMap) error
         setKey    func(slot reflect.Value, k *Key)
         newElem   func() reflect.Value
 }
 
-func (mat *multiArgType) GetKeysPMs(aid, ns string, slice reflect.Value, meta bool) ([]*Key, []PropertyMap, error) {

[dnj, 2016/05/25 05:27:16]

This has been moved to the bottom, and is now a function in metaMultiArg.

-        retKey := make([]*Key, slice.Len())
-        retPM := make([]PropertyMap, slice.Len())
-        getter := mat.getPM
-        if meta {
-                getter = func(slot reflect.Value) (PropertyMap, error) {
-                        return mat.getMetaPM(slot), nil
+// parseArg checks that et is of type S, *S, I, P or *P, for some
+// struct type S, for some interface type I, or some non-interface non-pointer
+// type P such that P or *P implements PropertyLoadSaver.
+func parseArg(et reflect.Type) *multiArgType {
+        if et.Kind() == reflect.Interface {
+                return multiArgTypeInterface()
+        }
+
+        // If a map type implements an interface, its pointer is also considered to

[dnj, 2016/05/25 05:27:16]

(This is why the map thing was so wonky. I didn't know this, but see:
https://play.golang.org/p/mV6UDNsF_0 )

+        // implement that interface.
+        //
+        // In this case, we have special pointer-to-map logic in multiArgTypePLS.
+        if et.Implements(typeOfPropertyLoadSaver) {
+                return multiArgTypePLS(et)
+        }
+        if reflect.PtrTo(et).Implements(typeOfPropertyLoadSaver) {
+                return multiArgTypePLSPtr(et)
+        }
+
+        switch et.Kind() {
+        case reflect.Ptr:
+                if et.Elem().Kind() == reflect.Struct {
+                        return multiArgTypeStructPtr(et)
                 }
+
+        case reflect.Struct:
+                return multiArgTypeStruct(et)
         }
-        lme := errors.NewLazyMultiError(len(retKey))
-        for i := range retKey {
-                key, err := mat.getKey(aid, ns, slice.Index(i))
-                if !lme.Assign(i, err) {
-                        retKey[i] = key
-                        pm, err := getter(slice.Index(i))
-                        if !lme.Assign(i, err) {
-                                retPM[i] = pm
-                        }
-                }
-        }
-        return retKey, retPM, lme.Get()
+
+        return nil
 }
 
 // parseMultiArg checks that v has type []S, []*S, []I, []P or []*P, for some
 // struct type S, for some interface type I, or some non-interface non-pointer
 // type P such that P or *P implements PropertyLoadSaver.
-func parseMultiArg(e reflect.Type) multiArgType {
-        if e.Kind() != reflect.Slice {
-                panic(fmt.Errorf("invalid argument type: expected slice, got %s", e))
+func mustParseMultiArg(et reflect.Type) *multiArgType {
+        if et.Kind() != reflect.Slice {
+                panic(fmt.Errorf("invalid argument type: expected slice, got %s", et))
         }
-        return parseArg(e.Elem(), true)
+        return mustParseArg(et.Elem())
 }
 
-// parseArg checks that et is of type S, *S, I, P or *P, for some
-// struct type S, for some interface type I, or some non-interface non-pointer
-// type P such that P or *P implements PropertyLoadSaver.
-func parseArg(et reflect.Type, multi bool) multiArgType {
-        if reflect.PtrTo(et).Implements(typeOfPropertyLoadSaver) {
-                return multiArgTypePLS(et)
+func mustParseArg(et reflect.Type) *multiArgType {
+        if mat := parseArg(et); mat != nil {
+                return mat
         }
-        if et.Implements(typeOfPropertyLoadSaver) && et.Kind() != reflect.Interface {
-                return multiArgTypePLSPtr(et.Elem())
-        }
-        switch et.Kind() {
-        case reflect.Struct:
-                return multiArgTypeStruct(et)
-        case reflect.Interface:
-                return multiArgTypeInterface()
-        case reflect.Ptr:
-                et = et.Elem()
-                if et.Kind() == reflect.Struct {
-                        return multiArgTypeStructPtr(et)
-                }
-        }
-        if multi {
-                panic(fmt.Errorf("invalid argument type: []%s", et))
-        }
-        panic(fmt.Errorf("invalid argument type: %s", et))
+        panic(fmt.Errorf("invalid argument type: %s is not a PLS or pointer-to-struct", et))
 }
 
-type newKeyFunc func(kind, sid string, iid int64, par Key) Key
-
-// multiArgTypePLS == []P
-// *P implements PropertyLoadSaver
-func multiArgTypePLS(et reflect.Type) multiArgType {
+// multiArgTypePLS handles the case where et implements PropertyLoadSaver.
+//
+// This handles the special case of pointer-to-map (see parseArg).
+func multiArgTypePLS(et reflect.Type) *multiArgType {
         ret := multiArgType{
                 getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
-                        return newKeyObjErr(aid, ns, slot.Addr().Interface())
-                },
-                getPM: func(slot reflect.Value) (PropertyMap, error) {
-                        return slot.Addr().Interface().(PropertyLoadSaver).Save(true)
-                },
-                getMetaPM: func(slot reflect.Value) PropertyMap {
-                        return getMGS(slot.Addr().Interface()).GetAllMeta()
-                },
-                setPM: func(slot reflect.Value, pm PropertyMap) error {
-                        return slot.Addr().Interface().(PropertyLoadSaver).Load(pm)
-                },
-                setKey: func(slot reflect.Value, k *Key) {
-                        PopulateKey(slot.Addr().Interface(), k)
-                },
-        }
-        if et.Kind() == reflect.Map {
-                ret.newElem = func() reflect.Value {
-                        // Create a *map so that way slot.Addr() works above when this is
-                        // called from Run(). Otherwise the map is 'unaddressable' according
-                        // to reflect.  ¯\_(ツ)_/¯
-                        ptr := reflect.New(et)
-                        ptr.Elem().Set(reflect.MakeMap(et))
-                        return ptr.Elem()
-                }
-        } else {
-                ret.newElem = func() reflect.Value {
-                        return reflect.New(et).Elem()
-                }
-        }
-        return ret
-}
-
-// multiArgTypePLSPtr == []*P
-// *P implements PropertyLoadSaver
-func multiArgTypePLSPtr(et reflect.Type) multiArgType {
-        ret := multiArgType{
-                getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
-                        return newKeyObjErr(aid, ns, slot.Interface())
+                        return newKeyObjErr(aid, ns, getMGS(slot.Interface()))
                 },
                 getPM: func(slot reflect.Value) (PropertyMap, error) {
                         return slot.Interface().(PropertyLoadSaver).Save(true)
                 },
                 getMetaPM: func(slot reflect.Value) PropertyMap {
                         return getMGS(slot.Interface()).GetAllMeta()
                 },
                 setPM: func(slot reflect.Value, pm PropertyMap) error {
                         return slot.Interface().(PropertyLoadSaver).Load(pm)
                 },
                 setKey: func(slot reflect.Value, k *Key) {
                         PopulateKey(slot.Interface(), k)
                 },
         }
-        if et.Kind() == reflect.Map {
+        switch et.Kind() {
+        case reflect.Map:
                 ret.newElem = func() reflect.Value {
-                        ptr := reflect.New(et)
-                        ptr.Elem().Set(reflect.MakeMap(et))
-                        return ptr
+                        return reflect.MakeMap(et)
                 }
-        } else {
-                ret.newElem = func() reflect.Value { return reflect.New(et) }
+
+        case reflect.Ptr:
+                mapElem := et.Elem()
+                if mapElem.Kind() == reflect.Map {
+                        ret.newElem = func() reflect.Value {
+                                ptr := reflect.New(mapElem)
+                                ptr.Elem().Set(reflect.MakeMap(mapElem))
+                                return ptr
+                        }
+                }
         }
-        return ret
+
+        if ret.newElem == nil {
+                ret.newElem = func() reflect.Value {
+                        return reflect.New(et.Elem())
+                }
+        }
+        return &ret
+}
+
+// multiArgTypePLSPtr handles the case where et doesn't implement
+// PropertyLoadSaver, but a pointer to et does.
+func multiArgTypePLSPtr(et reflect.Type) *multiArgType {
+        return &multiArgType{
+                getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
+                        return newKeyObjErr(aid, ns, getMGS(slot.Addr().Interface()))
+                },
+                getPM: func(slot reflect.Value) (PropertyMap, error) {
+                        return slot.Addr().Interface().(PropertyLoadSaver).Save(true)
+                },
+                getMetaPM: func(slot reflect.Value) PropertyMap {
+                        return getMGS(slot.Addr().Interface()).GetAllMeta()
+                },
+                setPM: func(slot reflect.Value, pm PropertyMap) error {
+                        return slot.Addr().Interface().(PropertyLoadSaver).Load(pm)
+                },
+                setKey: func(slot reflect.Value, k *Key) {
+                        PopulateKey(slot.Addr().Interface(), k)
+                },
+                newElem: func() reflect.Value {
+                        return reflect.New(et).Elem()
+                },
+        }
 }
 
 // multiArgTypeStruct == []S
-func multiArgTypeStruct(et reflect.Type) multiArgType {
+func multiArgTypeStruct(et reflect.Type) *multiArgType {
         cdc := getCodec(et)
         toPLS := func(slot reflect.Value) *structPLS {
                 return &structPLS{slot, cdc}
         }
-        return multiArgType{
+        return &multiArgType{
                 getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
-                        return newKeyObjErr(aid, ns, toPLS(slot))
+                        return newKeyObjErr(aid, ns, getMGS(slot.Addr().Interface()))

[dnj, 2016/05/25 05:27:16]

This was a bug. MGS overrides wouldn't be honored. New vararg tests this now.

                 },
                 getPM: func(slot reflect.Value) (PropertyMap, error) {
                         return toPLS(slot).Save(true)
                 },
                 getMetaPM: func(slot reflect.Value) PropertyMap {
-                        if slot.Type().Implements(typeOfMGS) {
-                                return slot.Interface().(MetaGetterSetter).GetAllMeta()
-                        }
-                        return toPLS(slot).GetAllMeta()
+                        return getMGS(slot.Addr().Interface()).GetAllMeta()
                 },
                 setPM: func(slot reflect.Value, pm PropertyMap) error {
                         return toPLS(slot).Load(pm)
                 },
                 setKey: func(slot reflect.Value, k *Key) {
                         PopulateKey(toPLS(slot), k)
                 },
                 newElem: func() reflect.Value {
                         return reflect.New(et).Elem()
                 },
         }
 }
 
 // multiArgTypeStructPtr == []*S
-func multiArgTypeStructPtr(et reflect.Type) multiArgType {
-        cdc := getCodec(et)
+func multiArgTypeStructPtr(et reflect.Type) *multiArgType {
+        cdc := getCodec(et.Elem())
         toPLS := func(slot reflect.Value) *structPLS {
                 return &structPLS{slot.Elem(), cdc}
         }
-        return multiArgType{
+        return &multiArgType{
                 getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
-                        return newKeyObjErr(aid, ns, toPLS(slot))
+                        return newKeyObjErr(aid, ns, getMGS(slot.Interface()))
                 },
                 getPM: func(slot reflect.Value) (PropertyMap, error) {
                         return toPLS(slot).Save(true)
                 },
                 getMetaPM: func(slot reflect.Value) PropertyMap {
-                        if slot.Elem().Type().Implements(typeOfMGS) {
-                                return getMGS(slot.Interface()).GetAllMeta()
-                        }
-                        return toPLS(slot).GetAllMeta()
+                        return getMGS(slot.Interface()).GetAllMeta()
                 },
                 setPM: func(slot reflect.Value, pm PropertyMap) error {
-                        return toPLS(slot).Load(pm)
+                        err := toPLS(slot).Load(pm)

[dnj, 2016/05/25 05:27:16]

I'll collapse this in the next PS.

+                        return err
                 },
                 setKey: func(slot reflect.Value, k *Key) {
                         PopulateKey(toPLS(slot), k)
                 },
                 newElem: func() reflect.Value {
-                        return reflect.New(et)
+                        return reflect.New(et.Elem())
                 },
         }
 }
 
 // multiArgTypeInterface == []I
-func multiArgTypeInterface() multiArgType {
-        return multiArgType{
+func multiArgTypeInterface() *multiArgType {
+        return &multiArgType{
                 getKey: func(aid, ns string, slot reflect.Value) (*Key, error) {
-                        return newKeyObjErr(aid, ns, slot.Elem().Interface())
+                        return newKeyObjErr(aid, ns, getMGS(slot.Elem().Interface()))

[dnj, 2016/05/25 05:27:16]

This, I think, was a bug.

                 },
                 getPM: func(slot reflect.Value) (PropertyMap, error) {
                         return mkPLS(slot.Elem().Interface()).Save(true)
                 },
                 getMetaPM: func(slot reflect.Value) PropertyMap {
                         return getMGS(slot.Elem().Interface()).GetAllMeta()
                 },
                 setPM: func(slot reflect.Value, pm PropertyMap) error {
                         return mkPLS(slot.Elem().Interface()).Load(pm)
                 },
                 setKey: func(slot reflect.Value, k *Key) {
                         PopulateKey(slot.Elem().Interface(), k)
                 },
         }
 }
 
-func newKeyObjErr(aid, ns string, src interface{}) (*Key, error) {
-        pls := getMGS(src)
-        if key, _ := GetMetaDefault(pls, "key", nil).(*Key); key != nil {
+func newKeyObjErr(aid, ns string, mgs MetaGetterSetter) (*Key, error) {
+        if key, _ := GetMetaDefault(mgs, "key", nil).(*Key); key != nil {
                 return key, nil
         }
 
         // get kind
-        kind := GetMetaDefault(pls, "kind", "").(string)
+        kind := GetMetaDefault(mgs, "kind", "").(string)
         if kind == "" {
-                return nil, fmt.Errorf("unable to extract $kind from %T", src)
+                return nil, errors.New("unable to extract $kind")
         }
 
         // get id - allow both to be default for default keys
-        sid := GetMetaDefault(pls, "id", "").(string)
-        iid := GetMetaDefault(pls, "id", 0).(int64)
+        sid := GetMetaDefault(mgs, "id", "").(string)
+        iid := GetMetaDefault(mgs, "id", 0).(int64)
 
         // get parent
-        par, _ := GetMetaDefault(pls, "parent", nil).(*Key)
+        par, _ := GetMetaDefault(mgs, "parent", nil).(*Key)
 
         return NewKey(aid, ns, kind, sid, iid, par), nil
 }
 
 func mkPLS(o interface{}) PropertyLoadSaver {
         if pls, ok := o.(PropertyLoadSaver); ok {
                 return pls
         }
         return GetPLS(o)
 }
+
+func isOKSingleType(t reflect.Type) error {
+        if t == nil {
+                return errors.New("no type information")
+        }
+        if t.Implements(typeOfPropertyLoadSaver) {
+                return nil
+        }
+        if t == typeOfKey {
+                return errors.New("not user datatype")
+        }
+        if t.Kind() != reflect.Ptr {
+                return errors.New("not a pointer")
+        }
+        if t.Elem().Kind() != reflect.Struct {
+                return errors.New("does not point to a struct")
+        }
+        return nil
+}
+
+type metaMultiArgElement struct {
+        arg  reflect.Value
+        mat  *multiArgType
+        size int // size is -1 if this element is not a slice.
+}
+
+type metaMultiArg struct {

[dnj, 2016/05/25 05:27:16]

*so meta*

But if there's a better name I might be open to it.

This should be pretty efficient.

+        elems []metaMultiArgElement
+
+        count int // total number of elements, flattening slices
+}
+
+func makeMetaMultiArg(args []interface{}) (*metaMultiArg, error) {
+        mma := metaMultiArg{
+                elems: make([]metaMultiArgElement, len(args)),
+        }
+
+        lme := errors.NewLazyMultiError(len(args))
+        for i, arg := range args {
+                if arg == nil {
+                        lme.Assign(i, errors.New("cannot use nil as single argument"))
+                        continue
+                }
+
+                v := reflect.ValueOf(arg)
+                vt := v.Type()
+                mma.elems[i].arg = v
+
+                // Try and treat the argument as a single-value first. This allows slices
+                // that implement PropertyLoadSaver to be properly treated as a single
+                // element.
+                var err error
+                isSlice := false
+                mat := parseArg(vt)
+                if mat == nil {
+                        // If this is a slice, treat it as a slice of arg candidates.
+                        if v.Kind() == reflect.Slice {
+                                isSlice = true
+                                mat = parseArg(vt.Elem())
+                        }
+                } else {
+                        // Single types need to be able to be assigned to.
+                        err = isOKSingleType(vt)
+                }
+                if mat == nil {
+                        err = errors.New("not a PLS or pointer-to-struct")
+                }
+                if err != nil {
+                        lme.Assign(i, fmt.Errorf("invalid input type (%T): %s", arg, err))
+                        continue
+                }
+
+                mma.elems[i].mat = mat
+                if isSlice {
+                        l := v.Len()
+                        mma.count += l
+                        mma.elems[i].size = l
+                } else {
+                        mma.count++
+                        mma.elems[i].size = -1
+                }
+        }
+        if err := lme.Get(); err != nil {
+                return nil, err
+        }
+
+        return &mma, nil
+}
+
+func (mma *metaMultiArg) iterator() *metaMultiArgIterator {
+        return &metaMultiArgIterator{
+                metaMultiArg: mma,
+
+                elemErrors: errors.NewLazyMultiError(len(mma.elems)),
+        }
+}
+
+// getKeysPMs returns the
+func (mma *metaMultiArg) getKeysPMs(aid, ns string, meta bool) ([]*Key, []PropertyMap, error) {
+        it := mma.iterator()
+
+        // Determine our flattened keys and property maps.
+        retKey := make([]*Key, mma.count)
+        retPM := make([]PropertyMap, mma.count)
+
+        for i := 0; i < mma.count; i++ {
+                it.next(func(mat *multiArgType, slot reflect.Value) error {
+                        key, err := mat.getKey(aid, ns, slot)
+                        if err != nil {
+                                return err
+                        }
+                        retKey[i] = key
+
+                        var pm PropertyMap
+                        if meta {
+                                pm = mat.getMetaPM(slot)
+                        } else {
+                                var err error
+                                if pm, err = mat.getPM(slot); err != nil {
+                                        return err
+                                }
+                        }
+                        retPM[i] = pm
+                        return nil
+                })
+        }
+        return retKey, retPM, it.error()
+}
+
+type metaMultiArgIterator struct {
+        *metaMultiArg
+
+        elemErrors  errors.LazyMultiError
+        sliceErrors map[int]errors.LazyMultiError
+
+        index   int // flattened index
+        elemIdx int // current index in slice
+        slotIdx int // current index within elemIdx element (0 if single)
+}
+
+func (mac *metaMultiArgIterator) next(fn func(*multiArgType, reflect.Value) error) {
+        if mac.remaining() <= 0 {
+                panic("out of bounds")
+        }
+
+        // Get the current element and slot.
+        elem := &mac.elems[mac.elemIdx]
+
+        // Get the current slot value.
+        slot := elem.arg
+        if elem.size >= 0 {
+                // slot is a slice type, get its member.
+                slot = slot.Index(mac.slotIdx)
+        }
+
+        // Execute our callback.
+        if err := fn(elem.mat, slot); err != nil {
+                // If this is a single element, assign the error directly.
+                if elem.size < 0 {
+                        mac.elemErrors.Assign(mac.elemIdx, err)
+                } else {
+                        // This is a slice element. Use a slice-sized MultiError for its element
+                        // error slot, then add this error to the inner MultiError's slot index.
+                        ilme := mac.sliceErrors[mac.elemIdx]
+                        if ilme == nil {
+                                ilme = errors.NewLazyMultiError(elem.size)
+
+                                if mac.sliceErrors == nil {
+                                        mac.sliceErrors = make(map[int]errors.LazyMultiError)
+                                }
+                                mac.sliceErrors[mac.elemIdx] = ilme
+                        }
+                        ilme.Assign(mac.slotIdx, err)
+                }
+        }
+
+        // Advance to the next element/slot.
+        mac.index++
+        mac.slotIdx++
+        if mac.slotIdx >= elem.size {
+                mac.elemIdx++
+                mac.slotIdx = 0
+        }
+}
+
+func (mac *metaMultiArgIterator) error() error {
+        lme := errors.NewLazyMultiError(len(mac.elems))
+        for i, elem := range mac.elems {
+                if elem.size < 0 {
+                        lme.Assign(i, mac.elemErrors.GetOne(i))
+                } else {
+                        if serr := mac.sliceErrors[i]; serr != nil {
+                                lme.Assign(i, serr.Get())
+                        }
+                }
+        }
+        return lme.Get()
+}
+
+func (mac *metaMultiArgIterator) remaining() int {
+        return mac.count - mac.index
+}