Collection test (#3624)

* loadLeafSeqs => loadLeafCollections

* Collection tests examine leaf node counts

go/types/blob.go

@@ -56,13 +56,12 @@ func (b Blob) ReadAt(p []byte, off int64) (n int, err error) {
 		return
 	}
 
-	root := b.sequence()
-	leaves, localStart := loadLeafSequences(root.valueReader(), []sequence{root}, startIdx, endIdx)
+	leaves, localStart := loadLeafNodes([]Collection{b}, startIdx, endIdx)
 	endIdx = localStart + endIdx - startIdx
 	startIdx = localStart
 
-	for _, s := range leaves {
-		bl := s.(blobLeafSequence)
+	for _, b := range leaves {
+		bl := b.sequence().(blobLeafSequence)
 
 		localEnd := endIdx
 		leafLength := uint64(len(bl.data))

go/types/blob_test.go

@@ -95,7 +95,7 @@ func newBlobTestSuite(size uint, expectChunkCount int, expectPrependChunkDiff in
 }
 
 func TestBlobSuite4K(t *testing.T) {
-	suite.Run(t, newBlobTestSuite(12, 0, 0, 0))
+	suite.Run(t, newBlobTestSuite(12, 1, 2, 2))
 }
 
 func TestBlobSuite64K(t *testing.T) {

go/types/collection_test.go

@@ -38,12 +38,11 @@ func (suite *collectionTestSuite) TestEquals() {
 }
 
 func (suite *collectionTestSuite) TestChunkCountAndType() {
-	chunks := getChunks(suite.col)
-	suite.Equal(suite.expectChunkCount, len(chunks), "chunk count")
+	suite.Equal(suite.expectChunkCount, leafCount(suite.col), "chunk count")
 	refType := MakeRefType(suite.expectType)
-	for _, r := range chunks {
+	suite.col.WalkRefs(func(r Ref) {
 		suite.True(refType.Equals(TypeOf(r)))
-	}
+	})
 }
 
 func (suite *collectionTestSuite) TestRoundTripAndValidate() {
@@ -57,12 +56,12 @@ func (suite *collectionTestSuite) TestRoundTripAndValidate() {
 
 func (suite *collectionTestSuite) TestPrependChunkDiff() {
 	v2 := suite.prependOne()
-	suite.Equal(suite.expectPrependChunkDiff, chunkDiffCount(getChunks(suite.col), getChunks(v2)), "prepend count")
+	suite.Equal(suite.expectPrependChunkDiff, leafDiffCount(suite.col, v2), "prepend count")
 }
 
 func (suite *collectionTestSuite) TestAppendChunkDiff() {
 	v2 := suite.appendOne()
-	suite.Equal(suite.expectAppendChunkDiff, chunkDiffCount(getChunks(suite.col), getChunks(v2)), "append count")
+	suite.Equal(suite.expectAppendChunkDiff, leafDiffCount(suite.col, v2), "append count")
 }
 
 func deriveCollectionHeight(c Collection) uint64 {

go/types/indexed_sequences.go

@@ -81,22 +81,25 @@ func orderedKeyFromSum(msd []metaTuple) orderedKey {
 	return orderedKeyFromUint64(sum)
 }
 
-// loads the set of leaf sequences which contain the items [startIdx -> endIdx).
-// Returns the set of sequences and the offset within the first sequence which corresponds to |startIdx|.
-func loadLeafSequences(vr ValueReader, seqs []sequence, startIdx, endIdx uint64) ([]sequence, uint64) {
-	if seqs[0].isLeaf() {
-		for _, s := range seqs {
-			d.PanicIfFalse(s.isLeaf())
+// loads the set of leaf nodes which contain the items [startIdx -> endIdx).
+// Returns the set of nodes and the offset within the first sequence which corresponds to |startIdx|.
+func loadLeafNodes(cols []Collection, startIdx, endIdx uint64) ([]Collection, uint64) {
+	vr := cols[0].sequence().valueReader()
+
+	if cols[0].sequence().isLeaf() {
+		for _, c := range cols {
+			d.PanicIfFalse(c.sequence().isLeaf())
 		}
 
-		return seqs, startIdx
+		return cols, startIdx
 	}
 
-	level := seqs[0].treeLevel()
+	level := cols[0].sequence().treeLevel()
 	childTuples := []metaTuple{}
 
 	cum := uint64(0)
-	for _, s := range seqs {
+	for _, c := range cols {
+		s := c.sequence()
 		d.PanicIfFalse(s.treeLevel() == level)
 		ms := s.(metaSequence)
 
@@ -125,7 +128,7 @@ func loadLeafSequences(vr ValueReader, seqs []sequence, startIdx, endIdx uint64)
 	}
 
 	// Fetch committed child sequences in a single batch
-	fetched := make(map[hash.Hash]sequence, len(hs))
+	fetched := make(map[hash.Hash]Collection, len(hs))
 	if len(hs) > 0 {
 		valueChan := make(chan Value, len(hs))
 		go func() {
@@ -134,19 +137,19 @@ func loadLeafSequences(vr ValueReader, seqs []sequence, startIdx, endIdx uint64)
 			close(valueChan)
 		}()
 		for value := range valueChan {
-			fetched[value.Hash()] = value.(Collection).sequence()
+			fetched[value.Hash()] = value.(Collection)
 		}
 	}
 
-	childSeqs := make([]sequence, len(childTuples))
+	childCols := make([]Collection, len(childTuples))
 	for i, mt := range childTuples {
 		if mt.child != nil {
-			childSeqs[i] = mt.child.sequence()
+			childCols[i] = mt.child.(Collection)
 			continue
 		}
 
-		childSeqs[i] = fetched[mt.ref.TargetHash()]
+		childCols[i] = fetched[mt.ref.TargetHash()]
 	}
 
-	return loadLeafSequences(vr, childSeqs, startIdx, endIdx)
+	return loadLeafNodes(childCols, startIdx, endIdx)
 }

go/types/util_test.go

@@ -75,27 +75,38 @@ func generateNumbersAsRefOfStructs(n int) []Value {
 	return nums
 }
 
-func chunkDiffCount(c1 []Ref, c2 []Ref) int {
+func leafCount(c Collection) int {
+	leaves, _ := loadLeafNodes([]Collection{c}, 0, c.Len())
+	return len(leaves)
+}
+
+func leafDiffCount(c1, c2 Collection) int {
 	count := 0
 	hashes := make(map[hash.Hash]int)
 
-	for _, r := range c1 {
-		hashes[r.TargetHash()]++
+	leaves1, _ := loadLeafNodes([]Collection{c1}, 0, c1.Len())
+	leaves2, _ := loadLeafNodes([]Collection{c2}, 0, c2.Len())
+
+	for _, l := range leaves1 {
+		hashes[l.Hash()]++
 	}
 
-	for _, r := range c2 {
-		if c, ok := hashes[r.TargetHash()]; ok {
+	for _, l := range leaves2 {
+		if c, ok := hashes[l.Hash()]; ok {
 			if c == 1 {
-				delete(hashes, r.TargetHash())
+				delete(hashes, l.Hash())
 			} else {
-				hashes[r.TargetHash()] = c - 1
+				hashes[l.Hash()] = c - 1
 			}
 		} else {
 			count++
 		}
 	}
 
-	count += len(hashes)
+	for _, c := range hashes {
+		count += c
+	}
+
 	return count
 }