add MarshalJSON and UnmarshalJSON to BitArray (#200)

See CHANGELOG

common/bit_array.go

@@ -3,6 +3,7 @@ package common
 import (
 	"encoding/binary"
 	"fmt"
+	"regexp"
 	"strings"
 	"sync"
 )
@@ -249,13 +250,14 @@ func (bA *BitArray) PickRandom() (int, bool) {
 	return 0, false
 }
 
+// String returns a string representation of BitArray: BA{<bit-string>},
+// where <bit-string> is a sequence of 'x' (1) and '_' (0).
+// The <bit-string> includes spaces and newlines to help people.
+// For a simple sequence of 'x' and '_' characters with no spaces or newlines,
+// see the MarshalJSON() method.
+// Example: "BA{_x_}" or "nil-BitArray" for nil.
 func (bA *BitArray) String() string {
-	if bA == nil {
-		return "nil-BitArray"
-	}
-	bA.mtx.Lock()
-	defer bA.mtx.Unlock()
-	return bA.stringIndented("")
+	return bA.StringIndented("")
 }
 
 func (bA *BitArray) StringIndented(indent string) string {
@@ -268,12 +270,11 @@ func (bA *BitArray) StringIndented(indent string) string {
 }
 
 func (bA *BitArray) stringIndented(indent string) string {
-
 	lines := []string{}
 	bits := ""
 	for i := 0; i < bA.Bits; i++ {
 		if bA.getIndex(i) {
-			bits += "X"
+			bits += "x"
 		} else {
 			bits += "_"
 		}
@@ -282,10 +283,10 @@ func (bA *BitArray) stringIndented(indent string) string {
 			bits = ""
 		}
 		if i%10 == 9 {
-			bits += " "
+			bits += indent
 		}
 		if i%50 == 49 {
-			bits += " "
+			bits += indent
 		}
 	}
 	if len(bits) > 0 {
@@ -320,3 +321,58 @@ func (bA *BitArray) Update(o *BitArray) {
 
 	copy(bA.Elems, o.Elems)
 }
+
+// MarshalJSON implements json.Marshaler interface by marshaling bit array
+// using a custom format: a string of '-' or 'x' where 'x' denotes the 1 bit.
+func (bA *BitArray) MarshalJSON() ([]byte, error) {
+	if bA == nil {
+		return []byte("null"), nil
+	}
+
+	bA.mtx.Lock()
+	defer bA.mtx.Unlock()
+
+	bits := `"`
+	for i := 0; i < bA.Bits; i++ {
+		if bA.getIndex(i) {
+			bits += `x`
+		} else {
+			bits += `_`
+		}
+	}
+	bits += `"`
+	return []byte(bits), nil
+}
+
+var bitArrayJSONRegexp = regexp.MustCompile(`\A"([_x]*)"\z`)
+
+// UnmarshalJSON implements json.Unmarshaler interface by unmarshaling a custom
+// JSON description.
+func (bA *BitArray) UnmarshalJSON(bz []byte) error {
+	b := string(bz)
+	if b == "null" {
+		// This is required e.g. for encoding/json when decoding
+		// into a pointer with pre-allocated BitArray.
+		bA.Bits = 0
+		bA.Elems = nil
+		return nil
+	}
+
+	// Validate 'b'.
+	match := bitArrayJSONRegexp.FindStringSubmatch(b)
+	if match == nil {
+		return fmt.Errorf("BitArray in JSON should be a string of format %q but got %s", bitArrayJSONRegexp.String(), b)
+	}
+	bits := match[1]
+
+	// Construct new BitArray and copy over.
+	numBits := len(bits)
+	bA2 := NewBitArray(numBits)
+	for i := 0; i < numBits; i++ {
+		if bits[i] == 'x' {
+			bA2.SetIndex(i, true)
+		}
+	}
+	*bA = *bA2
+	return nil
+}