GODRIVER-663 Extend the Marshal* and Unmarshal* family functions with *WithContext functions

Change-Id: I261c3276a78b897039d8e9629b6eab22d608eedf

Author: Isabella Siu

bson/decoder.go

@@ -28,21 +28,33 @@ var decPool = sync.Pool{
 // A Decoder reads and decodes BSON documents from a stream. It reads from a bsonrw.ValueReader as
 // the source of BSON data.
 type Decoder struct {
-	r  *bsoncodec.Registry
+	dc bsoncodec.DecodeContext
 	vr bsonrw.ValueReader
 }
 
-// NewDecoder returns a new decoder that uses Registry reg to read from r.
-func NewDecoder(r *bsoncodec.Registry, vr bsonrw.ValueReader) (*Decoder, error) {
-	if r == nil {
-		return nil, errors.New("cannot create a new Decoder with a nil Registry")
+// NewDecoder returns a new decoder that uses the DefaultRegistry to read from vr.
+func NewDecoder(vr bsonrw.ValueReader) (*Decoder, error) {
+	if vr == nil {
+		return nil, errors.New("cannot create a new Decoder with a nil ValueReader")
+	}
+
+	return &Decoder{
+		dc: bsoncodec.DecodeContext{Registry: DefaultRegistry},
+		vr: vr,
+	}, nil
+}
+
+// NewDecoderWithContext returns a new decoder that uses DecodeContext dc to read from vr.
+func NewDecoderWithContext(dc bsoncodec.DecodeContext, vr bsonrw.ValueReader) (*Decoder, error) {
+	if dc.Registry == nil {
+		dc.Registry = DefaultRegistry
 	}
 	if vr == nil {
 		return nil, errors.New("cannot create a new Decoder with a nil ValueReader")
 	}
 
 	return &Decoder{
-		r:  r,
+		dc: dc,
 		vr: vr,
 	}, nil
 }
@@ -67,22 +79,28 @@ func (d *Decoder) Decode(val interface{}) error {
 		return fmt.Errorf("argument to Decode must be a pointer to a type, but got %v", rval)
 	}
 	rval = rval.Elem()
-	decoder, err := d.r.LookupDecoder(rval.Type())
+	decoder, err := d.dc.LookupDecoder(rval.Type())
 	if err != nil {
 		return err
 	}
-	return decoder.DecodeValue(bsoncodec.DecodeContext{Registry: d.r}, d.vr, rval)
+	return decoder.DecodeValue(d.dc, d.vr, rval)
 }
 
-// Reset will reset the state of the decoder, using the same *Registry used in
-// the original construction but using r for reading.
+// Reset will reset the state of the decoder, using the same *DecodeContext used in
+// the original construction but using vr for reading.
 func (d *Decoder) Reset(vr bsonrw.ValueReader) error {
 	d.vr = vr
 	return nil
 }
 
 // SetRegistry replaces the current registry of the decoder with r.
 func (d *Decoder) SetRegistry(r *bsoncodec.Registry) error {
-	d.r = r
+	d.dc.Registry = r
+	return nil
+}
+
+// SetContext replaces the current registry of the decoder with dc.
+func (d *Decoder) SetContext(dc bsoncodec.DecodeContext) error {
+	d.dc = dc
 	return nil
 }

bson/decoder_test.go

@@ -51,7 +51,7 @@ func TestDecoderv2(t *testing.T) {
 				} else {
 					reg = DefaultRegistry
 				}
-				dec, err := NewDecoder(reg, vr)
+				dec, err := NewDecoderWithContext(bsoncodec.DecodeContext{Registry: reg}, vr)
 				noerr(t, err)
 				err = dec.Decode(got)
 				noerr(t, err)
@@ -64,7 +64,7 @@ func TestDecoderv2(t *testing.T) {
 		t.Run("lookup error", func(t *testing.T) {
 			type certainlydoesntexistelsewhereihope func(string, string) string
 			cdeih := func(string, string) string { return "certainlydoesntexistelsewhereihope" }
-			dec, err := NewDecoder(DefaultRegistry, bsonrw.NewBSONDocumentReader([]byte{}))
+			dec, err := NewDecoder(bsonrw.NewBSONDocumentReader([]byte{}))
 			noerr(t, err)
 			want := bsoncodec.ErrNoDecoder{Type: reflect.TypeOf(cdeih)}
 			got := dec.Decode(&cdeih)
@@ -102,7 +102,7 @@ func TestDecoderv2(t *testing.T) {
 			for _, tc := range testCases {
 				t.Run(tc.name, func(t *testing.T) {
 					unmarshaler := &testUnmarshaler{err: tc.err}
-					dec, err := NewDecoder(DefaultRegistry, tc.vr)
+					dec, err := NewDecoder(tc.vr)
 					noerr(t, err)
 					got := dec.Decode(unmarshaler)
 					want := tc.err
@@ -123,7 +123,7 @@ func TestDecoderv2(t *testing.T) {
 				want := bsoncore.BuildDocument(nil, bsoncore.AppendDoubleElement(nil, "pi", 3.14159))
 				unmarshaler := &testUnmarshaler{}
 				vr := bsonrw.NewBSONDocumentReader(want)
-				dec, err := NewDecoder(DefaultRegistry, vr)
+				dec, err := NewDecoder(vr)
 				noerr(t, err)
 				err = dec.Decode(unmarshaler)
 				noerr(t, err)
@@ -134,21 +134,39 @@ func TestDecoderv2(t *testing.T) {
 			})
 		})
 	})
-	t.Run("NewDecoderv2", func(t *testing.T) {
-		t.Run("errors", func(t *testing.T) {
-			_, got := NewDecoder(nil, bsonrw.ValueReader(nil))
-			want := errors.New("cannot create a new Decoder with a nil Registry")
+	t.Run("NewDecoder", func(t *testing.T) {
+		t.Run("error", func(t *testing.T) {
+			_, got := NewDecoder(nil)
+			want := errors.New("cannot create a new Decoder with a nil ValueReader")
 			if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 				t.Errorf("Was expecting error but got different error. got %v; want %v", got, want)
 			}
-			_, got = NewDecoder(DefaultRegistry, nil)
-			want = errors.New("cannot create a new Decoder with a nil ValueReader")
+		})
+		t.Run("success", func(t *testing.T) {
+			got, err := NewDecoder(bsonrw.NewBSONDocumentReader([]byte{}))
+			noerr(t, err)
+			if got == nil {
+				t.Errorf("Was expecting a non-nil Decoder, but got <nil>")
+			}
+		})
+	})
+	t.Run("NewDecoderWithContext", func(t *testing.T) {
+		t.Run("errors", func(t *testing.T) {
+			dc := bsoncodec.DecodeContext{Registry: DefaultRegistry}
+			_, got := NewDecoderWithContext(dc, nil)
+			want := errors.New("cannot create a new Decoder with a nil ValueReader")
 			if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
 				t.Errorf("Was expecting error but got different error. got %v; want %v", got, want)
 			}
 		})
 		t.Run("success", func(t *testing.T) {
-			got, err := NewDecoder(DefaultRegistry, bsonrw.NewBSONDocumentReader([]byte{}))
+			got, err := NewDecoderWithContext(bsoncodec.DecodeContext{}, bsonrw.NewBSONDocumentReader([]byte{}))
+			noerr(t, err)
+			if got == nil {
+				t.Errorf("Was expecting a non-nil Decoder, but got <nil>")
+			}
+			dc := bsoncodec.DecodeContext{Registry: DefaultRegistry}
+			got, err = NewDecoderWithContext(dc, bsonrw.NewBSONDocumentReader([]byte{}))
 			noerr(t, err)
 			if got == nil {
 				t.Errorf("Was expecting a non-nil Decoder, but got <nil>")
@@ -166,7 +184,7 @@ func TestDecoderv2(t *testing.T) {
 		got.Bonus = 2
 		data := docToBytes(bsonx.Doc{{"item", bsonx.String("canvas")}, {"qty", bsonx.Int32(4)}})
 		vr := bsonrw.NewBSONDocumentReader(data)
-		dec, err := NewDecoder(DefaultRegistry, vr)
+		dec, err := NewDecoder(vr)
 		noerr(t, err)
 		err = dec.Decode(&got)
 		noerr(t, err)
@@ -177,7 +195,8 @@ func TestDecoderv2(t *testing.T) {
 	})
 	t.Run("Reset", func(t *testing.T) {
 		vr1, vr2 := bsonrw.NewBSONDocumentReader([]byte{}), bsonrw.NewBSONDocumentReader([]byte{})
-		dec, err := NewDecoder(DefaultRegistry, vr1)
+		dc := bsoncodec.DecodeContext{Registry: DefaultRegistry}
+		dec, err := NewDecoderWithContext(dc, vr1)
 		noerr(t, err)
 		if dec.vr != vr1 {
 			t.Errorf("Decoder should use the value reader provided. got %v; want %v", dec.vr, vr1)
@@ -188,17 +207,33 @@ func TestDecoderv2(t *testing.T) {
 			t.Errorf("Decoder should use the value reader provided. got %v; want %v", dec.vr, vr2)
 		}
 	})
+	t.Run("SetContext", func(t *testing.T) {
+		dc1 := bsoncodec.DecodeContext{Registry: DefaultRegistry}
+		dc2 := bsoncodec.DecodeContext{Registry: NewRegistryBuilder().Build()}
+		dec, err := NewDecoderWithContext(dc1, bsonrw.NewBSONDocumentReader([]byte{}))
+		noerr(t, err)
+		if dec.dc != dc1 {
+			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc1)
+		}
+		err = dec.SetContext(dc2)
+		noerr(t, err)
+		if dec.dc != dc2 {
+			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc2)
+		}
+	})
 	t.Run("SetRegistry", func(t *testing.T) {
-		reg1, reg2 := DefaultRegistry, NewRegistryBuilder().Build()
-		dec, err := NewDecoder(reg1, bsonrw.NewBSONDocumentReader([]byte{}))
+		r1, r2 := DefaultRegistry, NewRegistryBuilder().Build()
+		dc1 := bsoncodec.DecodeContext{Registry: r1}
+		dc2 := bsoncodec.DecodeContext{Registry: r2}
+		dec, err := NewDecoder(bsonrw.NewBSONDocumentReader([]byte{}))
 		noerr(t, err)
-		if dec.r != reg1 {
-			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.r, reg1)
+		if dec.dc != dc1 {
+			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc1)
 		}
-		err = dec.SetRegistry(reg2)
+		err = dec.SetRegistry(r2)
 		noerr(t, err)
-		if dec.r != reg2 {
-			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.r, reg2)
+		if dec.dc != dc2 {
+			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc2)
 		}
 	})
 }

bson/encoder.go

@@ -27,21 +27,33 @@ var encPool = sync.Pool{
 // An Encoder writes a serialization format to an output stream. It writes to a bsonrw.ValueWriter
 // as the destination of BSON data.
 type Encoder struct {
-	r  *bsoncodec.Registry
+	ec bsoncodec.EncodeContext
 	vw bsonrw.ValueWriter
 }
 
-// NewEncoder returns a new encoder that uses Registry r to write to w.
-func NewEncoder(r *bsoncodec.Registry, vw bsonrw.ValueWriter) (*Encoder, error) {
-	if r == nil {
-		return nil, errors.New("cannot create a new Encoder with a nil Registry")
+// NewEncoder returns a new encoder that uses the DefaultRegistry to write to vw.
+func NewEncoder(vw bsonrw.ValueWriter) (*Encoder, error) {
+	if vw == nil {
+		return nil, errors.New("cannot create a new Encoder with a nil ValueWriter")
+	}
+
+	return &Encoder{
+		ec: bsoncodec.EncodeContext{Registry: DefaultRegistry},
+		vw: vw,
+	}, nil
+}
+
+// NewEncoderWithContext returns a new encoder that uses EncodeContext ec to write to vw.
+func NewEncoderWithContext(ec bsoncodec.EncodeContext, vw bsonrw.ValueWriter) (*Encoder, error) {
+	if ec.Registry == nil {
+		ec = bsoncodec.EncodeContext{Registry: DefaultRegistry}
 	}
 	if vw == nil {
 		return nil, errors.New("cannot create a new Encoder with a nil ValueWriter")
 	}
 
 	return &Encoder{
-		r:  r,
+		ec: ec,
 		vw: vw,
 	}, nil
 }
@@ -60,14 +72,14 @@ func (e *Encoder) Encode(val interface{}) error {
 		return bsonrw.Copier{}.CopyDocumentFromBytes(e.vw, buf)
 	}
 
-	encoder, err := e.r.LookupEncoder(reflect.TypeOf(val))
+	encoder, err := e.ec.LookupEncoder(reflect.TypeOf(val))
 	if err != nil {
 		return err
 	}
-	return encoder.EncodeValue(bsoncodec.EncodeContext{Registry: e.r}, e.vw, reflect.ValueOf(val))
+	return encoder.EncodeValue(e.ec, e.vw, reflect.ValueOf(val))
 }
 
-// Reset will reset the state of the encoder, using the same *Registry used in
+// Reset will reset the state of the encoder, using the same *EncodeContext used in
 // the original construction but using vw.
 func (e *Encoder) Reset(vw bsonrw.ValueWriter) error {
 	e.vw = vw
@@ -76,6 +88,12 @@ func (e *Encoder) Reset(vw bsonrw.ValueWriter) error {
 
 // SetRegistry replaces the current registry of the encoder with r.
 func (e *Encoder) SetRegistry(r *bsoncodec.Registry) error {
-	e.r = r
+	e.ec.Registry = r
+	return nil
+}
+
+// SetContext replaces the current EncodeContext of the encoder with er.
+func (e *Encoder) SetContext(ec bsoncodec.EncodeContext) error {
+	e.ec = ec
 	return nil
 }

bson/encoder_test.go

@@ -44,8 +44,7 @@ func TestEncoderEncode(t *testing.T) {
 			got := make(bsonrw.SliceWriter, 0, 1024)
 			vw, err := bsonrw.NewBSONValueWriter(&got)
 			noerr(t, err)
-			reg := DefaultRegistry
-			enc, err := NewEncoder(reg, vw)
+			enc, err := NewEncoder(vw)
 			noerr(t, err)
 			err = enc.Encode(tc.val)
 			noerr(t, err)
@@ -103,8 +102,7 @@ func TestEncoderEncode(t *testing.T) {
 					vw, err = bsonrw.NewBSONValueWriter(&b)
 					noerr(t, err)
 				}
-
-				enc, err := NewEncoder(DefaultRegistry, vw)
+				enc, err := NewEncoder(vw)
 				noerr(t, err)
 				got := enc.Encode(marshaler)
 				want := tc.wanterr

bson/marshal.go

@@ -53,10 +53,23 @@ func MarshalWithRegistry(r *bsoncodec.Registry, val interface{}) ([]byte, error)
 	return MarshalAppendWithRegistry(r, dst, val)
 }
 
+// MarshalWithContext returns the BSON encoding of val using EncodeContext ec.
+func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) {
+	dst := make([]byte, 0, 256) // TODO: make the default cap a constant
+	return MarshalAppendWithContext(ec, dst, val)
+}
+
 // MarshalAppendWithRegistry will append the BSON encoding of val to dst using
 // Registry r. If dst is not large enough to hold the BSON encoding of val, dst
 // will be grown.
 func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) {
+	return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val)
+}
+
+// MarshalAppendWithContext will append the BSON encoding of val to dst using
+// EncodeContext ec. If dst is not large enough to hold the BSON encoding of val, dst
+// will be grown.
+func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) {
 	sw := new(bsonrw.SliceWriter)
 	*sw = dst
 	vw := bvwPool.Get(sw)
@@ -69,7 +82,7 @@ func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{
 	if err != nil {
 		return nil, err
 	}
-	err = enc.SetRegistry(r)
+	err = enc.SetContext(ec)
 	if err != nil {
 		return nil, err
 	}
@@ -97,13 +110,26 @@ func MarshalExtJSONAppend(dst []byte, val interface{}, canonical, escapeHTML boo
 // MarshalExtJSONWithRegistry returns the extended JSON encoding of val using Registry r.
 func MarshalExtJSONWithRegistry(r *bsoncodec.Registry, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
 	dst := make([]byte, 0, defaultDstCap)
-	return MarshalExtJSONAppendWithRegistry(r, dst, val, canonical, escapeHTML)
+	return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML)
+}
+
+// MarshalExtJSONWithContext returns the extended JSON encoding of val using Registry r.
+func MarshalExtJSONWithContext(ec bsoncodec.EncodeContext, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
+	dst := make([]byte, 0, defaultDstCap)
+	return MarshalExtJSONAppendWithContext(ec, dst, val, canonical, escapeHTML)
 }
 
 // MarshalExtJSONAppendWithRegistry will append the extended JSON encoding of
 // val to dst using Registry r. If dst is not large enough to hold the BSON
 // encoding of val, dst will be grown.
 func MarshalExtJSONAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
+	return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML)
+}
+
+// MarshalExtJSONAppendWithContext will append the extended JSON encoding of
+// val to dst using Registry r. If dst is not large enough to hold the BSON
+// encoding of val, dst will be grown.
+func MarshalExtJSONAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
 	sw := new(bsonrw.SliceWriter)
 	*sw = dst
 	ejvw := extjPool.Get(sw, canonical, escapeHTML)
@@ -116,7 +142,7 @@ func MarshalExtJSONAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val int
 	if err != nil {
 		return nil, err
 	}
-	err = enc.SetRegistry(r)
+	err = enc.SetContext(ec)
 	if err != nil {
 		return nil, err
 	}