github.com/grailbio/base@v0.0.11/mapio/buf.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package mapio
     6  
     7  import (
     8  	"bytes"
     9  	"sort"
    10  )
    11  
    12  // A Buf is an unordered write buffer for maps. It holds entries
    13  // in memory; these are then sorted and written to a map.
    14  type Buf struct {
    15  	keys, values       [][]byte
    16  	keySize, valueSize int
    17  }
    18  
    19  // Append append the given entry to the buffer.
    20  func (b *Buf) Append(key, value []byte) {
    21  	keyCopy := make([]byte, len(key))
    22  	copy(keyCopy, key)
    23  	valueCopy := make([]byte, len(value))
    24  	copy(valueCopy, value)
    25  	b.keys = append(b.keys, keyCopy)
    26  	b.values = append(b.values, valueCopy)
    27  	b.keySize += len(keyCopy)
    28  	b.valueSize += len(valueCopy)
    29  }
    30  
    31  // Size returns the number size of this buffer in bytes.
    32  func (b *Buf) Size() int { return b.keySize + b.valueSize }
    33  
    34  // Len implements sort.Interface
    35  func (b *Buf) Len() int { return len(b.keys) }
    36  
    37  // Less implements sort.Interface
    38  func (b *Buf) Less(i, j int) bool { return bytes.Compare(b.keys[i], b.keys[j]) < 0 }
    39  
    40  // Swap implements sort.Interface
    41  func (b *Buf) Swap(i, j int) {
    42  	b.keys[i], b.keys[j] = b.keys[j], b.keys[i]
    43  	b.values[i], b.values[j] = b.values[j], b.values[i]
    44  }
    45  
    46  // WriteTo sorts and then writes all of the entries in this buffer to
    47  // the provided writer.
    48  func (b *Buf) WriteTo(w *Writer) error {
    49  	sort.Sort(b)
    50  	for i := range b.keys {
    51  		if err := w.Append(b.keys[i], b.values[i]); err != nil {
    52  			return err
    53  		}
    54  	}
    55  	return nil
    56  }