<commit_before>// Copyright 2016 Hajime Hoshi
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context) (*Player, error) {
	const bytesPerSample = 4 // TODO: This should be defined in audio package
	s, err := vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
	if err != nil {
		return nil, err
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Current Time: %s
Current Volume: %d/128`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128))
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer *Player
)

func update(screen *ebiten.Image) error {
	if err := musicPlayer.update(); err != nil {
		return err
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	musicPlayer.draw(screen)
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>examples/audio: Switchable to MP3 from Ogg<commit_after>// Copyright 2016 Hajime Hoshi
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build example jsgo

// This is an example to implement an audio player.
// See examples/wav for a simpler example to play a sound file.

package main

import (
	"fmt"
	"image/color"
	"io/ioutil"
	"log"
	"time"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/audio/mp3"
	"github.com/hajimehoshi/ebiten/audio/vorbis"
	"github.com/hajimehoshi/ebiten/audio/wav"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	raudio "github.com/hajimehoshi/ebiten/examples/resources/audio"
	"github.com/hajimehoshi/ebiten/inpututil"
)

const (
	screenWidth  = 320
	screenHeight = 240

	sampleRate = 22050
)

var (
	playerBarColor     = color.RGBA{0x80, 0x80, 0x80, 0xff}
	playerCurrentColor = color.RGBA{0xff, 0xff, 0xff, 0xff}
)

type musicType int

const (
	typeOgg musicType = iota
	typeMP3
)

func (t musicType) String() string {
	switch t {
	case typeOgg:
		return "Ogg"
	case typeMP3:
		return "MP3"
	default:
		panic("not reached")
	}
}

// Player represents the current audio state.
type Player struct {
	audioContext *audio.Context
	audioPlayer  *audio.Player
	current      time.Duration
	total        time.Duration
	seBytes      []byte
	seCh         chan []byte
	volume128    int
	musicType    musicType
}

func playerBarRect() (x, y, w, h int) {
	w, h = 300, 4
	x = (screenWidth - w) / 2
	y = screenHeight - h - 16
	return
}

func NewPlayer(audioContext *audio.Context, musicType musicType) (*Player, error) {
	type audioStream interface {
		audio.ReadSeekCloser
		Length() int64
	}

	const bytesPerSample = 4 // TODO: This should be defined in audio package

	var s audioStream

	switch musicType {
	case typeOgg:
		var err error
		s, err = vorbis.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Ragtime_ogg))
		if err != nil {
			return nil, err
		}
	case typeMP3:
		var err error
		s, err = mp3.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Classic_mp3))
		if err != nil {
			return nil, err
		}
	default:
		panic("not reached")
	}
	p, err := audio.NewPlayer(audioContext, s)
	if err != nil {
		return nil, err
	}
	player := &Player{
		audioContext: audioContext,
		audioPlayer:  p,
		total:        time.Second * time.Duration(s.Length()) / bytesPerSample / sampleRate,
		volume128:    128,
		seCh:         make(chan []byte),
		musicType:    musicType,
	}
	if player.total == 0 {
		player.total = 1
	}
	player.audioPlayer.Play()
	go func() {
		s, err := wav.Decode(audioContext, audio.BytesReadSeekCloser(raudio.Jab_wav))
		if err != nil {
			log.Fatal(err)
			return
		}
		b, err := ioutil.ReadAll(s)
		if err != nil {
			log.Fatal(err)
			return
		}
		player.seCh <- b
	}()
	return player, nil
}

func (p *Player) Close() error {
	return p.audioPlayer.Close()
}

func (p *Player) update() error {
	select {
	case p.seBytes = <-p.seCh:
		close(p.seCh)
		p.seCh = nil
	default:
	}

	if p.audioPlayer.IsPlaying() {
		p.current = p.audioPlayer.Current()
	}
	p.seekBarIfNeeded()
	p.switchPlayStateIfNeeded()
	p.playSEIfNeeded()
	p.updateVolumeIfNeeded()

	if inpututil.IsKeyJustPressed(ebiten.KeyB) {
		b := ebiten.IsRunnableInBackground()
		ebiten.SetRunnableInBackground(!b)
	}
	return nil
}

func (p *Player) playSEIfNeeded() {
	if p.seBytes == nil {
		// Bytes for the SE is not loaded yet.
		return
	}

	if !inpututil.IsKeyJustPressed(ebiten.KeyP) {
		return
	}
	sePlayer, _ := audio.NewPlayerFromBytes(p.audioContext, p.seBytes)
	sePlayer.Play()
}

func (p *Player) updateVolumeIfNeeded() {
	if ebiten.IsKeyPressed(ebiten.KeyZ) {
		p.volume128--
	}
	if ebiten.IsKeyPressed(ebiten.KeyX) {
		p.volume128++
	}
	if p.volume128 < 0 {
		p.volume128 = 0
	}
	if 128 < p.volume128 {
		p.volume128 = 128
	}
	p.audioPlayer.SetVolume(float64(p.volume128) / 128)
}

func (p *Player) switchPlayStateIfNeeded() {
	if !inpututil.IsKeyJustPressed(ebiten.KeyS) {
		return
	}
	if p.audioPlayer.IsPlaying() {
		p.audioPlayer.Pause()
		return
	}
	p.audioPlayer.Play()
}

func (p *Player) seekBarIfNeeded() {
	if !inpututil.IsMouseButtonJustPressed(ebiten.MouseButtonLeft) {
		return
	}

	// Calculate the next seeking position from the current cursor position.
	x, y := ebiten.CursorPosition()
	bx, by, bw, bh := playerBarRect()
	const padding = 4
	if y < by-padding || by+bh+padding <= y {
		return
	}
	if x < bx || bx+bw <= x {
		return
	}
	pos := time.Duration(x-bx) * p.total / time.Duration(bw)
	p.current = pos
	p.audioPlayer.Seek(pos)
}

func (p *Player) draw(screen *ebiten.Image) {
	// Draw the bar.
	x, y, w, h := playerBarRect()
	ebitenutil.DrawRect(screen, float64(x), float64(y), float64(w), float64(h), playerBarColor)

	// Draw the cursor on the bar.
	c := p.current
	cw, ch := 4, 10
	cx := int(time.Duration(w)*c/p.total) + x - cw/2
	cy := y - (ch-h)/2
	ebitenutil.DrawRect(screen, float64(cx), float64(cy), float64(cw), float64(ch), playerCurrentColor)

	// Compose the curren time text.
	m := (c / time.Minute) % 100
	s := (c / time.Second) % 60
	currentTimeStr := fmt.Sprintf("%02d:%02d", m, s)

	// Draw the debug message.
	msg := fmt.Sprintf(`TPS: %0.2f
Press S to toggle Play/Pause
Press P to play SE
Press Z or X to change volume of the music
Press B to switch the run-in-background state
Press A to switch Ogg and MP3
Current Time: %s
Current Volume: %d/128
Type: %s`, ebiten.CurrentTPS(), currentTimeStr, int(p.audioPlayer.Volume()*128), musicPlayer.musicType)
	ebitenutil.DebugPrint(screen, msg)
}

var (
	musicPlayer   *Player
	musicPlayerCh = make(chan *Player)
	errCh         = make(chan error)
)

func update(screen *ebiten.Image) error {
	select {
	case p := <-musicPlayerCh:
		musicPlayer = p
	case err := <-errCh:
		return err
	default:
	}

	if musicPlayer != nil && inpututil.IsKeyJustPressed(ebiten.KeyA) {
		var t musicType
		switch musicPlayer.musicType {
		case typeOgg:
			t = typeMP3
		case typeMP3:
			t = typeOgg
		default:
			panic("not reached")
		}

		musicPlayer.Close()
		musicPlayer = nil

		go func() {
			p, err := NewPlayer(audio.CurrentContext(), t)
			if err != nil {
				errCh <- err
				return
			}
			musicPlayerCh <- p
		}()
	}

	if musicPlayer != nil {
		if err := musicPlayer.update(); err != nil {
			return err
		}
	}

	if ebiten.IsDrawingSkipped() {
		return nil
	}

	if musicPlayer != nil {
		musicPlayer.draw(screen)
	}
	return nil
}

func main() {
	audioContext, err := audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}

	musicPlayer, err = NewPlayer(audioContext, typeOgg)
	if err != nil {
		log.Fatal(err)
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package floatgeom

import (
	"math"
	"math/rand"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestPointGreaterOf(t *testing.T) {
	a := Point2{0, 1}
	b := Point2{1, 0}
	assert.Equal(t, a.GreaterOf(b), Point2{1, 1})

	c := Point3{0, 1, 2}
	d := Point3{1, 2, 0}
	assert.Equal(t, c.GreaterOf(d), Point3{1, 2, 2})
}

func TestPointLesserOf(t *testing.T) {
	a := Point2{0, 1}
	b := Point2{1, 0}
	assert.Equal(t, a.LesserOf(b), Point2{0, 0})

	c := Point3{0, 1, 2}
	d := Point3{1, 2, 0}
	assert.Equal(t, c.LesserOf(d), Point3{0, 1, 0})
}

func TestPointAccess(t *testing.T) {
	a := Point3{0, 1, 2}
	assert.Equal(t, 0.0, a.Dim(0))
	assert.Equal(t, 1.0, a.Dim(1))
	assert.Equal(t, 2.0, a.Dim(2))
	assert.Equal(t, 0.0, a.X())
	assert.Equal(t, 1.0, a.Y())
	assert.Equal(t, 2.0, a.Z())

	b := Point2{0, 1}
	assert.Equal(t, 0.0, b.Dim(0))
	assert.Equal(t, 1.0, b.Dim(1))
	assert.Equal(t, 0.0, b.X())
	assert.Equal(t, 1.0, b.Y())
}

// Pattern here: there's a set of input pairs here
// each test takes these and has expected outputs for each pair index.
var (
	// Todo: add more test cases
	pt3cases = []struct{ x1, y1, z1, x2, y2, z2 float64 }{
		{0, 0, 0, 1, 1, 1},
	}
	pt2cases = []struct{ x1, y1, x2, y2 float64 }{
		{0, 0, 1, 1},
	}
)

func TestPointDistance3(t *testing.T) {

	expected := []float64{math.Sqrt(3)}

	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Distance(b), e)
	}
}

func TestPointDistance2(t *testing.T) {
	expected := []float64{math.Sqrt(2)}

	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Distance(b), e)
	}
}

func TestPointAdd3(t *testing.T) {
	expected := []Point3{
		{1, 1, 1},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Add(b), e)
	}
}

func TestPointAdd2(t *testing.T) {
	expected := []Point2{
		{1, 1},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Add(b), e)
	}
}

func TestPointSub3(t *testing.T) {
	expected := []Point3{
		{-1, -1, -1},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Sub(b), e)
	}
}

func TestPointSub2(t *testing.T) {
	expected := []Point2{
		{-1, -1},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Sub(b), e)
	}
}

func TestPointMul3(t *testing.T) {
	expected := []Point3{
		{0, 0, 0},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Mul(b), e)
	}
}

func TestPointMul2(t *testing.T) {
	expected := []Point2{
		{0, 0},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Mul(b), e)
	}
}

func TestPointDiv3(t *testing.T) {
	expected := []Point3{
		{0, 0, 0},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Div(b), e)
	}
}

func TestPointDiv2(t *testing.T) {
	expected := []Point2{
		{0, 0},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Div(b), e)
	}
}

var (
	randTests = 100
)

func TestPointToRec(t *testing.T) {
	for i := 0; i < randTests; i++ {
		span := rand.Float64()
		x := rand.Float64()
		y := rand.Float64()
		z := rand.Float64()

		expected3 := Rect3{Min: Point3{x, y, z}, Max: Point3{x + span, y + span, z + span}}
		expected2 := Rect2{Min: Point2{x, y}, Max: Point2{x + span, y + span}}

		assert.Equal(t,
			expected3,
			Point3{x, y, z}.ToRect(span),
		)

		assert.Equal(t,
			expected2,
			Point2{x, y}.ToRect(span),
		)
	}
}
<commit_msg>Additional floatgeom tests<commit_after>package floatgeom

import (
	"math"
	"math/rand"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func Seed() {
	rand.Seed(time.Now().UnixNano())
}

func TestPointNormalize(t *testing.T) {
	p1 := Point2{100, 200}.Normalize()
	p2 := Point3{100, 200, 300}.Normalize()

	assert.InEpsilon(t, p1.X(), 1/math.Sqrt(5), .0001)
	assert.InEpsilon(t, p1.Y(), 2/math.Sqrt(5), .0001)
	assert.InEpsilon(t, p2.X(), 1/math.Sqrt(14), .0001)
	assert.InEpsilon(t, p2.Y(), 2/math.Sqrt(14), .0001)
	assert.InEpsilon(t, p2.Z(), 3/math.Sqrt(14), .0001)
}

func TestPointProject(t *testing.T) {
	Seed()
	for i := 0; i < randTests; i++ {
		x, y, z := rand.Float64(), rand.Float64(), rand.Float64()
		p := Point3{x, y, z}
		assert.Equal(t, Point2{x, y}, p.ProjectZ())
		assert.Equal(t, Point2{x, z}, p.ProjectY())
		assert.Equal(t, Point2{y, z}, p.ProjectX())
	}
}

func TestPointConstMods(t *testing.T) {
	p1 := Point2{1, 1}
	p2 := Point3{1, 1, 1}
	assert.Equal(t, Point2{5, 5}, p1.MulConst(5))
	assert.Equal(t, Point3{100, 100, 100}, p2.MulConst(100))
	p3 := Point2{2, 2}
	p4 := Point3{2, 2, 2}
	assert.Equal(t, Point2{.5, .5}, p3.DivConst(4))
	assert.Equal(t, Point3{.25, .25, .25}, p4.DivConst(8))
}

func TestAnglePoints(t *testing.T) {
	p1 := Point2{1, 0}
	p2 := Point2{0, 1}

	assert.Equal(t, p1, AnglePoint(0))
	assert.Equal(t, p1, RadianPoint(0))
	p3 := AnglePoint(90)
	p4 := RadianPoint(math.Pi / 2)
	assert.InEpsilon(t, p2.Y(), p3.Y(), .0001)
	assert.InEpsilon(t, p2.Y(), p4.Y(), .0001)

	deg := p1.ToAngle()
	rds := p1.ToRadians()

	assert.Equal(t, 0.0, deg)
	assert.Equal(t, 0.0, rds)

	deg = p2.ToAngle()
	rds = p2.ToRadians()

	assert.InEpsilon(t, 90.0, deg, .0001)
	assert.InEpsilon(t, math.Pi/2, rds, .0001)

	p5 := Point2{-1, 0}

	assert.InEpsilon(t, 45.0, p2.AngleTo(p5), .0001)
	assert.InEpsilon(t, math.Pi/4, p2.RadiansTo(p5), .0001)
}

func TestPointGreaterOf(t *testing.T) {
	a := Point2{0, 1}
	b := Point2{1, 0}
	assert.Equal(t, a.GreaterOf(b), Point2{1, 1})

	c := Point3{0, 1, 2}
	d := Point3{1, 2, 0}
	assert.Equal(t, c.GreaterOf(d), Point3{1, 2, 2})
}

func TestPointLesserOf(t *testing.T) {
	a := Point2{0, 1}
	b := Point2{1, 0}
	assert.Equal(t, a.LesserOf(b), Point2{0, 0})

	c := Point3{0, 1, 2}
	d := Point3{1, 2, 0}
	assert.Equal(t, c.LesserOf(d), Point3{0, 1, 0})
}

func TestPointAccess(t *testing.T) {
	a := Point3{0, 1, 2}
	assert.Equal(t, 0.0, a.Dim(0))
	assert.Equal(t, 1.0, a.Dim(1))
	assert.Equal(t, 2.0, a.Dim(2))
	assert.Equal(t, 0.0, a.X())
	assert.Equal(t, 1.0, a.Y())
	assert.Equal(t, 2.0, a.Z())

	b := Point2{0, 1}
	assert.Equal(t, 0.0, b.Dim(0))
	assert.Equal(t, 1.0, b.Dim(1))
	assert.Equal(t, 0.0, b.X())
	assert.Equal(t, 1.0, b.Y())
}

// Pattern here: there's a set of input pairs here
// each test takes these and has expected outputs for each pair index.
var (
	// Todo: add more test cases
	pt3cases = []struct{ x1, y1, z1, x2, y2, z2 float64 }{
		{0, 0, 0, 1, 1, 1},
	}
	pt2cases = []struct{ x1, y1, x2, y2 float64 }{
		{0, 0, 1, 1},
	}
)

func TestPointDistance3(t *testing.T) {

	expected := []float64{math.Sqrt(3)}

	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Distance(b), e)
	}
}

func TestPointDistance2(t *testing.T) {
	expected := []float64{math.Sqrt(2)}

	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Distance(b), e)
	}
}

func TestPointAdd3(t *testing.T) {
	expected := []Point3{
		{1, 1, 1},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Add(b), e)
	}
}

func TestPointAdd2(t *testing.T) {
	expected := []Point2{
		{1, 1},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Add(b), e)
	}
}

func TestPointSub3(t *testing.T) {
	expected := []Point3{
		{-1, -1, -1},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Sub(b), e)
	}
}

func TestPointSub2(t *testing.T) {
	expected := []Point2{
		{-1, -1},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Sub(b), e)
	}
}

func TestPointMul3(t *testing.T) {
	expected := []Point3{
		{0, 0, 0},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Mul(b), e)
	}
}

func TestPointMul2(t *testing.T) {
	expected := []Point2{
		{0, 0},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Mul(b), e)
	}
}

func TestPointDiv3(t *testing.T) {
	expected := []Point3{
		{0, 0, 0},
	}
	for i, e := range expected {
		c := pt3cases[i]
		a := Point3{c.x1, c.y1, c.z1}
		b := Point3{c.x2, c.y2, c.z2}
		assert.Equal(t, a.Div(b), e)
	}
}

func TestPointDiv2(t *testing.T) {
	expected := []Point2{
		{0, 0},
	}
	for i, e := range expected {
		c := pt2cases[i]
		a := Point2{c.x1, c.y1}
		b := Point2{c.x2, c.y2}
		assert.Equal(t, a.Div(b), e)
	}
}

var (
	randTests = 100
)

func TestPointToRec(t *testing.T) {
	for i := 0; i < randTests; i++ {
		span := rand.Float64()
		x := rand.Float64()
		y := rand.Float64()
		z := rand.Float64()

		expected3 := Rect3{Min: Point3{x, y, z}, Max: Point3{x + span, y + span, z + span}}
		expected2 := Rect2{Min: Point2{x, y}, Max: Point2{x + span, y + span}}

		assert.Equal(t,
			expected3,
			Point3{x, y, z}.ToRect(span),
		)

		assert.Equal(t,
			expected2,
			Point2{x, y}.ToRect(span),
		)
	}
}
<|endoftext|>


<commit_before>package s3

import (
	"io"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// Emitter stores data in S3 bucket.
//
// The use of  this struct requires the configuration of an S3 bucket/endpoint. When the buffer is full, this
// struct's Emit method adds the contents of the buffer to S3 as one file. The filename is generated
// from the first and last sequence numbers of the records contained in that file separated by a
// dash. This struct requires the configuration of an S3 bucket and endpoint.
type Emitter struct {
	Bucket string
	Region string
}

// Emit is invoked when the buffer is full. This method emits the set of filtered records.
func (e Emitter) Emit(s3Key string, b io.ReadSeeker) error {
	svc := s3.New(
		session.New(aws.NewConfig().WithMaxRetries(10)),
		&aws.Config{
			Region: aws.String(e.Region),
		},
	)

	params := &s3.PutObjectInput{
		Body:        b,
		Bucket:      aws.String(e.Bucket),
		ContentType: aws.String("text/plain"),
		Key:         aws.String(s3Key),
	}

	_, err := svc.PutObject(params)

	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Update S3 emitter initialization<commit_after>package s3

import (
	"io"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// Emitter stores data in S3 bucket.
//
// The use of  this struct requires the configuration of an S3 bucket/endpoint. When the buffer is full, this
// struct's Emit method adds the contents of the buffer to S3 as one file. The filename is generated
// from the first and last sequence numbers of the records contained in that file separated by a
// dash. This struct requires the configuration of an S3 bucket and endpoint.
type Emitter struct {
	Bucket string
	Region string
}

// Emit is invoked when the buffer is full. This method emits the set of filtered records.
func (e Emitter) Emit(s3Key string, b io.ReadSeeker) error {
	svc := s3.New(
		session.New(aws.NewConfig().WithMaxRetries(10)),
		aws.NewConfig().WithRegion(e.Region),
	)

	params := &s3.PutObjectInput{
		Body:        b,
		Bucket:      aws.String(e.Bucket),
		ContentType: aws.String("text/plain"),
		Key:         aws.String(s3Key),
	}

	_, err := svc.PutObject(params)

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package apiserver

import (
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"golang.org/x/net/websocket"
)

type chanLevel struct {
	ch    chan struct{}
	level logrus.Level
}

type wsLogHook struct {
	sync.Mutex
	data   map[*websocket.Conn]chanLevel
	logger *logrus.Logger
}

func newWSLogHook(logger *logrus.Logger) *wsLogHook {
	return &wsLogHook{
		logger: logger,
		data:   map[*websocket.Conn]chanLevel{},
	}
}

func (hook *wsLogHook) addConn(ws *websocket.Conn, level logrus.Level) <-chan struct{} {
	hook.Lock()
	defer hook.Unlock()

	ch := make(chan struct{})
	hook.data[ws] = chanLevel{
		ch:    ch,
		level: level,
	}

	return ch
}

func (hook *wsLogHook) Levels() []logrus.Level {
	ret := make([]logrus.Level, logrus.DebugLevel+1)

	// add all levels
	for i := logrus.PanicLevel; i <= logrus.DebugLevel; i++ {
		ret[i] = i
	}

	return ret
}

func (hook *wsLogHook) deleteConnNoLock(ws *websocket.Conn) {
	cl, ok := hook.data[ws]
	if !ok {
		return
	}

	delete(hook.data, ws)

	// let the handler return
	select {
	case cl.ch <- struct{}{}:
	default:
	}
}

type wsMessage struct {
	Data    logrus.Fields
	Time    time.Time
	Level   logrus.Level
	Message string
}

func newWSMessage(entry *logrus.Entry) *wsMessage {
	return &wsMessage{
		Data:    entry.Data,
		Time:    entry.Time,
		Level:   entry.Level,
		Message: entry.Message,
	}
}

func (hook *wsLogHook) Fire(entry *logrus.Entry) error {
	go func() {
		hook.Lock()
		defer hook.Unlock()

		msg := newWSMessage(entry)

		for ws, cl := range hook.data {
			if ws == nil {
				hook.deleteConnNoLock(ws)
				continue
			}

			if cl.level < entry.Level {
				continue
			}

			if err := websocket.JSON.Send(ws, msg); err != nil {
				hook.deleteConnNoLock(ws)
			}
		}
	}()
	return nil
}
<commit_msg>ensure api logs are have strings for both keys and values<commit_after>package apiserver

import (
	"fmt"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"golang.org/x/net/websocket"
)

type chanLevel struct {
	ch    chan struct{}
	level logrus.Level
}

type wsLogHook struct {
	sync.Mutex
	data   map[*websocket.Conn]chanLevel
	logger *logrus.Logger
}

func newWSLogHook(logger *logrus.Logger) *wsLogHook {
	return &wsLogHook{
		logger: logger,
		data:   map[*websocket.Conn]chanLevel{},
	}
}

func (hook *wsLogHook) addConn(ws *websocket.Conn, level logrus.Level) <-chan struct{} {
	hook.Lock()
	defer hook.Unlock()

	ch := make(chan struct{})
	hook.data[ws] = chanLevel{
		ch:    ch,
		level: level,
	}

	return ch
}

func (hook *wsLogHook) Levels() []logrus.Level {
	ret := make([]logrus.Level, logrus.DebugLevel+1)

	// add all levels
	for i := logrus.PanicLevel; i <= logrus.DebugLevel; i++ {
		ret[i] = i
	}

	return ret
}

func (hook *wsLogHook) deleteConnNoLock(ws *websocket.Conn) {
	cl, ok := hook.data[ws]
	if !ok {
		return
	}

	delete(hook.data, ws)

	// let the handler return
	select {
	case cl.ch <- struct{}{}:
	default:
	}
}

type wsMessage struct {
	Data    map[string]string
	Time    time.Time
	Level   logrus.Level
	Message string
}

func newWSMessage(entry *logrus.Entry) *wsMessage {
	data := make(map[string]string, len(entry.Data))
	for key, value := range entry.Data {
		data[key] = fmt.Sprintf("%v", value)
	}
	return &wsMessage{
		Data:    data,
		Time:    entry.Time,
		Level:   entry.Level,
		Message: entry.Message,
	}
}

func (hook *wsLogHook) Fire(entry *logrus.Entry) error {
	go func() {
		hook.Lock()
		defer hook.Unlock()

		msg := newWSMessage(entry)

		for ws, cl := range hook.data {
			if ws == nil {
				hook.deleteConnNoLock(ws)
				continue
			}

			if cl.level < entry.Level {
				continue
			}

			if err := websocket.JSON.Send(ws, msg); err != nil {
				hook.deleteConnNoLock(ws)
			}
		}
	}()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015-2018 trivago N.V.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tio

import (
	"bytes"
	"encoding/binary"
	"io"
	"regexp"

	"github.com/trivago/tgo/tstrings"
)

// BufferedReaderFlags is an enum to configure a buffered reader
type BufferedReaderFlags byte

const (
	// BufferedReaderFlagDelimiter enables reading for a delimiter. This flag is
	// ignored if an MLE flag is set.
	BufferedReaderFlagDelimiter = BufferedReaderFlags(0)

	// BufferedReaderFlagMLE enables reading if length encoded messages.
	// Runlength is read as ASCII (to uint64) until the first byte (ASCII char)
	// of the delimiter string.
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE = BufferedReaderFlags(1)

	// BufferedReaderFlagMLE8 enables reading if length encoded messages.
	// Runlength is read as binary (to uint8).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE8 = BufferedReaderFlags(2)

	// BufferedReaderFlagMLE16 enables reading if length encoded messages.
	// Runlength is read as binary (to uint16).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE16 = BufferedReaderFlags(3)

	// BufferedReaderFlagMLE32 enables reading if length encoded messages.
	// Runlength is read as binary (to uint32).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE32 = BufferedReaderFlags(4)

	// BufferedReaderFlagMLE64 enables reading if length encoded messages.
	// Runlength is read as binary (to uint64).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE64 = BufferedReaderFlags(5)

	// BufferedReaderFlagMLEFixed enables reading messages with a fixed length.
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLEFixed = BufferedReaderFlags(6)

	// BufferedReaderFlagMaskMLE is a bitmask to mask out everything but MLE flags
	BufferedReaderFlagMaskMLE = BufferedReaderFlags(7)

	// BufferedReaderFlagBigEndian sets binary reading to big endian encoding.
	BufferedReaderFlagBigEndian = BufferedReaderFlags(8)

	// BufferedReaderFlagEverything will keep MLE and/or delimiters when
	// building a message.
	BufferedReaderFlagEverything = BufferedReaderFlags(16)

	// BufferedReaderFlagRegex enables reading messages delimited by regexp.
	// To enable multiline reading, prefix the regex with "(?m)"
	BufferedReaderFlagRegex = BufferedReaderFlags(32)

	// BufferedReaderFlagRegexStart switches regex reading to find the start
	// of the next message instead of looking for the end of the current.
	BufferedReaderFlagRegexStart = BufferedReaderFlags(64 + 32)
)

type bufferError string

func (b bufferError) Error() string {
	return string(b)
}

// BufferReadCallback defines the function signature for callbacks passed to
// ReadAll.
type BufferReadCallback func(msg []byte)

// BufferDataInvalid is returned when a parsing encounters an error
var BufferDataInvalid = bufferError("Invalid data")

// BufferedReader is a helper struct to read from any io.Reader into a byte
// slice. The data can arrive "in pieces" and will be assembled.
// A data "piece" is considered complete if a delimiter or a certain runlength
// has been reached. The latter has to be enabled by flag and will disable the
// default behavior, which is looking for a delimiter string.
type BufferedReader struct {
	data            []byte
	delimiter       []byte
	parse           func() ([]byte, int)
	paramMLE        int
	growSize        int
	end             int
	encoding        binary.ByteOrder
	flags           BufferedReaderFlags
	incomplete      bool
	delimiterRegexp *regexp.Regexp
}

// NewBufferedReader creates a new buffered reader that reads messages from a
// continuous stream of bytes.
// Messages can be separated from the stream by using common methods such as
// fixed size, encoded message length or delimiter string.
// The internal buffer is grown statically (by its original size) if necessary.
// bufferSize defines the initial size / grow size of the buffer.
// flags configures the parsing method.
// offsetOrLength sets either the runlength offset or fixed message size.
// delimiter defines the delimiter used for textual message parsing, i.e.
// BufferedReaderFlagDelimiter or BufferedReaderFlagRegex.
func NewBufferedReader(bufferSize int, flags BufferedReaderFlags, offsetOrLength int, delimiter string) *BufferedReader {
	buffer := BufferedReader{
		data:            make([]byte, bufferSize),
		delimiter:       []byte(delimiter),
		paramMLE:        offsetOrLength,
		encoding:        binary.LittleEndian,
		end:             0,
		flags:           flags,
		growSize:        bufferSize,
		incomplete:      true,
		delimiterRegexp: nil,
	}

	if flags&BufferedReaderFlagBigEndian != 0 {
		buffer.encoding = binary.BigEndian
	}

	if flags&BufferedReaderFlagRegex != 0 {
		buffer.parse = buffer.parseDelimiterRegex
		buffer.delimiterRegexp = regexp.MustCompile(delimiter)
		return &buffer
	}

	if flags&BufferedReaderFlagMaskMLE == 0 {
		buffer.parse = buffer.parseDelimiter
	} else {
		switch flags & BufferedReaderFlagMaskMLE {
		default:
			buffer.parse = buffer.parseMLEText
		case BufferedReaderFlagMLE8:
			buffer.parse = buffer.parseMLE8
		case BufferedReaderFlagMLE16:
			buffer.parse = buffer.parseMLE16
		case BufferedReaderFlagMLE32:
			buffer.parse = buffer.parseMLE32
		case BufferedReaderFlagMLE64:
			buffer.parse = buffer.parseMLE64
		case BufferedReaderFlagMLEFixed:
			buffer.parse = buffer.parseMLEFixed
		}
	}

	return &buffer
}

// Reset clears the buffer by resetting its internal state
func (buffer *BufferedReader) Reset(sequence uint64) {
	buffer.end = 0
	buffer.incomplete = true
}

// ResetGetIncomplete works like Reset but returns any incomplete contents
// left in the buffer. Note that the data in the buffer returned is overwritten
// with the next read call.
func (buffer *BufferedReader) ResetGetIncomplete() []byte {
	remains := buffer.data[:buffer.end]
	buffer.Reset(0)
	return remains
}

// HasIncompleteData returns true if there is unparsed data left in the buffer
func (buffer *BufferedReader) HasIncompleteData() bool {
	return buffer.end > 0
}

// general message extraction part of all parser methods
func (buffer *BufferedReader) extractMessage(messageLen int, msgStartIdx int) ([]byte, int) {
	nextMsgIdx := msgStartIdx + messageLen
	if nextMsgIdx > buffer.end {
		return nil, 0 // ### return, incomplete ###
	}
	if buffer.flags&BufferedReaderFlagEverything != 0 {
		msgStartIdx = 0
	}
	return buffer.data[msgStartIdx:nextMsgIdx], nextMsgIdx
}

// messages have a fixed size
func (buffer *BufferedReader) parseMLEFixed() ([]byte, int) {
	return buffer.extractMessage(buffer.paramMLE, 0)
}

// messages are separated by a delimiter string
func (buffer *BufferedReader) parseDelimiter() ([]byte, int) {
	delimiterIdx := bytes.Index(buffer.data[:buffer.end], buffer.delimiter)
	if delimiterIdx == -1 {
		return nil, 0 // ### return, incomplete ###
	}

	messageLen := delimiterIdx
	if buffer.flags&BufferedReaderFlagEverything != 0 {
		messageLen += len(buffer.delimiter)
	}

	data, nextMsgIdx := buffer.extractMessage(messageLen, 0)

	if data != nil && buffer.flags&BufferedReaderFlagEverything == 0 {
		nextMsgIdx += len(buffer.delimiter)
	}
	return data, nextMsgIdx
}

// messages are separeated length encoded by ASCII number and (an optional)
// delimiter.
func (buffer *BufferedReader) parseMLEText() ([]byte, int) {
	messageLen, msgStartIdx := tstrings.Btoi(buffer.data[buffer.paramMLE:buffer.end])
	if msgStartIdx == 0 {
		return nil, -1 // ### return, malformed ###
	}

	msgStartIdx += buffer.paramMLE
	// Read delimiter if necessary (check if valid runlength)
	if delimiterLen := len(buffer.delimiter); delimiterLen > 0 {
		msgStartIdx += delimiterLen
		if msgStartIdx >= buffer.end {
			return nil, 0 // ### return, incomplete ###
		}
		if !bytes.Equal(buffer.data[msgStartIdx-delimiterLen:msgStartIdx], buffer.delimiter) {
			return nil, -1 // ### return, malformed ###
		}
	}

	return buffer.extractMessage(int(messageLen), msgStartIdx)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE8() ([]byte, int) {
	var messageLen uint8
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+1)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE16() ([]byte, int) {
	var messageLen uint16
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+2)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE32() ([]byte, int) {
	var messageLen uint32
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+4)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE64() ([]byte, int) {
	var messageLen uint64
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+8)
}

// messages are separated by regexp
func (buffer *BufferedReader) parseDelimiterRegex() ([]byte, int) {
	startOffset := 0
	nextIdx := 0

	for {
		delimiterIdx := buffer.delimiterRegexp.FindIndex(buffer.data[startOffset:buffer.end])
		if delimiterIdx == nil {
			return nil, 0 // ### return, incomplete ###
		}

		// If we do end of message parsing, we're done
		if buffer.flags&BufferedReaderFlagRegexStart != BufferedReaderFlagRegexStart {
			nextIdx = delimiterIdx[1]
			break
		}

		// We're done if this is the second pass (start offset > 0)
		// We're done if we have data before the match (incomplete message)
		if startOffset > 0 || delimiterIdx[0] > 0 {
			nextIdx = delimiterIdx[0] + startOffset
			break
		}

		// Found delimiter at start of data, look for the start of a second message.
		// We don't need to add startOffset here as we never reach this if startOffset != 0
		startOffset = delimiterIdx[1]
	}

	return buffer.extractMessage(nextIdx, 0)
}

// ReadAll calls ReadOne as long as there are messages in the stream.
// Messages will be send to the given write callback.
// If callback is nil, data will be read and discarded.
func (buffer *BufferedReader) ReadAll(reader io.Reader, callback BufferReadCallback) error {
	for {
		data, more, err := buffer.ReadOne(reader)
		if data != nil && callback != nil {
			callback(data)
		}

		if err != nil {
			return err // ### return, error ###
		}

		if !more {
			return nil // ### return, done ###
		}
	}
}

// ReadOne reads the next message from the given stream (if possible) and
// generates a sequence number for this message.
// The more return parameter is set to true if there are still messages or parts
// of messages in the stream. Data and seq is only set if a complete message
// could be parsed.
// Errors are returned if reading from the stream failed or the parser
// encountered an error.
func (buffer *BufferedReader) ReadOne(reader io.Reader) (data []byte, more bool, err error) {
	if buffer.incomplete {
		bytesRead, err := reader.Read(buffer.data[buffer.end:])

		if err != nil && err != io.EOF {
			return nil, buffer.end > 0, err // ### return, error reading ###
		}

		if bytesRead == 0 {
			return nil, buffer.end > 0, err // ### return, no data ###
		}

		buffer.end += bytesRead
		buffer.incomplete = false
	}

	msgData, nextMsgIdx := buffer.parse()

	if nextMsgIdx == -1 {
		buffer.end = 0
		buffer.incomplete = true
		return nil, true, BufferDataInvalid // ### return, invalid data ###
	}

	if msgData == nil {
		// Check if buffer needs to be resized
		if len(buffer.data) == buffer.end {
			temp := buffer.data
			buffer.data = make([]byte, len(buffer.data)+buffer.growSize)
			copy(buffer.data, temp)
		}
		buffer.incomplete = true
		return nil, true, err // ### return, incomplete ###
	}

	msgDataCopy := make([]byte, len(msgData))
	copy(msgDataCopy, msgData)

	if nextMsgIdx < buffer.end {
		copy(buffer.data, buffer.data[nextMsgIdx:buffer.end])
		buffer.end -= nextMsgIdx
	} else {
		buffer.end = 0
		buffer.incomplete = true
	}

	return msgDataCopy, buffer.end > 0, err
}
<commit_msg>simplify loop body<commit_after>// Copyright 2015-2018 trivago N.V.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tio

import (
	"bytes"
	"encoding/binary"
	"io"
	"regexp"

	"github.com/trivago/tgo/tstrings"
)

// BufferedReaderFlags is an enum to configure a buffered reader
type BufferedReaderFlags byte

const (
	// BufferedReaderFlagDelimiter enables reading for a delimiter. This flag is
	// ignored if an MLE flag is set.
	BufferedReaderFlagDelimiter = BufferedReaderFlags(0)

	// BufferedReaderFlagMLE enables reading if length encoded messages.
	// Runlength is read as ASCII (to uint64) until the first byte (ASCII char)
	// of the delimiter string.
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE = BufferedReaderFlags(1)

	// BufferedReaderFlagMLE8 enables reading if length encoded messages.
	// Runlength is read as binary (to uint8).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE8 = BufferedReaderFlags(2)

	// BufferedReaderFlagMLE16 enables reading if length encoded messages.
	// Runlength is read as binary (to uint16).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE16 = BufferedReaderFlags(3)

	// BufferedReaderFlagMLE32 enables reading if length encoded messages.
	// Runlength is read as binary (to uint32).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE32 = BufferedReaderFlags(4)

	// BufferedReaderFlagMLE64 enables reading if length encoded messages.
	// Runlength is read as binary (to uint64).
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLE64 = BufferedReaderFlags(5)

	// BufferedReaderFlagMLEFixed enables reading messages with a fixed length.
	// Only one MLE flag is supported at a time.
	BufferedReaderFlagMLEFixed = BufferedReaderFlags(6)

	// BufferedReaderFlagMaskMLE is a bitmask to mask out everything but MLE flags
	BufferedReaderFlagMaskMLE = BufferedReaderFlags(7)

	// BufferedReaderFlagBigEndian sets binary reading to big endian encoding.
	BufferedReaderFlagBigEndian = BufferedReaderFlags(8)

	// BufferedReaderFlagEverything will keep MLE and/or delimiters when
	// building a message.
	BufferedReaderFlagEverything = BufferedReaderFlags(16)

	// BufferedReaderFlagRegex enables reading messages delimited by regexp.
	// To enable multiline reading, prefix the regex with "(?m)"
	BufferedReaderFlagRegex = BufferedReaderFlags(32)

	// BufferedReaderFlagRegexStart switches regex reading to find the start
	// of the next message instead of looking for the end of the current.
	BufferedReaderFlagRegexStart = BufferedReaderFlags(64 + 32)
)

type bufferError string

func (b bufferError) Error() string {
	return string(b)
}

// BufferReadCallback defines the function signature for callbacks passed to
// ReadAll.
type BufferReadCallback func(msg []byte)

// BufferDataInvalid is returned when a parsing encounters an error
var BufferDataInvalid = bufferError("Invalid data")

// BufferedReader is a helper struct to read from any io.Reader into a byte
// slice. The data can arrive "in pieces" and will be assembled.
// A data "piece" is considered complete if a delimiter or a certain runlength
// has been reached. The latter has to be enabled by flag and will disable the
// default behavior, which is looking for a delimiter string.
type BufferedReader struct {
	data            []byte
	delimiter       []byte
	parse           func() ([]byte, int)
	paramMLE        int
	growSize        int
	end             int
	encoding        binary.ByteOrder
	flags           BufferedReaderFlags
	incomplete      bool
	delimiterRegexp *regexp.Regexp
}

// NewBufferedReader creates a new buffered reader that reads messages from a
// continuous stream of bytes.
// Messages can be separated from the stream by using common methods such as
// fixed size, encoded message length or delimiter string.
// The internal buffer is grown statically (by its original size) if necessary.
// bufferSize defines the initial size / grow size of the buffer.
// flags configures the parsing method.
// offsetOrLength sets either the runlength offset or fixed message size.
// delimiter defines the delimiter used for textual message parsing, i.e.
// BufferedReaderFlagDelimiter or BufferedReaderFlagRegex.
func NewBufferedReader(bufferSize int, flags BufferedReaderFlags, offsetOrLength int, delimiter string) *BufferedReader {
	buffer := BufferedReader{
		data:            make([]byte, bufferSize),
		delimiter:       []byte(delimiter),
		paramMLE:        offsetOrLength,
		encoding:        binary.LittleEndian,
		end:             0,
		flags:           flags,
		growSize:        bufferSize,
		incomplete:      true,
		delimiterRegexp: nil,
	}

	if flags&BufferedReaderFlagBigEndian != 0 {
		buffer.encoding = binary.BigEndian
	}

	if flags&BufferedReaderFlagRegex != 0 {
		buffer.parse = buffer.parseDelimiterRegex
		buffer.delimiterRegexp = regexp.MustCompile(delimiter)
		return &buffer
	}

	if flags&BufferedReaderFlagMaskMLE == 0 {
		buffer.parse = buffer.parseDelimiter
	} else {
		switch flags & BufferedReaderFlagMaskMLE {
		default:
			buffer.parse = buffer.parseMLEText
		case BufferedReaderFlagMLE8:
			buffer.parse = buffer.parseMLE8
		case BufferedReaderFlagMLE16:
			buffer.parse = buffer.parseMLE16
		case BufferedReaderFlagMLE32:
			buffer.parse = buffer.parseMLE32
		case BufferedReaderFlagMLE64:
			buffer.parse = buffer.parseMLE64
		case BufferedReaderFlagMLEFixed:
			buffer.parse = buffer.parseMLEFixed
		}
	}

	return &buffer
}

// Reset clears the buffer by resetting its internal state
func (buffer *BufferedReader) Reset(sequence uint64) {
	buffer.end = 0
	buffer.incomplete = true
}

// ResetGetIncomplete works like Reset but returns any incomplete contents
// left in the buffer. Note that the data in the buffer returned is overwritten
// with the next read call.
func (buffer *BufferedReader) ResetGetIncomplete() []byte {
	remains := buffer.data[:buffer.end]
	buffer.Reset(0)
	return remains
}

// HasIncompleteData returns true if there is unparsed data left in the buffer
func (buffer *BufferedReader) HasIncompleteData() bool {
	return buffer.end > 0
}

// general message extraction part of all parser methods
func (buffer *BufferedReader) extractMessage(messageLen int, msgStartIdx int) ([]byte, int) {
	nextMsgIdx := msgStartIdx + messageLen
	if nextMsgIdx > buffer.end {
		return nil, 0 // ### return, incomplete ###
	}
	if buffer.flags&BufferedReaderFlagEverything != 0 {
		msgStartIdx = 0
	}
	return buffer.data[msgStartIdx:nextMsgIdx], nextMsgIdx
}

// messages have a fixed size
func (buffer *BufferedReader) parseMLEFixed() ([]byte, int) {
	return buffer.extractMessage(buffer.paramMLE, 0)
}

// messages are separated by a delimiter string
func (buffer *BufferedReader) parseDelimiter() ([]byte, int) {
	delimiterIdx := bytes.Index(buffer.data[:buffer.end], buffer.delimiter)
	if delimiterIdx == -1 {
		return nil, 0 // ### return, incomplete ###
	}

	messageLen := delimiterIdx
	if buffer.flags&BufferedReaderFlagEverything != 0 {
		messageLen += len(buffer.delimiter)
	}

	data, nextMsgIdx := buffer.extractMessage(messageLen, 0)

	if data != nil && buffer.flags&BufferedReaderFlagEverything == 0 {
		nextMsgIdx += len(buffer.delimiter)
	}
	return data, nextMsgIdx
}

// messages are separeated length encoded by ASCII number and (an optional)
// delimiter.
func (buffer *BufferedReader) parseMLEText() ([]byte, int) {
	messageLen, msgStartIdx := tstrings.Btoi(buffer.data[buffer.paramMLE:buffer.end])
	if msgStartIdx == 0 {
		return nil, -1 // ### return, malformed ###
	}

	msgStartIdx += buffer.paramMLE
	// Read delimiter if necessary (check if valid runlength)
	if delimiterLen := len(buffer.delimiter); delimiterLen > 0 {
		msgStartIdx += delimiterLen
		if msgStartIdx >= buffer.end {
			return nil, 0 // ### return, incomplete ###
		}
		if !bytes.Equal(buffer.data[msgStartIdx-delimiterLen:msgStartIdx], buffer.delimiter) {
			return nil, -1 // ### return, malformed ###
		}
	}

	return buffer.extractMessage(int(messageLen), msgStartIdx)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE8() ([]byte, int) {
	var messageLen uint8
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+1)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE16() ([]byte, int) {
	var messageLen uint16
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+2)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE32() ([]byte, int) {
	var messageLen uint32
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+4)
}

// messages are separated binary length encoded
func (buffer *BufferedReader) parseMLE64() ([]byte, int) {
	var messageLen uint64
	reader := bytes.NewReader(buffer.data[buffer.paramMLE:buffer.end])
	err := binary.Read(reader, buffer.encoding, &messageLen)
	if err != nil {
		return nil, -1 // ### return, malformed ###
	}
	return buffer.extractMessage(int(messageLen), buffer.paramMLE+8)
}

// messages are separated by regexp
func (buffer *BufferedReader) parseDelimiterRegex() ([]byte, int) {
	startOffset := 0

	for {
		delimiterIdx := buffer.delimiterRegexp.FindIndex(buffer.data[startOffset:buffer.end])
		if delimiterIdx == nil {
			return nil, 0 // ### return, incomplete ###
		}

		// If we do end of message parsing, we're done
		if buffer.flags&BufferedReaderFlagRegexStart != BufferedReaderFlagRegexStart {
			return buffer.extractMessage(delimiterIdx[1], 0) // ### done, end of line ###
		}

		// We're done if this is the second pass (start offset > 0)
		// We're done if we have data before the match (incomplete message)
		if startOffset > 0 || delimiterIdx[0] > 0 {
			return buffer.extractMessage(delimiterIdx[0]+startOffset, 0) // ### done, start of line ###
		}

		// Found delimiter at start of data, look for the start of a second message.
		// We don't need to add startOffset here as we never reach this if startOffset != 0
		startOffset = delimiterIdx[1]
	}
}

// ReadAll calls ReadOne as long as there are messages in the stream.
// Messages will be send to the given write callback.
// If callback is nil, data will be read and discarded.
func (buffer *BufferedReader) ReadAll(reader io.Reader, callback BufferReadCallback) error {
	for {
		data, more, err := buffer.ReadOne(reader)
		if data != nil && callback != nil {
			callback(data)
		}

		if err != nil {
			return err // ### return, error ###
		}

		if !more {
			return nil // ### return, done ###
		}
	}
}

// ReadOne reads the next message from the given stream (if possible) and
// generates a sequence number for this message.
// The more return parameter is set to true if there are still messages or parts
// of messages in the stream. Data and seq is only set if a complete message
// could be parsed.
// Errors are returned if reading from the stream failed or the parser
// encountered an error.
func (buffer *BufferedReader) ReadOne(reader io.Reader) (data []byte, more bool, err error) {
	if buffer.incomplete {
		bytesRead, err := reader.Read(buffer.data[buffer.end:])

		if err != nil && err != io.EOF {
			return nil, buffer.end > 0, err // ### return, error reading ###
		}

		if bytesRead == 0 {
			return nil, buffer.end > 0, err // ### return, no data ###
		}

		buffer.end += bytesRead
		buffer.incomplete = false
	}

	msgData, nextMsgIdx := buffer.parse()

	if nextMsgIdx == -1 {
		buffer.end = 0
		buffer.incomplete = true
		return nil, true, BufferDataInvalid // ### return, invalid data ###
	}

	if msgData == nil {
		// Check if buffer needs to be resized
		if len(buffer.data) == buffer.end {
			temp := buffer.data
			buffer.data = make([]byte, len(buffer.data)+buffer.growSize)
			copy(buffer.data, temp)
		}
		buffer.incomplete = true
		return nil, true, err // ### return, incomplete ###
	}

	msgDataCopy := make([]byte, len(msgData))
	copy(msgDataCopy, msgData)

	if nextMsgIdx < buffer.end {
		copy(buffer.data, buffer.data[nextMsgIdx:buffer.end])
		buffer.end -= nextMsgIdx
	} else {
		buffer.end = 0
		buffer.incomplete = true
	}

	return msgDataCopy, buffer.end > 0, err
}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"
	"os"
	"time"

	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/syndtr/goleveldb/leveldb/opt"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

func loadVolumeWithoutIndex(dirname string, collection string, id needle.VolumeId, needleMapKind NeedleMapType) (v *Volume, e error) {
	v = &Volume{dir: dirname, Collection: collection, Id: id}
	v.SuperBlock = SuperBlock{}
	v.needleMapKind = needleMapKind
	e = v.load(false, false, needleMapKind, 0)
	return
}

func (v *Volume) load(alsoLoadIndex bool, createDatIfMissing bool, needleMapKind NeedleMapType, preallocate int64) error {
	var e error
	fileName := v.FileName()
	alreadyHasSuperBlock := false

	if exists, canRead, canWrite, modifiedTime, fileSize := checkFile(fileName + ".dat"); exists {
		if !canRead {
			return fmt.Errorf("cannot read Volume Data file %s.dat", fileName)
		}
		if canWrite {
			v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
			v.lastModifiedTsSeconds = uint64(modifiedTime.Unix())
		} else {
			glog.V(0).Infoln("opening " + fileName + ".dat in READONLY mode")
			v.dataFile, e = os.Open(fileName + ".dat")
			v.readOnly = true
		}
		if fileSize >= _SuperBlockSize {
			alreadyHasSuperBlock = true
		}
	} else {
		if createDatIfMissing {
			v.dataFile, e = createVolumeFile(fileName+".dat", preallocate)
		} else {
			return fmt.Errorf("Volume Data file %s.dat does not exist.", fileName)
		}
	}

	if e != nil {
		if !os.IsPermission(e) {
			return fmt.Errorf("cannot load Volume Data %s.dat: %v", fileName, e)
		} else {
			return fmt.Errorf("load data file %s.dat: %v", fileName, e)
		}
	}

	if alreadyHasSuperBlock {
		e = v.readSuperBlock()
	} else {
		e = v.maybeWriteSuperBlock()
	}
	if e == nil && alsoLoadIndex {
		var indexFile *os.File
		if v.readOnly {
			glog.V(1).Infoln("open to read file", fileName+".idx")
			if indexFile, e = os.OpenFile(fileName+".idx", os.O_RDONLY, 0644); e != nil {
				return fmt.Errorf("cannot read Volume Index %s.idx: %v", fileName, e)
			}
		} else {
			glog.V(1).Infoln("open to write file", fileName+".idx")
			if indexFile, e = os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644); e != nil {
				return fmt.Errorf("cannot write Volume Index %s.idx: %v", fileName, e)
			}
		}
		if v.lastAppendAtNs, e = CheckVolumeDataIntegrity(v, indexFile); e != nil {
			v.readOnly = true
			glog.V(0).Infof("volumeDataIntegrityChecking failed %v", e)
		}
		switch needleMapKind {
		case NeedleMapInMemory:
			glog.V(0).Infoln("loading index", fileName+".idx", "to memory readonly", v.readOnly)
			if v.nm, e = LoadCompactNeedleMap(indexFile); e != nil {
				glog.V(0).Infof("loading index %s to memory error: %v", fileName+".idx", e)
			}
		case NeedleMapLevelDb:
			glog.V(0).Infoln("loading leveldb", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            2 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   1 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		case NeedleMapLevelDbMedium:
			glog.V(0).Infoln("loading leveldb medium", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            4 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   2 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		case NeedleMapLevelDbLarge:
			glog.V(0).Infoln("loading leveldb large", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            8 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   4 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		}
	}

	stats.VolumeServerVolumeCounter.WithLabelValues(v.Collection, "volume").Inc()

	return e
}

func checkFile(filename string) (exists, canRead, canWrite bool, modTime time.Time, fileSize int64) {
	exists = true
	fi, err := os.Stat(filename)
	if os.IsNotExist(err) {
		exists = false
		return
	}
	if fi.Mode()&0400 != 0 {
		canRead = true
	}
	if fi.Mode()&0200 != 0 {
		canWrite = true
	}
	modTime = fi.ModTime()
	fileSize = fi.Size()
	return
}
<commit_msg>volume: read dat files' last modified time no matter dat files are catWrite or readonly<commit_after>package storage

import (
	"fmt"
	"os"
	"time"

	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/syndtr/goleveldb/leveldb/opt"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

func loadVolumeWithoutIndex(dirname string, collection string, id needle.VolumeId, needleMapKind NeedleMapType) (v *Volume, e error) {
	v = &Volume{dir: dirname, Collection: collection, Id: id}
	v.SuperBlock = SuperBlock{}
	v.needleMapKind = needleMapKind
	e = v.load(false, false, needleMapKind, 0)
	return
}

func (v *Volume) load(alsoLoadIndex bool, createDatIfMissing bool, needleMapKind NeedleMapType, preallocate int64) error {
	var e error
	fileName := v.FileName()
	alreadyHasSuperBlock := false

	if exists, canRead, canWrite, modifiedTime, fileSize := checkFile(fileName + ".dat"); exists {
		if !canRead {
			return fmt.Errorf("cannot read Volume Data file %s.dat", fileName)
		}
		if canWrite {
			v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
		} else {
			glog.V(0).Infoln("opening " + fileName + ".dat in READONLY mode")
			v.dataFile, e = os.Open(fileName + ".dat")
			v.readOnly = true
		}
		v.lastModifiedTsSeconds = uint64(modifiedTime.Unix())
		if fileSize >= _SuperBlockSize {
			alreadyHasSuperBlock = true
		}
	} else {
		if createDatIfMissing {
			v.dataFile, e = createVolumeFile(fileName+".dat", preallocate)
		} else {
			return fmt.Errorf("Volume Data file %s.dat does not exist.", fileName)
		}
	}

	if e != nil {
		if !os.IsPermission(e) {
			return fmt.Errorf("cannot load Volume Data %s.dat: %v", fileName, e)
		} else {
			return fmt.Errorf("load data file %s.dat: %v", fileName, e)
		}
	}

	if alreadyHasSuperBlock {
		e = v.readSuperBlock()
	} else {
		e = v.maybeWriteSuperBlock()
	}
	if e == nil && alsoLoadIndex {
		var indexFile *os.File
		if v.readOnly {
			glog.V(1).Infoln("open to read file", fileName+".idx")
			if indexFile, e = os.OpenFile(fileName+".idx", os.O_RDONLY, 0644); e != nil {
				return fmt.Errorf("cannot read Volume Index %s.idx: %v", fileName, e)
			}
		} else {
			glog.V(1).Infoln("open to write file", fileName+".idx")
			if indexFile, e = os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644); e != nil {
				return fmt.Errorf("cannot write Volume Index %s.idx: %v", fileName, e)
			}
		}
		if v.lastAppendAtNs, e = CheckVolumeDataIntegrity(v, indexFile); e != nil {
			v.readOnly = true
			glog.V(0).Infof("volumeDataIntegrityChecking failed %v", e)
		}
		switch needleMapKind {
		case NeedleMapInMemory:
			glog.V(0).Infoln("loading index", fileName+".idx", "to memory readonly", v.readOnly)
			if v.nm, e = LoadCompactNeedleMap(indexFile); e != nil {
				glog.V(0).Infof("loading index %s to memory error: %v", fileName+".idx", e)
			}
		case NeedleMapLevelDb:
			glog.V(0).Infoln("loading leveldb", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            2 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   1 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		case NeedleMapLevelDbMedium:
			glog.V(0).Infoln("loading leveldb medium", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            4 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   2 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		case NeedleMapLevelDbLarge:
			glog.V(0).Infoln("loading leveldb large", fileName+".ldb")
			opts := &opt.Options{
				BlockCacheCapacity:            8 * 1024 * 1024, // default value is 8MiB
				WriteBuffer:                   4 * 1024 * 1024, // default value is 4MiB
				CompactionTableSizeMultiplier: 10,              // default value is 1
			}
			if v.nm, e = NewLevelDbNeedleMap(fileName+".ldb", indexFile, opts); e != nil {
				glog.V(0).Infof("loading leveldb %s error: %v", fileName+".ldb", e)
			}
		}
	}

	stats.VolumeServerVolumeCounter.WithLabelValues(v.Collection, "volume").Inc()

	return e
}

func checkFile(filename string) (exists, canRead, canWrite bool, modTime time.Time, fileSize int64) {
	exists = true
	fi, err := os.Stat(filename)
	if os.IsNotExist(err) {
		exists = false
		return
	}
	if fi.Mode()&0400 != 0 {
		canRead = true
	}
	if fi.Mode()&0200 != 0 {
		canWrite = true
	}
	modTime = fi.ModTime()
	fileSize = fi.Size()
	return
}
<|endoftext|>


<commit_before>package scenario

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"

	"github.com/catatsuy/isucon6-final/bench/action"
	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/seed"
	"github.com/catatsuy/isucon6-final/bench/session"
)

// 部屋を作って線を描くとトップページに出てくる
func StrokeReflectedToTop(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	seedStrokes := seed.GetStrokes("star")
	seedStroke := seed.FluctuateStroke(seedStrokes[0])
	_, ok = drawStroke(s1, token, room.ID, seedStroke)
	if !ok {
		fails.Critical("線の投稿に失敗しました", nil)
		return
	}

	// 描いた直後にトップページに表示される
	_ = action.Get(s2, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		imageUrls := extractImages(doc)

		found := false
		for _, img := range imageUrls {
			if img == "/img/"+strconv.FormatInt(room.ID, 10) {
				found = true
			}
		}
		if !found {
			l.Critical("投稿が反映されていません", nil)
			return false
		}
		return true
	}))
}

// 線の描かれてない部屋はトップページに並ばない
func RoomWithoutStrokeNotShownAtTop(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	_ = action.Get(s2, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		imageUrls := extractImages(doc)

		for _, img := range imageUrls {
			if img == "/img/"+strconv.FormatInt(room.ID, 10) {
				l.Critical("まだ線の無い部屋が表示されています", nil)
				return false
			}
		}
		return true
	}))
}

// 線がSVGに反映される
func StrokeReflectedToSVG(origins []string) {
	s1 := session.New(randomOrigin(origins))
	defer s1.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	seedStrokes := seed.GetStrokes("wwws")
	for _, seedStroke := range seedStrokes {
		stroke2 := seed.FluctuateStroke(seedStroke)
		stroke, ok := drawStroke(s1, token, room.ID, stroke2)
		if !ok {
			fails.Critical("線の投稿に失敗しました", nil)
			return
		}

		s2 := session.New(randomOrigin(origins))
		_ = checkStrokeReflectedToSVG(s2, room.ID, stroke.ID, stroke2)
		s2.Bye()
	}
}

// ページ内のCSRFトークンが毎回変わっている
func CSRFTokenRefreshed(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token1, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	token2, ok := fetchCSRFToken(s2, "/")
	if !ok {
		return
	}

	if token1 == token2 {
		fails.Critical("csrf_tokenが使いまわされています", nil)
	}
}

// 他人の作った部屋に最初の線を描けない
func CantDrawFirstStrokeOnSomeoneElsesRoom(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token1, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token1)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	token2, ok := fetchCSRFToken(s2, "/")
	if !ok {
		return
	}

	strokes := seed.GetStrokes("star")
	stroke := seed.FluctuateStroke(strokes[0])

	postBody, _ := json.Marshal(struct {
		RoomID int64 `json:"room_id"`
		seed.Stroke
	}{
		RoomID: room.ID,
		Stroke: stroke,
	})

	headers := map[string]string{
		"Content-Type": "application/json",
		"x-csrf-token": token2,
	}

	u := "/api/strokes/rooms/" + strconv.FormatInt(room.ID, 10)
	ok = action.Post(s2, u, postBody, headers, action.BadRequest(func(body io.Reader, l *fails.Logger) bool {
		// JSONも検証する？
		return true
	}))
	if !ok {
		fails.Critical("他人の作成した部屋に1画目を描くことができました", nil)
	}
}

// トップページの内容が正しいかをチェック
func TopPageContent(origins []string) {
	s := session.New(randomOrigin(origins))
	defer s.Bye()

	_ = action.Get(s, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		images := extractImages(doc)
		if len(images) < 100 {
			l.Critical("画像の枚数が少なすぎます", nil)
			return false
		}

		reactidNum := doc.Find("[data-reactid]").Length()
		expected := 1325
		if reactidNum != expected {
			l.Critical("トップページの内容が正しくありません",
				fmt.Errorf("data-reactidの数が一致しません (expected %d, actual %d)", expected, reactidNum))
			return false
		}
		return true
	}))
}

// 静的ファイルが正しいかをチェック
func CheckStaticFiles(origins []string) {
	s := session.New(randomOrigin(origins))
	defer s.Bye()

	ok := loadStaticFiles(s, true /*checkHash*/)
	if !ok {
		fails.Critical("静的ファイルが正しくありません", nil)
	}
}
<commit_msg>Check if created_at in response is in between request and response<commit_after>package scenario

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"time"

	"github.com/catatsuy/isucon6-final/bench/action"
	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/seed"
	"github.com/catatsuy/isucon6-final/bench/session"
)

// 部屋を作って線を描くとトップページに出てくる
func StrokeReflectedToTop(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	t1 := time.Now()

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	t2 := time.Now()
	if room.CreatedAt.After(t2) || room.CreatedAt.Before(t1) {
		fails.Critical("作成した部屋のcreated_atが正しくありません",
			fmt.Errorf("should be %s < %s < %s", t1.Format("2006-01-02-15:04:05.000"), room.CreatedAt.Format("2006-01-02-15:04:05.000"), t2.Format("2006-01-02-15:04:05.000")))
	}

	seedStrokes := seed.GetStrokes("star")
	seedStroke := seed.FluctuateStroke(seedStrokes[0])
	stroke, ok := drawStroke(s1, token, room.ID, seedStroke)
	if !ok {
		fails.Critical("線の投稿に失敗しました", nil)
		return
	}

	t3 := time.Now()
	if stroke.CreatedAt.After(t3) || stroke.CreatedAt.Before(t2) {
		fails.Critical("作成した部屋のcreated_atが正しくありません",
			fmt.Errorf("should be %s < %s < %s", t2.Format("2006-01-02-15:04:05.000"), stroke.CreatedAt.Format("2006-01-02-15:04:05.000"), t3.Format("2006-01-02-15:04:05.000")))
	}

	// 描いた直後にトップページに表示される
	_ = action.Get(s2, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		imageUrls := extractImages(doc)

		found := false
		for _, img := range imageUrls {
			if img == "/img/"+strconv.FormatInt(room.ID, 10) {
				found = true
			}
		}
		if !found {
			l.Critical("投稿が反映されていません", nil)
			return false
		}
		return true
	}))
}

// 線の描かれてない部屋はトップページに並ばない
func RoomWithoutStrokeNotShownAtTop(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	_ = action.Get(s2, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		imageUrls := extractImages(doc)

		for _, img := range imageUrls {
			if img == "/img/"+strconv.FormatInt(room.ID, 10) {
				l.Critical("まだ線の無い部屋が表示されています", nil)
				return false
			}
		}
		return true
	}))
}

// 線がSVGに反映される
func StrokeReflectedToSVG(origins []string) {
	s1 := session.New(randomOrigin(origins))
	defer s1.Bye()

	token, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	seedStrokes := seed.GetStrokes("wwws")
	for _, seedStroke := range seedStrokes {
		stroke2 := seed.FluctuateStroke(seedStroke)
		stroke, ok := drawStroke(s1, token, room.ID, stroke2)
		if !ok {
			fails.Critical("線の投稿に失敗しました", nil)
			return
		}

		s2 := session.New(randomOrigin(origins))
		_ = checkStrokeReflectedToSVG(s2, room.ID, stroke.ID, stroke2)
		s2.Bye()
	}
}

// ページ内のCSRFトークンが毎回変わっている
func CSRFTokenRefreshed(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token1, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	token2, ok := fetchCSRFToken(s2, "/")
	if !ok {
		return
	}

	if token1 == token2 {
		fails.Critical("csrf_tokenが使いまわされています", nil)
	}
}

// 他人の作った部屋に最初の線を描けない
func CantDrawFirstStrokeOnSomeoneElsesRoom(origins []string) {
	s1 := session.New(randomOrigin(origins))
	s2 := session.New(randomOrigin(origins))
	defer s1.Bye()
	defer s2.Bye()

	token1, ok := fetchCSRFToken(s1, "/")
	if !ok {
		return
	}

	room, ok := makeRoom(s1, token1)
	if !ok {
		fails.Critical("部屋の作成に失敗しました", nil)
		return
	}

	token2, ok := fetchCSRFToken(s2, "/")
	if !ok {
		return
	}

	strokes := seed.GetStrokes("star")
	stroke := seed.FluctuateStroke(strokes[0])

	postBody, _ := json.Marshal(struct {
		RoomID int64 `json:"room_id"`
		seed.Stroke
	}{
		RoomID: room.ID,
		Stroke: stroke,
	})

	headers := map[string]string{
		"Content-Type": "application/json",
		"x-csrf-token": token2,
	}

	u := "/api/strokes/rooms/" + strconv.FormatInt(room.ID, 10)
	ok = action.Post(s2, u, postBody, headers, action.BadRequest(func(body io.Reader, l *fails.Logger) bool {
		// JSONも検証する？
		return true
	}))
	if !ok {
		fails.Critical("他人の作成した部屋に1画目を描くことができました", nil)
	}
}

// トップページの内容が正しいかをチェック
func TopPageContent(origins []string) {
	s := session.New(randomOrigin(origins))
	defer s.Bye()

	_ = action.Get(s, "/", action.OK(func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}
		images := extractImages(doc)
		if len(images) < 100 {
			l.Critical("画像の枚数が少なすぎます", nil)
			return false
		}

		reactidNum := doc.Find("[data-reactid]").Length()
		expected := 1325
		if reactidNum != expected {
			l.Critical("トップページの内容が正しくありません",
				fmt.Errorf("data-reactidの数が一致しません (expected %d, actual %d)", expected, reactidNum))
			return false
		}
		return true
	}))
}

// 静的ファイルが正しいかをチェック
func CheckStaticFiles(origins []string) {
	s := session.New(randomOrigin(origins))
	defer s.Bye()

	ok := loadStaticFiles(s, true /*checkHash*/)
	if !ok {
		fails.Critical("静的ファイルが正しくありません", nil)
	}
}
<|endoftext|>


<commit_before>package stats

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/guigolab/bamstats/annotation"
	"github.com/guigolab/bamstats/sam"
)

// general element constants
const (
	Exon       = "exon"
	ExonIntron = "exonic_intronic"
	Intergenic = "intergenic"
	Intron     = "intron"
	Other      = "others"
	Total      = "total"
)

var (
	allElems = []string{
		ExonIntron,
		Intron,
		Exon,
		Intergenic,
		Other,
		Total,
	}
)

func getElemSet() map[string]struct{} {
	elemSet := make(map[string]struct{})
	for _, el := range allElems {
		elemSet[el] = struct{}{}
	}
	return elemSet
}

// ElementStats represents mappings statistics for genomic elements
type ElementStats map[string]uint64

// Keys returns map keys
func (s ElementStats) Keys() []string {
	var keys []string
	elemSet := getElemSet()
	for k := range s {
		if _, found := elemSet[k]; !found {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	return append(keys, allElems...)
}

// MergeKeys combine keys from two ElementsStats instances
func (s ElementStats) MergeKeys(other ElementStats) []string {
	elemSet := getElemSet()
	var keys []string
	for _, k := range s.Keys() {
		if _, found := elemSet[k]; !found {
			keys = append(keys, k)
		}
	}
	for _, k := range other.Keys() {
		if _, found := elemSet[k]; !found {
			if _, hasThis := s[k]; !hasThis {
				keys = append(keys, k)
			}
		}
	}
	sort.Strings(keys)
	return append(keys, allElems...)
}

// CoverageStats represents genome coverage statistics for continuos, split and total mapped reads.
type CoverageStats struct {
	Total      ElementStats `json:"total"`
	Continuous ElementStats `json:"continuous"`
	Split      ElementStats `json:"split"`
	Uniq       bool         `json:"-"`
	index      *annotation.RtreeMap
}

// Update updates all counts from a Stats instance.
func (s *CoverageStats) Update(other Stats) {
	if other, ok := other.(*CoverageStats); ok {
		s.Continuous.Update(other.Continuous)
		s.Split.Update(other.Split)
		s.Finalize()
	}
}

// Merge update counts from a channel of Stats instances.
func (s *CoverageStats) Merge(others chan Stats) {
	for other := range others {
		if other, ok := other.(*CoverageStats); ok {
			s.Update(other)
		}
	}
}

// Finalize updates dependent counts of a CoverageStats instance.
func (s *CoverageStats) Finalize() {
	for _, elem := range s.Continuous.MergeKeys(s.Split) {
		s.Total[elem] = s.Continuous[elem] + s.Split[elem]
	}
}

// Update updates all counts from another ElementsStats instance.
func (s ElementStats) Update(other ElementStats) {
	for k := range other {
		s[k] += other[k]
	}
}

// MarshalJSON implements JSON Marshaller interface
func (s ElementStats) MarshalJSON() ([]byte, error) {
	buf := &bytes.Buffer{}
	buf.Write([]byte{'{', '\n'})
	keys := s.Keys()
	l := len(keys)
	for i, k := range keys {
		// TODO: write 0 values elements? Update test output in case
		// if s[k] == 0 {
		// 	continue
		// }
		fmt.Fprintf(buf, "\t\"%s\": %d", k, s[k])
		if i < l-1 {
			buf.WriteByte(',')
		}
		buf.WriteByte('\n')
	}
	buf.Write([]byte{'}', '\n'})
	return buf.Bytes(), nil
}

func updateCount(elems map[string]uint8, st ElementStats) {
	if len(elems) == 0 {
		return
	}

	exons, hasExon := elems["exon"]
	introns, hasIntron := elems["intron"]
	st[Total]++
	if _, isIntergenic := elems["intergenic"]; isIntergenic {
		if len(elems) > 1 {
			st[Other]++
		} else {
			st[Intergenic]++
		}
		return
	}

	if hasExon && !hasIntron && exons > 0 {
		st[Exon]++
		return
	}
	if hasIntron && !hasExon && introns > 0 {
		st[Intron]++
		return
	}
	if hasIntron && hasExon && introns > 0 && exons > 0 {
		st[ExonIntron]++
		return
	}
	var keys []string
	for k := range elems {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var buf bytes.Buffer
	for i, e := range keys {
		if i > 0 {
			buf.WriteByte('_')
		}
		buf.WriteString(e)
	}
	st[buf.String()]++
}

// Collect collects genome coverage statistics from a sam.Record.
func (s *CoverageStats) Collect(record *sam.Record) {
	if s.index == nil || !record.IsPrimary() || record.IsUnmapped() {
		return
	}
	if s.Uniq && !record.IsUniq() {
		return
	}
	elements := map[string]uint8{}
	for _, mappingLocation := range record.GetBlocks() {
		rtree := s.index.Get(mappingLocation.Chrom())
		if rtree == nil || rtree.Size() == 0 {
			return
		}
		results := annotation.QueryIndex(rtree, mappingLocation.Start(), mappingLocation.End())
		mappingLocation.GetElements(&results, elements)
	}
	if record.IsSplit() {
		updateCount(elements, s.Split)
	} else {
		updateCount(elements, s.Continuous)
	}
}

// NewCoverageStats create a new instance of CoverageStats.
func NewCoverageStats(index *annotation.RtreeMap, uniq bool) *CoverageStats {
	return &CoverageStats{
		Total:      make(ElementStats),
		Continuous: make(ElementStats),
		Split:      make(ElementStats),
		Uniq:       uniq,
		index:      index,
	}
}
<commit_msg>[ci skip] Remove custom elements from converage stats<commit_after>package stats

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/guigolab/bamstats/annotation"
	"github.com/guigolab/bamstats/sam"
)

// general element constants
const (
	Exon       = "exon"
	ExonIntron = "exonic_intronic"
	Intergenic = "intergenic"
	Intron     = "intron"
	Other      = "others"
	Total      = "total"
)

var (
	allElems = []string{
		ExonIntron,
		Intron,
		Exon,
		Intergenic,
		Other,
		Total,
	}
)

func getElemSet() map[string]struct{} {
	elemSet := make(map[string]struct{})
	for _, el := range allElems {
		elemSet[el] = struct{}{}
	}
	return elemSet
}

// ElementStats represents mappings statistics for genomic elements
type ElementStats map[string]uint64

// Keys returns map keys
func (s ElementStats) Keys() []string {
	var keys []string
	elemSet := getElemSet()
	for k := range s {
		if _, found := elemSet[k]; !found {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	return append(keys, allElems...)
}

// MergeKeys combine keys from two ElementsStats instances
func (s ElementStats) MergeKeys(other ElementStats) []string {
	elemSet := getElemSet()
	var keys []string
	for _, k := range s.Keys() {
		if _, found := elemSet[k]; !found {
			keys = append(keys, k)
		}
	}
	for _, k := range other.Keys() {
		if _, found := elemSet[k]; !found {
			if _, hasThis := s[k]; !hasThis {
				keys = append(keys, k)
			}
		}
	}
	sort.Strings(keys)
	return append(keys, allElems...)
}

// CoverageStats represents genome coverage statistics for continuos, split and total mapped reads.
type CoverageStats struct {
	Total      ElementStats `json:"total"`
	Continuous ElementStats `json:"continuous"`
	Split      ElementStats `json:"split"`
	Uniq       bool         `json:"-"`
	index      *annotation.RtreeMap
}

// Update updates all counts from a Stats instance.
func (s *CoverageStats) Update(other Stats) {
	if other, ok := other.(*CoverageStats); ok {
		s.Continuous.Update(other.Continuous)
		s.Split.Update(other.Split)
		s.Finalize()
	}
}

// Merge update counts from a channel of Stats instances.
func (s *CoverageStats) Merge(others chan Stats) {
	for other := range others {
		if other, ok := other.(*CoverageStats); ok {
			s.Update(other)
		}
	}
}

// Finalize updates dependent counts of a CoverageStats instance.
func (s *CoverageStats) Finalize() {
	for _, elem := range s.Continuous.MergeKeys(s.Split) {
		s.Total[elem] = s.Continuous[elem] + s.Split[elem]
	}
}

// Update updates all counts from another ElementsStats instance.
func (s ElementStats) Update(other ElementStats) {
	for k := range other {
		s[k] += other[k]
	}
}

// MarshalJSON implements JSON Marshaller interface
func (s ElementStats) MarshalJSON() ([]byte, error) {
	buf := &bytes.Buffer{}
	buf.Write([]byte{'{', '\n'})
	keys := s.Keys()
	l := len(keys)
	for i, k := range keys {
		// TODO: write 0 values elements? Update test output in case
		// if s[k] == 0 {
		// 	continue
		// }
		fmt.Fprintf(buf, "\t\"%s\": %d", k, s[k])
		if i < l-1 {
			buf.WriteByte(',')
		}
		buf.WriteByte('\n')
	}
	buf.Write([]byte{'}', '\n'})
	return buf.Bytes(), nil
}

func updateCount(elems map[string]uint8, st ElementStats) {
	if len(elems) == 0 {
		return
	}

	exons, hasExon := elems["exon"]
	introns, hasIntron := elems["intron"]
	st[Total]++
	if _, isIntergenic := elems["intergenic"]; isIntergenic {
		if len(elems) > 1 {
			st[Other]++
		} else {
			st[Intergenic]++
		}
		return
	}

	if hasExon && !hasIntron && exons > 0 {
		st[Exon]++
		return
	}
	if hasIntron && !hasExon && introns > 0 {
		st[Intron]++
		return
	}
	if hasIntron && hasExon && introns > 0 && exons > 0 {
		st[ExonIntron]++
		return
	}
}

// Collect collects genome coverage statistics from a sam.Record.
func (s *CoverageStats) Collect(record *sam.Record) {
	if s.index == nil || !record.IsPrimary() || record.IsUnmapped() {
		return
	}
	if s.Uniq && !record.IsUniq() {
		return
	}
	elements := map[string]uint8{}
	for _, mappingLocation := range record.GetBlocks() {
		rtree := s.index.Get(mappingLocation.Chrom())
		if rtree == nil || rtree.Size() == 0 {
			return
		}
		results := annotation.QueryIndex(rtree, mappingLocation.Start(), mappingLocation.End())
		mappingLocation.GetElements(&results, elements)
	}
	if record.IsSplit() {
		updateCount(elements, s.Split)
	} else {
		updateCount(elements, s.Continuous)
	}
}

// NewCoverageStats create a new instance of CoverageStats.
func NewCoverageStats(index *annotation.RtreeMap, uniq bool) *CoverageStats {
	return &CoverageStats{
		Total:      make(ElementStats),
		Continuous: make(ElementStats),
		Split:      make(ElementStats),
		Uniq:       uniq,
		index:      index,
	}
}
<|endoftext|>