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  }