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

     1  /*
     2   * Atree - Scalable Arrays and Ordered Maps
     3   *
     4   * Copyright 2021 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 main
    20  
    21  import (
    22  	"flag"
    23  	"fmt"
    24  
    25  	"github.com/onflow/atree"
    26  
    27  	"github.com/fxamacker/cbor/v2"
    28  )
    29  
    30  const cborTagUInt64Value = 164
    31  
    32  type Uint64Value uint64
    33  
    34  var _ atree.Value = Uint64Value(0)
    35  var _ atree.Storable = Uint64Value(0)
    36  
    37  func (v Uint64Value) ChildStorables() []atree.Storable {
    38  	return nil
    39  }
    40  
    41  func (v Uint64Value) StoredValue(_ atree.SlabStorage) (atree.Value, error) {
    42  	return v, nil
    43  }
    44  
    45  func (v Uint64Value) Storable(_ atree.SlabStorage, _ atree.Address, _ uint64) (atree.Storable, error) {
    46  	return v, nil
    47  }
    48  
    49  // Encode encodes UInt64Value as
    50  //
    51  //	cbor.Tag{
    52  //			Number:  cborTagUInt64Value,
    53  //			Content: uint64(v),
    54  //	}
    55  func (v Uint64Value) Encode(enc *atree.Encoder) error {
    56  	err := enc.CBOR.EncodeRawBytes([]byte{
    57  		// tag number
    58  		0xd8, cborTagUInt64Value,
    59  	})
    60  	if err != nil {
    61  		return err
    62  	}
    63  	return enc.CBOR.EncodeUint64(uint64(v))
    64  }
    65  
    66  // TODO: cache size
    67  func (v Uint64Value) ByteSize() uint32 {
    68  	// tag number (2 bytes) + encoded content
    69  	return 2 + atree.GetUintCBORSize(uint64(v))
    70  }
    71  
    72  func (v Uint64Value) String() string {
    73  	return fmt.Sprintf("%d", uint64(v))
    74  }
    75  
    76  type testTypeInfo struct{}
    77  
    78  var _ atree.TypeInfo = testTypeInfo{}
    79  
    80  func (testTypeInfo) Encode(e *cbor.StreamEncoder) error {
    81  	return e.EncodeUint8(42)
    82  }
    83  
    84  func (i testTypeInfo) Equal(other atree.TypeInfo) bool {
    85  	_, ok := other.(testTypeInfo)
    86  	return ok
    87  }
    88  
    89  func decodeStorable(dec *cbor.StreamDecoder, _ atree.StorageID) (atree.Storable, error) {
    90  	tagNumber, err := dec.DecodeTagNumber()
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  
    95  	switch tagNumber {
    96  	case atree.CBORTagStorageID:
    97  		return atree.DecodeStorageIDStorable(dec)
    98  
    99  	case cborTagUInt64Value:
   100  		n, err := dec.DecodeUint64()
   101  		if err != nil {
   102  			return nil, err
   103  		}
   104  		return Uint64Value(n), nil
   105  
   106  	default:
   107  		return nil, fmt.Errorf("invalid tag number %d", tagNumber)
   108  	}
   109  }
   110  
   111  // TODO: implement different slab size for metadata slab and data slab.
   112  func main() {
   113  	var slabSize uint64
   114  	var numElements uint64
   115  	var verbose bool
   116  
   117  	flag.Uint64Var(&slabSize, "size", 1024, "slab size in bytes")
   118  	flag.Uint64Var(&numElements, "count", 500, "number of elements in array")
   119  	flag.BoolVar(&verbose, "verbose", false, "verbose output")
   120  
   121  	flag.Parse()
   122  
   123  	minThreshold, maxThreshold, _, _ := atree.SetThreshold(slabSize)
   124  
   125  	fmt.Printf(
   126  		"Inserting %d elements (uint64) into array with slab size %d, min size %d, and max size %d ...\n",
   127  		numElements,
   128  		slabSize,
   129  		minThreshold,
   130  		maxThreshold,
   131  	)
   132  
   133  	encMode, err := cbor.EncOptions{}.EncMode()
   134  	if err != nil {
   135  		fmt.Println(err)
   136  		return
   137  	}
   138  
   139  	decMode, err := cbor.DecOptions{}.DecMode()
   140  	if err != nil {
   141  		fmt.Println(err)
   142  		return
   143  	}
   144  
   145  	storage := atree.NewBasicSlabStorage(encMode, decMode, decodeStorable, decodeTypeInfo)
   146  
   147  	typeInfo := testTypeInfo{}
   148  
   149  	address := atree.Address{1, 2, 3, 4, 5, 6, 7, 8}
   150  
   151  	array, err := atree.NewArray(storage, address, typeInfo)
   152  
   153  	if err != nil {
   154  		fmt.Println(err)
   155  		return
   156  	}
   157  
   158  	for i := uint64(0); i < numElements; i++ {
   159  		err := array.Append(Uint64Value(i))
   160  		if err != nil {
   161  			fmt.Println(err)
   162  			return
   163  		}
   164  	}
   165  
   166  	stats, err := atree.GetArrayStats(array)
   167  	if err != nil {
   168  		fmt.Println(err)
   169  		return
   170  	}
   171  
   172  	fmt.Printf("%+v\n", stats)
   173  
   174  	if verbose {
   175  		fmt.Printf("\n\n=========== array layout ===========\n")
   176  		atree.PrintArray(array)
   177  	}
   178  }
   179  
   180  func decodeTypeInfo(_ *cbor.StreamDecoder) (atree.TypeInfo, error) {
   181  	return testTypeInfo{}, nil
   182  }