av/codec/h264/lex.go

/*
NAME
  lex.go

DESCRIPTION
  lex.go provides a lexer to lex h264 bytestream into access units.

AUTHOR
  Dan Kortschak <dan@ausocean.org>

LICENSE
  lex.go is Copyright (C) 2017 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package h264

import (
	"io"
	"time"
)

var noDelay = make(chan time.Time)

func init() {
	close(noDelay)
}

var h264Prefix = [...]byte{0x00, 0x00, 0x01, 0x09, 0xf0}

// H264 lexes H.264 NAL units read from src into separate writes to dst with
// successive writes being performed not earlier than the specified delay.
// NAL units are split after type 1 (Coded slice of a non-IDR picture), 5
// (Coded slice of a IDR picture) and 8 (Picture parameter set).
func Lex(dst io.Writer, src io.Reader, delay time.Duration) error {
	var tick <-chan time.Time
	if delay == 0 {
		tick = noDelay
	} else {
		ticker := time.NewTicker(delay)
		defer ticker.Stop()
		tick = ticker.C
	}

	const bufSize = 8 << 10

	c := newScanner(src, make([]byte, 4<<10)) // Standard file buffer size.

	buf := make([]byte, len(h264Prefix), bufSize)
	copy(buf, h264Prefix[:])
	writeOut := false
outer:
	for {
		var b byte
		var err error
		buf, b, err = c.scanUntilZeroInto(buf)
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}

		for n := 1; b == 0x0 && n < 4; n++ {
			b, err = c.readByte()
			if err != nil {
				if err != io.EOF {
					return err
				}
				break outer
			}
			buf = append(buf, b)

			if b != 0x1 || (n != 2 && n != 3) {
				continue
			}

			if writeOut {
				<-tick
				_, err := dst.Write(buf[:len(buf)-(n+1)])
				if err != nil {
					return err
				}
				buf = make([]byte, len(h264Prefix)+n, bufSize)
				copy(buf, h264Prefix[:])
				buf = append(buf, 1)
				writeOut = false
			}

			b, err = c.readByte()
			if err != nil {
				if err != io.EOF {
					return err
				}
				break outer
			}
			buf = append(buf, b)

			// http://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.264-200305-S!!PDF-E&type=items
			// Table 7-1 NAL unit type codes
			const (
				nonIdrPic   = 1
				idrPic      = 5
				suppEnhInfo = 6
				paramSet    = 8
			)
			switch nalTyp := b & 0x1f; nalTyp {
			case nonIdrPic, idrPic, paramSet, suppEnhInfo:
				writeOut = true
			}
		}
	}
	if len(buf) == len(h264Prefix) {
		return nil
	}
	<-tick
	_, err := dst.Write(buf)
	return err
}

// scanner is a byte scanner.
type scanner struct {
	buf []byte
	off int

	// r is the source of data for the scanner.
	r io.Reader
}

// newScanner returns a scanner initialised with an io.Reader and a read buffer.
func newScanner(r io.Reader, buf []byte) *scanner {
	return &scanner{r: r, buf: buf[:0]}
}

// scanUntilZeroInto scans the scanner's underlying io.Reader until a zero byte
// has been read, appending all read bytes to dst. The resulting appended data,
// the last read byte and whether the last read byte was zero are returned.
func (c *scanner) scanUntilZeroInto(dst []byte) (res []byte, b byte, err error) {
outer:
	for {
		var i int
		for i, b = range c.buf[c.off:] {
			if b != 0x0 {
				continue
			}
			dst = append(dst, c.buf[c.off:c.off+i+1]...)
			c.off += i + 1
			break outer
		}
		dst = append(dst, c.buf[c.off:]...)
		err = c.reload()
		if err != nil {
			break
		}
	}
	return dst, b, err
}

// readByte is an unexported ReadByte.
func (c *scanner) readByte() (byte, error) {
	if c.off >= len(c.buf) {
		err := c.reload()
		if err != nil {
			return 0, err
		}
	}
	b := c.buf[c.off]
	c.off++
	return b, nil
}

// reload re-fills the scanner's buffer.
func (c *scanner) reload() error {
	n, err := c.r.Read(c.buf[:cap(c.buf)])
	c.buf = c.buf[:n]
	if err != nil {
		if err != io.EOF {
			return err
		}
		if n == 0 {
			return io.EOF
		}
	}
	c.off = 0
	return nil
}