Ignite/decoder/libdav1d/libdav1d.go

//go:build cgo && !disable_codec_livdav1d

package libdav1d

/*
#cgo pkg-config: dav1d
#cgo LDFLAGS: -lm
#include "libdav1d.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"
	"golang.org/x/exp/constraints"
	"io"
	"runtime"
	"strconv"
	"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

	closer sync.Once

	pool *frame.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
	}

	maxThreads := int(getSettingUnsigned[uint64](settings, "threads", 0))
	if maxThreads == 0 {
		maxThreads = runtime.NumCPU()
	}

	//TODO settings: apply_grain, output_invisible_frames, decode_frame_type, inloop_filters

	C.dav1d_default_settings(&d.settings)
	C.set_threading(&d.settings, C.int(maxThreads))
	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 {
		d.Close()
		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() {
	d.closer.Do(func() {
		if d.data.sz > 0 {
			C.dav1d_data_unref(&d.data)
		}
		if d.ctx != nil {
			C.dav1d_close(&d.ctx)
		}
	})
}

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

func (d *Decoder) pictureToFrame() (f frame.Frame, err 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)
	//TODO check
	properties.FullColorRange = d.picture.seq_hdr.color_range == 1

	properties.ColorSpace.ChromaSamplePosition = color.ChromaSamplePositionUnspecified
	if d.picture.seq_hdr.chr == C.DAV1D_CHR_UNKNOWN {
		properties.ColorSpace.ChromaSamplePosition = color.ChromaSamplePositionCenter
	} else if d.picture.seq_hdr.chr == C.DAV1D_CHR_VERTICAL {
		properties.ColorSpace.ChromaSamplePosition = color.ChromaSamplePositionLeft
	} else if d.picture.seq_hdr.chr == C.DAV1D_CHR_COLOCATED {
		properties.ColorSpace.ChromaSamplePosition = color.ChromaSamplePositionTopLeft
	}

	defer C.dav1d_picture_unref(&d.picture)

	if d.pool == nil || properties != d.pool.Properties() {
		//initialize pool to known size
		d.pool, err = frame.NewPool(properties)
		if err != nil {
			return nil, err
		}
	}

	f = d.pool.Get(int64(d.picture.m.timestamp))

	var yData, uData, vData []byte

	if bitDepth > 8 { //16-bit
		yData = unsafe.Slice((*byte)(d.picture.data[planeY]), properties.Height*properties.Width*2)
		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)
		}
	} else {
		yData = unsafe.Slice((*byte)(d.picture.data[planeY]), properties.Height*properties.Width)
		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)
		}
	}

	copy(f.GetLuma(), yData)
	copy(f.GetCb(), uData)
	copy(f.GetCr(), vData)

	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()
}

func getSettingUnsigned[T constraints.Unsigned](m map[string]any, name string, fallback T) T {
	if v, ok := m[name]; ok {
		if val, ok := v.(string); ok {
			if intVal, err := strconv.ParseUint(val, 10, 0); err != nil {
				delete(m, name)
				return T(intVal)
			} else {
				return fallback
			}
		}
		if val, ok := v.(int); ok {
			return T(val)
		}
		if val, ok := v.(int64); ok {
			return T(val)
		}
		if val, ok := v.(uint); ok {
			return T(val)
		}
		if val, ok := v.(uint64); ok {
			return T(val)
		}
		if val, ok := v.(C.int); ok {
			return T(val)
		}
		if val, ok := v.(C.uint); ok {
			return T(val)
		}
		if val, ok := v.(bool); ok {
			if val {
				return 1
			} else {
				return 0
			}
		}
	}
	return fallback
}

func getSettingBool(m map[string]any, name string, fallback bool) bool {
	if v, ok := m[name]; ok {
		if val, ok := v.(string); ok {
			return val == "false" || val == "f" || val == "n"
		}
		if val, ok := v.(int); ok {
			return val != 0
		}
		if val, ok := v.(int64); ok {
			return val != 0
		}
		if val, ok := v.(uint); ok {
			return val != 0
		}
		if val, ok := v.(uint64); ok {
			return val != 0
		}

		return true
	}
	return fallback
}