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>
  Saxon Nelson-Milton <saxon@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 provides a h264 bytestream lexer and RTP H264 access unit
// extracter.
package h264

import (
	"io"
	"time"

	"bitbucket.org/ausocean/av/codec/codecutil"
)

var noDelay = make(chan time.Time)

func init() {
	close(noDelay)
}

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

// Lex 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, n int) 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 := codecutil.NewByteScanner(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.ScanUntil(buf, 0x00)
		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
}