METANOIA/metadata/hasher.go

package metadata

import (
	"encoding/binary"
	"git.gammaspectra.live/S.O.N.G/Hibiki/utilities/audio/format"
	"github.com/minio/sha256-simd"
	"hash"
	"hash/crc32"
	"sync"
	"sync/atomic"
	"time"
)

const chanBuf = 16

type HasherChannel chan *format.AnalyzerPacket

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

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

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

	return
}

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

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

func (c HasherChannel) SkipStartSamples(samples int) (channel HasherChannel) {
	channel = make(HasherChannel, 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 <- &format.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 HasherChannel) SkipEndSamples(samples int) (channel HasherChannel) {
	channel = make(HasherChannel, chanBuf)
	go func() {
		defer close(channel)

		var buffer []*format.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 <- &format.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 HasherChannel) SkipEndSamplesMultiple(wg *sync.WaitGroup, offset *uint32, samples int) (channel HasherChannel) {
	channel = make(HasherChannel, chanBuf)
	go func() {
		defer close(channel)

		var buffer []*format.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(atomic.LoadUint32(offset))

		if len(buffer) > 0 {
			p := &format.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 <- &format.AnalyzerPacket{
					Samples:    p.Samples[:endIndex],
					Channels:   p.Channels,
					SampleRate: p.SampleRate,
					BitDepth:   p.BitDepth,
				}
			}
		}
	}()

	return
}

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

	return
}

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

	go func() {
		defer close(channel)

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

	return
}

type HashType int

const (
	HashtypeCrc32 = HashType(iota)
	HashtypeSha256
	HashtypeAccurateRipV1
	HashtypeAccurateRipV1Start
	HashtypeAccurateRipV2
	HashtypeAccurateRipV2Start
)

type Hasher struct {
	hash       HashType
	hasher     hash.Hash
	result     []byte
	channel    HasherChannel
	wg         sync.WaitGroup
	samples    int
	duration   float64
	sampleRate int
	bitDepth   int
	channels   int
	buffer     [][]int32
}

func NewHasher(channel HasherChannel, hashType HashType) (h *Hasher) {
	h = &Hasher{
		hash:    hashType,
		channel: channel,
	}

	switch hashType {
	case HashtypeCrc32:
		h.hasher = crc32.NewIEEE()
	case HashtypeSha256:
		h.hasher = sha256.New()
	case HashtypeAccurateRipV1:
		h.hasher = NewAccurateRipV1(0)
	case HashtypeAccurateRipV1Start:
		h.hasher = NewAccurateRipV1(Int16SamplesPerSector*5 - 1)
	case HashtypeAccurateRipV2:
		h.hasher = NewAccurateRipV2(0)
	case HashtypeAccurateRipV2Start:
		h.hasher = NewAccurateRipV2(Int16SamplesPerSector*5 - 1)

	}

	h.startRoutine()

	return
}

func (h *Hasher) startRoutine() {
	h.wg.Add(1)
	go func() {
		defer h.wg.Done()

		for packet := range h.channel {
			h.handlePacket(packet)
		}

		h.result = h.hasher.Sum([]byte{})

	}()
}

func (h *Hasher) handlePacket(packet *format.AnalyzerPacket) {
	samples := len(packet.Samples) / packet.Channels

	h.samples += samples

	if h.sampleRate == 0 {
		h.sampleRate = packet.SampleRate
	} else if h.sampleRate != packet.SampleRate {
		h.sampleRate = -1
	}
	if h.bitDepth == 0 {
		h.bitDepth = packet.BitDepth
	} else if h.bitDepth != packet.BitDepth {
		h.bitDepth = -1
	}
	if h.channels == 0 {
		h.channels = packet.Channels
	} else if h.channels != packet.Channels {
		h.channels = -1
	}

	h.duration += float64(samples) / float64(packet.SampleRate)

	var buf []byte
	switch packet.BitDepth {
	case 8:
		buf = make([]byte, len(packet.Samples))
		for i := range packet.Samples {
			buf[i] = byte(packet.Samples[i])
		}
	case 16:
		buf = make([]byte, len(packet.Samples)*2)
		for i := range packet.Samples {
			binary.LittleEndian.PutUint16(buf[i*2:], uint16(int16(packet.Samples[i])))
		}
	case 24:
		buf = make([]byte, len(packet.Samples)*3)
		for i := range packet.Samples {
			buf[i*3] = byte((packet.Samples[i] >> 16) & 0xFF)
			buf[i*3+1] = byte((packet.Samples[i] >> 8) & 0xFF)
			buf[i*3+2] = byte(packet.Samples[i] & 0xFF)
		}
	default:
		buf = make([]byte, len(packet.Samples)*4)
		for i := range packet.Samples {
			binary.LittleEndian.PutUint32(buf[i*4:], uint32(packet.Samples[i]))
		}

	}

	h.hasher.Write(buf)
}

func (h *Hasher) GetSampleCount() int {
	return h.samples
}

func (h *Hasher) GetChannels() int {
	return h.channels
}

func (h *Hasher) GetSampleRate() int {
	return h.sampleRate
}

func (h *Hasher) GetHashType() HashType {
	return h.hash
}

func (h *Hasher) GetResult() []byte {
	return h.result
}

func (h *Hasher) GetDuration() time.Duration {
	if h.sampleRate > 0 {
		return time.Duration(float64(time.Second) * (float64(h.samples) / float64(h.sampleRate)))
	}

	//Fallback calculated duration
	return time.Duration(float64(time.Second) * h.duration)
}

func (h *Hasher) GetWaitGroup() *sync.WaitGroup {
	return &h.wg
}

func (h *Hasher) Wait() {
	h.wg.Wait()
}