1
0

encode.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356
  1. // Copyright 2011 The Snappy-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. package snappy
  5. import (
  6. "encoding/binary"
  7. "errors"
  8. "io"
  9. )
  10. // We limit how far copy back-references can go, the same as the C++ code.
  11. const maxOffset = 1 << 15
  12. // emitLiteral writes a literal chunk and returns the number of bytes written.
  13. func emitLiteral(dst, lit []byte) int {
  14. i, n := 0, uint(len(lit)-1)
  15. switch {
  16. case n < 60:
  17. dst[0] = uint8(n)<<2 | tagLiteral
  18. i = 1
  19. case n < 1<<8:
  20. dst[0] = 60<<2 | tagLiteral
  21. dst[1] = uint8(n)
  22. i = 2
  23. case n < 1<<16:
  24. dst[0] = 61<<2 | tagLiteral
  25. dst[1] = uint8(n)
  26. dst[2] = uint8(n >> 8)
  27. i = 3
  28. case n < 1<<24:
  29. dst[0] = 62<<2 | tagLiteral
  30. dst[1] = uint8(n)
  31. dst[2] = uint8(n >> 8)
  32. dst[3] = uint8(n >> 16)
  33. i = 4
  34. case int64(n) < 1<<32:
  35. dst[0] = 63<<2 | tagLiteral
  36. dst[1] = uint8(n)
  37. dst[2] = uint8(n >> 8)
  38. dst[3] = uint8(n >> 16)
  39. dst[4] = uint8(n >> 24)
  40. i = 5
  41. default:
  42. panic("snappy: source buffer is too long")
  43. }
  44. if copy(dst[i:], lit) != len(lit) {
  45. panic("snappy: destination buffer is too short")
  46. }
  47. return i + len(lit)
  48. }
  49. // emitCopy writes a copy chunk and returns the number of bytes written.
  50. func emitCopy(dst []byte, offset, length int) int {
  51. i := 0
  52. for length > 0 {
  53. x := length - 4
  54. if 0 <= x && x < 1<<3 && offset < 1<<11 {
  55. dst[i+0] = uint8(offset>>8)&0x07<<5 | uint8(x)<<2 | tagCopy1
  56. dst[i+1] = uint8(offset)
  57. i += 2
  58. break
  59. }
  60. x = length
  61. if x > 1<<6 {
  62. x = 1 << 6
  63. }
  64. dst[i+0] = uint8(x-1)<<2 | tagCopy2
  65. dst[i+1] = uint8(offset)
  66. dst[i+2] = uint8(offset >> 8)
  67. i += 3
  68. length -= x
  69. }
  70. return i
  71. }
  72. // Encode returns the encoded form of src. The returned slice may be a sub-
  73. // slice of dst if dst was large enough to hold the entire encoded block.
  74. // Otherwise, a newly allocated slice will be returned.
  75. // It is valid to pass a nil dst.
  76. func Encode(dst, src []byte) []byte {
  77. if n := MaxEncodedLen(len(src)); len(dst) < n {
  78. dst = make([]byte, n)
  79. }
  80. // The block starts with the varint-encoded length of the decompressed bytes.
  81. d := binary.PutUvarint(dst, uint64(len(src)))
  82. // Return early if src is short.
  83. if len(src) <= 4 {
  84. if len(src) != 0 {
  85. d += emitLiteral(dst[d:], src)
  86. }
  87. return dst[:d]
  88. }
  89. // Initialize the hash table. Its size ranges from 1<<8 to 1<<14 inclusive.
  90. const maxTableSize = 1 << 14
  91. shift, tableSize := uint(32-8), 1<<8
  92. for tableSize < maxTableSize && tableSize < len(src) {
  93. shift--
  94. tableSize *= 2
  95. }
  96. var table [maxTableSize]int
  97. // Iterate over the source bytes.
  98. var (
  99. s int // The iterator position.
  100. t int // The last position with the same hash as s.
  101. lit int // The start position of any pending literal bytes.
  102. )
  103. for uint(s+3) < uint(len(src)) { // The uint conversions catch overflow from the +3.
  104. // Update the hash table.
  105. b0, b1, b2, b3 := src[s], src[s+1], src[s+2], src[s+3]
  106. h := uint32(b0) | uint32(b1)<<8 | uint32(b2)<<16 | uint32(b3)<<24
  107. p := &table[(h*0x1e35a7bd)>>shift]
  108. // We need to to store values in [-1, inf) in table. To save
  109. // some initialization time, (re)use the table's zero value
  110. // and shift the values against this zero: add 1 on writes,
  111. // subtract 1 on reads.
  112. t, *p = *p-1, s+1
  113. // If t is invalid or src[s:s+4] differs from src[t:t+4], accumulate a literal byte.
  114. if t < 0 || s-t >= maxOffset || b0 != src[t] || b1 != src[t+1] || b2 != src[t+2] || b3 != src[t+3] {
  115. // Skip multiple bytes if the last match was >= 32 bytes prior.
  116. s += 1 + (s-lit)>>5
  117. continue
  118. }
  119. // Otherwise, we have a match. First, emit any pending literal bytes.
  120. if lit != s {
  121. d += emitLiteral(dst[d:], src[lit:s])
  122. }
  123. // Extend the match to be as long as possible.
  124. s0 := s
  125. s, t = s+4, t+4
  126. for s < len(src) && src[s] == src[t] {
  127. s++
  128. t++
  129. }
  130. // Emit the copied bytes.
  131. d += emitCopy(dst[d:], s-t, s-s0)
  132. lit = s
  133. }
  134. // Emit any final pending literal bytes and return.
  135. if lit != len(src) {
  136. d += emitLiteral(dst[d:], src[lit:])
  137. }
  138. return dst[:d]
  139. }
  140. // MaxEncodedLen returns the maximum length of a snappy block, given its
  141. // uncompressed length.
  142. func MaxEncodedLen(srcLen int) int {
  143. // Compressed data can be defined as:
  144. // compressed := item* literal*
  145. // item := literal* copy
  146. //
  147. // The trailing literal sequence has a space blowup of at most 62/60
  148. // since a literal of length 60 needs one tag byte + one extra byte
  149. // for length information.
  150. //
  151. // Item blowup is trickier to measure. Suppose the "copy" op copies
  152. // 4 bytes of data. Because of a special check in the encoding code,
  153. // we produce a 4-byte copy only if the offset is < 65536. Therefore
  154. // the copy op takes 3 bytes to encode, and this type of item leads
  155. // to at most the 62/60 blowup for representing literals.
  156. //
  157. // Suppose the "copy" op copies 5 bytes of data. If the offset is big
  158. // enough, it will take 5 bytes to encode the copy op. Therefore the
  159. // worst case here is a one-byte literal followed by a five-byte copy.
  160. // That is, 6 bytes of input turn into 7 bytes of "compressed" data.
  161. //
  162. // This last factor dominates the blowup, so the final estimate is:
  163. return 32 + srcLen + srcLen/6
  164. }
  165. var errClosed = errors.New("snappy: Writer is closed")
  166. // NewWriter returns a new Writer that compresses to w.
  167. //
  168. // The Writer returned does not buffer writes. There is no need to Flush or
  169. // Close such a Writer.
  170. //
  171. // Deprecated: the Writer returned is not suitable for many small writes, only
  172. // for few large writes. Use NewBufferedWriter instead, which is efficient
  173. // regardless of the frequency and shape of the writes, and remember to Close
  174. // that Writer when done.
  175. func NewWriter(w io.Writer) *Writer {
  176. return &Writer{
  177. w: w,
  178. obuf: make([]byte, obufLen),
  179. }
  180. }
  181. // NewBufferedWriter returns a new Writer that compresses to w, using the
  182. // framing format described at
  183. // https://github.com/google/snappy/blob/master/framing_format.txt
  184. //
  185. // The Writer returned buffers writes. Users must call Close to guarantee all
  186. // data has been forwarded to the underlying io.Writer. They may also call
  187. // Flush zero or more times before calling Close.
  188. func NewBufferedWriter(w io.Writer) *Writer {
  189. return &Writer{
  190. w: w,
  191. ibuf: make([]byte, 0, maxUncompressedChunkLen),
  192. obuf: make([]byte, obufLen),
  193. }
  194. }
  195. // Writer is an io.Writer than can write Snappy-compressed bytes.
  196. type Writer struct {
  197. w io.Writer
  198. err error
  199. // ibuf is a buffer for the incoming (uncompressed) bytes.
  200. //
  201. // Its use is optional. For backwards compatibility, Writers created by the
  202. // NewWriter function have ibuf == nil, do not buffer incoming bytes, and
  203. // therefore do not need to be Flush'ed or Close'd.
  204. ibuf []byte
  205. // obuf is a buffer for the outgoing (compressed) bytes.
  206. obuf []byte
  207. // wroteStreamHeader is whether we have written the stream header.
  208. wroteStreamHeader bool
  209. }
  210. // Reset discards the writer's state and switches the Snappy writer to write to
  211. // w. This permits reusing a Writer rather than allocating a new one.
  212. func (w *Writer) Reset(writer io.Writer) {
  213. w.w = writer
  214. w.err = nil
  215. if w.ibuf != nil {
  216. w.ibuf = w.ibuf[:0]
  217. }
  218. w.wroteStreamHeader = false
  219. }
  220. // Write satisfies the io.Writer interface.
  221. func (w *Writer) Write(p []byte) (nRet int, errRet error) {
  222. if w.ibuf == nil {
  223. // Do not buffer incoming bytes. This does not perform or compress well
  224. // if the caller of Writer.Write writes many small slices. This
  225. // behavior is therefore deprecated, but still supported for backwards
  226. // compatibility with code that doesn't explicitly Flush or Close.
  227. return w.write(p)
  228. }
  229. // The remainder of this method is based on bufio.Writer.Write from the
  230. // standard library.
  231. for len(p) > (cap(w.ibuf)-len(w.ibuf)) && w.err == nil {
  232. var n int
  233. if len(w.ibuf) == 0 {
  234. // Large write, empty buffer.
  235. // Write directly from p to avoid copy.
  236. n, _ = w.write(p)
  237. } else {
  238. n = copy(w.ibuf[len(w.ibuf):cap(w.ibuf)], p)
  239. w.ibuf = w.ibuf[:len(w.ibuf)+n]
  240. w.Flush()
  241. }
  242. nRet += n
  243. p = p[n:]
  244. }
  245. if w.err != nil {
  246. return nRet, w.err
  247. }
  248. n := copy(w.ibuf[len(w.ibuf):cap(w.ibuf)], p)
  249. w.ibuf = w.ibuf[:len(w.ibuf)+n]
  250. nRet += n
  251. return nRet, nil
  252. }
  253. func (w *Writer) write(p []byte) (nRet int, errRet error) {
  254. if w.err != nil {
  255. return 0, w.err
  256. }
  257. for len(p) > 0 {
  258. obufStart := len(magicChunk)
  259. if !w.wroteStreamHeader {
  260. w.wroteStreamHeader = true
  261. copy(w.obuf, magicChunk)
  262. obufStart = 0
  263. }
  264. var uncompressed []byte
  265. if len(p) > maxUncompressedChunkLen {
  266. uncompressed, p = p[:maxUncompressedChunkLen], p[maxUncompressedChunkLen:]
  267. } else {
  268. uncompressed, p = p, nil
  269. }
  270. checksum := crc(uncompressed)
  271. // Compress the buffer, discarding the result if the improvement
  272. // isn't at least 12.5%.
  273. compressed := Encode(w.obuf[obufHeaderLen:], uncompressed)
  274. chunkType := uint8(chunkTypeCompressedData)
  275. chunkLen := 4 + len(compressed)
  276. obufEnd := obufHeaderLen + len(compressed)
  277. if len(compressed) >= len(uncompressed)-len(uncompressed)/8 {
  278. chunkType = chunkTypeUncompressedData
  279. chunkLen = 4 + len(uncompressed)
  280. obufEnd = obufHeaderLen
  281. }
  282. // Fill in the per-chunk header that comes before the body.
  283. w.obuf[len(magicChunk)+0] = chunkType
  284. w.obuf[len(magicChunk)+1] = uint8(chunkLen >> 0)
  285. w.obuf[len(magicChunk)+2] = uint8(chunkLen >> 8)
  286. w.obuf[len(magicChunk)+3] = uint8(chunkLen >> 16)
  287. w.obuf[len(magicChunk)+4] = uint8(checksum >> 0)
  288. w.obuf[len(magicChunk)+5] = uint8(checksum >> 8)
  289. w.obuf[len(magicChunk)+6] = uint8(checksum >> 16)
  290. w.obuf[len(magicChunk)+7] = uint8(checksum >> 24)
  291. if _, err := w.w.Write(w.obuf[obufStart:obufEnd]); err != nil {
  292. w.err = err
  293. return nRet, err
  294. }
  295. if chunkType == chunkTypeUncompressedData {
  296. if _, err := w.w.Write(uncompressed); err != nil {
  297. w.err = err
  298. return nRet, err
  299. }
  300. }
  301. nRet += len(uncompressed)
  302. }
  303. return nRet, nil
  304. }
  305. // Flush flushes the Writer to its underlying io.Writer.
  306. func (w *Writer) Flush() error {
  307. if w.err != nil {
  308. return w.err
  309. }
  310. if len(w.ibuf) == 0 {
  311. return nil
  312. }
  313. w.write(w.ibuf)
  314. w.ibuf = w.ibuf[:0]
  315. return w.err
  316. }
  317. // Close calls Flush and then closes the Writer.
  318. func (w *Writer) Close() error {
  319. w.Flush()
  320. ret := w.err
  321. if w.err == nil {
  322. w.err = errClosed
  323. }
  324. return ret
  325. }