Merge pull request #5948 from onflow/fxamacker/add-payload-address-func

Add Payload.Address() to allow migrations, etc. to reduce overhead

cmd/util/ledger/util/payload_grouping.go

@@ -10,7 +10,6 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/ledger"
-	"github.com/onflow/flow-go/ledger/common/convert"
 	"github.com/onflow/flow-go/model/flow"
 )
 
@@ -56,7 +55,7 @@ func (g *PayloadAccountGrouping) Next() (*PayloadAccountGroup, error) {
 	}
 	g.current++
 
-	address, err := PayloadToAddress(g.payloads[accountStartIndex])
+	address, err := g.payloads[accountStartIndex].Address()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get address from payload: %w", err)
 	}
@@ -125,31 +124,6 @@ func GroupPayloadsByAccount(
 	}
 }
 
-// PayloadToAddress takes a payload and return:
-// - (address, nil) if the payload is for an account, the account address is returned
-// - (common.ZeroAddress, nil) if the payload is not for an account
-// - (common.ZeroAddress, err) if running into any exception
-// The zero address is used for global Payloads and is not an actual account
-func PayloadToAddress(p *ledger.Payload) (flow.Address, error) {
-	k, err := p.Key()
-	if err != nil {
-		return flow.EmptyAddress, fmt.Errorf("could not find key for payload: %w", err)
-	}
-
-	id, err := convert.LedgerKeyToRegisterID(k)
-	if err != nil {
-		return flow.EmptyAddress, fmt.Errorf("error converting key to register ID")
-	}
-
-	if len([]byte(id.Owner)) != flow.AddressLength {
-		return flow.EmptyAddress, nil
-	}
-
-	address := flow.BytesToAddress([]byte(id.Owner))
-
-	return address, nil
-}
-
 type sortablePayloads []*ledger.Payload
 
 func (s sortablePayloads) Len() int {

ledger/common/convert/convert.go

@@ -7,14 +7,6 @@ import (
 	"github.com/onflow/flow-go/model/flow"
 )
 
-const (
-	KeyPartOwner = uint16(0)
-	// Deprecated: KeyPartController was only used by the very first
-	// version of Cadence for access control, which was later retired
-	_          = uint16(1) // DO NOT REUSE
-	KeyPartKey = uint16(2)
-)
-
 // UnexpectedLedgerKeyFormat is returned when a ledger key is not in the expected format
 var UnexpectedLedgerKeyFormat = fmt.Errorf("unexpected ledger key format")
 
@@ -23,8 +15,8 @@ var UnexpectedLedgerKeyFormat = fmt.Errorf("unexpected ledger key format")
 func LedgerKeyToRegisterID(key ledger.Key) (flow.RegisterID, error) {
 	parts := key.KeyParts
 	if len(parts) != 2 ||
-		parts[0].Type != KeyPartOwner ||
-		parts[1].Type != KeyPartKey {
+		parts[0].Type != ledger.KeyPartOwner ||
+		parts[1].Type != ledger.KeyPartKey {
 		return flow.RegisterID{}, fmt.Errorf("ledger key %s: %w", key.String(), UnexpectedLedgerKeyFormat)
 	}
 
@@ -39,11 +31,11 @@ func RegisterIDToLedgerKey(registerID flow.RegisterID) ledger.Key {
 	return ledger.Key{
 		KeyParts: []ledger.KeyPart{
 			{
-				Type:  KeyPartOwner,
+				Type:  ledger.KeyPartOwner,
 				Value: []byte(registerID.Owner),
 			},
 			{
-				Type:  KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte(registerID.Key),
 			},
 		},

ledger/common/convert/convert_test.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/onflow/flow-go/cmd/util/ledger/util"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/convert"
 	"github.com/onflow/flow-go/model/flow"
@@ -18,11 +17,11 @@ func TestLedgerKeyToRegisterID(t *testing.T) {
 	key := ledger.Key{
 		KeyParts: []ledger.KeyPart{
 			{
-				Type:  convert.KeyPartOwner,
+				Type:  ledger.KeyPartOwner,
 				Value: []byte(expectedRegisterID.Owner),
 			},
 			{
-				Type:  convert.KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte("key"),
 			},
 		},
@@ -34,7 +33,7 @@ func TestLedgerKeyToRegisterID(t *testing.T) {
 
 	p := ledger.NewPayload(key, ledger.Value("value"))
 
-	address, err := util.PayloadToAddress(p)
+	address, err := p.Address()
 
 	require.NoError(t, err)
 	require.Equal(t, registerID.Owner, flow.AddressToRegisterOwner(address))
@@ -45,11 +44,11 @@ func TestLedgerKeyToRegisterID_Global(t *testing.T) {
 	key := ledger.Key{
 		KeyParts: []ledger.KeyPart{
 			{
-				Type:  convert.KeyPartOwner,
+				Type:  ledger.KeyPartOwner,
 				Value: []byte(""),
 			},
 			{
-				Type:  convert.KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte("uuid"),
 			},
 		},
@@ -62,7 +61,7 @@ func TestLedgerKeyToRegisterID_Global(t *testing.T) {
 
 	p := ledger.NewPayload(key, ledger.Value("value"))
 
-	address, err := util.PayloadToAddress(p)
+	address, err := p.Address()
 
 	require.NoError(t, err)
 	require.Equal(t, registerID.Owner, flow.AddressToRegisterOwner(address))
@@ -77,7 +76,7 @@ func TestLedgerKeyToRegisterID_Error(t *testing.T) {
 				Value: []byte("owner"),
 			},
 			{
-				Type:  convert.KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte("key"),
 			},
 		},
@@ -93,13 +92,13 @@ func TestRegisterIDToLedgerKey(t *testing.T) {
 	expectedKey := ledger.Key{
 		KeyParts: []ledger.KeyPart{
 			{
-				Type: convert.KeyPartOwner,
+				Type: ledger.KeyPartOwner,
 				// Note: the owner field is extended to address length during NewRegisterID
 				// so we have to do the same here
 				Value: []byte(registerID.Owner),
 			},
 			{
-				Type:  convert.KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte("key"),
 			},
 		},
@@ -114,11 +113,11 @@ func TestRegisterIDToLedgerKey_Global(t *testing.T) {
 	expectedKey := ledger.Key{
 		KeyParts: []ledger.KeyPart{
 			{
-				Type:  convert.KeyPartOwner,
+				Type:  ledger.KeyPartOwner,
 				Value: []byte(""),
 			},
 			{
-				Type:  convert.KeyPartKey,
+				Type:  ledger.KeyPartKey,
 				Value: []byte("uuid"),
 			},
 		},
@@ -135,8 +134,8 @@ func TestPayloadToRegister(t *testing.T) {
 		p := ledger.NewPayload(
 			ledger.NewKey(
 				[]ledger.KeyPart{
-					ledger.NewKeyPart(convert.KeyPartOwner, []byte(expected.Owner)),
-					ledger.NewKeyPart(convert.KeyPartKey, []byte(expected.Key)),
+					ledger.NewKeyPart(ledger.KeyPartOwner, []byte(expected.Owner)),
+					ledger.NewKeyPart(ledger.KeyPartKey, []byte(expected.Key)),
 				},
 			),
 			value,
@@ -152,8 +151,8 @@ func TestPayloadToRegister(t *testing.T) {
 		p := ledger.NewPayload(
 			ledger.NewKey(
 				[]ledger.KeyPart{
-					ledger.NewKeyPart(convert.KeyPartOwner, []byte("")),
-					ledger.NewKeyPart(convert.KeyPartKey, []byte("uuid")),
+					ledger.NewKeyPart(ledger.KeyPartOwner, []byte("")),
+					ledger.NewKeyPart(ledger.KeyPartKey, []byte("uuid")),
 				},
 			),
 			value,

ledger/trie.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/onflow/flow-go/ledger/common/bitutils"
 	"github.com/onflow/flow-go/ledger/common/hash"
+	"github.com/onflow/flow-go/model/flow"
 )
 
 // Path captures storage path of a payload;
@@ -237,6 +238,14 @@ func (k encKey) DeepCopy() encKey {
 	return newK
 }
 
+const (
+	KeyPartOwner = uint16(0)
+	// Deprecated: KeyPartController was only used by the very first
+	// version of Cadence for access control, which was later retired
+	_          = uint16(1) // DO NOT REUSE
+	KeyPartKey = uint16(2)
+)
+
 // Payload is the smallest immutable storable unit in ledger
 type Payload struct {
 	// encKey is key encoded using PayloadVersion.
@@ -330,6 +339,28 @@ func (p *Payload) EncodedKey() []byte {
 	return p.encKey
 }
 
+// Address returns:
+// - (address, nil) if the payload is for an account, the account address is returned
+// - (flow.EmptyAddress, nil) if the payload is not for an account (global register)
+// - (flow.EmptyAddress, err) if running into any exception
+// The zero address is used for global Payloads and is not an actual account
+func (p *Payload) Address() (flow.Address, error) {
+	if p == nil {
+		return flow.EmptyAddress, fmt.Errorf("failed to get payload address: payload is nil")
+	}
+	if len(p.encKey) == 0 {
+		return flow.EmptyAddress, fmt.Errorf("failed to get payload address: encoded key is empty")
+	}
+	b, found, err := decodeKeyPartValueByType(p.encKey, KeyPartOwner, true, PayloadVersion)
+	if err != nil {
+		return flow.EmptyAddress, err
+	}
+	if !found {
+		return flow.EmptyAddress, fmt.Errorf("failed to find address by type %d", KeyPartOwner)
+	}
+	return flow.BytesToAddress(b), nil
+}
+
 // Value returns payload value.
 // CAUTION: do not modify returned value because it shares underlying data with payload value.
 func (p *Payload) Value() Value {

ledger/trie_encoder.go

@@ -148,28 +148,28 @@ func DecodeKeyPart(encodedKeyPart []byte) (*KeyPart, error) {
 	}
 
 	// decode the key part content (zerocopy)
-	key, err := decodeKeyPart(rest, true, version)
+	kpt, kpv, err := decodeKeyPart(rest, true, version)
 	if err != nil {
 		return nil, fmt.Errorf("error decoding key part: %w", err)
 	}
 
-	return key, nil
+	return &KeyPart{Type: kpt, Value: kpv}, nil
 }
 
 // decodeKeyPart decodes inp into KeyPart. If zeroCopy is true, KeyPart
 // references data in inp.  Otherwise, it is copied.
-func decodeKeyPart(inp []byte, zeroCopy bool, _ uint16) (*KeyPart, error) {
+func decodeKeyPart(inp []byte, zeroCopy bool, _ uint16) (uint16, []byte, error) {
 	// read key part type and the rest is the key item part
 	kpt, kpv, err := utils.ReadUint16(inp)
 	if err != nil {
-		return nil, fmt.Errorf("error decoding key part (content): %w", err)
+		return 0, nil, fmt.Errorf("error decoding key part (content): %w", err)
 	}
 	if zeroCopy {
-		return &KeyPart{Type: kpt, Value: kpv}, nil
+		return kpt, kpv, nil
 	}
 	v := make([]byte, len(kpv))
 	copy(v, kpv)
-	return &KeyPart{Type: kpt, Value: v}, nil
+	return kpt, v, nil
 }
 
 // EncodeKey encodes a key into a byte slice
@@ -240,6 +240,63 @@ func DecodeKey(encodedKey []byte) (*Key, error) {
 	return key, nil
 }
 
+func decodeKeyPartValueByType(inp []byte, typ uint16, zeroCopy bool, version uint16) ([]byte, bool, error) {
+	// Read number of key parts
+	numOfParts, rest, err := utils.ReadUint16(inp)
+	if err != nil {
+		return nil, false, fmt.Errorf("error decoding number of key parts: %w", err)
+	}
+
+	for i := 0; i < int(numOfParts); i++ {
+		var kpt uint16
+		var kpv []byte
+
+		kpt, kpv, rest, err = decodeKeyPartWithEncodedSizeInfo(rest, zeroCopy, version)
+		if err != nil {
+			return nil, false, err
+		}
+		if kpt == typ {
+			return kpv, true, nil
+		}
+	}
+
+	return nil, false, nil
+}
+
+func decodeKeyPartWithEncodedSizeInfo(
+	inp []byte,
+	zeroCopy bool,
+	version uint16,
+) (
+	// kp KeyPart,
+	kpt uint16,
+	kpv []byte,
+	rest []byte,
+	err error,
+) {
+
+	// Read encoded key part size
+	kpEncSize, rest, err := utils.ReadUint32(inp)
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("error decoding key part: %w", err)
+	}
+
+	// Read encoded key part
+	var kpEnc []byte
+	kpEnc, rest, err = utils.ReadSlice(rest, int(kpEncSize))
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("error decoding key part: %w", err)
+	}
+
+	// Decode encoded key part
+	kpType, kpValue, err := decodeKeyPart(kpEnc, zeroCopy, version)
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("error decoding key part: %w", err)
+	}
+
+	return kpType, kpValue, rest, nil
+}
+
 // decodeKey decodes inp into Key. If zeroCopy is true, returned key
 // references data in inp.  Otherwise, it is copied.
 func decodeKey(inp []byte, zeroCopy bool, version uint16) (*Key, error) {
@@ -257,27 +314,15 @@ func decodeKey(inp []byte, zeroCopy bool, version uint16) (*Key, error) {
 	key.KeyParts = make([]KeyPart, numOfParts)
 
 	for i := 0; i < int(numOfParts); i++ {
-		var kpEncSize uint32
-		var kpEnc []byte
-		// read encoded key part size
-		kpEncSize, rest, err = utils.ReadUint32(rest)
-		if err != nil {
-			return nil, fmt.Errorf("error decoding key (content): %w", err)
-		}
-
-		// read encoded key part
-		kpEnc, rest, err = utils.ReadSlice(rest, int(kpEncSize))
-		if err != nil {
-			return nil, fmt.Errorf("error decoding key (content): %w", err)
-		}
+		var kpt uint16
+		var kpv []byte
 
-		// decode encoded key part
-		kp, err := decodeKeyPart(kpEnc, zeroCopy, version)
+		kpt, kpv, rest, err = decodeKeyPartWithEncodedSizeInfo(rest, zeroCopy, version)
 		if err != nil {
-			return nil, fmt.Errorf("error decoding key (content): %w", err)
+			return nil, err
 		}
 
-		key.KeyParts[i] = *kp
+		key.KeyParts[i] = KeyPart{kpt, kpv}
 	}
 	return key, nil
 }