From ab065006b477834cfcbf99952a0f46bfdaec04eb Mon Sep 17 00:00:00 2001
From: Tyler Sturos <55340199+tjsturos@users.noreply.github.com>
Date: Tue, 3 Sep 2024 14:40:50 -0800
Subject: [PATCH] Remove extra difficulty check for proofs (#289)

* Remove extra difficulty check for proofs

* remove need for underflow check

---------

Co-authored-by: Base Dev <gitlab.dollop533@passmail.net>
---
 node/crypto/wesolowski_frame_prover.go      | 20 ++++++++++++--------
 node/crypto/wesolowski_frame_prover_test.go | 21 +++++++++++++++++++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/node/crypto/wesolowski_frame_prover.go b/node/crypto/wesolowski_frame_prover.go
index 2d234cb..812b034 100644
--- a/node/crypto/wesolowski_frame_prover.go
+++ b/node/crypto/wesolowski_frame_prover.go
@@ -579,15 +579,22 @@ func (w *WesolowskiFrameProver) VerifyWeakRecursiveProof(
 	}
 }
 
+func (w *WesolowskiFrameProver) CalculateChallengeProofDifficulty(
+	increment uint32,
+) uint32 {
+	if increment >= 700000 {
+		return 25000
+	}
+
+	return 200000 - (increment / 4)
+}
+
 func (w *WesolowskiFrameProver) CalculateChallengeProof(
 	challenge []byte,
 	core uint32,
 	increment uint32,
 ) ([]byte, error) {
-	difficulty := 200000 - (increment / 4)
-	if difficulty < 25000 || increment > 800000 {
-		difficulty = 25000
-	}
+	difficulty := w.CalculateChallengeProofDifficulty(increment)
 
 	instanceInput := binary.BigEndian.AppendUint32([]byte{}, core)
 	instanceInput = append(instanceInput, challenge...)
@@ -606,10 +613,7 @@ func (w *WesolowskiFrameProver) VerifyChallengeProof(
 	core uint32,
 	proof []byte,
 ) bool {
-	difficulty := 200000 - (increment / 4)
-	if difficulty < 25000 || increment > 800000 {
-		difficulty = 25000
-	}
+	difficulty := w.CalculateChallengeProofDifficulty(increment)
 
 	if len(proof) != 516 {
 		return false
diff --git a/node/crypto/wesolowski_frame_prover_test.go b/node/crypto/wesolowski_frame_prover_test.go
index 4e19566..48b3772 100644
--- a/node/crypto/wesolowski_frame_prover_test.go
+++ b/node/crypto/wesolowski_frame_prover_test.go
@@ -34,3 +34,24 @@ func TestChallengeProof(t *testing.T) {
 	assert.NoError(t, err)
 	assert.True(t, w.VerifyChallengeProof([]byte{0x01, 0x02, 0x03}, 1, 0, proofs))
 }
+
+func TestCalculateChallengeProofDifficulty(t *testing.T) {
+	l, _ := zap.NewProduction()
+	w := crypto.NewWesolowskiFrameProver(l)
+
+	// At 0 increments, the difficulty should be 200,000
+	difficulty0 := w.CalculateChallengeProofDifficulty(0)
+	assert.Equal(t, 200000, difficulty0)
+
+	// At 100,000 increments, the difficulty should be 175,000
+	difficulty100k := w.CalculateChallengeProofDifficulty(100000)
+	assert.Equal(t, 175000, difficulty100k)
+
+	// At 700,000 increments, the difficulty should be 25,000
+	difficulty700k := w.CalculateChallengeProofDifficulty(700000)
+	assert.Equal(t, 25000, difficulty700k)
+
+	// At 800,000 increments, the difficulty should stay at 25,000
+	difficulty800k := w.CalculateChallengeProofDifficulty(800000)
+	assert.Equal(t, 25000, difficulty800k)
+}