github.com/onflow/atree@v0.6.0/mapcollision_bench_test.go (about)

     1  /*
     2   * Atree - Scalable Arrays and Ordered Maps
     3   *
     4   * Copyright 2022 Dapper Labs, Inc.
     5   *
     6   * Licensed under the Apache License, Version 2.0 (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   *   http://www.apache.org/licenses/LICENSE-2.0
    11   *
    12   * Unless required by applicable law or agreed to in writing, software
    13   * distributed under the License is distributed on an "AS IS" BASIS,
    14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15   * See the License for the specific language governing permissions and
    16   * limitations under the License.
    17   */
    18  
    19  package atree
    20  
    21  import (
    22  	"encoding/binary"
    23  	"fmt"
    24  	"testing"
    25  
    26  	"github.com/stretchr/testify/require"
    27  	"github.com/zeebo/blake3"
    28  )
    29  
    30  type collisionDigesterBuilder struct {
    31  	digest            uint64
    32  	collisionCount    uint32
    33  	maxCollisionCount uint32
    34  }
    35  
    36  var _ DigesterBuilder = &collisionDigesterBuilder{}
    37  
    38  func NewCollisionDigesterBuilder(maxCollisionLimitPerDigest uint32) DigesterBuilder {
    39  	return &collisionDigesterBuilder{
    40  		maxCollisionCount: maxCollisionLimitPerDigest + 1,
    41  	}
    42  }
    43  
    44  func (db *collisionDigesterBuilder) Digest(hip HashInputProvider, value Value) (Digester, error) {
    45  
    46  	if db.collisionCount < db.maxCollisionCount {
    47  		db.collisionCount++
    48  	} else {
    49  		db.digest++
    50  		db.collisionCount = 0
    51  	}
    52  	firstLevelHash := db.digest
    53  
    54  	var scratch [32]byte
    55  	msg, err := hip(value, scratch[:])
    56  	if err != nil {
    57  		return nil, err
    58  	}
    59  
    60  	return &collisionDigester{
    61  		firstLevelHash: firstLevelHash,
    62  		msg:            msg,
    63  	}, nil
    64  }
    65  
    66  func (db *collisionDigesterBuilder) SetSeed(k1 uint64, k2 uint64) {
    67  }
    68  
    69  type collisionDigester struct {
    70  	firstLevelHash uint64
    71  	blake3Hash     [4]uint64
    72  	msg            []byte
    73  }
    74  
    75  var _ Digester = &collisionDigester{}
    76  
    77  func (d *collisionDigester) Digest(level uint) (Digest, error) {
    78  	if level >= d.Levels() {
    79  		return Digest(0), fmt.Errorf("invalid digest level %d", level)
    80  	}
    81  
    82  	switch level {
    83  	case 0:
    84  		return Digest(d.firstLevelHash), nil
    85  	default:
    86  		if d.blake3Hash == emptyBlake3Hash {
    87  			sum := blake3.Sum256(d.msg)
    88  			d.blake3Hash[0] = binary.BigEndian.Uint64(sum[:])
    89  			d.blake3Hash[1] = binary.BigEndian.Uint64(sum[8:])
    90  			d.blake3Hash[2] = binary.BigEndian.Uint64(sum[16:])
    91  			d.blake3Hash[3] = binary.BigEndian.Uint64(sum[24:])
    92  		}
    93  		return Digest(d.blake3Hash[level-1]), nil
    94  	}
    95  }
    96  
    97  func (d *collisionDigester) DigestPrefix(level uint) ([]Digest, error) {
    98  	return nil, nil
    99  }
   100  
   101  func (d *collisionDigester) Levels() uint {
   102  	return 4
   103  }
   104  
   105  func (d *collisionDigester) Reset() {
   106  }
   107  
   108  func BenchmarkCollisionPerDigest(b *testing.B) {
   109  
   110  	savedMaxCollisionLimitPerDigest := MaxCollisionLimitPerDigest
   111  	defer func() {
   112  		MaxCollisionLimitPerDigest = savedMaxCollisionLimitPerDigest
   113  	}()
   114  
   115  	const mapCount = 1_000_000
   116  
   117  	collisionPerDigests := []uint32{0, 10, 255, 500, 1_000, 2_000, 5_000, 10_000}
   118  
   119  	for _, collisionPerDigest := range collisionPerDigests {
   120  
   121  		name := fmt.Sprintf("%d elements %d collision per digest", mapCount, collisionPerDigest)
   122  
   123  		b.Run(name, func(b *testing.B) {
   124  
   125  			MaxCollisionLimitPerDigest = collisionPerDigest
   126  
   127  			digesterBuilder := NewCollisionDigesterBuilder(collisionPerDigest)
   128  			keyValues := make(map[Value]Value, mapCount)
   129  			for i := uint64(0); i < mapCount; i++ {
   130  				k := Uint64Value(i)
   131  				v := Uint64Value(i)
   132  				keyValues[k] = v
   133  			}
   134  
   135  			typeInfo := testTypeInfo{42}
   136  			address := Address{1, 2, 3, 4, 5, 6, 7, 8}
   137  			storage := newTestPersistentStorage(b)
   138  
   139  			m, err := NewMap(storage, address, digesterBuilder, typeInfo)
   140  			require.NoError(b, err)
   141  
   142  			b.StartTimer()
   143  
   144  			for i := 0; i < b.N; i++ {
   145  				for k, v := range keyValues {
   146  					_, _ = m.Set(compare, hashInputProvider, k, v)
   147  				}
   148  			}
   149  		})
   150  	}
   151  }