Improve Set marshaling to add encoding support, and decoding to map (#2845)

The only support that marshal has for Set at the moment is decoding to
slice.

go/marshal/decode.go

@@ -27,21 +27,28 @@ import (
 // length to zero and then appends each element to the slice. If the Go slice
 // was nil a new slice is created.
 //
-// To unmarshal a Noms list or set into a Go array, Unmarshal decodes Noms list
+// To unmarshal a Noms list into a Go array, Unmarshal decodes Noms list
 // elements into corresponding Go array elements.
 //
 // To unmarshal a Noms map into a Go map, Unmarshal decodes Noms key and values
 // into corresponding Go array elements. If the Go map was nil a new map is
 // created.
 //
-// When unmarshalling onto `interface{}` the following rules are used:
-//  - `types.Bool` -> `bool`
-//  - `types.List` -> `[]T`, where `T` is determined recursively using the same rules.
-//  - `types.Set` -> same as `types.List`
-//  - `types.Map` -> `map[T]V`, where `T` and `V` is determined recursively using the same rules.
-//  - `types.Number` -> `float64`
-//  - `types.String` -> `string`
-//  - `types.Union` -> `interface`
+// To unmarshal Noms sets, it depends on the presence of a `noms:",set"` tag:
+//  - Without (default), Unmarshal decodes into corresponding Go list elements.
+//  - With, Unmarshal decodes into Go map keys corresponding to the set values,
+//    with values struct{}{}. Map values must have struct{} type.
+//
+// When unmarshalling onto interface{} the following rules are used:
+//  - types.Bool -> bool
+//  - types.List -> []T, where T is determined recursively using the same rules.
+//  - types.Set -> depends on `noms:",set"` annotation: without, same as
+//    types.List, with, same as types.Map.
+//  - types.Map -> map[T]V, where T and V is determined recursively using the
+//    same rules.
+//  - types.Number -> float64
+//  - types.String -> string
+//  - types.Union -> interface
 //  - Everything else an error
 //
 // Unmarshal returns an UnmarshalTypeMismatchError if:
@@ -66,7 +73,7 @@ func Unmarshal(v types.Value, out interface{}) (err error) {
 		return &InvalidUnmarshalError{reflect.TypeOf(out)}
 	}
 	rv = rv.Elem()
-	d := typeDecoder(rv.Type())
+	d := typeDecoder(rv.Type(), nomsTags{})
 	d(v, rv)
 	return
 }
@@ -112,7 +119,7 @@ func overflowError(v types.Number, t reflect.Type) *UnmarshalTypeMismatchError {
 
 type decoderFunc func(v types.Value, rv reflect.Value)
 
-func typeDecoder(t reflect.Type) decoderFunc {
+func typeDecoder(t reflect.Type, tags nomsTags) decoderFunc {
 	switch t.Kind() {
 	case reflect.Bool:
 		return boolDecoder
@@ -133,7 +140,10 @@ func typeDecoder(t reflect.Type) decoderFunc {
 	case reflect.Array:
 		return arrayDecoder(t)
 	case reflect.Map:
-		return mapDecoder(t)
+		if shouldMapDecodeFromSet(t, tags) {
+			return mapFromSetDecoder(t)
+		}
+		return mapDecoder(t, tags)
 	default:
 		panic(&UnsupportedTypeError{Type: t})
 	}
@@ -193,6 +203,10 @@ type decoderCacheT struct {
 
 var decoderCache = &decoderCacheT{}
 
+// Separate Set decoder cache because the same type with and without the
+// `noms:",set"` tag decode differently (Set vs Map).
+var setDecoderCache = &decoderCacheT{}
+
 func (c *decoderCacheT) get(t reflect.Type) decoderFunc {
 	c.RLock()
 	defer c.RUnlock()
@@ -236,7 +250,7 @@ func structDecoder(t reflect.Type) decoderFunc {
 
 		fields = append(fields, decField{
 			name:    tags.name,
-			decoder: typeDecoder(f.Type),
+			decoder: typeDecoder(f.Type, tags),
 			index:   i,
 		})
 	}
@@ -307,7 +321,7 @@ func sliceDecoder(t reflect.Type) decoderFunc {
 	}
 
 	decoderCache.set(t, d)
-	decoder = typeDecoder(t.Elem())
+	decoder = typeDecoder(t.Elem(), nomsTags{})
 	return d
 }
 
@@ -336,11 +350,43 @@ func arrayDecoder(t reflect.Type) decoderFunc {
 	}
 
 	decoderCache.set(t, d)
-	decoder = typeDecoder(t.Elem())
+	decoder = typeDecoder(t.Elem(), nomsTags{})
 	return d
 }
 
-func mapDecoder(t reflect.Type) decoderFunc {
+func mapFromSetDecoder(t reflect.Type) decoderFunc {
+	d := setDecoderCache.get(t)
+	if d != nil {
+		return d
+	}
+
+	var decoder decoderFunc
+
+	d = func(v types.Value, rv reflect.Value) {
+		m := rv
+		if m.IsNil() {
+			m = reflect.MakeMap(t)
+		}
+
+		nomsSet, ok := v.(types.Set)
+		if !ok {
+			panic(&UnmarshalTypeMismatchError{v, t, `, field has "set" tag`})
+		}
+
+		nomsSet.IterAll(func(v types.Value) {
+			keyRv := reflect.New(t.Key()).Elem()
+			decoder(v, keyRv)
+			m.SetMapIndex(keyRv, reflect.New(t.Elem()).Elem())
+		})
+		rv.Set(m)
+	}
+
+	setDecoderCache.set(t, d)
+	decoder = typeDecoder(t.Key(), nomsTags{})
+	return d
+}
+
+func mapDecoder(t reflect.Type, tags nomsTags) decoderFunc {
 	d := decoderCache.get(t)
 	if d != nil {
 		return d
@@ -355,6 +401,12 @@ func mapDecoder(t reflect.Type) decoderFunc {
 			m = reflect.MakeMap(t)
 		}
 
+		// Special case decoding failure if it looks like the "set" tag is missing,
+		// because it's helpful.
+		if _, ok := v.(types.Set); ok && !tags.set {
+			panic(&UnmarshalTypeMismatchError{v, t, `, field missing "set" tag`})
+		}
+
 		nomsMap, ok := v.(types.Map)
 		if !ok {
 			panic(&UnmarshalTypeMismatchError{v, t, ""})
@@ -371,8 +423,8 @@ func mapDecoder(t reflect.Type) decoderFunc {
 	}
 
 	decoderCache.set(t, d)
-	keyDecoder = typeDecoder(t.Key())
-	valueDecoder = typeDecoder(t.Elem())
+	keyDecoder = typeDecoder(t.Key(), nomsTags{})
+	valueDecoder = typeDecoder(t.Elem(), nomsTags{})
 	return d
 }
 
@@ -388,7 +440,7 @@ func interfaceDecoder(t reflect.Type) decoderFunc {
 	return func(v types.Value, rv reflect.Value) {
 		t := getGoTypeForNomsType(v.Type(), rv.Type(), v)
 		i := reflect.New(t).Elem()
-		typeDecoder(t)(v, i)
+		typeDecoder(t, nomsTags{})(v, i)
 		rv.Set(i)
 	}
 }
@@ -420,3 +472,10 @@ func getGoTypeForNomsType(nt *types.Type, rt reflect.Type, v types.Value) reflec
 		panic(&UnmarshalTypeMismatchError{Value: v, Type: rt})
 	}
 }
+
+func shouldMapDecodeFromSet(rt reflect.Type, tags nomsTags) bool {
+	// map[T]struct{} `noms:,"set"`
+	return tags.set &&
+		rt.Elem().Kind() == reflect.Struct &&
+		rt.Elem().NumField() == 0
+}

go/marshal/decode_test.go

@@ -639,3 +639,89 @@ func TestDecodeOntoInterfaceStruct(t *testing.T) {
 	var i interface{}
 	assertDecodeErrorMessage(t, types.NewStruct("", types.StructData{}), &i, "Cannot unmarshal struct {} into Go value of type interface {}")
 }
+
+func TestDecodeSet(t *testing.T) {
+	assert := assert.New(t)
+
+	type T struct {
+		A map[int]struct{} `noms:",set"`
+		B map[int]struct{}
+		C map[string]struct{} `noms:",set"`
+		D map[string]struct{}
+		E []int
+		F []int
+	}
+
+	ns := types.NewStruct("T", types.StructData{
+		"a": types.NewSet(types.Number(0), types.Number(1), types.Number(2)),
+		"b": types.NewMap(types.Number(3), types.EmptyStruct, types.Number(4), types.EmptyStruct, types.Number(5), types.EmptyStruct),
+		"c": types.NewSet(types.String("0"), types.String("1"), types.String("2")),
+		"d": types.NewMap(types.String("3"), types.EmptyStruct, types.String("4"), types.EmptyStruct, types.String("5"), types.EmptyStruct),
+		"e": types.NewSet(types.Number(6), types.Number(7), types.Number(8)),
+		"f": types.NewList(types.Number(9), types.Number(10), types.Number(11)),
+	})
+
+	gs := T{}
+	assert.NoError(Unmarshal(ns, &gs))
+	assert.Equal(T{
+		A: map[int]struct{}{0: {}, 1: {}, 2: {}},
+		B: map[int]struct{}{3: {}, 4: {}, 5: {}},
+		C: map[string]struct{}{"0": {}, "1": {}, "2": {}},
+		D: map[string]struct{}{"3": {}, "4": {}, "5": {}},
+		E: []int{6, 7, 8},
+		F: []int{9, 10, 11},
+	}, gs)
+}
+
+func TestDecodeNamedSet(t *testing.T) {
+	assert := assert.New(t)
+
+	type T struct {
+		A map[int]struct{} `noms:"foo,set"`
+	}
+
+	ns := types.NewStruct("T", types.StructData{
+		"a":   types.NewSet(types.Number(0)),
+		"foo": types.NewSet(types.Number(1)),
+	})
+
+	gs := T{}
+	assert.NoError(Unmarshal(ns, &gs))
+	assert.Equal(T{
+		map[int]struct{}{1: {}},
+	}, gs)
+}
+
+func TestDecodeSetWrongMapType(t *testing.T) {
+	assert := assert.New(t)
+
+	type T1 struct {
+		A map[int]int `noms:",set"`
+	}
+
+	err := Unmarshal(types.NewStruct("T1", types.StructData{
+		"a": types.NewSet(types.Number(0)),
+	}), &T1{})
+	assert.Error(err)
+	assert.Equal("Cannot unmarshal Set<Number> into Go value of type map[int]int", err.Error())
+
+	type T2 struct {
+		A map[int]struct{}
+	}
+
+	err = Unmarshal(types.NewStruct("T2", types.StructData{
+		"a": types.NewSet(types.Number(0)),
+	}), &T2{})
+	assert.Error(err)
+	assert.Equal(`Cannot unmarshal Set<Number> into Go value of type map[int]struct {}, field missing "set" tag`, err.Error())
+
+	type T3 struct {
+		A map[int]struct{} `noms:",set"`
+	}
+
+	err = Unmarshal(types.NewStruct("T3", types.StructData{
+		"a": types.NewMap(types.Number(0), types.EmptyStruct),
+	}), &T3{})
+	assert.Error(err)
+	assert.Equal(`Cannot unmarshal Map<Number, struct {}> into Go value of type map[int]struct {}, field has "set" tag`, err.Error())
+}

go/marshal/encode.go

@@ -33,7 +33,8 @@ import (
 //
 // Slices and arrays are encoded as Noms types.List.
 //
-// Maps are encoded as Noms types.Map.
+// Maps are encoded as Noms types.Map, or a types.Set if the value type is
+// struct{} and the field is tagged with `noms:"set"`.
 //
 // Struct values are encoded as Noms structs (types.Struct). Each exported Go
 // struct field becomes a member of the Noms struct unless
@@ -76,7 +77,7 @@ import (
 // Noms values (values implementing types.Value) are copied over without any
 // change.
 //
-// When marshalling `interface{}` the dynamic type is used.
+// When marshalling interface{} the dynamic type is used.
 //
 // Go pointers, complex, function are not supported. Attempting to encode such a
 // value causes Marshal to return an UnsupportedTypeError.
@@ -93,7 +94,7 @@ func Marshal(v interface{}) (nomsValue types.Value, err error) {
 		}
 	}()
 	rv := reflect.ValueOf(v)
-	encoder := typeEncoder(rv.Type(), nil)
+	encoder := typeEncoder(rv.Type(), nil, nomsTags{})
 	nomsValue = encoder(rv)
 	return
 }
@@ -134,6 +135,7 @@ func (e *InvalidTagError) Error() string {
 type nomsTags struct {
 	name      string
 	omitEmpty bool
+	set       bool
 	skip      bool
 }
 
@@ -166,7 +168,7 @@ func nomsValueEncoder(v reflect.Value) types.Value {
 	return v.Interface().(types.Value)
 }
 
-func typeEncoder(t reflect.Type, parentStructTypes []reflect.Type) encoderFunc {
+func typeEncoder(t reflect.Type, parentStructTypes []reflect.Type, tags nomsTags) encoderFunc {
 	switch t.Kind() {
 	case reflect.Bool:
 		return boolEncoder
@@ -183,12 +185,15 @@ func typeEncoder(t reflect.Type, parentStructTypes []reflect.Type) encoderFunc {
 	case reflect.Slice, reflect.Array:
 		return listEncoder(t, parentStructTypes)
 	case reflect.Map:
+		if shouldMapEncodeAsSet(t, tags) {
+			return setEncoder(t, parentStructTypes)
+		}
 		return mapEncoder(t, parentStructTypes)
 	case reflect.Interface:
 		return func(v reflect.Value) types.Value {
 			// Get the dynamic type.
 			v2 := reflect.ValueOf(v.Interface())
-			return typeEncoder(v2.Type(), parentStructTypes)(v2)
+			return typeEncoder(v2.Type(), parentStructTypes, tags)(v2)
 		}
 	default:
 		panic(&UnsupportedTypeError{Type: t})
@@ -277,6 +282,10 @@ type encoderCacheT struct {
 
 var encoderCache = &encoderCacheT{}
 
+// Separate Set encoder cache because the same type with and without the
+// `noms:",set"` tag encode differently (Set vs Map).
+var setEncoderCache = &encoderCacheT{}
+
 func (c *encoderCacheT) get(t reflect.Type) encoderFunc {
 	c.RLock()
 	defer c.RUnlock()
@@ -312,9 +321,15 @@ func getTags(f reflect.StructField) (tags nomsTags) {
 		panic(&InvalidTagError{"Invalid struct field name: " + tags.name})
 	}
 
-	if len(tagsSlice) > 1 {
-		// This is pretty simplistic but it is good enough for now.
-		tags.omitEmpty = tagsSlice[1] == "omitempty"
+	for i := 1; i < len(tagsSlice); i++ {
+		switch tag := tagsSlice[i]; tag {
+		case "omitempty":
+			tags.omitEmpty = true
+		case "set":
+			tags.set = true
+		default:
+			panic(&InvalidTagError{"Unrecognized tag: " + tag})
+		}
 	}
 	return
 }
@@ -349,7 +364,7 @@ func typeFields(t reflect.Type, parentStructTypes []reflect.Type) (fields fieldS
 
 		fields = append(fields, field{
 			name:      tags.name,
-			encoder:   typeEncoder(f.Type, parentStructTypes),
+			encoder:   typeEncoder(f.Type, parentStructTypes, tags),
 			index:     i,
 			nomsType:  nt,
 			omitEmpty: tags.omitEmpty,
@@ -437,7 +452,27 @@ func listEncoder(t reflect.Type, parentStructTypes []reflect.Type) encoderFunc {
 	}
 
 	encoderCache.set(t, e)
-	elemEncoder = typeEncoder(t.Elem(), parentStructTypes)
+	elemEncoder = typeEncoder(t.Elem(), parentStructTypes, nomsTags{})
+	return e
+}
+
+func setEncoder(t reflect.Type, parentStructTypes []reflect.Type) encoderFunc {
+	e := setEncoderCache.get(t)
+	if e != nil {
+		return e
+	}
+
+	var encoder encoderFunc
+	e = func(v reflect.Value) types.Value {
+		values := make([]types.Value, v.Len(), v.Len())
+		for i, k := range v.MapKeys() {
+			values[i] = encoder(k)
+		}
+		return types.NewSet(values...)
+	}
+
+	setEncoderCache.set(t, e)
+	encoder = typeEncoder(t.Key(), parentStructTypes, nomsTags{})
 	return e
 }
 
@@ -460,7 +495,15 @@ func mapEncoder(t reflect.Type, parentStructTypes []reflect.Type) encoderFunc {
 	}
 
 	encoderCache.set(t, e)
-	keyEncoder = typeEncoder(t.Key(), parentStructTypes)
-	valueEncoder = typeEncoder(t.Elem(), parentStructTypes)
+	keyEncoder = typeEncoder(t.Key(), parentStructTypes, nomsTags{})
+	valueEncoder = typeEncoder(t.Elem(), parentStructTypes, nomsTags{})
 	return e
 }
+
+func shouldMapEncodeAsSet(t reflect.Type, tags nomsTags) bool {
+	d.PanicIfFalse(t.Kind() == reflect.Map)
+	// map[T]struct{} `noms:,"set"`
+	return tags.set &&
+		t.Elem().Kind() == reflect.Struct &&
+		t.Elem().NumField() == 0
+}

go/marshal/encode_test.go

@@ -503,6 +503,109 @@ func TestEncodeInterface(t *testing.T) {
 	).Equals(v))
 }
 
+func TestEncodeSet(t *testing.T) {
+	assert := assert.New(t)
+
+	v, err := Marshal(struct {
+		A map[int]struct{} `noms:",set"`
+		B map[int]struct{}
+		C map[int]string      `noms:",set"`
+		D map[string]struct{} `noms:",set"`
+		E map[string]struct{}
+		F map[string]int `noms:",set"`
+		G []int          `noms:",set"`
+		H string         `noms:",set"`
+	}{
+		map[int]struct{}{0: {}, 1: {}, 2: {}},
+		map[int]struct{}{3: {}, 4: {}, 5: {}},
+		map[int]string{},
+		map[string]struct{}{"A": {}, "B": {}, "C": {}},
+		map[string]struct{}{"D": {}, "E": {}, "F": {}},
+		map[string]int{},
+		[]int{},
+		"",
+	})
+	assert.NoError(err)
+	s, ok := v.(types.Struct)
+	assert.True(ok)
+
+	expect := map[string]types.NomsKind{
+		"a": types.SetKind,
+		"b": types.MapKind,
+		"c": types.MapKind,
+		"d": types.SetKind,
+		"e": types.MapKind,
+		"f": types.MapKind,
+		"g": types.ListKind,
+		"h": types.StringKind,
+	}
+	for fieldName, kind := range expect {
+		assert.Equal(kind, s.Get(fieldName).Type().Kind())
+	}
+
+	// Test both the Set values are correct, and that the equivalent typed Map
+	// are correct in case the Set marshaling interferes with it.
+
+	a := s.Get("a").(types.Set)
+	assert.True(a.Has(types.Number(0)))
+	assert.True(a.Has(types.Number(1)))
+	assert.True(a.Has(types.Number(2)))
+
+	b := s.Get("b").(types.Map)
+	assert.True(b.Has(types.Number(3)))
+	assert.True(b.Has(types.Number(4)))
+	assert.True(b.Has(types.Number(5)))
+
+	d := s.Get("d").(types.Set)
+	assert.True(d.Has(types.String("A")))
+	assert.True(d.Has(types.String("B")))
+	assert.True(d.Has(types.String("C")))
+
+	e := s.Get("e").(types.Map)
+	assert.True(e.Has(types.String("D")))
+	assert.True(e.Has(types.String("E")))
+	assert.True(e.Has(types.String("F")))
+}
+
+func TestEncodeSetWithTags(t *testing.T) {
+	assert := assert.New(t)
+
+	v, err := Marshal(struct {
+		A map[int]struct{} `noms:"foo,set"`
+		B map[int]struct{} `noms:",omitempty,set"`
+		C map[int]struct{} `noms:"bar,omitempty,set"`
+	}{
+		A: map[int]struct{}{0: {}, 1: {}},
+		C: map[int]struct{}{2: {}, 3: {}},
+	})
+	assert.NoError(err)
+	s, ok := v.(types.Struct)
+	assert.True(ok)
+
+	_, ok = s.MaybeGet("a")
+	assert.False(ok)
+	_, ok = s.MaybeGet("b")
+	assert.False(ok)
+	_, ok = s.MaybeGet("c")
+	assert.False(ok)
+
+	foo, ok := s.Get("foo").(types.Set)
+	assert.True(ok)
+	assert.True(types.NewSet(types.Number(0), types.Number(1)).Equals(foo))
+
+	bar, ok := s.Get("bar").(types.Set)
+	assert.True(ok)
+	assert.True(types.NewSet(types.Number(2), types.Number(3)).Equals(bar))
+}
+
+func TestInvalidTag(t *testing.T) {
+	_, err := Marshal(struct {
+		F string `noms:",omitEmpty"`
+	}{"F"})
+	assert.Error(t, err)
+	assert.Equal(t, `Unrecognized tag: omitEmpty`, err.Error())
+}
+
 type TestInterface interface {
 	M()
 }