diff --git a/qkc/types/token.go b/qkc/types/token.go new file mode 100644 index 000000000000..44642b2946d4 --- /dev/null +++ b/qkc/types/token.go @@ -0,0 +1,343 @@ +// Copyright 2026-2027, QuarkChain. + +// Token balance encoding follows pyquarkchain-compatible QKC wire bytes. + +package types + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "math/big" + "sort" + "strconv" + "strings" + + "github.com/ethereum/go-ethereum/common" + "github.com/ethereum/go-ethereum/qkc/serialize" + "github.com/ethereum/go-ethereum/rlp" + "github.com/holiman/uint256" +) + +// TokenTrieThreshold is the maximum number of non-zero token balances supported +// by this list-only implementation. +const TokenTrieThreshold = 16 + +const ( + tokenBalanceListPrefix = byte(0) + tokenBalanceTriePrefix = byte(1) +) + +type TokenBalancePair struct { + TokenID uint64 + Balance *uint256.Int +} + +// TokenBalances holds QKC account multi-token balances. +// +// Serialization uses list format (0x00 prefix) for up to TokenTrieThreshold +// non-zero balances. Trie format (>16 non-zero balances) is not supported; +// SerializeToBytes returns an error instead. +type TokenBalances struct { + balances map[uint64]*uint256.Int +} + +type TokenBalancesAlias TokenBalances + +func (t *TokenBalances) MarshalJSON() ([]byte, error) { + keys := make([]uint64, 0, len(t.balances)) + for key := range t.balances { + keys = append(keys, key) + } + sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] }) + balances := make([]string, 0, len(keys)) + for _, key := range keys { + val := t.balances[key] + balance := "0" + if val != nil { + balance = val.Dec() + } + balances = append(balances, strconv.FormatUint(key, 10)+":"+balance) + } + jsoncfg := struct { + TokenBalancesAlias + Balances string `json:"balances"` + }{TokenBalancesAlias: TokenBalancesAlias(*t), Balances: strings.Join(balances, ",")} + return json.Marshal(jsoncfg) +} + +func (t *TokenBalances) UnmarshalJSON(input []byte) error { + var jsoncfg struct { + TokenBalancesAlias + Balances string `json:"balances"` + } + if err := json.Unmarshal(input, &jsoncfg); err != nil { + return err + } + *t = TokenBalances(jsoncfg.TokenBalancesAlias) + t.balances = make(map[uint64]*uint256.Int) + if jsoncfg.Balances == "" { + return nil + } + balList := strings.Split(jsoncfg.Balances, ",") + for _, val := range balList { + parts := strings.SplitN(val, ":", 2) + if len(parts) != 2 { + return fmt.Errorf("invalid token balance entry: %q", val) + } + key, err := strconv.ParseUint(parts[0], 10, 64) + if err != nil { + return fmt.Errorf("invalid token id %q: %w", parts[0], err) + } + balance, err := uint256.FromDecimal(parts[1]) + if err != nil { + return fmt.Errorf("invalid token balance %q: %w", parts[1], err) + } + t.balances[key] = balance + } + return nil +} + +func NewEmptyTokenBalances() *TokenBalances { + return &TokenBalances{ + balances: make(map[uint64]*uint256.Int), + } +} + +func NewTokenBalancesWithMap(data map[uint64]*uint256.Int) *TokenBalances { + t := &TokenBalances{ + balances: make(map[uint64]*uint256.Int, len(data)), + } + for tokenID, balance := range data { + if balance == nil { + continue + } + t.balances[tokenID] = new(uint256.Int).Set(balance) + } + return t +} + +func NewTokenBalances(data []byte) (*TokenBalances, error) { + tokenBalances := NewEmptyTokenBalances() + if len(data) == 0 { + return tokenBalances, nil + } + + switch data[0] { + case tokenBalanceListPrefix: + balanceList := make([]*TokenBalancePair, 0) + if err := rlp.DecodeBytes(data[1:], &balanceList); err != nil { + return nil, err + } + for _, v := range balanceList { + if v == nil || v.Balance == nil { + continue + } + tokenBalances.balances[v.TokenID] = new(uint256.Int).Set(v.Balance) + } + case tokenBalanceTriePrefix: + return nil, errors.New("token balance trie encoding is unsupported") + default: + return nil, fmt.Errorf("Unknown enum byte in token_balances:%v", data[0]) + + } + return tokenBalances, nil +} + +func (t *TokenBalances) Add(other map[uint64]*big.Int) error { + updates := make(map[uint64]*uint256.Int, len(other)) + for tokenID, delta := range other { + if delta == nil { + continue + } + balance := new(big.Int) + if current := t.balances[tokenID]; current != nil { + balance.Set(current.ToBig()) + } + balance.Add(balance, delta) + if balance.Sign() < 0 { + return fmt.Errorf("token balance underflow for token %d", tokenID) + } + updated, overflow := uint256.FromBig(balance) + if overflow { + return fmt.Errorf("token balance overflow for token %d", tokenID) + } + updates[tokenID] = updated + } + for tokenID, balance := range updates { + t.balances[tokenID] = balance + } + return nil +} + +func (t *TokenBalances) SetValue(amount *uint256.Int, tokenID uint64) { + if amount == nil { + amount = new(uint256.Int) + } + t.balances[tokenID] = new(uint256.Int).Set(amount) +} + +func (t *TokenBalances) GetTokenBalance(tokenID uint64) *uint256.Int { + data, ok := t.balances[tokenID] + if ok { + return new(uint256.Int).Set(data) + } + + return new(uint256.Int) +} + +func (t *TokenBalances) GetBalanceMap() map[uint64]*uint256.Int { + data := t.Copy() + return data.balances +} + +func (t *TokenBalances) Len() int { + return len(t.balances) +} + +func (t *TokenBalances) nonZeroEntriesInBalancesCache() int { + sum := 0 + for _, v := range t.balances { + if v != nil && !v.IsZero() { + sum++ + } + } + return sum +} + +func (t *TokenBalances) IsBlank() bool { + return t.nonZeroEntriesInBalancesCache() == 0 +} + +func (t *TokenBalances) Copy() *TokenBalances { + balancesCopy := make(map[uint64]*uint256.Int) + for k, v := range t.balances { + if v == nil { + continue + } + balancesCopy[k] = new(uint256.Int).Set(v) + } + return &TokenBalances{ + balances: balancesCopy, + } +} + +func (t *TokenBalances) SerializeToBytes() ([]byte, error) { + nonZeroEntries := t.nonZeroEntriesInBalancesCache() + if nonZeroEntries > TokenTrieThreshold { + return nil, fmt.Errorf("token balances exceed supported list size: %d > %d", nonZeroEntries, TokenTrieThreshold) + } + if t.Len() == 0 { + return nil, nil + } + list := make([]*TokenBalancePair, 0) + for k, v := range t.balances { + if v == nil || v.IsZero() { + continue + } + list = append(list, &TokenBalancePair{ + TokenID: k, + Balance: v, + }) + } + sort.Slice(list, func(i, j int) bool { return list[i].TokenID < (list[j].TokenID) }) + rlpData := new(bytes.Buffer) + rlpData.WriteByte(tokenBalanceListPrefix) + err := rlp.Encode(rlpData, list) + if err != nil { + return nil, err + } + return rlpData.Bytes(), nil +} + +func (t *TokenBalances) EncodeRLP(w io.Writer) error { + dataSer, err := t.SerializeToBytes() + if err != nil { + return err + } + dataRlp, err := rlp.EncodeToBytes(dataSer) + if err != nil { + return err + } + _, err = w.Write(dataRlp) + return err +} + +func (t *TokenBalances) DecodeRLP(s *rlp.Stream) error { + dataRawBytes, err := s.Raw() + if err != nil { + return err + } + dataRlp := new([]byte) + err = rlp.DecodeBytes(dataRawBytes, dataRlp) + if err != nil { + return err + } + b, err := NewTokenBalances(*dataRlp) + if err != nil { + return err + } + (*t).balances = (*b).balances + return err +} + +func (t *TokenBalances) Serialize(w *[]byte) error { + // Follow pyquarkchain (TokenBalanceMap's PrependedSizeMapSerializer skip_func): + // omit zero-balance entries from both the count and the body. This keeps + // Serialize/Deserialize symmetric and preserves pyquarkchain-compatible bytes. + keys := make([]uint64, 0, t.Len()) + for k, v := range t.balances { + if v == nil || v.IsZero() { + continue + } + keys = append(keys, k) + } + sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] }) + if err := serialize.Serialize(w, uint32(len(keys))); err != nil { + return err + } + for _, key := range keys { + v := t.balances[key] + if err := serialize.Serialize(w, new(big.Int).SetUint64(key)); err != nil { + return err + } + if err := serialize.Serialize(w, v.ToBig()); err != nil { + return err + } + } + return nil +} + +func (t *TokenBalances) Deserialize(bb *serialize.ByteBuffer) error { + t.balances = make(map[uint64]*uint256.Int) + num, err := bb.GetUInt32() + if err != nil { + return err + } + for index := uint32(0); index < num; index++ { + k := new(big.Int) + if err := serialize.Deserialize(bb, k); err != nil { + return err + } + v := new(big.Int) + + if err := serialize.Deserialize(bb, v); err != nil { + return err + } + if !k.IsUint64() { + return fmt.Errorf("token id overflows uint64: %v", k) + } + if v.Cmp(common.Big0) == 0 { + delete(t.balances, k.Uint64()) + continue + } + balance, overflow := uint256.FromBig(v) + if overflow { + return fmt.Errorf("token balance overflows uint256: %v", v) + } + t.balances[k.Uint64()] = balance + } + return nil +} diff --git a/qkc/types/token_test.go b/qkc/types/token_test.go new file mode 100644 index 000000000000..e80dd8f8af96 --- /dev/null +++ b/qkc/types/token_test.go @@ -0,0 +1,223 @@ +// Copyright 2026-2027, QuarkChain. + +// Token balance tests exercise pyquarkchain-compatible QKC wire bytes. + +package types + +import ( + "encoding/hex" + "encoding/json" + "math/big" + "reflect" + "testing" + + "github.com/ethereum/go-ethereum/qkc/serialize" + "github.com/ethereum/go-ethereum/rlp" + "github.com/holiman/uint256" + "github.com/stretchr/testify/assert" +) + +func testU256(v uint64) *uint256.Int { + return new(uint256.Int).SetUint64(v) +} + +func testU256Decimal(t *testing.T, v string) *uint256.Int { + t.Helper() + data, err := uint256.FromDecimal(v) + if err != nil { + t.Fatal(err) + } + return data +} + +func TestNewTokenBalanceMap(t *testing.T) { + m0 := make(map[uint64]*uint256.Int) + m0[3234] = testU256(1000) + m0[0] = testU256(0) + m0[3567] = testU256(0) + tb := NewTokenBalancesWithMap(m0) + t.Logf("token balance map:%v", tb.balances) +} + +func TestTokenBalancesJSONEmptyRoundTrip(t *testing.T) { + encoded, err := json.Marshal(NewEmptyTokenBalances()) + assert.NoError(t, err) + + var decoded TokenBalances + assert.NoError(t, json.Unmarshal(encoded, &decoded)) + assert.Equal(t, 0, decoded.Len()) + assert.True(t, decoded.IsBlank()) +} + +func TestTokenBalancesJSONLargeValueRoundTrip(t *testing.T) { + const ( + maxTokenID = ^uint64(0) + largeValue = "1208925819614629174706176" + ) + original := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + maxTokenID: testU256Decimal(t, largeValue), + }) + + encoded, err := json.Marshal(original) + assert.NoError(t, err) + assert.Contains(t, string(encoded), "18446744073709551615:"+largeValue) + + var decoded TokenBalances + assert.NoError(t, json.Unmarshal(encoded, &decoded)) + assert.Equal(t, original.GetTokenBalance(maxTokenID), decoded.GetTokenBalance(maxTokenID)) +} + +func TestTokenBalances_Add(t *testing.T) { + check := func(f string, got, want interface{}) { + if !reflect.DeepEqual(got, want) { + t.Errorf("%s mismatch: got %v, want %v", f, got, want) + } + } + m0 := make(map[uint64]*uint256.Int) + m0[3567] = testU256(0) + tb := NewTokenBalancesWithMap(m0) + assert.NoError(t, tb.Add(map[uint64]*big.Int{3234: big.NewInt(10)})) + m3 := make(map[uint64]*uint256.Int) + m3[3567] = testU256(0) + m3[3234] = testU256(10) + check("", tb.balances, m3) +} + +func TestTokenBalancesAddSupportsNegativeDelta(t *testing.T) { + tb := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + 1: testU256(10), + }) + + assert.NoError(t, tb.Add(map[uint64]*big.Int{1: big.NewInt(-4)})) + assert.Equal(t, testU256(6), tb.GetTokenBalance(1)) +} + +func TestTokenBalancesAddRejectsUnderflow(t *testing.T) { + tb := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + 1: testU256(10), + }) + + err := tb.Add(map[uint64]*big.Int{1: big.NewInt(-11)}) + assert.ErrorContains(t, err, "underflow") + assert.Equal(t, testU256(10), tb.GetTokenBalance(1)) +} + +func TestTokenBalancesAddRejectsOverflow(t *testing.T) { + max := new(uint256.Int).SetAllOne() + tb := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + 1: max, + }) + + err := tb.Add(map[uint64]*big.Int{1: big.NewInt(1)}) + assert.ErrorContains(t, err, "overflow") + assert.Equal(t, max, tb.GetTokenBalance(1)) +} + +func TestTokenBalancesUsesListEncodingUpToThreshold(t *testing.T) { + mapping := make(map[uint64]*uint256.Int, 0) + for index := 0; index < TokenTrieThreshold; index++ { + mapping[uint64(index+1)] = testU256(uint64(index*1000 + 42)) + } + + b0 := NewTokenBalancesWithMap(mapping) + data, err := b0.SerializeToBytes() + assert.NoError(t, err) + assert.Equal(t, tokenBalanceListPrefix, data[0]) + + b1, err := NewTokenBalances(data) + assert.NoError(t, err) + assert.Equal(t, len(mapping), b1.Len()) + for tokenID, balance := range mapping { + assert.Equal(t, balance, b1.GetTokenBalance(tokenID)) + } +} + +func TestTokenBalancesRejectsMoreThanThreshold(t *testing.T) { + mapping := make(map[uint64]*uint256.Int, 0) + for index := 0; index < TokenTrieThreshold+1; index++ { + mapping[uint64(index+1)] = testU256(uint64(index*1000 + 42)) + } + + b0 := NewTokenBalancesWithMap(mapping) + _, err := b0.SerializeToBytes() + assert.ErrorContains(t, err, "exceed") +} + +func TestTokenBalancesRejectsTrieEncoding(t *testing.T) { + data := make([]byte, 33) + data[0] = tokenBalanceTriePrefix + + _, err := NewTokenBalances(data) + assert.ErrorContains(t, err, "trie") +} + +func TestTokenBalancesRLPRoundTripWithoutDB(t *testing.T) { + mapping := make(map[uint64]*uint256.Int, 0) + mapping[1] = testU256(42) + mapping[3] = testU256(99) + + b := NewTokenBalancesWithMap(mapping) + encoded, err := rlp.EncodeToBytes(b) + assert.NoError(t, err) + + var decoded TokenBalances + assert.NoError(t, rlp.DecodeBytes(encoded, &decoded)) + assert.Equal(t, testU256(42), decoded.GetTokenBalance(1)) + assert.Equal(t, testU256(99), decoded.GetTokenBalance(3)) + assert.Equal(t, new(uint256.Int), decoded.GetTokenBalance(2)) +} + +func TestTokenBalancesZeroOnlyStateEncodingPythonGolden(t *testing.T) { + empty := NewEmptyTokenBalances() + emptyInner, err := empty.SerializeToBytes() + assert.NoError(t, err) + assert.Empty(t, emptyInner) + emptyOuter, err := rlp.EncodeToBytes(empty) + assert.NoError(t, err) + assert.Equal(t, []byte{0x80}, emptyOuter) + + zeroOnly := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + 1: testU256(0), + }) + zeroOnlyInner, err := zeroOnly.SerializeToBytes() + assert.NoError(t, err) + assert.Equal(t, []byte{0x00, 0xc0}, zeroOnlyInner) + zeroOnlyOuter, err := rlp.EncodeToBytes(zeroOnly) + assert.NoError(t, err) + assert.Equal(t, []byte{0x82, 0x00, 0xc0}, zeroOnlyOuter) +} + +func TestTokenBalancesQKCSerializePythonGolden(t *testing.T) { + golden, err := hex.DecodeString("00000002020ca0016d020def030f423f") + assert.NoError(t, err) + + tb := NewTokenBalancesWithMap(map[uint64]*uint256.Int{ + 3232: testU256(109), + 0: testU256(0), + 3567: testU256(999999), + }) + encoded, err := serialize.SerializeToBytes(tb) + assert.NoError(t, err) + assert.Equal(t, golden, encoded) + + var decoded TokenBalances + assert.NoError(t, serialize.DeserializeFromBytes(golden, &decoded)) + assert.Equal(t, 2, decoded.Len()) + assert.Equal(t, testU256(109), decoded.GetTokenBalance(3232)) + assert.Equal(t, testU256(999999), decoded.GetTokenBalance(3567)) + assert.Equal(t, new(uint256.Int), decoded.GetTokenBalance(0)) +} + +func TestTokenBalancesQKCDeserializeDuplicatePythonGolden(t *testing.T) { + golden, err := hex.DecodeString("0000000201010101010100") + assert.NoError(t, err) + + var decoded TokenBalances + assert.NoError(t, serialize.DeserializeFromBytes(golden, &decoded)) + assert.Equal(t, 0, decoded.Len()) + assert.Equal(t, new(uint256.Int), decoded.GetTokenBalance(1)) + + encoded, err := serialize.SerializeToBytes(&decoded) + assert.NoError(t, err) + assert.Equal(t, []byte{0x00, 0x00, 0x00, 0x00}, encoded) +}