lattigo/utils/structs/map.go

package structs

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"

	"github.com/tuneinsight/lattigo/v4/utils"
	"github.com/tuneinsight/lattigo/v4/utils/buffer"
	"golang.org/x/exp/constraints"
)

// Map is a struct storing a map of any element indexed by an Integer.
type Map[V constraints.Integer, T any] map[V]*T

// CopyNew creates a copy of the oject.
func (m Map[V, T]) CopyNew() *Map[V, T] {

	var mcpy = make(Map[V, T])

	codec := Codec[T]{}

	var err error
	for key, object := range m {
		if mcpy[key], err = codec.CopynewWrapper(object); err != nil {
			panic(err)
		}
	}

	return &mcpy
}

// MarshalBinary encodes the object into a binary form on a newly allocated slice of bytes.
func (m *Map[V, T]) MarshalBinary() (p []byte, err error) {
	buf := bytes.NewBuffer([]byte{})
	_, err = m.WriteTo(buf)
	return buf.Bytes(), err
}

// UnmarshalBinary decodes a slice of bytes generated by
// MarshalBinary or WriteTo on the object.
func (m *Map[V, T]) UnmarshalBinary(p []byte) (err error) {
	_, err = m.ReadFrom(bytes.NewBuffer(p))
	return
}

// WriteTo writes the object on an io.Writer.
// To ensure optimal efficiency and minimal allocations, the user is encouraged
// to provide a struct implementing the interface buffer.Writer, which defines
// a subset of the method of the bufio.Writer.
// If w is not compliant to the buffer.Writer interface, it will be wrapped in
// a new bufio.Writer.
// For additional information, see lattigo/utils/buffer/writer.go.
func (m *Map[V, T]) WriteTo(w io.Writer) (n int64, err error) {

	switch w := w.(type) {
	case buffer.Writer:

		mi := *m

		var inc1 int

		if inc1, err = buffer.WriteUint32(w, uint32(len(mi))); err != nil {
			return n + int64(inc1), err
		}

		n += int64(inc1)

		codec := Codec[T]{}

		for _, key := range utils.GetSortedKeys(mi) {

			if inc1, err = buffer.WriteUint64(w, uint64(key)); err != nil {
				return n + int64(inc1), err
			}

			n += int64(inc1)

			var inc2 int64
			if inc2, err = codec.WriteToWrapper(w, mi[key]); err != nil {
				return n + inc2, err
			}

			n += inc2
		}

		return

	default:
		return m.WriteTo(bufio.NewWriter(w))
	}
}

// ReadFrom reads on the object from an io.Writer.
// To ensure optimal efficiency and minimal allocations, the user is encouraged
// to provide a struct implementing the interface buffer.Reader, which defines
// a subset of the method of the bufio.Reader.
// If r is not compliant to the buffer.Reader interface, it will be wrapped in
// a new bufio.Reader.
// For additional information, see lattigo/utils/buffer/reader.go.
func (m *Map[V, T]) ReadFrom(r io.Reader) (n int64, err error) {
	switch r := r.(type) {
	case buffer.Reader:

		mi := *m

		var inc1 int

		var size uint32
		if inc1, err = buffer.ReadUint32(r, &size); err != nil {
			return n + int64(inc1), err
		}

		n += int64(inc1)

		codec := Codec[T]{}

		for i := 0; i < int(size); i++ {

			var key uint64

			if inc1, err = buffer.ReadUint64(r, &key); err != nil {
				return n + int64(inc1), err
			}

			n += int64(inc1)

			if mi[V(key)] == nil {
				mi[V(key)] = new(T)
			}

			var inc2 int64
			if inc2, err = codec.ReadFromWrapper(r, mi[V(key)]); err != nil {
				return n + inc2, err
			}

			n += inc2
		}

		return

	default:
		return m.ReadFrom(bufio.NewReader(r))
	}
}

// BinarySize returns the size in bytes of the object
// when encoded using Encode.
func (m Map[V, T]) BinarySize() (size int) {
	size = 4 // #Ct

	codec := Codec[T]{}

	var inc int
	var err error
	for _, v := range m {

		size += 8

		if inc, err = codec.BinarySizeWrapper(v); err != nil {
			panic(err)
		}

		size += inc
	}

	return
}

// Encode encodes the object into a binary form on a preallocated slice of bytes
// and returns the number of bytes written.
func (m *Map[V, T]) Encode(p []byte) (n int, err error) {

	if len(p) < m.BinarySize() {
		return n, fmt.Errorf("cannot Encode: len(p)=%d < %d", len(p), m.BinarySize())
	}

	codec := Codec[T]{}

	mi := *m

	binary.LittleEndian.PutUint32(p[n:], uint32(len(mi)))
	n += 4

	for _, key := range utils.GetSortedKeys(mi) {

		binary.LittleEndian.PutUint64(p[n:], uint64(key))
		n += 8

		var inc int
		if inc, err = codec.EncodeWrapper(p[n:], mi[key]); err != nil {
			return n + inc, err
		}

		n += inc
	}

	return
}

// Decode decodes a slice of bytes generated by Encode
// on the object and returns the number of bytes read.
func (m *Map[V, T]) Decode(p []byte) (n int, err error) {

	mi := *m

	size := int(binary.LittleEndian.Uint32(p[n:]))
	n += 4

	codec := Codec[T]{}

	for i := 0; i < size; i++ {

		idx := V(binary.LittleEndian.Uint64(p[n:]))
		n += 8

		if mi[idx] == nil {
			mi[idx] = new(T)
		}

		var inc int
		if inc, err = codec.DecodeWrapper(p[n:], mi[idx]); err != nil {
			return n + inc, err
		}

		n += inc
	}

	return
}