From 162f2bad346238488d8652bdf46292cb947b19b7 Mon Sep 17 00:00:00 2001
From: Vlad <13818348+walldiss@users.noreply.github.com>
Date: Fri, 5 Jan 2024 20:38:45 +0700
Subject: [PATCH] Allow reconstructSome to reconstruct parity shards (#267)

* Allow reconstructSome to reconstruct parity shards
---
 leopard.go          |  5 ++++-
 leopard8.go         |  5 ++++-
 reedsolomon.go      | 19 ++++++++++++-------
 reedsolomon_test.go | 22 ++++++++++++++++++++--
 4 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/leopard.go b/leopard.go
index b55eddd5..6b4c8018 100644
--- a/leopard.go
+++ b/leopard.go
@@ -333,7 +333,10 @@ func (r *leopardFF16) Split(data []byte) ([][]byte, error) {
 }
 
 func (r *leopardFF16) ReconstructSome(shards [][]byte, required []bool) error {
-	return r.ReconstructData(shards)
+	if len(required) == r.totalShards {
+		return r.reconstruct(shards, true)
+	}
+	return r.reconstruct(shards, false)
 }
 
 func (r *leopardFF16) Reconstruct(shards [][]byte) error {
diff --git a/leopard8.go b/leopard8.go
index ae297672..cd863a13 100644
--- a/leopard8.go
+++ b/leopard8.go
@@ -369,7 +369,10 @@ func (r *leopardFF8) Split(data []byte) ([][]byte, error) {
 }
 
 func (r *leopardFF8) ReconstructSome(shards [][]byte, required []bool) error {
-	return r.ReconstructData(shards)
+	if len(required) == r.totalShards {
+		return r.reconstruct(shards, true)
+	}
+	return r.reconstruct(shards, false)
 }
 
 func (r *leopardFF8) Reconstruct(shards [][]byte) error {
diff --git a/reedsolomon.go b/reedsolomon.go
index bd82015f..55b56650 100644
--- a/reedsolomon.go
+++ b/reedsolomon.go
@@ -77,11 +77,12 @@ type Encoder interface {
 	// calling the Verify function is likely to fail.
 	ReconstructData(shards [][]byte) error
 
-	// ReconstructSome will recreate only requested data shards, if possible.
+	// ReconstructSome will recreate only requested shards, if possible.
 	//
 	// Given a list of shards, some of which contain data, fills in the
-	// data shards indicated by true values in the "required" parameter.
-	// The length of "required" array must be equal to DataShards.
+	// shards indicated by true values in the "required" parameter.
+	// The length of the "required" array must be equal to either Shards or DataShards.
+	// If the length is equal to DataShards, the reconstruction of parity shards will be ignored.
 	//
 	// The length of "shards" array must be equal to Shards.
 	// You indicate that a shard is missing by setting it to nil or zero-length.
@@ -1402,13 +1403,14 @@ func (r *reedSolomon) ReconstructData(shards [][]byte) error {
 	return r.reconstruct(shards, true, nil)
 }
 
-// ReconstructSome will recreate only requested data shards, if possible.
+// ReconstructSome will recreate only requested shards, if possible.
 //
 // Given a list of shards, some of which contain data, fills in the
-// data shards indicated by true values in the "required" parameter.
-// The length of "required" array must be equal to dataShards.
+// shards indicated by true values in the "required" parameter.
+// The length of the "required" array must be equal to either Shards or DataShards.
+// If the length is equal to DataShards, the reconstruction of parity shards will be ignored.
 //
-// The length of "shards" array must be equal to shards.
+// The length of "shards" array must be equal to Shards.
 // You indicate that a shard is missing by setting it to nil or zero-length.
 // If a shard is zero-length but has sufficient capacity, that memory will
 // be used, otherwise a new []byte will be allocated.
@@ -1419,6 +1421,9 @@ func (r *reedSolomon) ReconstructData(shards [][]byte) error {
 // As the reconstructed shard set may contain missing parity shards,
 // calling the Verify function is likely to fail.
 func (r *reedSolomon) ReconstructSome(shards [][]byte, required []bool) error {
+	if len(required) == r.totalShards {
+		return r.reconstruct(shards, false, required)
+	}
 	return r.reconstruct(shards, true, required)
 }
 
diff --git a/reedsolomon_test.go b/reedsolomon_test.go
index 2932787c..a1745a35 100644
--- a/reedsolomon_test.go
+++ b/reedsolomon_test.go
@@ -826,8 +826,26 @@ func testReconstructData(t *testing.T, o ...Option) {
 		t.Fatal(err)
 	}
 
-	// Reconstruct 3 shards with 3 data and 5 parity shards
+	// Reconstruct parity shards from data
 	shardsCopy := make([][]byte, 13)
+	for i := 0; i < 8; i++ {
+		shardsCopy[i] = shards[i]
+	}
+
+	shardsRequired := make([]bool, 13)
+	shardsRequired[10] = true
+
+	err = r.ReconstructSome(shardsCopy, shardsRequired)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(shardsCopy[10], shards[10]) {
+		t.Fatal("ReconstructSome did not reconstruct required shards correctly")
+	}
+
+	// Reconstruct 3 shards with 3 data and 5 parity shards
+	shardsCopy = make([][]byte, 13)
 	copy(shardsCopy, shards)
 	shardsCopy[2] = nil
 	shardsCopy[3] = nil
@@ -835,7 +853,7 @@ func testReconstructData(t *testing.T, o ...Option) {
 	shardsCopy[5] = nil
 	shardsCopy[6] = nil
 
-	shardsRequired := make([]bool, 8)
+	shardsRequired = make([]bool, 8)
 	shardsRequired[3] = true
 	shardsRequired[4] = true
 	err = r.ReconstructSome(shardsCopy, shardsRequired)