<commit_before>package main

import (
	"errors"
	"sync"
	"time"
)

const (
	// a special value for a tag's length than means "until the next tag of same level".
	ebmlIndeterminate = 0xFFFFFFFFFFFFFF
	// some of the possible tag ids.
	// https://www.matroska.org/technical/specs/index.html
	ebmlTagVoid            = 0xEC
	ebmlTagEBML            = 0x1A45DFA3
	ebmlTagSegment         = 0x18538067
	ebmlTagSeekHead        = 0x114D9B74
	ebmlTagInfo            = 0x1549A966
	ebmlTagTimecodeScale   = 0x2AD7B1
	ebmlTagDuration        = 0x4489
	ebmlTagDateUTC         = 0x4461
	ebmlTagMuxingApp       = 0x4D80
	ebmlTagWritingApp      = 0x5741
	ebmlTagTracks          = 0x1654AE6B
	ebmlTagTrackEntry      = 0xAE
	ebmlTagTrackNumber     = 0xD7
	ebmlTagTrackUID        = 0x73C5
	ebmlTagTrackType       = 0x83
	ebmlTagFlagEnabled     = 0xB9
	ebmlTagFlagDefault     = 0x88
	ebmlTagFlagForced      = 0x55AA
	ebmlTagFlagLacing      = 0x9C
	ebmlTagDefaultDuration = 0x23E383
	ebmlTagName            = 0x536E
	ebmlTagCodecID         = 0x86
	ebmlTagCodecName       = 0x228688
	ebmlTagVideo           = 0xE0
	ebmlTagPixelWidth      = 0xB0
	ebmlTagPixelHeight     = 0xBA
	ebmlTagAudio           = 0xE1
	ebmlTagCluster         = 0x1F43B675
	ebmlTagTimecode        = 0xE7
	ebmlTagPrevSize        = 0xAB
	ebmlTagSimpleBlock     = 0xA3
	ebmlTagBlockGroup      = 0xA0
	ebmlTagBlock           = 0xA1
	ebmlTagBlockDuration   = 0x9B
	ebmlTagReferenceBlock  = 0xFB
	ebmlTagDiscardPadding  = 0x75A2
	ebmlTagCues            = 0x1C53BB6B
	ebmlTagChapters        = 0x1043A770
	ebmlTagTags            = 0x1254C367
	ebmlTagTag             = 0x7373
	ebmlTagTargets         = 0x63C0
	ebmlTagTargetType      = 0x63CA
	ebmlTagTagTrackUID     = 0x63C5
	ebmlTagSimpleTag       = 0x67C8
	ebmlTagTagName         = 0x45A3
	ebmlTagTagLanguage     = 0x447A
	ebmlTagTagDefault      = 0x4484
	ebmlTagTagString       = 0x4487
	ebmlTagTagBinary       = 0x4485
)

var ebmlIndeterminateCoding = [...]uint64{
	0, // these values in the "length" field all decode to `ebmlIndeterminate`.
	0xFF,
	0x7FFF,
	0x3FFFFF,
	0x1FFFFFFF,
	0x0FFFFFFFFF,
	0x07FFFFFFFFFF,
	0x03FFFFFFFFFFFF,
	0x01FFFFFFFFFFFFFF,
}

func fixedUint(data []byte) uint64 {
	var x uint64 = 0
	for _, b := range data {
		x = x<<8 | uint64(b)
	}
	return x
}

func ebmlTagID(data []byte) (uint64, int) {
	if len(data) != 0 && data[0] != 0 {
		// 1xxxxxxx
		// 01xxxxxx xxxxxxxx
		// 001xxxxx xxxxxxxx xxxxxxxx
		// ...
		// 00000001 xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx
		//        ^---- this length marker is included in tag ids but not in other ints
		consumed := 9
		for b := data[0]; b != 0; b >>= 1 {
			consumed -= 1
		}
		if len(data) >= consumed {
			return fixedUint(data[:consumed]), consumed
		}
	}
	return 0, 0
}

func ebmlUint(data []byte) (uint64, int) {
	id, consumed := ebmlTagID(data)
	if ebmlIndeterminateCoding[consumed] == id {
		return ebmlIndeterminate, consumed
	}
	return id & ^(1 << uint(7*consumed)), consumed
}

type ebmlTag struct {
	Consumed int
	ID       uint
	Length   uint64
}

func ebmlParseTagIncomplete(data []byte) ebmlTag {
	if id, off := ebmlTagID(data); off != 0 {
		if length, off2 := ebmlUint(data[off:]); off2 != 0 {
			return ebmlTag{off + off2, uint(id), length}
		}
	}
	return ebmlTag{0, 0, 0}
}

func ebmlParseTag(data []byte) ebmlTag {
	if tag := ebmlParseTagIncomplete(data); tag.Length+uint64(tag.Consumed) <= uint64(len(data)) {
		return tag
	}
	return ebmlTag{0, 0, 0}
}

func (t ebmlTag) Contents(data []byte) []byte {
	return data[t.Consumed : uint64(t.Consumed)+t.Length]
}

func (t ebmlTag) Skip(data []byte) []byte {
	return data[uint64(t.Consumed)+t.Length:]
}

type frame struct {
	buf   []byte // Either a Block(Group) or a Cluster.
	track uint64 // 64 for a Cluster (track masks are 32-bit, so streams with a real 64-th track are rejected)
	key   bool
}

type framebuffer struct {
	data        []frame
	start       int
	headCluster []byte // Last Cluster popped off the ring. Blocks at the head still belong to it.
}

func (fb *framebuffer) PushCluster(buf []byte) {
	fb.PushFrame(frame{buf, 64, false})
}

func (fb *framebuffer) PushFrame(packed frame) {
	if len(fb.data) == cap(fb.data) {
		if fb.data[fb.start].track == 64 {
			fb.headCluster = fb.data[fb.start].buf
		}
		fb.data[fb.start] = packed
		fb.start = (fb.start + 1) % len(fb.data)
	} else {
		fb.data = append(fb.data, packed)
	}
}

func (fb *framebuffer) Read(cb func(cluster []byte, forceCluster bool, packed frame)) {
	cluster, forceCluster := fb.headCluster, true
	for i, s, n := 0, fb.start, len(fb.data); i < n; i++ {
		f := fb.data[(i+s)%n]
		if f.track == 64 {
			cluster = f.buf
			forceCluster = true
		} else if cluster != nil {
			cb(cluster, forceCluster, f)
			forceCluster = false
		}
	}
}

type viewer struct {
	// This function may return `false` to signal that it cannot write any more data.
	// The stream will resynchronize at next keyframe.
	write func(data []byte) bool
	// Viewers may hop between streams, but should only receive headers once.
	// This includes track info, as codecs must stay the same between segments.
	skipHeaders bool
	// We group blocks into indeterminate-length clusters. So long as
	// the cluster's timecode has not changed, there's no need to start a new one.
	skipCluster bool
	// Bit vector of tracks for which the viewer has both reference frames
	// (the previous frame and the last keyframe.)
	seenKeyframes uint32
}

func (cb *viewer) WriteFrame(cluster []byte, forceCluster bool, packed frame) {
	trackMask := uint32(1) << packed.track
	if packed.key {
		cb.seenKeyframes |= trackMask
	}
	if cb.seenKeyframes&trackMask != 0 {
		if !cb.skipCluster || forceCluster {
			cb.skipCluster = cb.write(cluster)
		}
		if !cb.skipCluster || !cb.write(packed.buf) {
			cb.seenKeyframes &= ^trackMask
		}
	}
}

type BroadcastSet struct {
	mutex   sync.Mutex
	streams map[string]*Broadcast
	// How long to keep a stream alive after a call to `Close`.
	Timeout time.Duration
	// Called right after a stream is destroyed. (`Timeout` seconds after a `Close`.)
	OnStreamClose     func(id string)
	OnStreamTrackInfo func(id string, info *StreamTrackInfo)
}

type Broadcast struct {
	StreamTrackInfo
	closing time.Duration
	Closed  bool
	dirty   bool // (Has unseen data in `StreamTrackInfo`.)
	buffer  []byte
	header  []byte // The EBML (DocType) tag.
	tracks  []byte // The beginning of the Segment (Tracks + Info).
	frames  framebuffer
	// outbound blocks must have monotonically increasing timecodes even if the inbound
	// stream restarts from the beginning, so we'll shift clusters' timecodes.
	blockTimecode       uint64
	recvClusterTimecode uint64
	sentClusterTimecode uint64
	timecodeShift       uint64
	// these values are for the whole stream, so they include audio and muxing overhead.
	// the latter is negligible, however, and the former is normally about 64k,
	// so also negligible. or at least predictable.
	rateUnit float64
	RateMean float64
	RateVar  float64

	vlock   sync.Mutex
	viewers map[chan<- []byte]*viewer
}

func (ctx *BroadcastSet) Readable(id string) (*Broadcast, bool) {
	if ctx.streams == nil {
		return nil, false
	}
	ctx.mutex.Lock()
	cast, ok := ctx.streams[id]
	ctx.mutex.Unlock()
	return cast, ok
}

func (ctx *BroadcastSet) Writable(id string) (*Broadcast, bool) {
	ctx.mutex.Lock()
	defer ctx.mutex.Unlock()
	if ctx.streams == nil {
		ctx.streams = make(map[string]*Broadcast)
	}
	if cast, ok := ctx.streams[id]; ok {
		if cast.closing == -1 {
			return nil, false
		}
		cast.closing = -1
		return cast, true
	}
	cast := Broadcast{
		closing:             -1,
		frames:              framebuffer{make([]frame, 0, 120), 0, nil},
		viewers:             make(map[chan<- []byte]*viewer),
		sentClusterTimecode: 0xFFFFFFFFFFFFFFFF,
	}
	ctx.streams[id] = &cast
	go func() {
		ticker := time.NewTicker(time.Second)
		for range ticker.C {
			if cast.dirty {
				cast.dirty = false
				ctx.OnStreamTrackInfo(id, &cast.StreamTrackInfo)
			}
			if cast.closing >= 0 {
				if cast.closing += time.Second; cast.closing > ctx.Timeout {
					break
				}
			}
			// exponentially weighted moving moments at a = 0.5
			//     avg[n] = a * x + (1 - a) * avg[n - 1]
			//     var[n] = a * (x - avg[n]) ** 2 / (1 - a) + (1 - a) * var[n - 1]
			cast.RateMean += cast.rateUnit / 2
			cast.RateVar += cast.rateUnit*cast.rateUnit - cast.RateVar/2
			cast.rateUnit = -cast.RateMean
		}
		ticker.Stop()

		ctx.mutex.Lock()
		delete(ctx.streams, id)
		ctx.mutex.Unlock()
		cast.Closed = true
		cast.vlock.Lock()
		for _, cb := range cast.viewers {
			cb.write([]byte{})
		}
		cast.vlock.Unlock()
		if ctx.OnStreamClose != nil {
			ctx.OnStreamClose(id)
		}
	}()
	return &cast, true
}

func (cast *Broadcast) Close() error {
	cast.closing = 0
	return nil
}

func (cast *Broadcast) Connect(ch chan<- []byte, skipHeaders bool) {
	write := func(data []byte) bool {
		// `Broadcast.Write` emits data in block-sized chunks.
		// Thus the buffer size is measured in frames, not bytes.
		if len(ch) == cap(ch) {
			return false
		}
		ch <- data
		return true
	}

	cast.vlock.Lock()
	cast.viewers[ch] = &viewer{write: write}
	cast.vlock.Unlock()
}

func (cast *Broadcast) Disconnect(ch chan<- []byte) {
	cast.vlock.Lock()
	delete(cast.viewers, ch)
	cast.vlock.Unlock()
}

func (cast *Broadcast) Reset() {
	cast.buffer = nil<commit_msg>Wait until buffer level drops to 50% before resyncing to keyframe.<commit_after>package main

import (
	"errors"
	"sync"
	"time"
)

const (
	// a special value for a tag's length than means "until the next tag of same level".
	ebmlIndeterminate = 0xFFFFFFFFFFFFFF
	// some of the possible tag ids.
	// https://www.matroska.org/technical/specs/index.html
	ebmlTagVoid            = 0xEC
	ebmlTagEBML            = 0x1A45DFA3
	ebmlTagSegment         = 0x18538067
	ebmlTagSeekHead        = 0x114D9B74
	ebmlTagInfo            = 0x1549A966
	ebmlTagTimecodeScale   = 0x2AD7B1
	ebmlTagDuration        = 0x4489
	ebmlTagDateUTC         = 0x4461
	ebmlTagMuxingApp       = 0x4D80
	ebmlTagWritingApp      = 0x5741
	ebmlTagTracks          = 0x1654AE6B
	ebmlTagTrackEntry      = 0xAE
	ebmlTagTrackNumber     = 0xD7
	ebmlTagTrackUID        = 0x73C5
	ebmlTagTrackType       = 0x83
	ebmlTagFlagEnabled     = 0xB9
	ebmlTagFlagDefault     = 0x88
	ebmlTagFlagForced      = 0x55AA
	ebmlTagFlagLacing      = 0x9C
	ebmlTagDefaultDuration = 0x23E383
	ebmlTagName            = 0x536E
	ebmlTagCodecID         = 0x86
	ebmlTagCodecName       = 0x228688
	ebmlTagVideo           = 0xE0
	ebmlTagPixelWidth      = 0xB0
	ebmlTagPixelHeight     = 0xBA
	ebmlTagAudio           = 0xE1
	ebmlTagCluster         = 0x1F43B675
	ebmlTagTimecode        = 0xE7
	ebmlTagPrevSize        = 0xAB
	ebmlTagSimpleBlock     = 0xA3
	ebmlTagBlockGroup      = 0xA0
	ebmlTagBlock           = 0xA1
	ebmlTagBlockDuration   = 0x9B
	ebmlTagReferenceBlock  = 0xFB
	ebmlTagDiscardPadding  = 0x75A2
	ebmlTagCues            = 0x1C53BB6B
	ebmlTagChapters        = 0x1043A770
	ebmlTagTags            = 0x1254C367
	ebmlTagTag             = 0x7373
	ebmlTagTargets         = 0x63C0
	ebmlTagTargetType      = 0x63CA
	ebmlTagTagTrackUID     = 0x63C5
	ebmlTagSimpleTag       = 0x67C8
	ebmlTagTagName         = 0x45A3
	ebmlTagTagLanguage     = 0x447A
	ebmlTagTagDefault      = 0x4484
	ebmlTagTagString       = 0x4487
	ebmlTagTagBinary       = 0x4485
)

var ebmlIndeterminateCoding = [...]uint64{
	0, // these values in the "length" field all decode to `ebmlIndeterminate`.
	0xFF,
	0x7FFF,
	0x3FFFFF,
	0x1FFFFFFF,
	0x0FFFFFFFFF,
	0x07FFFFFFFFFF,
	0x03FFFFFFFFFFFF,
	0x01FFFFFFFFFFFFFF,
}

func fixedUint(data []byte) uint64 {
	var x uint64 = 0
	for _, b := range data {
		x = x<<8 | uint64(b)
	}
	return x
}

func ebmlTagID(data []byte) (uint64, int) {
	if len(data) != 0 && data[0] != 0 {
		// 1xxxxxxx
		// 01xxxxxx xxxxxxxx
		// 001xxxxx xxxxxxxx xxxxxxxx
		// ...
		// 00000001 xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx
		//        ^---- this length marker is included in tag ids but not in other ints
		consumed := 9
		for b := data[0]; b != 0; b >>= 1 {
			consumed -= 1
		}
		if len(data) >= consumed {
			return fixedUint(data[:consumed]), consumed
		}
	}
	return 0, 0
}

func ebmlUint(data []byte) (uint64, int) {
	id, consumed := ebmlTagID(data)
	if ebmlIndeterminateCoding[consumed] == id {
		return ebmlIndeterminate, consumed
	}
	return id & ^(1 << uint(7*consumed)), consumed
}

type ebmlTag struct {
	Consumed int
	ID       uint
	Length   uint64
}

func ebmlParseTagIncomplete(data []byte) ebmlTag {
	if id, off := ebmlTagID(data); off != 0 {
		if length, off2 := ebmlUint(data[off:]); off2 != 0 {
			return ebmlTag{off + off2, uint(id), length}
		}
	}
	return ebmlTag{0, 0, 0}
}

func ebmlParseTag(data []byte) ebmlTag {
	if tag := ebmlParseTagIncomplete(data); tag.Length+uint64(tag.Consumed) <= uint64(len(data)) {
		return tag
	}
	return ebmlTag{0, 0, 0}
}

func (t ebmlTag) Contents(data []byte) []byte {
	return data[t.Consumed : uint64(t.Consumed)+t.Length]
}

func (t ebmlTag) Skip(data []byte) []byte {
	return data[uint64(t.Consumed)+t.Length:]
}

type frame struct {
	buf   []byte // Either a Block(Group) or a Cluster.
	track uint64 // 64 for a Cluster (track masks are 32-bit, so streams with a real 64-th track are rejected)
	key   bool
}

type framebuffer struct {
	data        []frame
	start       int
	headCluster []byte // Last Cluster popped off the ring. Blocks at the head still belong to it.
}

func (fb *framebuffer) PushCluster(buf []byte) {
	fb.PushFrame(frame{buf, 64, false})
}

func (fb *framebuffer) PushFrame(packed frame) {
	if len(fb.data) == cap(fb.data) {
		if fb.data[fb.start].track == 64 {
			fb.headCluster = fb.data[fb.start].buf
		}
		fb.data[fb.start] = packed
		fb.start = (fb.start + 1) % len(fb.data)
	} else {
		fb.data = append(fb.data, packed)
	}
}

func (fb *framebuffer) Read(cb func(cluster []byte, forceCluster bool, packed frame)) {
	cluster, forceCluster := fb.headCluster, true
	for i, s, n := 0, fb.start, len(fb.data); i < n; i++ {
		f := fb.data[(i+s)%n]
		if f.track == 64 {
			cluster = f.buf
			forceCluster = true
		} else if cluster != nil {
			cb(cluster, forceCluster, f)
			forceCluster = false
		}
	}
}

type viewer struct {
	// This function may return `false` to signal that it cannot write any more data.
	// The stream will resynchronize at next keyframe.
	write func(data []byte) bool
	// Viewers may hop between streams, but should only receive headers once.
	// This includes track info, as codecs must stay the same between segments.
	skipHeaders bool
	// We group blocks into indeterminate-length clusters. So long as
	// the cluster's timecode has not changed, there's no need to start a new one.
	skipCluster bool
	// Bit vector of tracks for which the viewer has both reference frames
	// (the previous frame and the last keyframe.)
	seenKeyframes uint32
}

func (cb *viewer) WriteFrame(cluster []byte, forceCluster bool, packed frame) {
	trackMask := uint32(1) << packed.track
	if packed.key {
		cb.seenKeyframes |= trackMask
	}
	if cb.seenKeyframes&trackMask != 0 {
		if !cb.skipCluster || forceCluster {
			cb.skipCluster = cb.write(cluster)
		}
		if !cb.skipCluster || !cb.write(packed.buf) {
			cb.seenKeyframes &= ^trackMask
		}
	}
}

type BroadcastSet struct {
	mutex   sync.Mutex
	streams map[string]*Broadcast
	// How long to keep a stream alive after a call to `Close`.
	Timeout time.Duration
	// Called right after a stream is destroyed. (`Timeout` seconds after a `Close`.)
	OnStreamClose     func(id string)
	OnStreamTrackInfo func(id string, info *StreamTrackInfo)
}

type Broadcast struct {
	StreamTrackInfo
	closing time.Duration
	Closed  bool
	dirty   bool // (Has unseen data in `StreamTrackInfo`.)
	buffer  []byte
	header  []byte // The EBML (DocType) tag.
	tracks  []byte // The beginning of the Segment (Tracks + Info).
	frames  framebuffer
	// outbound blocks must have monotonically increasing timecodes even if the inbound
	// stream restarts from the beginning, so we'll shift clusters' timecodes.
	blockTimecode       uint64
	recvClusterTimecode uint64
	sentClusterTimecode uint64
	timecodeShift       uint64
	// these values are for the whole stream, so they include audio and muxing overhead.
	// the latter is negligible, however, and the former is normally about 64k,
	// so also negligible. or at least predictable.
	rateUnit float64
	RateMean float64
	RateVar  float64

	vlock   sync.Mutex
	viewers map[chan<- []byte]*viewer
}

func (ctx *BroadcastSet) Readable(id string) (*Broadcast, bool) {
	if ctx.streams == nil {
		return nil, false
	}
	ctx.mutex.Lock()
	cast, ok := ctx.streams[id]
	ctx.mutex.Unlock()
	return cast, ok
}

func (ctx *BroadcastSet) Writable(id string) (*Broadcast, bool) {
	ctx.mutex.Lock()
	defer ctx.mutex.Unlock()
	if ctx.streams == nil {
		ctx.streams = make(map[string]*Broadcast)
	}
	if cast, ok := ctx.streams[id]; ok {
		if cast.closing == -1 {
			return nil, false
		}
		cast.closing = -1
		return cast, true
	}
	cast := Broadcast{
		closing:             -1,
		frames:              framebuffer{make([]frame, 0, 120), 0, nil},
		viewers:             make(map[chan<- []byte]*viewer),
		sentClusterTimecode: 0xFFFFFFFFFFFFFFFF,
	}
	ctx.streams[id] = &cast
	go func() {
		ticker := time.NewTicker(time.Second)
		for range ticker.C {
			if cast.dirty {
				cast.dirty = false
				ctx.OnStreamTrackInfo(id, &cast.StreamTrackInfo)
			}
			if cast.closing >= 0 {
				if cast.closing += time.Second; cast.closing > ctx.Timeout {
					break
				}
			}
			// exponentially weighted moving moments at a = 0.5
			//     avg[n] = a * x + (1 - a) * avg[n - 1]
			//     var[n] = a * (x - avg[n]) ** 2 / (1 - a) + (1 - a) * var[n - 1]
			cast.RateMean += cast.rateUnit / 2
			cast.RateVar += cast.rateUnit*cast.rateUnit - cast.RateVar/2
			cast.rateUnit = -cast.RateMean
		}
		ticker.Stop()

		ctx.mutex.Lock()
		delete(ctx.streams, id)
		ctx.mutex.Unlock()
		cast.Closed = true
		cast.vlock.Lock()
		for _, cb := range cast.viewers {
			cb.write([]byte{})
		}
		cast.vlock.Unlock()
		if ctx.OnStreamClose != nil {
			ctx.OnStreamClose(id)
		}
	}()
	return &cast, true
}

func (cast *Broadcast) Close() error {
	cast.closing = 0
	return nil
}

func (cast *Broadcast) Connect(ch chan<- []byte, skipHeaders bool) {
	blocked := false
	write := func(data []byte) bool {
		// `Broadcast.Write` emits data in block-sized chunks.
		// Thus the buffer size is measured in frames, not bytes.
		blocked = len(ch) == cap(ch) || (blocked && len(ch)*2 >= cap(ch))
		if !blocked {
			ch <- data
		}
		return !blocked
	}

	cast.vlock.Lock()
	cast.viewers[ch] = &viewer{write: write}
	cast.vlock.Unlock()
}

func (cast *Broadcast) Disconnect(ch chan<- []byte) {
	cast.vlock.Lock()
	delete(cast.viewers, ch)
	cast.vlock.Unlock()
}

func (cast *Broadcast) Reset() {
	cast.buffer = nil
{'fraction_non_alphanumeric': 0.111690800681431, 'fraction_numerical': 0.035615417376490634, 'mean_word_length': 3.0389163620726727, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 14, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 38, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3335899', 'n_tokens_mistral': 7576, 'n_tokens_neox': 6785, 'n_words': 1870}