~niemeyer/strepr/trunk

// Package strepr is an implementation of the strepr stable representation algorithm.
//
// For details on strepr, see the specification at:
//
//     http://blog.labix.org/2013/06/25/strepr-v1
//
package strepr

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"reflect"
	"sort"
)

// Encoder streams the strepr stable representation for provided values
// into an io.Writer.
type Encoder struct {
	w io.Writer
	b [148]byte
}

// NewEncoder returns a new strepr encoder for writing the stable
// representation for provided values into w.
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w: w}
}

// Encode serializes v's stable representation into e's writer.
func (e *Encoder) Encode(v interface{}) error {
	var err error
	var kind byte
	var uval uint64
	switch v := v.(type) {
	case nil:
		_, err = e.w.Write([]byte{'z'})
	case bool:
		if v {
			_, err = e.w.Write([]byte{'t'})
		} else {
			_, err = e.w.Write([]byte{'f'})
		}
	case string:
		err = e.varint('s', uint64(len(v)))
		if err == nil {
			_, err = e.w.Write([]byte(v))
		}
	case []byte:
		err = e.varint('s', uint64(len(v)))
		if err == nil {
			_, err = e.w.Write(v)
		}
	case int:
		if v < 0 {
			kind = 'n'
			uval = uint64(-v)
		} else {
			kind = 'p'
			uval = uint64(v)
		}
	case int8:
		if v < 0 {
			kind = 'n'
			uval = uint64(-v)
		} else {
			kind = 'p'
			uval = uint64(v)
		}
	case int16:
		if v < 0 {
			kind = 'n'
			uval = uint64(-v)
		} else {
			kind = 'p'
			uval = uint64(v)
		}
	case int32:
		if v < 0 {
			kind = 'n'
			uval = uint64(-v)
		} else {
			kind = 'p'
			uval = uint64(v)
		}
	case int64:
		if v < 0 {
			kind = 'n'
			uval = uint64(-v)
		} else {
			kind = 'p'
			uval = uint64(v)
		}
	case uint:
		kind = 'p'
		uval = uint64(v)
	case uint8:
		kind = 'p'
		uval = uint64(v)
	case uint16:
		kind = 'p'
		uval = uint64(v)
	case uint32:
		kind = 'p'
		uval = uint64(v)
	case uint64:
		kind = 'p'
		uval = uint64(v)
	case float32:
		err = e.float(float64(v))
	case float64:
		err = e.float(v)
	default:
		// Maps and slices may have keys and values with any of the above types.
		rv := reflect.ValueOf(v)
		switch rv.Kind() {
		case reflect.Array:
			if rv.Type().Elem() == byteType {
				return fmt.Errorf("byte arrays are not supported at this time")
			}
			return e.list(rv)
		case reflect.Slice:
			return e.list(rv)
		case reflect.Map:
			return e.dict(rv)
		default:
			return fmt.Errorf("strepr cannot encode %T", v)
		}
	}
	if kind != 0 {
		err = e.varint(kind, uval)
	}
	return err
}

var byteType = reflect.TypeOf(byte(0))

func (e *Encoder) varint(kind byte, v uint64) error {
	i := len(e.b) - 1
	e.b[i] = byte(v & 0x7f)
	v >>= 7
	for v > 0 {
		i--
		e.b[i] = byte(0x80 | (v & 0x7f))
		v >>= 7
	}
	i--
	e.b[i] = kind
	_, err := e.w.Write(e.b[i:])
	return err
}

const (
	mask = 0x7ff
	bias = 1023
	unan = 0x7ff8000000000000
)

func (e *Encoder) float(f float64) error {
	if f == 0 {
		_, err := e.w.Write([]byte{'p', 0})
		return err
	}
	if math.IsNaN(f) {
		f = math.Float64frombits(unan)
	}

	bits := math.Float64bits(f)
	sign := bits & (1 << 63)
	frac := (bits & ((1 << 52) - 1)) | (1 << 52)
	exp := int((bits>>52)&0x7ff) - bias - 52

	if math.IsNaN(f) || math.IsInf(f, 0) || exp < -52 || exp < 0 && (frac&(1<<uint64(-exp)-1)) != 0 {
		// It's not an integer.
		for i := 8; i > 0; i-- {
			e.b[i] = byte(bits & 0xff)
			bits >>= 8
		}
		e.b[0] = 'd'
		_, err := e.w.Write(e.b[:9])
		return err
	}

	// It's an integer. Encode as such.
	if exp > 0 {
		// Having exp as a multiple of 7 simplifies the logic below.
		frac <<= uint64(exp % 7)
		exp -= exp % 7
	}
	i := len(e.b) - 1
	if exp <= 0 {
		frac >>= uint64(-exp)
		e.b[i] = byte(frac & 0x7f)
		frac >>= 7
	} else {
		e.b[i] = 0
		exp -= 7
		for exp > 0 {
			i--
			e.b[i] = 0x80
			exp -= 7
		}
	}
	for frac > 0 {
		i--
		e.b[i] = byte(0x80 | (frac & 0x7f))
		frac >>= 7
	}
	i--
	if sign == 0 {
		e.b[i] = 'p'
	} else {
		e.b[i] = 'n'
	}
	_, err := e.w.Write(e.b[i:])
	return err
}

func (e *Encoder) list(v reflect.Value) error {
	l := v.Len()
	err := e.varint('l', uint64(l))
	if err != nil {
		return err
	}
	for i := 0; i < l; i++ {
		err := e.Encode(v.Index(i).Interface())
		if err != nil {
			return err
		}
	}
	return nil
}

func (e *Encoder) dict(v reflect.Value) error {
	buf := bytes.Buffer{}
	encoder := Encoder{w: &buf}
	ksorter := keySorter{keys: v.MapKeys()}
	ksorter.bufs = make([][]byte, len(ksorter.keys))

	err := e.varint('m', uint64(len(ksorter.keys)))
	if err != nil {
		return err
	}

	// Write all keys in a single buffer.
	pos := 0
	for i := range ksorter.keys {
		err := encoder.Encode(ksorter.keys[i].Interface())
		if err != nil {
			return err
		}
		ksorter.bufs[i] = buf.Bytes()[pos:buf.Len()]
		pos = buf.Len()
	}

	// Sort just the indexes.
	sort.Sort(&ksorter)

	// Write each previously serialized key along with the
	// encoding of its respective value.
	for i := range ksorter.keys {
		_, err := e.w.Write(ksorter.bufs[i])
		if err != nil {
			return err
		}
		err = e.Encode(v.MapIndex(ksorter.keys[i]).Interface())
		if err != nil {
			return err
		}
	}
	return nil
}

type keySorter struct {
	keys []reflect.Value
	bufs [][]byte
}

func (ks *keySorter) Len() int { return len(ks.keys) }

func (ks *keySorter) Swap(i, j int) {
	ks.keys[i], ks.keys[j] = ks.keys[j], ks.keys[i]
	ks.bufs[i], ks.bufs[j] = ks.bufs[j], ks.bufs[i]
}

func (ks *keySorter) Less(i, j int) bool {
	return bytes.Compare(ks.bufs[i], ks.bufs[j]) == -1
}