lattigo/utils/structs/matrix.go

package structs

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

	"github.com/tuneinsight/lattigo/v4/utils/buffer"
)

// Matrix is a struct storing a vector of Vector.
type Matrix[T any] [][]*T

func (m Matrix[T]) CopyNew() *Matrix[T] {
	mcpy := Matrix[T](make([][]*T, len(m)))

	var err error

	codec := Codec[T]{}

	for i := range m {

		mcpy[i] = make([]*T, len(m[i]))

		for j := range m[i] {
			if mcpy[i][j], err = codec.CopynewWrapper(m[i][j]); err != nil {
				panic(err)
			}
		}
	}

	return &mcpy
}

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

// UnmarshalBinary decodes a slice of bytes generated by
// MarshalBinary or WriteTo on the object.
func (m *Matrix[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 Matrix[T]) WriteTo(w io.Writer) (int64, error) {
	switch w := w.(type) {
	case buffer.Writer:

		var err error
		var n int64

		var inc int
		if inc, err = buffer.WriteInt(w, len(m)); err != nil {
			return int64(inc), err
		}

		n += int64(inc)

		codec := Codec[T]{}

		for _, v := range m {

			var inc int
			if inc, err = buffer.WriteInt(w, len(v)); err != nil {
				return int64(inc), err
			}

			n += int64(inc)

			for _, vi := range v {
				var inc int64
				if inc, err = codec.WriteToWrapper(w, vi); err != nil {
					return n + inc, err
				}

				n += inc
			}
		}

		return n, nil

	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 *Matrix[T]) ReadFrom(r io.Reader) (int64, error) {
	switch r := r.(type) {
	case buffer.Reader:

		var err error
		var size, n int

		if n, err = buffer.ReadInt(r, &size); err != nil {
			return int64(n), fmt.Errorf("cannot ReadFrom: size: %w", err)
		}

		if len(*m) != size {
			*m = make([][]*T, size)
		}

		mi := *m

		codec := Codec[T]{}

		for i := range mi {

			var inc int
			if inc, err = buffer.ReadInt(r, &size); err != nil {
				return int64(n), fmt.Errorf("cannot ReadFrom: size: %w", err)
			}

			n += inc

			if len(mi[i]) != size {
				mi[i] = make([]*T, size)
			}

			for j := range mi[i] {

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

				var inc int64
				if inc, err = codec.ReadFromWrapper(r, mi[i][j]); err != nil {
					return int64(n) + inc, err
				}

				n += int(inc)
			}
		}

		return int64(n), nil

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

// BinarySize returns the size in bytes of the object
// when encoded using Encode.
func (m Matrix[T]) BinarySize() (size int) {
	size += 8
	var err error
	var inc int

	codec := Codec[T]{}

	for _, v := range m {
		size += 8
		for _, vi := range v {
			if inc, err = codec.BinarySizeWrapper(vi); 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 Matrix[T]) Encode(b []byte) (n int, err error) {

	binary.LittleEndian.PutUint64(b[n:], uint64(len(m)))
	n += 8

	codec := Codec[T]{}

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

		binary.LittleEndian.PutUint64(b[n:], uint64(len(v)))
		n += 8

		for _, vi := range v {

			if inc, err = codec.EncodeWrapper(b[n:], vi); 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 *Matrix[T]) Decode(p []byte) (n int, err error) {

	size := int(binary.LittleEndian.Uint64(p[n:]))
	n += 8

	if len(*m) != size {
		*m = make([][]*T, size)
	}

	mi := *m

	codec := Codec[T]{}

	var inc int
	for i := range mi {

		size := int(binary.LittleEndian.Uint64(p[n:]))
		n += 8

		if len(mi[i]) != size {
			mi[i] = make([]*T, size)
		}

		for j := range mi[i] {

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

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

			n += inc
		}
	}

	return
}