go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/service/datastore/dropped_arg_tracker.go

// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package datastore

import (
	"sort"

	"go.chromium.org/luci/common/errors"
)

// DroppedArgTracker is used to track dropping items from Keys as well as meta
// and/or PropertyMap arrays from one layer of the RawInterface to the next.
//
// If you're not writing a datastore backend implementation (like
// "go.chromium.org/luci/gae/impl/*"), then you can ignore this type.
//
// For example, say your GetMulti method was passed 4 arguments, but one of them
// was bad. DroppedArgTracker would allow you to "drop" the bad entry, and then
// synthesize new keys/meta/values arrays excluding the bad entry. You could
// then map from the new arrays back to the indexes of the original arrays.
//
// This DroppedArgTracker will do no allocations if you don't end up dropping
// any arguments (so in the 'good' case, there are zero allocations).
//
// Example:
//
//	  Say we're given a list of arguments which look like ("_" means a bad value
//	  that we drop):
//
//	   input: A B _ C D _ _ E
//	    Idxs: 0 1 2 3 4 5 6 7
//	 dropped:     2     5 6
//
//	DropKeys(input): A B C D E
//	                 0 1 2 3 4
//
//	OriginalIndex(0) -> 0
//	OriginalIndex(1) -> 1
//	OriginalIndex(2) -> 3
//	OriginalIndex(3) -> 4
//	OriginalIndex(4) -> 7
//
// Methods on this type are NOT goroutine safe.
type DroppedArgTracker []int

// MarkForRemoval tracks `originalIndex` for removal when `Drop*` methods
// are called.
//
// N is a size hint for the maximum number of entries that `dat` could have. If
// `dat` has a capacity of < N, it will be allocated to N.
//
// If called with N == len(args) and originalIndex is always increasing, then
// this will only do one allocation for the life of this DroppedArgTracker, and
// each MarkForRemoval will only cost a single slice append. If called out of
// order, or with a bad value of N, this will do more allocations and will do
// a binary search on each call.
func (dat *DroppedArgTracker) MarkForRemoval(originalIndex, N int) {
	datLen := len(*dat)

	if cap(*dat) < N {
		newDat := make([]int, datLen, N)
		copy(newDat, *dat)
		*dat = newDat
	}

	// most uses will insert linearly, so do a quick check of the max element to
	// see if originalIndex is larger and then do a simple append.
	if datLen == 0 || originalIndex > (*dat)[datLen-1] {
		*dat = append(*dat, originalIndex)
		return
	}

	// Otherwise, search for the correct location and insert it
	insIdx := sort.SearchInts(*dat, originalIndex)
	if insIdx < datLen && (*dat)[insIdx] == originalIndex {
		return
	}
	*dat = append(*dat, 0)
	copy((*dat)[insIdx+1:], (*dat)[insIdx:])
	(*dat)[insIdx] = originalIndex
}

// MarkNilKeys is a helper method which calls MarkForRemoval for each nil key.
func (dat *DroppedArgTracker) MarkNilKeys(keys []*Key) {
	for idx, k := range keys {
		if k == nil {
			dat.MarkForRemoval(idx, len(keys))
		}
	}
}

// MarkNilKeysMeta is a helper method which calls MarkForRemoval for each nil
// key or meta.
func (dat *DroppedArgTracker) MarkNilKeysMeta(keys []*Key, meta MultiMetaGetter) {
	for idx, k := range keys {
		if k == nil || meta[idx] == nil {
			dat.MarkForRemoval(idx, len(keys))
		}
	}
}

// MarkNilKeysVals is a helper method which calls MarkForRemoval for each nil
// key or value.
func (dat *DroppedArgTracker) MarkNilKeysVals(keys []*Key, vals []PropertyMap) {
	for idx, k := range keys {
		if k == nil || vals[idx] == nil {
			dat.MarkForRemoval(idx, len(keys))
		}
	}
}

// If `dat` has a positive length, this will invoke `init` once, followed by
// `include` for every non-overlapping (i, j) range less than N which doesn't
// include any elements indicated with MarkForRemoval.
//
// If `dat` contains a removed index larger than N, this panics.
func (dat DroppedArgTracker) mustCompress(N int, init func(), include func(i, j int)) DroppedArgLookup {
	if len(dat) == 0 || N == 0 {
		return nil
	}

	if largestDropIdx := dat[len(dat)-1]; largestDropIdx >= N {
		panic(errors.Reason(
			"DroppedArgTracker has out of bound index: %d >= %d ",
			largestDropIdx, N,
		).Err())
	}

	// dal may have len < len(dat) in the event that multiple dat entries are
	// contiguous (they'll be compressed into a single entry in dal).
	dal := make(DroppedArgLookup, 0, len(dat))

	init()
	nextPotentialOriginalInclusion := 0

	for numDroppedSoFar, originalIndexToDrop := range dat {
		if originalIndexToDrop > nextPotentialOriginalInclusion {
			include(nextPotentialOriginalInclusion, originalIndexToDrop)
		}

		nextPotentialOriginalInclusion = originalIndexToDrop + 1
		reducedIndex := originalIndexToDrop - numDroppedSoFar

		if len(dal) != 0 && dal[len(dal)-1].reducedIndex == reducedIndex {
			// If the user drops multiple original indices in a row, we need to
			// reflect them in a single entry in dal.
			dal[len(dal)-1].originalIndex++
		} else {
			// otherwise, we make a new entry.
			dal = append(dal, idxPair{
				reducedIndex,
				nextPotentialOriginalInclusion,
			})
		}
	}
	include(nextPotentialOriginalInclusion, N)

	return dal
}

// DropKeys returns a compressed version of `keys`, dropping all elements which
// were marked with MarkForRemoval.
func (dat DroppedArgTracker) DropKeys(keys []*Key) ([]*Key, DroppedArgLookup) {
	newKeys := keys

	init := func() {
		newKeys = make([]*Key, 0, len(keys)-len(dat))
	}
	include := func(i, j int) {
		newKeys = append(newKeys, keys[i:j]...)
	}

	dal := dat.mustCompress(len(keys), init, include)
	return newKeys, dal
}

// DropKeysAndMeta returns a compressed version of `keys` and `meta`, dropping
// all elements which were marked with MarkForRemoval.
//
// `keys` and `meta` must have the same lengths.
func (dat DroppedArgTracker) DropKeysAndMeta(keys []*Key, meta MultiMetaGetter) ([]*Key, MultiMetaGetter, DroppedArgLookup) {
	newKeys := keys
	newMeta := meta

	// MultiMetaGetter is special and frequently is len 0 with non-nil keys, so we
	// just keep it empty.

	init := func() {
		newKeys = make([]*Key, 0, len(keys)-len(dat))
		if len(meta) > 0 {
			newMeta = make(MultiMetaGetter, 0, len(keys)-len(dat))
		}
	}
	include := func(i, j int) {
		newKeys = append(newKeys, keys[i:j]...)
		if len(meta) > 0 {
			newMeta = append(newMeta, meta[i:j]...)
		}
	}

	dal := dat.mustCompress(len(keys), init, include)
	return newKeys, newMeta, dal
}

// DropKeysAndVals returns a compressed version of `keys` and `vals`, dropping
// all elements which were marked with MarkForRemoval.
//
// `keys` and `vals` must have the same lengths.
func (dat DroppedArgTracker) DropKeysAndVals(keys []*Key, vals []PropertyMap) ([]*Key, []PropertyMap, DroppedArgLookup) {
	newKeys := keys
	newVals := vals

	if len(keys) != len(vals) {
		panic(errors.Reason(
			"DroppedArgTracker.DropKeysAndVals: mismatched lengths: %d vs %d",
			len(keys), len(vals),
		).Err())
	}

	init := func() {
		newKeys = make([]*Key, 0, len(keys)-len(dat))
		newVals = make([]PropertyMap, 0, len(keys)-len(dat))
	}
	include := func(i, j int) {
		newKeys = append(newKeys, keys[i:j]...)
		newVals = append(newVals, vals[i:j]...)
	}

	dal := dat.mustCompress(len(keys), init, include)
	return newKeys, newVals, dal
}

type idxPair struct {
	reducedIndex  int
	originalIndex int
}

// DroppedArgLookup is returned from using a DroppedArgTracker.
//
// It can be used to recover the index from the original slice by providing the
// reduced slice index.
type DroppedArgLookup []idxPair

// OriginalIndex maps from an index into the array(s) returned from MustDrop
// back to the corresponding index in the original arrays.
func (dal DroppedArgLookup) OriginalIndex(reducedIndex int) int {
	if len(dal) == 0 {
		return reducedIndex
	}

	// Search for the idxPair whose reducedIndex is LARGER than what we want.
	dalInsertIdx := sort.Search(len(dal), func(dalIdx int) bool {
		return dal[dalIdx].reducedIndex > reducedIndex
	})
	// If search told us that it was "0" it means that no entry in `dal` includes
	// our reducedIndex.
	if dalInsertIdx == 0 {
		return reducedIndex
	}
	// Now look up the idxPair before what search returned. This will have
	// a reducedIndex which is <= reducedIndex.
	entry := dal[dalInsertIdx-1]
	if entry.reducedIndex > reducedIndex {
		return reducedIndex
	}
	// Finally, return the originalIndex and add the difference between our user's
	// reducedIndex and the one in this entry.
	return entry.originalIndex + (reducedIndex - entry.reducedIndex)
}