1// Copyright (c) 2017 Couchbase, Inc. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package zap 16 17import ( 18 "bytes" 19 "encoding/binary" 20 "io" 21 22 "github.com/Smerity/govarint" 23) 24 25type chunkedIntCoder struct { 26 final []byte 27 maxDocNum uint64 28 chunkSize uint64 29 chunkBuf bytes.Buffer 30 encoder *govarint.Base128Encoder 31 chunkLens []uint64 32 currChunk uint64 33 34 buf []byte 35} 36 37// newChunkedIntCoder returns a new chunk int coder which packs data into 38// chunks based on the provided chunkSize and supports up to the specified 39// maxDocNum 40func newChunkedIntCoder(chunkSize uint64, maxDocNum uint64) *chunkedIntCoder { 41 total := maxDocNum/chunkSize + 1 42 rv := &chunkedIntCoder{ 43 chunkSize: chunkSize, 44 maxDocNum: maxDocNum, 45 chunkLens: make([]uint64, total), 46 final: make([]byte, 0, 64), 47 } 48 rv.encoder = govarint.NewU64Base128Encoder(&rv.chunkBuf) 49 50 return rv 51} 52 53// Reset lets you reuse this chunked int coder. buffers are reset and reused 54// from previous use. you cannot change the chunk size or max doc num. 55func (c *chunkedIntCoder) Reset() { 56 c.final = c.final[:0] 57 c.chunkBuf.Reset() 58 c.currChunk = 0 59 for i := range c.chunkLens { 60 c.chunkLens[i] = 0 61 } 62} 63 64// Add encodes the provided integers into the correct chunk for the provided 65// doc num. You MUST call Add() with increasing docNums. 66func (c *chunkedIntCoder) Add(docNum uint64, vals ...uint64) error { 67 chunk := docNum / c.chunkSize 68 if chunk != c.currChunk { 69 // starting a new chunk 70 if c.encoder != nil { 71 // close out last 72 c.Close() 73 c.chunkBuf.Reset() 74 } 75 c.currChunk = chunk 76 } 77 78 for _, val := range vals { 79 _, err := c.encoder.PutU64(val) 80 if err != nil { 81 return err 82 } 83 } 84 85 return nil 86} 87 88// Close indicates you are done calling Add() this allows the final chunk 89// to be encoded. 90func (c *chunkedIntCoder) Close() { 91 c.encoder.Close() 92 encodingBytes := c.chunkBuf.Bytes() 93 c.chunkLens[c.currChunk] = uint64(len(encodingBytes)) 94 c.final = append(c.final, encodingBytes...) 95} 96 97// Write commits all the encoded chunked integers to the provided writer. 98func (c *chunkedIntCoder) Write(w io.Writer) (int, error) { 99 bufNeeded := binary.MaxVarintLen64 * (1 + len(c.chunkLens)) 100 if len(c.buf) < bufNeeded { 101 c.buf = make([]byte, bufNeeded) 102 } 103 buf := c.buf 104 105 // write out the number of chunks & each chunkLen 106 n := binary.PutUvarint(buf, uint64(len(c.chunkLens))) 107 for _, chunkLen := range c.chunkLens { 108 n += binary.PutUvarint(buf[n:], uint64(chunkLen)) 109 } 110 111 tw, err := w.Write(buf[:n]) 112 if err != nil { 113 return tw, err 114 } 115 116 // write out the data 117 nw, err := w.Write(c.final) 118 tw += nw 119 if err != nil { 120 return tw, err 121 } 122 return tw, nil 123} 124