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