From ae74c56a08db46b2585cfdd6e657c03a9db69584 Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <57538841+WeebDataHoarder@users.noreply.github.com>
Date: Sun, 30 Jul 2023 20:45:02 +0200
Subject: [PATCH] Encode types.Difficulty as uint64 if Difficulty.Hi is zero

---
 types/difficulty.go | 50 +++++++++++++++++++++++++++++++--------------
 utils/number.go     | 40 ++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/types/difficulty.go b/types/difficulty.go
index 32859f4..e2f5515 100644
--- a/types/difficulty.go
+++ b/types/difficulty.go
@@ -8,10 +8,12 @@ import (
 	"git.gammaspectra.live/P2Pool/p2pool-observer/utils"
 	"github.com/holiman/uint256"
 	fasthex "github.com/tmthrgd/go-hex"
+	"io"
 	"lukechampine.com/uint128"
 	"math"
 	"math/big"
 	"math/bits"
+	"strconv"
 	"strings"
 )
 
@@ -189,6 +191,10 @@ func (d Difficulty) Big() *big.Int {
 }
 
 func (d Difficulty) MarshalJSON() ([]byte, error) {
+	if d.Hi == 0 {
+		return []byte(strconv.FormatUint(d.Lo, 10)), nil
+	}
+
 	var encodeBuf [DifficultySize]byte
 	d.PutBytesBE(encodeBuf[:])
 
@@ -243,28 +249,42 @@ func DifficultyFrom64(v uint64) Difficulty {
 }
 
 func (d *Difficulty) UnmarshalJSON(b []byte) error {
-	var s string
-	if err := utils.UnmarshalJSON(b, &s); err != nil {
-		return err
+	if len(b) == 0 {
+		return io.ErrUnexpectedEOF
 	}
 
-	if len(b) == DifficultySize*2+2 {
-		// fast path
-		var buf [DifficultySize]byte
-		if _, err := fasthex.Decode(buf[:], b[1:len(b)-1]); err != nil {
+	if b[0] == '"' {
+		if len(b) < 2 || (len(b)%2) != 0 || b[len(b)-1] != '"' {
+			return errors.New("invalid bytes")
+		}
+
+		if len(b) == DifficultySize*2+2 {
+			// fast path
+			var buf [DifficultySize]byte
+			if _, err := fasthex.Decode(buf[:], b[1:len(b)-1]); err != nil {
+				return err
+			} else {
+				*d = DifficultyFromBytes(buf[:])
+				return nil
+			}
+		}
+
+		if diff, err := DifficultyFromString(string(b[1 : len(b)-1])); err != nil {
 			return err
 		} else {
-			*d = DifficultyFromBytes(buf[:])
+			*d = diff
+
 			return nil
 		}
-	}
-
-	if diff, err := DifficultyFromString(s); err != nil {
-		return err
 	} else {
-		*d = diff
-
-		return nil
+		// Difficulty as uint64
+		var err error
+		if d.Lo, err = utils.ParseUint64(b); err != nil {
+			return err
+		} else {
+			d.Hi = 0
+			return nil
+		}
 	}
 }
 
diff --git a/utils/number.go b/utils/number.go
index 01b5cc1..8ed46f1 100644
--- a/utils/number.go
+++ b/utils/number.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/hex"
+	"fmt"
 	"github.com/jxskiss/base62"
 	fasthex "github.com/tmthrgd/go-hex"
 	"math/bits"
@@ -162,3 +163,42 @@ func UVarInt64Size[T uint64 | int](v T) (n int) {
 		return binary.MaxVarintLen64
 	}
 }
+
+// ParseUint64 parses uint64 from s.
+//
+// It is equivalent to strconv.ParseUint(s, 10, 64), but is faster.
+//
+// From https://github.com/valyala/fastjson
+func ParseUint64(s []byte) (uint64, error) {
+	if len(s) == 0 {
+		return 0, fmt.Errorf("cannot parse uint64 from empty string")
+	}
+	i := uint(0)
+	d := uint64(0)
+	j := i
+	for i < uint(len(s)) {
+		if s[i] >= '0' && s[i] <= '9' {
+			d = d*10 + uint64(s[i]-'0')
+			i++
+			if i > 18 {
+				// The integer part may be out of range for uint64.
+				// Fall back to slow parsing.
+				dd, err := strconv.ParseUint(string(s), 10, 64)
+				if err != nil {
+					return 0, err
+				}
+				return dd, nil
+			}
+			continue
+		}
+		break
+	}
+	if i <= j {
+		return 0, fmt.Errorf("cannot parse uint64 from %q", s)
+	}
+	if i < uint(len(s)) {
+		// Unparsed tail left.
+		return 0, fmt.Errorf("unparsed tail left after parsing uint64 from %q: %q", s, s[i:])
+	}
+	return d, nil
+}