github.com/chain5j/chain5j-pkg@v1.0.7/collection/maps/treemap/treemap.go (about)

     1  // Package treemap
     2  //
     3  // @author: xwc1125
     4  package treemap
     5  
     6  import (
     7  	"encoding/json"
     8  	"io"
     9  
    10  	"github.com/chain5j/chain5j-pkg/codec/rlp"
    11  	"github.com/chain5j/chain5j-pkg/collection/maps/treemap/v2"
    12  	"golang.org/x/exp/constraints"
    13  )
    14  
    15  // TreeMap based on red-black tree, alias of RedBlackTree.
    16  type TreeMap[K constraints.Ordered, V any] struct {
    17  	*treemap.TreeMap[K, V]
    18  }
    19  
    20  func NewKvs[K constraints.Ordered, V any](t *treemap.TreeMap[K, V]) *TreeMap[K, V] {
    21  	return &TreeMap[K, V]{
    22  		TreeMap: t,
    23  	}
    24  }
    25  
    26  func (k TreeMap[K, V]) EncodeRLP(w io.Writer) error {
    27  	// [xwc1125] 如果使用，encoding/json返回的是64进制
    28  	// 如果使用，github.com/json-iterator/go返回的是16进制
    29  	json, err := MarshalJSON(k.TreeMap)
    30  	if err != nil {
    31  		return err
    32  	}
    33  	return rlp.Encode(w, json)
    34  }
    35  
    36  func (k *TreeMap[K, V]) DecodeRLP(s *rlp.Stream) error {
    37  	var jsonBytes []byte
    38  	err := s.Decode(&jsonBytes)
    39  	if err != nil {
    40  		return err
    41  	}
    42  	var treeMap1 = treemap.New[K, V]()
    43  	err = UnmarshalJSON(jsonBytes, treeMap1)
    44  	if err != nil {
    45  		return err
    46  	}
    47  	*k.TreeMap = *treeMap1
    48  	return err
    49  }
    50  
    51  func (k *TreeMap[K, V]) Clone() *TreeMap[K, V] {
    52  	clone := Clone(k.TreeMap)
    53  	return &TreeMap[K, V]{
    54  		TreeMap: clone,
    55  	}
    56  }
    57  
    58  func (k *TreeMap[K, V]) MarshalJSON() ([]byte, error) {
    59  	return MarshalJSON(k.TreeMap)
    60  }
    61  
    62  func MarshalJSON[K constraints.Ordered, V any](t *treemap.TreeMap[K, V]) ([]byte, error) {
    63  	return json.Marshal(Map(t))
    64  }
    65  
    66  func Map[K constraints.Ordered, V any](t *treemap.TreeMap[K, V]) map[K]V {
    67  	m := make(map[K]V, t.Len())
    68  	for it := t.Iterator(); it.Valid(); it.Next() {
    69  		m[it.Key()] = it.Value()
    70  	}
    71  	return m
    72  }
    73  
    74  func (k *TreeMap[K, V]) UnmarshalJSON(data []byte) error {
    75  	if k.TreeMap == nil {
    76  		k.TreeMap = treemap.New[K, V]()
    77  	}
    78  	return UnmarshalJSON(data, k.TreeMap)
    79  }
    80  
    81  func UnmarshalJSON[K constraints.Ordered, V any](b []byte, t *treemap.TreeMap[K, V]) error {
    82  	var data map[K]V
    83  	if err := json.Unmarshal(b, &data); err != nil {
    84  		return err
    85  	}
    86  	for k, v := range data {
    87  		t.Set(k, v)
    88  	}
    89  	return nil
    90  }
    91  
    92  func Clone[K constraints.Ordered, V any](t *treemap.TreeMap[K, V]) *treemap.TreeMap[K, V] {
    93  	newTree := treemap.NewWithKeyCompare[K, V](t.KeyCompare())
    94  	m := Map(t)
    95  	for k, v := range m {
    96  		newTree.Set(k, v)
    97  	}
    98  	return newTree
    99  }