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