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