Ignite/decoder/libdav1d/libdav1d.go

//go:build cgo && !disable_codec_livdav1d

package libdav1d

/*
#cgo pkg-config: dav1d
#include <dav1d/dav1d.h>
*/
import "C"
import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/color"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities/ivfreader"
	"io"
	"runtime"
	"sync"
	"sync/atomic"
	"unsafe"
)

type Decoder struct {
	r       *ivfreader.IVFReader
	h       *ivfreader.IVFFileHeader
	cleaned atomic.Bool

	properties frame.StreamProperties
	firstFrame frame.Frame

	settings C.Dav1dSettings
	ctx      *C.Dav1dContext
	picture  C.Dav1dPicture
	data     C.Dav1dData

	bufPool sync.Pool
}

var dav1dVersion = fmt.Sprintf("dav1d %s", C.GoString(C.dav1d_version()))

const (
	ErrEAGAIN = 11
)

func Version() string {
	return dav1dVersion
}

func NewDecoder(r io.Reader, settings map[string]any) (d *Decoder, err error) {

	d = &Decoder{}
	d.properties.PixelAspectRatio = utilities.Ratio{Numerator: 1, Denominator: 1}

	if d.r, d.h, err = ivfreader.NewWith(r); err != nil {
		return nil, err
	}

	C.dav1d_default_settings(&d.settings)
	if ret := C.dav1d_open(&d.ctx, &d.settings); ret != 0 {
		return nil, fmt.Errorf("error %d", ret)
	}
	if d.firstFrame, err = d.Decode(); err != nil {
		return nil, err
	}

	fP := d.firstFrame.Properties()
	d.properties = frame.StreamProperties{
		Width:            fP.Width,
		Height:           fP.Height,
		PixelAspectRatio: fP.PixelAspectRatio,
		ColorSpace:       fP.ColorSpace,
		FrameRate: utilities.Ratio{
			Numerator:   int(d.h.TimebaseNumerator),
			Denominator: int(d.h.TimebaseDenominator),
		}.Reciprocal(),
		FullColorRange: fP.FullColorRange,
	}

	return d, nil
}

func (d *Decoder) readIvf() error {
	if buf, header, err := d.r.ParseNextFrame(); err != nil {
		return err
	} else {
		if ptr := C.dav1d_data_create(&d.data, C.size_t(len(buf))); ptr == nil {
			return errors.New("could not allocate data")
		} else {
			copy(unsafe.Slice((*byte)(ptr), len(buf)), buf)
			d.data.m.timestamp = C.long(header.Timestamp)
		}
	}
	return nil
}

func (d *Decoder) Properties() frame.StreamProperties {
	return d.properties
}

func (d *Decoder) decodeFrame() C.int {
	var res C.int

	if res = d.flushPicture(); res < 0 {
		if res != -ErrEAGAIN {
			return res
		}
	} else {
		return res
	}

	if d.data.sz == 0 {
		if err := d.readIvf(); err != nil {
			if err == io.EOF {
				//try flush?
				if res = d.flushPicture(); res != -ErrEAGAIN {
					return res
				}
			}
			res = -1
			return res
		}
	}

	if res = C.dav1d_send_data(d.ctx, &d.data); res < 0 {
		if res != -ErrEAGAIN {
			return res
		}
	}

	if res = d.flushPicture(); res < 0 {
		if res != -ErrEAGAIN {
			return res
		}
	}
	return res
}

func (d *Decoder) flushPicture() C.int {
	//todo memset picture to 0?
	res := C.dav1d_get_picture(d.ctx, &d.picture)

	return res
}

func (d *Decoder) Close() {
	if d.data.sz > 0 {
		C.dav1d_data_unref(&d.data)
	}
	//TODO: close other context

}

const (
	planeY = 0
	planeU = 1
	planeV = 2
)

func (d *Decoder) pictureToFrame() (frame.Frame, error) {

	properties := d.properties.FrameProperties()

	bitDepth := int(d.picture.p.bpc)
	properties.Height = int(d.picture.p.h)
	properties.Width = int(d.picture.p.w)

	chromaWidth := 0
	chromaHeight := 0

	switch d.picture.p.layout {
	case C.DAV1D_PIXEL_LAYOUT_I400:
		properties.ColorSpace.ChromaSampling = color.ChromaSampling{J: 4, A: 0, B: 0}
	case C.DAV1D_PIXEL_LAYOUT_I420:
		properties.ColorSpace.ChromaSampling = color.ChromaSampling{J: 4, A: 2, B: 0}
	case C.DAV1D_PIXEL_LAYOUT_I422:
		properties.ColorSpace.ChromaSampling = color.ChromaSampling{J: 4, A: 2, B: 2}
	case C.DAV1D_PIXEL_LAYOUT_I444:
		properties.ColorSpace.ChromaSampling = color.ChromaSampling{J: 4, A: 4, B: 4}
	}

	if d.picture.p.layout != C.DAV1D_PIXEL_LAYOUT_I400 {
		ssVer := 0
		if d.picture.p.layout == C.DAV1D_PIXEL_LAYOUT_I420 {
			ssVer = 1
		}
		ssHor := 0
		if d.picture.p.layout != C.DAV1D_PIXEL_LAYOUT_I444 {
			ssHor = 1
		}

		chromaWidth = (properties.Width + ssHor) >> ssHor
		chromaHeight = (properties.Height + ssVer) >> ssVer
	}

	properties.ColorSpace.BitDepth = byte(bitDepth)
	properties.FullColorRange = false

	if d.picture.seq_hdr.color_range == 1 {
		//TODO check
		properties.FullColorRange = true
	}

	defer C.dav1d_picture_unref(&d.picture)

	if bitDepth > 8 { //16-bit

		yData := unsafe.Slice((*byte)(d.picture.data[planeY]), properties.Height*properties.Width*2)
		var uData, vData []byte

		if d.picture.p.layout != C.DAV1D_PIXEL_LAYOUT_I400 {
			uData = unsafe.Slice((*byte)(d.picture.data[planeU]), chromaHeight*chromaWidth*2)
			vData = unsafe.Slice((*byte)(d.picture.data[planeV]), chromaHeight*chromaWidth*2)

		}

		n := len(yData) + len(uData) + len(vData)

		var buf []byte
		if b := d.bufPool.Get(); b != nil && len(b.([]byte)) == n {
			buf = b.([]byte)
		} else {
			buf = make([]byte, n)
		}

		copy(buf, yData)
		copy(buf[len(yData):], uData)
		copy(buf[len(yData)+len(uData):], vData)

		if f, err := frame.NewUint16FrameFromBytes(properties, int64(d.picture.m.timestamp), buf); err != nil {
			return nil, err
		} else {
			d.bufPool.Put(buf)
			return f, nil
		}
	} else {
		yData := unsafe.Slice((*byte)(d.picture.data[planeY]), properties.Height*properties.Width)
		var uData, vData []byte

		if d.picture.p.layout != C.DAV1D_PIXEL_LAYOUT_I400 {
			uData = unsafe.Slice((*byte)(d.picture.data[planeU]), chromaHeight*chromaWidth)
			vData = unsafe.Slice((*byte)(d.picture.data[planeV]), chromaHeight*chromaWidth)
		}

		n := len(yData) + len(uData) + len(vData)

		var buf []byte
		if b := d.bufPool.Get(); b != nil && len(b.([]byte)) == n {
			buf = b.([]byte)
		} else {
			buf = make([]byte, n)
		}

		copy(buf, yData)
		copy(buf[len(yData):], uData)
		copy(buf[len(yData)+len(uData):], vData)

		if f, err := frame.NewUint8FrameFromBytes(properties, int64(d.picture.m.timestamp), buf); err != nil {
			return nil, err
		} else {
			runtime.SetFinalizer(f, func(frameUint8 *frame.FrameUint8) {
				d.bufPool.Put(buf)
			})
			return f, nil
		}
	}
}

func (d *Decoder) DecodeStream() *frame.Stream {
	stream, channel := frame.NewStream(d.properties)
	go func() {
		defer close(channel)
		for {
			if f, err := d.Decode(); err != nil {
				return
			} else {
				channel <- f
			}
		}
	}()
	return stream
}

func (d *Decoder) Decode() (frame.Frame, error) {
	if f := d.firstFrame; f != nil {
		d.firstFrame = nil
		return f, nil
	}
	var ret C.int
	for {
		if ret = d.decodeFrame(); ret == -ErrEAGAIN {
			continue
		}
		if ret == 0 {
			return d.pictureToFrame()
		}

		if ret == -1 {
			return nil, io.EOF
		}

		return nil, fmt.Errorf("error %d", ret)
	}
}

func (d *Decoder) Version() string {
	return Version()
}