improve NamespacedData comment

extract switch
simplify AxisHalf

share/new_eds/nd.go

@@ -9,7 +9,8 @@ import (
 )
 
 // NamespacedData extracts shares for a specific namespace from an EDS, considering
-// each row independently.
+// each row independently. It uses root to determine which rows to extract data from,
+// avoiding the need to recalculate the row roots for each row.
 func NamespacedData(
 	ctx context.Context,
 	root *share.Root,

share/new_eds/rsmt2d.go

@@ -42,19 +42,11 @@ func (eds Rsmt2D) SampleForProofAxis(
 	rowIdx, colIdx int,
 	proofType rsmt2d.Axis,
 ) (shwap.Sample, error) {
-	var axisIdx, shrIdx int
-	switch proofType {
-	case rsmt2d.Row:
-		axisIdx, shrIdx = rowIdx, colIdx
-	case rsmt2d.Col:
-		axisIdx, shrIdx = colIdx, rowIdx
-	default:
-		return shwap.Sample{}, fmt.Errorf("invalid proof type: %d", proofType)
-	}
-	shrs := getAxis(eds.ExtendedDataSquare, proofType, axisIdx)
+	axisIdx, shrIdx := relativeIndexes(rowIdx, colIdx, proofType)
+	shares := getAxis(eds.ExtendedDataSquare, proofType, axisIdx)
 
 	tree := wrapper.NewErasuredNamespacedMerkleTree(uint64(eds.Width()/2), uint(axisIdx))
-	for _, shr := range shrs {
+	for _, shr := range shares {
 		err := tree.Push(shr)
 		if err != nil {
 			return shwap.Sample{}, fmt.Errorf("while pushing shares to NMT: %w", err)
@@ -67,16 +59,18 @@ func (eds Rsmt2D) SampleForProofAxis(
 	}
 
 	return shwap.Sample{
-		Share:     shrs[shrIdx],
+		Share:     shares[shrIdx],
 		Proof:     &prf,
 		ProofType: proofType,
 	}, nil
 }
 
 // AxisHalf returns Shares for the first half of the axis of the given type and index.
 func (eds Rsmt2D) AxisHalf(_ context.Context, axisType rsmt2d.Axis, axisIdx int) (AxisHalf, error) {
+	shares := getAxis(eds.ExtendedDataSquare, axisType, axisIdx)
+	halfShares := shares[:eds.Size()/2]
 	return AxisHalf{
-		Shares:   getAxis(eds.ExtendedDataSquare, axisType, axisIdx)[:eds.Size()/2],
+		Shares:   halfShares,
 		IsParity: false,
 	}, nil
 }
@@ -98,9 +92,9 @@ func (eds Rsmt2D) RowNamespaceData(
 	return shwap.RowNamespaceDataFromShares(shares, namespace, rowIdx)
 }
 
-// Flattened returns data shares extracted from the EDS.
+// Flattened returns data shares extracted from the EDS. It returns new copy of the shares each
+// time.
 func (eds Rsmt2D) Flattened(_ context.Context) ([]share.Share, error) {
-	// Flattened returns copy of shares, which is not efficient.
 	return eds.ExtendedDataSquare.Flattened(), nil
 }
 
@@ -114,3 +108,14 @@ func getAxis(eds *rsmt2d.ExtendedDataSquare, axisType rsmt2d.Axis, axisIdx int)
 		panic("unknown axis")
 	}
 }
+
+func relativeIndexes(rowIdx, colIdx int, axisType rsmt2d.Axis) (axisIdx, shrIdx int) {
+	switch axisType {
+	case rsmt2d.Row:
+		return rowIdx, colIdx
+	case rsmt2d.Col:
+		return colIdx, rowIdx
+	default:
+		panic(fmt.Sprintf("invalid proof type: %d", axisType))
+	}
+}