github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/enum/enum.go

// Copyright 2020 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package enum

import "bytes"

// Note that while maxToken is outside the range of a single
// byte, we never actually insert it in GenByteStringBetween.
// We only use it to perform computation using the value 256.
const minToken int = 0
const maxToken int = 256
const midToken = maxToken / 2
const shiftInterval = 5

// ByteSpacing is a type that controls what distribution of generated
// bytes strings is created with calls to GenByteStringBetween.
type ByteSpacing int

const (
	// PackedSpacing is used when the generated bytes are intended to be "close"
	// together in the generated key space.
	PackedSpacing ByteSpacing = iota
	// SpreadSpacing is used when the generated bytes are intended to be evenly
	// spaced out within the generated key space.
	SpreadSpacing
)

func (s ByteSpacing) String() string {
	switch s {
	case PackedSpacing:
		return "packed"
	case SpreadSpacing:
		return "spread"
	}
	panic("unknown spacing type")
}

// GenByteStringBetween generates a byte string that sorts
// between the two input strings. If prev is length 0, it is
// treated as negative infinity. If next is length 0, it is
// treated as positive infinity. Importantly, the input strings
// cannot end with minToken.
func GenByteStringBetween(prev []byte, next []byte, spacing ByteSpacing) []byte {
	result := make([]byte, 0)
	if len(prev) == 0 && len(next) == 0 {
		// If both prev and next are unbounded, return the midpoint.
		return append(result, byte(midToken))
	}
	maxLen := len(prev)
	if len(next) > maxLen {
		maxLen = len(next)
	}

	// Construct the prefix of prev and next.
	pos := 0
	for ; pos < maxLen; pos++ {
		p, n := get(prev, pos, minToken), get(next, pos, maxToken)
		if p != n {
			break
		}
		result = append(result, byte(p))
	}

	// We've found an index where prev and next disagree. So, it's time to start
	// trying to construct a value between prev and next.
	p, n := get(prev, pos, minToken), get(next, pos, maxToken)
	var mid int
	switch spacing {
	case PackedSpacing:
		mid = byteBetweenPacked(p, n)
	case SpreadSpacing:
		mid = byteBetweenSpread(p, n)
	}

	// If mid == p, then we know there is no more room between
	// prev and next at this index. So, we can append p to result
	// to ensure that it is less than next. To generate the rest
	// of the string, we try to find a string that fits between
	// the remainder of prev and posinf.
	if mid == p {
		result = append(result, byte(p))
		rest := GenByteStringBetween(slice(prev, pos+1), nil, spacing)
		return append(result, rest...)
	}

	// If mid != p, then there is room between prev and next at this index.
	// So, occupy that spot and return.
	result = append(result, byte(mid))
	return result
}

// Utility functions for GenByteStringBetween.

func get(arr []byte, idx int, def int) int {
	if idx >= len(arr) {
		return def
	}
	return int(arr[idx])
}

func slice(arr []byte, idx int) []byte {
	if idx > len(arr) {
		return []byte(nil)
	}
	return arr[idx:]
}

// byteBetweenPacked generates a byte value between hi and lo, inclusive.
// Additionally, it optimizes for adding values at the beginning and end of the
// enum byte sequence by returning a value moved by a small constant
// rather than in the middle of the range when the upper and lower
// bounds are the min or max token.
func byteBetweenPacked(lo int, hi int) int {
	switch {
	case lo == minToken && hi == maxToken:
		return lo + (hi-lo)/2
	case lo == minToken && hi-lo > shiftInterval:
		return hi - shiftInterval
	case hi == maxToken && hi-lo > shiftInterval:
		return lo + shiftInterval
	default:
		return lo + (hi-lo)/2
	}
}

// byteBetweenSpread returns a byte value between hi and lo, inclusive.
func byteBetweenSpread(lo int, hi int) int {
	return lo + (hi-lo)/2
}

// GenerateNEvenlySpacedBytes returns an array of n byte slices that
// evenly split the key space into n pieces.
func GenerateNEvenlySpacedBytes(n int) [][]byte {
	result := make([][]byte, n)
	genEvenlySpacedHelper(result, 0, n, nil, nil)
	return result
}

// genEvenlySpacedHelper fills in result with a byte value between bot and top
// at the index between lo and hi.
func genEvenlySpacedHelper(result [][]byte, lo, hi int, bot, top []byte) {
	if lo == hi {
		return
	}
	mid := lo + (hi-lo)/2
	midBytes := GenByteStringBetween(bot, top, SpreadSpacing)
	result[mid] = midBytes
	genEvenlySpacedHelper(result, lo, mid, bot, midBytes)
	genEvenlySpacedHelper(result, mid+1, hi, midBytes, top)
}

// enumBytesAreLess is a utility function for comparing byte values generated
// for enum's physical representation. It adjusts bytes.Compare to treat b=nil
// as positive infinity.
func enumBytesAreLess(a []byte, b []byte) bool {
	if len(b) == 0 {
		return true
	}
	return bytes.Compare(a, b) == -1
}