Change signature generation API to allow caller to provide storage and control iteration

tools/rsync/algorithm.go

@@ -126,12 +126,13 @@ type BlockHash struct {
 	StrongHash uint64
 }
 
-func (self BlockHash) Serialize() []byte {
-	ans := make([]byte, 20)
-	bin.PutUint64(ans, self.Index)
-	bin.PutUint32(ans[8:], self.WeakHash)
-	bin.PutUint64(ans[12:], self.StrongHash)
-	return ans
+const BlockHashSize = 20
+
+// Put the serialization of this BlockHash to output
+func (self BlockHash) Serialize(output []byte) {
+	bin.PutUint64(output, self.Index)
+	bin.PutUint32(output[8:], self.WeakHash)
+	bin.PutUint64(output[12:], self.StrongHash)
 }
 
 func (self *BlockHash) Unserialize(data []byte) (err error) {
@@ -145,7 +146,6 @@ func (self *BlockHash) Unserialize(data []byte) (err error) {
 }
 
 // Write signatures as they are generated.
-type SignatureWriter func(bl BlockHash) error
 type OperationWriter func(op Operation) error
 
 // Properties to use while working with the rsync algorithm.
@@ -176,43 +176,40 @@ func (r *rsync) BlockHashCount(targetLength int64) (count int64) {
 	return
 }
 
+type signature_iterator struct {
+	hasher hash.Hash64
+	buffer []byte
+	src    io.Reader
+	rc     rolling_checksum
+	index  uint64
+}
+
+// ans is valid only iff err == nil
+func (self *signature_iterator) next() (ans BlockHash, err error) {
+	n, err := io.ReadAtLeast(self.src, self.buffer, cap(self.buffer))
+	switch err {
+	case io.ErrUnexpectedEOF, io.EOF, nil:
+		err = nil
+	default:
+		return
+	}
+	if n == 0 {
+		return ans, io.EOF
+	}
+	b := self.buffer[:n]
+	self.hasher.Reset()
+	self.hasher.Write(b)
+	ans = BlockHash{Index: self.index, WeakHash: self.rc.full(b), StrongHash: self.hasher.Sum64()}
+	self.index++
+	return
+
+}
+
 // Calculate the signature of target.
-func (r *rsync) CreateSignature(target io.Reader, sw SignatureWriter) error {
-	var err error
-	var n int
-
-	minBufferSize := r.BlockSize
-	if len(r.buffer) < minBufferSize {
-		r.buffer = make([]byte, minBufferSize)
-	}
-	buffer := r.buffer
-
-	var block []byte
-	loop := true
-	var index uint64
-	rc := rolling_checksum{}
-	for loop {
-		n, err = io.ReadAtLeast(target, buffer, r.BlockSize)
-		if err != nil {
-			// n == 0.
-			if err == io.EOF {
-				return nil
-			}
-			if err != io.ErrUnexpectedEOF {
-				return err
-			}
-			// n > 0.
-			loop = false
-		}
-		block = buffer[:n]
-		weak := rc.full(block)
-		err = sw(BlockHash{StrongHash: r.hash(block), WeakHash: weak, Index: index})
-		if err != nil {
-			return err
-		}
-		index++
-	}
-	return nil
+func (r *rsync) CreateSignatureIterator(target io.Reader) func() (BlockHash, error) {
+	return (&signature_iterator{
+		hasher: r.hasher_constructor(), buffer: make([]byte, r.BlockSize), src: target,
+	}).next
 }
 
 // Apply the difference to the target.

tools/rsync/api.go

@@ -163,27 +163,49 @@ func (self *Patcher) FinishDelta() (err error) {
 	return
 }
 
+func ensure_size(output []byte, count int) ([]byte, []byte) {
+	if cap(output)-len(output) < count {
+		t := make([]byte, len(output), utils.Max(8192, cap(output)*2))
+		copy(t, output)
+		output = t
+	}
+	return output[:len(output)+count], output[len(output) : len(output)+count]
+}
+
+func write_block_hash(output []byte, bl BlockHash) []byte {
+	output, b := ensure_size(output, BlockHashSize)
+	bl.Serialize(b)
+	return output
+}
+
 // Create a signature for the data source in src
-func (self *Patcher) CreateSignature(src io.Reader, callback func([]byte) error) (err error) {
-	header := make([]byte, 12)
-	bin.PutUint16(header[4:], uint16(self.Strong_hash_type))
-	bin.PutUint16(header[6:], uint16(self.Weak_hash_type))
-	bin.PutUint32(header[8:], uint32(self.rsync.BlockSize))
-	if err = callback(header); err != nil {
-		return err
-	}
-	if self.expected_input_size_for_signature_generation > 0 {
-		self.signature = make([]BlockHash, 0, self.rsync.BlockHashCount(self.expected_input_size_for_signature_generation))
-	} else {
-		self.signature = make([]BlockHash, 0, 1024)
-	}
-	return self.rsync.CreateSignature(src, func(bl BlockHash) error {
-		if err = callback(bl.Serialize()); err != nil {
-			return err
+func (self *Patcher) CreateSignatureIterator(src io.Reader) func([]byte) ([]byte, error) {
+	var it func() (BlockHash, error)
+	finished := false
+	return func(output []byte) ([]byte, error) {
+		if finished {
+			return output, nil
 		}
-		self.signature = append(self.signature, bl)
-		return nil
-	})
+		var b []byte
+		if it == nil {
+			it = self.rsync.CreateSignatureIterator(src)
+			output, b = ensure_size(output, 12)
+			bin.PutUint32(b, 0)
+			bin.PutUint16(b[4:], uint16(self.Strong_hash_type))
+			bin.PutUint16(b[6:], uint16(self.Weak_hash_type))
+			bin.PutUint32(b[8:], uint32(self.rsync.BlockSize))
+		}
+		bl, err := it()
+		switch err {
+		case io.EOF:
+			finished = true
+			return output, io.EOF
+		case nil:
+			return write_block_hash(output, bl), nil
+		default:
+			return nil, err
+		}
+	}
 }
 
 type DeltaIterator = func() ([]byte, error)

tools/rsync/api_test.go

@@ -5,6 +5,7 @@ package rsync
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 	"testing"
@@ -44,10 +45,17 @@ func run_roundtrip_test(t *testing.T, src_data, changed []byte, num_of_patches,
 	// first try just the engine without serialization
 	p := NewPatcher(int64(len(src_data)))
 	signature := make([]BlockHash, 0, 128)
-	p.rsync.CreateSignature(bytes.NewReader(changed), func(s BlockHash) error {
-		signature = append(signature, s)
-		return nil
-	})
+	s_it := p.rsync.CreateSignatureIterator(bytes.NewReader(changed))
+	for {
+		s, err := s_it()
+		if err == nil {
+			signature = append(signature, s)
+		} else if err == io.EOF {
+			break
+		} else {
+			t.Fatal(err)
+		}
+	}
 
 	total_data_in_delta := 0
 	apply_delta := func(signature []BlockHash) []byte {
@@ -73,12 +81,19 @@ func run_roundtrip_test(t *testing.T, src_data, changed []byte, num_of_patches,
 	// Now try with serialization
 	using_serialization = true
 	p = NewPatcher(int64(len(src_data)))
-	signature_of_changed := bytes.Buffer{}
-	if err := p.CreateSignature(bytes.NewReader(changed), func(p []byte) error { _, err := signature_of_changed.Write(p); return err }); err != nil {
-		t.Fatal(err)
+	var signature_of_changed []byte
+	ss_it := p.CreateSignatureIterator(bytes.NewReader(changed))
+	var err error
+	for {
+		signature_of_changed, err = ss_it(signature_of_changed)
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			t.Fatal(err)
+		}
 	}
 	d := NewDiffer()
-	if err := d.AddSignatureData(signature_of_changed.Bytes()); err != nil {
+	if err := d.AddSignatureData(signature_of_changed); err != nil {
 		t.Fatal(err)
 	}
 	deltabuf := bytes.Buffer{}