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