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 r[i+len(v)] == value originally at r[i]. 132 // Insert panics if i is out of range. 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 // The s[:0:0] preserves nil in case it matters. 350 return append(s[:0:0], s...) 351 } 352 353 // Compact replaces consecutive runs of equal elements with a single copy. 354 // This is like the uniq command found on Unix. 355 // Compact modifies the contents of the slice s and returns the modified slice, 356 // which may have a smaller length. 357 // Compact zeroes the elements between the new length and the original length. 358 func Compact[S ~[]E, E comparable](s S) S { 359 if len(s) < 2 { 360 return s 361 } 362 for k := 1; k < len(s); k++ { 363 if s[k] == s[k-1] { 364 s2 := s[k:] 365 for k2 := 1; k2 < len(s2); k2++ { 366 if s2[k2] != s2[k2-1] { 367 s[k] = s2[k2] 368 k++ 369 } 370 } 371 372 clear(s[k:]) // zero/nil out the obsolete elements, for GC 373 return s[:k] 374 } 375 } 376 return s 377 } 378 379 // CompactFunc is like [Compact] but uses an equality function to compare elements. 380 // For runs of elements that compare equal, CompactFunc keeps the first one. 381 // CompactFunc zeroes the elements between the new length and the original length. 382 func CompactFunc[S ~[]E, E any](s S, eq func(E, E) bool) S { 383 if len(s) < 2 { 384 return s 385 } 386 for k := 1; k < len(s); k++ { 387 if eq(s[k], s[k-1]) { 388 s2 := s[k:] 389 for k2 := 1; k2 < len(s2); k2++ { 390 if !eq(s2[k2], s2[k2-1]) { 391 s[k] = s2[k2] 392 k++ 393 } 394 } 395 396 clear(s[k:]) // zero/nil out the obsolete elements, for GC 397 return s[:k] 398 } 399 } 400 return s 401 } 402 403 // Grow increases the slice's capacity, if necessary, to guarantee space for 404 // another n elements. After Grow(n), at least n elements can be appended 405 // to the slice without another allocation. If n is negative or too large to 406 // allocate the memory, Grow panics. 407 func Grow[S ~[]E, E any](s S, n int) S { 408 if n < 0 { 409 panic("cannot be negative") 410 } 411 if n -= cap(s) - len(s); n > 0 { 412 s = append(s[:cap(s)], make([]E, n)...)[:len(s)] 413 } 414 return s 415 } 416 417 // Clip removes unused capacity from the slice, returning s[:len(s):len(s)]. 418 func Clip[S ~[]E, E any](s S) S { 419 return s[:len(s):len(s)] 420 } 421 422 // TODO: There are other rotate algorithms. 423 // This algorithm has the desirable property that it moves each element at most twice. 424 // The follow-cycles algorithm can be 1-write but it is not very cache friendly. 425 426 // rotateLeft rotates s left by r spaces. 427 // s_final[i] = s_orig[i+r], wrapping around. 428 func rotateLeft[E any](s []E, r int) { 429 Reverse(s[:r]) 430 Reverse(s[r:]) 431 Reverse(s) 432 } 433 func rotateRight[E any](s []E, r int) { 434 rotateLeft(s, len(s)-r) 435 } 436 437 // overlaps reports whether the memory ranges a[0:len(a)] and b[0:len(b)] overlap. 438 func overlaps[E any](a, b []E) bool { 439 if len(a) == 0 || len(b) == 0 { 440 return false 441 } 442 elemSize := unsafe.Sizeof(a[0]) 443 if elemSize == 0 { 444 return false 445 } 446 // TODO: use a runtime/unsafe facility once one becomes available. See issue 12445. 447 // Also see crypto/internal/alias/alias.go:AnyOverlap 448 return uintptr(unsafe.Pointer(&a[0])) <= uintptr(unsafe.Pointer(&b[len(b)-1]))+(elemSize-1) && 449 uintptr(unsafe.Pointer(&b[0])) <= uintptr(unsafe.Pointer(&a[len(a)-1]))+(elemSize-1) 450 } 451 452 // startIdx returns the index in haystack where the needle starts. 453 // prerequisite: the needle must be aliased entirely inside the haystack. 454 func startIdx[E any](haystack, needle []E) int { 455 p := &needle[0] 456 for i := range haystack { 457 if p == &haystack[i] { 458 return i 459 } 460 } 461 // TODO: what if the overlap is by a non-integral number of Es? 462 panic("needle not found") 463 } 464 465 // Reverse reverses the elements of the slice in place. 466 func Reverse[S ~[]E, E any](s S) { 467 for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 { 468 s[i], s[j] = s[j], s[i] 469 } 470 } 471 472 // Concat returns a new slice concatenating the passed in slices. 473 func Concat[S ~[]E, E any](slices ...S) S { 474 size := 0 475 for _, s := range slices { 476 size += len(s) 477 if size < 0 { 478 panic("len out of range") 479 } 480 } 481 newslice := Grow[S](nil, size) 482 for _, s := range slices { 483 newslice = append(newslice, s...) 484 } 485 return newslice 486 } 487 488 // Repeat returns a new slice that repeats the provided slice the given number of times. 489 // The result has length and capacity (len(x) * count). 490 // The result is never nil. 491 // Repeat panics if count is negative or if the result of (len(x) * count) 492 // overflows. 493 func Repeat[S ~[]E, E any](x S, count int) S { 494 if count < 0 { 495 panic("cannot be negative") 496 } 497 498 const maxInt = ^uint(0) >> 1 499 if hi, lo := bits.Mul(uint(len(x)), uint(count)); hi > 0 || lo > maxInt { 500 panic("the result of (len(x) * count) overflows") 501 } 502 503 newslice := make(S, len(x)*count) 504 n := copy(newslice, x) 505 for n < len(newslice) { 506 n += copy(newslice[n:], newslice[:n]) 507 } 508 return newslice 509 } 510