From 70c08083214cd99a7dbb24ce3b2d8687fe96ceaf Mon Sep 17 00:00:00 2001
From: Aman Sanghi <asanghi@offchainlabs.com>
Date: Mon, 13 May 2024 18:23:59 +0530
Subject: [PATCH] snap-sync consensus 0.2: start reading parent chain from the
 right block

---
 arbnode/node.go          | 81 ++++++++++++++++++++++++++++++++++------
 staker/rollup_watcher.go | 39 ++++++++++++++++++-
 2 files changed, 108 insertions(+), 12 deletions(-)

diff --git a/arbnode/node.go b/arbnode/node.go
index c346a38e14..a209e81e3a 100644
--- a/arbnode/node.go
+++ b/arbnode/node.go
@@ -277,19 +277,21 @@ type Node struct {
 }
 
 type SnapSyncConfig struct {
-	Enabled               bool
-	PrevBatchMessageCount uint64
-	PrevDelayedRead       uint64
-	BatchCount            uint64
-	DelayedCount          uint64
+	Enabled                   bool
+	PrevBatchMessageCount     uint64
+	PrevDelayedRead           uint64
+	BatchCount                uint64
+	DelayedCount              uint64
+	ParentChainAssertionBlock uint64
 }
 
 var DefaultSnapSyncConfig = SnapSyncConfig{
-	Enabled:               false,
-	PrevBatchMessageCount: 0,
-	BatchCount:            0,
-	DelayedCount:          0,
-	PrevDelayedRead:       0,
+	Enabled:                   false,
+	PrevBatchMessageCount:     0,
+	PrevDelayedRead:           0,
+	BatchCount:                0,
+	DelayedCount:              0,
+	ParentChainAssertionBlock: 0,
 }
 
 type ConfigFetcher interface {
@@ -564,7 +566,25 @@ func createNodeImpl(
 	if err != nil {
 		return nil, err
 	}
-	inboxReader, err := NewInboxReader(inboxTracker, l1client, l1Reader, new(big.Int).SetUint64(deployInfo.DeployedAt), delayedBridge, sequencerInbox, func() *InboxReaderConfig { return &configFetcher.Get().InboxReader })
+	firstMessageBlock := new(big.Int).SetUint64(deployInfo.DeployedAt)
+	if config.SnapSyncTest.Enabled {
+		firstMessageToRead := config.SnapSyncTest.DelayedCount
+		if firstMessageToRead > config.SnapSyncTest.BatchCount {
+			firstMessageToRead = config.SnapSyncTest.BatchCount
+		}
+		if firstMessageToRead > 0 {
+			firstMessageToRead--
+		}
+		// Find the first block containing the first message to read
+		// Subtract 1 to get the block before the first message to read,
+		// this is done to fetch previous batch metadata needed for snap sync.
+		block, err := FindBlockContainingBatch(ctx, deployInfo.Rollup, l1client, config.SnapSyncTest.ParentChainAssertionBlock, firstMessageToRead-1)
+		if err != nil {
+			return nil, err
+		}
+		firstMessageBlock.SetUint64(block)
+	}
+	inboxReader, err := NewInboxReader(inboxTracker, l1client, l1Reader, firstMessageBlock, delayedBridge, sequencerInbox, func() *InboxReaderConfig { return &configFetcher.Get().InboxReader })
 	if err != nil {
 		return nil, err
 	}
@@ -740,6 +760,45 @@ func createNodeImpl(
 	}, nil
 }
 
+func FindBlockContainingBatch(ctx context.Context, rollupAddress common.Address, l1Client arbutil.L1Interface, parentChainAssertionBlock uint64, batch uint64) (uint64, error) {
+	callOpts := bind.CallOpts{Context: ctx}
+	rollup, err := staker.NewRollupWatcher(rollupAddress, l1Client, callOpts)
+	if err != nil {
+		return 0, err
+	}
+	high := parentChainAssertionBlock
+	low := high / 2
+	// Exponentially reduce high and low by a factor of 2 until lowNode.InboxMaxCount < batch
+	// This will give us a range (low to high) of blocks that contain the batch
+	for low > 0 {
+		lowNode, err := rollup.LookupNodeByBlockNumber(ctx, low)
+		if err != nil {
+			return 0, err
+		}
+		if lowNode.InboxMaxCount.Uint64() > batch {
+			high = low
+			low = low / 2
+		} else {
+			break
+		}
+	}
+	// Then binary search between low and high to find the block containing the batch
+	for low < high {
+		mid := low + (high-low)/2
+
+		midNode, err := rollup.LookupNodeByBlockNumber(ctx, mid)
+		if err != nil {
+			return 0, err
+		}
+		if midNode.InboxMaxCount.Uint64() < batch {
+			low = mid + 1
+		} else {
+			high = mid
+		}
+	}
+	return low, nil
+}
+
 func (n *Node) OnConfigReload(_ *Config, _ *Config) error {
 	// TODO
 	return nil
diff --git a/staker/rollup_watcher.go b/staker/rollup_watcher.go
index 118ce15b44..0cc5b43999 100644
--- a/staker/rollup_watcher.go
+++ b/staker/rollup_watcher.go
@@ -57,7 +57,6 @@ func NewRollupWatcher(address common.Address, client arbutil.L1Interface, callOp
 	if err != nil {
 		return nil, err
 	}
-
 	return &RollupWatcher{
 		address:         address,
 		client:          client,
@@ -165,6 +164,44 @@ func (r *RollupWatcher) LookupNode(ctx context.Context, number uint64) (*NodeInf
 	}, nil
 }
 
+func (r *RollupWatcher) LookupNodeByBlockNumber(ctx context.Context, blockNumber uint64) (*NodeInfo, error) {
+	var query = ethereum.FilterQuery{
+		FromBlock: big.NewInt(int64(blockNumber)),
+		ToBlock:   big.NewInt(int64(blockNumber)),
+		Addresses: []common.Address{r.address},
+		Topics:    [][]common.Hash{{nodeCreatedID}},
+	}
+	logs, err := r.client.FilterLogs(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+	if len(logs) == 0 {
+		return nil, fmt.Errorf("couldn't find node at the request blockNumber %v", blockNumber)
+	}
+	if len(logs) > 1 {
+		return nil, fmt.Errorf("found multiple instances of node at the requested blockNumber %v", blockNumber)
+	}
+	ethLog := logs[0]
+	parsedLog, err := r.ParseNodeCreated(ethLog)
+	if err != nil {
+		return nil, err
+	}
+	l1BlockProposed, err := arbutil.CorrespondingL1BlockNumber(ctx, r.client, ethLog.BlockNumber)
+	if err != nil {
+		return nil, err
+	}
+	return &NodeInfo{
+		NodeNum:                  parsedLog.NodeNum,
+		L1BlockProposed:          l1BlockProposed,
+		ParentChainBlockProposed: ethLog.BlockNumber,
+		Assertion:                NewAssertionFromSolidity(parsedLog.Assertion),
+		InboxMaxCount:            parsedLog.InboxMaxCount,
+		AfterInboxBatchAcc:       parsedLog.AfterInboxBatchAcc,
+		NodeHash:                 parsedLog.NodeHash,
+		WasmModuleRoot:           parsedLog.WasmModuleRoot,
+	}, nil
+}
+
 func (r *RollupWatcher) LookupNodeChildren(ctx context.Context, nodeNum uint64, nodeHash common.Hash) ([]*NodeInfo, error) {
 	node, err := r.RollupUserLogic.GetNode(r.getCallOpts(ctx), nodeNum)
 	if err != nil {