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