Kirika/audio/format/analyzer.go

package format

import (
	"git.gammaspectra.live/S.O.N.G/Kirika/audio"
	"io"
	"runtime"
	"sync"
	"sync/atomic"
	"unsafe"
)

const chanBuf = 16

type AnalyzerPacket struct {
	//Samples interleaved samples
	Samples    []int32
	Channels   int
	SampleRate int
	BitDepth   int
}

type AnalyzerChannel chan *AnalyzerPacket

type AnalyzerDecoder interface {
	Decoder
	// OpenAnalyzer Opens a stream and decodes it into an audio.Source, and additionally copy AnalyzerPacket back
	OpenAnalyzer(r io.ReadSeekCloser) (audio.Source, AnalyzerChannel, error)
}

func NewAnalyzerChannel(source audio.Source, err error) (audio.Source, AnalyzerChannel, error) {
	if source == nil {
		return nil, nil, err
	}

	sources := source.Split(2)

	channel := make(AnalyzerChannel)

	go func() {
		defer close(channel)

		if float32Source, ok := sources[1].(audio.TypedSource[float32]); ok {
			for block := range float32Source.GetBlocks() {
				//Convert float32 to int32 representation with 32 bit depth, so it hashes the same way
				buf := make([]int32, len(block))
				copy(buf, unsafe.Slice((*int32)(unsafe.Pointer(unsafe.SliceData(block))), len(block)))
				runtime.KeepAlive(block)
				channel <- &AnalyzerPacket{
					Samples:    buf,
					Channels:   source.GetChannels(),
					SampleRate: source.GetSampleRate(),
					BitDepth:   32,
				}
			}
		} else {
			for block := range sources[1].ToInt32(source.GetBitDepth()).GetBlocks() {
				channel <- &AnalyzerPacket{
					Samples:    block,
					Channels:   source.GetChannels(),
					SampleRate: source.GetSampleRate(),
					BitDepth:   source.GetBitDepth(),
				}
			}
		}
	}()

	return sources[0], channel, err
}

func (c AnalyzerChannel) Split(n int) (channels []AnalyzerChannel) {
	channels = make([]AnalyzerChannel, n)
	for i := range channels {
		channels[i] = make(AnalyzerChannel, chanBuf)
	}

	go func() {
		defer func() {
			for _, channel := range channels {
				close(channel)
			}
		}()

		for packet := range c {
			for _, channel := range channels {
				channel <- packet
			}
		}
	}()

	return
}

func (c AnalyzerChannel) PrependGap(samples, sampleRate, channels, bitDepth int) (channel AnalyzerChannel) {
	return MergeHasherChannels(NewHasherAudioGap(samples, sampleRate, channels, bitDepth), c)
}

func (c AnalyzerChannel) AppendGap(samples, sampleRate, channels, bitDepth int) (channel AnalyzerChannel) {
	return MergeHasherChannels(c, NewHasherAudioGap(samples, sampleRate, channels, bitDepth))
}

func (c AnalyzerChannel) SkipStartSamples(samples int) (channel AnalyzerChannel) {
	channel = make(AnalyzerChannel, chanBuf)
	go func() {
		defer close(channel)

		for samples > 0 {
			packet, ok := <-c
			if !ok {
				return
			}

			if len(packet.Samples)/packet.Channels > samples {
				startIndex := samples * packet.Channels
				channel <- &AnalyzerPacket{
					Samples:    packet.Samples[startIndex:],
					Channels:   packet.Channels,
					SampleRate: packet.SampleRate,
					BitDepth:   packet.BitDepth,
				}
				samples = 0
				break
			} else {
				samples -= len(packet.Samples) / packet.Channels
				continue
			}
		}

		for packet := range c {
			channel <- packet
		}
	}()

	return
}

func (c AnalyzerChannel) SkipEndSamples(samples int) (channel AnalyzerChannel) {
	channel = make(AnalyzerChannel, chanBuf)
	go func() {
		defer close(channel)

		var buffer []*AnalyzerPacket
		bufferSamples := 0

		for packet := range c {
			for len(buffer) > 0 && (bufferSamples-len(buffer[0].Samples)/buffer[0].Channels) > samples {
				channel <- buffer[0]
				bufferSamples -= len(buffer[0].Samples) / buffer[0].Channels
				buffer = buffer[1:]
			}

			bufferSamples += len(packet.Samples) / packet.Channels
			buffer = append(buffer, packet)
		}

		for _, packet := range buffer {
			//TODO: check this
			leftSamples := bufferSamples - len(packet.Samples)/packet.Channels

			if leftSamples <= samples {
				endIndex := len(packet.Samples) - (samples-leftSamples)*packet.Channels
				channel <- &AnalyzerPacket{
					Samples:    packet.Samples[:endIndex],
					Channels:   packet.Channels,
					SampleRate: packet.SampleRate,
					BitDepth:   packet.BitDepth,
				}
				samples = 0
				break
			} else {
				channel <- packet
				bufferSamples -= len(packet.Samples) / packet.Channels
			}
		}
	}()

	return
}

func (c AnalyzerChannel) SkipEndSamplesMultiple(wg *sync.WaitGroup, offset *atomic.Uint32, samples int) (channel AnalyzerChannel) {
	channel = make(AnalyzerChannel, chanBuf)
	go func() {
		defer close(channel)

		var buffer []*AnalyzerPacket
		bufferSamples := 0

		maxSamples := samples * 2

		samplesRead := 0
		for packet := range c {
			for len(buffer) > 0 && (bufferSamples-len(buffer[0].Samples)/buffer[0].Channels) > maxSamples {
				channel <- buffer[0]
				samplesRead += len(buffer[0].Samples) / buffer[0].Channels
				bufferSamples -= len(buffer[0].Samples) / buffer[0].Channels
				buffer = buffer[1:]
			}

			bufferSamples += len(packet.Samples) / packet.Channels
			buffer = append(buffer, packet)
		}

		wg.Wait()
		totalSampleOffset := samplesRead + int(offset.Load())

		if len(buffer) > 0 {
			p := &AnalyzerPacket{
				Channels:   buffer[0].Channels,
				SampleRate: buffer[0].SampleRate,
				BitDepth:   buffer[0].BitDepth,
			}
			for _, packet := range buffer {
				p.Samples = append(p.Samples, packet.Samples...)
			}
			nsamples := samples + (((len(p.Samples) / p.Channels) + totalSampleOffset) % samples)

			if len(p.Samples)/p.Channels > nsamples {
				endIndex := len(p.Samples) - nsamples*p.Channels
				channel <- &AnalyzerPacket{
					Samples:    p.Samples[:endIndex],
					Channels:   p.Channels,
					SampleRate: p.SampleRate,
					BitDepth:   p.BitDepth,
				}
			}
		}
	}()

	return
}

func NewHasherAudioGap(samples, sampleRate, channels, bitDepth int) (channel AnalyzerChannel) {
	channel = make(AnalyzerChannel, 1)
	defer close(channel)
	channel <- &AnalyzerPacket{
		Samples:    make([]int32, samples*channels),
		Channels:   channels,
		SampleRate: sampleRate,
		BitDepth:   bitDepth,
	}

	return
}

func MergeHasherChannels(channels ...AnalyzerChannel) (channel AnalyzerChannel) {
	channel = make(AnalyzerChannel, chanBuf)

	go func() {
		defer close(channel)

		for _, c := range channels {
			for packet := range c {
				channel <- packet
			}
		}
	}()

	return
}