| // 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/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) |
| } |