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