varint.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. package varint
  2. import (
  3. "encoding/binary"
  4. "errors"
  5. "io"
  6. "math/bits"
  7. )
  8. var (
  9. ErrOverflow = errors.New("varints larger than uint63 not supported")
  10. ErrUnderflow = errors.New("varints malformed, could not reach the end")
  11. ErrNotMinimal = errors.New("varint not minimally encoded")
  12. )
  13. const (
  14. // MaxLenUvarint63 is the maximum number of bytes representing an uvarint in
  15. // this encoding, supporting a maximum value of 2^63 (uint63), aka
  16. // MaxValueUvarint63.
  17. MaxLenUvarint63 = 9
  18. // MaxValueUvarint63 is the maximum encodable uint63 value.
  19. MaxValueUvarint63 = (1 << 63) - 1
  20. )
  21. // UvarintSize returns the size (in bytes) of `num` encoded as a unsigned varint.
  22. //
  23. // This may return a size greater than MaxUvarintLen63, which would be an
  24. // illegal value, and would be rejected by readers.
  25. func UvarintSize(num uint64) int {
  26. bits := bits.Len64(num)
  27. q, r := bits/7, bits%7
  28. size := q
  29. if r > 0 || size == 0 {
  30. size++
  31. }
  32. return size
  33. }
  34. // ToUvarint converts an unsigned integer to a varint-encoded []byte
  35. func ToUvarint(num uint64) []byte {
  36. buf := make([]byte, UvarintSize(num))
  37. n := binary.PutUvarint(buf, uint64(num))
  38. return buf[:n]
  39. }
  40. // FromUvarint reads an unsigned varint from the beginning of buf, returns the
  41. // varint, and the number of bytes read.
  42. func FromUvarint(buf []byte) (uint64, int, error) {
  43. // Modified from the go standard library. Copyright the Go Authors and
  44. // released under the BSD License.
  45. var x uint64
  46. var s uint
  47. for i, b := range buf {
  48. if (i == 8 && b >= 0x80) || i >= MaxLenUvarint63 {
  49. // this is the 9th and last byte we're willing to read, but it
  50. // signals there's more (1 in MSB).
  51. // or this is the >= 10th byte, and for some reason we're still here.
  52. return 0, 0, ErrOverflow
  53. }
  54. if b < 0x80 {
  55. if b == 0 && s > 0 {
  56. return 0, 0, ErrNotMinimal
  57. }
  58. return x | uint64(b)<<s, i + 1, nil
  59. }
  60. x |= uint64(b&0x7f) << s
  61. s += 7
  62. }
  63. return 0, 0, ErrUnderflow
  64. }
  65. // ReadUvarint reads a unsigned varint from the given reader.
  66. func ReadUvarint(r io.ByteReader) (uint64, error) {
  67. // Modified from the go standard library. Copyright the Go Authors and
  68. // released under the BSD License.
  69. var x uint64
  70. var s uint
  71. for i := 0; ; i++ {
  72. b, err := r.ReadByte()
  73. if err != nil {
  74. if err == io.EOF && i != 0 {
  75. // "eof" will look like a success.
  76. // If we've read part of a value, this is not a
  77. // success.
  78. err = io.ErrUnexpectedEOF
  79. }
  80. return 0, err
  81. }
  82. if (i == 8 && b >= 0x80) || i >= MaxLenUvarint63 {
  83. // this is the 9th and last byte we're willing to read, but it
  84. // signals there's more (1 in MSB).
  85. // or this is the >= 10th byte, and for some reason we're still here.
  86. return 0, ErrOverflow
  87. }
  88. if b < 0x80 {
  89. if b == 0 && s > 0 {
  90. return 0, ErrNotMinimal
  91. }
  92. return x | uint64(b)<<s, nil
  93. }
  94. x |= uint64(b&0x7f) << s
  95. s += 7
  96. }
  97. }
  98. // PutUvarint is an alias for binary.PutUvarint.
  99. //
  100. // This is provided for convenience so users of this library can avoid built-in
  101. // varint functions and easily audit code for uses of those functions.
  102. //
  103. // Make sure that x is smaller or equal to MaxValueUvarint63, otherwise this
  104. // function will produce values that may be rejected by readers.
  105. func PutUvarint(buf []byte, x uint64) int {
  106. return binary.PutUvarint(buf, x)
  107. }