srs/trunk/3rdparty/srs-bench/vendor/github.com/yapingcat/gomedia/mpeg2/ts-muxer.go

package mpeg2

import (
    "errors"

    "github.com/yapingcat/gomedia/codec"
)

type pes_stream struct {
    pid        uint16
    cc         uint8
    streamtype TS_STREAM_TYPE
}

func NewPESStream(pid uint16, cid TS_STREAM_TYPE) *pes_stream {
    return &pes_stream{
        pid:        pid,
        cc:         0,
        streamtype: cid,
    }
}

type table_pmt struct {
    pid            uint16
    cc             uint8
    pcr_pid        uint16
    version_number uint8
    pm             uint16
    streams        []*pes_stream
}

func NewTablePmt() *table_pmt {
    return &table_pmt{
        pid:            0,
        cc:             0,
        pcr_pid:        0,
        version_number: 0,
        pm:             0,
        streams:        make([]*pes_stream, 0, 2),
    }
}

type table_pat struct {
    cc             uint8
    version_number uint8
    pmts           []*table_pmt
}

func NewTablePat() *table_pat {
    return &table_pat{
        cc:             0,
        version_number: 0,
        pmts:           make([]*table_pmt, 0, 8),
    }
}

type TSMuxer struct {
    pat        *table_pat
    stream_pid uint16
    pmt_pid    uint16
    pat_period uint64
    OnPacket   func(pkg []byte)
}

func NewTSMuxer() *TSMuxer {
    return &TSMuxer{
        pat:        NewTablePat(),
        stream_pid: 0x100,
        pmt_pid:    0x200,
        pat_period: 0,
        OnPacket:   nil,
    }
}

func (mux *TSMuxer) AddStream(cid TS_STREAM_TYPE) uint16 {
    if mux.pat == nil {
        mux.pat = NewTablePat()
    }
    if len(mux.pat.pmts) == 0 {
        tmppmt := NewTablePmt()
        tmppmt.pid = mux.pmt_pid
        tmppmt.pm = 1
        mux.pmt_pid++
        mux.pat.pmts = append(mux.pat.pmts, tmppmt)
    }
    sid := mux.stream_pid
    tmpstream := NewPESStream(sid, cid)
    mux.stream_pid++
    mux.pat.pmts[0].streams = append(mux.pat.pmts[0].streams, tmpstream)
    return sid
}

/// Muxer audio/video stream data
/// pid: stream id by AddStream
/// pts: audio/video stream timestamp in ms
/// dts: audio/video stream timestamp in ms
func (mux *TSMuxer) Write(pid uint16, data []byte, pts uint64, dts uint64) error {
    var whichpmt *table_pmt = nil
    var whichstream *pes_stream = nil
    for _, pmt := range mux.pat.pmts {
        for _, stream := range pmt.streams {
            if stream.pid == pid {
                whichpmt = pmt
                whichstream = stream
                break
            }
        }
    }
    if whichpmt == nil || whichstream == nil {
        return errors.New("not Found pid stream")
    }
    if whichpmt.pcr_pid == 0 || (findPESIDByStreamType(whichstream.streamtype) == PES_STREAM_VIDEO && whichpmt.pcr_pid != pid) {
        whichpmt.pcr_pid = pid
    }

    var withaud bool = false

    if whichstream.streamtype == TS_STREAM_H264 || whichstream.streamtype == TS_STREAM_H265 {
        codec.SplitFrame(data, func(nalu []byte) bool {
            if whichstream.streamtype == TS_STREAM_H264 {
                nalu_type := codec.H264NaluTypeWithoutStartCode(nalu)
                if nalu_type == codec.H264_NAL_AUD {
                    withaud = true
                    return false
                } else if codec.IsH264VCLNaluType(nalu_type) {
                    return false
                }
                return true
            } else {
                nalu_type := codec.H265NaluTypeWithoutStartCode(nalu)
                if nalu_type == codec.H265_NAL_AUD {
                    withaud = true
                    return false
                } else if codec.IsH265VCLNaluType(nalu_type) {
                    return false
                }
                return true
            }
        })
    }

    if mux.pat_period == 0 || mux.pat_period+400 < dts {
        mux.pat_period = dts
        if mux.pat_period == 0 {
            mux.pat_period = 1 //avoid write pat twice
        }
        tmppat := NewPat()
        tmppat.Version_number = mux.pat.version_number
        for _, pmt := range mux.pat.pmts {
            tmppm := PmtPair{
                Program_number: pmt.pm,
                PID:            pmt.pid,
            }
            tmppat.Pmts = append(tmppat.Pmts, tmppm)
        }
        mux.writePat(tmppat)

        for _, pmt := range mux.pat.pmts {
            tmppmt := NewPmt()
            tmppmt.Program_number = pmt.pm
            tmppmt.Version_number = pmt.version_number
            tmppmt.PCR_PID = pmt.pcr_pid
            for _, stream := range pmt.streams {
                var sp StreamPair
                sp.StreamType = uint8(stream.streamtype)
                sp.Elementary_PID = stream.pid
                sp.ES_Info_Length = 0
                tmppmt.Streams = append(tmppmt.Streams, sp)
            }
            mux.writePmt(tmppmt, pmt)
        }
    }

    flag := false
    switch whichstream.streamtype {
    case TS_STREAM_H264:
        flag = codec.IsH264IDRFrame(data)
    case TS_STREAM_H265:
        flag = codec.IsH265IDRFrame(data)
    }

    mux.writePES(whichstream, whichpmt, data, pts*90, dts*90, flag, withaud)
    return nil
}

func (mux *TSMuxer) writePat(pat *Pat) {
    var tshdr TSPacket
    tshdr.Payload_unit_start_indicator = 1
    tshdr.PID = 0
    tshdr.Adaptation_field_control = 0x01
    tshdr.Continuity_counter = mux.pat.cc
    mux.pat.cc++
    mux.pat.cc = (mux.pat.cc + 1) % 16
    bsw := codec.NewBitStreamWriter(TS_PAKCET_SIZE)
    tshdr.EncodeHeader(bsw)
    bsw.PutByte(0x00) //pointer
    pat.Encode(bsw)
    bsw.FillRemainData(0xff)
    if mux.OnPacket != nil {
        mux.OnPacket(bsw.Bits())
    }
}

func (mux *TSMuxer) writePmt(pmt *Pmt, t_pmt *table_pmt) {
    var tshdr TSPacket
    tshdr.Payload_unit_start_indicator = 1
    tshdr.PID = t_pmt.pid
    tshdr.Adaptation_field_control = 0x01
    tshdr.Continuity_counter = t_pmt.cc
    t_pmt.cc = (t_pmt.cc + 1) % 16
    bsw := codec.NewBitStreamWriter(TS_PAKCET_SIZE)
    tshdr.EncodeHeader(bsw)
    bsw.PutByte(0x00) //pointer
    pmt.Encode(bsw)
    bsw.FillRemainData(0xff)
    if mux.OnPacket != nil {
        mux.OnPacket(bsw.Bits())
    }
}

func (mux *TSMuxer) writePES(pes *pes_stream, pmt *table_pmt, data []byte, pts uint64, dts uint64, idr_flag bool, withaud bool) {
    var firstPesPacket bool = true
    bsw := codec.NewBitStreamWriter(TS_PAKCET_SIZE)
    for {
        bsw.Reset()
        var tshdr TSPacket
        if firstPesPacket {
            tshdr.Payload_unit_start_indicator = 1
        }
        tshdr.PID = pes.pid
        tshdr.Adaptation_field_control = 0x01
        tshdr.Continuity_counter = pes.cc
        headlen := 4
        pes.cc = (pes.cc + 1) % 16
        var adaptation *Adaptation_field = nil
        if firstPesPacket && idr_flag {
            adaptation = new(Adaptation_field)
            tshdr.Adaptation_field_control = tshdr.Adaptation_field_control | 0x20
            adaptation.Random_access_indicator = 1
            headlen += 2
        }

        if firstPesPacket && pes.pid == pmt.pcr_pid {
            if adaptation == nil {
                adaptation = new(Adaptation_field)
                headlen += 2
            }
            tshdr.Adaptation_field_control = tshdr.Adaptation_field_control | 0x20
            adaptation.PCR_flag = 1
            var pcr_base uint64 = 0
            var pcr_ext uint16 = 0
            if dts == 0 {
                pcr_base = pts * 300 / 300
                pcr_ext = uint16(pts * 300 % 300)
            } else {
                pcr_base = dts * 300 / 300
                pcr_ext = uint16(dts * 300 % 300)
            }
            adaptation.Program_clock_reference_base = pcr_base
            adaptation.Program_clock_reference_extension = pcr_ext
            headlen += 6
        }

        var payload []byte
        var pespkg *PesPacket = nil
        if firstPesPacket {
            oldheadlen := headlen
            headlen += 19
            if !withaud && pes.streamtype == TS_STREAM_H264 {
                headlen += 6
                payload = append(payload, H264_AUD_NALU...)
            } else if !withaud && pes.streamtype == TS_STREAM_H265 {
                payload = append(payload, H265_AUD_NALU...)
                headlen += 7
            }
            pespkg = NewPesPacket()
            pespkg.PTS_DTS_flags = 0x03
            pespkg.PES_header_data_length = 10
            pespkg.Pts = pts
            pespkg.Dts = dts
            pespkg.Stream_id = uint8(findPESIDByStreamType(pes.streamtype))
            if idr_flag {
                pespkg.Data_alignment_indicator = 1
            }
            if headlen-oldheadlen-6+len(data) > 0xFFFF {
                pespkg.PES_packet_length = 0
            } else {
                pespkg.PES_packet_length = uint16(len(data) + headlen - oldheadlen - 6)
            }

        }

        if len(data)+headlen < TS_PAKCET_SIZE {
            if adaptation == nil {
                adaptation = new(Adaptation_field)
                headlen += 1
                if TS_PAKCET_SIZE-len(data)-headlen >= 1 {
                    headlen += 1
                } else {
                    adaptation.SingleStuffingByte = true
                }
            }
            adaptation.Stuffing_byte = uint8(TS_PAKCET_SIZE - len(data) - headlen)
            payload = append(payload, data...)
            data = data[:0]
        } else {
            payload = append(payload, data[0:TS_PAKCET_SIZE-headlen]...)
            data = data[TS_PAKCET_SIZE-headlen:]
        }

        if adaptation != nil {
            tshdr.Field = adaptation
            tshdr.Adaptation_field_control |= 0x02
        }
        tshdr.EncodeHeader(bsw)
        if pespkg != nil {
            pespkg.Pes_payload = payload
            pespkg.Encode(bsw)
        } else {
            bsw.PutBytes(payload)
        }
        firstPesPacket = false
        if mux.OnPacket != nil {
            if len(bsw.Bits()) != TS_PAKCET_SIZE {
                panic("packet ts packet failed")
            }
            mux.OnPacket(bsw.Bits())
        }
        if len(data) == 0 {
            break
        }
    }
}