src/slices/slices.go

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