<commit_before>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package span

import (
	"encoding/json"
	"fmt"
)

// Span represents a source code range in standardized form.
type Span struct {
	v span
}

// Point represents a single point within a file.
// In general this should only be used as part of a Span, as on its own it
// does not carry enough information.
type Point struct {
	v point
}

type span struct {
	URI   URI   `json:"uri"`
	Start point `json:"start"`
	End   point `json:"end"`
}

type point struct {
	Line   int `json:"line"`
	Column int `json:"column"`
	Offset int `json:"offset"`
}

var invalidPoint = Point{v: point{Line: 0, Column: 0, Offset: -1}}

// Converter is the interface to an object that can convert between line:column
// and offset forms for a single file.
type Converter interface {
	//ToPosition converts from an offset to a line:column pair.
	ToPosition(offset int) (int, int, error)
	//ToOffset converts from a line:column pair to an offset.
	ToOffset(line, col int) (int, error)
}

func New(uri URI, start Point, end Point) Span {
	s := Span{v: span{URI: uri, Start: start.v, End: end.v}}
	s.v.clean()
	return s
}

func NewPoint(line, col, offset int) Point {
	p := Point{v: point{Line: line, Column: col, Offset: offset}}
	p.v.clean()
	return p
}

func (s Span) HasPosition() bool             { return s.v.Start.hasPosition() }
func (s Span) HasOffset() bool               { return s.v.Start.hasOffset() }
func (s Span) IsValid() bool                 { return s.v.Start.isValid() }
func (s Span) IsPoint() bool                 { return s.v.Start == s.v.End }
func (s Span) URI() URI                      { return s.v.URI }
func (s Span) Start() Point                  { return Point{s.v.Start} }
func (s Span) End() Point                    { return Point{s.v.End} }
func (s *Span) MarshalJSON() ([]byte, error) { return json.Marshal(&s.v) }
func (s *Span) UnmarshalJSON(b []byte) error { return json.Unmarshal(b, &s.v) }

func (p Point) HasPosition() bool             { return p.v.hasPosition() }
func (p Point) HasOffset() bool               { return p.v.hasOffset() }
func (p Point) IsValid() bool                 { return p.v.isValid() }
func (p *Point) MarshalJSON() ([]byte, error) { return json.Marshal(&p.v) }
func (p *Point) UnmarshalJSON(b []byte) error { return json.Unmarshal(b, &p.v) }
func (p Point) Line() int {
	if !p.v.hasPosition() {
		panic(fmt.Errorf("position not set in %v", p.v))
	}
	return p.v.Line
}
func (p Point) Column() int {
	if !p.v.hasPosition() {
		panic(fmt.Errorf("position not set in %v", p.v))
	}
	return p.v.Column
}
func (p Point) Offset() int {
	if !p.v.hasOffset() {
		panic(fmt.Errorf("offset not set in %v", p.v))
	}
	return p.v.Offset
}

func (p point) hasPosition() bool { return p.Line > 0 }
func (p point) hasOffset() bool   { return p.Offset >= 0 }
func (p point) isValid() bool     { return p.hasPosition() || p.hasOffset() }
func (p point) isZero() bool {
	return (p.Line == 1 && p.Column == 1) || (!p.hasPosition() && p.Offset == 0)
}

func (s *span) clean() {
	//this presumes the points are already clean
	if !s.End.isValid() || (s.End == point{}) {
		s.End = s.Start
	}
}

func (p *point) clean() {
	if p.Line < 0 {
		p.Line = 0
	}
	if p.Column <= 0 {
		if p.Line > 0 {
			p.Column = 1
		} else {
			p.Column = 0
		}
	}
	if p.Offset == 0 && (p.Line > 1 || p.Column > 1) {
		p.Offset = -1
	}
}

// Format implements fmt.Formatter to print the Location in a standard form.
// The format produced is one that can be read back in using Parse.
func (s Span) Format(f fmt.State, c rune) {
	fullForm := f.Flag('+')
	preferOffset := f.Flag('#')
	// we should always have a uri, simplify if it is file format
	//TODO: make sure the end of the uri is unambiguous
	uri := string(s.v.URI)
	if !fullForm {
		if filename, err := s.v.URI.Filename(); err == nil {
			uri = filename
		}
	}
	fmt.Fprint(f, uri)
	if !s.IsValid() || (!fullForm && s.v.Start.isZero() && s.v.End.isZero()) {
		return
	}
	// see which bits of start to write
	printOffset := s.HasOffset() && (fullForm || preferOffset || !s.HasPosition())
	printLine := s.HasPosition() && (fullForm || !printOffset)
	printColumn := printLine && (fullForm || (s.v.Start.Column > 1 || s.v.End.Column > 1))
	fmt.Fprint(f, ":")
	if printLine {
		fmt.Fprintf(f, "%d", s.v.Start.Line)
	}
	if printColumn {
		fmt.Fprintf(f, ":%d", s.v.Start.Column)
	}
	if printOffset {
		fmt.Fprintf(f, "#%d", s.v.Start.Offset)
	}
	// start is written, do we need end?
	if s.IsPoint() {
		return
	}
	// we don't print the line if it did not change
	printLine = fullForm || (printLine && s.v.End.Line > s.v.Start.Line)
	fmt.Fprint(f, "-")
	if printLine {
		fmt.Fprintf(f, "%d", s.v.End.Line)
	}
	if printColumn {
		if printLine {
			fmt.Fprint(f, ":")
		}
		fmt.Fprintf(f, "%d", s.v.End.Column)
	}
	if printOffset {
		fmt.Fprintf(f, "#%d", s.v.End.Offset)
	}
}

func (s Span) WithPosition(c Converter) (Span, error) {
	if err := s.update(c, true, false); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s Span) WithOffset(c Converter) (Span, error) {
	if err := s.update(c, false, true); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s Span) WithAll(c Converter) (Span, error) {
	if err := s.update(c, true, true); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s *Span) update(c Converter, withPos, withOffset bool) error {
	if !s.IsValid() {
		return fmt.Errorf("cannot add information to an invalid span")
	}
	if withPos && !s.HasPosition() {
		if err := s.v.Start.updatePosition(c); err != nil {
			return err
		}
		if s.v.End.Offset == s.v.Start.Offset {
			s.v.End = s.v.Start
		} else if err := s.v.End.updatePosition(c); err != nil {
			return err
		}
	}
	if withOffset && !s.HasOffset() {
		if err := s.v.Start.updateOffset(c); err != nil {
			return err
		}
		if s.v.End.Line == s.v.Start.Line && s.v.End.Column == s.v.Start.Column {
			s.v.End.Offset = s.v.Start.Offset
		} else if err := s.v.End.updateOffset(c); err != nil {
			return err
		}
	}
	return nil
}

func (p *point) updatePosition(c Converter) error {
	line, col, err := c.ToPosition(p.Offset)
	if err != nil {
		return err
	}
	p.Line = line
	p.Column = col
	return nil
}

func (p *point) updateOffset(c Converter) error {
	offset, err := c.ToOffset(p.Line, p.Column)
	if err != nil {
		return err
	}
	p.Offset = offset
	return nil
}
<commit_msg>internal/lsp: add compare functions for spans<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package span

import (
	"encoding/json"
	"fmt"
	"strings"
)

// Span represents a source code range in standardized form.
type Span struct {
	v span
}

// Point represents a single point within a file.
// In general this should only be used as part of a Span, as on its own it
// does not carry enough information.
type Point struct {
	v point
}

type span struct {
	URI   URI   `json:"uri"`
	Start point `json:"start"`
	End   point `json:"end"`
}

type point struct {
	Line   int `json:"line"`
	Column int `json:"column"`
	Offset int `json:"offset"`
}

var invalidPoint = Point{v: point{Line: 0, Column: 0, Offset: -1}}

// Converter is the interface to an object that can convert between line:column
// and offset forms for a single file.
type Converter interface {
	//ToPosition converts from an offset to a line:column pair.
	ToPosition(offset int) (int, int, error)
	//ToOffset converts from a line:column pair to an offset.
	ToOffset(line, col int) (int, error)
}

func New(uri URI, start Point, end Point) Span {
	s := Span{v: span{URI: uri, Start: start.v, End: end.v}}
	s.v.clean()
	return s
}

func NewPoint(line, col, offset int) Point {
	p := Point{v: point{Line: line, Column: col, Offset: offset}}
	p.v.clean()
	return p
}

func Compare(a, b Span) int {
	if r := strings.Compare(string(a.v.URI), string(b.v.URI)); r != 0 {
		return r
	}
	if r := comparePoint(a.v.Start, b.v.Start); r != 0 {
		return r
	}
	return comparePoint(a.v.End, b.v.End)
}

func ComparePoint(a, b Point) int {
	return comparePoint(a.v, b.v)
}

func comparePoint(a, b point) int {
	if !a.hasPosition() {
		if a.Offset < b.Offset {
			return -1
		}
		if a.Offset > b.Offset {
			return 1
		}
		return 0
	}
	if a.Line < b.Line {
		return -1
	}
	if a.Line > b.Line {
		return 1
	}
	if a.Column < b.Column {
		return -1
	}
	if a.Column > b.Column {
		return 1
	}
	return 0
}

func (s Span) HasPosition() bool             { return s.v.Start.hasPosition() }
func (s Span) HasOffset() bool               { return s.v.Start.hasOffset() }
func (s Span) IsValid() bool                 { return s.v.Start.isValid() }
func (s Span) IsPoint() bool                 { return s.v.Start == s.v.End }
func (s Span) URI() URI                      { return s.v.URI }
func (s Span) Start() Point                  { return Point{s.v.Start} }
func (s Span) End() Point                    { return Point{s.v.End} }
func (s *Span) MarshalJSON() ([]byte, error) { return json.Marshal(&s.v) }
func (s *Span) UnmarshalJSON(b []byte) error { return json.Unmarshal(b, &s.v) }

func (p Point) HasPosition() bool             { return p.v.hasPosition() }
func (p Point) HasOffset() bool               { return p.v.hasOffset() }
func (p Point) IsValid() bool                 { return p.v.isValid() }
func (p *Point) MarshalJSON() ([]byte, error) { return json.Marshal(&p.v) }
func (p *Point) UnmarshalJSON(b []byte) error { return json.Unmarshal(b, &p.v) }
func (p Point) Line() int {
	if !p.v.hasPosition() {
		panic(fmt.Errorf("position not set in %v", p.v))
	}
	return p.v.Line
}
func (p Point) Column() int {
	if !p.v.hasPosition() {
		panic(fmt.Errorf("position not set in %v", p.v))
	}
	return p.v.Column
}
func (p Point) Offset() int {
	if !p.v.hasOffset() {
		panic(fmt.Errorf("offset not set in %v", p.v))
	}
	return p.v.Offset
}

func (p point) hasPosition() bool { return p.Line > 0 }
func (p point) hasOffset() bool   { return p.Offset >= 0 }
func (p point) isValid() bool     { return p.hasPosition() || p.hasOffset() }
func (p point) isZero() bool {
	return (p.Line == 1 && p.Column == 1) || (!p.hasPosition() && p.Offset == 0)
}

func (s *span) clean() {
	//this presumes the points are already clean
	if !s.End.isValid() || (s.End == point{}) {
		s.End = s.Start
	}
}

func (p *point) clean() {
	if p.Line < 0 {
		p.Line = 0
	}
	if p.Column <= 0 {
		if p.Line > 0 {
			p.Column = 1
		} else {
			p.Column = 0
		}
	}
	if p.Offset == 0 && (p.Line > 1 || p.Column > 1) {
		p.Offset = -1
	}
}

// Format implements fmt.Formatter to print the Location in a standard form.
// The format produced is one that can be read back in using Parse.
func (s Span) Format(f fmt.State, c rune) {
	fullForm := f.Flag('+')
	preferOffset := f.Flag('#')
	// we should always have a uri, simplify if it is file format
	//TODO: make sure the end of the uri is unambiguous
	uri := string(s.v.URI)
	if !fullForm {
		if filename, err := s.v.URI.Filename(); err == nil {
			uri = filename
		}
	}
	fmt.Fprint(f, uri)
	if !s.IsValid() || (!fullForm && s.v.Start.isZero() && s.v.End.isZero()) {
		return
	}
	// see which bits of start to write
	printOffset := s.HasOffset() && (fullForm || preferOffset || !s.HasPosition())
	printLine := s.HasPosition() && (fullForm || !printOffset)
	printColumn := printLine && (fullForm || (s.v.Start.Column > 1 || s.v.End.Column > 1))
	fmt.Fprint(f, ":")
	if printLine {
		fmt.Fprintf(f, "%d", s.v.Start.Line)
	}
	if printColumn {
		fmt.Fprintf(f, ":%d", s.v.Start.Column)
	}
	if printOffset {
		fmt.Fprintf(f, "#%d", s.v.Start.Offset)
	}
	// start is written, do we need end?
	if s.IsPoint() {
		return
	}
	// we don't print the line if it did not change
	printLine = fullForm || (printLine && s.v.End.Line > s.v.Start.Line)
	fmt.Fprint(f, "-")
	if printLine {
		fmt.Fprintf(f, "%d", s.v.End.Line)
	}
	if printColumn {
		if printLine {
			fmt.Fprint(f, ":")
		}
		fmt.Fprintf(f, "%d", s.v.End.Column)
	}
	if printOffset {
		fmt.Fprintf(f, "#%d", s.v.End.Offset)
	}
}

func (s Span) WithPosition(c Converter) (Span, error) {
	if err := s.update(c, true, false); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s Span) WithOffset(c Converter) (Span, error) {
	if err := s.update(c, false, true); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s Span) WithAll(c Converter) (Span, error) {
	if err := s.update(c, true, true); err != nil {
		return Span{}, err
	}
	return s, nil
}

func (s *Span) update(c Converter, withPos, withOffset bool) error {
	if !s.IsValid() {
		return fmt.Errorf("cannot add information to an invalid span")
	}
	if withPos && !s.HasPosition() {
		if err := s.v.Start.updatePosition(c); err != nil {
			return err
		}
		if s.v.End.Offset == s.v.Start.Offset {
			s.v.End = s.v.Start
		} else if err := s.v.End.updatePosition(c); err != nil {
			return err
		}
	}
	if withOffset && !s.HasOffset() {
		if err := s.v.Start.updateOffset(c); err != nil {
			return err
		}
		if s.v.End.Line == s.v.Start.Line && s.v.End.Column == s.v.Start.Column {
			s.v.End.Offset = s.v.Start.Offset
		} else if err := s.v.End.updateOffset(c); err != nil {
			return err
		}
	}
	return nil
}

func (p *point) updatePosition(c Converter) error {
	line, col, err := c.ToPosition(p.Offset)
	if err != nil {
		return err
	}
	p.Line = line
	p.Column = col
	return nil
}

func (p *point) updateOffset(c Converter) error {
	offset, err := c.ToOffset(p.Line, p.Column)
	if err != nil {
		return err
	}
	p.Offset = offset
	return nil
}
<|endoftext|>


<commit_before>package goriacache

import "testing"

func TestGoriaCache(t *testing.T) {
	cache, err := New("sample", 256, true)

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i := 0; i < 256; i++ {
		cache.Put(i, i)
	}

	if cache.IsStatsEnabled() {
		if cache.Stats().Items != 256 {
			t.Fatalf("Wrong Items stats %v", cache.Stats().Items)
		}
	}

	v, ok := cache.Get(255)
	if v != 255 && !ok {
		t.Fatalf("255 should be in the cache with key 255")
	}

	keyValuesSet := map[interface{}]interface{}{
		10: 10,
		20: 20,
		30: 30,
		40: 40,
	}

	returnedKeyValuesSet := make(map[interface{}]interface{})
	returnedKeyValuesSet = cache.GetAll(keyValuesSet)

	for k, v := range keyValuesSet {
		if returnedKeyValuesSet[k] != v {
			t.Fatalf("key %v should have value %v", k, v)
		}
	}

	keyValuesSet = map[interface{}]interface{}{
		10: 11,
		20: 22,
		30: 33,
		40: 44,
	}

	cache.PutAll(keyValuesSet)

	returnedKeyValuesSet = cache.GetAll(keyValuesSet)

	for k, v := range keyValuesSet {
		if returnedKeyValuesSet[k] != v {
			t.Fatalf("key %v should have value %v", k, v)
		}
	}
}
<commit_msg>Improved Tests<commit_after>package goriacache

import "testing"

func TestGoriaCache(t *testing.T) {
	cache, err := New("sample", 256, true)

	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i := 0; i < 256; i++ {
		cache.Put(i, i)
	}

	if cache.IsStatsEnabled() {
		if cache.Stats().Items != 256 {
			t.Fatalf("Wrong Items stats %v", cache.Stats().Items)
		}
	}

	v, ok := cache.Get(255)
	if v != 255 && !ok {
		t.Fatalf("255 should be in the cache with key 255")
	}

	keyValuesSet := map[interface{}]interface{}{
		10: 10,
		20: 20,
		30: 30,
		40: 40,
	}

	returnedKeyValuesSet := make(map[interface{}]interface{})
	returnedKeyValuesSet = cache.GetAll(keyValuesSet)

	for k, v := range keyValuesSet {
		if returnedKeyValuesSet[k] != v {
			t.Fatalf("key %v should have value %v", k, v)
		}
	}

	keyValuesSet = map[interface{}]interface{}{
		10: 11,
		20: 22,
		30: 33,
		40: 44,
	}

	cache.PutAll(keyValuesSet)

	returnedKeyValuesSet = cache.GetAll(keyValuesSet)

	for k, v := range keyValuesSet {
		if returnedKeyValuesSet[k] != v {
			t.Fatalf("key %v should have value %v", k, v)
		}
	}

	ok = cache.Replace(40, 44, 40)
	if !ok {
		t.Fatalf("err: %v", err)
	}

	v, ok = cache.Get(40)
	if !ok && v != 40 {
		t.Fatalf("key %v should have value %v", 40, 40)
	}
}
<|endoftext|>


<commit_before>package uploader

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"

	"github.com/dustin/go-humanize"
	"github.com/itchio/go-itchio"
	"github.com/itchio/wharf/counter"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/splitfunc"
)

var seed = 0

// ResumableUpload keeps track of an upload and reports back on its progress
type ResumableUpload struct {
	c *itchio.Client

	TotalBytes    int64
	UploadedBytes int64
	OnProgress    func()

	// resumable URL as per GCS
	uploadURL string

	// where data is written so we can update counts
	writeCounter io.Writer

	// need to flush to squeeze all the data out
	bufferedWriter *bufio.Writer

	// need to close so reader end of pipe gets EOF
	pipeWriter io.Closer

	id       int
	consumer *pwr.StateConsumer
}

// Close flushes all intermediary buffers and closes the connection
func (ru *ResumableUpload) Close() error {
	var err error

	ru.Debugf("flushing buffered writer, %d written", ru.TotalBytes)

	err = ru.bufferedWriter.Flush()
	if err != nil {
		return err
	}

	ru.Debugf("closing pipe writer")

	err = ru.pipeWriter.Close()
	if err != nil {
		return err
	}

	ru.Debugf("closed pipe writer")
	ru.Debugf("everything closed! uploadedbytes = %d, totalbytes = %d", ru.UploadedBytes, ru.TotalBytes)

	return nil
}

// Write is our implementation of io.Writer
func (ru *ResumableUpload) Write(p []byte) (int, error) {
	return ru.writeCounter.Write(p)
}

func NewResumableUpload(uploadURL string, done chan bool, errs chan error, consumer *pwr.StateConsumer) (*ResumableUpload, error) {
	ru := &ResumableUpload{}
	ru.uploadURL = uploadURL
	ru.id = seed
	seed++
	ru.consumer = consumer
	ru.c = itchio.ClientWithKey("x")

	pipeR, pipeW := io.Pipe()

	ru.pipeWriter = pipeW

	// TODO: make configurable?
	const bufferSize = 32 * 1024 * 1024

	bufferedWriter := bufio.NewWriterSize(pipeW, bufferSize)
	ru.bufferedWriter = bufferedWriter

	onWrite := func(count int64) {
		ru.Debugf("onwrite %d", count)
		ru.TotalBytes = count
		if ru.OnProgress != nil {
			ru.OnProgress()
		}
	}
	ru.writeCounter = counter.NewWriterCallback(onWrite, bufferedWriter)

	go ru.uploadChunks(pipeR, done, errs)

	return ru, nil
}

func (ru *ResumableUpload) Debugf(f string, args ...interface{}) {
	ru.consumer.Debugf("[upload %d] %s", ru.id, fmt.Sprintf(f, args...))
}

const minBlockSize = 256 * 1024 // 256KB

func (ru *ResumableUpload) uploadChunks(reader io.Reader, done chan bool, errs chan error) {
	var offset int64 = 0

	sendBytes := func(buf []byte, isEnd bool) error {
		buflen := int64(len(buf))
		ru.Debugf("received %d bytes", buflen)

		body := bytes.NewReader(buf)
		countingReader := counter.NewReaderCallback(func(count int64) {
			ru.UploadedBytes = offset + count
			if ru.OnProgress != nil {
				ru.OnProgress()
			}
		}, body)

		req, err := http.NewRequest("PUT", ru.uploadURL, countingReader)
		if err != nil {
			return err
		}

		start := offset
		end := start + buflen - 1
		contentRange := fmt.Sprintf("bytes %d-%d/*", offset, end)

		if isEnd {
			contentRange = fmt.Sprintf("bytes %d-%d/%d", offset, end, offset+buflen)
		}

		req.Header.Set("content-range", contentRange)

		res, err := ru.c.Do(req)
		if err != nil {
			return err
		}

		if res.StatusCode != 200 && res.StatusCode != 308 {
			ru.Debugf("uh oh, got HTTP %s", res.Status)
			resb, _ := ioutil.ReadAll(res.Body)
			ru.Debugf("server said %s", string(resb))
			return fmt.Errorf("HTTP %d while uploading", res.StatusCode)
		}

		offset += buflen
		ru.Debugf("%s uploaded, at %s", humanize.Bytes(uint64(offset)), res.Status)
		return nil
	}

	splitSize := 4 * minBlockSize

	s := bufio.NewScanner(reader)
	s.Buffer(make([]byte, splitSize), 0)
	s.Split(splitfunc.New(splitSize))

	buf1 := make([]byte, 0, splitSize)
	buf2 := make([]byte, 0, splitSize)

	for s.Scan() {
		buf2 = append(buf2[:0], buf1...)
		buf1 = append(buf1[:0], s.Bytes()...)

		if len(buf2) > 0 {
			ru.Debugf("sending %d block", len(buf2))
			err := sendBytes(buf2, false)
			if err != nil {
				errs <- err
				return
			}
		}
	}

	err := s.Err()
	if err != nil {
		ru.Debugf("scanner error :(")
		errs <- err
		return
	}

	ru.Debugf("sending last block, %d bytes", len(buf1))
	err = sendBytes(buf1, true)
	if err != nil {
		errs <- err
		return
	}

	done <- true
	ru.Debugf("done sent!")
}
<commit_msg>Better resumable upload performance<commit_after>package uploader

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"

	"github.com/dustin/go-humanize"
	"github.com/itchio/go-itchio"
	"github.com/itchio/wharf/counter"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/splitfunc"
)

var seed = 0

// ResumableUpload keeps track of an upload and reports back on its progress
type ResumableUpload struct {
	c *itchio.Client

	TotalBytes    int64
	UploadedBytes int64
	OnProgress    func()

	// resumable URL as per GCS
	uploadURL string

	// where data is written so we can update counts
	writeCounter io.Writer

	// need to flush to squeeze all the data out
	bufferedWriter *bufio.Writer

	// need to close so reader end of pipe gets EOF
	pipeWriter io.Closer

	id       int
	consumer *pwr.StateConsumer
}

// Close flushes all intermediary buffers and closes the connection
func (ru *ResumableUpload) Close() error {
	var err error

	ru.Debugf("flushing buffered writer, %d written", ru.TotalBytes)

	err = ru.bufferedWriter.Flush()
	if err != nil {
		return err
	}

	ru.Debugf("closing pipe writer")

	err = ru.pipeWriter.Close()
	if err != nil {
		return err
	}

	ru.Debugf("closed pipe writer")
	ru.Debugf("everything closed! uploadedbytes = %d, totalbytes = %d", ru.UploadedBytes, ru.TotalBytes)

	return nil
}

// Write is our implementation of io.Writer
func (ru *ResumableUpload) Write(p []byte) (int, error) {
	return ru.writeCounter.Write(p)
}

func NewResumableUpload(uploadURL string, done chan bool, errs chan error, consumer *pwr.StateConsumer) (*ResumableUpload, error) {
	ru := &ResumableUpload{}
	ru.uploadURL = uploadURL
	ru.id = seed
	seed++
	ru.consumer = consumer
	ru.c = itchio.ClientWithKey("x")

	pipeR, pipeW := io.Pipe()

	ru.pipeWriter = pipeW

	// TODO: make configurable?
	const bufferSize = 32 * 1024 * 1024

	bufferedWriter := bufio.NewWriterSize(pipeW, bufferSize)
	ru.bufferedWriter = bufferedWriter

	onWrite := func(count int64) {
		// ru.Debugf("onwrite %d", count)
		ru.TotalBytes = count
		if ru.OnProgress != nil {
			ru.OnProgress()
		}
	}
	ru.writeCounter = counter.NewWriterCallback(onWrite, bufferedWriter)

	go ru.uploadChunks(pipeR, done, errs)

	return ru, nil
}

func (ru *ResumableUpload) Debugf(f string, args ...interface{}) {
	ru.consumer.Debugf("[upload %d] %s", ru.id, fmt.Sprintf(f, args...))
}

const minChunkSize = 256 * 1024 // 256KB
const maxChunkGroup = 64
const maxSendBuf = maxChunkGroup * minChunkSize // 16MB

type blockItem struct {
	buf    []byte
	isLast bool
}

func (ru *ResumableUpload) uploadChunks(reader io.Reader, done chan bool, errs chan error) {
	var offset int64 = 0

	sendBuf := make([]byte, 0, maxSendBuf)
	reqBlocks := make(chan blockItem, maxChunkGroup)

	canceller := make(chan bool)

	sendBytes := func(buf []byte, isLast bool) {
		reqBlocks <- blockItem{buf: append([]byte{}, buf...), isLast: isLast}
	}

	doSendBytes := func(buf []byte, isLast bool) error {
		buflen := int64(len(sendBuf))
		ru.Debugf("uploading chunk of %d bytes", buflen)

		body := bytes.NewReader(buf)
		countingReader := counter.NewReaderCallback(func(count int64) {
			ru.UploadedBytes = offset + count
			if ru.OnProgress != nil {
				ru.OnProgress()
			}
		}, body)

		req, err := http.NewRequest("PUT", ru.uploadURL, countingReader)
		if err != nil {
			return err
		}

		start := offset
		end := start + buflen - 1
		contentRange := fmt.Sprintf("bytes %d-%d/*", offset, end)

		if isLast {
			contentRange = fmt.Sprintf("bytes %d-%d/%d", offset, end, offset+buflen)
		}

		req.Header.Set("content-range", contentRange)

		res, err := ru.c.Do(req)
		if err != nil {
			return err
		}

		if res.StatusCode != 200 && res.StatusCode != 308 {
			ru.Debugf("uh oh, got HTTP %s", res.Status)
			resb, _ := ioutil.ReadAll(res.Body)
			ru.Debugf("server said %s", string(resb))
			return fmt.Errorf("HTTP %d while uploading", res.StatusCode)
		}

		offset += buflen
		ru.Debugf("%s uploaded, at %s", humanize.Bytes(uint64(offset)), res.Status)
		return nil
	}

	s := bufio.NewScanner(reader)
	s.Buffer(make([]byte, minChunkSize), 0)
	s.Split(splitfunc.New(minChunkSize))

	// we need two buffers to know when we're at EOF,
	// for sizes that are an exact multiple of minChunkSize
	buf1 := make([]byte, 0, minChunkSize)
	buf2 := make([]byte, 0, minChunkSize)

	subDone := make(chan bool)
	subErrs := make(chan error)

	ru.Debugf("kicking off sender")

	go func() {
		isLast := false

		for !isLast {
			sendBuf = sendBuf[:0]

			for len(sendBuf) < maxSendBuf && !isLast {
				var item blockItem
				if len(sendBuf) == 0 {
					ru.Debugf("sender blocking receive")
					select {
					case item = <-reqBlocks:
						// cool
					case <-canceller:
						ru.Debugf("send cancelled")
						return
					}
				} else {
					ru.Debugf("sender non-blocking receive")
					select {
					case item = <-reqBlocks:
						// cool
					case <-canceller:
						ru.Debugf("send cancelled")
						return
					default:
						ru.Debugf("sent faster than scanned, uploading smaller chunk")
						break
					}
				}

				if item.isLast {
					isLast = true
				}

				sendBuf = append(sendBuf, item.buf...)
			}

			if len(sendBuf) > 0 {
				err := doSendBytes(sendBuf, isLast)
				if err != nil {
					ru.Debugf("send error, bailing out")
					subErrs <- err
					return
				}
			}
		}

		subDone <- true
		ru.Debugf("sender done")
	}()

	// break patch into chunks of minChunkSize, signal last block
	go func() {
		for s.Scan() {
			buf2 = append(buf2[:0], buf1...)
			buf1 = append(buf1[:0], s.Bytes()...)

			if len(buf2) > 0 {
				sendBytes(buf2, false)
			}

			select {
			case <-canceller:
				ru.Debugf("scan cancelled")
				return
			default:
				// okay cool let's go c'mon
			}
		}

		err := s.Err()
		if err != nil {
			ru.Debugf("scanner error :(")
			subErrs <- err
			return
		}

		sendBytes(buf1, true)

		subDone <- true
		ru.Debugf("scanner done")
	}()

	for i := 0; i < 2; i++ {
		select {
		case <-subDone:
			// woo!
		case err := <-subErrs:
			ru.Debugf("got sub error: %s, bailing", err.Error())
			close(canceller)
			errs <- err
			return
		}
	}

	done <- true
	ru.Debugf("done sent!")
}
<|endoftext|>


<commit_before>// Forex data API

// http://finance.yahoo.com/webservice/v1/symbols/CNY=X/quote?format=json

package forex

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"
)

// Forex data API URL
const (
	DATAURL = "http://finance.yahoo.com/webservice/v1/symbols/"
)

// Quote forex info
type Quote struct {
	Price  float64
	Symbol string
	Error  error
}

// CommunicateFX sends the latest FX quote to the supplied channel
func CommunicateFX(symbol string, fxChan chan<- Quote, doneChan <-chan bool) Quote {
	// Initial quote to return
	quote := getQuote(symbol)

	// Run read loop in new goroutine
	go runLoop(symbol, fxChan, doneChan)

	return quote
}

// HTTP read loop
func runLoop(symbol string, fxChan chan<- Quote, doneChan <-chan bool) {
	ticker := time.NewTicker(15 * time.Second)

	for {
		select {
		case <-doneChan:
			ticker.Stop()
			return
		case <-ticker.C:
			fxChan <- getQuote(symbol)
		}
	}
}

// Returns quote for requested instrument
func getQuote(symbol string) Quote {
	tmp := struct {
		List struct {
			Resources []struct {
				Resource struct {
					Fields struct {
						Price float64 `json:"price,string"`
					} `json:"fields"`
				} `json:"resource"`
			} `json:"resources"`
		} `json:"list"`
	}{}

	url := fmt.Sprintf("%s=x/quote?format=json", symbol)

	data, err := get(url)
	if err != nil {
		return Quote{Error: fmt.Errorf("Forex error %s", err)}
	}

	if err = json.Unmarshal(data, &tmp); err != nil {
		return Quote{Error: fmt.Errorf("Forex error %s", err)}
	}

	price := tmp.List.Resources[0].Resource.Fields.Price
	if price < .000001 {
		return Quote{Error: fmt.Errorf("Forex zero price error")}
	}

	return Quote{
		Price:  price,
		Symbol: symbol,
		Error:  nil,
	}
}

// unauthenticated GET
func get(url string) ([]byte, error) {
	resp, err := http.Get(DATAURL + url)
	if err != nil {
		return []byte{}, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}
<commit_msg>clean up code a bit<commit_after>// Forex data API
// Currently using yahoo finance
// http://finance.yahoo.com/webservice/v1/symbols/CNY=X/quote?format=json

package forex

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"
)

// Forex data API URL
const DATAURL = "http://finance.yahoo.com/webservice/v1/symbols/"

// Quote contains forex quote information
type Quote struct {
	Price  float64
	Symbol string
	Error  error
}

// CommunicateFX sends the latest FX quote to the supplied channel
func CommunicateFX(symbol string, fxChan chan<- Quote, doneChan <-chan bool) Quote {
	// Initial quote to return
	quote := getQuote(symbol)

	// Run read loop in new goroutine
	go runLoop(symbol, fxChan, doneChan)

	return quote
}

// HTTP read loop
func runLoop(symbol string, fxChan chan<- Quote, doneChan <-chan bool) {
	ticker := time.NewTicker(15 * time.Second)

	for {
		select {
		case <-doneChan:
			ticker.Stop()
			return
		case <-ticker.C:
			fxChan <- getQuote(symbol)
		}
	}
}

// Returns quote for requested currency
func getQuote(symbol string) Quote {
	// Get data
	url := fmt.Sprintf("%s%s=x/quote?format=json", DATAURL, symbol)
	data, err := get(url)
	if err != nil {
		return Quote{Error: fmt.Errorf("Forex error %s", err)}
	}

	// Unmarshal
	response := struct {
		List struct {
			Resources []struct {
				Resource struct {
					Fields struct {
						Price float64 `json:"price,string"`
					} `json:"fields"`
				} `json:"resource"`
			} `json:"resources"`
		} `json:"list"`
	}{}
	if err = json.Unmarshal(data, &response); err != nil {
		return Quote{Error: fmt.Errorf("Forex error %s", err)}
	}

	// Pull out price
	price := response.List.Resources[0].Resource.Fields.Price
	if price < .000001 {
		return Quote{Error: fmt.Errorf("Forex zero price error")}
	}

	return Quote{
		Price:  price,
		Symbol: symbol,
		Error:  nil,
	}
}

// Unauthenticated GET
func get(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return []byte{}, err
	}
	if resp.StatusCode != 200 {
		return []byte{}, fmt.Errorf(resp.Status)
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 format

import (
	"bytes"
	"fmt"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"sync"
	"time"
)

type jsonReaderState int

const (
	jsonReadArrayEnd    = jsonReaderState(iota)
	jsonReadObjectEnd   = jsonReaderState(iota)
	jsonReadObject      = jsonReaderState(iota)
	jsonReadKey         = jsonReaderState(iota)
	jsonReadValue       = jsonReaderState(iota)
	jsonReadArray       = jsonReaderState(iota)
	jsonReadArrayAppend = jsonReaderState(iota)
)

// JSON formatter plugin
// JSON is a formatter that passes a message encapsulated as JSON in the form
// {"message":"..."}. The actual message is formatted by a nested formatter and
// HTML escaped.
// Configuration example
//
//  - "stream.Broadcast":
//    Formatter: "format.JSON"
//    JSONStartState: "findKey"
//    JSONDirectives:
//	    - 'findKey :":  key     ::'
//	    - 'findKey :}:          : pop  : end'
//	    - 'key     :":  findVal :      : key'
//	    - 'findVal :\:: value   ::'
//
// JSONStartState defines the initial parser state when parsing a message.
// By default this is set to "" which will fall back to the first state used in
// the JSONDirectives array.
//
// JSONTimestampRead defines the go timestamp format expected from fields that
// are parsed as "dat". By default this is set to "20060102150405".
//
// JSONTimestampWrite defines the go timestamp format that "dat" fields will be
// converted to. By default this is set to "2006-01-02 15:04:05 MST".
//
// JSONDirectives defines an array of parser directives.
// This setting is mandatory and has no default value.
// Each string must be of the following format: "State:Token:NextState:Flags:Function".
// Spaces will be stripped from all fields but Token. If a fields requires a
// colon it has to be escaped with a backslash. Other escape characters
// supported are \n, \r and \t.
//
// Flags (JSONDirectives) can be a comma separated set of the following flags.
//  * continue -> Prepend the token to the next match.
//  * append   -> Append the token to the current match and continue reading.
//  * include  -> Append the token to the current match.
//  * push     -> Push the current state to the stack.
//  * pop      -> Pop the stack and use the returned state if possible.
//
// Function (JSONDirectives) can hold one of the following names.
//  * key     -> Write the current match as a key.
//  * val     -> Write the current match as a value without quotes.
//  * esc     -> Write the current match as a escaped string value.
//  * dat     -> Write the current match as a timestamp value.
//  * arr     -> Start a new array.
//  * obj     -> Start a new object.
//  * end     -> Close an array or object.
//  * arr+val -> arr followed by val.
//  * arr+esc -> arr followed by esc.
//  * arr+dat -> arr followed by dat.
//  * val+end -> val followed by end.
//  * esc+end -> esc followed by end.
//  * dat+end -> dat followed by end.
//
// Rules for storage (JSONDirectives): if a value is written without a previous key write, a key will be auto
// generated from the current parser state name. This does not happen when
// inside an array.
// If key is written without a previous value write, a null value will be
// written. This does not happen after an object has been started.
// A key write inside an array will cause the array to be closed. If the array
// is nested, all arrays will be closed.
type JSON struct {
	message   *bytes.Buffer
	parser    shared.TransitionParser
	state     jsonReaderState
	stack     []jsonReaderState
	parseLock *sync.Mutex
	initState string
	timeRead  string
	timeWrite string
}

func init() {
	shared.TypeRegistry.Register(JSON{})
}

// Configure initializes this formatter with values from a plugin config.
func (format *JSON) Configure(conf core.PluginConfig) error {
	format.parser = shared.NewTransitionParser()
	format.state = jsonReadObject
	format.initState = conf.GetString("JSONStartState", "")
	format.timeRead = conf.GetString("JSONTimestampRead", "20060102150405")
	format.timeWrite = conf.GetString("JSONTimestampWrite", "2006-01-02 15:04:05 MST")
	format.parseLock = new(sync.Mutex)

	if !conf.HasValue("JSONDirectives") {
		Log.Warning.Print("JSON formatter has no JSONDirectives setting")
		return nil // ### return, no directives ###
	}

	directiveStrings := conf.GetStringArray("JSONDirectives", []string{})
	if len(directiveStrings) == 0 {
		Log.Warning.Print("JSON formatter has no directives")
		return nil // ### return, no directives ###
	}

	// Parse directives

	parserFunctions := make(map[string]shared.ParsedFunc)
	parserFunctions["key"] = format.readKey
	parserFunctions["val"] = format.readValue
	parserFunctions["esc"] = format.readEscaped
	parserFunctions["dat"] = format.readDate
	parserFunctions["arr"] = format.readArray
	parserFunctions["obj"] = format.readObject
	parserFunctions["end"] = format.readEnd
	parserFunctions["arr+val"] = format.readArrayValue
	parserFunctions["arr+esc"] = format.readArrayEscaped
	parserFunctions["val+end"] = format.readValueEnd
	parserFunctions["esc+end"] = format.readEscapedEnd
	parserFunctions["dat+end"] = format.readDateEnd

	directives := []shared.TransitionDirective{}
	for _, dirString := range directiveStrings {
		directive, err := shared.ParseTransitionDirective(dirString, parserFunctions)
		if err != nil {
			return fmt.Errorf("%s: %s", err.Error(), dirString) // ### return, malformed directive ###
		}
		if format.initState == "" {
			format.initState = directive.State
		}
		directives = append(directives, directive)
	}

	format.parser.AddDirectives(directives)
	return nil
}

func (format *JSON) writeKey(key []byte) {
	// Make sure we are not in an array anymore
	for format.state == jsonReadArray || format.state == jsonReadArrayAppend {
		format.readEnd(nil, 0)
	}

	// If no value was written, write null
	if format.state > jsonReadKey {
		format.message.WriteString("null")
	}

	// Prepend a comma except after an object has started
	if format.state != jsonReadObject {
		format.message.WriteByte(',')
	}

	format.message.WriteByte('"')
	format.message.Write(key)
	format.message.WriteString(`":`)
}

func (format *JSON) readKey(data []byte, state shared.ParserStateID) {
	format.writeKey(data)
	format.state = jsonReadValue
}

func (format *JSON) readValue(data []byte, state shared.ParserStateID) {
	switch format.state {
	default:
		format.writeKey([]byte(format.parser.GetStateName(state)))
		fallthrough

	case jsonReadValue:
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadKey

	case jsonReadArray:
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadArrayAppend

	case jsonReadArrayAppend:
		format.message.WriteByte(',')
		format.message.Write(bytes.TrimSpace(data))
	}
}

func (format *JSON) readEscaped(data []byte, state shared.ParserStateID) {
	switch format.state {
	default:
		format.writeKey([]byte(format.parser.GetStateName(state)))
		fallthrough

	case jsonReadValue:
		format.message.WriteByte('"')
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadKey

	case jsonReadArray:
		format.message.WriteByte('"')
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadArrayAppend

	case jsonReadArrayAppend:
		format.message.WriteString(`,"`)
		format.message.Write(bytes.TrimSpace(data))
	}
	format.message.WriteByte('"')
}

func (format *JSON) readDate(data []byte, state shared.ParserStateID) {
	date, _ := time.Parse(format.timeRead, string(bytes.TrimSpace(data)))
	formattedDate := date.Format(format.timeWrite)
	format.readEscaped([]byte(formattedDate), state)
}

func (format *JSON) readValueEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readValue(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readEscapedEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readEscaped(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readDateEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readDate(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readArrayValue(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readValue(data, state)
}

func (format *JSON) readArrayEscaped(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readEscaped(data, state)
}

func (format *JSON) readArrayDate(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readDate(data, state)
}

func (format *JSON) readArray(data []byte, state shared.ParserStateID) {
	if format.state == jsonReadArrayAppend {
		format.message.WriteString(",[")
	} else {
		format.message.WriteByte('[')
	}
	format.stack = append(format.stack, format.state)
	format.state = jsonReadArray
}

func (format *JSON) readObject(data []byte, state shared.ParserStateID) {
	if format.state == jsonReadArrayAppend {
		format.message.WriteString(",{")
	} else {
		format.message.WriteByte('{')
	}
	format.stack = append(format.stack, format.state)
	format.state = jsonReadObject
}

func (format *JSON) readEnd(data []byte, state shared.ParserStateID) {
	stackSize := len(format.stack)

	if stackSize > 0 {
		switch format.state {
		case jsonReadArray, jsonReadArrayAppend:
			format.message.WriteByte(']')
		default:
			format.message.WriteByte('}')
		}
	}

	if stackSize > 1 {
		format.state = format.stack[stackSize-1]
		format.stack = format.stack[:stackSize-1] // Pop the stack
	} else {
		format.stack = format.stack[:0] // Clear the stack
		format.state = jsonReadValue
	}
}

// Format parses the incoming message and generates JSON from it.
// This function is mutex locked.
func (format *JSON) Format(msg core.Message) ([]byte, core.MessageStreamID) {
	// The internal state is not threadsafe so we need to lock here
	format.parseLock.Lock()
	defer format.parseLock.Unlock()

	format.message = bytes.NewBuffer(nil)
	format.state = jsonReadObject

	format.message.WriteString("{")
	remains, state := format.parser.Parse(msg.Data, format.initState)

	// Write remains as string value
	if remains != nil {
		format.readEscaped(remains, state)
	}

	// Close any open tags
	if format.message.Len() > 1 {
		for format.state == jsonReadArray || format.state == jsonReadArrayAppend || format.state == jsonReadObject {
			format.readEnd(nil, 0)
		}
	}

	format.message.WriteString("}\n")
	return bytes.TrimSpace(format.message.Bytes()), msg.StreamID
}
<commit_msg>support unix dates in the JSON formatter<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 format

import (
	"bytes"
	"fmt"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"strconv"
	"sync"
	"time"
)

type jsonReaderState int

const (
	jsonReadArrayEnd    = jsonReaderState(iota)
	jsonReadObjectEnd   = jsonReaderState(iota)
	jsonReadObject      = jsonReaderState(iota)
	jsonReadKey         = jsonReaderState(iota)
	jsonReadValue       = jsonReaderState(iota)
	jsonReadArray       = jsonReaderState(iota)
	jsonReadArrayAppend = jsonReaderState(iota)
)

// JSON formatter plugin
// JSON is a formatter that passes a message encapsulated as JSON in the form
// {"message":"..."}. The actual message is formatted by a nested formatter and
// HTML escaped.
// Configuration example
//
//  - "stream.Broadcast":
//    Formatter: "format.JSON"
//    JSONStartState: "findKey"
//    JSONDirectives:
//	    - 'findKey :":  key     ::'
//	    - 'findKey :}:          : pop  : end'
//	    - 'key     :":  findVal :      : key'
//	    - 'findVal :\:: value   ::'
//
// JSONStartState defines the initial parser state when parsing a message.
// By default this is set to "" which will fall back to the first state used in
// the JSONDirectives array.
//
// JSONTimestampRead defines the go timestamp format expected from fields that
// are parsed as "dat". By default this is set to "20060102150405".
//
// JSONTimestampWrite defines the go timestamp format that "dat" fields will be
// converted to. By default this is set to "2006-01-02 15:04:05 MST".
//
// JSONDirectives defines an array of parser directives.
// This setting is mandatory and has no default value.
// Each string must be of the following format: "State:Token:NextState:Flags:Function".
// Spaces will be stripped from all fields but Token. If a fields requires a
// colon it has to be escaped with a backslash. Other escape characters
// supported are \n, \r and \t.
//
// Flags (JSONDirectives) can be a comma separated set of the following flags.
//  * continue -> Prepend the token to the next match.
//  * append   -> Append the token to the current match and continue reading.
//  * include  -> Append the token to the current match.
//  * push     -> Push the current state to the stack.
//  * pop      -> Pop the stack and use the returned state if possible.
//
// Function (JSONDirectives) can hold one of the following names.
//  * key     -> Write the current match as a key.
//  * val     -> Write the current match as a value without quotes.
//  * esc     -> Write the current match as a escaped string value.
//  * dat     -> Write the current match as a timestamp value.
//  * arr     -> Start a new array.
//  * obj     -> Start a new object.
//  * end     -> Close an array or object.
//  * arr+val -> arr followed by val.
//  * arr+esc -> arr followed by esc.
//  * arr+dat -> arr followed by dat.
//  * val+end -> val followed by end.
//  * esc+end -> esc followed by end.
//  * dat+end -> dat followed by end.
//
// Rules for storage (JSONDirectives): if a value is written without a previous key write, a key will be auto
// generated from the current parser state name. This does not happen when
// inside an array.
// If key is written without a previous value write, a null value will be
// written. This does not happen after an object has been started.
// A key write inside an array will cause the array to be closed. If the array
// is nested, all arrays will be closed.
type JSON struct {
	message   *bytes.Buffer
	parser    shared.TransitionParser
	state     jsonReaderState
	stack     []jsonReaderState
	parseLock *sync.Mutex
	initState string
	timeRead  string
	timeWrite string
	timeParse func(string, string) (time.Time, error)
}

func init() {
	shared.TypeRegistry.Register(JSON{})
}

func parseUnix(layout, value string) (time.Time, error) {
	s, ns := int64(0), int64(0)
	switch layout {
	case "s", "sec":
		valueInt, err := strconv.ParseInt(value, 10, 64)
		if err != nil { return time.Time{}, err }
		s = valueInt
	case "ms", "msec":
		valueInt, err := strconv.ParseInt(value, 10, 64)
		if err != nil { return time.Time{}, err }
		ns = valueInt*int64(time.Millisecond)
	case "ns", "nsec":
		valueInt, err := strconv.ParseInt(value, 10, 64)
		if err != nil { return time.Time{}, err }
		ns = valueInt
	}
	return time.Unix(s, ns), nil
}

// Configure initializes this formatter with values from a plugin config.
func (format *JSON) Configure(conf core.PluginConfig) error {
	format.parser = shared.NewTransitionParser()
	format.state = jsonReadObject
	format.initState = conf.GetString("JSONStartState", "")
	format.timeRead = conf.GetString("JSONTimestampRead", "20060102150405")
	format.timeWrite = conf.GetString("JSONTimestampWrite", "2006-01-02 15:04:05 MST")
	format.parseLock = new(sync.Mutex)

	switch format.timeRead {
	case "s", "sec", "ms", "msec", "ns", "nsec":
		format.timeParse = parseUnix
	default:
		format.timeParse = time.Parse
	}

	if !conf.HasValue("JSONDirectives") {
		Log.Warning.Print("JSON formatter has no JSONDirectives setting")
		return nil // ### return, no directives ###
	}

	directiveStrings := conf.GetStringArray("JSONDirectives", []string{})
	if len(directiveStrings) == 0 {
		Log.Warning.Print("JSON formatter has no directives")
		return nil // ### return, no directives ###
	}

	// Parse directives

	parserFunctions := make(map[string]shared.ParsedFunc)
	parserFunctions["key"] = format.readKey
	parserFunctions["val"] = format.readValue
	parserFunctions["esc"] = format.readEscaped
	parserFunctions["dat"] = format.readDate
	parserFunctions["arr"] = format.readArray
	parserFunctions["obj"] = format.readObject
	parserFunctions["end"] = format.readEnd
	parserFunctions["arr+val"] = format.readArrayValue
	parserFunctions["arr+esc"] = format.readArrayEscaped
	parserFunctions["val+end"] = format.readValueEnd
	parserFunctions["esc+end"] = format.readEscapedEnd
	parserFunctions["dat+end"] = format.readDateEnd

	directives := []shared.TransitionDirective{}
	for _, dirString := range directiveStrings {
		directive, err := shared.ParseTransitionDirective(dirString, parserFunctions)
		if err != nil {
			return fmt.Errorf("%s: %s", err.Error(), dirString) // ### return, malformed directive ###
		}
		if format.initState == "" {
			format.initState = directive.State
		}
		directives = append(directives, directive)
	}

	format.parser.AddDirectives(directives)
	return nil
}

func (format *JSON) writeKey(key []byte) {
	// Make sure we are not in an array anymore
	for format.state == jsonReadArray || format.state == jsonReadArrayAppend {
		format.readEnd(nil, 0)
	}

	// If no value was written, write null
	if format.state > jsonReadKey {
		format.message.WriteString("null")
	}

	// Prepend a comma except after an object has started
	if format.state != jsonReadObject {
		format.message.WriteByte(',')
	}

	format.message.WriteByte('"')
	format.message.Write(key)
	format.message.WriteString(`":`)
}

func (format *JSON) readKey(data []byte, state shared.ParserStateID) {
	format.writeKey(data)
	format.state = jsonReadValue
}

func (format *JSON) readValue(data []byte, state shared.ParserStateID) {
	switch format.state {
	default:
		format.writeKey([]byte(format.parser.GetStateName(state)))
		fallthrough

	case jsonReadValue:
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadKey

	case jsonReadArray:
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadArrayAppend

	case jsonReadArrayAppend:
		format.message.WriteByte(',')
		format.message.Write(bytes.TrimSpace(data))
	}
}

func (format *JSON) readEscaped(data []byte, state shared.ParserStateID) {
	switch format.state {
	default:
		format.writeKey([]byte(format.parser.GetStateName(state)))
		fallthrough

	case jsonReadValue:
		format.message.WriteByte('"')
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadKey

	case jsonReadArray:
		format.message.WriteByte('"')
		format.message.Write(bytes.TrimSpace(data))
		format.state = jsonReadArrayAppend

	case jsonReadArrayAppend:
		format.message.WriteString(`,"`)
		format.message.Write(bytes.TrimSpace(data))
	}
	format.message.WriteByte('"')
}

func (format *JSON) readDate(data []byte, state shared.ParserStateID) {
	date, _ := format.timeParse(format.timeRead, string(bytes.TrimSpace(data)))
	formattedDate := date.Format(format.timeWrite)
	format.readEscaped([]byte(formattedDate), state)
}

func (format *JSON) readValueEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readValue(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readEscapedEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readEscaped(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readDateEnd(data []byte, state shared.ParserStateID) {
	formatState := format.state
	format.readDate(data, state)
	format.state = formatState
	format.readEnd(data, state)
}

func (format *JSON) readArrayValue(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readValue(data, state)
}

func (format *JSON) readArrayEscaped(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readEscaped(data, state)
}

func (format *JSON) readArrayDate(data []byte, state shared.ParserStateID) {
	format.readArray(data, state)
	format.readDate(data, state)
}

func (format *JSON) readArray(data []byte, state shared.ParserStateID) {
	if format.state == jsonReadArrayAppend {
		format.message.WriteString(",[")
	} else {
		format.message.WriteByte('[')
	}
	format.stack = append(format.stack, format.state)
	format.state = jsonReadArray
}

func (format *JSON) readObject(data []byte, state shared.ParserStateID) {
	if format.state == jsonReadArrayAppend {
		format.message.WriteString(",{")
	} else {
		format.message.WriteByte('{')
	}
	format.stack = append(format.stack, format.state)
	format.state = jsonReadObject
}

func (format *JSON) readEnd(data []byte, state shared.ParserStateID) {
	stackSize := len(format.stack)

	if stackSize > 0 {
		switch format.state {
		case jsonReadArray, jsonReadArrayAppend:
			format.message.WriteByte(']')
		default:
			format.message.WriteByte('}')
		}
	}

	if stackSize > 1 {
		format.state = format.stack[stackSize-1]
		format.stack = format.stack[:stackSize-1] // Pop the stack
	} else {
		format.stack = format.stack[:0] // Clear the stack
		format.state = jsonReadValue
	}
}

// Format parses the incoming message and generates JSON from it.
// This function is mutex locked.
func (format *JSON) Format(msg core.Message) ([]byte, core.MessageStreamID) {
	// The internal state is not threadsafe so we need to lock here
	format.parseLock.Lock()
	defer format.parseLock.Unlock()

	format.message = bytes.NewBuffer(nil)
	format.state = jsonReadObject

	format.message.WriteString("{")
	remains, state := format.parser.Parse(msg.Data, format.initState)

	// Write remains as string value
	if remains != nil {
		format.readEscaped(remains, state)
	}

	// Close any open tags
	if format.message.Len() > 1 {
		for format.state == jsonReadArray || format.state == jsonReadArrayAppend || format.state == jsonReadObject {
			format.readEnd(nil, 0)
		}
	}

	format.message.WriteString("}\n")
	return bytes.TrimSpace(format.message.Bytes()), msg.StreamID
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file provides an implementation of the FileSystem
// interface based on the contents of a .zip file.
//
// Assumptions:
//
// - The file paths stored in the zip file must use a slash ('/') as path
//   separator; and they must be relative (i.e., they must not start with
//   a '/' - this is usually the case if the file was created w/o special
//   options).
// - The zip file system treats the file paths found in the zip internally
//   like absolute paths w/o a leading '/'; i.e., the paths are considered
//   relative to the root of the file system.
// - All path arguments to file system methods must be absolute paths.

package main

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
)

// zipFI is the zip-file based implementation of FileInfo
type zipFI struct {
	name string    // directory-local name
	file *zip.File // nil for a directory
}

func (fi zipFI) Name() string {
	return fi.name
}

func (fi zipFI) Size() int64 {
	if f := fi.file; f != nil {
		return int64(f.UncompressedSize)
	}
	return 0 // directory
}

func (fi zipFI) Mtime_ns() int64 {
	if f := fi.file; f != nil {
		return f.Mtime_ns()
	}
	return 0 // directory has no modified time entry
}

func (fi zipFI) IsDirectory() bool {
	return fi.file == nil
}

func (fi zipFI) IsRegular() bool {
	return fi.file != nil
}

// zipFS is the zip-file based implementation of FileSystem
type zipFS struct {
	*zip.ReadCloser
	list zipList
}

func (fs *zipFS) Close() os.Error {
	fs.list = nil
	return fs.ReadCloser.Close()
}

func zipPath(name string) string {
	name = path.Clean(name)
	if !path.IsAbs(name) {
		panic(fmt.Sprintf("stat: not an absolute path: %s", name))
	}
	return name[1:] // strip leading '/'
}

func (fs *zipFS) stat(abspath string) (int, zipFI, os.Error) {
	i, exact := fs.list.lookup(abspath)
	if i < 0 {
		// abspath has leading '/' stripped - print it explicitly
		return -1, zipFI{}, fmt.Errorf("file not found: /%s", abspath)
	}
	_, name := path.Split(abspath)
	var file *zip.File
	if exact {
		file = fs.list[i] // exact match found - must be a file
	}
	return i, zipFI{name, file}, nil
}

func (fs *zipFS) Open(abspath string) (io.ReadCloser, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	if err != nil {
		return nil, err
	}
	if fi.IsDirectory() {
		return nil, fmt.Errorf("Open: %s is a directory", abspath)
	}
	return fi.file.Open()
}

func (fs *zipFS) Lstat(abspath string) (FileInfo, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	return fi, err
}

func (fs *zipFS) Stat(abspath string) (FileInfo, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	return fi, err
}

func (fs *zipFS) ReadDir(abspath string) ([]FileInfo, os.Error) {
	path := zipPath(abspath)
	i, fi, err := fs.stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDirectory() {
		return nil, fmt.Errorf("ReadDir: %s is not a directory", abspath)
	}

	var list []FileInfo
	dirname := path + "/"
	prevname := ""
	for _, e := range fs.list[i:] {
		if !strings.HasPrefix(e.Name, dirname) {
			break // not in the same directory anymore
		}
		name := e.Name[len(dirname):] // local name
		file := e
		if i := strings.IndexRune(name, '/'); i >= 0 {
			// We infer directories from files in subdirectories.
			// If we have x/y, return a directory entry for x.
			name = name[0:i] // keep local directory name only
			file = nil
		}
		// If we have x/y and x/z, don't return two directory entries for x.
		// TODO(gri): It should be possible to do this more efficiently
		// by determining the (fs.list) range of local directory entries
		// (via two binary searches).
		if name != prevname {
			list = append(list, zipFI{name, file})
			prevname = name
		}
	}

	return list, nil
}

func (fs *zipFS) ReadFile(abspath string) ([]byte, os.Error) {
	rc, err := fs.Open(abspath)
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(rc)
}

func NewZipFS(rc *zip.ReadCloser) FileSystem {
	list := make(zipList, len(rc.File))
	copy(list, rc.File) // sort a copy of rc.File
	sort.Sort(list)
	return &zipFS{rc, list}
}

type zipList []*zip.File

// zipList implements sort.Interface
func (z zipList) Len() int           { return len(z) }
func (z zipList) Less(i, j int) bool { return z[i].Name < z[j].Name }
func (z zipList) Swap(i, j int)      { z[i], z[j] = z[j], z[i] }

// lookup returns the smallest index of an entry with an exact match
// for name, or an inexact match starting with name/. If there is no
// such entry, the result is -1, false.
func (z zipList) lookup(name string) (index int, exact bool) {
	// look for exact match first (name comes before name/ in z)
	i := sort.Search(len(z), func(i int) bool {
		return name <= z[i].Name
	})
	if i < 0 {
		return -1, false
	}
	if z[i].Name == name {
		return i, true
	}

	// look for inexact match (must be in z[i:], if present)
	z = z[i:]
	name += "/"
	j := sort.Search(len(z), func(i int) bool {
		return name <= z[i].Name
	})
	if j < 0 {
		return -1, false
	}
	if strings.HasPrefix(z[j].Name, name) {
		return i + j, false
	}

	return -1, false
}
<commit_msg>godoc: fix bug in zip.go<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file provides an implementation of the FileSystem
// interface based on the contents of a .zip file.
//
// Assumptions:
//
// - The file paths stored in the zip file must use a slash ('/') as path
//   separator; and they must be relative (i.e., they must not start with
//   a '/' - this is usually the case if the file was created w/o special
//   options).
// - The zip file system treats the file paths found in the zip internally
//   like absolute paths w/o a leading '/'; i.e., the paths are considered
//   relative to the root of the file system.
// - All path arguments to file system methods must be absolute paths.

package main

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
)

// zipFI is the zip-file based implementation of FileInfo
type zipFI struct {
	name string    // directory-local name
	file *zip.File // nil for a directory
}

func (fi zipFI) Name() string {
	return fi.name
}

func (fi zipFI) Size() int64 {
	if f := fi.file; f != nil {
		return int64(f.UncompressedSize)
	}
	return 0 // directory
}

func (fi zipFI) Mtime_ns() int64 {
	if f := fi.file; f != nil {
		return f.Mtime_ns()
	}
	return 0 // directory has no modified time entry
}

func (fi zipFI) IsDirectory() bool {
	return fi.file == nil
}

func (fi zipFI) IsRegular() bool {
	return fi.file != nil
}

// zipFS is the zip-file based implementation of FileSystem
type zipFS struct {
	*zip.ReadCloser
	list zipList
}

func (fs *zipFS) Close() os.Error {
	fs.list = nil
	return fs.ReadCloser.Close()
}

func zipPath(name string) string {
	name = path.Clean(name)
	if !path.IsAbs(name) {
		panic(fmt.Sprintf("stat: not an absolute path: %s", name))
	}
	return name[1:] // strip leading '/'
}

func (fs *zipFS) stat(abspath string) (int, zipFI, os.Error) {
	i, exact := fs.list.lookup(abspath)
	if i < 0 {
		// abspath has leading '/' stripped - print it explicitly
		return -1, zipFI{}, fmt.Errorf("file not found: /%s", abspath)
	}
	_, name := path.Split(abspath)
	var file *zip.File
	if exact {
		file = fs.list[i] // exact match found - must be a file
	}
	return i, zipFI{name, file}, nil
}

func (fs *zipFS) Open(abspath string) (io.ReadCloser, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	if err != nil {
		return nil, err
	}
	if fi.IsDirectory() {
		return nil, fmt.Errorf("Open: %s is a directory", abspath)
	}
	return fi.file.Open()
}

func (fs *zipFS) Lstat(abspath string) (FileInfo, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	return fi, err
}

func (fs *zipFS) Stat(abspath string) (FileInfo, os.Error) {
	_, fi, err := fs.stat(zipPath(abspath))
	return fi, err
}

func (fs *zipFS) ReadDir(abspath string) ([]FileInfo, os.Error) {
	path := zipPath(abspath)
	i, fi, err := fs.stat(path)
	if err != nil {
		return nil, err
	}
	if !fi.IsDirectory() {
		return nil, fmt.Errorf("ReadDir: %s is not a directory", abspath)
	}

	var list []FileInfo
	dirname := path + "/"
	prevname := ""
	for _, e := range fs.list[i:] {
		if !strings.HasPrefix(e.Name, dirname) {
			break // not in the same directory anymore
		}
		name := e.Name[len(dirname):] // local name
		file := e
		if i := strings.IndexRune(name, '/'); i >= 0 {
			// We infer directories from files in subdirectories.
			// If we have x/y, return a directory entry for x.
			name = name[0:i] // keep local directory name only
			file = nil
		}
		// If we have x/y and x/z, don't return two directory entries for x.
		// TODO(gri): It should be possible to do this more efficiently
		// by determining the (fs.list) range of local directory entries
		// (via two binary searches).
		if name != prevname {
			list = append(list, zipFI{name, file})
			prevname = name
		}
	}

	return list, nil
}

func (fs *zipFS) ReadFile(abspath string) ([]byte, os.Error) {
	rc, err := fs.Open(abspath)
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(rc)
}

func NewZipFS(rc *zip.ReadCloser) FileSystem {
	list := make(zipList, len(rc.File))
	copy(list, rc.File) // sort a copy of rc.File
	sort.Sort(list)
	return &zipFS{rc, list}
}

type zipList []*zip.File

// zipList implements sort.Interface
func (z zipList) Len() int           { return len(z) }
func (z zipList) Less(i, j int) bool { return z[i].Name < z[j].Name }
func (z zipList) Swap(i, j int)      { z[i], z[j] = z[j], z[i] }

// lookup returns the smallest index of an entry with an exact match
// for name, or an inexact match starting with name/. If there is no
// such entry, the result is -1, false.
func (z zipList) lookup(name string) (index int, exact bool) {
	// look for exact match first (name comes before name/ in z)
	i := sort.Search(len(z), func(i int) bool {
		return name <= z[i].Name
	})
	if i >= len(z) {
		return -1, false
	}
	// 0 <= i < len(z)
	if z[i].Name == name {
		return i, true
	}

	// look for inexact match (must be in z[i:], if present)
	z = z[i:]
	name += "/"
	j := sort.Search(len(z), func(i int) bool {
		return name <= z[i].Name
	})
	if j >= len(z) {
		return -1, false
	}
	// 0 <= j < len(z)
	if strings.HasPrefix(z[j].Name, name) {
		return i + j, false
	}

	return -1, false
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 bgp

import (
	"encoding/binary"
	"fmt"
	"net"
)

const (
	RPKI_DEFAULT_PORT = 323
)

const (
	RTR_SERIAL_NOTIFY = iota
	RTR_SERIAL_QUERY
	RTR_RESET_QUERY
	RTR_CACHE_RESPONSE
	RTR_IPV4_PREFIX
	_
	RTR_IPV6_PREFIX
	RTR_END_OF_DATA
	RTR_CACHE_RESET
	_
	RTR_ERROR_REPORT
)

const (
	RTR_SERIAL_NOTIFY_LEN  = 12
	RTR_SERIAL_QUERY_LEN   = 12
	RTR_RESET_QUERY_LEN    = 8
	RTR_CACHE_RESPONSE_LEN = 8
	RTR_IPV4_PREFIX_LEN    = 20
	RTR_IPV6_PREFIX_LEN    = 32
	RTR_END_OF_DATA_LEN    = 12
	RTR_CACHE_RESET_LEN    = 8
	RTR_MIN_LEN            = 8
)

type RTRMessage interface {
	DecodeFromBytes([]byte) error
	Serialize() ([]byte, error)
}

type RTRCommon struct {
	Version      uint8
	Type         uint8
	SessionID    uint16
	Len          uint32
	SerialNumber uint32
}

func (m *RTRCommon) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.SerialNumber = binary.BigEndian.Uint32(data[8:12])
	return nil
}

func (m *RTRCommon) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	binary.BigEndian.PutUint32(data[8:12], m.SerialNumber)
	return data, nil
}

type RTRSerialNotify struct {
	RTRCommon
}

type RTRSerialQuery struct {
	RTRCommon
}

type RTRReset struct {
	Version uint8
	Type    uint8
	Len     uint32
}

func (m *RTRReset) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.Len = binary.BigEndian.Uint32(data[4:8])
	return nil
}

func (m *RTRReset) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

type RTRResetQuery struct {
	RTRReset
}

func (m *RTRResetQuery) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

func NewRTRResetQuery() *RTRResetQuery {
	return &RTRResetQuery{
		RTRReset{
			Type: RTR_RESET_QUERY,
			Len:  RTR_RESET_QUERY_LEN,
		},
	}
}

type RTRCacheResponse struct {
	Version   uint8
	Type      uint8
	SessionID uint16
	Len       uint32
}

func (m *RTRCacheResponse) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	return nil
}

func (m *RTRCacheResponse) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

type RTRIPPrefix struct {
	Version   uint8
	Type      uint8
	SessionID uint16
	Len       uint32
	Flags     uint8
	PrefixLen uint8
	MaxLen    uint8
	Prefix    net.IP
	AS        uint32
}

func (m *RTRIPPrefix) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.Flags = data[8]
	m.PrefixLen = data[9]
	m.MaxLen = data[10]
	if m.Type == RTR_IPV4_PREFIX {
		m.Prefix = net.IP(data[12:16]).To4()
		m.AS = binary.BigEndian.Uint32(data[16:20])
	} else {
		m.Prefix = net.IP(data[12:28]).To16()
		m.AS = binary.BigEndian.Uint32(data[28:32])
	}
	return nil
}

func (m *RTRIPPrefix) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	data[8] = m.Flags
	data[9] = m.PrefixLen
	data[10] = m.MaxLen
	if m.Type == RTR_IPV4_PREFIX {
		copy(data[12:16], m.Prefix.To4())
		binary.BigEndian.PutUint32(data[16:20], m.AS)
	} else {
		copy(data[12:28], m.Prefix.To16())
		binary.BigEndian.PutUint32(data[28:32], m.AS)
	}
	return data, nil
}

type RTREndOfData struct {
	RTRCommon
}

type RTRCacheReset struct {
	RTRReset
}

type RTRErrorReport struct {
	Version   uint8
	Type      uint8
	SessionID uint16
	Len       uint32
	PDULen    uint32
	PDU       []byte
	TextLen   uint32
	Text      []byte
}

func (m *RTRErrorReport) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.PDULen = binary.BigEndian.Uint32(data[8:12])
	m.PDU = make([]byte, m.PDULen)
	copy(m.PDU, data[12:12+m.PDULen])
	m.TextLen = binary.BigEndian.Uint32(data[12+m.PDULen : 16+m.PDULen])
	m.PDU = make([]byte, m.TextLen)
	copy(m.Text, data[16+m.PDULen:])
	return nil
}

func (m *RTRErrorReport) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	binary.BigEndian.PutUint32(data[8:12], m.PDULen)
	copy(data[12:], m.PDU)
	binary.BigEndian.PutUint32(data[12+m.PDULen:16+m.PDULen], m.TextLen)
	copy(data[16+m.PDULen:], m.Text)
	return data, nil
}

func SplitRTR(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 || len(data) < RTR_MIN_LEN {
		return 0, nil, nil
	}

	totalLen := binary.BigEndian.Uint32(data[4:8])
	if totalLen < RTR_MIN_LEN {
		return 0, nil, fmt.Errorf("Invalid length: %d", totalLen)
	}
	if uint32(len(data)) < totalLen {
		return 0, nil, nil
	}
	return int(totalLen), data[0:totalLen], nil
}

func ParseRTR(data []byte) (RTRMessage, error) {
	var msg RTRMessage
	switch data[1] {
	case RTR_SERIAL_NOTIFY:
		msg = &RTRSerialNotify{}
	case RTR_SERIAL_QUERY:
		msg = &RTRSerialQuery{}
	case RTR_RESET_QUERY:
		msg = &RTRResetQuery{}
	case RTR_CACHE_RESPONSE:
		msg = &RTRCacheResponse{}
	case RTR_IPV4_PREFIX:
		msg = &RTRIPPrefix{}
	case RTR_IPV6_PREFIX:
		msg = &RTRIPPrefix{}
	case RTR_END_OF_DATA:
		msg = &RTREndOfData{}
	case RTR_CACHE_RESET:
		msg = &RTRCacheReset{}
	case RTR_ERROR_REPORT:
		msg = &RTRErrorReport{}
	default:
		return nil, fmt.Errorf("unknown RTR message type %d:", data[1])
	}
	err := msg.DecodeFromBytes(data)
	return msg, err
}
<commit_msg>packet: remove SessionID field from RTRIPPrefix<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 bgp

import (
	"encoding/binary"
	"fmt"
	"net"
)

const (
	RPKI_DEFAULT_PORT = 323
)

const (
	RTR_SERIAL_NOTIFY = iota
	RTR_SERIAL_QUERY
	RTR_RESET_QUERY
	RTR_CACHE_RESPONSE
	RTR_IPV4_PREFIX
	_
	RTR_IPV6_PREFIX
	RTR_END_OF_DATA
	RTR_CACHE_RESET
	_
	RTR_ERROR_REPORT
)

const (
	RTR_SERIAL_NOTIFY_LEN  = 12
	RTR_SERIAL_QUERY_LEN   = 12
	RTR_RESET_QUERY_LEN    = 8
	RTR_CACHE_RESPONSE_LEN = 8
	RTR_IPV4_PREFIX_LEN    = 20
	RTR_IPV6_PREFIX_LEN    = 32
	RTR_END_OF_DATA_LEN    = 12
	RTR_CACHE_RESET_LEN    = 8
	RTR_MIN_LEN            = 8
)

type RTRMessage interface {
	DecodeFromBytes([]byte) error
	Serialize() ([]byte, error)
}

type RTRCommon struct {
	Version      uint8
	Type         uint8
	SessionID    uint16
	Len          uint32
	SerialNumber uint32
}

func (m *RTRCommon) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.SerialNumber = binary.BigEndian.Uint32(data[8:12])
	return nil
}

func (m *RTRCommon) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	binary.BigEndian.PutUint32(data[8:12], m.SerialNumber)
	return data, nil
}

type RTRSerialNotify struct {
	RTRCommon
}

type RTRSerialQuery struct {
	RTRCommon
}

type RTRReset struct {
	Version uint8
	Type    uint8
	Len     uint32
}

func (m *RTRReset) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.Len = binary.BigEndian.Uint32(data[4:8])
	return nil
}

func (m *RTRReset) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

type RTRResetQuery struct {
	RTRReset
}

func (m *RTRResetQuery) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

func NewRTRResetQuery() *RTRResetQuery {
	return &RTRResetQuery{
		RTRReset{
			Type: RTR_RESET_QUERY,
			Len:  RTR_RESET_QUERY_LEN,
		},
	}
}

type RTRCacheResponse struct {
	Version   uint8
	Type      uint8
	SessionID uint16
	Len       uint32
}

func (m *RTRCacheResponse) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	return nil
}

func (m *RTRCacheResponse) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	return data, nil
}

type RTRIPPrefix struct {
	Version   uint8
	Type      uint8
	Len       uint32
	Flags     uint8
	PrefixLen uint8
	MaxLen    uint8
	Prefix    net.IP
	AS        uint32
}

func (m *RTRIPPrefix) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.Flags = data[8]
	m.PrefixLen = data[9]
	m.MaxLen = data[10]
	if m.Type == RTR_IPV4_PREFIX {
		m.Prefix = net.IP(data[12:16]).To4()
		m.AS = binary.BigEndian.Uint32(data[16:20])
	} else {
		m.Prefix = net.IP(data[12:28]).To16()
		m.AS = binary.BigEndian.Uint32(data[28:32])
	}
	return nil
}

func (m *RTRIPPrefix) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	data[8] = m.Flags
	data[9] = m.PrefixLen
	data[10] = m.MaxLen
	if m.Type == RTR_IPV4_PREFIX {
		copy(data[12:16], m.Prefix.To4())
		binary.BigEndian.PutUint32(data[16:20], m.AS)
	} else {
		copy(data[12:28], m.Prefix.To16())
		binary.BigEndian.PutUint32(data[28:32], m.AS)
	}
	return data, nil
}

type RTREndOfData struct {
	RTRCommon
}

type RTRCacheReset struct {
	RTRReset
}

type RTRErrorReport struct {
	Version   uint8
	Type      uint8
	SessionID uint16
	Len       uint32
	PDULen    uint32
	PDU       []byte
	TextLen   uint32
	Text      []byte
}

func (m *RTRErrorReport) DecodeFromBytes(data []byte) error {
	m.Version = data[0]
	m.Type = data[1]
	m.SessionID = binary.BigEndian.Uint16(data[2:4])
	m.Len = binary.BigEndian.Uint32(data[4:8])
	m.PDULen = binary.BigEndian.Uint32(data[8:12])
	m.PDU = make([]byte, m.PDULen)
	copy(m.PDU, data[12:12+m.PDULen])
	m.TextLen = binary.BigEndian.Uint32(data[12+m.PDULen : 16+m.PDULen])
	m.PDU = make([]byte, m.TextLen)
	copy(m.Text, data[16+m.PDULen:])
	return nil
}

func (m *RTRErrorReport) Serialize() ([]byte, error) {
	data := make([]byte, m.Len)
	data[0] = m.Version
	data[1] = m.Type
	binary.BigEndian.PutUint16(data[2:4], m.SessionID)
	binary.BigEndian.PutUint32(data[4:8], m.Len)
	binary.BigEndian.PutUint32(data[8:12], m.PDULen)
	copy(data[12:], m.PDU)
	binary.BigEndian.PutUint32(data[12+m.PDULen:16+m.PDULen], m.TextLen)
	copy(data[16+m.PDULen:], m.Text)
	return data, nil
}

func SplitRTR(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 || len(data) < RTR_MIN_LEN {
		return 0, nil, nil
	}

	totalLen := binary.BigEndian.Uint32(data[4:8])
	if totalLen < RTR_MIN_LEN {
		return 0, nil, fmt.Errorf("Invalid length: %d", totalLen)
	}
	if uint32(len(data)) < totalLen {
		return 0, nil, nil
	}
	return int(totalLen), data[0:totalLen], nil
}

func ParseRTR(data []byte) (RTRMessage, error) {
	var msg RTRMessage
	switch data[1] {
	case RTR_SERIAL_NOTIFY:
		msg = &RTRSerialNotify{}
	case RTR_SERIAL_QUERY:
		msg = &RTRSerialQuery{}
	case RTR_RESET_QUERY:
		msg = &RTRResetQuery{}
	case RTR_CACHE_RESPONSE:
		msg = &RTRCacheResponse{}
	case RTR_IPV4_PREFIX:
		msg = &RTRIPPrefix{}
	case RTR_IPV6_PREFIX:
		msg = &RTRIPPrefix{}
	case RTR_END_OF_DATA:
		msg = &RTREndOfData{}
	case RTR_CACHE_RESET:
		msg = &RTRCacheReset{}
	case RTR_ERROR_REPORT:
		msg = &RTRErrorReport{}
	default:
		return nil, fmt.Errorf("unknown RTR message type %d:", data[1])
	}
	err := msg.DecodeFromBytes(data)
	return msg, err
}
<|endoftext|>


<commit_before>package rtp

import (
	"math/rand"
	"time"
)

// Payloader payloads a byte array for use as rtp.Packet payloads
type Payloader interface {
	Payload(mtu int, payload []byte) [][]byte
}

// Packetizer packetizes a payload
type Packetizer interface {
	Packetize(payload []byte, samples uint32) []*Packet
}

type packetizer struct {
	MTU         int
	PayloadType uint8
	SSRC        uint32
	Payloader   Payloader
	Sequencer   Sequencer
	Timestamp   uint32
	ClockRate   uint32
}

// NewPacketizer returns a new instance of a Packetizer for a specific payloader
func NewPacketizer(mtu int, pt uint8, ssrc uint32, payloader Payloader, sequencer Sequencer, clockRate uint32) Packetizer {
	rs := rand.NewSource(time.Now().UnixNano())
	r := rand.New(rs)

	return &packetizer{
		MTU:         mtu,
		PayloadType: pt,
		SSRC:        ssrc,
		Payloader:   payloader,
		Sequencer:   sequencer,
		Timestamp:   r.Uint32(),
		ClockRate:   clockRate,
	}
}

// Packetize packetizes the payload of an RTP packet and returns one or more RTP packets
func (p *packetizer) Packetize(payload []byte, samples uint32) []*Packet {
	payloads := p.Payloader.Payload(p.MTU-12, payload)
	packets := make([]*Packet, len(payloads))

	for i, pp := range payloads {
		packets[i] = &Packet{
			Version:        2,
			Padding:        false,
			Extension:      false,
			Marker:         i == len(payloads)-1,
			PayloadType:    p.PayloadType,
			SequenceNumber: p.Sequencer.NextSequenceNumber(),
			Timestamp:      p.Timestamp, // Figure out how to do timestamps
			SSRC:           p.SSRC,
			Payload:        pp,
		}
	}
	p.Timestamp += samples

	return packets
}
<commit_msg>Handle empty payload in packetizer.Packetize()<commit_after>package rtp

import (
	"math/rand"
	"time"
)

// Payloader payloads a byte array for use as rtp.Packet payloads
type Payloader interface {
	Payload(mtu int, payload []byte) [][]byte
}

// Packetizer packetizes a payload
type Packetizer interface {
	Packetize(payload []byte, samples uint32) []*Packet
}

type packetizer struct {
	MTU         int
	PayloadType uint8
	SSRC        uint32
	Payloader   Payloader
	Sequencer   Sequencer
	Timestamp   uint32
	ClockRate   uint32
}

// NewPacketizer returns a new instance of a Packetizer for a specific payloader
func NewPacketizer(mtu int, pt uint8, ssrc uint32, payloader Payloader, sequencer Sequencer, clockRate uint32) Packetizer {
	rs := rand.NewSource(time.Now().UnixNano())
	r := rand.New(rs)

	return &packetizer{
		MTU:         mtu,
		PayloadType: pt,
		SSRC:        ssrc,
		Payloader:   payloader,
		Sequencer:   sequencer,
		Timestamp:   r.Uint32(),
		ClockRate:   clockRate,
	}
}

// Packetize packetizes the payload of an RTP packet and returns one or more RTP packets
func (p *packetizer) Packetize(payload []byte, samples uint32) []*Packet {
	// Guard against an empty payload
	if len(payload) == 0 {
		return nil
	}

	payloads := p.Payloader.Payload(p.MTU-12, payload)
	packets := make([]*Packet, len(payloads))

	for i, pp := range payloads {
		packets[i] = &Packet{
			Version:        2,
			Padding:        false,
			Extension:      false,
			Marker:         i == len(payloads)-1,
			PayloadType:    p.PayloadType,
			SequenceNumber: p.Sequencer.NextSequenceNumber(),
			Timestamp:      p.Timestamp, // Figure out how to do timestamps
			SSRC:           p.SSRC,
			Payload:        pp,
		}
	}
	p.Timestamp += samples

	return packets
}
<|endoftext|>


<commit_before>// A Naive Bayesian Classifier
// Jake Brukhman <jbrukh@gmail.com>

//
// BAYESIAN CLASSIFICATION REFRESHER: suppose you have a set
// of classes (e.g. categories) C := {C_1, ..., C_n}, and a
// document D consisting of words D := {W_1, ..., W_k}.
// We wish to ascertain the probability that the document
// belongs to some class C_j given some set of training data
// associating documents and classes.
//
// By Bayes Theorem, we have that
//
//    P(C_j|D) = P(D|C_j)*P(C_j)/P(D).
//
// The LHS is the probability that the document belongs to class
// C_j given the document itself (by which is meant, in practice,
// the word frequencies occurring in this document), and our program
// will calculate this probability for each j and spit out the
// most likely class for this document.
//
// P(C_j) is referred to as the "prior" probability, or the
// probability that a document belongs to C_j in general, without
// seeing the document first. P(D|C_j) is the probability of seeing
// such a document, given that it belongs to C_j. Here, by assuming
// that words appear independently in documents (this being the
// "naive" assumption), we can estimate
//
//    P(D|C_j) ~= P(W_1|C_j)*...*P(W_k|C_j)
//
// where P(W_i|C_j) is the probability of seeing the given word
// in a document of the given class. Finally, P(D) can be seen as
// merely a scaling factor and is not strictly relevant to
// classificiation, unless you want to normalize the resulting
// scores and actually see probabilities. In this case, note that
//
//    P(D) = SUM_j(P(D|C_j)*P(C_j))
//
// End of refresher.
//
package bayesian

// defaultProb is the tiny non-zero probability that a word
// we have not seen before appears in the class.
const defaultProb = 0.001

// the type of float we use here
type float float64

// This type defines a set of classes that the classifier will
// filter: C = {C_1, ..., C_n}. You should define your classes
// as a set of constants, for example as follows:
//
//    const (
//        Good Class = "Good"
//        Bad Class = "Bad
//    )
//
type Class string

// Classifier implements the Naive Bayesian Classifier.
type Classifier struct {
    Classes []Class
    datas map[Class]*classData
}

// classData holds the frequency data for words in a
// particular class. In the future, we may replace this
// structure with a trie-like structure for more
// efficient storage.
type classData struct {
    freqs map[string]int
    total int
}

// newClassData creates a new empty class data node.
func newClassData() *classData {
    return &classData{
        freqs: make(map[string]int),
    }
}

// P(W|Cj) -- the probability of seeing a particular word
// in a document of this class.
func (this *classData) getWordProb(word string) float {
    value, ok := this.freqs[word]
    if !ok {
        return defaultProb
    }
    return float(value)/float(this.total)
}

// P(D|C_j) -- the probability of seeing this set of words
// in a document of this class. Note that words should not
// be empty.
func (this *classData) getWordsProb(words []string) (prob float) {
    prob = 1
    for _, word := range words {
        prob *= this.getWordProb(word)
    }
    return
}

// New creates a new Classifier.
func NewClassifier(classes ...Class) (inst *Classifier) {
    if len(classes) < 2 {
        panic("provide at least two classes")
    }
    inst = &Classifier{
            classes,
            make(map[Class]*classData),
    }
    for _, class := range classes {
        inst.datas[class] = newClassData()
    }
    return
}

// getPriors returns the prior probabilities for the
// classes provided -- P(C_i). There is a way to
// smooth priors, currently not implemented here.
func (this *Classifier) getPriors() (priors []float) {
    n := len(this.Classes)
    priors = make([]float, n, n)
    sum := 0
    for index, class := range this.Classes {
        total := this.datas[class].total;
        priors[index] = float(total)
        sum += total
    }
    if sum != 0 {
        for i := 0; i < n; i++ {
            priors[i] /= float(sum)
        }
    }
    return
}

// Learn will train the classifier on the provided data.
func (this *Classifier) Learn(words []string, which Class) {
    data := this.datas[which]
    for _, word := range words {
        data.freqs[word]++
        data.total++
    }
}

// Score will produce an array of probabilities that correspond
// to its opinion on the document in question, and whether it
// belongs to the given class. The order of the probabilities
// in the return values follows the order of the inital array
// of Class objects parameterized to the New() function. If no
// training data has been provided, this will return a 0 array.
//
// Additionally, this function will return the index of the
// maximum probability. The value of this number is given by
// scores[inx]. The class of that corresponds to this number
// is classifier.Classes[inx]. If more than one of the
// returned probabilities has the maximum values, then
// strict is false.
func (this *Classifier) Score(words []string) (scores []float, inx int, strict bool) {
    n := len(this.Classes)
    scores = make([]float, n, n)
    priors := this.getPriors()
    sum := float(0)
    for index, class := range this.Classes {
        data := this.datas[class]
        score := priors[index]*data.getWordsProb(words)
        scores[index] = score
        sum += score
    }
    inx = 0
    strict = true
    for i := 0; i < n; i++ {
        scores[i] /= sum
        if scores[inx] < scores[i] {
            inx = i
            strict = true
        } else if scores[inx] == scores[i] && i != 0 {
            strict = false
        }
    }
    return
}
<commit_msg>Using logs for score.<commit_after>// A Naive Bayesian Classifier
// Jake Brukhman <jbrukh@gmail.com>

//
// BAYESIAN CLASSIFICATION REFRESHER: suppose you have a set
// of classes (e.g. categories) C := {C_1, ..., C_n}, and a
// document D consisting of words D := {W_1, ..., W_k}.
// We wish to ascertain the probability that the document
// belongs to some class C_j given some set of training data
// associating documents and classes.
//
// By Bayes Theorem, we have that
//
//    P(C_j|D) = P(D|C_j)*P(C_j)/P(D).
//
// The LHS is the probability that the document belongs to class
// C_j given the document itself (by which is meant, in practice,
// the word frequencies occurring in this document), and our program
// will calculate this probability for each j and spit out the
// most likely class for this document.
//
// P(C_j) is referred to as the "prior" probability, or the
// probability that a document belongs to C_j in general, without
// seeing the document first. P(D|C_j) is the probability of seeing
// such a document, given that it belongs to C_j. Here, by assuming
// that words appear independently in documents (this being the
// "naive" assumption), we can estimate
//
//    P(D|C_j) ~= P(W_1|C_j)*...*P(W_k|C_j)
//
// where P(W_i|C_j) is the probability of seeing the given word
// in a document of the given class. Finally, P(D) can be seen as
// merely a scaling factor and is not strictly relevant to
// classificiation, unless you want to normalize the resulting
// scores and actually see probabilities. In this case, note that
//
//    P(D) = SUM_j(P(D|C_j)*P(C_j))
//
// One practical issue with performing these calculations is the
// possibility of float underflow when calculating P(D|C_j), as
// individual word probabilities can be arbitrarily small, and
// a document can have an arbitrarily large number of them. A
// typical method for dealing with this case is to transform the
// probability to the log domain and perform additions instead
// of multiplications:
//
//   log P(C_j|D) ~ log(P(C_j)) + SUM_i(log P(W_i|C_j))
//
// where i = 1, ..., k. Note that by doing this, we are discarding
// the scaling factor P(D) and our scores are no longer
// probabilities.
//
package bayesian

import "math"

// defaultProb is the tiny non-zero probability that a word
// we have not seen before appears in the class.
const defaultProb = 0.001

// the type of float we use here
type float float64

// This type defines a set of classes that the classifier will
// filter: C = {C_1, ..., C_n}. You should define your classes
// as a set of constants, for example as follows:
//
//    const (
//        Good Class = "Good"
//        Bad Class = "Bad
//    )
//
type Class string

// Classifier implements the Naive Bayesian Classifier.
type Classifier struct {
    Classes []Class
    datas map[Class]*classData
}

// classData holds the frequency data for words in a
// particular class. In the future, we may replace this
// structure with a trie-like structure for more
// efficient storage.
type classData struct {
    freqs map[string]int
    total int
}

// newClassData creates a new empty class data node.
func newClassData() *classData {
    return &classData{
        freqs: make(map[string]int),
    }
}

// P(W|Cj) -- the probability of seeing a particular word
// in a document of this class.
func (this *classData) getWordProb(word string) float {
    value, ok := this.freqs[word]
    if !ok {
        return defaultProb
    }
    return float(value)/float(this.total)
}

// P(D|C_j) -- the probability of seeing this set of words
// in a document of this class.
//
// Note that words should not be empty, and this method of
// calulation is prone to underflow if there are many words
// and their individual probabilties are small.
func (this *classData) getWordsProb(words []string) (prob float) {
    prob = 1
    for _, word := range words {
        prob *= this.getWordProb(word)
    }
    return
}

// New creates a new Classifier.
func NewClassifier(classes ...Class) (inst *Classifier) {
    if len(classes) < 2 {
        panic("provide at least two classes")
    }
    inst = &Classifier{
            classes,
            make(map[Class]*classData),
    }
    for _, class := range classes {
        inst.datas[class] = newClassData()
    }
    return
}

// getPriors returns the prior probabilities for the
// classes provided -- P(C_i). There is a way to
// smooth priors, currently not implemented here.
func (this *Classifier) getPriors() (priors []float) {
    n := len(this.Classes)
    priors = make([]float, n, n)
    sum := 0
    for index, class := range this.Classes {
        total := this.datas[class].total;
        priors[index] = float(total)
        sum += total
    }
    if sum != 0 {
        for i := 0; i < n; i++ {
            priors[i] /= float(sum)
        }
    }
    return
}

// Learn will train the classifier on the provided data.
func (this *Classifier) Learn(words []string, which Class) {
    data := this.datas[which]
    for _, word := range words {
        data.freqs[word]++
        data.total++
    }
}

// Score will produce an array of scores that correspond
// to its opinion on the document in question, and whether it
// belongs to the given class. The order of the scores
// in the return values follows the order of the inital array
// of Class objects parameterized to the NewClassifier() function.
// If no training data has been provided, this will return
// a 0 array.
//
// The value of the score is proportional to the likelihood,
// even if the score is negative, so that the score with the
// greatest value corresponds to the most likely class.
//
// Additionally, this function will return the index of the
// maximum probability. The value of this number is given by
// scores[inx]. The class of that corresponds to this number
// is classifier.Classes[inx]. If more than one of the
// returned probabilities has the maximum values, then
// strict is false.
func (this *Classifier) Score(words []string) (scores []float, inx int, strict bool) {
    n := len(this.Classes)
    scores = make([]float, n, n)
    priors := this.getPriors()

    // calculate the score for each class
    for index, class := range this.Classes {
        data := this.datas[class]
        // this is the sum of the logarithms
        // as outlined in the refresher
        score := math.Log(float64(priors[index]))
        for _, word := range words {
            score += math.Log(float64(data.getWordProb(word)))
        }
        scores[index] = float(score)
    }

    // calculate the index of the maximum score
    inx = 0
    strict = true
    for i := 1; i < n; i++ {
        if scores[inx] < scores[i] {
            inx = i
            strict = true
        } else if scores[inx] == scores[i] {
            strict = false
        }
    }
    return
}
<|endoftext|>