github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/util/slices.go (about) 1 // Copyright 2021 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package util 12 13 import ( 14 "golang.org/x/exp/constraints" 15 "golang.org/x/exp/slices" 16 ) 17 18 // CombineUnique merges two ordered slices. If both slices have unique elements 19 // then so does the resulting slice. More generally, each element is present 20 // max(timesInA, timesInB) times. 21 // 22 // Takes ownership of both slices, and uses the longer one to store the result. 23 // 24 // This function is used to combine slices where one of the slices is small or 25 // has mostly the same elements as the other. If the two slices are large and 26 // don't have many duplicates, this function should be avoided, because of the 27 // usage of `copy` that can increase CPU. 28 func CombineUnique[T constraints.Ordered](a, b []T) []T { 29 // We want b to be the smaller slice, so there are fewer elements to be added. 30 if len(b) > len(a) { 31 b, a = a, b 32 } 33 aIter, bIter := 0, 0 34 for aIter < len(a) && bIter < len(b) { 35 if a[aIter] == b[bIter] { 36 aIter++ 37 bIter++ 38 } else if a[aIter] < b[bIter] { 39 aIter++ 40 } else { 41 var zero T 42 a = append(a, zero) 43 copy(a[aIter+1:], a[aIter:]) 44 a[aIter] = b[bIter] 45 aIter++ 46 bIter++ 47 } 48 } 49 if bIter < len(b) { 50 a = append(a, b[bIter:]...) 51 } 52 return a 53 } 54 55 // Filter returns a new slice that only contains elements from collection that 56 // satisfy predicate. 57 // 58 // // Filter in place 59 // numbers = Filter(numbers, isEven) 60 // // Filter into a new slice 61 // odds := Filter(numbers, isEven) 62 func Filter[T any](collection []T, predicate func(T) bool) []T { 63 i := 0 64 out := make([]T, len(collection)) 65 for j := range collection { 66 if predicate(collection[j]) { 67 out[i] = collection[j] 68 i++ 69 } 70 } 71 return slices.Clip(out[:i]) 72 } 73 74 // Map returns a new slice containing the results of fn for each element within 75 // collection. Usage: 76 // 77 // Map([]int{1, 2, 3}, func(i int) int { 78 // return i 79 // }) 80 func Map[T, K any](collection []T, fn func(T) K) []K { 81 out := make([]K, len(collection)) 82 for i, el := range collection { 83 out[i] = fn(el) 84 } 85 return out 86 } 87 88 // MapFrom returns a map populated with keys and values returned by fn. 89 // Usage: 90 // 91 // // Construct a set. 92 // MapFrom(numbers, func(i int) (int, struct{}) { 93 // return i, struct{}{} 94 // }) 95 // 96 // // Construct a map of numbers to their square. 97 // MapFrom(numbers, func(i int) (int, int) { 98 // return i, i * i 99 // }) 100 func MapFrom[T any, K comparable, V any](collection []T, fn func(T) (K, V)) map[K]V { 101 out := make(map[K]V, len(collection)) 102 for _, el := range collection { 103 key, value := fn(el) 104 out[key] = value 105 } 106 return out 107 }