diff --git a/migration/trie/codec.go b/migration/trie/codec.go new file mode 100644 index 0000000000..33b12253c3 --- /dev/null +++ b/migration/trie/codec.go @@ -0,0 +1,123 @@ +package trie + +import ( + "github.com/NethermindEth/juno/core/crypto" + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/core/trie" + "github.com/NethermindEth/juno/core/trie2/trieutils" + "github.com/NethermindEth/juno/db" +) + +const ( + binaryNodeTag byte = 0x01 + edgeNodeTag byte = 0x02 + + valueNodeBlobSize = felt.Bytes + binaryNodeBlobSize = 1 + 2*felt.Bytes + edgeNodeMinSize = 1 + felt.Bytes + 1 + edgeNodeMaxSize = 1 + felt.Bytes + trieutils.MaxBitArraySize + maxNodeKeySize = 1 + felt.Bytes + 1 + trieutils.MaxBitArraySize + + nonLeafByte byte = 1 + leafByte byte = 2 +) + +// +// --- Encoding related helpers --- +// + +// encodeBinaryNode writes a binary-node blob into dst. +// dst must have at least binaryNodeBlobSize bytes of capacity. +func encodeBinaryNode(dst []byte, leftEdgeHash, rightEdgeHash *felt.Felt) int { + dst[0] = binaryNodeTag + lb := leftEdgeHash.Bytes() + rb := rightEdgeHash.Bytes() + copy(dst[1:], lb[:]) + copy(dst[1+felt.Bytes:], rb[:]) + return binaryNodeBlobSize +} + +// encodeEdgeNode writes an edge-node blob into dst. +// dst must have at least edgeNodeMaxSize bytes of capacity. +func encodeEdgeNode(dst []byte, childHash *felt.Felt, pathSeg *trieutils.Path) int { + encoded := pathSeg.EncodedBytes() + dst[0] = edgeNodeTag + h := childHash.Bytes() + copy(dst[1:], h[:]) + copy(dst[1+felt.Bytes:], encoded) + return 1 + felt.Bytes + len(encoded) +} + +func encodeNodeKey( + dst []byte, + bucket db.Bucket, + owner *felt.Address, + path *trieutils.Path, + isLeaf bool, +) int { + n := 0 + dst[n] = byte(bucket) + n++ + + if !felt.IsZero(owner) { + ownerBytes := owner.Bytes() + copy(dst[n:], ownerBytes[:]) + n += 32 + } + + if isLeaf { + dst[n] = leafByte + } else { + dst[n] = nonLeafByte + } + n++ + + pathBytes := path.EncodedBytes() + copy(dst[n:], pathBytes) + n += len(pathBytes) + + return n +} + +// +// --- Path related helpers --- +// + +func parseDeprecatedPath(val []byte) (trie.BitArray, error) { + if len(val) == 0 { + return trie.BitArray{}, nil + } + var ba trie.BitArray + if err := ba.UnmarshalBinary(val); err != nil { + return trie.BitArray{}, err + } + return ba, nil +} + +func toNewPath(old *trie.BitArray) trieutils.Path { + b := old.Bytes() + var p trieutils.Path + p.SetBytes(old.Len(), b[:]) + return p +} + +func compressedSegment(childFullPath *trie.BitArray, parentLen uint8) trieutils.Path { + var seg trie.BitArray + seg.LSBs(childFullPath, parentLen+1) + return toNewPath(&seg) +} + +// +// --- Hash related helpers --- +// + +func computeEdgeHash(childHash *felt.Felt, path *trieutils.Path, hashFn crypto.HashFn) felt.Felt { + if path.Len() == 0 { + return *childHash + } + pathFelt := path.Felt() + h := hashFn(childHash, &pathFelt) + lenFelt := felt.FromUint64[felt.Felt](uint64(path.Len())) + h.Add(&h, &lenFelt) + return h +} diff --git a/migration/trie/committer.go b/migration/trie/committer.go new file mode 100644 index 0000000000..4086e3ff7e --- /dev/null +++ b/migration/trie/committer.go @@ -0,0 +1,39 @@ +package trie + +import ( + "fmt" + + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/migration/semaphore" + "github.com/NethermindEth/juno/utils/log" +) + +type committer struct { + batchSem semaphore.ResourceSemaphore[db.Batch] + counter counter +} + +func newCommitter( + logger log.StructuredLogger, + batchSem semaphore.ResourceSemaphore[db.Batch], + allTries, allNodes uint64, +) *committer { + return &committer{ + batchSem: batchSem, + counter: newCounter(logger, timeLogRate, allTries, allNodes), + } +} + +func (c *committer) Run(_ int, t task, _ chan<- struct{}) error { + byteSize := uint64(t.batch.Size()) + if err := t.batch.Write(); err != nil { + return fmt.Errorf("trie migration: batch write failed: %w", err) + } + c.counter.log(byteSize, t.tries, t.nodes) + c.batchSem.Put() + return nil +} + +func (c *committer) Done(int, chan<- struct{}) error { + return nil +} diff --git a/migration/trie/counter.go b/migration/trie/counter.go new file mode 100644 index 0000000000..88cdd47d19 --- /dev/null +++ b/migration/trie/counter.go @@ -0,0 +1,75 @@ +package trie + +import ( + "fmt" + "time" + + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/utils/log" + "go.uber.org/zap" +) + +type counter struct { + logger log.StructuredLogger + timeLogRate time.Duration + migrationStart time.Time + allTries uint64 + allNodes uint64 + totalTries uint64 + totalNodes uint64 + start time.Time + size uint64 + tries uint64 + nodes uint64 +} + +func newCounter( + logger log.StructuredLogger, + timeLogRate time.Duration, + allTries, allNodes uint64, +) counter { + now := time.Now() + return counter{ + logger: logger, + timeLogRate: timeLogRate, + migrationStart: now, + start: now, + allTries: allTries, + allNodes: allNodes, + } +} + +func (c *counter) log(byteSize uint64, tries, nodes int) { + c.size += byteSize + c.tries += uint64(tries) + c.nodes += uint64(nodes) + c.totalTries += uint64(tries) + c.totalNodes += uint64(nodes) + + now := time.Now() + elapsed := now.Sub(c.start).Seconds() + if elapsed > c.timeLogRate.Seconds() { + mbs := float64(c.size) / float64(db.Megabyte) + c.logger.Info( + "write speed", + zap.Float64("MB", mbs), + zap.Float64("MB/s", mbs/elapsed), + zap.Float64("nodes/s", float64(c.nodes)/elapsed), + zap.Float64("tries/s", float64(c.tries)/elapsed), + zap.String("tries_processed", fmtPercent(c.totalTries, c.allTries)), + zap.String("nodes_processed", fmtPercent(c.totalNodes, c.allNodes)), + zap.Float64("totalTime", now.Sub(c.migrationStart).Seconds()), + ) + c.start = now + c.size = 0 + c.tries = 0 + c.nodes = 0 + } +} + +func fmtPercent(done, total uint64) string { + if total == 0 { + return "100.0%" + } + return fmt.Sprintf("%.1f%%", 100.0*float64(done)/float64(total)) +} diff --git a/migration/trie/hashpool.go b/migration/trie/hashpool.go new file mode 100644 index 0000000000..ca08b338c1 --- /dev/null +++ b/migration/trie/hashpool.go @@ -0,0 +1,61 @@ +package trie + +import ( + "sync" + + "github.com/NethermindEth/juno/core/crypto" + "github.com/NethermindEth/juno/core/felt" +) + +type hashWork struct { + hashFn crypto.HashFn + jobs []edgeHashJob + results []felt.Felt + wg *sync.WaitGroup +} + +type hashWorkerPool struct { + work chan hashWork + n int +} + +func newHashWorkerPool() *hashWorkerPool { + p := &hashWorkerPool{ + work: make(chan hashWork, IngestorCount*2), + n: IngestorCount, + } + for range IngestorCount { + go func() { + for w := range p.work { + for i := range w.jobs { + w.results[2*i] = computeEdgeHash(&w.jobs[i].leftChildHash, &w.jobs[i].leftSeg, w.hashFn) + w.results[2*i+1] = computeEdgeHash(&w.jobs[i].rightChildHash, &w.jobs[i].rightSeg, w.hashFn) + } + w.wg.Done() + } + }() + } + return p +} + +func (p *hashWorkerPool) submit( + hashFn crypto.HashFn, + jobs []edgeHashJob, + results []felt.Felt, +) <-chan struct{} { + done := make(chan struct{}) + go func() { + var wg sync.WaitGroup + chunkSize := max(1, (len(jobs)+p.n-1)/p.n) + for i := 0; i < len(jobs); i += chunkSize { + end := min(i+chunkSize, len(jobs)) + wg.Add(1) + p.work <- hashWork{hashFn, jobs[i:end], results[2*i : 2*end], &wg} + } + wg.Wait() + close(done) + }() + return done +} + +func (p *hashWorkerPool) close() { close(p.work) } diff --git a/migration/trie/hashworker.go b/migration/trie/hashworker.go new file mode 100644 index 0000000000..0fafa08920 --- /dev/null +++ b/migration/trie/hashworker.go @@ -0,0 +1,166 @@ +package trie + +import ( + "github.com/NethermindEth/juno/core/crypto" + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/core/trie2/trieutils" + "github.com/NethermindEth/juno/db" +) + +type edgeHashJob struct { + leftChildHash, rightChildHash felt.Felt + leftSeg, rightSeg trieutils.Path + parentPath trieutils.Path +} + +type inFlightBatch struct { + jobs []edgeHashJob + results []felt.Felt + done <-chan struct{} +} + +type hashScheduler struct { + hashFn crypto.HashFn + parallel bool + bucket db.Bucket + owner felt.Address + pool *hashWorkerPool + + jobs []edgeHashJob + altJobs []edgeHashJob + results []felt.Felt + inFlightBuf inFlightBatch + hasInFlight bool +} + +func newHashScheduler( + hashFn crypto.HashFn, + parallel bool, + bucket db.Bucket, + owner felt.Address, + pool *hashWorkerPool, +) *hashScheduler { + s := &hashScheduler{ + hashFn: hashFn, + parallel: parallel, + bucket: bucket, + owner: owner, + pool: pool, + } + if parallel { + s.jobs = make([]edgeHashJob, 0, parallelHashBatchSize) + s.altJobs = make([]edgeHashJob, 0, parallelHashBatchSize) + s.results = make([]felt.Felt, 2*parallelHashBatchSize) + } + return s +} + +func (s *hashScheduler) schedule(job *edgeHashJob, batch db.Batch) error { + if !s.parallel { + leftEdge := computeEdgeHash(&job.leftChildHash, &job.leftSeg, s.hashFn) + rightEdge := computeEdgeHash(&job.rightChildHash, &job.rightSeg, s.hashFn) + return s.writeBinaryAndEdges(job, &leftEdge, &rightEdge, batch) + } + s.jobs = append(s.jobs, *job) + if len(s.jobs) >= parallelHashBatchSize { + return s.fire(batch) + } + return nil +} + +func (s *hashScheduler) fire(batch db.Batch) error { + if err := s.drainInFlight(batch); err != nil { + return err + } + results := s.results[:2*len(s.jobs)] + s.inFlightBuf = inFlightBatch{ + jobs: s.jobs, + results: results, + done: s.pool.submit(s.hashFn, s.jobs, results), + } + s.hasInFlight = true + s.jobs, s.altJobs = s.altJobs[:0], s.jobs + return nil +} + +func (s *hashScheduler) drainInFlight(batch db.Batch) error { + if !s.hasInFlight { + return nil + } + <-s.inFlightBuf.done + for i := range s.inFlightBuf.jobs { + err := s.writeBinaryAndEdges( + &s.inFlightBuf.jobs[i], + &s.inFlightBuf.results[2*i], + &s.inFlightBuf.results[2*i+1], + batch, + ) + if err != nil { + return err + } + } + s.hasInFlight = false + return nil +} + +func (s *hashScheduler) sync(batch db.Batch) error { + if !s.parallel { + return nil + } + if err := s.drainInFlight(batch); err != nil { + return err + } + if len(s.jobs) > 0 { + results := s.results[:2*len(s.jobs)] + <-s.pool.submit(s.hashFn, s.jobs, results) + for i := range s.jobs { + if err := s.writeBinaryAndEdges( + &s.jobs[i], + &results[2*i], + &results[2*i+1], + batch, + ); err != nil { + return err + } + } + s.jobs = s.jobs[:0] + } + return nil +} + +func (s *hashScheduler) writeBinaryAndEdges( + job *edgeHashJob, + leftEdge, + rightEdge *felt.Felt, + batch db.Batch, +) error { + var buf [maxNodeKeySize + binaryNodeBlobSize]byte + keyLen := encodeNodeKey(buf[:], s.bucket, &s.owner, &job.parentPath, false) + blobLen := encodeBinaryNode(buf[keyLen:], leftEdge, rightEdge) + if err := batch.Put(buf[:keyLen], buf[keyLen:keyLen+blobLen]); err != nil { + return err + } + + if err := s.writeEdge(&job.parentPath, 0, &job.leftChildHash, &job.leftSeg, batch); err != nil { + return err + } + return s.writeEdge(&job.parentPath, 1, &job.rightChildHash, &job.rightSeg, batch) +} + +func (s *hashScheduler) writeEdge( + parentPath *trieutils.Path, + bit uint8, + childHash *felt.Felt, + seg *trieutils.Path, + batch db.Batch, +) error { + if seg.Len() == 0 { + return nil + } + var edgePath trieutils.Path + edgePath.AppendBit(parentPath, bit) + var buf [maxNodeKeySize + edgeNodeMaxSize]byte + keyLen := encodeNodeKey(buf[:], s.bucket, &s.owner, &edgePath, false) + blobLen := encodeEdgeNode(buf[keyLen:], childHash, seg) + return batch.Put(buf[:keyLen], buf[keyLen:keyLen+blobLen]) +} diff --git a/migration/trie/ingestor.go b/migration/trie/ingestor.go new file mode 100644 index 0000000000..76d85f56e0 --- /dev/null +++ b/migration/trie/ingestor.go @@ -0,0 +1,364 @@ +package trie + +import ( + "context" + "fmt" + + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/core/trie" + "github.com/NethermindEth/juno/core/trie2/trieutils" + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/migration/semaphore" +) + +type task struct { + batch db.Batch + tries int + nodes int +} + +type ingestor struct { + ctx context.Context + database db.KeyValueReader + batchSemaphore semaphore.ResourceSemaphore[db.Batch] + pool *hashWorkerPool + tasks [IngestorCount]task +} + +func newIngestor( + ctx context.Context, + database db.KeyValueReader, + batchSemaphore semaphore.ResourceSemaphore[db.Batch], + pool *hashWorkerPool, +) *ingestor { + in := &ingestor{ + ctx: ctx, + database: database, + batchSemaphore: batchSemaphore, + pool: pool, + } + for i := range IngestorCount { + in.tasks[i].batch = batchSemaphore.GetBlocking() + } + return in +} + +func (i *ingestor) Run(index int, desc TrieDesc, outputs chan<- task) error { + done, err := rootProcessed(i.database, desc.TrieBucket, &desc.Owner) + if err != nil { + return fmt.Errorf("rootProcessed(%v, %x): %w", desc.TrieBucket, desc.Owner, err) + } + + t := &i.tasks[index] + if done { + // Already migrated — credit the counts so progress display reaches 100% on resume. + t.tries++ + t.nodes += desc.NodeCount + return i.flush(t, outputs) + } + + if err := i.migrateTrie(t, desc, outputs); err != nil { + return err + } + + t.tries++ + return i.flush(t, outputs) +} + +func (i *ingestor) Done(index int, outputs chan<- task) error { + select { + case <-i.ctx.Done(): + return i.ctx.Err() + case outputs <- i.tasks[index]: + } + return nil +} + +func (i *ingestor) flush(t *task, outputs chan<- task) error { + if t.batch.Size() < targetBatchByteSize { + return nil + } + select { + case <-i.ctx.Done(): + return i.ctx.Err() + case outputs <- task{batch: t.batch, tries: t.tries, nodes: t.nodes}: + } + t.tries = 0 + t.nodes = 0 + t.batch = i.batchSemaphore.GetBlocking() + return nil +} + +// migrateTrie reads one deprecated trie and writes its equivalent into the +// new layout. Three things differ between the formats: how nodes are keyed +// on disk, how nodes are encoded, and how path compression is expressed. +// +// On-disk keying +// -------------- +// Both layouts share a common prefix; only the suffix differs: +// +// common (both) suffix +// ───────────── ───────────────────────────────────────── +// bucket [|| owner] → path-length-byte || path-bytes (deprecated) +// → nodeType-byte || path-length-byte || path-bytes +// (new) +// +// The owner is present only for storage tries. The new layout's extra +// nodeType byte splits leaves from internal nodes into two index slices +// within the same bucket — the new-state lookups use this to short-circuit +// between leaf reads and internal-node traversals. +// +// Node encoding +// ------------- +// Both layouts are byte streams. The deprecated format keeps each node +// self-contained — internal binary nodes embed the compressed paths to +// their children inline: +// +// leaf value +// binary value || left-child-path || right-child-path +// [|| left-hash || right-hash, optional cache, ignored here] +// +// "value" is the node's own Starknet trie hash, or the stored value when +// the node is a leaf. +// +// The new format gives each node an explicit type tag and moves path +// compression into separate edge nodes: +// +// value value +// binary 0x01 || left-edge-hash || right-edge-hash +// edge 0x02 || child-hash || encoded-path-segment +// +// Path compression +// ---------------- +// This is the key structural change. The deprecated format compresses +// paths inside the parent binary node (via its embedded child-path +// fields). The new format moves compression into dedicated edge nodes +// sitting between binary nodes and their children: +// +// old: binary ──────── child-path ────────► child +// new: binary ──► edge ──► child +// +// The deprecated root marker — a single entry at the bare bucket prefix +// recording the root's path — disappears in the new layout. Whatever the +// deprecated root embedded becomes either a direct binary/leaf at the +// empty path or, when the deprecated root path was itself non-empty, an +// edge node at the empty path that points "down" to the real root. +// +// Traversal +// --------- +// The migrator walks the deprecated trie depth-first, decoding one node at +// a time. A leaf becomes a value node at the same path. An internal binary +// node, after both subtrees have been visited, becomes a binary node plus +// up to two edge nodes (one per non-empty child segment). If the trie's +// stored root path is itself non-empty — meaning the deprecated root +// embeds a compression — a single edge node at the empty path is written +// after the traversal completes, replacing the root marker. +// +// Hashes +// ------ +// Starknet trie hashes: +// +// leaf value +// binary hashFn(left-edge-hash, right-edge-hash) +// edge hashFn(child-hash, path-segment-as-felt) + segment-length +// +// Zero-length edges short-circuit to the bare child-hash — the convention +// for absent edges. Class tries hash with Poseidon; contract and storage +// tries with Pedersen. +// +// For small tries every edge hash is computed inline. Above the threshold, +// edge-hash jobs are batched and dispatched to a worker pool for parallel +// computation; the scheduler preserves the original job order so the +// persisted bytes are byte-identical to a natively-built trie2. +// +// In-flight batches flush at target size; cancellation is observed at +// every flush and every channel send. +func (i *ingestor) migrateTrie(t *task, desc TrieDesc, outputs chan<- task) error { + if desc.NodeCount == 0 { + return nil + } + parallelDispatch := desc.NodeCount >= SmallTrieThreshold + prefix := deprecatedTriePrefix(desc) + sched := newHashScheduler(desc.HashFn, parallelDispatch, desc.TrieBucket, desc.Owner, i.pool) + + rootHash, err := i.traverse(t, outputs, prefix, *desc.RootPath, sched) + if err != nil { + return err + } + if err := sched.sync(t.batch); err != nil { + return err + } + if desc.RootPath.Len() > 0 { + if err := writeRootEdgeNode(desc.RootPath, rootHash, sched, t.batch); err != nil { + return err + } + } + return nil +} + +func (i *ingestor) traverse( + t *task, + outputs chan<- task, + prefix []byte, + oldPath trie.BitArray, + sched *hashScheduler, +) (felt.Felt, error) { + parsed, err := readNode(i.database, prefix, &oldPath) + if err != nil { + return felt.Felt{}, err + } + t.nodes++ + + if parsed.isLeaf { + newPath := toNewPath(&oldPath) + if err := writeLeafNode(newPath, &parsed.value, sched, t.batch); err != nil { + return felt.Felt{}, err + } + if err := i.flush(t, outputs); err != nil { + return felt.Felt{}, err + } + return parsed.value, nil + } + + leftHash, err := i.traverse(t, outputs, prefix, parsed.left, sched) + if err != nil { + return felt.Felt{}, err + } + rightHash, err := i.traverse(t, outputs, prefix, parsed.right, sched) + if err != nil { + return felt.Felt{}, err + } + + newPath := toNewPath(&oldPath) + if err := processBinary( + newPath, &parsed.left, &parsed.right, leftHash, rightHash, sched, t.batch, + ); err != nil { + return felt.Felt{}, err + } + if err := i.flush(t, outputs); err != nil { + return felt.Felt{}, err + } + return parsed.value, nil +} + +func rootProcessed(r db.KeyValueReader, newBucket db.Bucket, owner *felt.Address) (bool, error) { + var emptyPath trieutils.Path + var buf [maxNodeKeySize]byte + + n := encodeNodeKey(buf[:], newBucket, owner, &emptyPath, false) + if exists, err := r.Has(buf[:n]); err != nil || exists { + return exists, err + } + n = encodeNodeKey(buf[:], newBucket, owner, &emptyPath, true) + return r.Has(buf[:n]) +} + +func encodeOldPath(path *trie.BitArray, dst []byte) int { + pathLen := path.Len() + b := path.Bytes() + activeBytes := (uint(pathLen) + 7) / 8 + dst[0] = pathLen + copy(dst[1:], b[32-activeBytes:]) + return int(activeBytes) + 1 +} + +type parsedNode struct { + value felt.Felt + left trie.BitArray + right trie.BitArray + isLeaf bool +} + +func readNode(r db.KeyValueReader, prefix []byte, oldPath *trie.BitArray) (parsedNode, error) { + var arr [maxNodeKeySize]byte + n := copy(arr[:], prefix) + n += encodeOldPath(oldPath, arr[n:]) + var node parsedNode + err := r.Get(arr[:n], node.UnmarshalBinary) + return node, err +} + +func (n *parsedNode) UnmarshalBinary(data []byte) error { + if len(data) < felt.Bytes { + return fmt.Errorf("trie: node data too short (%d bytes)", len(data)) + } + n.value = felt.FromBytes[felt.Felt](data[:felt.Bytes]) + data = data[felt.Bytes:] + if len(data) == 0 { + n.isLeaf = true + return nil + } + if err := n.left.UnmarshalBinary(data); err != nil { + return fmt.Errorf("trie: unmarshalling left path: %w", err) + } + data = data[n.left.EncodedLen():] + if err := n.right.UnmarshalBinary(data); err != nil { + return fmt.Errorf("trie: unmarshalling right path: %w", err) + } + return nil +} + +func writeLeafNode( + path trieutils.Path, + value *felt.Felt, + sched *hashScheduler, + batch db.Batch, +) error { + var buf [maxNodeKeySize + valueNodeBlobSize]byte + keyLen := encodeNodeKey(buf[:], sched.bucket, &sched.owner, &path, true) + blob := value.Bytes() + copy(buf[keyLen:], blob[:]) + return batch.Put(buf[:keyLen], buf[keyLen:keyLen+valueNodeBlobSize]) +} + +func processBinary( + parentPath trieutils.Path, + left, right *trie.BitArray, + leftChildHash, rightChildHash felt.Felt, + sched *hashScheduler, + batch db.Batch, +) error { + leftSeg := compressedSegment(left, parentPath.Len()) + rightSeg := compressedSegment(right, parentPath.Len()) + return sched.schedule(&edgeHashJob{ + leftChildHash: leftChildHash, + leftSeg: leftSeg, + rightChildHash: rightChildHash, + rightSeg: rightSeg, + parentPath: parentPath, + }, batch) +} + +func writeRootEdgeNode( + rootPath *trie.BitArray, + childHash felt.Felt, + sched *hashScheduler, + batch db.Batch, +) error { + seg := toNewPath(rootPath) + var buf [edgeNodeMaxSize]byte + n := encodeEdgeNode(buf[:], &childHash, &seg) + return trieutils.WriteNodeByPath( + batch, + sched.bucket, + &sched.owner, + &trieutils.Path{}, + false, + buf[:n], + ) +} + +func deprecatedTriePrefix(desc TrieDesc) []byte { + switch desc.DeprecatedTrieBucket { + case db.ClassesTrie, db.StateTrie: + return desc.DeprecatedTrieBucket.Key() + case db.ContractStorage: + ownerBytes := desc.Owner.Bytes() + return desc.DeprecatedTrieBucket.Key(ownerBytes[:]) + default: + panic(fmt.Sprintf( + "unexpected deprecated trie bucket %v", + desc.DeprecatedTrieBucket, + )) + } +} diff --git a/migration/trie/trie.go b/migration/trie/trie.go new file mode 100644 index 0000000000..415ca2a5d6 --- /dev/null +++ b/migration/trie/trie.go @@ -0,0 +1,299 @@ +package trie + +import ( + "bytes" + "context" + "errors" + "fmt" + "time" + + "github.com/NethermindEth/juno/blockchain/networks" + "github.com/NethermindEth/juno/core/crypto" + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/core/trie" + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/db/dbutils" + "github.com/NethermindEth/juno/migration" + "github.com/NethermindEth/juno/migration/pipeline" + "github.com/NethermindEth/juno/migration/semaphore" + "github.com/NethermindEth/juno/utils/log" +) + +const ( + batchByteSize = 128 * db.Megabyte + targetBatchByteSize = 96 * db.Megabyte + timeLogRate = 5 * time.Second + SmallTrieThreshold = 100_000 + parallelHashBatchSize = 16384 + IngestorCount = 4 +) + +var ( + shouldRerun = []byte{} + shouldNotRerun []byte +) + +var deprecatedTrieBuckets = []db.Bucket{db.ClassesTrie, db.StateTrie, db.ContractStorage} + +// Migrator converts every deprecated Starknet trie on disk into the +// equivalent trie2 layout used by the new state: +// +// ClassesTrie ─→ ClassTrie (Poseidon) +// StateTrie ─→ ContractTrieContract (Pedersen) +// ContractStorage ─→ ContractTrieStorage (Pedersen, per contract owner) +// +// Each deprecated trie is enumerated (its bucket holds one root-path entry +// plus N node entries), then walked in DFS order; every visited node is +// re-encoded into the new format and written to its destination bucket in +// the same batch. After every trie completes successfully, the three +// deprecated buckets are wiped via DeleteRange. +// +// The pipeline runs IngestorCount worker goroutines, each pulling one trie +// at a time from the enumeration source, plus a single committer that +// flushes filled batches to disk. A semaphore caps in-flight batches at +// IngestorCount * 2. See migrateTrie for the per-trie traversal. +// +// Re-run safe: every trie's first action checks for its new-format root +// key; if present, the trie is treated as already migrated and skipped. +// A subsequent run after a crash picks up where the previous one stopped — +// partially migrated tries either have a root key (skipped on the next +// pass) or don't (re-migrated from scratch; the deprecated source data is +// still present because the trailing wipe runs only on full success). +// +// Cancellation: every flush and every channel send checks ctx.Done. On +// cancel, Migrate returns the shouldRerun sentinel with ctx.Err(); the +// migration runner re-invokes on the next process start. +type Migrator struct{} + +var _ migration.Migration = (*Migrator)(nil) + +func (*Migrator) Before([]byte) error { return nil } + +func (*Migrator) Migrate( + ctx context.Context, + database db.KeyValueStore, + _ *networks.Network, + logger log.StructuredLogger, +) ([]byte, error) { + needed, err := needsMigration(database) + if err != nil { + return shouldRerun, err + } + if !needed { + logger.Info("trie migration: no old-format data found, marking applied") + return shouldNotRerun, nil + } + + return runMigration(ctx, database, logger) +} + +func needsMigration(r db.KeyValueReader) (bool, error) { + for _, bucket := range deprecatedTrieBuckets { + prefix := bucket.Key() + iter, err := r.NewIterator(prefix, true) + if err != nil { + return false, err + } + hasKeys := iter.First() + if err := iter.Close(); err != nil { + return false, err + } + if hasKeys { + return true, nil + } + } + return false, nil +} + +func runMigration( + ctx context.Context, + database db.KeyValueStore, + logger log.StructuredLogger, +) ([]byte, error) { + batchSem := semaphore.New(IngestorCount*2, func() db.Batch { + return database.NewBatchWithSize(batchByteSize) + }) + + pool := newHashWorkerPool() + defer pool.close() + + ing := newIngestor(ctx, database, batchSem, pool) + + tries, err := enumerateTries(database) + if err != nil { + return shouldRerun, err + } + + var allNodes uint64 + for _, d := range tries { + allNodes += uint64(d.NodeCount) + } + allTries := uint64(len(tries)) + + src := pipeline.Source(func(yield func(TrieDesc) bool) { + for _, d := range tries { + if !yield(d) { + return + } + } + }) + ingested := pipeline.New(src, IngestorCount, ing) + committed := pipeline.New( + ingested, + 1, + newCommitter(logger, batchSem, allTries, allNodes), + ) + + _, wait := committed.Run(ctx) + res := wait() + if res.Err != nil { + return shouldRerun, res.Err + } + if !res.IsDone { + if ctxErr := ctx.Err(); ctxErr != nil { + return shouldRerun, ctxErr + } + return shouldRerun, errors.New("trie migration: pipeline did not complete") + } + + if err := wipeDeprecatedBuckets(database); err != nil { + return shouldRerun, err + } + logger.Info("trie migration: source buckets deleted") + + return shouldNotRerun, nil +} + +func wipeDeprecatedBuckets(database db.KeyValueRangeDeleter) error { + for _, bucket := range deprecatedTrieBuckets { + start := bucket.Key() + end := dbutils.UpperBound(start) + if err := database.DeleteRange(start, end); err != nil { + return fmt.Errorf("trie migration: cleanup DeleteRange for %v: %w", bucket, err) + } + } + return nil +} + +type TrieDesc struct { + DeprecatedTrieBucket db.Bucket + TrieBucket db.Bucket + Owner felt.Address + HashFn crypto.HashFn + NodeCount int + RootPath *trie.BitArray +} + +func enumerateTries(r db.KeyValueReader) ([]TrieDesc, error) { + var descs []TrieDesc + + for _, spec := range []struct { + oldBucket, newBucket db.Bucket + hashFn crypto.HashFn + }{ + {db.ClassesTrie, db.ClassTrie, crypto.Poseidon}, + {db.StateTrie, db.ContractTrieContract, crypto.Pedersen}, + } { + desc, err := enumerateGlobalTrie(r, spec.oldBucket, spec.newBucket, spec.hashFn) + if err != nil { + return nil, err + } + descs = append(descs, desc) + } + + storageDescs, err := enumerateStorageTries(r) + if err != nil { + return nil, err + } + descs = append(descs, storageDescs...) + + return descs, nil +} + +func enumerateGlobalTrie( + r db.KeyValueReader, + oldBucket, newBucket db.Bucket, + hashFn crypto.HashFn, +) (TrieDesc, error) { + prefix := oldBucket.Key() + it, err := r.NewIterator(prefix, true) + if err != nil { + return TrieDesc{}, fmt.Errorf("opening iterator for bucket %v: %w", oldBucket, err) + } + defer it.Close() + it.First() + + rootPath, count, err := scanTrie(it, prefix) + if err != nil { + return TrieDesc{}, fmt.Errorf("enumerating bucket %v: %w", oldBucket, err) + } + return TrieDesc{ + DeprecatedTrieBucket: oldBucket, + TrieBucket: newBucket, + HashFn: hashFn, + NodeCount: count, + RootPath: &rootPath, + }, nil +} + +func enumerateStorageTries(r db.KeyValueReader) ([]TrieDesc, error) { + it, err := r.NewIterator(db.ContractStorage.Key(), true) + if err != nil { + return nil, fmt.Errorf("opening storage iterator: %w", err) + } + defer it.Close() + it.First() + + var descs []TrieDesc + for it.Valid() { + key := it.Key() + if len(key) < 1+felt.Bytes { + it.Next() + continue + } + owner := felt.FromBytes[felt.Address](key[1 : 1+felt.Bytes]) + ownerBytes := owner.Bytes() + ownerPrefix := db.ContractStorage.Key(ownerBytes[:]) + + rootPath, count, err := scanTrie(it, ownerPrefix) + if err != nil { + return nil, fmt.Errorf("enumerating storage owner %s: %w", &owner, err) + } + descs = append(descs, TrieDesc{ + DeprecatedTrieBucket: db.ContractStorage, + TrieBucket: db.ContractTrieStorage, + Owner: owner, + HashFn: crypto.Pedersen, + NodeCount: count, + RootPath: &rootPath, + }) + // scanTrie leaves the iterator positioned past this owner's range. + } + return descs, nil +} + +func scanTrie(it db.Iterator, prefix []byte) (trie.BitArray, int, error) { + var rootPath trie.BitArray + count := 0 + for it.Valid() { + key := it.Key() + if !bytes.HasPrefix(key, prefix) { + return rootPath, count, nil + } + if len(key) == len(prefix) { + val, err := it.Value() + if err != nil { + return trie.BitArray{}, 0, err + } + parsedRootPath, err := parseDeprecatedPath(val) + if err != nil { + return trie.BitArray{}, 0, err + } + rootPath = parsedRootPath + } else { + count++ + } + it.Next() + } + return rootPath, count, nil +} diff --git a/migration/trie/trie_test.go b/migration/trie/trie_test.go new file mode 100644 index 0000000000..c2d08b9df0 --- /dev/null +++ b/migration/trie/trie_test.go @@ -0,0 +1,441 @@ +package trie_test + +import ( + "context" + "testing" + + "github.com/NethermindEth/juno/core/crypto" + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/core/trie" + "github.com/NethermindEth/juno/core/trie2" + "github.com/NethermindEth/juno/core/trie2/triedb/rawdb" + "github.com/NethermindEth/juno/core/trie2/trienode" + "github.com/NethermindEth/juno/core/trie2/trieutils" + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/db/memory" + trielib "github.com/NethermindEth/juno/migration/trie" + "github.com/NethermindEth/juno/utils/log" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type leafMap map[felt.Felt]felt.Felt + +type trieCase struct { + name string + oldBucket db.Bucket + newBucket db.Bucket + owner felt.Address + oldBuildPrefix func(owner *felt.Address) []byte + newTrieID func(owner *felt.Address) trieutils.TrieID + hashFn crypto.HashFn + //nolint:staticcheck // Necessary for old state + buildOldFn func(db.IndexedBatch, []byte, uint8) (*trie.Trie, error) +} + +var trieCases = []trieCase{ + { + name: "ClassTrie", + oldBucket: db.ClassesTrie, + newBucket: db.ClassTrie, + oldBuildPrefix: func(_ *felt.Address) []byte { return []byte{byte(db.ClassesTrie)} }, + newTrieID: func(_ *felt.Address) trieutils.TrieID { + return trieutils.NewClassTrieID(felt.StateRootHash(felt.One)) + }, + hashFn: crypto.Poseidon, + buildOldFn: trie.NewTriePoseidon, + }, + { + name: "ContractTrie", + oldBucket: db.StateTrie, + newBucket: db.ContractTrieContract, + oldBuildPrefix: func(_ *felt.Address) []byte { return []byte{byte(db.StateTrie)} }, + newTrieID: func(_ *felt.Address) trieutils.TrieID { + return trieutils.NewContractTrieID(felt.StateRootHash(felt.One)) + }, + hashFn: crypto.Pedersen, + buildOldFn: trie.NewTriePedersen, + }, + { + name: "StorageTrie", + oldBucket: db.ContractStorage, + newBucket: db.ContractTrieStorage, + owner: felt.FromUint64[felt.Address](42), + oldBuildPrefix: func(owner *felt.Address) []byte { + ownerBytes := owner.Bytes() + return db.ContractStorage.Key(ownerBytes[:]) + }, + newTrieID: func(owner *felt.Address) trieutils.TrieID { + return trieutils.NewContractStorageTrieID(felt.StateRootHash(felt.One), *owner) + }, + hashFn: crypto.Pedersen, + buildOldFn: trie.NewTriePedersen, + }, +} + +var transcoderCases = []struct { + name string + leaves leafMap +}{ + {"empty trie", nil}, + {"single leaf", leafMap{ + felt.FromUint64[felt.Felt](1): felt.FromUint64[felt.Felt](100), + }}, + {"deep split", leafMap{ + felt.FromUint64[felt.Felt](2): felt.FromUint64[felt.Felt](10), + felt.FromUint64[felt.Felt](3): felt.FromUint64[felt.Felt](20), + }}, + {"left right split", leafMap{ + felt.FromUint64[felt.Felt](1): felt.FromUint64[felt.Felt](10), + felt.FromBytes[felt.Felt]([]byte{0x04}): felt.FromUint64[felt.Felt](20), // 2^250 + }}, + {"full depth 2 tree", leafMap{ + felt.FromUint64[felt.Felt](1): felt.FromUint64[felt.Felt](10), // 00... + felt.FromBytes[felt.Felt]([]byte{0x02}): felt.FromUint64[felt.Felt](20), // 01... + felt.FromBytes[felt.Felt]([]byte{0x04}): felt.FromUint64[felt.Felt](30), // 10... + felt.FromBytes[felt.Felt]([]byte{0x06}): felt.FromUint64[felt.Felt](40), // 11... + }}, + {"hundred sequential leaves", func() leafMap { + leaves := make(leafMap, 100) + for i := 1; i <= 100; i++ { + leaves[felt.FromUint64[felt.Felt](uint64(i))] = felt.FromUint64[felt.Felt](uint64(i) * 7) + } + return leaves + }()}, + {"random 1000 leaves", randomLeaves(1000)}, +} + +func TestMigrate_FreshDBIsNoOp(t *testing.T) { + memDB := memory.New() + + state, err := (&trielib.Migrator{}).Migrate( + context.Background(), + memDB, + nil, + log.NewNopZapLogger(), + ) + require.NoError(t, err) + assert.Nil( + t, + state, + "fresh DB must mark migration applied (nil intermediate state) without doing work", + ) +} + +func TestMigrate_RunsWhenOldDataPresent(t *testing.T) { + leaves := randomLeaves(100) + memDB := buildFullDB(t, leaves) + + require.True(t, bucketHasKeys(t, memDB, db.ClassesTrie), "precondition: DB has old-format data") + + state, err := (&trielib.Migrator{}).Migrate( + context.Background(), + memDB, + nil, + log.NewNopZapLogger(), + ) + require.NoError(t, err) + assert.Nil(t, state, "completed migration must return nil intermediate state") + + for _, bucket := range []db.Bucket{db.ClassesTrie, db.StateTrie, db.ContractStorage} { + assert.False(t, + bucketHasKeys(t, memDB, bucket), + "old-format bucket %v should be empty after migration", bucket) + } +} + +// TestMigrationEndToEnd verifies that the migration produces byte-for-byte +// identical DB output to a natively-built trie2 for all three trie types and +// all leaf counts. Catches encoding bugs that root-hash comparison cannot. +func TestMigrationEndToEnd(t *testing.T) { + type testCase struct { + name string + tc trieCase + leaves leafMap + } + + var cases []testCase + for _, tc := range trieCases { + for _, lc := range transcoderCases { + cases = append(cases, testCase{ + name: tc.name + "/" + lc.name, + tc: tc, + leaves: lc.leaves, + }) + } + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + prefix := c.tc.oldBuildPrefix(&c.tc.owner) + + migratedDB := memory.New() + buildDeprecatedTrie(t, migratedDB, c.leaves, c.tc.buildOldFn, prefix) + _, err := (&trielib.Migrator{}).Migrate( + context.Background(), migratedDB, nil, log.NewNopZapLogger(), + ) + require.NoError(t, err) + + nativeDB := memory.New() + buildTrie(t, nativeDB, c.leaves, + c.tc.newTrieID(&c.tc.owner), c.tc.hashFn, c.tc.newBucket) + + assert.Equal(t, + allKeysUnder(t, nativeDB, c.tc.newBucket), + allKeysUnder(t, migratedDB, c.tc.newBucket)) + }) + } +} + +// TestMigrationIsResumable verifies that re-running migration over a DB +// whose class-trie destination root is already present skips the class +// migration and only finishes the contract trie. The "partial state" is +// faked by copying the reference DB's new-format class-trie keys into the +// partial DB before running migration. +func TestMigrationIsResumable(t *testing.T) { + leaves := randomLeaves(1000) + + // Reference: full migration from scratch. + refDB := buildFullDB(t, leaves) + _, err := (&trielib.Migrator{}).Migrate(context.Background(), refDB, nil, log.NewNopZapLogger()) + require.NoError(t, err) + + // Partial DB: both tries in old format initially. + partialDB := buildFullDB(t, leaves) + + // Fake a prior successful class-trie migration by copying refDB's + // new-format class-trie keys directly into partialDB. + refClassKeys := allKeysUnder(t, refDB, db.ClassTrie) + require.NotEmpty(t, refClassKeys, "reference class trie should be non-empty") + for k, v := range refClassKeys { + require.NoError(t, partialDB.Put([]byte(k), v)) + } + + // Resume: migration should skip the class trie (its dest root is present) + // and complete only the contract trie. + _, err = (&trielib.Migrator{}).Migrate(context.Background(), partialDB, nil, log.NewNopZapLogger()) + require.NoError(t, err) + + // Final state must match the reference full-run output for every new-format bucket. + for _, bucket := range []db.Bucket{db.ClassTrie, db.ContractTrieContract} { + refKeys := allKeysUnder(t, refDB, bucket) + resumedKeys := allKeysUnder(t, partialDB, bucket) + assert.Equal(t, refKeys, resumedKeys, + "resumed migration result differs from full run for bucket %v", bucket) + } +} + +// TestMigrationMultiStorageOwners exercises enumerateStorageTries across +// multiple owners (scanTrie's prefix-leave path) and keeps all 4 ingestor +// workers busy by giving them 7 tries to chew through (2 global + 5 storage). +func TestMigrationMultiStorageOwners(t *testing.T) { + leaves := randomLeaves(50) + + migratedDB := memory.New() + buildDeprecatedTrie(t, migratedDB, leaves, trie.NewTriePoseidon, db.ClassesTrie.Key()) + buildDeprecatedTrie(t, migratedDB, leaves, trie.NewTriePedersen, db.StateTrie.Key()) + + owners := []felt.Address{ + felt.FromUint64[felt.Address](1), + felt.FromUint64[felt.Address](2), + felt.FromUint64[felt.Address](3), + felt.FromUint64[felt.Address](42), + felt.FromUint64[felt.Address](999), + } + for _, owner := range owners { + ownerBytes := owner.Bytes() + buildDeprecatedTrie(t, migratedDB, leaves, trie.NewTriePedersen, + db.ContractStorage.Key(ownerBytes[:])) + } + + _, err := (&trielib.Migrator{}).Migrate( + context.Background(), + migratedDB, + nil, + log.NewNopZapLogger(), + ) + require.NoError(t, err) + + // Per-owner native build → assert every native key is present (with the + migratedAll := allKeysUnder(t, migratedDB, db.ContractTrieStorage) + for _, owner := range owners { + nativeDB := memory.New() + id := trieutils.NewContractStorageTrieID(felt.StateRootHash(felt.One), owner) + buildTrie(t, nativeDB, leaves, id, crypto.Pedersen, db.ContractTrieStorage) + for k, v := range allKeysUnder(t, nativeDB, db.ContractTrieStorage) { + gotV, ok := migratedAll[k] + require.True(t, ok, "owner %v missing key", owner) + assert.Equal(t, v, gotV, "owner %v value differs at key", owner) + } + } + + // Old buckets fully drained. + for _, bucket := range []db.Bucket{db.ClassesTrie, db.StateTrie, db.ContractStorage} { + assert.False(t, bucketHasKeys(t, migratedDB, bucket), + "old bucket %v should be drained", bucket) + } +} + +// TestMigrationIdempotent verifies that a successful migration is a no-op on +// a second run: needsMigration sees the wiped deprecated buckets and returns +// early without touching the migrated state. +func TestMigrationIsNoopOnSecondRun(t *testing.T) { + leaves := randomLeaves(100) + memDB := buildFullDB(t, leaves) + + state, err := (&trielib.Migrator{}).Migrate( + context.Background(), + memDB, + nil, + log.NewNopZapLogger(), + ) + require.NoError(t, err) + require.Nil(t, state) + snapshot := snapshotAllBuckets(t, memDB, + db.ClassTrie, db.ContractTrieContract, db.ContractTrieStorage) + + state, err = (&trielib.Migrator{}).Migrate(context.Background(), memDB, nil, log.NewNopZapLogger()) + require.NoError(t, err) + require.Nil(t, state) + require.Equal(t, snapshot, + snapshotAllBuckets(t, memDB, + db.ClassTrie, db.ContractTrieContract, db.ContractTrieStorage), + "second Migrate call must not change state") +} + +// TestMigrationCancelledContext verifies that a pre-cancelled ctx surfaces +// context.Canceled with the shouldRerun sentinel, and that a fresh ctx +// completes the migration normally afterwards. +func TestMigrationCancelledContext(t *testing.T) { + leaves := randomLeaves(100) + memDB := buildFullDB(t, leaves) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + state, err := (&trielib.Migrator{}).Migrate(ctx, memDB, nil, log.NewNopZapLogger()) + require.ErrorIs(t, err, context.Canceled) + require.NotNil(t, state, "shouldRerun sentinel must not be nil") + require.Empty(t, state, "shouldRerun is a non-nil empty slice") + + state, err = (&trielib.Migrator{}).Migrate(context.Background(), memDB, nil, log.NewNopZapLogger()) + require.NoError(t, err) + require.Nil(t, state) +} + +// buildFullDB creates an old-format DB populated with a class, a contract, and +// one storage trie, all built from the same leaf set. +func buildFullDB(t *testing.T, leaves leafMap) db.KeyValueStore { + t.Helper() + memDB := memory.New() + + owner := felt.FromUint64[felt.Address](42) + ownerBytes := owner.Bytes() + storagePrefix := db.ContractStorage.Key(ownerBytes[:]) + + buildDeprecatedTrie(t, memDB, leaves, trie.NewTriePoseidon, db.ClassesTrie.Key()) + buildDeprecatedTrie(t, memDB, leaves, trie.NewTriePedersen, db.StateTrie.Key()) + buildDeprecatedTrie(t, memDB, leaves, trie.NewTriePedersen, storagePrefix) + + return memDB +} + +func buildDeprecatedTrie( + t *testing.T, + database db.KeyValueStore, + leaves leafMap, + //nolint:staticcheck // Necessary for old state + trieFn func(db.IndexedBatch, []byte, uint8) (*trie.Trie, error), + prefix []byte, +) felt.Felt { + t.Helper() + //nolint:staticcheck // Necessary for old state + txn := database.NewIndexedBatch() + tr, err := trieFn(txn, prefix, 251) + require.NoError(t, err) + for key, value := range leaves { + _, err := tr.Put(&key, &value) + require.NoError(t, err) + } + root, err := tr.Root() + require.NoError(t, err) + require.NoError(t, tr.Commit()) + require.NoError(t, txn.Write()) + return root +} + +func buildTrie( + t *testing.T, + kvStore db.KeyValueStore, + leaves leafMap, + id trieutils.TrieID, + hashFn crypto.HashFn, + newBucket db.Bucket, +) { + t.Helper() + rawDB := rawdb.New(kvStore) + tr, err := trie2.New(id, 251, hashFn, rawDB) + require.NoError(t, err) + for key, value := range leaves { + require.NoError(t, tr.Update(&key, &value)) + } + root, nodes := tr.Commit() + if nodes == nil { + return // empty trie — nothing to persist + } + mergeSet := trienode.NewMergeNodeSet(nodes) + var zero felt.StateRootHash + stateRoot := felt.StateRootHash(root) + batch := kvStore.NewBatch() + if newBucket == db.ClassTrie { + require.NoError(t, rawDB.Update(&stateRoot, &zero, 0, mergeSet, nil, batch)) + } else { + require.NoError(t, rawDB.Update(&stateRoot, &zero, 0, nil, mergeSet, batch)) + } + require.NoError(t, batch.Write()) +} + +func randomLeaves(n int) leafMap { + leaves := make(leafMap, n) + for len(leaves) < n { + var k, v felt.Felt + k.SetRandom() + v.SetRandom() + leaves[k] = v + } + return leaves +} + +func snapshotAllBuckets(t *testing.T, r db.KeyValueReader, buckets ...db.Bucket) map[string][]byte { + t.Helper() + out := make(map[string][]byte) + for _, b := range buckets { + for k, v := range allKeysUnder(t, r, b) { + out[k] = v + } + } + return out +} + +func allKeysUnder(t *testing.T, r db.KeyValueReader, bucket db.Bucket) map[string][]byte { + t.Helper() + prefix := bucket.Key() + iter, err := r.NewIterator(prefix, true) + require.NoError(t, err) + defer iter.Close() + out := make(map[string][]byte) + for ok := iter.First(); ok; ok = iter.Next() { + val, err := iter.Value() + require.NoError(t, err) + out[string(iter.Key())] = val + } + return out +} + +func bucketHasKeys(t *testing.T, r db.KeyValueReader, bucket db.Bucket) bool { + t.Helper() + it, err := r.NewIterator(bucket.Key(), true) + require.NoError(t, err) + defer it.Close() + return it.First() +} diff --git a/node/migration.go b/node/migration.go index 57a53e9eac..2ca79e2332 100644 --- a/node/migration.go +++ b/node/migration.go @@ -14,6 +14,7 @@ import ( "github.com/NethermindEth/juno/migration/headstate" "github.com/NethermindEth/juno/migration/historyprunner" "github.com/NethermindEth/juno/migration/statehistory" + "github.com/NethermindEth/juno/migration/trie" "github.com/NethermindEth/juno/utils/log" ) @@ -29,7 +30,8 @@ func registerMigrations(cfg *Config) *migration.Registry { PruneModeFlag, ). WithOptional(&headstate.Migrator{}, cfg.NewState, "new-state"). - WithOptional(&statehistory.Migrator{}, cfg.NewState, "new-state") + WithOptional(&statehistory.Migrator{}, cfg.NewState, "new-state"). + WithOptional(&trie.Migrator{}, cfg.NewState, "new-state") return registry }