Merge pull request #6050 from onflow/fxamacker/port-pr-5942

Optimize atree migration by porting and modifying PR 5942 to feature/atree-inlining-cadence-v0.42

Author: fxamacker

cmd/util/cmd/execution-state-extract/execution_state_extract.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 	"go.uber.org/atomic"
 	"golang.org/x/sync/errgroup"
 
@@ -547,10 +548,21 @@ func newByAccountRegistersFromPayloadAccountGrouping(
 		for accountRegisters := range results {
 			oldAccountRegisters := registersByAccount.SetAccountRegisters(accountRegisters)
 			if oldAccountRegisters != nil {
-				return fmt.Errorf(
-					"duplicate account registers for account %s",
+				// Account grouping should never create multiple groups for an account.
+				// In case it does anyway, merge the groups together,
+				// by merging the existing registers into the new ones.
+
+				log.Warn().Msgf(
+					"account registers already exist for account %x. merging %d existing registers into %d new",
 					accountRegisters.Owner(),
+					oldAccountRegisters.Count(),
+					accountRegisters.Count(),
 				)
+
+				err := accountRegisters.Merge(oldAccountRegisters)
+				if err != nil {
+					return fmt.Errorf("failed to merge account registers: %w", err)
+				}
 			}
 		}
 

cmd/util/ledger/util/payload_grouping.go

@@ -13,13 +13,6 @@ import (
 	"github.com/onflow/flow-go/model/flow"
 )
 
-// encodedKeyAddressPrefixLength is the length of the address prefix in the encoded key
-// 2 for uint16 of number of key parts
-// 4 for uint32 of the length of the first key part
-// 2 for uint16 of the key part type
-// 8 for the address which is the actual length of the first key part
-const encodedKeyAddressPrefixLength = 2 + 4 + 2 + flow.AddressLength
-
 // minSizeForSplitSortingIntoGoroutines below this size, no need to split
 // the sorting into goroutines
 const minSizeForSplitSortingIntoGoroutines = 100_000
@@ -135,11 +128,17 @@ func (s sortablePayloads) Less(i, j int) bool {
 }
 
 func (s sortablePayloads) Compare(i, j int) int {
-	// sort descending to force one of the big accounts to be more at the beginning
-	return bytes.Compare(
-		s[j].EncodedKey()[:encodedKeyAddressPrefixLength],
-		s[i].EncodedKey()[:encodedKeyAddressPrefixLength],
-	)
+	a, err := s[i].Address()
+	if err != nil {
+		panic(err)
+	}
+
+	b, err := s[j].Address()
+	if err != nil {
+		panic(err)
+	}
+
+	return bytes.Compare(a[:], b[:])
 }
 
 func (s sortablePayloads) Swap(i, j int) {

cmd/util/ledger/util/registers/registers.go

@@ -263,6 +263,17 @@ func (a *AccountRegisters) insertPayloads(payloads []*ledger.Payload) {
 	}
 }
 
+func (a *AccountRegisters) Merge(other *AccountRegisters) error {
+	for key, value := range other.registers {
+		_, ok := a.registers[key]
+		if ok {
+			return fmt.Errorf("key already exists: %s", key)
+		}
+		a.registers[key] = value
+	}
+	return nil
+}
+
 func NewAccountRegistersFromPayloads(owner string, payloads []*ledger.Payload) (*AccountRegisters, error) {
 	accountRegisters := NewAccountRegisters(owner)