encode.go

package bencode

import (
    "io"
    "math/big"
    "reflect"
    "runtime"
    "sort"
    "strconv"
    "sync"
)

type Encoder struct {
    w       io.Writer
    scratch [64]byte
}

func (e *Encoder) Encode(v any) (err error) {
    if v == nil {
        return
    }

    defer func() {
        if e := recover(); e != nil {
            if _, ok := e.(runtime.Error); ok {
                panic(e)
            }

            var ok bool
            err, ok = e.(error)
            if !ok {
                panic(e)
            }
        }
    }()

    e.reflectValue(reflect.ValueOf(v))
    return nil
}

type stringValues []reflect.Value

func (sv stringValues) Len() int           { return len(sv) }
func (sv stringValues) Swap(i, j int)      { sv[i], sv[j] = sv[j], sv[i] }
func (sv stringValues) Less(i, j int) bool { return sv.get(i) < sv.get(j) }
func (sv stringValues) get(i int) string   { return sv[i].String() }

func (e *Encoder) write(s []byte) {
    _, err := e.w.Write(s)
    if err != nil {
        panic(err)
    }
}

func (e *Encoder) writeString(s string) {
    for s != "" {
        n := copy(e.scratch[:], s)
        s = s[n:]
        e.write(e.scratch[:n])
    }
}

func (e *Encoder) reflectString(s string) {
    e.writeStringPrefix(int64(len(s)))
    e.writeString(s)
}

func (e *Encoder) writeStringPrefix(l int64) {
    b := strconv.AppendInt(e.scratch[:0], l, 10)
    e.write(b)
    e.writeString(":")
}

func (e *Encoder) reflectByteSlice(s []byte) {
    e.writeStringPrefix(int64(len(s)))
    e.write(s)
}

// Returns true if the value implements Marshaler interface and marshaling was
// done successfully.
func (e *Encoder) reflectMarshaler(v reflect.Value) bool {
    if !v.Type().Implements(marshalerType) {
        if v.Kind() != reflect.Ptr && v.CanAddr() && v.Addr().Type().Implements(marshalerType) {
            v = v.Addr()
        } else {
            return false
        }
    }
    m := v.Interface().(Marshaler)
    data, err := m.MarshalBencode()
    if err != nil {
        panic(&MarshalerError{v.Type(), err})
    }
    e.write(data)
    return true
}

var bigIntType = reflect.TypeOf((*big.Int)(nil)).Elem()

func (e *Encoder) reflectValue(v reflect.Value) {
    if e.reflectMarshaler(v) {
        return
    }

    if v.Type() == bigIntType {
        e.writeString("i")
        bi := v.Interface().(big.Int)
        e.writeString(bi.String())
        e.writeString("e")
        return
    }

    switch v.Kind() {
    case reflect.Bool:
        if v.Bool() {
            e.writeString("i1e")
        } else {
            e.writeString("i0e")
        }
    case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
        e.writeString("i")
        b := strconv.AppendInt(e.scratch[:0], v.Int(), 10)
        e.write(b)
        e.writeString("e")
    case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
        e.writeString("i")
        b := strconv.AppendUint(e.scratch[:0], v.Uint(), 10)
        e.write(b)
        e.writeString("e")
    case reflect.String:
        e.reflectString(v.String())
    case reflect.Struct:
        e.writeString("d")
        for _, ef := range getEncodeFields(v.Type()) {
            fieldValue := ef.i(v)
            if !fieldValue.IsValid() {
                continue
            }
            if ef.omitEmpty && IsEmptyValue(fieldValue) {
                continue
            }
            e.reflectString(ef.tag)
            e.reflectValue(fieldValue)
        }
        e.writeString("e")
    case reflect.Map:
        if v.Type().Key().Kind() != reflect.String {
            panic(&MarshalTypeError{v.Type()})
        }
        if v.IsNil() {
            e.writeString("de")
            break
        }
        e.writeString("d")
        sv := stringValues(v.MapKeys())
        sort.Sort(sv)
        for _, key := range sv {
            e.reflectString(key.String())
            e.reflectValue(v.MapIndex(key))
        }
        e.writeString("e")
    case reflect.Slice, reflect.Array:
        e.reflectSequence(v)
    case reflect.Interface:
        e.reflectValue(v.Elem())
    case reflect.Ptr:
        if v.IsNil() {
            v = reflect.Zero(v.Type().Elem())
        } else {
            v = v.Elem()
        }
        e.reflectValue(v)
    default:
        panic(&MarshalTypeError{v.Type()})
    }
}

func (e *Encoder) reflectSequence(v reflect.Value) {
    // Use bencode string-type
    if v.Type().Elem().Kind() == reflect.Uint8 {
        if v.Kind() != reflect.Slice {
            // Can't use []byte optimization
            if !v.CanAddr() {
                e.writeStringPrefix(int64(v.Len()))
                for i := 0; i < v.Len(); i++ {
                    var b [1]byte
                    b[0] = byte(v.Index(i).Uint())
                    e.write(b[:])
                }
                return
            }
            v = v.Slice(0, v.Len())
        }
        s := v.Bytes()
        e.reflectByteSlice(s)
        return
    }
    if v.IsNil() {
        e.writeString("le")
        return
    }
    e.writeString("l")
    for i, n := 0, v.Len(); i < n; i++ {
        e.reflectValue(v.Index(i))
    }
    e.writeString("e")
}

type encodeField struct {
    i         func(v reflect.Value) reflect.Value
    tag       string
    omitEmpty bool
}

type encodeFieldsSortType []encodeField

func (ef encodeFieldsSortType) Len() int           { return len(ef) }
func (ef encodeFieldsSortType) Swap(i, j int)      { ef[i], ef[j] = ef[j], ef[i] }
func (ef encodeFieldsSortType) Less(i, j int) bool { return ef[i].tag < ef[j].tag }

var (
    typeCacheLock     sync.RWMutex
    encodeFieldsCache = make(map[reflect.Type][]encodeField)
)

func getEncodeFields(t reflect.Type) []encodeField {
    typeCacheLock.RLock()
    fs, ok := encodeFieldsCache[t]
    typeCacheLock.RUnlock()
    if ok {
        return fs
    }
    fs = makeEncodeFields(t)
    typeCacheLock.Lock()
    defer typeCacheLock.Unlock()
    encodeFieldsCache[t] = fs
    return fs
}

func makeEncodeFields(t reflect.Type) (fs []encodeField) {
    for _i, n := 0, t.NumField(); _i < n; _i++ {
        i := _i
        f := t.Field(i)
        if f.PkgPath != "" {
            continue
        }
        if f.Anonymous {
            t := f.Type
            if t.Kind() == reflect.Ptr {
                t = t.Elem()
            }
            anonEFs := makeEncodeFields(t)
            for aefi := range anonEFs {
                anonEF := anonEFs[aefi]
                bottomField := anonEF
                bottomField.i = func(v reflect.Value) reflect.Value {
                    v = v.Field(i)
                    if v.Kind() == reflect.Ptr {
                        if v.IsNil() {
                            // This will skip serializing this value.
                            return reflect.Value{}
                        }
                        v = v.Elem()
                    }
                    return anonEF.i(v)
                }
                fs = append(fs, bottomField)
            }
            continue
        }
        var ef encodeField
        ef.i = func(v reflect.Value) reflect.Value {
            return v.Field(i)
        }
        ef.tag = f.Name

        tv := getTag(f.Tag)
        if tv.Ignore() {
            continue
        }
        if tv.Key() != "" {
            ef.tag = tv.Key()
        }
        ef.omitEmpty = tv.OmitEmpty()
        fs = append(fs, ef)
    }
    fss := encodeFieldsSortType(fs)
    sort.Sort(fss)
    return fs
}