feat(shwap): Add bufferisation for writing ods/Q1Q4 files

store/file/ods.go

@@ -1,6 +1,7 @@
 package file
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
@@ -108,22 +109,50 @@ func writeODSFile(w io.Writer, eds *rsmt2d.ExtendedDataSquare, axisRoots *share.
 		return nil, fmt.Errorf("writing axis roots: %w", err)
 	}
 
-	for _, shr := range eds.FlattenedODS() {
-		if _, err := w.Write(shr); err != nil {
-			return nil, fmt.Errorf("writing shares: %w", err)
-		}
+	// write quadrants
+	err = writeQuadrant(w, eds, int(h.shareSize), 0)
+	if err != nil {
+		return nil, fmt.Errorf("writing Q1: %w", err)
 	}
 	return h, nil
 }
 
+// writeQuadrant writes the quadrant of the square to the writer. it writes the quadrant in row-major
+// order. It uses buffer to write the shares in bulk (row by row), which improves the write performance.
+func writeQuadrant(w io.Writer, eds *rsmt2d.ExtendedDataSquare, shareSize, quadrantIdx int) error {
+	fromRow := quadrantIdx / 2 * int(eds.Width()) / 2
+	toRow := fromRow + int(eds.Width())/2
+
+	fromCol := quadrantIdx % 2 * int(eds.Width()) / 2
+	toCol := fromCol + int(eds.Width())/2
+
+	buf := bufio.NewWriterSize(w, shareSize*int(eds.Width()))
+	for rowIdx := fromRow; rowIdx < toRow; rowIdx++ {
+		for colIdx := fromCol; colIdx < toCol; colIdx++ {
+			share := eds.GetCell(uint(rowIdx), uint(colIdx))
+			_, err := buf.Write(share)
+			if err != nil {
+				return fmt.Errorf("writing shares: %w", err)
+			}
+		}
+	}
+	err := buf.Flush()
+	if err != nil {
+		return fmt.Errorf("flushing buffer: %w", err)
+	}
+	return nil
+}
+
 func writeAxisRoots(w io.Writer, roots *share.AxisRoots) error {
+	buf := make([]byte, 0, share.AxisRootSize*len(roots.RowRoots))
 	for _, roots := range [][][]byte{roots.RowRoots, roots.ColumnRoots} {
 		for _, root := range roots {
-			if _, err := w.Write(root); err != nil {
-				return fmt.Errorf("writing axis root: %w", err)
-			}
+			buf = append(buf, root...)
 		}
 	}
+	if _, err := w.Write(buf); err != nil {
+		return fmt.Errorf("writing axis roots: %w", err)
+	}
 	return nil
 }
 

store/file/q1q4_file.go

@@ -39,7 +39,7 @@ func CreateQ1Q4File(path string, roots *share.AxisRoots, eds *rsmt2d.ExtendedDat
 		return nil, err
 	}
 
-	err = writeQ4(ods.fl, eds)
+	err = writeQuadrant(ods.fl, eds, int(ods.hdr.shareSize), 3)
 	if err != nil {
 		return nil, fmt.Errorf("writing Q4: %w", err)
 	}
@@ -113,19 +113,6 @@ func (f *Q1Q4File) Close() error {
 	return f.ods.Close()
 }
 
-func writeQ4(w io.Writer, eds *rsmt2d.ExtendedDataSquare) error {
-	odsLn := int(eds.Width()) / 2
-	for x := odsLn; x < int(eds.Width()); x++ {
-		for y := odsLn; y < int(eds.Width()); y++ {
-			_, err := w.Write(eds.GetCell(uint(x), uint(y)))
-			if err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}
-
 func (f *Q1Q4File) readAxisHalfFromQ4(axisType rsmt2d.Axis, axisIdx int) (eds.AxisHalf, error) {
 	q4idx := axisIdx - f.ods.size()/2
 	if q4idx < 0 {

store/store_test.go

@@ -254,7 +254,7 @@ func BenchmarkStore(b *testing.B) {
 	eds := edstest.RandEDS(b, 128)
 	require.NoError(b, err)
 
-	// BenchmarkStore/bench_put_128-10         	      27	  79025268 ns/op (~79ms)
+	// BenchmarkStore/bench_put_128-10         	      27	  19209780 ns/op (~19ms)
 	b.Run("put 128", func(b *testing.B) {
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {