Improve map descriptor handling

plenccodec/descriptor.go

@@ -88,9 +88,16 @@ func (d *Descriptor) read(out Outputter, data []byte) (n int, err error) {
 		return n, err
 
 	case FieldTypeFlatInt:
-		var v int64
-		n, err = FlatIntCodec[uint64]{}.Read(data, unsafe.Pointer(&v), plenccore.WTVarInt)
-		out.Int64(v)
+		switch d.LogicalType {
+		case LogicalTypeTimestamp:
+			var v time.Time
+			n, err = BQTimestampCodec{}.Read(data, unsafe.Pointer(&v), plenccore.WTVarInt)
+			out.Time(v)
+		default:
+			var v int64
+			n, err = FlatIntCodec[uint64]{}.Read(data, unsafe.Pointer(&v), plenccore.WTVarInt)
+			out.Int64(v)
+		}
 		return n, err
 
 	case FieldTypeUint:
@@ -130,11 +137,19 @@ func (d *Descriptor) read(out Outputter, data []byte) (n int, err error) {
 		return n, err
 
 	case FieldTypeSlice:
-		out.StartArray()
-		defer out.EndArray()
+		if d.isValidJSONMap() {
+			out.StartObject()
+			defer out.EndObject()
+		} else {
+			out.StartArray()
+			defer out.EndArray()
+		}
 		return d.readAsSlice(out, data)
 
 	case FieldTypeStruct:
+		if d.isValidJSONMapEntry() {
+			return d.readAsMapEntry(out, data)
+		}
 		out.StartObject()
 		defer out.EndObject()
 		return d.readAsStruct(out, data)
@@ -153,6 +168,27 @@ func (d *Descriptor) read(out Outputter, data []byte) (n int, err error) {
 	return 0, fmt.Errorf("unrecognised field type %s", d.Type)
 }
 
+func (d *Descriptor) isValidJSONMap() bool {
+	if d.Type != FieldTypeSlice || d.LogicalType != LogicalTypeMap {
+		return false
+	}
+	if len(d.Elements) != 1 {
+		return false
+	}
+	return d.Elements[0].isValidJSONMapEntry()
+}
+
+func (d *Descriptor) isValidJSONMapEntry() bool {
+	if d.Type != FieldTypeStruct || d.LogicalType != LogicalTypeMapEntry {
+		return false
+	}
+	if len(d.Elements) != 2 {
+		return false
+	}
+	key := &d.Elements[0]
+	return key.Type == FieldTypeString
+}
+
 func (d *Descriptor) readAsSlice(out Outputter, data []byte) (n int, err error) {
 	elt := &d.Elements[0]
 	switch elt.Type {
@@ -206,6 +242,64 @@ func (d *Descriptor) readAsSlice(out Outputter, data []byte) (n int, err error)
 	}
 }
 
+func (d *Descriptor) readAsMapEntry(out Outputter, data []byte) (n int, err error) {
+	if d.Elements[0].Type != FieldTypeString {
+		// map keys have to be strings to be valid JSON. So we'll output as a
+		// struct instead
+		return
+	}
+
+	l := len(data)
+
+	var offset int
+	for offset < l {
+		wt, index, n := plenccore.ReadTag(data[offset:])
+		offset += n
+
+		var elt *Descriptor
+		for i := range d.Elements {
+			candidate := &d.Elements[i]
+			if candidate.Index == index {
+				elt = candidate
+				break
+			}
+		}
+
+		if elt == nil {
+			// Field corresponding to index does not exist
+			n, err := plenccore.Skip(data[offset:], wt)
+			if err != nil {
+				return 0, fmt.Errorf("failed to skip field %d in %s: %w", index, d.Name, err)
+			}
+			offset += n
+			continue
+		}
+
+		fl := l
+		if wt == plenccore.WTLength {
+			// For WTLength types we read out the length and ensure the data we
+			// read the field from is the right length
+			v, n := plenccore.ReadVarUint(data[offset:])
+			if n <= 0 {
+				return 0, fmt.Errorf("varuint overflow reading field %d of %s", index, d.Name)
+			}
+			offset += n
+			fl = int(v) + offset
+			if fl > l {
+				return 0, fmt.Errorf("length %d of field %d of %s exceeds data length", fl, index, d.Name)
+			}
+		}
+
+		n, err := elt.read(out, data[offset:fl])
+		if err != nil {
+			return 0, fmt.Errorf("failed reading field %d(%s) of %s. %w", index, elt.Name, d.Name, err)
+		}
+		offset += n
+	}
+
+	return offset, nil
+}
+
 func (d *Descriptor) readAsStruct(out Outputter, data []byte) (n int, err error) {
 	l := len(data)
 

plenccodec/descriptor_test.go

@@ -47,10 +47,13 @@ func TestDescriptor(t *testing.T) {
 		O bool                `plenc:"15" json:"elephant"`
 		P map[string]any      `plenc:"16"`
 		Q int32               `plenc:"17,flat"`
+		R time.Time           `plenc:"18,flattime"`
+		S map[string]int      `plenc:"19"`
 	}
 
 	plenc.RegisterCodec(reflect.TypeOf(map[string]any{}), plenccodec.JSONMapCodec{})
 	plenc.RegisterCodec(reflect.TypeOf([]any{}), plenccodec.JSONArrayCodec{})
+	plenc.RegisterCodecWithTag(reflect.TypeOf(time.Time{}), "flattime", plenccodec.BQTimestampCodec{})
 
 	c, err := plenc.CodecForType(reflect.TypeOf(my{}))
 	if err != nil {
@@ -108,6 +111,11 @@ func TestDescriptor(t *testing.T) {
 			"array": []any{1, 1.3, "cheese", json.Number("1337")},
 		},
 		Q: 123,
+		R: time.Date(1970, 3, 15, 0, 0, 0, 1337e5, time.UTC),
+		S: map[string]int{
+			"one": 1,
+			"two": 2,
+		},
 	}
 
 	data, err := plenc.Marshal(nil, in)

plenccodec/descriptor_test.json

@@ -68,5 +68,10 @@
       1337
     ]
   },
-  "Q": 123
+  "Q": 123,
+  "R": "1970-03-15T00:00:00.1337Z",
+  "S": {
+    "one":     1,
+    "two":     2
+  }
 }

plenccodec/output.go

@@ -29,6 +29,20 @@ type JSONOutput struct {
 	data    []byte
 	depth   int
 	inField bool
+
+	stack []stackEntry
+}
+
+type state int
+
+const (
+	stateValue state = iota
+	stateKey
+	stateObjValue
+)
+
+type stackEntry struct {
+	state state
 }
 
 func (j *JSONOutput) Done() []byte {
@@ -40,6 +54,7 @@ func (j *JSONOutput) Reset() {
 	j.data = j.data[:0]
 	j.depth = 0
 	j.inField = false
+	j.stack = j.stack[:0]
 }
 
 func (j *JSONOutput) prefix() {
@@ -53,7 +68,12 @@ func (j *JSONOutput) prefix() {
 }
 
 func (j *JSONOutput) end() {
+	if j.depth == 0 {
+		j.data = append(j.data, '\n')
+		return
+	}
 	j.depth--
+	j.stack = j.stack[:len(j.stack)-1]
 	l := len(j.data)
 	if l < 2 {
 		return
@@ -64,83 +84,104 @@ func (j *JSONOutput) end() {
 	}
 }
 
+func (j *JSONOutput) punctuate() {
+	if len(j.stack) == 0 {
+		return
+	}
+	s := &j.stack[len(j.stack)-1]
+	switch s.state {
+	case stateKey:
+		j.data = append(j.data, ": "...)
+		s.state = stateObjValue
+	case stateObjValue:
+		j.data = append(j.data, ",\n"...)
+		s.state = stateKey
+	case stateValue:
+		j.data = append(j.data, ",\n"...)
+	}
+}
+
 func (j *JSONOutput) StartObject() {
 	j.prefix()
 	j.data = append(j.data, "{\n"...)
 	j.depth++
+	j.stack = append(j.stack, stackEntry{state: stateKey})
 }
 
 func (j *JSONOutput) EndObject() {
 	j.end()
 	j.prefix()
-	j.data = append(j.data, "},\n"...)
+	j.data = append(j.data, '}')
+	j.punctuate()
 }
 
 func (j *JSONOutput) StartArray() {
 	j.prefix()
 	j.data = append(j.data, "[\n"...)
+	j.stack = append(j.stack, stackEntry{state: stateValue})
 	j.depth++
 }
 
 func (j *JSONOutput) EndArray() {
 	j.end()
 	j.prefix()
-	j.data = append(j.data, "],\n"...)
+	j.data = append(j.data, ']')
+	j.punctuate()
 }
 
 func (j *JSONOutput) NameField(name string) {
 	j.prefix()
 	j.inField = true
 	j.data = j.appendString(j.data, name)
-	j.data = append(j.data, ": "...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Int64(v int64) {
 	j.prefix()
 	j.data = strconv.AppendInt(j.data, v, 10)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Uint64(v uint64) {
 	j.prefix()
 	j.data = strconv.AppendUint(j.data, v, 10)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Float64(v float64) {
 	j.prefix()
 	j.data = strconv.AppendFloat(j.data, v, 'g', -1, 64)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Float32(v float32) {
 	j.prefix()
 	j.data = strconv.AppendFloat(j.data, float64(v), 'g', -1, 64)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) String(v string) {
 	j.prefix()
 	j.data = j.appendString(j.data, v)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Raw(v string) {
 	j.prefix()
 	j.data = append(j.data, v...)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Bool(v bool) {
 	j.prefix()
 	j.data = strconv.AppendBool(j.data, v)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 func (j *JSONOutput) Time(t time.Time) {
 	j.prefix()
 	j.data = t.AppendFormat(j.data, `"`+time.RFC3339Nano+`"`)
-	j.data = append(j.data, ",\n"...)
+	j.punctuate()
 }
 
 const hex = "0123456789abcdef"