feat(migration): trie migration

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
+}

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
+}

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))
+}

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) }