vosk-recognize-dev/pkg/riff/parser.go

package riff

import (
    "encoding/binary"
    "errors"
    "fmt"
    "io"
    "sync"
    "time"
)

type Parser struct {
    r io.Reader
    Chan chan *Chunk
    ChunkParserTimeout time.Duration
    Wg *sync.WaitGroup
    ID [4]byte
    Size uint32
    Format [4]byte
    wavHeaderSize uint32
    WavAudioFormat uint16
    NumChannels uint16
    SampleRate uint32
    AvgBytesPerSec uint32
    BlockAlign uint16
    BitsPerSample uint16
}

func (c *Parser) ParseHeaders() error {
    id, size, err := c.IDnSize()
    if err != nil {
        return err
    }
    c.ID = id
    if c.ID != RiffID {
        return fmt.Errorf("%s - %s", c.ID, ErrFmtNotSupported)
    }
    c.Size = size
    if err := binary.Read(c.r, binary.BigEndian, &c.Format); err != nil {
        return err
    }
    return nil
}

func (c *Parser) Duration() (time.Duration, error) {
    if c == nil {
        return 0, errors.New("can't calculate the duration of a nil pointer")
    }
    if c.ID == [4]byte{} || c.AvgBytesPerSec == 0 {
        err := c.Parse()
        if err != nil {
            return 0, nil
        }
    }
    switch c.Format {
        case WavFormatID:
            return c.wavDuration()
        default:
            return 0, ErrFmtNotSupported
    }
}

func (c *Parser) String() string {
    out := fmt.Sprintf("Format: %s - ", c.Format)
    if c.Format == WavFormatID {
        out += fmt.Sprintf("%d channels @ %d / %d bits - ", c.NumChannels, c.SampleRate, c.BitsPerSample)
        d, _ := c.Duration()
        out += fmt.Sprintf("Duration: %f seconds", d.Seconds())
    }
    return out
}

func (c *Parser) NextChunk() (*Chunk, error) {
    if c == nil {
        return nil, errors.New("can't calculate the duration of a nil pointer")
    }
    id, size, err := c.IDnSize()
    if err != nil {
        return nil, err
    }
    if size%2 == 1 {
        size++
    }
    ch := &Chunk{
        ID:   id,
        Size: int(size),
        R:    c.r,
    }
    return ch, nil
}

func (c *Parser) IDnSize() ([4]byte, uint32, error) {
    var ID [4]byte
    var blockSize uint32
    if err := binary.Read(c.r, binary.BigEndian, &ID); err != nil {
        return ID, blockSize, err
    }
    if err := binary.Read(c.r, binary.LittleEndian, &blockSize); err != err {
        return ID, blockSize, err
    }
    return ID, blockSize, nil
}

func (p *Parser) Parse() error {
    if p == nil {
        return errors.New("can't calculate the wav duration of a nil pointer")
    }
    if p.Size == 0 {
        id, size, err := p.IDnSize()
        if err != nil {
            return err
        }
        p.ID = id
        if p.ID != RiffID {
            return fmt.Errorf("%s - %s", p.ID, ErrFmtNotSupported)
        }
        p.Size = size
        if err := binary.Read(p.r, binary.BigEndian, &p.Format); err != nil {
            return err
        }
    }
    var chunk *Chunk
    var err error
    for err == nil {
        chunk, err = p.NextChunk()
        if err != nil {
            break
        }
        if chunk.ID == FmtID {
            chunk.DecodeWavHeader(p)
        } else {
            if p.Chan != nil {
                if chunk.Wg == nil {
                    chunk.Wg = p.Wg
                }
                chunk.Wg.Add(1)
                p.Chan <- chunk
                chunk.Wg.Wait()
            }
        }
        if !chunk.IsFullyRead() {
            chunk.Drain()
        }
    }
    if p.Wg != nil {
        p.Wg.Wait()
    }
    if p.Chan != nil {
        close(p.Chan)
    }
    if err == io.EOF {
        return nil
    }
    return err
}

func (p *Parser) wavDuration() (time.Duration, error) {
    if p.Size == 0 || p.AvgBytesPerSec == 0 {
        return 0, fmt.Errorf("can't extract the duration due to the file not properly parsed")
    }
    duration := time.Duration((float64(p.Size) / float64(p.AvgBytesPerSec)) * float64(time.Second))
    return duration, nil
}

func (p *Parser) jumpTo(bytesAhead int) error {
    var err error
    for bytesAhead > 0 {
        readSize := bytesAhead
        if readSize > 4000 {
            readSize = 4000
        }
        buf := make([]byte, readSize)
        err = binary.Read(p.r, binary.LittleEndian, &buf)
        if err != nil {
            return nil
        }
        bytesAhead -= readSize
    }
    return nil
}