Use pull-style iterators to build index and tree

ca/ca.go

@@ -739,19 +739,44 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 		}
 	}
 
+	// Prepare index builder
+	indexPath := gopath.Join(dir, "index")
+	indexW, err := os.Create(indexPath)
+	if err != nil {
+		return fmt.Errorf("creating %s: %w", indexPath, err)
+	}
+
+	defer indexW.Close()
+	ib := internal.NewIndexBuilder(indexW)
+
+	// Prepare tree builder
+	tb := batch.NewTreeBuilder()
+
 	if !empty {
+		entryOffset := 0
+		evidenceOffset := 0
+		var entryKey [mtc.HashLen]byte
+
 		err = h.walkQueue(func(ar mtc.AssertionRequest) error {
+			oldEvidenceOffset := evidenceOffset
+			oldEntryOffset := entryOffset
+
 			// Skip assertions that are already expired.
 			if start, _ := batch.ValidityInterval(); ar.NotAfter.Before(start) {
 				return nil
 			}
 
 			be := mtc.NewBatchEntry(ar.Assertion, ar.NotAfter)
+			if err := be.Key(entryKey[:]); err != nil {
+				return fmt.Errorf("Computing key for %x: %w", ar.Checksum, err)
+			}
+
 			buf, err := be.MarshalBinary()
 			if err != nil {
 				return fmt.Errorf("Marshalling assertion %x: %w", ar.Checksum, err)
 			}
 
+			// Write out BatchEntry
 			_, err = besBW.Write(buf)
 			if err != nil {
 				return fmt.Errorf(
@@ -761,7 +786,10 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 					err,
 				)
 			}
+			entryOffset += len(buf)
 
+			// Prepare evidence when applicable: for instance by  deduplicating
+			// intermediates in umbilical chains.
 			evs := ar.Evidence
 			if ucBuilder != nil {
 				for i := range len(evs) {
@@ -777,6 +805,7 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 				}
 			}
 
+			// Write out Evidence
 			buf, err = evs.MarshalBinary()
 			if err != nil {
 				return fmt.Errorf("Marshalling evidence %x: %w", ar.Checksum, err)
@@ -791,6 +820,21 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 					err,
 				)
 			}
+			evidenceOffset += len(buf)
+
+			// Feed entry to tree builder, and entry and evidence to
+			// index builder.
+			if err := tb.Push(&be); err != nil {
+				return fmt.Errorf("Building tree: %w", err)
+			}
+
+			if err := ib.Push(internal.IndexBuildEntry{
+				EvidenceOffset: uint64(oldEvidenceOffset),
+				Offset:         uint64(oldEntryOffset),
+				Key:            entryKey,
+			}); err != nil {
+				return fmt.Errorf("Building index: %w", err)
+			}
 
 			return nil
 		})
@@ -808,11 +852,6 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 	if err != nil {
 		return fmt.Errorf("closing %s: %w", besPath, err)
 	}
-	besR, err := os.OpenFile(besPath, os.O_RDONLY, 0)
-	if err != nil {
-		return fmt.Errorf("opening %s: %w", besPath, err)
-	}
-	defer besR.Close()
 
 	err = evBW.Flush()
 	if err != nil {
@@ -823,11 +862,6 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 	if err != nil {
 		return fmt.Errorf("closing %s: %w", evPath, err)
 	}
-	evR, err := os.OpenFile(evPath, os.O_RDONLY, 0)
-	if err != nil {
-		return fmt.Errorf("opening %s: %w", evPath, err)
-	}
-	defer evR.Close()
 
 	if ucBuilder != nil {
 		err = ucBuilder.Finish()
@@ -841,9 +875,9 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 	}
 
 	// Compute tree
-	tree, err := batch.ComputeTree(bufio.NewReader(besR))
+	tree, err := tb.Finish()
 	if err != nil {
-		return fmt.Errorf("computing tree: %w", err)
+		return fmt.Errorf("finishing tree: %w", err)
 	}
 
 	treePath := gopath.Join(dir, "tree")
@@ -865,22 +899,9 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 	}
 
 	// Compute index
-	_, err = besR.Seek(0, io.SeekStart)
-	if err != nil {
-		return fmt.Errorf("seeking %s to start: %w", besPath, err)
-	}
-
-	indexPath := gopath.Join(dir, "index")
-	indexW, err := os.Create(indexPath)
-	if err != nil {
-		return fmt.Errorf("creating %s: %w", indexPath, err)
-	}
-
-	defer indexW.Close()
-
-	err = internal.ComputeIndex(besR, evR, indexW)
+	err = ib.Finish()
 	if err != nil {
-		return fmt.Errorf("computing %s to start: %w", indexPath, err)
+		return fmt.Errorf("finishing index: %w", err)
 	}
 
 	// Sign validity window
@@ -899,6 +920,7 @@ func (h *Handle) issueBatchTo(dir string, batch mtc.Batch, empty bool) error {
 	if err != nil {
 		return fmt.Errorf("writing to %s: %w", wPath, err)
 	}
+
 	return nil
 }
 

cmd/mtc/main.go

@@ -961,9 +961,10 @@ func handleInspectEvidence(cc *cli.Context) error {
 	defer r.Close()
 
 	count := 0
-	err = mtc.UnmarshalEvidenceLists(
-		bufio.NewReader(r),
-		func(_ int, el *mtc.EvidenceList) error {
+
+	err = mtc.ForEach(
+		mtc.UnmarshalEvidenceLists(bufio.NewReader(r)),
+		func(el *mtc.EvidenceList) error {
 			count++
 			w := tabwriter.NewWriter(os.Stdout, 1, 1, 1, ' ', 0)
 			err := writeEvidenceList(w, *el)
@@ -990,9 +991,9 @@ func handleInspectEntries(cc *cli.Context) error {
 	defer r.Close()
 
 	count := 0
-	err = mtc.UnmarshalBatchEntries(
-		bufio.NewReader(r),
-		func(_ int, be *mtc.BatchEntry) error {
+	err = mtc.ForEach(
+		mtc.UnmarshalBatchEntries(bufio.NewReader(r)),
+		func(be *mtc.BatchEntry) error {
 			count++
 			cs := be.Claims
 			subj := be.Subject

internal/index.go

@@ -44,6 +44,12 @@ type IndexSearchResult struct {
 	EvidenceOffset uint64
 }
 
+type IndexBuildEntry struct {
+	Key            [mtc.HashLen]byte
+	Offset         uint64
+	EvidenceOffset uint64
+}
+
 // Opens an index
 func OpenIndex(path string) (*Index, error) {
 	r, err := mmap.Open(path)
@@ -163,64 +169,60 @@ type indexEntry struct {
 	evidenceOffset uint64
 }
 
-// Reads a streams of BatchEntry Evidence from beReader and evReader,
-// and writes the index to w.
-func ComputeIndex(beReader, evReader io.Reader, w io.Writer) error {
-	// First compute keys
-	seqno := uint64(0)
-	entries := []indexEntry{}
-
-	var key [mtc.HashLen]byte
-	err := mtc.UnmarshalBatchEntries(beReader, func(offset int,
-		be *mtc.BatchEntry) error {
-		err := be.Key(key[:])
-		if err != nil {
-			return err
-		}
-		entries = append(entries, indexEntry{
-			seqno:  seqno,
-			key:    key,
-			offset: uint64(offset),
-		})
-		seqno++
-		return nil
-	})
+type IndexBuilder struct {
+	err     error
+	w       io.Writer
+	seqno   uint64
+	entries []indexEntry
+}
 
-	seqno = uint64(0)
-	err = mtc.UnmarshalEvidenceLists(evReader, func(offset int, _ *mtc.EvidenceList) error {
-		entries[seqno].evidenceOffset = uint64(offset)
-		seqno++
-		return nil
-	})
+func NewIndexBuilder(w io.Writer) *IndexBuilder {
+	return &IndexBuilder{
+		w:       w,
+		entries: []indexEntry{},
+	}
+}
 
-	if err != nil {
-		return fmt.Errorf("computing keys: %w", err)
+func (b *IndexBuilder) Push(in IndexBuildEntry) error {
+	if b.err != nil {
+		return b.err
 	}
 
+	b.entries = append(b.entries, indexEntry{
+		seqno:          b.seqno,
+		key:            in.Key,
+		offset:         in.Offset,
+		evidenceOffset: in.EvidenceOffset,
+	})
+
+	b.seqno++
+	return nil
+}
+
+func (b *IndexBuilder) Finish() error {
 	// Sort by key
-	slices.SortFunc(entries, func(a, b indexEntry) int {
+	slices.SortFunc(b.entries, func(a, b indexEntry) int {
 		return bytes.Compare(a.key[:], b.key[:])
 	})
 
 	// Write out
-	bw := bufio.NewWriter(w)
+	bw := bufio.NewWriter(b.w)
 	var lastKey [mtc.HashLen]byte
-	for _, entry := range entries {
+	for _, entry := range b.entries {
 		if lastKey == entry.key {
 			// skip duplicate entries
 			continue
 		}
 
 		lastKey = entry.key
-		var b cryptobyte.Builder
-		b.AddBytes(entry.key[:])
-		b.AddUint64(entry.seqno)
-		b.AddUint64(entry.offset)
-		b.AddUint64(entry.evidenceOffset)
-		buf, _ := b.Bytes()
-
-		_, err = bw.Write(buf)
-		if err != nil {
+		var cb cryptobyte.Builder
+		cb.AddBytes(entry.key[:])
+		cb.AddUint64(entry.seqno)
+		cb.AddUint64(entry.offset)
+		cb.AddUint64(entry.evidenceOffset)
+		buf, _ := cb.Bytes()
+
+		if _, err := bw.Write(buf); err != nil {
 			return fmt.Errorf("writing index: %w", err)
 		}
 	}

mirror/mirror.go

@@ -303,15 +303,35 @@ func (h *Handle) fetchBatch(number uint32) error {
 	}
 
 	// Recompute tree
-	aasR, err := os.OpenFile(besPath, os.O_RDONLY, 0)
+	tb := batch.NewTreeBuilder()
+	besR, err := os.OpenFile(besPath, os.O_RDONLY, 0)
 	if err != nil {
 		return fmt.Errorf("opening %s: %w", besPath, err)
 	}
-	defer aasR.Close()
+	defer besR.Close()
 
-	tree, err := batch.ComputeTree(bufio.NewReader(aasR))
+	besC := mtc.UnmarshalBatchEntries(bufio.NewReader(besR))
+	defer besC.Close()
+
+	var be mtc.BatchEntry
+	for {
+		err := besC.Pull(&be)
+		if err == mtc.EOF {
+			break
+		}
+
+		if err != nil {
+			return fmt.Errorf("reading %s: %w", besPath, err)
+		}
+
+		if err := tb.Push(&be); err != nil {
+			return fmt.Errorf("building tree: %w", err)
+		}
+	}
+
+	tree, err := tb.Finish()
 	if err != nil {
-		return fmt.Errorf("computing tree: %w", err)
+		return fmt.Errorf("finishing tree: %w", err)
 	}
 
 	treePath := gopath.Join(dir, "tree")