Add bsoncore.DocumentSequence for BatchCursor

Add the bsoncore.DocumentSequence type and use it in BatchCursor.

GODRIVER-800

Change-Id: I3b95cce0e7693559d9b79541156622cb1ab3b977

Author: skriptble

mongo/batch_cursor.go

@@ -2,6 +2,8 @@ package mongo
 
 import (
 	"context"
+
+	"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
 )
 
 // batchCursor is the interface implemented by types that can provide batches of document results.
@@ -13,14 +15,9 @@ type batchCursor interface {
 	// Next returns true if there is a batch available.
 	Next(context.Context) bool
 
-	// Batch appends the current batch of documents to dst. RequiredBytes can be used to determine
-	// the length of the current batch of documents.
-	//
-	// If there is no batch available, this method should do nothing.
-	Batch(dst []byte) []byte
-
-	// RequiredBytes returns the number of bytes required fo rthe current batch.
-	RequiredBytes() int
+	// Batch will return a DocumentSequence for the current batch of documents. The returned
+	// DocumentSequence is only valid until the next call to Next or Close.
+	Batch() *bsoncore.DocumentSequence
 
 	// Err returns the last error encountered.
 	Err() error

mongo/cursor.go

@@ -9,9 +9,11 @@ package mongo
 import (
 	"context"
 	"errors"
+	"io"
 
 	"github.com/mongodb/mongo-go-driver/bson"
 	"github.com/mongodb/mongo-go-driver/bson/bsoncodec"
+	"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
 	"github.com/mongodb/mongo-go-driver/x/mongo/driver"
 )
 
@@ -44,8 +46,7 @@ type Cursor struct {
 	Current bson.Raw
 
 	bc       batchCursor
-	pos      int
-	batch    []byte
+	batch    *bsoncore.DocumentSequence
 	registry *bsoncodec.Registry
 
 	err error
@@ -58,7 +59,7 @@ func newCursor(bc batchCursor, registry *bsoncodec.Registry) (*Cursor, error) {
 	if bc == nil {
 		return nil, errors.New("batch cursor must not be nil")
 	}
-	return &Cursor{bc: bc, pos: 0, batch: make([]byte, 0, 256), registry: registry}, nil
+	return &Cursor{bc: bc, registry: registry}, nil
 }
 
 func newEmptyCursor() *Cursor {
@@ -68,42 +69,26 @@ func newEmptyCursor() *Cursor {
 // ID returns the ID of this cursor.
 func (c *Cursor) ID() int64 { return c.bc.ID() }
 
-func (c *Cursor) advanceCurrentDocument() bool {
-	if len(c.batch[c.pos:]) < 4 {
-		c.err = errors.New("could not read next document: insufficient bytes")
-		return false
-	}
-	length := (int(c.batch[c.pos]) | int(c.batch[c.pos+1])<<8 | int(c.batch[c.pos+2])<<16 | int(c.batch[c.pos+3])<<24)
-	if len(c.batch[c.pos:]) < length {
-		c.err = errors.New("could not read next document: insufficient bytes")
-		return false
-	}
-	if len(c.Current) > 4 {
-		c.Current[0], c.Current[1], c.Current[2], c.Current[3] = 0x00, 0x00, 0x00, 0x00 // Invalidate the current document
-	}
-	c.Current = c.batch[c.pos : c.pos+length]
-	c.pos += length
-	return true
-}
-
 // Next gets the next result from this cursor. Returns true if there were no errors and the next
 // result is available for decoding.
 func (c *Cursor) Next(ctx context.Context) bool {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	if c.pos < len(c.batch) {
-		return c.advanceCurrentDocument()
+	doc, err := c.batch.Next()
+	switch err {
+	case nil:
+		c.Current = bson.Raw(doc)
+		return true
+	case io.EOF: // Need to do a getMore
+	default:
+		c.err = err
+		return false
 	}
 
-	// clear the batch
-	c.batch = c.batch[:0]
-	c.pos = 0
-	c.Current = c.Current[:0]
-
 	// call the Next method in a loop until at least one document is returned in the next batch or
 	// the context times out.
-	for len(c.batch) == 0 {
+	for {
 		// If we don't have a next batch
 		if !c.bc.Next(ctx) {
 			// Do we have an error? If so we return false.
@@ -119,10 +104,18 @@ func (c *Cursor) Next(ctx context.Context) bool {
 			continue
 		}
 
-		c.batch = c.bc.Batch(c.batch[:0])
+		c.batch = c.bc.Batch()
+		doc, err = c.batch.Next()
+		switch err {
+		case nil:
+			c.Current = bson.Raw(doc)
+			return true
+		case io.EOF: // Empty batch so we continue
+		default:
+			c.err = err
+			return false
+		}
 	}
-
-	return c.advanceCurrentDocument()
 }
 
 // Decode will decode the current document into val.

x/bsonx/bsoncore/document_sequence.go

@@ -0,0 +1,167 @@
+package bsoncore
+
+import (
+	"errors"
+	"io"
+
+	"github.com/mongodb/mongo-go-driver/bson/bsontype"
+)
+
+// DocumentSequenceStyle is used to represent how a document sequence is laid out in a slice of
+// bytes.
+type DocumentSequenceStyle uint32
+
+// These constants are the valid styles for a DocumentSequence.
+const (
+	_ DocumentSequenceStyle = iota
+	SequenceStyle
+	ArrayStyle
+)
+
+// DocumentSequence represents a sequence of documents. The Style field indicates how the documents
+// are laid out inside of the Data field.
+type DocumentSequence struct {
+	Style DocumentSequenceStyle
+	Data  []byte
+	Pos   int
+}
+
+// ErrCorruptedDocument is returned when a full document couldn't be read from the sequence.
+var ErrCorruptedDocument = errors.New("invalid DocumentSequence: corrupted document")
+
+// ErrNonDocument is returned when a DocumentSequence contains a non-document BSON value.
+var ErrNonDocument = errors.New("invalid DocumentSequence: a non-document value was found in sequence")
+
+// ErrInvalidDocumentSequenceStyle is returned when an unknown DocumentSequenceStyle is set on a
+// DocumentSequence.
+var ErrInvalidDocumentSequenceStyle = errors.New("invalid DocumentSequenceStyle")
+
+// DocumentCount returns the number of documents in the sequence.
+func (ds *DocumentSequence) DocumentCount() int {
+	if ds == nil {
+		return 0
+	}
+	switch ds.Style {
+	case SequenceStyle:
+		var count int
+		var ok bool
+		rem := ds.Data
+		for len(rem) > 0 {
+			_, rem, ok = ReadDocument(rem)
+			if !ok {
+				return 0
+			}
+			count++
+		}
+		return count
+	case ArrayStyle:
+		_, rem, ok := ReadLength(ds.Data)
+		if !ok {
+			return 0
+		}
+
+		var count int
+		for len(rem) > 1 {
+			_, rem, ok = ReadElement(rem)
+			if !ok {
+				return 0
+			}
+			count++
+		}
+		return count
+	default:
+		return 0
+	}
+}
+
+//ResetIterator resets the iteration point for the Next method to the beginning of the document
+//sequence.
+func (ds *DocumentSequence) ResetIterator() {
+	if ds == nil {
+		return
+	}
+	ds.Pos = 0
+}
+
+// documents returns a slice of the documents. If nil either the Data field is also nil or could not
+// be properly read.
+func (ds *DocumentSequence) documents() ([]Document, error) {
+	if ds == nil {
+		return nil, nil
+	}
+	switch ds.Style {
+	case SequenceStyle:
+		rem := ds.Data
+		var docs []Document
+		var doc Document
+		var ok bool
+		for {
+			doc, rem, ok = ReadDocument(rem)
+			if !ok {
+				if len(rem) == 0 {
+					break
+				}
+				return nil, ErrCorruptedDocument
+			}
+			docs = append(docs, doc)
+		}
+		return docs, nil
+	case ArrayStyle:
+		if len(ds.Data) == 0 {
+			return nil, nil
+		}
+		vals, err := Document(ds.Data).Values()
+		if err != nil {
+			return nil, ErrCorruptedDocument
+		}
+		docs := make([]Document, 0, len(vals))
+		for _, v := range vals {
+			if v.Type != bsontype.EmbeddedDocument {
+				return nil, ErrNonDocument
+			}
+			docs = append(docs, v.Data)
+		}
+		return docs, nil
+	default:
+		return nil, ErrInvalidDocumentSequenceStyle
+	}
+}
+
+// Next retrieves the next document from this sequence and returns it. This method will return
+// io.EOF when it has reached the end of the sequence.
+func (ds *DocumentSequence) Next() (Document, error) {
+	if ds == nil || ds.Pos >= len(ds.Data) {
+		return nil, io.EOF
+	}
+	switch ds.Style {
+	case SequenceStyle:
+		doc, _, ok := ReadDocument(ds.Data[ds.Pos:])
+		if !ok {
+			return nil, ErrCorruptedDocument
+		}
+		ds.Pos += len(doc)
+		return doc, nil
+	case ArrayStyle:
+		if ds.Pos < 4 {
+			if len(ds.Data) < 4 {
+				return nil, ErrCorruptedDocument
+			}
+			ds.Pos = 4 // Skip the length of the document
+		}
+		if len(ds.Data[ds.Pos:]) == 1 && ds.Data[ds.Pos] == 0x00 {
+			return nil, io.EOF // At the end of the document
+		}
+		elem, _, ok := ReadElement(ds.Data[ds.Pos:])
+		if !ok {
+			return nil, ErrCorruptedDocument
+		}
+		ds.Pos += len(elem)
+		val := elem.Value()
+		if val.Type != bsontype.EmbeddedDocument {
+			return nil, ErrNonDocument
+		}
+		return val.Data, nil
+	default:
+		return nil, ErrInvalidDocumentSequenceStyle
+	}
+}