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  	"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  

View as plain text