Source file 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 "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)], make([]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