From 8504b654b18f0eb74dd06a5c9687b855975927b5 Mon Sep 17 00:00:00 2001
From: Cui Xin <steven.cuixin@qq.com>
Date: Fri, 17 Nov 2017 23:00:43 +0800
Subject: [PATCH] add support customize binaryOrder & base64.StdEncoding

---
 bitset.go      | 24 ++++++++++++++++++------
 bitset_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/bitset.go b/bitset.go
index dbd1708..fcbb493 100644
--- a/bitset.go
+++ b/bitset.go
@@ -58,6 +58,18 @@ const log2WordSize = uint(6)
 // allBits has every bit set
 const allBits uint64 = 0xffffffffffffffff
 
+// default binary BigEndian
+var binaryOrder binary.ByteOrder = binary.BigEndian
+
+// default json encoding base64.URLEncoding
+var base64Encoding *base64.Encoding = base64.URLEncoding
+
+// Marshal/Unmarshal BitSet with base64.StdEncoding(Default: base64.URLEncoding)
+func Base64StdEncoding() { base64Encoding = base64.StdEncoding }
+
+// Marshal/Unmarshal Binary as Little Endian(Default: binary.BigEndian)
+func LittleEndian() { binaryOrder = binary.LittleEndian }
+
 // A BitSet is a set of bits. The zero value of a BitSet is an empty set of length 0.
 type BitSet struct {
 	length uint
@@ -617,13 +629,13 @@ func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
 	length := uint64(b.length)
 
 	// Write length
-	err := binary.Write(stream, binary.BigEndian, length)
+	err := binary.Write(stream, binaryOrder, length)
 	if err != nil {
 		return 0, err
 	}
 
 	// Write set
-	err = binary.Write(stream, binary.BigEndian, b.set)
+	err = binary.Write(stream, binaryOrder, b.set)
 	return int64(b.BinaryStorageSize()), err
 }
 
@@ -632,7 +644,7 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
 	var length uint64
 
 	// Read length first
-	err := binary.Read(stream, binary.BigEndian, &length)
+	err := binary.Read(stream, binaryOrder, &length)
 	if err != nil {
 		return 0, err
 	}
@@ -643,7 +655,7 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
 	}
 
 	// Read remaining bytes as set
-	err = binary.Read(stream, binary.BigEndian, newset.set)
+	err = binary.Read(stream, binaryOrder, newset.set)
 	if err != nil {
 		return 0, err
 	}
@@ -686,7 +698,7 @@ func (b *BitSet) MarshalJSON() ([]byte, error) {
 	}
 
 	// URLEncode all bytes
-	return json.Marshal(base64.URLEncoding.EncodeToString(buffer.Bytes()))
+	return json.Marshal(base64Encoding.EncodeToString(buffer.Bytes()))
 }
 
 // UnmarshalJSON unmarshals a BitSet from JSON created using MarshalJSON
@@ -699,7 +711,7 @@ func (b *BitSet) UnmarshalJSON(data []byte) error {
 	}
 
 	// URLDecode string
-	buf, err := base64.URLEncoding.DecodeString(s)
+	buf, err := base64Encoding.DecodeString(s)
 	if err != nil {
 		return err
 	}
diff --git a/bitset_test.go b/bitset_test.go
index bc419a0..8a2297e 100644
--- a/bitset_test.go
+++ b/bitset_test.go
@@ -883,6 +883,20 @@ func TestMarshalUnmarshalBinary(t *testing.T) {
 	}
 }
 
+func TestMarshalUnmarshalBinaryByLittleEndian(t *testing.T) {
+	LittleEndian()
+	a := New(1010).Set(10).Set(1001)
+	b := new(BitSet)
+
+	copyBinary(t, a, b)
+
+	// BitSets must be equal after marshalling and unmarshalling
+	if !a.Equal(b) {
+		t.Error("Bitsets are not equal:\n\t", a.DumpAsBits(), "\n\t", b.DumpAsBits())
+		return
+	}
+}
+
 func copyBinary(t *testing.T, from encoding.BinaryMarshaler, to encoding.BinaryUnmarshaler) {
 	data, err := from.MarshalBinary()
 	if err != nil {
@@ -919,6 +933,29 @@ func TestMarshalUnmarshalJSON(t *testing.T) {
 	}
 }
 
+func TestMarshalUnmarshalJSONByStdEncoding(t *testing.T) {
+	Base64StdEncoding()
+	a := New(1010).Set(10).Set(1001)
+	data, err := json.Marshal(a)
+	if err != nil {
+		t.Errorf(err.Error())
+		return
+	}
+
+	b := new(BitSet)
+	err = json.Unmarshal(data, b)
+	if err != nil {
+		t.Errorf(err.Error())
+		return
+	}
+
+	// Bitsets must be equal after marshalling and unmarshalling
+	if !a.Equal(b) {
+		t.Error("Bitsets are not equal:\n\t", a.DumpAsBits(), "\n\t", b.DumpAsBits())
+		return
+	}
+}
+
 func TestSafeSet(t *testing.T) {
 	b := new(BitSet)
 	c := b.safeSet()