github.com/aergoio/aergo@v1.3.1/contract/system/voteresult.go (about)

     1  package system
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"errors"
     7  	"math/big"
     8  	"sort"
     9  
    10  	"github.com/aergoio/aergo/internal/enc"
    11  	"github.com/aergoio/aergo/state"
    12  	"github.com/aergoio/aergo/types"
    13  	"github.com/mr-tron/base58"
    14  )
    15  
    16  type VoteResult struct {
    17  	rmap map[string]*big.Int
    18  	key  []byte
    19  	ex   bool
    20  }
    21  
    22  func newVoteResult(key []byte) *VoteResult {
    23  	voteResult := &VoteResult{}
    24  	voteResult.rmap = map[string]*big.Int{}
    25  	if bytes.Equal(key, defaultVoteKey) {
    26  		voteResult.ex = false
    27  	} else {
    28  		voteResult.ex = true
    29  	}
    30  	voteResult.key = key
    31  	return voteResult
    32  }
    33  
    34  func (voteResult *VoteResult) SubVote(vote *types.Vote) error {
    35  	if voteResult.ex {
    36  		if vote.Candidate != nil {
    37  			var args []string
    38  			err := json.Unmarshal(vote.Candidate, &args)
    39  			if err != nil {
    40  				return err
    41  			}
    42  			for _, v := range args {
    43  				voteResult.rmap[v] = new(big.Int).Sub(voteResult.rmap[v], vote.GetAmountBigInt())
    44  			}
    45  		}
    46  	} else {
    47  		for offset := 0; offset < len(vote.Candidate); offset += PeerIDLength {
    48  			peer := vote.Candidate[offset : offset+PeerIDLength]
    49  			pkey := base58.Encode(peer)
    50  			voteResult.rmap[pkey] = new(big.Int).Sub(voteResult.rmap[pkey], vote.GetAmountBigInt())
    51  		}
    52  	}
    53  	return nil
    54  }
    55  
    56  func (voteResult *VoteResult) AddVote(vote *types.Vote) error {
    57  	if voteResult.ex {
    58  		var args []string
    59  		err := json.Unmarshal(vote.Candidate, &args)
    60  		if err != nil {
    61  			return err
    62  		}
    63  		for _, v := range args {
    64  			if voteResult.rmap[v] == nil {
    65  				voteResult.rmap[v] = new(big.Int).SetUint64(0)
    66  			}
    67  			voteResult.rmap[v] = new(big.Int).Add(voteResult.rmap[v], vote.GetAmountBigInt())
    68  		}
    69  	} else {
    70  		for offset := 0; offset < len(vote.Candidate); offset += PeerIDLength {
    71  			key := vote.Candidate[offset : offset+PeerIDLength]
    72  			if voteResult.rmap[base58.Encode(key)] == nil {
    73  				voteResult.rmap[base58.Encode(key)] = new(big.Int).SetUint64(0)
    74  			}
    75  			voteResult.rmap[base58.Encode(key)] = new(big.Int).Add(voteResult.rmap[base58.Encode(key)], vote.GetAmountBigInt())
    76  		}
    77  	}
    78  	return nil
    79  }
    80  
    81  func (vr *VoteResult) buildVoteList() *types.VoteList {
    82  	var voteList types.VoteList
    83  	for k, v := range vr.rmap {
    84  		vote := &types.Vote{
    85  			Amount: v.Bytes(),
    86  		}
    87  		if vr.ex {
    88  			vote.Candidate = []byte(k)
    89  		} else {
    90  			vote.Candidate, _ = enc.ToBytes(k)
    91  		}
    92  		voteList.Votes = append(voteList.Votes, vote)
    93  	}
    94  	sort.Sort(sort.Reverse(voteList))
    95  
    96  	return &voteList
    97  }
    98  
    99  func (vr *VoteResult) Sync(scs *state.ContractState) error {
   100  	return scs.SetData(append(sortKey, vr.key...), serializeVoteList(vr.buildVoteList(), vr.ex))
   101  }
   102  
   103  func loadVoteResult(scs *state.ContractState, key []byte) (*VoteResult, error) {
   104  	data, err := scs.GetData(append(sortKey, key...))
   105  	if err != nil {
   106  		return nil, err
   107  	}
   108  	voteResult := newVoteResult(key)
   109  	if len(data) != 0 {
   110  		voteList := deserializeVoteList(data, voteResult.ex)
   111  		if voteList != nil {
   112  			for _, v := range voteList.GetVotes() {
   113  				if voteResult.ex {
   114  					voteResult.rmap[string(v.Candidate)] = v.GetAmountBigInt()
   115  				} else {
   116  					voteResult.rmap[base58.Encode(v.Candidate)] = v.GetAmountBigInt()
   117  				}
   118  			}
   119  		}
   120  	}
   121  	return voteResult, nil
   122  }
   123  
   124  func InitVoteResult(scs *state.ContractState, voteResult map[string]*big.Int) error {
   125  	if voteResult == nil {
   126  		return errors.New("Invalid argument : voteReult should not nil")
   127  	}
   128  	res := newVoteResult(defaultVoteKey)
   129  	res.rmap = voteResult
   130  	return res.Sync(scs)
   131  }
   132  
   133  func getVoteResult(scs *state.ContractState, key []byte, n int) (*types.VoteList, error) {
   134  	data, err := scs.GetData(append(sortKey, key...))
   135  	if err != nil {
   136  		return nil, err
   137  	}
   138  	var ex bool
   139  	if bytes.Equal(key, defaultVoteKey) {
   140  		ex = false
   141  	} else {
   142  		ex = true
   143  	}
   144  	voteList := deserializeVoteList(data, ex)
   145  	if n < len(voteList.Votes) {
   146  		voteList.Votes = voteList.Votes[:n]
   147  	}
   148  	return voteList, nil
   149  }
   150  
   151  func GetVoteResultEx(ar AccountStateReader, key []byte, n int) (*types.VoteList, error) {
   152  	scs, err := ar.GetSystemAccountState()
   153  	if err != nil {
   154  		return nil, err
   155  	}
   156  	return getVoteResult(scs, key, n)
   157  }