github.com/ledgerwatch/erigon-lib@v1.0.0/bptree/key_factory.go (about)

     1  /*
     2     Copyright 2022 Erigon contributors
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package bptree
    18  
    19  import (
    20  	"bufio"
    21  	"encoding/binary"
    22  	"fmt"
    23  	"io"
    24  	"sort"
    25  )
    26  
    27  type KeyFactory interface {
    28  	NewUniqueKeyValues(reader *bufio.Reader) KeyValues
    29  	NewUniqueKeys(reader *bufio.Reader) Keys
    30  }
    31  
    32  type KeyBinaryFactory struct {
    33  	keySize int
    34  }
    35  
    36  func NewKeyBinaryFactory(keySize int) KeyFactory {
    37  	return &KeyBinaryFactory{keySize: keySize}
    38  }
    39  
    40  func (factory *KeyBinaryFactory) NewUniqueKeyValues(reader *bufio.Reader) KeyValues {
    41  	kvPairs := factory.readUniqueKeyValues(reader)
    42  	sort.Sort(kvPairs)
    43  	return kvPairs
    44  }
    45  
    46  func (factory *KeyBinaryFactory) NewUniqueKeys(reader *bufio.Reader) Keys {
    47  	keys := factory.readUniqueKeys(reader)
    48  	sort.Sort(keys)
    49  	return keys
    50  }
    51  
    52  func (factory *KeyBinaryFactory) readUniqueKeyValues(reader *bufio.Reader) KeyValues {
    53  	kvPairs := KeyValues{make([]*Felt, 0), make([]*Felt, 0)}
    54  	keyRegistry := make(map[Felt]bool)
    55  	buffer := make([]byte, BufferSize)
    56  	for {
    57  		bytesRead, err := reader.Read(buffer)
    58  		ensure(err == nil || err == io.EOF, fmt.Sprintf("readUniqueKeyValues: read error %s\n", err))
    59  		if err == io.EOF {
    60  			break
    61  		}
    62  		keyBytesCount := factory.keySize * (bytesRead / factory.keySize)
    63  		duplicatedKeys := 0
    64  		for i := 0; i < keyBytesCount; i += factory.keySize {
    65  			key := factory.readKey(buffer, i)
    66  			if _, duplicated := keyRegistry[key]; duplicated {
    67  				duplicatedKeys++
    68  				continue
    69  			}
    70  			keyRegistry[key] = true
    71  			value := key // Shortcut: value equal to key
    72  			kvPairs.keys = append(kvPairs.keys, &key)
    73  			kvPairs.values = append(kvPairs.values, &value)
    74  		}
    75  	}
    76  	return kvPairs
    77  }
    78  
    79  func (factory *KeyBinaryFactory) readUniqueKeys(reader *bufio.Reader) Keys {
    80  	keys := make(Keys, 0)
    81  	keyRegistry := make(map[Felt]bool)
    82  	buffer := make([]byte, BufferSize)
    83  	for {
    84  		bytesRead, err := reader.Read(buffer)
    85  		if err == io.EOF {
    86  			break
    87  		}
    88  		keyBytesCount := factory.keySize * (bytesRead / factory.keySize)
    89  		duplicatedKeys := 0
    90  		for i := 0; i < keyBytesCount; i += factory.keySize {
    91  			key := factory.readKey(buffer, i)
    92  			if _, duplicated := keyRegistry[key]; duplicated {
    93  				duplicatedKeys++
    94  				continue
    95  			}
    96  			keyRegistry[key] = true
    97  			keys = append(keys, key)
    98  		}
    99  	}
   100  	return keys
   101  }
   102  
   103  func (factory *KeyBinaryFactory) readKey(buffer []byte, offset int) Felt {
   104  	keySlice := buffer[offset : offset+factory.keySize]
   105  	switch factory.keySize {
   106  	case 1:
   107  		return Felt(keySlice[0])
   108  	case 2:
   109  		return Felt(binary.BigEndian.Uint16(keySlice))
   110  	case 4:
   111  		return Felt(binary.BigEndian.Uint32(keySlice))
   112  	default:
   113  		return Felt(binary.BigEndian.Uint64(keySlice))
   114  	}
   115  }