1 // Copyright 2021 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 5 // Package slices defines various functions useful with slices of any type.
6 package slices
7 8 import (
9 "cmp"
10 "math/bits"
11 "unsafe"
12 )
13 14 // Equal reports whether two slices are equal: the same length and all
15 // elements equal. If the lengths are different, Equal returns false.
16 // Otherwise, the elements are compared in increasing index order, and the
17 // comparison stops at the first unequal pair.
18 // Empty and nil slices are considered equal.
19 // Floating point NaNs are not considered equal.
20 func Equal[S ~[]E, E comparable](s1, s2 S) bool {
21 if len(s1) != len(s2) {
22 return false
23 }
24 for i := range s1 {
25 if s1[i] != s2[i] {
26 return false
27 }
28 }
29 return true
30 }
31 32 // EqualFunc reports whether two slices are equal using an equality
33 // function on each pair of elements. If the lengths are different,
34 // EqualFunc returns false. Otherwise, the elements are compared in
35 // increasing index order, and the comparison stops at the first index
36 // for which eq returns false.
37 func EqualFunc[S1 ~[]E1, S2 ~[]E2, E1, E2 any](s1 S1, s2 S2, eq func(E1, E2) bool) bool {
38 if len(s1) != len(s2) {
39 return false
40 }
41 for i, v1 := range s1 {
42 v2 := s2[i]
43 if !eq(v1, v2) {
44 return false
45 }
46 }
47 return true
48 }
49 50 // Compare compares the elements of s1 and s2, using [cmp.Compare] on each pair
51 // of elements. The elements are compared sequentially, starting at index 0,
52 // until one element is not equal to the other.
53 // The result of comparing the first non-matching elements is returned.
54 // If both slices are equal until one of them ends, the shorter slice is
55 // considered less than the longer one.
56 // The result is 0 if s1 == s2, -1 if s1 < s2, and +1 if s1 > s2.
57 func Compare[S ~[]E, E cmp.Ordered](s1, s2 S) int {
58 for i, v1 := range s1 {
59 if i >= len(s2) {
60 return +1
61 }
62 v2 := s2[i]
63 if c := cmp.Compare(v1, v2); c != 0 {
64 return c
65 }
66 }
67 if len(s1) < len(s2) {
68 return -1
69 }
70 return 0
71 }
72 73 // CompareFunc is like [Compare] but uses a custom comparison function on each
74 // pair of elements.
75 // The result is the first non-zero result of cmp; if cmp always
76 // returns 0 the result is 0 if len(s1) == len(s2), -1 if len(s1) < len(s2),
77 // and +1 if len(s1) > len(s2).
78 func CompareFunc[S1 ~[]E1, S2 ~[]E2, E1, E2 any](s1 S1, s2 S2, cmp func(E1, E2) int) int {
79 for i, v1 := range s1 {
80 if i >= len(s2) {
81 return +1
82 }
83 v2 := s2[i]
84 if c := cmp(v1, v2); c != 0 {
85 return c
86 }
87 }
88 if len(s1) < len(s2) {
89 return -1
90 }
91 return 0
92 }
93 94 // Index returns the index of the first occurrence of v in s,
95 // or -1 if not present.
96 func Index[S ~[]E, E comparable](s S, v E) int {
97 for i := range s {
98 if v == s[i] {
99 return i
100 }
101 }
102 return -1
103 }
104 105 // IndexFunc returns the first index i satisfying f(s[i]),
106 // or -1 if none do.
107 func IndexFunc[S ~[]E, E any](s S, f func(E) bool) int {
108 for i := range s {
109 if f(s[i]) {
110 return i
111 }
112 }
113 return -1
114 }
115 116 // Contains reports whether v is present in s.
117 func Contains[S ~[]E, E comparable](s S, v E) bool {
118 return Index(s, v) >= 0
119 }
120 121 // ContainsFunc reports whether at least one
122 // element e of s satisfies f(e).
123 func ContainsFunc[S ~[]E, E any](s S, f func(E) bool) bool {
124 return IndexFunc(s, f) >= 0
125 }
126 127 // Insert inserts the values v... into s at index i,
128 // returning the modified slice.
129 // The elements at s[i:] are shifted up to make room.
130 // In the returned slice r, r[i] == v[0],
131 // and, if i < len(s), r[i+len(v)] == value originally at r[i].
132 // Insert panics if i > len(s).
133 // This function is O(len(s) + len(v)).
134 // If the result is empty, it has the same nilness as s.
135 func Insert[S ~[]E, E any](s S, i int, v ...E) S {
136 _ = s[i:] // bounds check
137 138 m := len(v)
139 if m == 0 {
140 return s
141 }
142 n := len(s)
143 if i == n {
144 return append(s, v...)
145 }
146 if n+m > cap(s) {
147 // Use append rather than make so that we bump the size of
148 // the slice up to the next storage class.
149 // This is what Grow does but we don't call Grow because
150 // that might copy the values twice.
151 s2 := append(s[:i], make(S, n+m-i)...)
152 copy(s2[i:], v)
153 copy(s2[i+m:], s[i:])
154 return s2
155 }
156 s = s[:n+m]
157 158 // before:
159 // s: aaaaaaaabbbbccccccccdddd
160 // ^ ^ ^ ^
161 // i i+m n n+m
162 // after:
163 // s: aaaaaaaavvvvbbbbcccccccc
164 // ^ ^ ^ ^
165 // i i+m n n+m
166 //
167 // a are the values that don't move in s.
168 // v are the values copied in from v.
169 // b and c are the values from s that are shifted up in index.
170 // d are the values that get overwritten, never to be seen again.
171 172 if !overlaps(v, s[i+m:]) {
173 // Easy case - v does not overlap either the c or d regions.
174 // (It might be in some of a or b, or elsewhere entirely.)
175 // The data we copy up doesn't write to v at all, so just do it.
176 177 copy(s[i+m:], s[i:])
178 179 // Now we have
180 // s: aaaaaaaabbbbbbbbcccccccc
181 // ^ ^ ^ ^
182 // i i+m n n+m
183 // Note the b values are duplicated.
184 185 copy(s[i:], v)
186 187 // Now we have
188 // s: aaaaaaaavvvvbbbbcccccccc
189 // ^ ^ ^ ^
190 // i i+m n n+m
191 // That's the result we want.
192 return s
193 }
194 195 // The hard case - v overlaps c or d. We can't just shift up
196 // the data because we'd move or clobber the values we're trying
197 // to insert.
198 // So instead, write v on top of d, then rotate.
199 copy(s[n:], v)
200 201 // Now we have
202 // s: aaaaaaaabbbbccccccccvvvv
203 // ^ ^ ^ ^
204 // i i+m n n+m
205 206 rotateRight(s[i:], m)
207 208 // Now we have
209 // s: aaaaaaaavvvvbbbbcccccccc
210 // ^ ^ ^ ^
211 // i i+m n n+m
212 // That's the result we want.
213 return s
214 }
215 216 // Delete removes the elements s[i:j] from s, returning the modified slice.
217 // Delete panics if j > len(s) or s[i:j] is not a valid slice of s.
218 // Delete is O(len(s)-i), so if many items must be deleted, it is better to
219 // make a single call deleting them all together than to delete one at a time.
220 // Delete zeroes the elements s[len(s)-(j-i):len(s)].
221 // If the result is empty, it has the same nilness as s.
222 func Delete[S ~[]E, E any](s S, i, j int) S {
223 _ = s[i:j:len(s)] // bounds check
224 225 if i == j {
226 return s
227 }
228 229 oldlen := len(s)
230 s = append(s[:i], s[j:]...)
231 clear(s[len(s):oldlen]) // zero/nil out the obsolete elements, for GC
232 return s
233 }
234 235 // DeleteFunc removes any elements from s for which del returns true,
236 // returning the modified slice.
237 // DeleteFunc zeroes the elements between the new length and the original length.
238 // If the result is empty, it has the same nilness as s.
239 func DeleteFunc[S ~[]E, E any](s S, del func(E) bool) S {
240 i := IndexFunc(s, del)
241 if i == -1 {
242 return s
243 }
244 // Don't start copying elements until we find one to delete.
245 for j := i + 1; j < len(s); j++ {
246 if v := s[j]; !del(v) {
247 s[i] = v
248 i++
249 }
250 }
251 clear(s[i:]) // zero/nil out the obsolete elements, for GC
252 return s[:i]
253 }
254 255 // Replace replaces the elements s[i:j] by the given v, and returns the
256 // modified slice.
257 // Replace panics if j > len(s) or s[i:j] is not a valid slice of s.
258 // When len(v) < (j-i), Replace zeroes the elements between the new length and the original length.
259 // If the result is empty, it has the same nilness as s.
260 func Replace[S ~[]E, E any](s S, i, j int, v ...E) S {
261 _ = s[i:j] // bounds check
262 263 if i == j {
264 return Insert(s, i, v...)
265 }
266 if j == len(s) {
267 s2 := append(s[:i], v...)
268 if len(s2) < len(s) {
269 clear(s[len(s2):]) // zero/nil out the obsolete elements, for GC
270 }
271 return s2
272 }
273 274 tot := len(s[:i]) + len(v) + len(s[j:])
275 if tot > cap(s) {
276 // Too big to fit, allocate and copy over.
277 s2 := append(s[:i], make(S, tot-i)...) // See Insert
278 copy(s2[i:], v)
279 copy(s2[i+len(v):], s[j:])
280 return s2
281 }
282 283 r := s[:tot]
284 285 if i+len(v) <= j {
286 // Easy, as v fits in the deleted portion.
287 copy(r[i:], v)
288 copy(r[i+len(v):], s[j:])
289 clear(s[tot:]) // zero/nil out the obsolete elements, for GC
290 return r
291 }
292 293 // We are expanding (v is bigger than j-i).
294 // The situation is something like this:
295 // (example has i=4,j=8,len(s)=16,len(v)=6)
296 // s: aaaaxxxxbbbbbbbbyy
297 // ^ ^ ^ ^
298 // i j len(s) tot
299 // a: prefix of s
300 // x: deleted range
301 // b: more of s
302 // y: area to expand into
303 304 if !overlaps(r[i+len(v):], v) {
305 // Easy, as v is not clobbered by the first copy.
306 copy(r[i+len(v):], s[j:])
307 copy(r[i:], v)
308 return r
309 }
310 311 // This is a situation where we don't have a single place to which
312 // we can copy v. Parts of it need to go to two different places.
313 // We want to copy the prefix of v into y and the suffix into x, then
314 // rotate |y| spots to the right.
315 //
316 // v[2:] v[:2]
317 // | |
318 // s: aaaavvvvbbbbbbbbvv
319 // ^ ^ ^ ^
320 // i j len(s) tot
321 //
322 // If either of those two destinations don't alias v, then we're good.
323 y := len(v) - (j - i) // length of y portion
324 325 if !overlaps(r[i:j], v) {
326 copy(r[i:j], v[y:])
327 copy(r[len(s):], v[:y])
328 rotateRight(r[i:], y)
329 return r
330 }
331 if !overlaps(r[len(s):], v) {
332 copy(r[len(s):], v[:y])
333 copy(r[i:j], v[y:])
334 rotateRight(r[i:], y)
335 return r
336 }
337 338 // Now we know that v overlaps both x and y.
339 // That means that the entirety of b is *inside* v.
340 // So we don't need to preserve b at all; instead we
341 // can copy v first, then copy the b part of v out of
342 // v to the right destination.
343 k := startIdx(v, s[j:])
344 copy(r[i:], v)
345 copy(r[i+len(v):], r[i+k:])
346 return r
347 }
348 349 // Clone returns a copy of the slice.
350 // The elements are copied using assignment, so this is a shallow clone.
351 // The result may have additional unused capacity.
352 // The result preserves the nilness of s.
353 func Clone[S ~[]E, E any](s S) S {
354 // Preserve nilness in case it matters.
355 if s == nil {
356 return nil
357 }
358 // Avoid s[:0:0] as it leads to unwanted liveness when cloning a
359 // zero-length slice of a large array; see https://go.dev/issue/68488.
360 return append(S{}, s...)
361 }
362 363 // Compact replaces consecutive runs of equal elements with a single copy.
364 // This is like the uniq command found on Unix.
365 // Compact modifies the contents of the slice s and returns the modified slice,
366 // which may have a smaller length.
367 // Compact zeroes the elements between the new length and the original length.
368 // The result preserves the nilness of s.
369 func Compact[S ~[]E, E comparable](s S) S {
370 if len(s) < 2 {
371 return s
372 }
373 for k := 1; k < len(s); k++ {
374 if s[k] == s[k-1] {
375 s2 := s[k:]
376 for k2 := 1; k2 < len(s2); k2++ {
377 if s2[k2] != s2[k2-1] {
378 s[k] = s2[k2]
379 k++
380 }
381 }
382 383 clear(s[k:]) // zero/nil out the obsolete elements, for GC
384 return s[:k]
385 }
386 }
387 return s
388 }
389 390 // CompactFunc is like [Compact] but uses an equality function to compare elements.
391 // For runs of elements that compare equal, CompactFunc keeps the first one.
392 // CompactFunc zeroes the elements between the new length and the original length.
393 // The result preserves the nilness of s.
394 func CompactFunc[S ~[]E, E any](s S, eq func(E, E) bool) S {
395 if len(s) < 2 {
396 return s
397 }
398 for k := 1; k < len(s); k++ {
399 if eq(s[k], s[k-1]) {
400 s2 := s[k:]
401 for k2 := 1; k2 < len(s2); k2++ {
402 if !eq(s2[k2], s2[k2-1]) {
403 s[k] = s2[k2]
404 k++
405 }
406 }
407 408 clear(s[k:]) // zero/nil out the obsolete elements, for GC
409 return s[:k]
410 }
411 }
412 return s
413 }
414 415 // Grow increases the slice's capacity, if necessary, to guarantee space for
416 // another n elements. After Grow(n), at least n elements can be appended
417 // to the slice without another allocation. If n is negative or too large to
418 // allocate the memory, Grow panics.
419 // The result preserves the nilness of s.
420 func Grow[S ~[]E, E any](s S, n int) S {
421 if n < 0 {
422 panic("cannot be negative")
423 }
424 if n -= cap(s) - len(s); n > 0 {
425 // This expression allocates only once (see test).
426 s = append(s[:cap(s)], []E{:n}...)[:len(s)]
427 }
428 return s
429 }
430 431 // Clip removes unused capacity from the slice, returning s[:len(s):len(s)].
432 // The result preserves the nilness of s.
433 func Clip[S ~[]E, E any](s S) S {
434 return s[:len(s):len(s)]
435 }
436 437 // TODO: There are other rotate algorithms.
438 // This algorithm has the desirable property that it moves each element at most twice.
439 // The follow-cycles algorithm can be 1-write but it is not very cache friendly.
440 441 // rotateLeft rotates s left by r spaces.
442 // s_final[i] = s_orig[i+r], wrapping around.
443 func rotateLeft[E any](s []E, r int) {
444 Reverse(s[:r])
445 Reverse(s[r:])
446 Reverse(s)
447 }
448 func rotateRight[E any](s []E, r int) {
449 rotateLeft(s, len(s)-r)
450 }
451 452 // overlaps reports whether the memory ranges a[:len(a)] and b[:len(b)] overlap.
453 func overlaps[E any](a, b []E) bool {
454 if len(a) == 0 || len(b) == 0 {
455 return false
456 }
457 elemSize := unsafe.Sizeof(a[0])
458 if elemSize == 0 {
459 return false
460 }
461 // TODO: use a runtime/unsafe facility once one becomes available. See issue 12445.
462 // Also see crypto/internal/fips140/alias/alias.go:AnyOverlap
463 return uintptr(unsafe.Pointer(&a[0])) <= uintptr(unsafe.Pointer(&b[len(b)-1]))+(elemSize-1) &&
464 uintptr(unsafe.Pointer(&b[0])) <= uintptr(unsafe.Pointer(&a[len(a)-1]))+(elemSize-1)
465 }
466 467 // startIdx returns the index in haystack where the needle starts.
468 // prerequisite: the needle must be aliased entirely inside the haystack.
469 func startIdx[E any](haystack, needle []E) int {
470 p := &needle[0]
471 for i := range haystack {
472 if p == &haystack[i] {
473 return i
474 }
475 }
476 // TODO: what if the overlap is by a non-integral number of Es?
477 panic("needle not found")
478 }
479 480 // Reverse reverses the elements of the slice in place.
481 func Reverse[S ~[]E, E any](s S) {
482 for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
483 s[i], s[j] = s[j], s[i]
484 }
485 }
486 487 // Concat returns a new slice concatenating the passed in slices.
488 // If the concatenation is empty, the result is nil.
489 func Concat[S ~[]E, E any](slices ...S) S {
490 size := 0
491 for _, s := range slices {
492 size += len(s)
493 if size < 0 {
494 panic("len out of range")
495 }
496 }
497 // Use Grow, not make, to round up to the size class:
498 // the extra space is otherwise unused and helps
499 // callers that append a few elements to the result.
500 newslice := Grow[S](nil, size)
501 for _, s := range slices {
502 newslice = append(newslice, s...)
503 }
504 return newslice
505 }
506 507 // Repeat returns a new slice that repeats the provided slice the given number of times.
508 // The result has length and capacity (len(x) * count).
509 // The result is never nil.
510 // Repeat panics if count is negative or if the result of (len(x) * count)
511 // overflows.
512 func Repeat[S ~[]E, E any](x S, count int) S {
513 if count < 0 {
514 panic("cannot be negative")
515 }
516 517 const maxInt = ^uint(0) >> 1
518 hi, lo := bits.Mul(uint(len(x)), uint(count))
519 if hi > 0 || lo > maxInt {
520 panic("the result of (len(x) * count) overflows")
521 }
522 523 newslice := make(S, int(lo)) // lo = len(x) * count
524 n := copy(newslice, x)
525 for n < len(newslice) {
526 n += copy(newslice[n:], newslice[:n])
527 }
528 return newslice
529 }
530