From a32c8408bf2d971f5e0b2218dae2c8b8947a59c4 Mon Sep 17 00:00:00 2001
From: Jared Wasinger <j-wasinger@hotmail.com>
Date: Sat, 16 Dec 2023 00:07:00 +0800
Subject: [PATCH] stateless: decode witness+block

---
 cmd/stateless/main.go       |  4 +---
 core/state/state_witness.go | 42 ++++++++++++++++++++++++++++++-------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/cmd/stateless/main.go b/cmd/stateless/main.go
index 1ef49041d4..d305b820ff 100644
--- a/cmd/stateless/main.go
+++ b/cmd/stateless/main.go
@@ -55,13 +55,11 @@ func stateless(ctx *cli.Context) error {
 		panic("block witness required")
 	}
 
-	f, err := os.Open(blockWitnessPath)
+	b, err := os.ReadFile(blockWitnessPath)
 	if err != nil {
 		panic(err)
 	}
 
-	var b []byte
-	f.Read(b)
 	block, witness, err := state.DecodeWitnessRLP(b)
 	if err != nil {
 		panic(err)
diff --git a/core/state/state_witness.go b/core/state/state_witness.go
index 6705a80b52..21018e2c3a 100644
--- a/core/state/state_witness.go
+++ b/core/state/state_witness.go
@@ -40,7 +40,7 @@ type encodedWitness struct {
 }
 
 func (e *encodedWitness) ToWitness() *Witness {
-	var res Witness
+	res := NewWitness()
 	res.root = e.root
 	for i := 0; i < len(e.codes); i++ {
 		res.codes[e.codeHashes[i]] = e.codes[i]
@@ -55,13 +55,43 @@ func (e *encodedWitness) ToWitness() *Witness {
 	for i, blockNum := range e.blockNums {
 		res.blockHashes[blockNum] = e.blockHashes[i]
 	}
-	return &res
+	return res
 }
 
 func DecodeWitnessRLP(b []byte) (*types.Block, *Witness, error) {
 	var res encodedWitness
-	if err := rlp.DecodeBytes(b, &res); err != nil {
-		return nil, nil, err
+	stream := rlp.NewStream(bytes.NewBuffer(b), 1_000_000)
+	_, err := stream.List()
+	if err != nil {
+		panic(err)
+	}
+	err = res.block.DecodeRLP(stream)
+	if err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.root); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.owners); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.allPaths); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.allNodes); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.blockNums); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.blockHashes); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.codeHashes); err != nil {
+		panic(err)
+	}
+	if err = stream.Decode(&res.codes); err != nil {
+		panic(err)
 	}
 	return &res.block, res.ToWitness(), nil
 }
@@ -103,9 +133,7 @@ func (w *Witness) EncodeRLP() []byte {
 	}
 	l := eb.List()
 	w.block.EncodeRLP(eb)
-	if err := rlp.Encode(eb, root); err != nil {
-		panic(err)
-	}
+	eb.WriteBytes(root[:])
 	if err := rlp.Encode(eb, owners); err != nil {
 		panic(err)
 	}