<commit_before>package html

import (
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"github.com/slyrz/newscat/util"
	"io"
	"unicode"
)

// Errors returned during Document parsing.
var (
	ErrNoHTML = errors.New("missing html element")
	ErrNoHead = errors.New("missing head element")
	ErrNoBody = errors.New("missing body element")
)

// Document is a parsed HTML document that extracts the document title and
// holds unexported pointers to the html, head and body nodes.
type Document struct {
	Title  *util.Text // the <title>...</title> text.
	Chunks []*Chunk   // all chunks found in this document.

	// Unexported fields.
	html *html.Node // the ... part
	head *html.Node // the ... part
	body *html.Node // the ... part

	// State variables used during parsing.
	ancestors int                // bitmask to track specific ancestor types
	linkText  map[*html.Node]int // length of text inside  tags
	normText  map[*html.Node]int // length of text outside  tags
}

// NewDocument parses the HTML data provided through an io.Reader interface.
func NewDocument(r io.Reader) (*Document, error) {
	doc := new(Document)
	if err := doc.init(r); err != nil {
		return nil, err
	}
	return doc, nil
}

func (doc *Document) init(r io.Reader) error {
	doc.Title = util.NewText()

	root, err := html.Parse(r)
	if err != nil {
		return err
	}

	// Assign the fields html, head and body from the HTML page.
	iterateNode(root, func(n *html.Node) int {
		switch n.DataAtom {
		case atom.Html:
			doc.html = n
			return IterNext
		case atom.Body:
			doc.body = n
			return IterSkip
		case atom.Head:
			doc.head = n
			return IterSkip
		}
		// Keep going as long as we're missing some nodes.
		return IterNext
	})

	switch {
	case doc.html == nil:
		return ErrNoHTML
	case doc.head == nil:
		return ErrNoHead
	case doc.body == nil:
		return ErrNoBody
	}

	// Detect the document title.
	iterateNode(doc.head, func(n *html.Node) int {
		if n.Type == html.ElementNode && n.DataAtom == atom.Title {
			iterateText(n, doc.Title.WriteString)
			return IterStop
		}
		return IterNext
	})

	doc.Chunks = make([]*Chunk, 0, 512)
	doc.linkText = make(map[*html.Node]int)
	doc.normText = make(map[*html.Node]int)

	doc.cleanBody(doc.body, 0)
	doc.countText(doc.body, false)
	doc.parseBody(doc.body)

	// Now we link the chunks.
	min, max := 0, len(doc.Chunks)-1
	for i := range doc.Chunks {
		if i > min {
			doc.Chunks[i].Prev = doc.Chunks[i-1]
		}
		if i < max {
			doc.Chunks[i].Next = doc.Chunks[i+1]
		}
	}
	return nil
}

const (
	// We remember a few special node types when descending into their
	// children.
	AncestorArticle = 1 << iota
	AncestorAside
	AncestorBlockquote
	AncestorList
)

// countText counts the text inside of links and the text outside of links
// per html.Node. Counting is done cumulative, so the umbers of a parent node
// include the numbers of it's child nodes.
func (doc *Document) countText(n *html.Node, insideLink bool) (linkText int, normText int) {
	linkText = 0
	normText = 0
	if n.Type == html.ElementNode && n.DataAtom == atom.A {
		insideLink = true
	}
	for s := n.FirstChild; s != nil; s = s.NextSibling {
		linkTextChild, normTextChild := doc.countText(s, insideLink)
		linkText += linkTextChild
		normText += normTextChild
	}
	if n.Type == html.TextNode {
		count := 0
		for _, rune := range n.Data {
			if unicode.IsLetter(rune) {
				count += 1
			}
		}
		if insideLink {
			linkText += count
		} else {
			normText += count
		}
	}
	doc.linkText[n] = linkText
	doc.normText[n] = normText
	return
}

var removeElements = map[atom.Atom]bool{
	atom.Address:    true,
	atom.Audio:      true,
	atom.Button:     true,
	atom.Canvas:     true,
	atom.Caption:    true,
	atom.Fieldset:   true,
	atom.Figcaption: true,
	atom.Figure:     true,
	atom.Footer:     true,
	atom.Form:       true,
	atom.Frame:      true,
	atom.Iframe:     true,
	atom.Map:        true,
	atom.Menu:       true,
	atom.Nav:        true,
	atom.Noscript:   true,
	atom.Object:     true,
	atom.Option:     true,
	atom.Output:     true,
	atom.Script:     true,
	atom.Select:     true,
	atom.Style:      true,
	atom.Svg:        true,
	atom.Textarea:   true,
	atom.Video:      true,
}

// cleanBody removes unwanted HTML elements from the HTML body.
func (doc *Document) cleanBody(n *html.Node, level int) {
	// removeNode returns true if a node should be removed from HTML document.
	removeNode := func(c *html.Node, level int) bool {
		if c.DataAtom == atom.Table {
			return level > 5
		}
		return removeElements[c.DataAtom]
	}

	var curr *html.Node = n.FirstChild
	var next *html.Node = nil
	for ; curr != nil; curr = next {
		// We have to remember the next sibling here because calling RemoveChild
		// sets curr's NextSibling pointer to nil and we would quit the loop
		// prematurely.
		next = curr.NextSibling
		if curr.Type == html.ElementNode {
			if removeNode(curr, level) {
				n.RemoveChild(curr)
			} else {
				doc.cleanBody(curr, level+1)
			}
		}
	}
}

var (
	ignoreNames = util.NewRegexFromWords(
		"breadcrumb",
		"byline",
		"caption",
		"comment",
		"community",
		"credit",
		"description",
		"email",
		"footer",
		"gallery",
		"hide",
		"infotext",
		"photo",
		"related",
		"shares",
		"social",
		"story[-_]?bar",
		"story[-_]?feature",
	)
	ignoreStyle = util.NewRegex(`(?i)display:\s*none`)
)

// parseBody parses the ... part of the HTML page. It creates
// Chunks for every html.TextNode found in the body.
func (doc *Document) parseBody(n *html.Node) {
	switch n.Type {
	case html.ElementNode:
		// We ignore the node if it has some nasty classes/ids/itemprops or if
		// its style attribute contains "display: none".
		for _, attr := range n.Attr {
			switch attr.Key {
			case "id", "class", "itemprop":
				if ignoreNames.In(attr.Val) {
					return
				}
			case "style":
				if ignoreStyle.In(attr.Val) {
					return
				}
			}
		}
		ancestorMask := 0
		switch n.DataAtom {
		// We convert headings and links to text immediately. This is easier
		// and feasible because headings and links don't contain many children.
		// Descending into these children and handling every TextNode separately
		// would make things unnecessary complicated and our results noisy.
		case atom.H1, atom.H2, atom.H3, atom.H4, atom.H5, atom.H6, atom.A:
			if chunk, err := NewChunk(doc, n); err == nil {
				doc.Chunks = append(doc.Chunks, chunk)
			}
			return
		// Now mask the element type, but only if it isn't already set.
		// If we mask a bit which was already set by one of our callers, we'd also
		// clear it at the end of this function, though it actually should be cleared
		// by the caller.
		case atom.Article:
			ancestorMask = AncestorArticle &^ doc.ancestors
		case atom.Aside:
			ancestorMask = AncestorAside &^ doc.ancestors
		case atom.Blockquote:
			ancestorMask = AncestorBlockquote &^ doc.ancestors
		case atom.Ul, atom.Ol:
			ancestorMask = AncestorList &^ doc.ancestors
		}
		// Add our mask to the ancestor bitmask.
		doc.ancestors |= ancestorMask
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			doc.parseBody(c)
		}
		// Remove our mask from the ancestor bitmask.
		doc.ancestors &^= ancestorMask
	case html.TextNode:
		if chunk, err := NewChunk(doc, n); err == nil {
			doc.Chunks = append(doc.Chunks, chunk)
		}
	}
}

// TextStat contains the number of words and sentences found in text.
type TextStat struct {
	Words     int // total number of words
	Sentences int // total number of sentences
	Count     int // number of texts used to calculate this stats
}

// GetClassStats groups the document chunks by their classes (defined by the
// class attribute of HTML nodes) and calculates TextStats for each class.
func (doc *Document) GetClassStats() map[string]*TextStat {
	result := make(map[string]*TextStat)
	for _, chunk := range doc.Chunks {
		for _, class := range chunk.Classes {
			if stat, ok := result[class]; ok {
				stat.Words += chunk.Text.Words
				stat.Sentences += chunk.Text.Sentences
				stat.Count += 1
			} else {
				result[class] = &TextStat{chunk.Text.Words, chunk.Text.Sentences, 1}
			}
		}
	}
	return result
}

// GetClusterStats groups the document chunks by common ancestors and
// calculates TextStats for each group of chunks.
func (doc *Document) GetClusterStats() map[*Chunk]*TextStat {
	// Don't ascend further than this.
	const maxAncestors = 3

	// Count TextStats for Chunk ancestors.
	ancestorStat := make(map[*html.Node]*TextStat)
	for _, chunk := range doc.Chunks {
		node, count := chunk.Block, 0
		for node != nil && count < maxAncestors {
			if stat, ok := ancestorStat[node]; ok {
				stat.Words += chunk.Text.Words
				stat.Sentences += chunk.Text.Sentences
				stat.Count += 1
			} else {
				ancestorStat[node] = &TextStat{chunk.Text.Words, chunk.Text.Sentences, 1}
			}
			node, count = node.Parent, count+1
		}
	}

	// Generate result. For each chunk pick the best TextStats from its ancestors.
	result := make(map[*Chunk]*TextStat)
	for _, chunk := range doc.Chunks {
		node := chunk.Block
		if node == nil {
			continue
		}
		// Start with the parent's TextStat. Then ascend and check if the
		// current chunk has an ancestor with better stats. Use the best stat
		// as result.
		stat := ancestorStat[node]
		for {
			if node = node.Parent; node == nil {
				break
			}
			if statPrev, ok := ancestorStat[node]; ok {
				if stat.Count < statPrev.Count {
					stat = statPrev
				}
			} else {
				break
			}
		}
		result[chunk] = stat
	}
	return result
}
<commit_msg>fixed spelling mistake in comment<commit_after>package html

import (
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"github.com/slyrz/newscat/util"
	"io"
	"unicode"
)

// Errors returned during Document parsing.
var (
	ErrNoHTML = errors.New("missing html element")
	ErrNoHead = errors.New("missing head element")
	ErrNoBody = errors.New("missing body element")
)

// Document is a parsed HTML document that extracts the document title and
// holds unexported pointers to the html, head and body nodes.
type Document struct {
	Title  *util.Text // the <title>...</title> text.
	Chunks []*Chunk   // all chunks found in this document.

	// Unexported fields.
	html *html.Node // the ... part
	head *html.Node // the ... part
	body *html.Node // the ... part

	// State variables used during parsing.
	ancestors int                // bitmask to track specific ancestor types
	linkText  map[*html.Node]int // length of text inside  tags
	normText  map[*html.Node]int // length of text outside  tags
}

// NewDocument parses the HTML data provided through an io.Reader interface.
func NewDocument(r io.Reader) (*Document, error) {
	doc := new(Document)
	if err := doc.init(r); err != nil {
		return nil, err
	}
	return doc, nil
}

func (doc *Document) init(r io.Reader) error {
	doc.Title = util.NewText()

	root, err := html.Parse(r)
	if err != nil {
		return err
	}

	// Assign the fields html, head and body from the HTML page.
	iterateNode(root, func(n *html.Node) int {
		switch n.DataAtom {
		case atom.Html:
			doc.html = n
			return IterNext
		case atom.Body:
			doc.body = n
			return IterSkip
		case atom.Head:
			doc.head = n
			return IterSkip
		}
		// Keep going as long as we're missing some nodes.
		return IterNext
	})

	switch {
	case doc.html == nil:
		return ErrNoHTML
	case doc.head == nil:
		return ErrNoHead
	case doc.body == nil:
		return ErrNoBody
	}

	// Detect the document title.
	iterateNode(doc.head, func(n *html.Node) int {
		if n.Type == html.ElementNode && n.DataAtom == atom.Title {
			iterateText(n, doc.Title.WriteString)
			return IterStop
		}
		return IterNext
	})

	doc.Chunks = make([]*Chunk, 0, 512)
	doc.linkText = make(map[*html.Node]int)
	doc.normText = make(map[*html.Node]int)

	doc.cleanBody(doc.body, 0)
	doc.countText(doc.body, false)
	doc.parseBody(doc.body)

	// Now we link the chunks.
	min, max := 0, len(doc.Chunks)-1
	for i := range doc.Chunks {
		if i > min {
			doc.Chunks[i].Prev = doc.Chunks[i-1]
		}
		if i < max {
			doc.Chunks[i].Next = doc.Chunks[i+1]
		}
	}
	return nil
}

const (
	// We remember a few special node types when descending into their
	// children.
	AncestorArticle = 1 << iota
	AncestorAside
	AncestorBlockquote
	AncestorList
)

// countText counts the text inside of links and the text outside of links
// per html.Node. Counting is done cumulative, so the numbers of a parent node
// include the numbers of its child nodes.
func (doc *Document) countText(n *html.Node, insideLink bool) (linkText int, normText int) {
	linkText = 0
	normText = 0
	if n.Type == html.ElementNode && n.DataAtom == atom.A {
		insideLink = true
	}
	for s := n.FirstChild; s != nil; s = s.NextSibling {
		linkTextChild, normTextChild := doc.countText(s, insideLink)
		linkText += linkTextChild
		normText += normTextChild
	}
	if n.Type == html.TextNode {
		count := 0
		for _, rune := range n.Data {
			if unicode.IsLetter(rune) {
				count += 1
			}
		}
		if insideLink {
			linkText += count
		} else {
			normText += count
		}
	}
	doc.linkText[n] = linkText
	doc.normText[n] = normText
	return
}

var removeElements = map[atom.Atom]bool{
	atom.Address:    true,
	atom.Audio:      true,
	atom.Button:     true,
	atom.Canvas:     true,
	atom.Caption:    true,
	atom.Fieldset:   true,
	atom.Figcaption: true,
	atom.Figure:     true,
	atom.Footer:     true,
	atom.Form:       true,
	atom.Frame:      true,
	atom.Iframe:     true,
	atom.Map:        true,
	atom.Menu:       true,
	atom.Nav:        true,
	atom.Noscript:   true,
	atom.Object:     true,
	atom.Option:     true,
	atom.Output:     true,
	atom.Script:     true,
	atom.Select:     true,
	atom.Style:      true,
	atom.Svg:        true,
	atom.Textarea:   true,
	atom.Video:      true,
}

// cleanBody removes unwanted HTML elements from the HTML body.
func (doc *Document) cleanBody(n *html.Node, level int) {
	// removeNode returns true if a node should be removed from HTML document.
	removeNode := func(c *html.Node, level int) bool {
		if c.DataAtom == atom.Table {
			return level > 5
		}
		return removeElements[c.DataAtom]
	}

	var curr *html.Node = n.FirstChild
	var next *html.Node = nil
	for ; curr != nil; curr = next {
		// We have to remember the next sibling here because calling RemoveChild
		// sets curr's NextSibling pointer to nil and we would quit the loop
		// prematurely.
		next = curr.NextSibling
		if curr.Type == html.ElementNode {
			if removeNode(curr, level) {
				n.RemoveChild(curr)
			} else {
				doc.cleanBody(curr, level+1)
			}
		}
	}
}

var (
	ignoreNames = util.NewRegexFromWords(
		"breadcrumb",
		"byline",
		"caption",
		"comment",
		"community",
		"credit",
		"description",
		"email",
		"footer",
		"gallery",
		"hide",
		"infotext",
		"photo",
		"related",
		"shares",
		"social",
		"story[-_]?bar",
		"story[-_]?feature",
	)
	ignoreStyle = util.NewRegex(`(?i)display:\s*none`)
)

// parseBody parses the ... part of the HTML page. It creates
// Chunks for every html.TextNode found in the body.
func (doc *Document) parseBody(n *html.Node) {
	switch n.Type {
	case html.ElementNode:
		// We ignore the node if it has some nasty classes/ids/itemprops or if
		// its style attribute contains "display: none".
		for _, attr := range n.Attr {
			switch attr.Key {
			case "id", "class", "itemprop":
				if ignoreNames.In(attr.Val) {
					return
				}
			case "style":
				if ignoreStyle.In(attr.Val) {
					return
				}
			}
		}
		ancestorMask := 0
		switch n.DataAtom {
		// We convert headings and links to text immediately. This is easier
		// and feasible because headings and links don't contain many children.
		// Descending into these children and handling every TextNode separately
		// would make things unnecessary complicated and our results noisy.
		case atom.H1, atom.H2, atom.H3, atom.H4, atom.H5, atom.H6, atom.A:
			if chunk, err := NewChunk(doc, n); err == nil {
				doc.Chunks = append(doc.Chunks, chunk)
			}
			return
		// Now mask the element type, but only if it isn't already set.
		// If we mask a bit which was already set by one of our callers, we'd also
		// clear it at the end of this function, though it actually should be cleared
		// by the caller.
		case atom.Article:
			ancestorMask = AncestorArticle &^ doc.ancestors
		case atom.Aside:
			ancestorMask = AncestorAside &^ doc.ancestors
		case atom.Blockquote:
			ancestorMask = AncestorBlockquote &^ doc.ancestors
		case atom.Ul, atom.Ol:
			ancestorMask = AncestorList &^ doc.ancestors
		}
		// Add our mask to the ancestor bitmask.
		doc.ancestors |= ancestorMask
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			doc.parseBody(c)
		}
		// Remove our mask from the ancestor bitmask.
		doc.ancestors &^= ancestorMask
	case html.TextNode:
		if chunk, err := NewChunk(doc, n); err == nil {
			doc.Chunks = append(doc.Chunks, chunk)
		}
	}
}

// TextStat contains the number of words and sentences found in text.
type TextStat struct {
	Words     int // total number of words
	Sentences int // total number of sentences
	Count     int // number of texts used to calculate this stats
}

// GetClassStats groups the document chunks by their classes (defined by the
// class attribute of HTML nodes) and calculates TextStats for each class.
func (doc *Document) GetClassStats() map[string]*TextStat {
	result := make(map[string]*TextStat)
	for _, chunk := range doc.Chunks {
		for _, class := range chunk.Classes {
			if stat, ok := result[class]; ok {
				stat.Words += chunk.Text.Words
				stat.Sentences += chunk.Text.Sentences
				stat.Count += 1
			} else {
				result[class] = &TextStat{chunk.Text.Words, chunk.Text.Sentences, 1}
			}
		}
	}
	return result
}

// GetClusterStats groups the document chunks by common ancestors and
// calculates TextStats for each group of chunks.
func (doc *Document) GetClusterStats() map[*Chunk]*TextStat {
	// Don't ascend further than this.
	const maxAncestors = 3

	// Count TextStats for Chunk ancestors.
	ancestorStat := make(map[*html.Node]*TextStat)
	for _, chunk := range doc.Chunks {
		node, count := chunk.Block, 0
		for node != nil && count < maxAncestors {
			if stat, ok := ancestorStat[node]; ok {
				stat.Words += chunk.Text.Words
				stat.Sentences += chunk.Text.Sentences
				stat.Count += 1
			} else {
				ancestorStat[node] = &TextStat{chunk.Text.Words, chunk.Text.Sentences, 1}
			}
			node, count = node.Parent, count+1
		}
	}

	// Generate result. For each chunk pick the best TextStats from its ancestors.
	result := make(map[*Chunk]*TextStat)
	for _, chunk := range doc.Chunks {
		node := chunk.Block
		if node == nil {
			continue
		}
		// Start with the parent's TextStat. Then ascend and check if the
		// current chunk has an ancestor with better stats. Use the best stat
		// as result.
		stat := ancestorStat[node]
		for {
			if node = node.Parent; node == nil {
				break
			}
			if statPrev, ok := ancestorStat[node]; ok {
				if stat.Count < statPrev.Count {
					stat = statPrev
				}
			} else {
				break
			}
		}
		result[chunk] = stat
	}
	return result
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"os"
	"runtime"
	"time"
)

var (
	counter     int
	userlist    map[int]*User
	chanlist    map[string]*Channel
	maxUsers    int
	maxRoutines int
	epoch       time.Time
	listeners   []net.Listener
)

func main() {
	epoch = time.Now()
	SetupNumerics()
	userlist = make(map[int]*User)
	chanlist = make(map[string]*Channel)
	// Listen for incoming connections.
	for _, LISTENING_IP := range listen_ips {
		for _, LISTENING_PORT := range listen_ports {
			l, err := net.Listen("tcp", LISTENING_IP+":"+LISTENING_PORT)
			if err != nil {
				log.Println("Error listening:", err.Error())
				os.Exit(1)
			} else {
				listeners = append(listeners, l)
				log.Println("Listening on " + LISTENING_IP + ":" + LISTENING_PORT)
			}
		}
	}
	// Close the listener when the application closes.
	for _, l := range listeners {
		defer l.Close()
	}
	go periodicStatusUpdate()
	for _, l := range listeners {
		go listenerthing(l)
	}
}

func listenerthing(l net.Listener) {
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println("Error accepting: ", err.Error())
		} else {
			user := NewUser()
			user.SetConn(conn)
			go user.HandleRequests()
		}
	}
}

func checkMaxUsers() {
	if len(userlist) > maxUsers {
		maxUsers = len(userlist)
	}
}

//periodicStatusUpdate shows information about the ircd every 5 seconds or so,
//as well as updating the max users, and goroutines numbers. Since these are
//only ran every 5 seconds or so, it may not be 100% accurate, but who cares
func periodicStatusUpdate() {
	for {
		checkMaxUsers()
		gor := runtime.NumGoroutine()
		if gor > maxRoutines {
			maxRoutines = gor
		}
		log.Println("Status:", len(userlist), "current users")
		log.Println("Status:", len(chanlist), "current channels")
		log.Println("Status:", gor, "current Goroutines")
		log.Println("Status:", maxRoutines, "max Goroutines")
		time.Sleep(stattime * time.Second)

	}
}
<commit_msg>user periodic status update to keep server from exiting<commit_after>package main

import (
	"log"
	"net"
	"os"
	"runtime"
	"time"
)

var (
	counter     int
	userlist    map[int]*User
	chanlist    map[string]*Channel
	maxUsers    int
	maxRoutines int
	epoch       time.Time
	listeners   []net.Listener
)

func main() {
	epoch = time.Now()
	SetupNumerics()
	userlist = make(map[int]*User)
	chanlist = make(map[string]*Channel)
	// Listen for incoming connections.
	for _, LISTENING_IP := range listen_ips {
		for _, LISTENING_PORT := range listen_ports {
			l, err := net.Listen("tcp", LISTENING_IP+":"+LISTENING_PORT)
			if err != nil {
				log.Println("Error listening:", err.Error())
				os.Exit(1)
			} else {
				listeners = append(listeners, l)
				log.Println("Listening on " + LISTENING_IP + ":" + LISTENING_PORT)
			}
		}
	}
	// Close the listener when the application closes.
	for _, l := range listeners {
		defer l.Close()
	}
	for _, l := range listeners {
		go listenerthing(l)
	}
	periodicStatusUpdate()
}

func listenerthing(l net.Listener) {
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println("Error accepting: ", err.Error())
		} else {
			user := NewUser()
			user.SetConn(conn)
			go user.HandleRequests()
		}
	}
}

func checkMaxUsers() {
	if len(userlist) > maxUsers {
		maxUsers = len(userlist)
	}
}

//periodicStatusUpdate shows information about the ircd every 5 seconds or so,
//as well as updating the max users, and goroutines numbers. Since these are
//only ran every 5 seconds or so, it may not be 100% accurate, but who cares
func periodicStatusUpdate() {
	for {
		checkMaxUsers()
		gor := runtime.NumGoroutine()
		if gor > maxRoutines {
			maxRoutines = gor
		}
		log.Println("Status:", len(userlist), "current users")
		log.Println("Status:", len(chanlist), "current channels")
		log.Println("Status:", gor, "current Goroutines")
		log.Println("Status:", maxRoutines, "max Goroutines")
		time.Sleep(stattime * time.Second)

	}
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	defaultRefreshRate = 150 * time.Millisecond
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	shutdown     chan struct{}
	cancel       func()
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	refreshRate        time.Duration
	idCount            int
	reqWidth           int
	popPriority        int
	popCompleted       bool
	outputDiscarded    bool
	disableAutoRefresh bool
	manualRefresh      chan interface{}
	renderDelay        <-chan struct{}
	shutdownNotifier   chan struct{}
	queueBars          map[*Bar]*Bar
	output             io.Writer
	debugOut           io.Writer
	uwg                *sync.WaitGroup
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		rows:          make([]io.Reader, 0, 64),
		pool:          make([]*Bar, 0, 64),
		refreshRate:   defaultRefreshRate,
		popPriority:   math.MinInt32,
		manualRefresh: make(chan interface{}),
		queueBars:     make(map[*Bar]*Bar),
		output:        os.Stdout,
		debugOut:      io.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	ctx, cancel := context.WithCancel(ctx)
	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
		cancel:       cancel,
	}

	if s.shutdownNotifier != nil {
		p.shutdown = s.shutdownNotifier
		s.shutdownNotifier = nil
	} else {
		p.shutdown = make(chan struct{})
	}

	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar by calling `Build` method on provided `BarFillerBuilder`.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.AddFiller(total, builder.Build(), options...)
}

// AddFiller creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) AddFiller(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediately, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container. After
// this method has been called, there is no way to reuse (*Progress) instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.bwg.Wait()
	p.Shutdown()
}

// Shutdown cancels any running bar immediately and then shutdowns (*Progress)
// instance. Normally this method shouldn't be called unless you know what you
// are doing. Proper way to shutdown is to call (*Progress).Wait() instead.
func (p *Progress) Shutdown() {
	p.cancel()
	<-p.shutdown
}

func (p *Progress) newTicker(s *pState) chan time.Time {
	ch := make(chan time.Time)
	go func() {
		var autoRefresh <-chan time.Time
		if !s.disableAutoRefresh && !s.outputDiscarded {
			if s.renderDelay != nil {
				<-s.renderDelay
			}
			ticker := time.NewTicker(s.refreshRate)
			defer ticker.Stop()
			autoRefresh = ticker.C
		}
		for {
			select {
			case t := <-autoRefresh:
				ch <- t
			case x := <-s.manualRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-p.ctx.Done():
				close(p.done)
				return
			}
		}
	}()
	return ch
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer close(p.shutdown)

	render := func() error {
		if s.bHeap.Len() == 0 {
			return nil
		}
		return s.render(cw)
	}

	refreshCh := p.newTicker(s)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-refreshCh:
			err := render()
			if err != nil {
				s.heapUpdated = false
				render = func() error { return nil }
				_, _ = fmt.Fprintln(s.debugOut, err.Error())
				p.cancel() // cancel all bars
			}
		case <-p.done:
			for s.heapUpdated {
				err := render()
				if err != nil {
					_, _ = fmt.Fprintln(s.debugOut, err.Error())
					return
				}
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	var wg sync.WaitGroup
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(&wg, s.pMatrix)
	syncWidth(&wg, s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	err = s.flush(&wg, cw, height)
	wg.Wait()
	return err
}

func (s *pState) flush(wg *sync.WaitGroup, cw *cwriter.Writer, height int) error {
	var popCount int

	for s.bHeap.Len() > 0 {
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.err != nil {
			return frame.err
		}
		var usedRows int
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				s.heapUpdated = true
				continue
			}
			if s.popCompleted && !b.bs.noPop {
				switch b.bs.shutdown++; b.bs.shutdown {
				case 1:
					b.priority = s.popPriority
					s.popPriority++
				default:
					if b.bs.dropOnComplete {
						popCount += usedRows
						s.heapUpdated = true
						continue
					}
				}
			} else if b.bs.dropOnComplete {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	switch len(s.pool) {
	case 0:
		if s.heapUpdated {
			s.updateSyncMatrix()
			s.heapUpdated = false
		}
	case 1:
		heap.Push(&s.bHeap, s.pool[0])
		s.pool = s.pool[:0]
	default:
		wg.Add(1)
		go func() {
			for _, b := range s.pool {
				heap.Push(&s.bHeap, b)
			}
			s.pool = s.pool[:0]
			wg.Done()
		}()
	}

	for i := len(s.rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			return err
		}
	}

	err := cw.Flush(len(s.rows) - popCount)
	s.rows = s.rows[:0]
	return err
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:            s.idCount,
		priority:      s.idCount,
		reqWidth:      s.reqWidth,
		total:         total,
		filler:        filler,
		manualRefresh: s.manualRefresh,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(wg *sync.WaitGroup, matrix map[int][]chan int) {
	for _, column := range matrix {
		wg.Add(1)
		go maxWidthDistributor(wg, column)
	}
}

func maxWidthDistributor(wg *sync.WaitGroup, column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
	wg.Done()
}
<commit_msg>comment why early return is ok<commit_after>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	defaultRefreshRate = 150 * time.Millisecond
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	shutdown     chan struct{}
	cancel       func()
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	refreshRate        time.Duration
	idCount            int
	reqWidth           int
	popPriority        int
	popCompleted       bool
	outputDiscarded    bool
	disableAutoRefresh bool
	manualRefresh      chan interface{}
	renderDelay        <-chan struct{}
	shutdownNotifier   chan struct{}
	queueBars          map[*Bar]*Bar
	output             io.Writer
	debugOut           io.Writer
	uwg                *sync.WaitGroup
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		rows:          make([]io.Reader, 0, 64),
		pool:          make([]*Bar, 0, 64),
		refreshRate:   defaultRefreshRate,
		popPriority:   math.MinInt32,
		manualRefresh: make(chan interface{}),
		queueBars:     make(map[*Bar]*Bar),
		output:        os.Stdout,
		debugOut:      io.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	ctx, cancel := context.WithCancel(ctx)
	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
		cancel:       cancel,
	}

	if s.shutdownNotifier != nil {
		p.shutdown = s.shutdownNotifier
		s.shutdownNotifier = nil
	} else {
		p.shutdown = make(chan struct{})
	}

	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar by calling `Build` method on provided `BarFillerBuilder`.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.AddFiller(total, builder.Build(), options...)
}

// AddFiller creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) AddFiller(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediately, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container. After
// this method has been called, there is no way to reuse (*Progress) instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.bwg.Wait()
	p.Shutdown()
}

// Shutdown cancels any running bar immediately and then shutdowns (*Progress)
// instance. Normally this method shouldn't be called unless you know what you
// are doing. Proper way to shutdown is to call (*Progress).Wait() instead.
func (p *Progress) Shutdown() {
	p.cancel()
	<-p.shutdown
}

func (p *Progress) newTicker(s *pState) chan time.Time {
	ch := make(chan time.Time)
	go func() {
		var autoRefresh <-chan time.Time
		if !s.disableAutoRefresh && !s.outputDiscarded {
			if s.renderDelay != nil {
				<-s.renderDelay
			}
			ticker := time.NewTicker(s.refreshRate)
			defer ticker.Stop()
			autoRefresh = ticker.C
		}
		for {
			select {
			case t := <-autoRefresh:
				ch <- t
			case x := <-s.manualRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-p.ctx.Done():
				close(p.done)
				return
			}
		}
	}()
	return ch
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer close(p.shutdown)

	render := func() error {
		if s.bHeap.Len() == 0 {
			return nil
		}
		return s.render(cw)
	}

	refreshCh := p.newTicker(s)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-refreshCh:
			err := render()
			if err != nil {
				s.heapUpdated = false
				render = func() error { return nil }
				_, _ = fmt.Fprintln(s.debugOut, err.Error())
				p.cancel() // cancel all bars
			}
		case <-p.done:
			for s.heapUpdated {
				err := render()
				if err != nil {
					_, _ = fmt.Fprintln(s.debugOut, err.Error())
					return
				}
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	var wg sync.WaitGroup
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(&wg, s.pMatrix)
	syncWidth(&wg, s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	err = s.flush(&wg, cw, height)
	wg.Wait()
	return err
}

func (s *pState) flush(wg *sync.WaitGroup, cw *cwriter.Writer, height int) error {
	var popCount int

	for s.bHeap.Len() > 0 {
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.err != nil {
			// b.frameCh is buffered it's ok to return here
			return frame.err
		}
		var usedRows int
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				s.heapUpdated = true
				continue
			}
			if s.popCompleted && !b.bs.noPop {
				switch b.bs.shutdown++; b.bs.shutdown {
				case 1:
					b.priority = s.popPriority
					s.popPriority++
				default:
					if b.bs.dropOnComplete {
						popCount += usedRows
						s.heapUpdated = true
						continue
					}
				}
			} else if b.bs.dropOnComplete {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	switch len(s.pool) {
	case 0:
		if s.heapUpdated {
			s.updateSyncMatrix()
			s.heapUpdated = false
		}
	case 1:
		heap.Push(&s.bHeap, s.pool[0])
		s.pool = s.pool[:0]
	default:
		wg.Add(1)
		go func() {
			for _, b := range s.pool {
				heap.Push(&s.bHeap, b)
			}
			s.pool = s.pool[:0]
			wg.Done()
		}()
	}

	for i := len(s.rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			return err
		}
	}

	err := cw.Flush(len(s.rows) - popCount)
	s.rows = s.rows[:0]
	return err
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:            s.idCount,
		priority:      s.idCount,
		reqWidth:      s.reqWidth,
		total:         total,
		filler:        filler,
		manualRefresh: s.manualRefresh,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(wg *sync.WaitGroup, matrix map[int][]chan int) {
	for _, column := range matrix {
		wg.Add(1)
		go maxWidthDistributor(wg, column)
	}
}

func maxWidthDistributor(wg *sync.WaitGroup, column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
	wg.Done()
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All rights reserved.
//
// 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 blueprint

import (
	"crypto/md5"
	"fmt"
	"sort"
	"strings"

	"github.com/google/blueprint/pathtools"
)

type GlobPath struct {
	Pattern  string
	Excludes []string
	Files    []string
	Deps     []string
	Name     string
}

func verifyGlob(fileName, pattern string, excludes []string, g GlobPath) {
	if pattern != g.Pattern {
		panic(fmt.Errorf("Mismatched patterns %q and %q for glob file %q", pattern, g.Pattern, fileName))
	}
	if len(excludes) != len(g.Excludes) {
		panic(fmt.Errorf("Mismatched excludes %v and %v for glob file %q", excludes, g.Excludes, fileName))
	}

	for i := range excludes {
		if g.Excludes[i] != excludes[i] {
			panic(fmt.Errorf("Mismatched excludes %v and %v for glob file %q", excludes, g.Excludes, fileName))
		}
	}
}

func (c *Context) glob(pattern string, excludes []string) ([]string, error) {
	fileName := globToFileName(pattern, excludes)

	// Try to get existing glob from the stored results
	c.globLock.Lock()
	g, exists := c.globs[fileName]
	c.globLock.Unlock()

	if exists {
		// Glob has already been done, double check it is identical
		verifyGlob(fileName, pattern, excludes, g)
		return g.Files, nil
	}

	// Get a globbed file list
	files, deps, err := c.fs.Glob(pattern, excludes, pathtools.FollowSymlinks)
	if err != nil {
		return nil, err
	}

	// Store the results
	c.globLock.Lock()
	if g, exists = c.globs[fileName]; !exists {
		c.globs[fileName] = GlobPath{pattern, excludes, files, deps, fileName}
	}
	c.globLock.Unlock()

	// Getting the list raced with another goroutine, throw away the results and use theirs
	if exists {
		verifyGlob(fileName, pattern, excludes, g)
		return g.Files, nil
	}

	return files, nil
}

func (c *Context) Globs() []GlobPath {
	fileNames := make([]string, 0, len(c.globs))
	for k := range c.globs {
		fileNames = append(fileNames, k)
	}
	sort.Strings(fileNames)

	globs := make([]GlobPath, len(fileNames))
	for i, fileName := range fileNames {
		globs[i] = c.globs[fileName]
	}

	return globs
}

func globToString(pattern string) string {
	ret := ""
	for _, c := range pattern {
		switch {
		case c >= 'a' && c <= 'z',
			c >= 'A' && c <= 'Z',
			c >= '0' && c <= '9',
			c == '_', c == '-', c == '/':
			ret += string(c)
		default:
			ret += "_"
		}
	}

	return ret
}

func globToFileName(pattern string, excludes []string) string {
	name := globToString(pattern)
	excludeName := ""
	for _, e := range excludes {
		excludeName += "__" + globToString(e)
	}

	// Prevent file names from reaching ninja's path component limit
	if strings.Count(name, "/")+strings.Count(excludeName, "/") > 30 {
		excludeName = fmt.Sprintf("___%x", md5.Sum([]byte(excludeName)))
	}

	return name + excludeName + ".glob"
}
<commit_msg>Return a copy of glob lists<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// 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 blueprint

import (
	"crypto/md5"
	"fmt"
	"sort"
	"strings"

	"github.com/google/blueprint/pathtools"
)

type GlobPath struct {
	Pattern  string
	Excludes []string
	Files    []string
	Deps     []string
	Name     string
}

func verifyGlob(fileName, pattern string, excludes []string, g GlobPath) {
	if pattern != g.Pattern {
		panic(fmt.Errorf("Mismatched patterns %q and %q for glob file %q", pattern, g.Pattern, fileName))
	}
	if len(excludes) != len(g.Excludes) {
		panic(fmt.Errorf("Mismatched excludes %v and %v for glob file %q", excludes, g.Excludes, fileName))
	}

	for i := range excludes {
		if g.Excludes[i] != excludes[i] {
			panic(fmt.Errorf("Mismatched excludes %v and %v for glob file %q", excludes, g.Excludes, fileName))
		}
	}
}

func (c *Context) glob(pattern string, excludes []string) ([]string, error) {
	fileName := globToFileName(pattern, excludes)

	// Try to get existing glob from the stored results
	c.globLock.Lock()
	g, exists := c.globs[fileName]
	c.globLock.Unlock()

	if exists {
		// Glob has already been done, double check it is identical
		verifyGlob(fileName, pattern, excludes, g)
		// Return a copy so that modifications don't affect the cached value.
		return append([]string(nil), g.Files...), nil
	}

	// Get a globbed file list
	files, deps, err := c.fs.Glob(pattern, excludes, pathtools.FollowSymlinks)
	if err != nil {
		return nil, err
	}

	// Store the results
	c.globLock.Lock()
	if g, exists = c.globs[fileName]; !exists {
		c.globs[fileName] = GlobPath{pattern, excludes, files, deps, fileName}
	}
	c.globLock.Unlock()

	// Getting the list raced with another goroutine, throw away the results and use theirs
	if exists {
		verifyGlob(fileName, pattern, excludes, g)
		// Return a copy so that modifications don't affect the cached value.
		return append([]string(nil), g.Files...), nil
	}

	// Return a copy so that modifications don't affect the cached value.
	return append([]string(nil), files...), nil
}

func (c *Context) Globs() []GlobPath {
	fileNames := make([]string, 0, len(c.globs))
	for k := range c.globs {
		fileNames = append(fileNames, k)
	}
	sort.Strings(fileNames)

	globs := make([]GlobPath, len(fileNames))
	for i, fileName := range fileNames {
		globs[i] = c.globs[fileName]
	}

	return globs
}

func globToString(pattern string) string {
	ret := ""
	for _, c := range pattern {
		switch {
		case c >= 'a' && c <= 'z',
			c >= 'A' && c <= 'Z',
			c >= '0' && c <= '9',
			c == '_', c == '-', c == '/':
			ret += string(c)
		default:
			ret += "_"
		}
	}

	return ret
}

func globToFileName(pattern string, excludes []string) string {
	name := globToString(pattern)
	excludeName := ""
	for _, e := range excludes {
		excludeName += "__" + globToString(e)
	}

	// Prevent file names from reaching ninja's path component limit
	if strings.Count(name, "/")+strings.Count(excludeName, "/") > 30 {
		excludeName = fmt.Sprintf("___%x", md5.Sum([]byte(excludeName)))
	}

	return name + excludeName + ".glob"
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors.

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 remotecommand

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	v1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/httpstream"
	"k8s.io/apimachinery/pkg/util/httpstream/spdy"
	remotecommandconsts "k8s.io/apimachinery/pkg/util/remotecommand"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
)

type AttachFunc func(in io.Reader, out, err io.WriteCloser, tty bool, resize <-chan TerminalSize) error
type streamContext struct {
	conn         io.Closer
	stdinStream  io.ReadCloser
	stdoutStream io.WriteCloser
	stderrStream io.WriteCloser
	writeStatus  func(status *apierrors.StatusError) error
}

type streamAndReply struct {
	httpstream.Stream
	replySent <-chan struct{}
}

type fakeEmptyDataPty struct {
}

func (s *fakeEmptyDataPty) Read(p []byte) (int, error) {
	return len(p), nil
}

func (s *fakeEmptyDataPty) Write(p []byte) (int, error) {
	return len(p), nil
}

type fakeMassiveDataPty struct{}

func (s *fakeMassiveDataPty) Read(p []byte) (int, error) {
	time.Sleep(time.Duration(1) * time.Second)
	return copy(p, []byte{}), errors.New("client crashed after 1 second")
}

func (s *fakeMassiveDataPty) Write(p []byte) (int, error) {
	time.Sleep(time.Duration(1) * time.Second)
	return len(p), errors.New("return err")
}

func fakeMassiveDataAttacher(stdin io.Reader, stdout, stderr io.WriteCloser, tty bool, resize <-chan TerminalSize) error {

	copyDone := make(chan struct{}, 3)

	if stdin == nil {
		return errors.New("stdin is requested") // we need stdin to notice the conn break
	}

	go func() {
		io.Copy(io.Discard, stdin)
		copyDone <- struct{}{}
	}()

	go func() {
		if stdout == nil {
			return
		}
		copyDone <- writeMassiveData(stdout)
	}()

	go func() {
		if stderr == nil {
			return
		}
		copyDone <- writeMassiveData(stderr)
	}()

	select {
	case <-copyDone:
		return nil
	}
}

func writeMassiveData(stdStream io.Writer) struct{} { // write to stdin or stdout
	for {
		_, err := io.Copy(stdStream, strings.NewReader("something"))
		if err != nil && err.Error() != "EOF" {
			break
		}
	}
	return struct{}{}
}

func TestSPDYExecutorStream(t *testing.T) {
	tests := []struct {
		timeout     time.Duration
		name        string
		options     StreamOptions
		expectError string
		attacher    AttachFunc
	}{
		{
			name: "stdoutBlockTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stdout: &fakeMassiveDataPty{},
			},
			expectError: "",
			attacher:    fakeMassiveDataAttacher,
		},
		{
			name: "stderrBlockTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stderr: &fakeMassiveDataPty{},
			},
			expectError: "",
			attacher:    fakeMassiveDataAttacher,
		},
		{
			timeout: 500 * time.Millisecond,
			name:    "timeoutTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stderr: &fakeMassiveDataPty{},
			},
			expectError: context.DeadlineExceeded.Error(),
			attacher:    fakeMassiveDataAttacher,
		},
	}

	for _, test := range tests {
		server := newTestHTTPServer(test.attacher, &test.options)

		ctx, cancelFn := context.Background(), func() {}
		if test.timeout > 0 {
			ctx, cancelFn = context.WithTimeout(ctx, test.timeout)
		}

		err := func(ctx context.Context, cancel context.CancelFunc) error {
			defer cancelFn()
			return attach2Server(ctx, server.URL, test.options)
		}(ctx, cancelFn)

		gotError := ""
		if err != nil {
			gotError = err.Error()
		}
		if test.expectError != gotError {
			t.Errorf("%s: expected [%v], got [%v]", test.name, test.expectError, gotError)
		}

		server.Close()
	}

}

func newTestHTTPServer(f AttachFunc, options *StreamOptions) *httptest.Server {
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		ctx, err := createHTTPStreams(writer, request, options)
		if err != nil {
			return
		}
		defer ctx.conn.Close()

		// handle input output
		err = f(ctx.stdinStream, ctx.stdoutStream, ctx.stderrStream, false, nil)
		if err != nil {
			ctx.writeStatus(apierrors.NewInternalError(err))
		} else {
			ctx.writeStatus(&apierrors.StatusError{ErrStatus: metav1.Status{
				Status: metav1.StatusSuccess,
			}})
		}
	}))
	return server
}

func attach2Server(ctx context.Context, rawURL string, options StreamOptions) error {
	uri, _ := url.Parse(rawURL)
	exec, err := NewSPDYExecutor(&rest.Config{Host: uri.Host}, "POST", uri)
	if err != nil {
		return err
	}

	e := make(chan error, 1)
	go func(e chan error) {
		e <- exec.StreamWithContext(ctx, options)
	}(e)
	select {
	case err := <-e:
		return err
	case <-time.After(wait.ForeverTestTimeout):
		return errors.New("execute timeout")
	}
}

// simplify createHttpStreams , only support StreamProtocolV4Name
func createHTTPStreams(w http.ResponseWriter, req *http.Request, opts *StreamOptions) (*streamContext, error) {
	_, err := httpstream.Handshake(req, w, []string{remotecommandconsts.StreamProtocolV4Name})
	if err != nil {
		return nil, err
	}

	upgrader := spdy.NewResponseUpgrader()
	streamCh := make(chan streamAndReply)
	conn := upgrader.UpgradeResponse(w, req, func(stream httpstream.Stream, replySent <-chan struct{}) error {
		streamCh <- streamAndReply{Stream: stream, replySent: replySent}
		return nil
	})
	ctx := &streamContext{
		conn: conn,
	}

	// wait for stream
	replyChan := make(chan struct{}, 4)
	defer close(replyChan)
	receivedStreams := 0
	expectedStreams := 1
	if opts.Stdout != nil {
		expectedStreams++
	}
	if opts.Stdin != nil {
		expectedStreams++
	}
	if opts.Stderr != nil {
		expectedStreams++
	}
WaitForStreams:
	for {
		select {
		case stream := <-streamCh:
			streamType := stream.Headers().Get(v1.StreamType)
			switch streamType {
			case v1.StreamTypeError:
				replyChan <- struct{}{}
				ctx.writeStatus = v4WriteStatusFunc(stream)
			case v1.StreamTypeStdout:
				replyChan <- struct{}{}
				ctx.stdoutStream = stream
			case v1.StreamTypeStdin:
				replyChan <- struct{}{}
				ctx.stdinStream = stream
			case v1.StreamTypeStderr:
				replyChan <- struct{}{}
				ctx.stderrStream = stream
			default:
				// add other stream ...
				return nil, errors.New("unimplemented stream type")
			}
		case <-replyChan:
			receivedStreams++
			if receivedStreams == expectedStreams {
				break WaitForStreams
			}
		}
	}

	return ctx, nil
}

func v4WriteStatusFunc(stream io.Writer) func(status *apierrors.StatusError) error {
	return func(status *apierrors.StatusError) error {
		bs, err := json.Marshal(status.Status())
		if err != nil {
			return err
		}
		_, err = stream.Write(bs)
		return err
	}
}

// writeDetector provides a helper method to block until the underlying writer written.
type writeDetector struct {
	written chan bool
	closed  bool
	io.Writer
}

func newWriterDetector(w io.Writer) *writeDetector {
	return &writeDetector{
		written: make(chan bool),
		Writer:  w,
	}
}

func (w *writeDetector) BlockUntilWritten() {
	<-w.written
}

func (w *writeDetector) Write(p []byte) (n int, err error) {
	if !w.closed {
		close(w.written)
		w.closed = true
	}
	return w.Writer.Write(p)
}

// `Executor.StreamWithContext` starts a goroutine in the background to do the streaming
// and expects the deferred close of the connection leads to the exit of the goroutine on cancellation.
// This test verifies that works.
func TestStreamExitsAfterConnectionIsClosed(t *testing.T) {
	writeDetector := newWriterDetector(&fakeEmptyDataPty{})
	options := StreamOptions{
		Stdin:  &fakeEmptyDataPty{},
		Stdout: writeDetector,
	}
	server := newTestHTTPServer(fakeMassiveDataAttacher, &options)

	ctx, cancelFn := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancelFn()

	uri, _ := url.Parse(server.URL)
	exec, err := NewSPDYExecutor(&rest.Config{Host: uri.Host}, "POST", uri)
	if err != nil {
		t.Fatal(err)
	}
	streamExec := exec.(*streamExecutor)

	conn, streamer, err := streamExec.newConnectionAndStream(ctx, options)
	if err != nil {
		t.Fatal(err)
	}

	errorChan := make(chan error)
	go func() {
		errorChan <- streamer.stream(conn)
	}()

	// Wait until stream goroutine starts.
	writeDetector.BlockUntilWritten()

	// Close the connection
	conn.Close()

	select {
	case <-time.After(1 * time.Second):
		t.Fatalf("expect stream to be closed after connection is closed.")
	case <-errorChan:
		return
	}
}
<commit_msg>use subtests and defer in TestSPDYExecutorStream<commit_after>/*
Copyright 2020 The Kubernetes Authors.

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 remotecommand

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	v1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/httpstream"
	"k8s.io/apimachinery/pkg/util/httpstream/spdy"
	remotecommandconsts "k8s.io/apimachinery/pkg/util/remotecommand"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
)

type AttachFunc func(in io.Reader, out, err io.WriteCloser, tty bool, resize <-chan TerminalSize) error
type streamContext struct {
	conn         io.Closer
	stdinStream  io.ReadCloser
	stdoutStream io.WriteCloser
	stderrStream io.WriteCloser
	writeStatus  func(status *apierrors.StatusError) error
}

type streamAndReply struct {
	httpstream.Stream
	replySent <-chan struct{}
}

type fakeEmptyDataPty struct {
}

func (s *fakeEmptyDataPty) Read(p []byte) (int, error) {
	return len(p), nil
}

func (s *fakeEmptyDataPty) Write(p []byte) (int, error) {
	return len(p), nil
}

type fakeMassiveDataPty struct{}

func (s *fakeMassiveDataPty) Read(p []byte) (int, error) {
	time.Sleep(time.Duration(1) * time.Second)
	return copy(p, []byte{}), errors.New("client crashed after 1 second")
}

func (s *fakeMassiveDataPty) Write(p []byte) (int, error) {
	time.Sleep(time.Duration(1) * time.Second)
	return len(p), errors.New("return err")
}

func fakeMassiveDataAttacher(stdin io.Reader, stdout, stderr io.WriteCloser, tty bool, resize <-chan TerminalSize) error {

	copyDone := make(chan struct{}, 3)

	if stdin == nil {
		return errors.New("stdin is requested") // we need stdin to notice the conn break
	}

	go func() {
		io.Copy(io.Discard, stdin)
		copyDone <- struct{}{}
	}()

	go func() {
		if stdout == nil {
			return
		}
		copyDone <- writeMassiveData(stdout)
	}()

	go func() {
		if stderr == nil {
			return
		}
		copyDone <- writeMassiveData(stderr)
	}()

	select {
	case <-copyDone:
		return nil
	}
}

func writeMassiveData(stdStream io.Writer) struct{} { // write to stdin or stdout
	for {
		_, err := io.Copy(stdStream, strings.NewReader("something"))
		if err != nil && err.Error() != "EOF" {
			break
		}
	}
	return struct{}{}
}

func TestSPDYExecutorStream(t *testing.T) {
	tests := []struct {
		timeout     time.Duration
		name        string
		options     StreamOptions
		expectError string
		attacher    AttachFunc
	}{
		{
			name: "stdoutBlockTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stdout: &fakeMassiveDataPty{},
			},
			expectError: "",
			attacher:    fakeMassiveDataAttacher,
		},
		{
			name: "stderrBlockTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stderr: &fakeMassiveDataPty{},
			},
			expectError: "",
			attacher:    fakeMassiveDataAttacher,
		},
		{
			timeout: 500 * time.Millisecond,
			name:    "timeoutTest",
			options: StreamOptions{
				Stdin:  &fakeMassiveDataPty{},
				Stderr: &fakeMassiveDataPty{},
			},
			expectError: context.DeadlineExceeded.Error(),
			attacher:    fakeMassiveDataAttacher,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			server := newTestHTTPServer(test.attacher, &test.options)
			defer server.Close()

			ctx, cancel := context.Background(), func() {}
			if test.timeout > 0 {
				ctx, cancel = context.WithTimeout(ctx, test.timeout)
			}
			defer cancel()

			err := attach2Server(ctx, server.URL, test.options)

			gotError := ""
			if err != nil {
				gotError = err.Error()
			}
			if test.expectError != gotError {
				t.Errorf("%s: expected [%v], got [%v]", test.name, test.expectError, gotError)
			}
		})
	}
}

func newTestHTTPServer(f AttachFunc, options *StreamOptions) *httptest.Server {
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		ctx, err := createHTTPStreams(writer, request, options)
		if err != nil {
			return
		}
		defer ctx.conn.Close()

		// handle input output
		err = f(ctx.stdinStream, ctx.stdoutStream, ctx.stderrStream, false, nil)
		if err != nil {
			ctx.writeStatus(apierrors.NewInternalError(err))
		} else {
			ctx.writeStatus(&apierrors.StatusError{ErrStatus: metav1.Status{
				Status: metav1.StatusSuccess,
			}})
		}
	}))
	return server
}

func attach2Server(ctx context.Context, rawURL string, options StreamOptions) error {
	uri, _ := url.Parse(rawURL)
	exec, err := NewSPDYExecutor(&rest.Config{Host: uri.Host}, "POST", uri)
	if err != nil {
		return err
	}

	e := make(chan error, 1)
	go func(e chan error) {
		e <- exec.StreamWithContext(ctx, options)
	}(e)
	select {
	case err := <-e:
		return err
	case <-time.After(wait.ForeverTestTimeout):
		return errors.New("execute timeout")
	}
}

// simplify createHttpStreams , only support StreamProtocolV4Name
func createHTTPStreams(w http.ResponseWriter, req *http.Request, opts *StreamOptions) (*streamContext, error) {
	_, err := httpstream.Handshake(req, w, []string{remotecommandconsts.StreamProtocolV4Name})
	if err != nil {
		return nil, err
	}

	upgrader := spdy.NewResponseUpgrader()
	streamCh := make(chan streamAndReply)
	conn := upgrader.UpgradeResponse(w, req, func(stream httpstream.Stream, replySent <-chan struct{}) error {
		streamCh <- streamAndReply{Stream: stream, replySent: replySent}
		return nil
	})
	ctx := &streamContext{
		conn: conn,
	}

	// wait for stream
	replyChan := make(chan struct{}, 4)
	defer close(replyChan)
	receivedStreams := 0
	expectedStreams := 1
	if opts.Stdout != nil {
		expectedStreams++
	}
	if opts.Stdin != nil {
		expectedStreams++
	}
	if opts.Stderr != nil {
		expectedStreams++
	}
WaitForStreams:
	for {
		select {
		case stream := <-streamCh:
			streamType := stream.Headers().Get(v1.StreamType)
			switch streamType {
			case v1.StreamTypeError:
				replyChan <- struct{}{}
				ctx.writeStatus = v4WriteStatusFunc(stream)
			case v1.StreamTypeStdout:
				replyChan <- struct{}{}
				ctx.stdoutStream = stream
			case v1.StreamTypeStdin:
				replyChan <- struct{}{}
				ctx.stdinStream = stream
			case v1.StreamTypeStderr:
				replyChan <- struct{}{}
				ctx.stderrStream = stream
			default:
				// add other stream ...
				return nil, errors.New("unimplemented stream type")
			}
		case <-replyChan:
			receivedStreams++
			if receivedStreams == expectedStreams {
				break WaitForStreams
			}
		}
	}

	return ctx, nil
}

func v4WriteStatusFunc(stream io.Writer) func(status *apierrors.StatusError) error {
	return func(status *apierrors.StatusError) error {
		bs, err := json.Marshal(status.Status())
		if err != nil {
			return err
		}
		_, err = stream.Write(bs)
		return err
	}
}

// writeDetector provides a helper method to block until the underlying writer written.
type writeDetector struct {
	written chan bool
	closed  bool
	io.Writer
}

func newWriterDetector(w io.Writer) *writeDetector {
	return &writeDetector{
		written: make(chan bool),
		Writer:  w,
	}
}

func (w *writeDetector) BlockUntilWritten() {
	<-w.written
}

func (w *writeDetector) Write(p []byte) (n int, err error) {
	if !w.closed {
		close(w.written)
		w.closed = true
	}
	return w.Writer.Write(p)
}

// `Executor.StreamWithContext` starts a goroutine in the background to do the streaming
// and expects the deferred close of the connection leads to the exit of the goroutine on cancellation.
// This test verifies that works.
func TestStreamExitsAfterConnectionIsClosed(t *testing.T) {
	writeDetector := newWriterDetector(&fakeEmptyDataPty{})
	options := StreamOptions{
		Stdin:  &fakeEmptyDataPty{},
		Stdout: writeDetector,
	}
	server := newTestHTTPServer(fakeMassiveDataAttacher, &options)

	ctx, cancelFn := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancelFn()

	uri, _ := url.Parse(server.URL)
	exec, err := NewSPDYExecutor(&rest.Config{Host: uri.Host}, "POST", uri)
	if err != nil {
		t.Fatal(err)
	}
	streamExec := exec.(*streamExecutor)

	conn, streamer, err := streamExec.newConnectionAndStream(ctx, options)
	if err != nil {
		t.Fatal(err)
	}

	errorChan := make(chan error)
	go func() {
		errorChan <- streamer.stream(conn)
	}()

	// Wait until stream goroutine starts.
	writeDetector.BlockUntilWritten()

	// Close the connection
	conn.Close()

	select {
	case <-time.After(1 * time.Second):
		t.Fatalf("expect stream to be closed after connection is closed.")
	case <-errorChan:
		return
	}
}
<|endoftext|>


<commit_before>package video

import "errors"
import "github.com/32bitkid/bitreader"

var ErrMissingMarkerBit = errors.New("missing marker bit")

type SequenceHeader struct {
	horizontal_size_value       uint32
	vertical_size_value         uint32
	aspect_ratio_information    uint32
	frame_rate_code             uint32
	bit_rate_value              uint32
	vbv_buffer_size_value       uint32
	constrained_parameters_flag bool

	load_intra_quantiser_matrix     bool
	load_non_intra_quantiser_matrix bool

	intra_quantiser_matrix     quantisationMatrix
	non_intra_quantiser_matrix quantisationMatrix
}

func sequence_header(br bitreader.BitReader) (*SequenceHeader, error) {

	var err error

	err = SequenceHeaderStartCode.Assert(br)
	if err != nil {
		return nil, err
	}

	sh := SequenceHeader{}

	if sh.horizontal_size_value, err = br.Read32(12); err != nil {
		return nil, err
	}

	if sh.vertical_size_value, err = br.Read32(12); err != nil {
		return nil, err
	}

	if sh.aspect_ratio_information, err = br.Read32(4); err != nil {
		return nil, err
	}

	if sh.frame_rate_code, err = br.Read32(4); err != nil {
		return nil, err
	}

	if sh.bit_rate_value, err = br.Read32(18); err != nil {
		return nil, err
	}

	err = marker_bit(br)
	if err != nil {
		return nil, err
	}

	if sh.vbv_buffer_size_value, err = br.Read32(10); err != nil {
		return nil, err
	}

	if sh.constrained_parameters_flag, err = br.ReadBit(); err != nil {
		return nil, err
	}

	sh.load_intra_quantiser_matrix, err = br.ReadBit()
	if err != nil {
		return nil, err
	}
	if sh.load_intra_quantiser_matrix {
		for v := 0; v < 8; v++ {
			for u := 0; u < 8; u++ {
				if val, err := br.Read32(8); err != nil {
					return nil, err
				} else {
					sh.intra_quantiser_matrix[v][u] = uint8(val)
				}
			}
		}
	}

	sh.load_non_intra_quantiser_matrix, err = br.ReadBit()
	if err != nil {
		return nil, err
	}
	if sh.load_non_intra_quantiser_matrix {
		for v := 0; v < 8; v++ {
			for u := 0; u < 8; u++ {
				if val, err := br.Read32(8); err != nil {
					return nil, err
				} else {
					sh.non_intra_quantiser_matrix[v][u] = uint8(val)
				}
			}
		}
	}

	return &sh, next_start_code(br)
}

func (vs *VideoSequence) sequence_header() (err error) {

	sh, err := sequence_header(vs)
	if err != nil {
		return err
	}

	if sh.load_intra_quantiser_matrix {
		vs.quantisationMatricies[0] = sh.intra_quantiser_matrix
		vs.quantisationMatricies[2] = sh.intra_quantiser_matrix
	} else {
		vs.quantisationMatricies[0] = defaultQuantisationMatrices.Intra
		vs.quantisationMatricies[2] = defaultQuantisationMatrices.Intra
	}

	if sh.load_non_intra_quantiser_matrix {
		vs.quantisationMatricies[1] = sh.non_intra_quantiser_matrix
		vs.quantisationMatricies[3] = sh.non_intra_quantiser_matrix
	} else {
		vs.quantisationMatricies[1] = defaultQuantisationMatrices.NonIntra
		vs.quantisationMatricies[3] = defaultQuantisationMatrices.NonIntra
	}

	vs.SequenceHeader = sh

	return nil
}
<commit_msg>more documentation<commit_after>package video

import "errors"
import "github.com/32bitkid/bitreader"

var ErrMissingMarkerBit = errors.New("missing marker bit")

type SequenceHeader struct {
	horizontal_size_value       uint32
	vertical_size_value         uint32
	aspect_ratio_information    uint32
	frame_rate_code             uint32
	bit_rate_value              uint32
	vbv_buffer_size_value       uint32
	constrained_parameters_flag bool

	load_intra_quantiser_matrix     bool
	load_non_intra_quantiser_matrix bool

	intra_quantiser_matrix     quantisationMatrix
	non_intra_quantiser_matrix quantisationMatrix
}

// ReadSequenceHeader reads a sequence header from the bit stream.
func ReadSequenceHeader(br bitreader.BitReader) (*SequenceHeader, error) {

	var err error

	err = SequenceHeaderStartCode.Assert(br)
	if err != nil {
		return nil, err
	}

	sh := SequenceHeader{}

	if sh.horizontal_size_value, err = br.Read32(12); err != nil {
		return nil, err
	}

	if sh.vertical_size_value, err = br.Read32(12); err != nil {
		return nil, err
	}

	if sh.aspect_ratio_information, err = br.Read32(4); err != nil {
		return nil, err
	}

	if sh.frame_rate_code, err = br.Read32(4); err != nil {
		return nil, err
	}

	if sh.bit_rate_value, err = br.Read32(18); err != nil {
		return nil, err
	}

	if err := marker_bit(br); err != nil {
		return nil, err
	}

	if sh.vbv_buffer_size_value, err = br.Read32(10); err != nil {
		return nil, err
	}

	if sh.constrained_parameters_flag, err = br.ReadBit(); err != nil {
		return nil, err
	}

	sh.load_intra_quantiser_matrix, err = br.ReadBit()
	if err != nil {
		return nil, err
	}
	if sh.load_intra_quantiser_matrix {
		for v := 0; v < 8; v++ {
			for u := 0; u < 8; u++ {
				if val, err := br.Read32(8); err != nil {
					return nil, err
				} else {
					sh.intra_quantiser_matrix[v][u] = uint8(val)
				}
			}
		}
	}

	sh.load_non_intra_quantiser_matrix, err = br.ReadBit()
	if err != nil {
		return nil, err
	}
	if sh.load_non_intra_quantiser_matrix {
		for v := 0; v < 8; v++ {
			for u := 0; u < 8; u++ {
				if val, err := br.Read32(8); err != nil {
					return nil, err
				} else {
					sh.non_intra_quantiser_matrix[v][u] = uint8(val)
				}
			}
		}
	}

	return &sh, next_start_code(br)
}

func (vs *VideoSequence) sequence_header() (err error) {

	sh, err := ReadSequenceHeader(vs)
	if err != nil {
		return err
	}

	if sh.load_intra_quantiser_matrix {
		vs.quantisationMatricies[0] = sh.intra_quantiser_matrix
		vs.quantisationMatricies[2] = sh.intra_quantiser_matrix
	} else {
		vs.quantisationMatricies[0] = defaultQuantisationMatrices.Intra
		vs.quantisationMatricies[2] = defaultQuantisationMatrices.Intra
	}

	if sh.load_non_intra_quantiser_matrix {
		vs.quantisationMatricies[1] = sh.non_intra_quantiser_matrix
		vs.quantisationMatricies[3] = sh.non_intra_quantiser_matrix
	} else {
		vs.quantisationMatricies[1] = defaultQuantisationMatrices.NonIntra
		vs.quantisationMatricies[3] = defaultQuantisationMatrices.NonIntra
	}

	vs.SequenceHeader = sh

	return nil
}
<|endoftext|>


<commit_before>package gnuplot

import (
	"fmt"
	"github.com/yassu/gnup/conf"
	"github.com/yassu/gnup/utils"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
)

// Plotter
type Plotter struct {
	configures []*conf.Configure
}

func NewPlotter() *Plotter {
	plotter := new(Plotter)
	return plotter
}

func (p *Plotter) Configure(conf *conf.Configure) {
	for j := range p.configures {
		if p.configures[j].GetKey() == conf.GetKey() {
			p.configures[j].SetVals(conf.GetVals())
			return
		}
	}
	p.configures = append(p.configures, conf)
}

func (p *Plotter) GetC(key string) []string {
	for j := range p.configures {
		if p.configures[j].GetKey() == key {
			return p.configures[j].GetVals()
		}
	}
	return []string{}
}

type PlotElement interface {
	GetData() [][2]float64
	getGnuData() string
	gnuplot(filename string) string
}

// Function2d
const DefaultFunction2dSplitNum int = 1000

type Function2d struct {
	plotter  Plotter
	splitNum int
	f        func(float64) float64
}

func NewFunction2d() *Function2d {
	fun := new(Function2d)
	fun.splitNum = DefaultFunction2dSplitNum
	fun.setConfigure()
	return fun
}

func (fun *Function2d) setConfigure() {
	for _, conf := range conf.Function2dConfs() {
		fun.plotter.Configure(conf)
	}
}

func (fun *Function2d) Configure(key string, vals []string) {
	for j, conf := range fun.plotter.configures {
		if utils.InStr(key, conf.AliasedKeys()) {
			fun.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (fun *Function2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		fun.Configure(key, vals)
	}
}

func (fun Function2d) GetData() [][2]float64 { // TODO: テスト書く
	xMin, _ := strconv.ParseFloat(fun.plotter.GetC("_xMin")[0], 32)
	xMax, _ := strconv.ParseFloat(fun.plotter.GetC("_xMax")[0], 32)
	var sep = float64(xMax-xMin) / float64(fun.splitNum-1)

	var a [][2]float64
	for j := 0; j < fun.splitNum; j++ {
		t := xMin + float64(j)*sep
		y := fun.f(t)
		a = append(a, [2]float64{t, y})
	}
	return a
}

func (fun Function2d) getGnuData() string {
	var s string
	for _, xs := range fun.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (fun *Function2d) SetF(_f func(float64) float64) {
	fun.f = _f
}

func (fun Function2d) gnuplot(filename string) string {
	title := fun.plotter.GetC("_title")
	var s = fmt.Sprintf("\"%v\"", filename)
	if !isDummyVal(title) {
		s += fmt.Sprintf(" title \"%v\"", title[0])
	}

	for _, conf := range fun.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += fmt.Sprintf(" %v ", conf.GetKey())
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			for _, val := range vals {
				s += fmt.Sprintf(" %v", val)
			}
		}
	}
	return s
}

func (fun *Function2d) writeIntoGnufile(f os.File) {
	f.WriteString(fun.getGnuData())
}

// Curve2d
const DefaultCurve2dSplitNum int = 1000

type Curve2d struct {
	plotter  Plotter
	splitNum int
	c        func(float64) [2]float64
}

func NewCurve2d() *Curve2d {
	c := new(Curve2d)
	c.splitNum = DefaultCurve2dSplitNum
	c.setConfigure()
	return c
}

func (c *Curve2d) setConfigure() {
	for _, conf := range conf.Curve2dConfs() {
		c.plotter.Configure(conf)
	}
}

func (c *Curve2d) Configure(key string, vals []string) {
	for j, conf := range c.plotter.configures {
		if utils.InStr(key, conf.AliasedKeys()) {
			c.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (c *Curve2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		c.Configure(key, vals)
	}
}

func (c Curve2d) GetData() [][2]float64 { // TODO: test
	tMin, _ := strconv.ParseFloat(c.plotter.GetC("_tMin")[0], 32)
	tMax, _ := strconv.ParseFloat(c.plotter.GetC("_tMax")[0], 32)
	var sep = float64(tMax-tMin) / float64(c.splitNum-1)

	var a [][2]float64
	for j := 0; j < c.splitNum; j++ {
		cs := c.c(tMin + float64(j)*sep)
		a = append(a, [2]float64{cs[0], cs[1]})
	}
	return a
}

func (c Curve2d) getGnuData() string {
	var s string
	for _, xs := range c.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (c *Curve2d) SetC(_c func(float64) [2]float64) {
	c.c = _c
}

func (c Curve2d) gnuplot(fileName string) string {
	var s = fmt.Sprintf("\"%v\" ", fileName)
	for _, conf := range c.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += fmt.Sprintf(" %v ", conf.GetKey())
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			for _, val := range vals {
				s += fmt.Sprintf(" %v", val)
			}
		}
	}
	return s
}

func isDummyVal(vals []string) bool {
	return len(vals) == 0
}

// Graph
type Graph2d struct {
	plotter Plotter
	pElems  []PlotElement
}

func NewGraph2d() *Graph2d {
	g := new(Graph2d)
	g.setConfigure()
	return g
}

func (g *Graph2d) setConfigure() {
	for _, conf := range conf.GraphConfs() {
		g.plotter.Configure(conf)
	}
}

func (g *Graph2d) Configure(key string, vals []string) {
	for j, conf := range g.plotter.configures {
		if conf.GetKey() == key {
			g.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (g *Graph2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		g.Configure(key, vals)
	}
}

func (g *Graph2d) AppendPElem(p PlotElement) {
	g.pElems = append(g.pElems, p)
}

func (g Graph2d) writeIntoFile(data string, f *os.File) {
	f.WriteString(data)
}

func (g Graph2d) gnuplot(elemFilenames []string) string {
	var s string

	for _, conf := range g.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += "set "
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			s += conf.GetKey()
			for _, val := range vals {
				s += fmt.Sprintf(" %v ", val)
			}
			s += ";\n"
		}
	}

	s += "plot "
	for j, _ := range g.pElems {
		s += g.pElems[j].gnuplot(elemFilenames[j])
		if j != len(g.pElems)-1 {
			s += ", "
		}
	}
	s += ";\n"

	s += "pause -1;\n"
	return s
}

func (g *Graph2d) Run() {
	tmpDir := os.TempDir() + "/gnup/"
	// TODO: tmpDirがなければ作る
	// execFilename := tmpDir + "exec.gnu"
	execFilename := "exec.gnu"

	// それぞれのcurveのdataをtempファイルに書き込む
	// また, それらのファイルの名前を curve_filenames []stringに格納する
	var plotElemFilenames []string
	for _, p := range g.pElems {
		file, _ := ioutil.TempFile(tmpDir, "")
		defer func() {
			file.Close()
		}()
		g.writeIntoFile(p.getGnuData(), file)
		plotElemFilenames = append(plotElemFilenames, file.Name())
	}

	// 実行するgnuplotの実行ファイルをtempファイルに書き込む
	os.Remove(execFilename)
	execFile, err := os.OpenFile(execFilename, os.O_CREATE|os.O_WRONLY, 0666)
	defer func() {
		execFile.Close()
	}()
	if err != nil {
		fmt.Println(err)
	} else {
		execFile.WriteString(g.gnuplot(plotElemFilenames))
	}
}
<commit_msg>PlotElement => PlotElement2d<commit_after>package gnuplot

import (
	"fmt"
	"github.com/yassu/gnup/conf"
	"github.com/yassu/gnup/utils"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
)

// Plotter
type Plotter struct {
	configures []*conf.Configure
}

func NewPlotter() *Plotter {
	plotter := new(Plotter)
	return plotter
}

func (p *Plotter) Configure(conf *conf.Configure) {
	for j := range p.configures {
		if p.configures[j].GetKey() == conf.GetKey() {
			p.configures[j].SetVals(conf.GetVals())
			return
		}
	}
	p.configures = append(p.configures, conf)
}

func (p *Plotter) GetC(key string) []string {
	for j := range p.configures {
		if p.configures[j].GetKey() == key {
			return p.configures[j].GetVals()
		}
	}
	return []string{}
}

type PlotElement2d interface {
	GetData() [][2]float64
	getGnuData() string
	gnuplot(filename string) string
}

// Function2d
const DefaultFunction2dSplitNum int = 1000

type Function2d struct {
	plotter  Plotter
	splitNum int
	f        func(float64) float64
}

func NewFunction2d() *Function2d {
	fun := new(Function2d)
	fun.splitNum = DefaultFunction2dSplitNum
	fun.setConfigure()
	return fun
}

func (fun *Function2d) setConfigure() {
	for _, conf := range conf.Function2dConfs() {
		fun.plotter.Configure(conf)
	}
}

func (fun *Function2d) Configure(key string, vals []string) {
	for j, conf := range fun.plotter.configures {
		if utils.InStr(key, conf.AliasedKeys()) {
			fun.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (fun *Function2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		fun.Configure(key, vals)
	}
}

func (fun Function2d) GetData() [][2]float64 { // TODO: テスト書く
	xMin, _ := strconv.ParseFloat(fun.plotter.GetC("_xMin")[0], 32)
	xMax, _ := strconv.ParseFloat(fun.plotter.GetC("_xMax")[0], 32)
	var sep = float64(xMax-xMin) / float64(fun.splitNum-1)

	var a [][2]float64
	for j := 0; j < fun.splitNum; j++ {
		t := xMin + float64(j)*sep
		y := fun.f(t)
		a = append(a, [2]float64{t, y})
	}
	return a
}

func (fun Function2d) getGnuData() string {
	var s string
	for _, xs := range fun.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (fun *Function2d) SetF(_f func(float64) float64) {
	fun.f = _f
}

func (fun Function2d) gnuplot(filename string) string {
	title := fun.plotter.GetC("_title")
	var s = fmt.Sprintf("\"%v\"", filename)
	if !isDummyVal(title) {
		s += fmt.Sprintf(" title \"%v\"", title[0])
	}

	for _, conf := range fun.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += fmt.Sprintf(" %v ", conf.GetKey())
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			for _, val := range vals {
				s += fmt.Sprintf(" %v", val)
			}
		}
	}
	return s
}

func (fun *Function2d) writeIntoGnufile(f os.File) {
	f.WriteString(fun.getGnuData())
}

// Curve2d
const DefaultCurve2dSplitNum int = 1000

type Curve2d struct {
	plotter  Plotter
	splitNum int
	c        func(float64) [2]float64
}

func NewCurve2d() *Curve2d {
	c := new(Curve2d)
	c.splitNum = DefaultCurve2dSplitNum
	c.setConfigure()
	return c
}

func (c *Curve2d) setConfigure() {
	for _, conf := range conf.Curve2dConfs() {
		c.plotter.Configure(conf)
	}
}

func (c *Curve2d) Configure(key string, vals []string) {
	for j, conf := range c.plotter.configures {
		if utils.InStr(key, conf.AliasedKeys()) {
			c.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (c *Curve2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		c.Configure(key, vals)
	}
}

func (c Curve2d) GetData() [][2]float64 { // TODO: test
	tMin, _ := strconv.ParseFloat(c.plotter.GetC("_tMin")[0], 32)
	tMax, _ := strconv.ParseFloat(c.plotter.GetC("_tMax")[0], 32)
	var sep = float64(tMax-tMin) / float64(c.splitNum-1)

	var a [][2]float64
	for j := 0; j < c.splitNum; j++ {
		cs := c.c(tMin + float64(j)*sep)
		a = append(a, [2]float64{cs[0], cs[1]})
	}
	return a
}

func (c Curve2d) getGnuData() string {
	var s string
	for _, xs := range c.GetData() {
		s += fmt.Sprintf("%f %f\n", xs[0], xs[1])
	}
	return s
}

func (c *Curve2d) SetC(_c func(float64) [2]float64) {
	c.c = _c
}

func (c Curve2d) gnuplot(fileName string) string {
	var s = fmt.Sprintf("\"%v\" ", fileName)
	for _, conf := range c.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += fmt.Sprintf(" %v ", conf.GetKey())
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			for _, val := range vals {
				s += fmt.Sprintf(" %v", val)
			}
		}
	}
	return s
}

func isDummyVal(vals []string) bool {
	return len(vals) == 0
}

// Graph
type Graph2d struct {
	plotter Plotter
	pElems  []PlotElement2d
}

func NewGraph2d() *Graph2d {
	g := new(Graph2d)
	g.setConfigure()
	return g
}

func (g *Graph2d) setConfigure() {
	for _, conf := range conf.GraphConfs() {
		g.plotter.Configure(conf)
	}
}

func (g *Graph2d) Configure(key string, vals []string) {
	for j, conf := range g.plotter.configures {
		if conf.GetKey() == key {
			g.plotter.configures[j].SetVals(vals)
			return
		}
	}
	panic(fmt.Sprintf("%v is not a key.", key))
}

func (g *Graph2d) Configures(sconf map[string][]string) {
	for key, vals := range sconf {
		g.Configure(key, vals)
	}
}

func (g *Graph2d) AppendPElem(p PlotElement2d) {
	g.pElems = append(g.pElems, p)
}

func (g Graph2d) writeIntoFile(data string, f *os.File) {
	f.WriteString(data)
}

func (g Graph2d) gnuplot(elemFilenames []string) string {
	var s string

	for _, conf := range g.plotter.configures {
		if !strings.HasPrefix(conf.GetKey(), "_") && !isDummyVal(conf.GetVals()) {
			vals := conf.GetVals()
			s += "set "
			if vals[len(vals)-1] == "true" {
				vals = vals[:len(vals)-1]
			} else if vals[len(vals)-1] == "false" {
				vals = vals[:len(vals)-1]
				s += "no"
			}
			s += conf.GetKey()
			for _, val := range vals {
				s += fmt.Sprintf(" %v ", val)
			}
			s += ";\n"
		}
	}

	s += "plot "
	for j, _ := range g.pElems {
		s += g.pElems[j].gnuplot(elemFilenames[j])
		if j != len(g.pElems)-1 {
			s += ", "
		}
	}
	s += ";\n"

	s += "pause -1;\n"
	return s
}

func (g *Graph2d) Run() {
	tmpDir := os.TempDir() + "/gnup/"
	// TODO: tmpDirがなければ作る
	// execFilename := tmpDir + "exec.gnu"
	execFilename := "exec.gnu"

	// それぞれのcurveのdataをtempファイルに書き込む
	// また, それらのファイルの名前を curve_filenames []stringに格納する
	var plotElemFilenames []string
	for _, p := range g.pElems {
		file, _ := ioutil.TempFile(tmpDir, "")
		defer func() {
			file.Close()
		}()
		g.writeIntoFile(p.getGnuData(), file)
		plotElemFilenames = append(plotElemFilenames, file.Name())
	}

	// 実行するgnuplotの実行ファイルをtempファイルに書き込む
	os.Remove(execFilename)
	execFile, err := os.OpenFile(execFilename, os.O_CREATE|os.O_WRONLY, 0666)
	defer func() {
		execFile.Close()
	}()
	if err != nil {
		fmt.Println(err)
	} else {
		execFile.WriteString(g.gnuplot(plotElemFilenames))
	}
}
<|endoftext|>


<commit_before>package chart

const (
	// DefaultEMASigma is the default exponential smoothing factor.
	DefaultEMASigma = 0.25
)

// EMASeries is a computed series.
type EMASeries struct {
	Name  string
	Style Style
	YAxis YAxisType

	Period      int
	InnerSeries ValueProvider
}

// GetName returns the name of the time series.
func (ema EMASeries) GetName() string {
	return ema.Name
}

// GetStyle returns the line style.
func (ema EMASeries) GetStyle() Style {
	return ema.Style
}

// GetYAxis returns which YAxis the series draws on.
func (ema EMASeries) GetYAxis() YAxisType {
	return ema.YAxis
}

// GetPeriod returns the window size.
func (ema EMASeries) GetPeriod(defaults ...int) int {
	if ema.Period == 0 {
		if len(defaults) > 0 {
			return defaults[0]
		}
		return ema.InnerSeries.Len()
	}
	return ema.Period
}

// Len returns the number of elements in the series.
func (ema EMASeries) Len() int {
	return ema.InnerSeries.Len()
}

// GetSigma returns the smoothing factor for the serise.
func (ema EMASeries) GetSigma() float64 {
	return 2.0 / (float64(ema.Period) + 1)
}

// GetValue gets a value at a given index.
func (ema EMASeries) GetValue(index int) (x, y float64) {
	if ema.InnerSeries == nil {
		return
	}
	vx, _ := ema.InnerSeries.GetValue(index)
	x = vx
	y = ema.compute(ema.GetPeriod(), index)
	return
}

// GetLastValue computes the last moving average value but walking back window size samples,
// and recomputing the last moving average chunk.
func (ema EMASeries) GetLastValue() (x, y float64) {
	if ema.InnerSeries == nil {
		return
	}
	lastIndex := ema.InnerSeries.Len() - 1
	x, _ = ema.InnerSeries.GetValue(lastIndex)
	y = ema.compute(ema.GetPeriod(), lastIndex)
	return
}

func (ema EMASeries) compute(period, index int) float64 {
	_, v := ema.InnerSeries.GetValue(index)
	if index == 0 {
		return v
	}
	previousEMA := ema.compute(period-1, index-1)
	return ((v - previousEMA) * ema.GetSigma()) + previousEMA
}

// Render renders the series.
func (ema EMASeries) Render(r Renderer, canvasBox Box, xrange, yrange Range, defaults Style) {
	style := ema.Style.WithDefaultsFrom(defaults)
	DrawLineSeries(r, canvasBox, xrange, yrange, style, ema)
}
<commit_msg>changing ema defaults.<commit_after>package chart

const (
	// DefaultEMAPeriod is the default EMA period used in the sigma calculation.
	DefaultEMAPeriod = 12
)

// EMASeries is a computed series.
type EMASeries struct {
	Name  string
	Style Style
	YAxis YAxisType

	Period      int
	InnerSeries ValueProvider
}

// GetName returns the name of the time series.
func (ema EMASeries) GetName() string {
	return ema.Name
}

// GetStyle returns the line style.
func (ema EMASeries) GetStyle() Style {
	return ema.Style
}

// GetYAxis returns which YAxis the series draws on.
func (ema EMASeries) GetYAxis() YAxisType {
	return ema.YAxis
}

// GetPeriod returns the window size.
func (ema EMASeries) GetPeriod(defaults ...int) int {
	if ema.Period == 0 {
		if len(defaults) > 0 {
			return defaults[0]
		}
		return DefaultEMAPeriod
	}
	return ema.Period
}

// Len returns the number of elements in the series.
func (ema EMASeries) Len() int {
	return ema.InnerSeries.Len()
}

// GetSigma returns the smoothing factor for the serise.
func (ema EMASeries) GetSigma() float64 {
	return 2.0 / (float64(ema.Period) + 1)
}

// GetValue gets a value at a given index.
func (ema EMASeries) GetValue(index int) (x, y float64) {
	if ema.InnerSeries == nil {
		return
	}
	vx, _ := ema.InnerSeries.GetValue(index)
	x = vx
	y = ema.compute(ema.GetPeriod(), index)
	return
}

// GetLastValue computes the last moving average value but walking back window size samples,
// and recomputing the last moving average chunk.
func (ema EMASeries) GetLastValue() (x, y float64) {
	if ema.InnerSeries == nil {
		return
	}
	lastIndex := ema.InnerSeries.Len() - 1
	x, _ = ema.InnerSeries.GetValue(lastIndex)
	y = ema.compute(ema.GetPeriod(), lastIndex)
	return
}

func (ema EMASeries) compute(period, index int) float64 {
	_, v := ema.InnerSeries.GetValue(index)
	if index == 0 {
		return v
	}
	previousEMA := ema.compute(period-1, index-1)
	return ((v - previousEMA) * ema.GetSigma()) + previousEMA
}

// Render renders the series.
func (ema EMASeries) Render(r Renderer, canvasBox Box, xrange, yrange Range, defaults Style) {
	style := ema.Style.WithDefaultsFrom(defaults)
	DrawLineSeries(r, canvasBox, xrange, yrange, style, ema)
}
<|endoftext|>


<commit_before>package gohr

import (
	"fmt"
	"log"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

// getWidth gets number of width of terminal from crypto subdirectory ssh/terminal
func getTerminalWidth() (int, error) {
	w, _, err := terminal.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return -1, err
	}
	return w, nil
}

// Draw fills a row with '#' by default (if no arguments are provided) or takes arguments and prints each pattern on a new line.
func Draw(patterns ...string) {
	w, err := getTerminalWidth()
	if err != nil {
		log.Fatalf("Error getting terminal width: %s\n", err)
	}

	if len(patterns) == 0 {
		for i := 0; i < w; i++ {
			fmt.Printf("#")
		}
		fmt.Printf("\n")
	} else {
		for _, pattern := range patterns {
			l := len(pattern)
			for i := 0; i < w/l; i++ {
				fmt.Printf(pattern)
			}
			// Fills up the remaining columns in the row with part of the pattern
			fmt.Printf("%s\n", pattern[:w%l])
		}
	}
}
<commit_msg>Separate pattern drawing logic out of Draw(patterns ...string)<commit_after>package gohr

import (
	"fmt"
	"io"
	"log"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

// getWidth gets number of width of terminal from crypto subdirectory ssh/terminal
func getWidth() (int, error) {
	w, _, err := terminal.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return -1, err
	}
	return w, nil
}

func drawPatterns(width int, patterns []string, wr io.Writer) {
	if len(patterns) == 0 {
		for i := 0; i < width; i++ {
			fmt.Printf("#")
		}
		fmt.Printf("\n")
	} else {
		for _, pattern := range patterns {
			patternLength := len(pattern)
			for i := 0; i < width/patternLength; i++ {
				fmt.Printf(pattern)
			}
			// Fills up the remaining columns in the row with part of the pattern
			fmt.Printf("%s\n", pattern[:width%patternLength])
		}
	}
}

// Draw fills a row with '#' by default (if no arguments are provided) or takes arguments and prints each pattern on a new line.
func Draw(patterns ...string) {
	w, err := getWidth()
	if err != nil {
		log.Fatalf("Error getting terminal width: %s\n", err)
	}

	var p []string

	for _, pattern := range patterns {
		p = append(p, pattern)
	}

	drawPatterns(w, p, os.Stdout)
}
<|endoftext|>