bytebuf.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. // Copyright 2019+ Klaus Post. All rights reserved.
  2. // License information can be found in the LICENSE file.
  3. // Based on work by Yann Collet, released under BSD License.
  4. package zstd
  5. import (
  6. "fmt"
  7. "io"
  8. "io/ioutil"
  9. )
  10. type byteBuffer interface {
  11. // Read up to 8 bytes.
  12. // Returns io.ErrUnexpectedEOF if this cannot be satisfied.
  13. readSmall(n int) ([]byte, error)
  14. // Read >8 bytes.
  15. // MAY use the destination slice.
  16. readBig(n int, dst []byte) ([]byte, error)
  17. // Read a single byte.
  18. readByte() (byte, error)
  19. // Skip n bytes.
  20. skipN(n int) error
  21. }
  22. // in-memory buffer
  23. type byteBuf []byte
  24. func (b *byteBuf) readSmall(n int) ([]byte, error) {
  25. if debugAsserts && n > 8 {
  26. panic(fmt.Errorf("small read > 8 (%d). use readBig", n))
  27. }
  28. bb := *b
  29. if len(bb) < n {
  30. return nil, io.ErrUnexpectedEOF
  31. }
  32. r := bb[:n]
  33. *b = bb[n:]
  34. return r, nil
  35. }
  36. func (b *byteBuf) readBig(n int, dst []byte) ([]byte, error) {
  37. bb := *b
  38. if len(bb) < n {
  39. return nil, io.ErrUnexpectedEOF
  40. }
  41. r := bb[:n]
  42. *b = bb[n:]
  43. return r, nil
  44. }
  45. func (b *byteBuf) remain() []byte {
  46. return *b
  47. }
  48. func (b *byteBuf) readByte() (byte, error) {
  49. bb := *b
  50. if len(bb) < 1 {
  51. return 0, nil
  52. }
  53. r := bb[0]
  54. *b = bb[1:]
  55. return r, nil
  56. }
  57. func (b *byteBuf) skipN(n int) error {
  58. bb := *b
  59. if len(bb) < n {
  60. return io.ErrUnexpectedEOF
  61. }
  62. *b = bb[n:]
  63. return nil
  64. }
  65. // wrapper around a reader.
  66. type readerWrapper struct {
  67. r io.Reader
  68. tmp [8]byte
  69. }
  70. func (r *readerWrapper) readSmall(n int) ([]byte, error) {
  71. if debugAsserts && n > 8 {
  72. panic(fmt.Errorf("small read > 8 (%d). use readBig", n))
  73. }
  74. n2, err := io.ReadFull(r.r, r.tmp[:n])
  75. // We only really care about the actual bytes read.
  76. if err != nil {
  77. if err == io.EOF {
  78. return nil, io.ErrUnexpectedEOF
  79. }
  80. if debugDecoder {
  81. println("readSmall: got", n2, "want", n, "err", err)
  82. }
  83. return nil, err
  84. }
  85. return r.tmp[:n], nil
  86. }
  87. func (r *readerWrapper) readBig(n int, dst []byte) ([]byte, error) {
  88. if cap(dst) < n {
  89. dst = make([]byte, n)
  90. }
  91. n2, err := io.ReadFull(r.r, dst[:n])
  92. if err == io.EOF && n > 0 {
  93. err = io.ErrUnexpectedEOF
  94. }
  95. return dst[:n2], err
  96. }
  97. func (r *readerWrapper) readByte() (byte, error) {
  98. n2, err := r.r.Read(r.tmp[:1])
  99. if err != nil {
  100. return 0, err
  101. }
  102. if n2 != 1 {
  103. return 0, io.ErrUnexpectedEOF
  104. }
  105. return r.tmp[0], nil
  106. }
  107. func (r *readerWrapper) skipN(n int) error {
  108. n2, err := io.CopyN(ioutil.Discard, r.r, int64(n))
  109. if n2 != int64(n) {
  110. err = io.ErrUnexpectedEOF
  111. }
  112. return err
  113. }