github.com/balzaczyy/golucene@v0.0.0-20151210033525-d0be9ee89713/core/util/packed/writer.go (about)

     1  package packed
     2  
     3  import (
     4  	"errors"
     5  	"github.com/balzaczyy/golucene/core/codec"
     6  )
     7  
     8  // util/packed/PackedInts.java#Writer
     9  
    10  /* A write-once Writer. */
    11  type Writer interface {
    12  	writeHeader() error
    13  	// Add a value to the stream.
    14  	Add(v int64) error
    15  	// The number of bits per value.
    16  	BitsPerValue() int
    17  	// Perform end-of-stream operations.
    18  	Finish() error
    19  }
    20  
    21  type WriterImpl struct {
    22  	out          DataOutput
    23  	valueCount   int
    24  	bitsPerValue int
    25  	format       PackedFormat
    26  }
    27  
    28  func newWriter(out DataOutput, valueCount, bitsPerValue int,
    29  	format PackedFormat) *WriterImpl {
    30  	assert(bitsPerValue <= 64)
    31  	assert(valueCount >= 0 || valueCount == -1)
    32  	return &WriterImpl{out, valueCount, bitsPerValue, format}
    33  }
    34  
    35  func (w *WriterImpl) writeHeader() error {
    36  	assert(w.valueCount != -1)
    37  	err := codec.WriteHeader(w.out, PACKED_CODEC_NAME, VERSION_CURRENT)
    38  	if err == nil {
    39  		err = w.out.WriteVInt(int32(w.bitsPerValue))
    40  		if err == nil {
    41  			err = w.out.WriteVInt(int32(w.valueCount))
    42  			if err == nil {
    43  				err = w.out.WriteVInt(int32(PackedFormat(w.format).Id()))
    44  			}
    45  		}
    46  	}
    47  	return err
    48  }
    49  
    50  func (w *WriterImpl) BitsPerValue() int {
    51  	return w.bitsPerValue
    52  }
    53  
    54  // util/packed/PackedWriter.java
    55  
    56  type PackedWriter struct {
    57  	*WriterImpl
    58  	finished   bool
    59  	format     PackedFormat
    60  	encoder    BulkOperation
    61  	nextBlocks []byte
    62  	nextValues []int64
    63  	iterations int
    64  	off        int
    65  	written    int
    66  }
    67  
    68  func newPackedWriter(format PackedFormat, out DataOutput,
    69  	valueCount, bitsPerValue, mem int) *PackedWriter {
    70  
    71  	encoder := newBulkOperation(format, uint32(bitsPerValue))
    72  	iterations := encoder.computeIterations(valueCount, mem)
    73  	return &PackedWriter{
    74  		WriterImpl: newWriter(out, valueCount, bitsPerValue, format),
    75  		format:     format,
    76  		encoder:    encoder,
    77  		iterations: iterations,
    78  		nextBlocks: make([]byte, iterations*encoder.ByteBlockCount()),
    79  		nextValues: make([]int64, iterations*encoder.ByteValueCount()),
    80  	}
    81  }
    82  
    83  func (w *PackedWriter) Add(v int64) error {
    84  	assert(UnsignedBitsRequired(v) <= w.bitsPerValue)
    85  	assert(!w.finished)
    86  	if w.valueCount != -1 && w.written >= w.valueCount {
    87  		return errors.New("Writing past end of stream")
    88  	}
    89  	w.nextValues[w.off] = v
    90  	if w.off++; w.off == len(w.nextValues) {
    91  		err := w.flush()
    92  		if err != nil {
    93  			return err
    94  		}
    95  	}
    96  	w.written++
    97  	return nil
    98  }
    99  
   100  func (w *PackedWriter) Finish() error {
   101  	assert(!w.finished)
   102  	var err error
   103  	if w.valueCount != -1 {
   104  		for w.written < w.valueCount && err == nil {
   105  			err = w.Add(0)
   106  		}
   107  	}
   108  	if err == nil {
   109  		err = w.flush()
   110  	}
   111  	w.finished = err == nil
   112  	return err
   113  }
   114  
   115  func (w *PackedWriter) flush() error {
   116  	w.encoder.encodeLongToByte(w.nextValues, w.nextBlocks, w.iterations)
   117  	blockCount := int(w.format.ByteCount(VERSION_CURRENT,
   118  		int32(w.off), uint32(w.bitsPerValue)))
   119  	err := w.out.WriteBytes(w.nextBlocks[:blockCount])
   120  	if err != nil {
   121  		return err
   122  	}
   123  	for i, _ := range w.nextValues {
   124  		w.nextValues[i] = 0
   125  	}
   126  	w.off = 0
   127  	return nil
   128  }