<commit_before>package xmpp

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

const (
	NsMUC     = "http://jabber.org/protocol/muc"
	NsMUCUser = "http://jabber.org/protocol/muc#user"
)

type Status struct {
	XMLName xml.Name `xml:"status"`
	Code    string   `xml:"code,attr"`
}

type Item struct {
	XMLName xml.Name `xml:"item"`
	// owner, admin, member, outcast, none
	Affiliation string `xml:"affiliation,attr,omitempty"`
	// moderator, participant, visitor, none
	Role string `xml:"role,attr,omitempty"`
}

type X struct {
	XMLName  xml.Name `xml:"http://jabber.org/protocol/muc#user x"`
	Items    []Item   `xml:"item,omitempty"`
	Statuses []Status `xml:"status,omitempty"`
}

type Photo struct {
	XMLName xml.Name `xml:"vcard-temp:x:update x"`
	Photo   string   `xml:"photo,omitempty"`
}

// http://xmpp.org/extensions/xep-0045.html
// <presence />
type MUCPresence struct {
	XMLName xml.Name `xml:"presence"`
	Lang    string   `xml:"lang,attr,omitempty"`
	From    string   `xml:"from,attr,omitempty"`
	To      string   `xml:"to,attr,omitempty"`
	Id      string   `xml:"id.attr,omitempty"`
	Type    string   `xml:"type,attr,omitempty"`

	X     []X   `xml:"http://jabber.org/protocol/muc#user x,omitempty"`
	Photo Photo `xml:"vcard-temp:x:update x"` // http://xmpp.org/extensions/xep-0153.html

	Show     string       `xml:"show,omitempty"`   // away, chat, dnd, xa
	Status   string       `xml:"status,omitempty"` // sb []clientText
	Priority string       `xml:"priority,omitempty"`
	Caps     *ClientCaps  `xml:"c"`
	Error    *ClientError `xml:"error"`
}

// JoinMUC joins to a given conference with nick and optional password
// http://xmpp.org/extensions/xep-0045.html#bizrules-presence
func (c *Conn) JoinMUC(to, nick, password string) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	// remove resource from jid
	parts := strings.SplitN(to, "/", 2)
	to = parts[0]
	if len(nick) == 0 {
		if len(parts) == 2 {
			nick = parts[1]
		} else { // if nick empty & bare jid, set nick to login
			nick = strings.SplitN(c.jid, "@", 2)[0]
		}
	}
	var pass string
	if len(password) > 0 {
		pass = "<password>" + password + "</password>"
	}
	stanza := fmt.Sprintf("<presence from='%s' to='%s/%s'>"+
		"\n  <x xmlns='%s'>"+
		"\n    <history maxchars='0' />"+
		"\n    "+pass+
		"\n  </x>"+
		"\n</presence>",
		xmlEscape(c.jid), xmlEscape(to), xmlEscape(nick), NsMUC)
	_, err := fmt.Fprintf(c.out, stanza)
	return err
}

// SendMUC sends a message to the given conference with specified type (chat or groupchat).
func (c *Conn) SendMUC(to, typ, msg string) error {
	if typ == "" {
		typ = "groupchat"
	}
	cookie := c.getCookie()
	stanza := fmt.Sprintf("<message from='%s' to='%s' type='%s' id='%x'>%s</message>",
		xmlEscape(c.jid), xmlEscape(to), xmlEscape(typ), cookie, xmlEscape(msg))
	_, err := fmt.Fprintf(c.out, stanza)
	return err
}
<commit_msg>MUCPrensence->X->Item: add JID attribute<commit_after>package xmpp

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

const (
	NsMUC     = "http://jabber.org/protocol/muc"
	NsMUCUser = "http://jabber.org/protocol/muc#user"
)

type Status struct {
	XMLName xml.Name `xml:"status"`
	Code    string   `xml:"code,attr"`
}

type Item struct {
	XMLName xml.Name `xml:"item"`
	// owner, admin, member, outcast, none
	Affiliation string `xml:"affiliation,attr,omitempty"`
	// moderator, participant, visitor, none
	Role string `xml:"role,attr,omitempty"`
	JID  string `xml:"jid,attr,omitempty"`
}

type X struct {
	XMLName  xml.Name `xml:"http://jabber.org/protocol/muc#user x"`
	Items    []Item   `xml:"item,omitempty"`
	Statuses []Status `xml:"status,omitempty"`
}

type Photo struct {
	XMLName xml.Name `xml:"vcard-temp:x:update x"`
	Photo   string   `xml:"photo,omitempty"`
}

// http://xmpp.org/extensions/xep-0045.html
// <presence />
type MUCPresence struct {
	XMLName xml.Name `xml:"presence"`
	Lang    string   `xml:"lang,attr,omitempty"`
	From    string   `xml:"from,attr,omitempty"`
	To      string   `xml:"to,attr,omitempty"`
	Id      string   `xml:"id.attr,omitempty"`
	Type    string   `xml:"type,attr,omitempty"`

	X     []X   `xml:"http://jabber.org/protocol/muc#user x,omitempty"`
	Photo Photo `xml:"vcard-temp:x:update x"` // http://xmpp.org/extensions/xep-0153.html

	Show     string       `xml:"show,omitempty"`   // away, chat, dnd, xa
	Status   string       `xml:"status,omitempty"` // sb []clientText
	Priority string       `xml:"priority,omitempty"`
	Caps     *ClientCaps  `xml:"c"`
	Error    *ClientError `xml:"error"`
}

// JoinMUC joins to a given conference with nick and optional password
// http://xmpp.org/extensions/xep-0045.html#bizrules-presence
func (c *Conn) JoinMUC(to, nick, password string) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	// remove resource from jid
	parts := strings.SplitN(to, "/", 2)
	to = parts[0]
	if len(nick) == 0 {
		if len(parts) == 2 {
			nick = parts[1]
		} else { // if nick empty & bare jid, set nick to login
			nick = strings.SplitN(c.jid, "@", 2)[0]
		}
	}
	var pass string
	if len(password) > 0 {
		pass = "<password>" + password + "</password>"
	}
	stanza := fmt.Sprintf("<presence from='%s' to='%s/%s'>"+
		"\n  <x xmlns='%s'>"+
		"\n    <history maxchars='0' />"+
		"\n    "+pass+
		"\n  </x>"+
		"\n</presence>",
		xmlEscape(c.jid), xmlEscape(to), xmlEscape(nick), NsMUC)
	_, err := fmt.Fprintf(c.out, stanza)
	return err
}

// SendMUC sends a message to the given conference with specified type (chat or groupchat).
func (c *Conn) SendMUC(to, typ, msg string) error {
	if typ == "" {
		typ = "groupchat"
	}
	cookie := c.getCookie()
	stanza := fmt.Sprintf("<message from='%s' to='%s' type='%s' id='%x'>%s</message>",
		xmlEscape(c.jid), xmlEscape(to), xmlEscape(typ), cookie, xmlEscape(msg))
	_, err := fmt.Fprintf(c.out, stanza)
	return err
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

// Package mux 功能完备的路由中间件
package mux

import (
	"net/http"
	"strings"

	"github.com/issue9/mux/v4/internal/handlers"
	"github.com/issue9/mux/v4/internal/syntax"
	"github.com/issue9/mux/v4/params"
)

var (
	defaultNotFound = func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
	}

	defaultMethodNotAllowed = func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
	}
)

// Mux 提供了强大的路由匹配功能
//
// 可以对路径按正则或是请求方法进行匹配。用法如下：
//  m := mux.Default()
//  router, ok := m.NewRouter("default", group.Any)
//  router.Get("/abc/h1", h1).
//      Post("/abc/h2", h2).
//      Handle("/api/{version:\\d+}",h3, http.MethodGet, http.MethodPost) // 只匹配 GET 和 POST
//  http.ListenAndServe(m)
type Mux struct {
	routers []*Router

	notFound,
	methodNotAllowed http.HandlerFunc

	disableOptions,
	disableHead,
	skipCleanPath bool

	middlewares *Middlewares
}

// Default New 的默认参数版本
func Default() *Mux { return New(false, false, false, nil, nil) }

// New 声明一个新的 Mux
//
// disableOptions 是否禁用自动生成 OPTIONS 功能；
// disableHead 是否禁用根据 Get 请求自动生成 HEAD 请求；
// skipCleanPath 是否不对访问路径作处理，比如 "//api" ==> "/api"；
// notFound 404 页面的处理方式，为 nil 时会调用默认的方式进行处理；
// methodNotAllowed 405 页面的处理方式，为 nil 时会调用默认的方式进行处理；
func New(disableOptions, disableHead, skipCleanPath bool, notFound, methodNotAllowed http.HandlerFunc) *Mux {
	if notFound == nil {
		notFound = defaultNotFound
	}
	if methodNotAllowed == nil {
		methodNotAllowed = defaultMethodNotAllowed
	}

	mux := &Mux{
		routers: make([]*Router, 0, 1),

		notFound:         notFound,
		methodNotAllowed: methodNotAllowed,

		disableOptions: disableOptions,
		disableHead:    disableHead,
		skipCleanPath:  skipCleanPath,
	}
	mux.middlewares = NewMiddlewares(http.HandlerFunc(mux.serveHTTP))

	return mux
}

func (mux *Mux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	mux.middlewares.ServeHTTP(w, r)
}

func (mux *Mux) serveHTTP(w http.ResponseWriter, r *http.Request) {
	if !mux.skipCleanPath {
		r.URL.Path = cleanPath(r.URL.Path)
	}

	for _, router := range mux.routers {
		if req, ok := router.matcher.Match(r); ok {
			router.middlewares.ServeHTTP(w, req)
			return
		}
	}
	mux.notFound(w, r)
}

// Params 获取路由的参数集合
func Params(r *http.Request) params.Params { return params.Get(r) }

// IsWell 语法格式是否正确
//
// 如果出错，则会返回具体的错误信息。
func IsWell(pattern string) error {
	_, err := syntax.Split(pattern)
	return err
}

// Methods 返回所有支持的请求方法
func Methods() []string {
	methods := make([]string, len(handlers.Methods))
	copy(methods, handlers.Methods)
	return methods
}

// 清除路径中的重复的 / 字符
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}

	var b strings.Builder
	b.Grow(len(p) + 1)

	if p[0] != '/' {
		b.WriteByte('/')
	}

	index := strings.Index(p, "//")
	if index == -1 {
		b.WriteString(p)
		return b.String()
	}

	b.WriteString(p[:index+1])

	slash := true
	for i := index + 2; i < len(p); i++ {
		if p[i] == '/' {
			if slash {
				continue
			}
			slash = true
		} else {
			slash = false
		}
		b.WriteByte(p[i])
	}

	return b.String()
}
<commit_msg>doc: 添加有关路由语法的描述<commit_after>// SPDX-License-Identifier: MIT

// Package mux 功能完备的路由中间件
//
// 语法
//
// 路由支持以 {} 的形式包含参数，比如：/posts/{id}.html，id 在解析时会解析任意字符。
// 也可以在 {} 中约束参数的范围，比如 /posts/{id:\\d+}.html，表示 id 只能匹配数字。
// 路由地址可以是 ascii 字符，但是参数名称如果是非 ascii，在正则表达式中无法使用。
package mux

import (
	"net/http"
	"strings"

	"github.com/issue9/mux/v4/internal/handlers"
	"github.com/issue9/mux/v4/internal/syntax"
	"github.com/issue9/mux/v4/params"
)

var (
	defaultNotFound = func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
	}

	defaultMethodNotAllowed = func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
	}
)

// Mux 提供了强大的路由匹配功能
//
// 可以对路径按正则或是请求方法进行匹配。用法如下：
//  m := mux.Default()
//  router, ok := m.NewRouter("default", group.Any)
//  router.Get("/abc/h1", h1).
//      Post("/abc/h2", h2).
//      Handle("/api/{version:\\d+}",h3, http.MethodGet, http.MethodPost) // 只匹配 GET 和 POST
//  http.ListenAndServe(m)
type Mux struct {
	routers []*Router

	notFound,
	methodNotAllowed http.HandlerFunc

	disableOptions,
	disableHead,
	skipCleanPath bool

	middlewares *Middlewares
}

// Default New 的默认参数版本
func Default() *Mux { return New(false, false, false, nil, nil) }

// New 声明一个新的 Mux
//
// disableOptions 是否禁用自动生成 OPTIONS 功能；
// disableHead 是否禁用根据 Get 请求自动生成 HEAD 请求；
// skipCleanPath 是否不对访问路径作处理，比如 "//api" ==> "/api"；
// notFound 404 页面的处理方式，为 nil 时会调用默认的方式进行处理；
// methodNotAllowed 405 页面的处理方式，为 nil 时会调用默认的方式进行处理；
func New(disableOptions, disableHead, skipCleanPath bool, notFound, methodNotAllowed http.HandlerFunc) *Mux {
	if notFound == nil {
		notFound = defaultNotFound
	}
	if methodNotAllowed == nil {
		methodNotAllowed = defaultMethodNotAllowed
	}

	mux := &Mux{
		routers: make([]*Router, 0, 1),

		notFound:         notFound,
		methodNotAllowed: methodNotAllowed,

		disableOptions: disableOptions,
		disableHead:    disableHead,
		skipCleanPath:  skipCleanPath,
	}
	mux.middlewares = NewMiddlewares(http.HandlerFunc(mux.serveHTTP))

	return mux
}

func (mux *Mux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	mux.middlewares.ServeHTTP(w, r)
}

func (mux *Mux) serveHTTP(w http.ResponseWriter, r *http.Request) {
	if !mux.skipCleanPath {
		r.URL.Path = cleanPath(r.URL.Path)
	}

	for _, router := range mux.routers {
		if req, ok := router.matcher.Match(r); ok {
			router.middlewares.ServeHTTP(w, req)
			return
		}
	}
	mux.notFound(w, r)
}

// Params 获取路由的参数集合
func Params(r *http.Request) params.Params { return params.Get(r) }

// IsWell 语法格式是否正确
//
// 如果出错，则会返回具体的错误信息。
func IsWell(pattern string) error {
	_, err := syntax.Split(pattern)
	return err
}

// Methods 返回所有支持的请求方法
func Methods() []string {
	methods := make([]string, len(handlers.Methods))
	copy(methods, handlers.Methods)
	return methods
}

// 清除路径中的重复的 / 字符
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}

	var b strings.Builder
	b.Grow(len(p) + 1)

	if p[0] != '/' {
		b.WriteByte('/')
	}

	index := strings.Index(p, "//")
	if index == -1 {
		b.WriteString(p)
		return b.String()
	}

	b.WriteString(p[:index+1])

	slash := true
	for i := index + 2; i < len(p); i++ {
		if p[i] == '/' {
			if slash {
				continue
			}
			slash = true
		} else {
			slash = false
		}
		b.WriteByte(p[i])
	}

	return b.String()
}
<|endoftext|>


<commit_before>package memberlist

import (
	"bytes"
	"fmt"
	"github.com/ugorji/go/codec"
	"io"
	"log"
	"net"
)

const (
	pingMsg = iota
	indirectPingMsg
	ackRespMsg
	suspectMsg
	aliveMsg
	deadMsg
	pushPullMsg
	compoundMsg
	userMsg // User mesg, not handled by us
)

const (
	udpRecvBuf             = 2 * 1024 * 1024
	udpBufSize             = 65536
	udpSendBuf             = 1400
	compoundHeaderOverhead = 2 // Assumed header overhead
	compoundOverhead       = 2 // Assumed overhead per entry in compoundHeader
	userMsgOverhead        = 1
	metaMaxSize            = 128 // Maximum size for nod emeta data
)

// ping request sent directly to node
type ping struct {
	SeqNo uint32
}

// indirect ping sent to an indirect ndoe
type indirectPingReq struct {
	SeqNo  uint32
	Target []byte
}

// ack response is sent for a ping
type ackResp struct {
	SeqNo uint32
}

// suspect is broadcast when we suspect a node is dead
type suspect struct {
	Incarnation uint32
	Node        string
}

// alive is broadcast when we know a node is alive.
// Overloaded for nodes joining
type alive struct {
	Incarnation uint32
	Node        string
	Addr        []byte
	Meta        []byte
}

// dead is broadcast when we confirm a node is dead
// Overloaded for nodes leaving
type dead struct {
	Incarnation uint32
	Node        string
}

// pushPullHeader is used to inform the
// otherside how many states we are transfering
type pushPullHeader struct {
	Nodes        int
	UserStateLen int // Encodes the byte lengh of user state
}

// pushNodeState is used for pushPullReq when we are
// transfering out node states
type pushNodeState struct {
	Name        string
	Addr        []byte
	Meta        []byte
	Incarnation uint32
	State       int
}

// setUDPRecvBuf is used to resize the UDP receive window
func setUDPRecvBuf(c *net.UDPConn) {
	size := udpRecvBuf
	for {
		if err := c.SetReadBuffer(size); err == nil {
			break
		}
		size = size / 2
	}
}

// tcpListen listens for and handles incoming connections
func (m *Memberlist) tcpListen() {
	for {
		conn, err := m.tcpListener.AcceptTCP()
		if err != nil {
			if m.shutdown {
				break
			}
			log.Printf("[ERR] Error accepting TCP connection: %s", err)
			continue
		}
		go m.handleConn(conn)
	}
}

// handleConn handles a single incoming TCP connection
func (m *Memberlist) handleConn(conn *net.TCPConn) {
	defer conn.Close()

	remoteNodes, userState, err := readRemoteState(conn)
	if err != nil {
		log.Printf("[ERR] Failed to receive remote state: %s", err)
		return
	}

	if err := m.sendLocalState(conn); err != nil {
		log.Printf("[ERR] Failed to push local state: %s", err)
	}

	// Merge the membership state
	m.mergeState(remoteNodes)

	// Invoke the delegate for user state
	if m.config.UserDelegate != nil {
		m.config.UserDelegate.MergeRemoteState(userState)
	}
}

// udpListen listens for and handles incoming UDP packets
func (m *Memberlist) udpListen() {
	mainBuf := make([]byte, udpBufSize)
	var n int
	var addr net.Addr
	var err error
	for {
		// Reset buffer
		buf := mainBuf[0:udpBufSize]

		// Read a packet
		n, addr, err = m.udpListener.ReadFrom(buf)
		if err != nil {
			if m.shutdown {
				break
			}
			log.Printf("[ERR] Error reading UDP packet: %s", err)
			continue
		}

		// Check the length
		if n < 1 {
			log.Printf("[ERR] UDP packet too short (%d bytes). From: %s", len(buf), addr)
			continue
		}

		// Handle the command
		m.handleCommand(buf[:n], addr)
	}
}

func (m *Memberlist) handleCommand(buf []byte, from net.Addr) {
	// Decode the message type
	msgType := uint8(buf[0])
	buf = buf[1:]

	// Switch on the msgType
	switch msgType {
	case compoundMsg:
		m.handleCompound(buf, from)
	case pingMsg:
		m.handlePing(buf, from)
	case indirectPingMsg:
		m.handleIndirectPing(buf, from)
	case ackRespMsg:
		m.handleAck(buf, from)
	case suspectMsg:
		m.handleSuspect(buf, from)
	case aliveMsg:
		m.handleAlive(buf, from)
	case deadMsg:
		m.handleDead(buf, from)
	case userMsg:
		m.handleUser(buf, from)
	default:
		log.Printf("[ERR] UDP msg type (%d) not supported. From: %s", msgType, from)
	}
}

func (m *Memberlist) handleCompound(buf []byte, from net.Addr) {
	// Decode the parts
	trunc, parts, err := decodeCompoundMessage(buf)
	if err != nil {
		log.Printf("[ERR] Failed to decode compound request: %s", err)
		return
	}

	// Log any truncation
	if trunc > 0 {
		log.Printf("[WARN] Compound request had %d truncated messages", trunc)
	}

	// Handle each message
	for _, part := range parts {
		m.handleCommand(part, from)
	}
}

func (m *Memberlist) handlePing(buf []byte, from net.Addr) {
	var p ping
	if err := decode(buf, &p); err != nil {
		log.Printf("[ERR] Failed to decode ping request: %s", err)
		return
	}
	ack := ackResp{p.SeqNo}
	if err := m.encodeAndSendMsg(from, ackRespMsg, &ack); err != nil {
		log.Printf("[ERR] Failed to send ack: %s", err)
	}
}

func (m *Memberlist) handleIndirectPing(buf []byte, from net.Addr) {
	var ind indirectPingReq
	if err := decode(buf, &ind); err != nil {
		log.Printf("[ERR] Failed to decode indirect ping request: %s", err)
		return
	}

	// Send a ping to the correct host
	localSeqNo := m.nextSeqNo()
	ping := ping{SeqNo: localSeqNo}
	destAddr := &net.UDPAddr{IP: ind.Target, Port: m.config.UDPPort}

	// Setup a response handler to relay the ack
	respHandler := func() {
		ack := ackResp{ind.SeqNo}
		if err := m.encodeAndSendMsg(from, ackRespMsg, &ack); err != nil {
			log.Printf("[ERR] Failed to forward ack: %s", err)
		}
	}
	m.setAckHandler(localSeqNo, respHandler, m.config.RTT)

	// Send the ping
	if err := m.encodeAndSendMsg(destAddr, pingMsg, &ping); err != nil {
		log.Printf("[ERR] Failed to send ping: %s", err)
	}
}

func (m *Memberlist) handleAck(buf []byte, from net.Addr) {
	var ack ackResp
	if err := decode(buf, &ack); err != nil {
		log.Printf("[ERR] Failed to decode ack response: %s", err)
		return
	}
	m.invokeAckHandler(ack.SeqNo)
}

func (m *Memberlist) handleSuspect(buf []byte, from net.Addr) {
	var sus suspect
	if err := decode(buf, &sus); err != nil {
		log.Printf("[ERR] Failed to decode suspect message: %s", err)
		return
	}
	m.suspectNode(&sus)
}

func (m *Memberlist) handleAlive(buf []byte, from net.Addr) {
	var live alive
	if err := decode(buf, &live); err != nil {
		log.Printf("[ERR] Failed to decode alive message: %s", err)
		return
	}
	m.aliveNode(&live)
}

func (m *Memberlist) handleDead(buf []byte, from net.Addr) {
	var d dead
	if err := decode(buf, &d); err != nil {
		log.Printf("[ERR] Failed to decode dead message: %s", err)
		return
	}
	m.deadNode(&d)
}

// handleUser is used to notify channels of incoming user data
func (m *Memberlist) handleUser(buf []byte, from net.Addr) {
	d := m.config.UserDelegate
	if d != nil {
		d.NotifyMsg(buf)
	}
}

// encodeAndSendMsg is used to combine the encoding and sending steps
func (m *Memberlist) encodeAndSendMsg(to net.Addr, msgType int, msg interface{}) error {
	out, err := encode(msgType, msg)
	if err != nil {
		return err
	}
	if err := m.sendMsg(to, out.Bytes()); err != nil {
		return err
	}
	return nil
}

// sendMsg is used to send a UDP message to another host. It will opportunistically
// create a compoundMsg and piggy back other broadcasts
func (m *Memberlist) sendMsg(to net.Addr, msg []byte) error {
	// Check if we can piggy back any messages
	bytesAvail := udpSendBuf - len(msg) - compoundHeaderOverhead
	extra := m.getBroadcasts(compoundOverhead, bytesAvail)

	// Fast path if nothing to piggypack
	if len(extra) == 0 {
		return m.rawSendMsg(to, msg)
	}

	// Join all the messages
	msgs := make([][]byte, 0, 1+len(extra))
	msgs = append(msgs, msg)
	msgs = append(msgs, extra...)

	// Create a compound message
	compound := makeCompoundMessage(msgs)

	// Send the message
	return m.rawSendMsg(to, compound.Bytes())
}

// rawSendMsg is used to send a UDP message to another host without modification
func (m *Memberlist) rawSendMsg(to net.Addr, msg []byte) error {
	_, err := m.udpListener.WriteTo(msg, to)
	return err
}

// sendState is used to initiate a push/pull over TCP with a remote node
func (m *Memberlist) sendAndReceiveState(addr []byte) ([]pushNodeState, []byte, error) {
	// Attempt to connect
	dialer := net.Dialer{Timeout: m.config.TCPTimeout}
	dest := net.TCPAddr{IP: addr, Port: m.config.TCPPort}
	conn, err := dialer.Dial("tcp", dest.String())
	if err != nil {
		return nil, nil, err
	}

	// Send our state
	if err := m.sendLocalState(conn); err != nil {
		return nil, nil, err
	}

	// Read remote state
	remote, userState, err := readRemoteState(conn)
	if err != nil {
		return nil, nil, err
	}

	// Return the remote state
	return remote, userState, nil
}

// sendLocalState is invoked to send our local state over a tcp connection
func (m *Memberlist) sendLocalState(conn net.Conn) error {
	// Prepare the local node state
	m.nodeLock.RLock()
	localNodes := make([]pushNodeState, len(m.nodes))
	for idx, n := range m.nodes {
		localNodes[idx].Name = n.Name
		localNodes[idx].Addr = n.Addr
		localNodes[idx].Incarnation = n.Incarnation
		localNodes[idx].State = n.State
	}
	m.nodeLock.RUnlock()

	// Get the delegate state
	var userData []byte
	if m.config.UserDelegate != nil {
		userData = m.config.UserDelegate.LocalState()
	}

	// Send our node state
	header := pushPullHeader{Nodes: len(localNodes), UserStateLen: len(userData)}
	hd := codec.MsgpackHandle{}
	enc := codec.NewEncoder(conn, &hd)

	// Begin state push
	conn.Write([]byte{pushPullMsg})

	if err := enc.Encode(&header); err != nil {
		return err
	}
	for i := 0; i < header.Nodes; i++ {
		if err := enc.Encode(&localNodes[i]); err != nil {
			return err
		}
	}

	// Write the user state as well
	conn.Write(userData)
	return nil
}

// recvRemoteState is used to read the remote state from a connection
func readRemoteState(conn net.Conn) ([]pushNodeState, []byte, error) {
	// Read the message type
	buf := []byte{0}
	if _, err := conn.Read(buf); err != nil {
		return nil, nil, err
	}
	msgType := uint8(buf[0])

	// Quit if not push/pull
	if msgType != pushPullMsg {
		err := fmt.Errorf("received invalid msgType (%d)", msgType)
		return nil, nil, err
	}

	// Read the push/pull header
	var header pushPullHeader
	hd := codec.MsgpackHandle{}
	dec := codec.NewDecoder(conn, &hd)
	if err := dec.Decode(&header); err != nil {
		return nil, nil, err
	}

	// Allocate space for the transfer
	remoteNodes := make([]pushNodeState, header.Nodes)

	// Try to decode all the states
	for i := 0; i < header.Nodes; i++ {
		if err := dec.Decode(&remoteNodes[i]); err != nil {
			return remoteNodes, nil, err
		}
	}

	// Read the remote user state into a buffer
	userBuf := &bytes.Buffer{}
	_, err := io.CopyN(userBuf, conn, int64(header.UserStateLen))
	if err != nil {
		return remoteNodes, nil, err
	}

	return remoteNodes, userBuf.Bytes(), nil
}
<commit_msg>Optimize user state transfer<commit_after>package memberlist

import (
	"fmt"
	"github.com/ugorji/go/codec"
	"log"
	"net"
)

const (
	pingMsg = iota
	indirectPingMsg
	ackRespMsg
	suspectMsg
	aliveMsg
	deadMsg
	pushPullMsg
	compoundMsg
	userMsg // User mesg, not handled by us
)

const (
	udpRecvBuf             = 2 * 1024 * 1024
	udpBufSize             = 65536
	udpSendBuf             = 1400
	compoundHeaderOverhead = 2 // Assumed header overhead
	compoundOverhead       = 2 // Assumed overhead per entry in compoundHeader
	userMsgOverhead        = 1
	metaMaxSize            = 128 // Maximum size for nod emeta data
)

// ping request sent directly to node
type ping struct {
	SeqNo uint32
}

// indirect ping sent to an indirect ndoe
type indirectPingReq struct {
	SeqNo  uint32
	Target []byte
}

// ack response is sent for a ping
type ackResp struct {
	SeqNo uint32
}

// suspect is broadcast when we suspect a node is dead
type suspect struct {
	Incarnation uint32
	Node        string
}

// alive is broadcast when we know a node is alive.
// Overloaded for nodes joining
type alive struct {
	Incarnation uint32
	Node        string
	Addr        []byte
	Meta        []byte
}

// dead is broadcast when we confirm a node is dead
// Overloaded for nodes leaving
type dead struct {
	Incarnation uint32
	Node        string
}

// pushPullHeader is used to inform the
// otherside how many states we are transfering
type pushPullHeader struct {
	Nodes        int
	UserStateLen int // Encodes the byte lengh of user state
}

// pushNodeState is used for pushPullReq when we are
// transfering out node states
type pushNodeState struct {
	Name        string
	Addr        []byte
	Meta        []byte
	Incarnation uint32
	State       int
}

// setUDPRecvBuf is used to resize the UDP receive window
func setUDPRecvBuf(c *net.UDPConn) {
	size := udpRecvBuf
	for {
		if err := c.SetReadBuffer(size); err == nil {
			break
		}
		size = size / 2
	}
}

// tcpListen listens for and handles incoming connections
func (m *Memberlist) tcpListen() {
	for {
		conn, err := m.tcpListener.AcceptTCP()
		if err != nil {
			if m.shutdown {
				break
			}
			log.Printf("[ERR] Error accepting TCP connection: %s", err)
			continue
		}
		go m.handleConn(conn)
	}
}

// handleConn handles a single incoming TCP connection
func (m *Memberlist) handleConn(conn *net.TCPConn) {
	defer conn.Close()

	remoteNodes, userState, err := readRemoteState(conn)
	if err != nil {
		log.Printf("[ERR] Failed to receive remote state: %s", err)
		return
	}

	if err := m.sendLocalState(conn); err != nil {
		log.Printf("[ERR] Failed to push local state: %s", err)
	}

	// Merge the membership state
	m.mergeState(remoteNodes)

	// Invoke the delegate for user state
	if m.config.UserDelegate != nil {
		m.config.UserDelegate.MergeRemoteState(userState)
	}
}

// udpListen listens for and handles incoming UDP packets
func (m *Memberlist) udpListen() {
	mainBuf := make([]byte, udpBufSize)
	var n int
	var addr net.Addr
	var err error
	for {
		// Reset buffer
		buf := mainBuf[0:udpBufSize]

		// Read a packet
		n, addr, err = m.udpListener.ReadFrom(buf)
		if err != nil {
			if m.shutdown {
				break
			}
			log.Printf("[ERR] Error reading UDP packet: %s", err)
			continue
		}

		// Check the length
		if n < 1 {
			log.Printf("[ERR] UDP packet too short (%d bytes). From: %s", len(buf), addr)
			continue
		}

		// Handle the command
		m.handleCommand(buf[:n], addr)
	}
}

func (m *Memberlist) handleCommand(buf []byte, from net.Addr) {
	// Decode the message type
	msgType := uint8(buf[0])
	buf = buf[1:]

	// Switch on the msgType
	switch msgType {
	case compoundMsg:
		m.handleCompound(buf, from)
	case pingMsg:
		m.handlePing(buf, from)
	case indirectPingMsg:
		m.handleIndirectPing(buf, from)
	case ackRespMsg:
		m.handleAck(buf, from)
	case suspectMsg:
		m.handleSuspect(buf, from)
	case aliveMsg:
		m.handleAlive(buf, from)
	case deadMsg:
		m.handleDead(buf, from)
	case userMsg:
		m.handleUser(buf, from)
	default:
		log.Printf("[ERR] UDP msg type (%d) not supported. From: %s", msgType, from)
	}
}

func (m *Memberlist) handleCompound(buf []byte, from net.Addr) {
	// Decode the parts
	trunc, parts, err := decodeCompoundMessage(buf)
	if err != nil {
		log.Printf("[ERR] Failed to decode compound request: %s", err)
		return
	}

	// Log any truncation
	if trunc > 0 {
		log.Printf("[WARN] Compound request had %d truncated messages", trunc)
	}

	// Handle each message
	for _, part := range parts {
		m.handleCommand(part, from)
	}
}

func (m *Memberlist) handlePing(buf []byte, from net.Addr) {
	var p ping
	if err := decode(buf, &p); err != nil {
		log.Printf("[ERR] Failed to decode ping request: %s", err)
		return
	}
	ack := ackResp{p.SeqNo}
	if err := m.encodeAndSendMsg(from, ackRespMsg, &ack); err != nil {
		log.Printf("[ERR] Failed to send ack: %s", err)
	}
}

func (m *Memberlist) handleIndirectPing(buf []byte, from net.Addr) {
	var ind indirectPingReq
	if err := decode(buf, &ind); err != nil {
		log.Printf("[ERR] Failed to decode indirect ping request: %s", err)
		return
	}

	// Send a ping to the correct host
	localSeqNo := m.nextSeqNo()
	ping := ping{SeqNo: localSeqNo}
	destAddr := &net.UDPAddr{IP: ind.Target, Port: m.config.UDPPort}

	// Setup a response handler to relay the ack
	respHandler := func() {
		ack := ackResp{ind.SeqNo}
		if err := m.encodeAndSendMsg(from, ackRespMsg, &ack); err != nil {
			log.Printf("[ERR] Failed to forward ack: %s", err)
		}
	}
	m.setAckHandler(localSeqNo, respHandler, m.config.RTT)

	// Send the ping
	if err := m.encodeAndSendMsg(destAddr, pingMsg, &ping); err != nil {
		log.Printf("[ERR] Failed to send ping: %s", err)
	}
}

func (m *Memberlist) handleAck(buf []byte, from net.Addr) {
	var ack ackResp
	if err := decode(buf, &ack); err != nil {
		log.Printf("[ERR] Failed to decode ack response: %s", err)
		return
	}
	m.invokeAckHandler(ack.SeqNo)
}

func (m *Memberlist) handleSuspect(buf []byte, from net.Addr) {
	var sus suspect
	if err := decode(buf, &sus); err != nil {
		log.Printf("[ERR] Failed to decode suspect message: %s", err)
		return
	}
	m.suspectNode(&sus)
}

func (m *Memberlist) handleAlive(buf []byte, from net.Addr) {
	var live alive
	if err := decode(buf, &live); err != nil {
		log.Printf("[ERR] Failed to decode alive message: %s", err)
		return
	}
	m.aliveNode(&live)
}

func (m *Memberlist) handleDead(buf []byte, from net.Addr) {
	var d dead
	if err := decode(buf, &d); err != nil {
		log.Printf("[ERR] Failed to decode dead message: %s", err)
		return
	}
	m.deadNode(&d)
}

// handleUser is used to notify channels of incoming user data
func (m *Memberlist) handleUser(buf []byte, from net.Addr) {
	d := m.config.UserDelegate
	if d != nil {
		d.NotifyMsg(buf)
	}
}

// encodeAndSendMsg is used to combine the encoding and sending steps
func (m *Memberlist) encodeAndSendMsg(to net.Addr, msgType int, msg interface{}) error {
	out, err := encode(msgType, msg)
	if err != nil {
		return err
	}
	if err := m.sendMsg(to, out.Bytes()); err != nil {
		return err
	}
	return nil
}

// sendMsg is used to send a UDP message to another host. It will opportunistically
// create a compoundMsg and piggy back other broadcasts
func (m *Memberlist) sendMsg(to net.Addr, msg []byte) error {
	// Check if we can piggy back any messages
	bytesAvail := udpSendBuf - len(msg) - compoundHeaderOverhead
	extra := m.getBroadcasts(compoundOverhead, bytesAvail)

	// Fast path if nothing to piggypack
	if len(extra) == 0 {
		return m.rawSendMsg(to, msg)
	}

	// Join all the messages
	msgs := make([][]byte, 0, 1+len(extra))
	msgs = append(msgs, msg)
	msgs = append(msgs, extra...)

	// Create a compound message
	compound := makeCompoundMessage(msgs)

	// Send the message
	return m.rawSendMsg(to, compound.Bytes())
}

// rawSendMsg is used to send a UDP message to another host without modification
func (m *Memberlist) rawSendMsg(to net.Addr, msg []byte) error {
	_, err := m.udpListener.WriteTo(msg, to)
	return err
}

// sendState is used to initiate a push/pull over TCP with a remote node
func (m *Memberlist) sendAndReceiveState(addr []byte) ([]pushNodeState, []byte, error) {
	// Attempt to connect
	dialer := net.Dialer{Timeout: m.config.TCPTimeout}
	dest := net.TCPAddr{IP: addr, Port: m.config.TCPPort}
	conn, err := dialer.Dial("tcp", dest.String())
	if err != nil {
		return nil, nil, err
	}

	// Send our state
	if err := m.sendLocalState(conn); err != nil {
		return nil, nil, err
	}

	// Read remote state
	remote, userState, err := readRemoteState(conn)
	if err != nil {
		return nil, nil, err
	}

	// Return the remote state
	return remote, userState, nil
}

// sendLocalState is invoked to send our local state over a tcp connection
func (m *Memberlist) sendLocalState(conn net.Conn) error {
	// Prepare the local node state
	m.nodeLock.RLock()
	localNodes := make([]pushNodeState, len(m.nodes))
	for idx, n := range m.nodes {
		localNodes[idx].Name = n.Name
		localNodes[idx].Addr = n.Addr
		localNodes[idx].Incarnation = n.Incarnation
		localNodes[idx].State = n.State
	}
	m.nodeLock.RUnlock()

	// Get the delegate state
	var userData []byte
	if m.config.UserDelegate != nil {
		userData = m.config.UserDelegate.LocalState()
	}

	// Send our node state
	header := pushPullHeader{Nodes: len(localNodes), UserStateLen: len(userData)}
	hd := codec.MsgpackHandle{}
	enc := codec.NewEncoder(conn, &hd)

	// Begin state push
	conn.Write([]byte{pushPullMsg})

	if err := enc.Encode(&header); err != nil {
		return err
	}
	for i := 0; i < header.Nodes; i++ {
		if err := enc.Encode(&localNodes[i]); err != nil {
			return err
		}
	}

	// Write the user state as well
	conn.Write(userData)
	return nil
}

// recvRemoteState is used to read the remote state from a connection
func readRemoteState(conn net.Conn) ([]pushNodeState, []byte, error) {
	// Read the message type
	buf := []byte{0}
	if _, err := conn.Read(buf); err != nil {
		return nil, nil, err
	}
	msgType := uint8(buf[0])

	// Quit if not push/pull
	if msgType != pushPullMsg {
		err := fmt.Errorf("received invalid msgType (%d)", msgType)
		return nil, nil, err
	}

	// Read the push/pull header
	var header pushPullHeader
	hd := codec.MsgpackHandle{}
	dec := codec.NewDecoder(conn, &hd)
	if err := dec.Decode(&header); err != nil {
		return nil, nil, err
	}

	// Allocate space for the transfer
	remoteNodes := make([]pushNodeState, header.Nodes)

	// Try to decode all the states
	for i := 0; i < header.Nodes; i++ {
		if err := dec.Decode(&remoteNodes[i]); err != nil {
			return remoteNodes, nil, err
		}
	}

	// Read the remote user state into a buffer
	var userBuf []byte
	if header.UserStateLen > 0 {
		userBuf = make([]byte, header.UserStateLen)
		bytes, err := conn.Read(userBuf)
		if err == nil && bytes != header.UserStateLen {
			err = fmt.Errorf("Failed to read full user state (%d / %d)", bytes, header.UserStateLen)
		}
		if err != nil {
			return remoteNodes, nil, err
		}
	}

	return remoteNodes, userBuf, nil
}
<|endoftext|>


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

package main

// Now tells what time it is in other time zones. The first argument
// identifies a time zone either by shorthand (EST, NYC) or by time zone
// file base name, such as Yellowknife or Paris.
//
// % now Paris
// Thu Apr 12 15:55:55 CEST 2012 Paris
// % now Adelaide
// Thu Apr 12 23:26:14 CST 2012 Adelaide
// %

import (
	"fmt"
	"os"
	"path/filepath"
	"time"
)

func main() {
	zone := ""
	t := time.Now()
	if len(os.Args) > 1 {
		zone = os.Args[1]
		if tz, ok := timeZone[zone]; ok {
			zone = tz
		}
		t = t.In(loadZone(zone))
	}
	fmt.Printf("%s %s\n", t.Format(time.UnixDate), zone)
}

func loadZone(zone string) *time.Location {
	loc, err := time.LoadLocation(zone)
	if err == nil {
		return loc
	}
	// Pure ASCII, but OK. Allow us to say "paris" as well as "Paris".
	if len(zone) > 0 && 'a' <= zone[0] && zone[0] <= 'z' {
		zone = string(zone[0]+'A'-'a') + string(zone[1:])
	}
	// See if there's a file with that name in /usr/share/zoneinfo
	files, _ := filepath.Glob("/usr/share/zoneinfo/*/" + zone)
	if len(files) >= 1 {
		if len(files) > 1 {
			fmt.Fprintf(os.Stderr, "now: multiple time zones; using first of %v\n", files)
		}
		loc, err = time.LoadLocation(files[0][len("/usr/share/zoneinfo/"):])
		if err == nil {
			return loc
		}
	}
	fmt.Fprintf(os.Stderr, "now: %s\n", err)
	os.Exit(1)
	return nil

}

// from /usr/share/zoneinfo
var timeZone = map[string]string{
	"GMT":     "Europe/London",
	"BST":     "Europe/London",
	"BSDT":    "Europe/London",
	"CET":     "Europe/Paris",
	"UTC":     "",
	"PST":     "America/Los_Angeles",
	"PDT":     "America/Los_Angeles",
	"LA":      "America/Los_Angeles",
	"LAX":     "America/Los_Angeles",
	"MST":     "America/Denver",
	"MDT":     "America/Denver",
	"CST":     "America/Chicago",
	"CDT":     "America/Chicago",
	"Chicago": "America/Chicago",
	"EST":     "America/New_York",
	"EDT":     "America/New_York",
	"NYC":     "America/New_York",
	"NY":      "America/New_York",
	"AEST":    "Australia/Sydney",
	"AEDT":    "Australia/Sydney",
	"AWST":    "Australia/Perth",
	"AWDT":    "Australia/Perth",
	"ACST":    "Australia/Adelaide",
	"ACDT":    "Australia/Adelaide",
}
<commit_msg>now: more casing tricks<commit_after>// Copyright 2012 The rspace Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

// Now tells what time it is in other time zones. The first argument
// identifies a time zone either by shorthand (EST, NYC) or by time zone
// file base name, such as Yellowknife or Paris.
//
// % now Paris
// Thu Apr 12 15:55:55 CEST 2012 Paris
// % now Adelaide
// Thu Apr 12 23:26:14 CST 2012 Adelaide
// %

import (
	"fmt"
	"os"
	"path/filepath"
	"time"
)

func main() {
	zone := ""
	t := time.Now()
	if len(os.Args) > 1 {
		zone = os.Args[1]
		if tz, ok := timeZone[zone]; ok {
			zone = tz
		} else if tz, ok = timeZone[toUpper(zone)]; ok {
			zone = tz
		}
		t = t.In(loadZone(zone))
	}
	fmt.Printf("%s %s\n", t.Format(time.UnixDate), zone)
}

func loadZone(zone string) *time.Location {
	loc, err := time.LoadLocation(zone)
	if err == nil {
		return loc
	}
	// Pure ASCII, but OK. Allow us to say "paris" as well as "Paris".
	if len(zone) > 0 && 'a' <= zone[0] && zone[0] <= 'z' {
		zone = string(zone[0]+'A'-'a') + string(zone[1:])
	}
	// See if there's a file with that name in /usr/share/zoneinfo
	files, _ := filepath.Glob("/usr/share/zoneinfo/*/" + zone)
	if len(files) >= 1 {
		if len(files) > 1 {
			fmt.Fprintf(os.Stderr, "now: multiple time zones; using first of %v\n", files)
		}
		loc, err = time.LoadLocation(files[0][len("/usr/share/zoneinfo/"):])
		if err == nil {
			return loc
		}
	}
	fmt.Fprintf(os.Stderr, "now: %s\n", err)
	os.Exit(1)
	return nil

}

// Pure ASCII
func toUpper(s string) string {
	var b = make([]byte, len(s))
	for i := range b {
		c := s[i]
		if 'a' <= c && c <= 'z' {
			c -= ' '
		}
		b[i] = c
	}
	return string(b)
}

// from /usr/share/zoneinfo
var timeZone = map[string]string{
	"GMT":     "Europe/London",
	"BST":     "Europe/London",
	"BSDT":    "Europe/London",
	"CET":     "Europe/Paris",
	"UTC":     "",
	"PST":     "America/Los_Angeles",
	"PDT":     "America/Los_Angeles",
	"LA":      "America/Los_Angeles",
	"LAX":     "America/Los_Angeles",
	"MST":     "America/Denver",
	"MDT":     "America/Denver",
	"CST":     "America/Chicago",
	"CDT":     "America/Chicago",
	"Chicago": "America/Chicago",
	"EST":     "America/New_York",
	"EDT":     "America/New_York",
	"NYC":     "America/New_York",
	"NY":      "America/New_York",
	"AEST":    "Australia/Sydney",
	"AEDT":    "Australia/Sydney",
	"AWST":    "Australia/Perth",
	"AWDT":    "Australia/Perth",
	"ACST":    "Australia/Adelaide",
	"ACDT":    "Australia/Adelaide",
}
<|endoftext|>


<commit_before>package markov
import (
	"strings"
	"math/rand"
	"time"
)

var random *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))

type Word string

type Link struct {
	weight int
	target Word
}

type adjacencyList map[Word][]Link

type Markov struct {
	chain 				 adjacencyList
	linkCount, wordCount int
}

func New() *Markov {
	return &Markov{chain: make(adjacencyList)}
}

// find a word in the linked list chain by value, returning true and
// the element if found, and false and nil or the last element otherwise
func find(theChain []Link, word Word) (bool, int) {
	for i := range theChain {
		if theChain[i].target == word {
			return true, i
		}
	}
	return false, -1
}

// AddLink creates a link from 'from' and to 'to'.
// from and to are added to the markov chain if they are not already in it
// return the new weight of the link
func (m *Markov) AddLink(from, to Word) int {
	m.AddWord(from)
	m.AddWord(to)
	fromLinks := m.chain[from]
	found, index := find(fromLinks, to)
	if found {
		// increase the weight of this link
		fromLinks[index].weight++
		return fromLinks[index].weight
	} else {
		// create a new link with a weight of 1
		m.chain[from] = append(fromLinks, Link{target: to, weight: 1})
		m.linkCount++
		return 1
    }
}

// AddWord adds a word to the markov chain with no links.
// Returns true if the word was a new word we've never seen and false if it was
// already in the chain
func (m *Markov) AddWord(word Word) bool {
	// if there's no chain for this word
	if m.chain[word] == nil {
		// create a new chain
		m.chain[word] = []Link{}
		m.wordCount++
		// and return true
		return true
	}
	// if we've already seen it return false
	return false
}

// Generate creates a string (if one exists) starting with the start word and
// continuing for at most maxWords (to avoid loops). The returned string will
// have the first word capitalized.
func (m *Markov) Generate(start string, maxWords int) string {
	startingWord := Word(start)
	if m.chain[startingWord] == nil {
		return ""
	}
	markovChain := []string{strings.Title(start)}

	links := m.chain[startingWord]
	for wordCount := 1; wordCount <= maxWords; wordCount++ {
		if len(links) == 0 {
			break
		}
		linkWeight := totalWeight(links)
		targetWeight := random.Intn(linkWeight)
		var seenWeight int
		for i := range links {
			seenWeight += links[i].weight
			if seenWeight > targetWeight {
				target := links[i].target
				links = m.chain[target]
				markovChain = append(markovChain, string(target))
				break
			}
		}
	}

	return strings.Join(markovChain, " ")
}

// GetWordCount returns the total number of unique words in the markov chain.
func (m *Markov) GetWordCount() int {
	return m.wordCount
}

// GetLinkCount returns the total number of links between words in the chain.
// This is likely greater than the number of unique words (but far less than
// the number of total words) due to the interconnectedness.
func (m *Markov) GetLinkCount() int {
	return m.linkCount
}

func totalWeight(theChain []Link) int {
	var sum int
	for i := range theChain {
		sum += theChain[i].weight
	}
	return sum
}<commit_msg>change names in structs to export so json will serialize correctly<commit_after>package markov
import (
	"strings"
	"math/rand"
	"time"
)

var random *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))

type Word string

type Link struct {
	Weight int	`json:"weight"`
	Target Word	`json:"target"`
}

type adjacencyList map[Word][]Link

type Markov struct {
	Chain                adjacencyList	`json:"chain"`
	LinkCount, WordCount int
}

func New() *Markov {
	return &Markov{Chain: make(adjacencyList)}
}

// find a word in the linked list chain by value, returning true and
// the element if found, and false and nil or the last element otherwise
func find(theChain []Link, word Word) (bool, int) {
	for i := range theChain {
		if theChain[i].Target == word {
			return true, i
		}
	}
	return false, -1
}

// AddLink creates a link from 'from' and to 'to'.
// from and to are added to the markov chain if they are not already in it
// return the new weight of the link
func (m *Markov) AddLink(from, to Word) int {
	m.AddWord(from)
	m.AddWord(to)
	fromLinks := m.Chain[from]
	found, index := find(fromLinks, to)
	if found {
		// increase the weight of this link
		fromLinks[index].Weight++
		return fromLinks[index].Weight
	} else {
		// create a new link with a weight of 1
		m.Chain[from] = append(fromLinks, Link{Target: to, Weight: 1})
		m.LinkCount++
		return 1
    }
}

// AddWord adds a word to the markov chain with no links.
// Returns true if the word was a new word we've never seen and false if it was
// already in the chain
func (m *Markov) AddWord(word Word) bool {
	// if there's no chain for this word
	if m.Chain[word] == nil {
		// create a new chain
		m.Chain[word] = []Link{}
		m.WordCount++
		// and return true
		return true
	}
	// if we've already seen it return false
	return false
}

// Generate creates a string (if one exists) starting with the start word and
// continuing for at most maxWords (to avoid loops). The returned string will
// have the first word capitalized.
func (m *Markov) Generate(start string, maxWords int) string {
	startingWord := Word(start)
	if m.Chain[startingWord] == nil {
		return ""
	}
	markovChain := []string{strings.Title(start)}

	links := m.Chain[startingWord]
	for wordCount := 1; wordCount <= maxWords; wordCount++ {
		if len(links) == 0 {
			break
		}
		linkWeight := totalWeight(links)
		targetWeight := random.Intn(linkWeight)
		var seenWeight int
		for i := range links {
			seenWeight += links[i].Weight
			if seenWeight > targetWeight {
				target := links[i].Target
				links = m.Chain[target]
				markovChain = append(markovChain, string(target))
				break
			}
		}
	}

	return strings.Join(markovChain, " ")
}

// GetWordCount returns the total number of unique words in the markov chain.
func (m *Markov) GetWordCount() int {
	return m.WordCount
}

// GetLinkCount returns the total number of links between words in the chain.
// This is likely greater than the number of unique words (but far less than
// the number of total words) due to the interconnectedness.
func (m *Markov) GetLinkCount() int {
	return m.LinkCount
}

func totalWeight(theChain []Link) int {
	var sum int
	for i := range theChain {
		sum += theChain[i].Weight
	}
	return sum
}<|endoftext|>


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

package big

import "math"

var (
	half  = NewFloat(0.5)
	two   = NewFloat(2.0)
	three = NewFloat(3.0)
)

// Sqrt sets z to the rounded square root of x, and returns it.
//
// If z's precision is 0, it is changed to x's precision before the
// operation. Rounding is performed according to z's precision and
// rounding mode.
//
// The function panics if z < 0. The value of z is undefined in that
// case.
func (z *Float) Sqrt(x *Float) *Float {
	if debugFloat {
		x.validate()
	}

	if z.prec == 0 {
		z.prec = x.prec
	}

	if x.Sign() == -1 {
		// following IEEE754-2008 (section 7.2)
		panic(ErrNaN{"square root of negative operand"})
	}

	// handle ±0 and +∞
	if x.form != finite {
		z.acc = Exact
		z.form = x.form
		z.neg = x.neg // IEEE754-2008 requires √±0 = ±0
		return z
	}

	// MantExp sets the argument's precision to the receiver's, and
	// when z.prec > x.prec this will lower z.prec. Restore it after
	// the MantExp call.
	prec := z.prec
	b := x.MantExp(z)
	z.prec = prec

	// Compute √(z·2**b) as
	//   √( z)·2**(½b)     if b is even
	//   √(2z)·2**(⌊½b⌋)   if b > 0 is odd
	//   √(½z)·2**(⌈½b⌉)   if b < 0 is odd
	switch b % 2 {
	case 0:
		// nothing to do
	case 1:
		z.Mul(two, z)
	case -1:
		z.Mul(half, z)
	}
	// 0.25 <= z < 2.0

	// Solving x² - z = 0 directly requires a Quo call, but it's
	// faster for small precisions.
	//
	// Solving 1/x² - z = 0 avoids the Quo call and is much faster for
	// high precisions.
	//
	// 128bit precision is an empirically chosen threshold.
	if z.prec <= 128 {
		z.sqrtDirect(z)
	} else {
		z.sqrtInverse(z)
	}

	// re-attach halved exponent
	return z.SetMantExp(z, b/2)
}

// Compute √x (up to prec 128) by solving
//   t² - x = 0
// for t, starting with a 53 bits precision guess from math.Sqrt and
// then using at most two iterations of Newton's method.
func (z *Float) sqrtDirect(x *Float) {
	// let
	//   f(t) = t² - x
	// then
	//   g(t) = f(t)/f'(t) = ½(t² - x)/t
	// and the next guess is given by
	//   t2 = t - g(t) = ½(t² + x)/t
	u := new(Float)
	ng := func(t *Float) *Float {
		u.prec = t.prec
		u.Mul(t, t)        // u = t²
		u.Add(u, x)        //   = t² + x
		u.Mul(half, u)     //   = ½(t² + x)
		return t.Quo(u, t) //   = ½(t² + x)/t
	}

	xf, _ := x.Float64()
	sq := NewFloat(math.Sqrt(xf))

	switch {
	case z.prec > 128:
		panic("sqrtDirect: only for z.prec <= 128")
	case z.prec > 64:
		sq.prec *= 2
		sq = ng(sq)
		fallthrough
	default:
		sq.prec *= 2
		sq = ng(sq)
	}

	z.Set(sq)
}

// Compute √x (to z.prec precision) by solving
//   1/t² - x = 0
// for t (using Newton's method), and then inverting.
func (z *Float) sqrtInverse(x *Float) {
	// let
	//   f(t) = 1/t² - x
	// then
	//   g(t) = f(t)/f'(t) = -½t(1 - xt²)
	// and the next guess is given by
	//   t2 = t - g(t) = ½t(3 - xt²)
	u := newFloat(z.prec)
	v := newFloat(z.prec)
	ng := func(t *Float) *Float {
		u.prec = t.prec
		v.prec = t.prec
		u.Mul(t, t)           // u = t²
		u.Mul(x, u)           //   = xt²
		v.Sub(three, u)       // v = 3 - xt²
		u.Mul(t, v)           // u = t(3 - xt²)
		return t.Mul(half, u) //   = ½t(3 - xt²)
	}

	xf, _ := x.Float64()
	sqi := newFloat(z.prec)
	sqi.SetFloat64(1 / math.Sqrt(xf))
	for prec := z.prec + 32; sqi.prec < prec; {
		sqi.prec *= 2
		sqi = ng(sqi)
	}
	// sqi = 1/√x

	// x/√x = √x
	z.Mul(x, sqi)
}

// newFloat returns a new *Float with space for twice the given
// precision.
func newFloat(prec2 uint32) *Float {
	z := new(Float)
	// nat.make ensures the slice length is > 0
	z.mant = z.mant.make(int(prec2/_W) * 2)
	return z
}
<commit_msg>math/big: optimize multiplication by 2 and 1/2 in float Sqrt<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package big

import "math"

var (
	three = NewFloat(3.0)
)

// Sqrt sets z to the rounded square root of x, and returns it.
//
// If z's precision is 0, it is changed to x's precision before the
// operation. Rounding is performed according to z's precision and
// rounding mode.
//
// The function panics if z < 0. The value of z is undefined in that
// case.
func (z *Float) Sqrt(x *Float) *Float {
	if debugFloat {
		x.validate()
	}

	if z.prec == 0 {
		z.prec = x.prec
	}

	if x.Sign() == -1 {
		// following IEEE754-2008 (section 7.2)
		panic(ErrNaN{"square root of negative operand"})
	}

	// handle ±0 and +∞
	if x.form != finite {
		z.acc = Exact
		z.form = x.form
		z.neg = x.neg // IEEE754-2008 requires √±0 = ±0
		return z
	}

	// MantExp sets the argument's precision to the receiver's, and
	// when z.prec > x.prec this will lower z.prec. Restore it after
	// the MantExp call.
	prec := z.prec
	b := x.MantExp(z)
	z.prec = prec

	// Compute √(z·2**b) as
	//   √( z)·2**(½b)     if b is even
	//   √(2z)·2**(⌊½b⌋)   if b > 0 is odd
	//   √(½z)·2**(⌈½b⌉)   if b < 0 is odd
	switch b % 2 {
	case 0:
		// nothing to do
	case 1:
		z.exp++
	case -1:
		z.exp--
	}
	// 0.25 <= z < 2.0

	// Solving x² - z = 0 directly requires a Quo call, but it's
	// faster for small precisions.
	//
	// Solving 1/x² - z = 0 avoids the Quo call and is much faster for
	// high precisions.
	//
	// 128bit precision is an empirically chosen threshold.
	if z.prec <= 128 {
		z.sqrtDirect(z)
	} else {
		z.sqrtInverse(z)
	}

	// re-attach halved exponent
	return z.SetMantExp(z, b/2)
}

// Compute √x (up to prec 128) by solving
//   t² - x = 0
// for t, starting with a 53 bits precision guess from math.Sqrt and
// then using at most two iterations of Newton's method.
func (z *Float) sqrtDirect(x *Float) {
	// let
	//   f(t) = t² - x
	// then
	//   g(t) = f(t)/f'(t) = ½(t² - x)/t
	// and the next guess is given by
	//   t2 = t - g(t) = ½(t² + x)/t
	u := new(Float)
	ng := func(t *Float) *Float {
		u.prec = t.prec
		u.Mul(t, t)        // u = t²
		u.Add(u, x)        //   = t² + x
		u.exp--            //   = ½(t² + x)
		return t.Quo(u, t) //   = ½(t² + x)/t
	}

	xf, _ := x.Float64()
	sq := NewFloat(math.Sqrt(xf))

	switch {
	case z.prec > 128:
		panic("sqrtDirect: only for z.prec <= 128")
	case z.prec > 64:
		sq.prec *= 2
		sq = ng(sq)
		fallthrough
	default:
		sq.prec *= 2
		sq = ng(sq)
	}

	z.Set(sq)
}

// Compute √x (to z.prec precision) by solving
//   1/t² - x = 0
// for t (using Newton's method), and then inverting.
func (z *Float) sqrtInverse(x *Float) {
	// let
	//   f(t) = 1/t² - x
	// then
	//   g(t) = f(t)/f'(t) = -½t(1 - xt²)
	// and the next guess is given by
	//   t2 = t - g(t) = ½t(3 - xt²)
	u := newFloat(z.prec)
	v := newFloat(z.prec)
	ng := func(t *Float) *Float {
		u.prec = t.prec
		v.prec = t.prec
		u.Mul(t, t)     // u = t²
		u.Mul(x, u)     //   = xt²
		v.Sub(three, u) // v = 3 - xt²
		u.Mul(t, v)     // u = t(3 - xt²)
		u.exp--         //   = ½t(3 - xt²)
		return t.Set(u)

	}

	xf, _ := x.Float64()
	sqi := newFloat(z.prec)
	sqi.SetFloat64(1 / math.Sqrt(xf))
	for prec := z.prec + 32; sqi.prec < prec; {
		sqi.prec *= 2
		sqi = ng(sqi)
	}
	// sqi = 1/√x

	// x/√x = √x
	z.Mul(x, sqi)
}

// newFloat returns a new *Float with space for twice the given
// precision.
func newFloat(prec2 uint32) *Float {
	z := new(Float)
	// nat.make ensures the slice length is > 0
	z.mant = z.mant.make(int(prec2/_W) * 2)
	return z
}
<|endoftext|>


<commit_before>// +build linux

package mount

// MakeShared ensures a mounted filesystem has the SHARED mount option enabled.
// See the supported options in flags.go for further reference.
func MakeShared(mountPoint string) error {
	return ensureMountedAs(mountPoint, "shared")
}

// MakeRShared ensures a mounted filesystem has the RSHARED mount option enabled.
// See the supported options in flags.go for further reference.
func MakeRShared(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rshared")
}

// MakePrivate ensures a mounted filesystem has the PRIVATE mount option enabled.
// See the supported options in flags.go for further reference.
func MakePrivate(mountPoint string) error {
	return ensureMountedAs(mountPoint, "private")
}

// MakeRPrivate ensures a mounted filesystem has the RPRIVATE mount option
// enabled. See the supported options in flags.go for further reference.
func MakeRPrivate(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rprivate")
}

// MakeSlave ensures a mounted filesystem has the SLAVE mount option enabled.
// See the supported options in flags.go for further reference.
func MakeSlave(mountPoint string) error {
	return ensureMountedAs(mountPoint, "slave")
}

// MakeRSlave ensures a mounted filesystem has the RSLAVE mount option enabled.
// See the supported options in flags.go for further reference.
func MakeRSlave(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rslave")
}

// MakeUnbindable ensures a mounted filesystem has the UNBINDABLE mount option
// enabled. See the supported options in flags.go for further reference.
func MakeUnbindable(mountPoint string) error {
	return ensureMountedAs(mountPoint, "unbindable")
}

// MakeRUnbindable ensures a mounted filesystem has the RUNBINDABLE mount
// option enabled. See the supported options in flags.go for further reference.
func MakeRUnbindable(mountPoint string) error {
	return ensureMountedAs(mountPoint, "runbindable")
}

func ensureMountedAs(mountPoint, options string) error {
	mounted, err := Mounted(mountPoint)
	if err != nil {
		return err
	}

	if !mounted {
		if err := Mount(mountPoint, mountPoint, "none", "bind,rw"); err != nil {
			return err
		}
	}
	mounted, err = Mounted(mountPoint)
	if err != nil {
		return err
	}

	return ForceMount("", mountPoint, "none", options)
}
<commit_msg>fix variables that werent being called<commit_after>// +build linux

package mount

// MakeShared ensures a mounted filesystem has the SHARED mount option enabled.
// See the supported options in flags.go for further reference.
func MakeShared(mountPoint string) error {
	return ensureMountedAs(mountPoint, "shared")
}

// MakeRShared ensures a mounted filesystem has the RSHARED mount option enabled.
// See the supported options in flags.go for further reference.
func MakeRShared(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rshared")
}

// MakePrivate ensures a mounted filesystem has the PRIVATE mount option enabled.
// See the supported options in flags.go for further reference.
func MakePrivate(mountPoint string) error {
	return ensureMountedAs(mountPoint, "private")
}

// MakeRPrivate ensures a mounted filesystem has the RPRIVATE mount option
// enabled. See the supported options in flags.go for further reference.
func MakeRPrivate(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rprivate")
}

// MakeSlave ensures a mounted filesystem has the SLAVE mount option enabled.
// See the supported options in flags.go for further reference.
func MakeSlave(mountPoint string) error {
	return ensureMountedAs(mountPoint, "slave")
}

// MakeRSlave ensures a mounted filesystem has the RSLAVE mount option enabled.
// See the supported options in flags.go for further reference.
func MakeRSlave(mountPoint string) error {
	return ensureMountedAs(mountPoint, "rslave")
}

// MakeUnbindable ensures a mounted filesystem has the UNBINDABLE mount option
// enabled. See the supported options in flags.go for further reference.
func MakeUnbindable(mountPoint string) error {
	return ensureMountedAs(mountPoint, "unbindable")
}

// MakeRUnbindable ensures a mounted filesystem has the RUNBINDABLE mount
// option enabled. See the supported options in flags.go for further reference.
func MakeRUnbindable(mountPoint string) error {
	return ensureMountedAs(mountPoint, "runbindable")
}

func ensureMountedAs(mountPoint, options string) error {
	mounted, err := Mounted(mountPoint)
	if err != nil {
		return err
	}

	if !mounted {
		if err := Mount(mountPoint, mountPoint, "none", "bind,rw"); err != nil {
			return err
		}
	}
	if _, err = Mounted(mountPoint); err != nil {
		return err
	}

	return ForceMount("", mountPoint, "none", options)
}
<|endoftext|>


<commit_before>// Ogo

package ogo

import (
	"fmt"
	"time"
)

/* {{{ func OmqSet(key, value string, expire int) error
 *
 */
func OmqSet(key, value string, expire int) error {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "SET", "redis", key, value, expire); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] == "OK" {
				return nil
			}
		} else {
			Info("set %s error: %s", key, e)
			return e
		}
	} else {
		return e
	}
	return fmt.Errorf("set failed")
}

/* }}} */

/* {{{ func OmqGet(key string) (string, error)
 *
 */
func OmqGet(key string) (string, error) {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "GET", "redis", key); e == nil {
			Debug("Received: %s", reply)
			if reply[0] == "OK" {
				return reply[1], nil
			}
		} else {
			Info("get %s error: %s", key, e)
			return "", e
		}
	} else {
		return "", e
	}
	return "", fmt.Errorf("get %s failed", key)
}

/* }}} */

/* {{{ func OmqDel(key string) (string, error)
 *
 */
func OmqDel(key string) (string, error) {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "DEL", "redis", key); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] != "" {
				return reply[0], nil
			}
		} else {
			Info("del %s error: %s", key, e)
			return "", e
		}
	} else {
		return "", e
	}
	return "", fmt.Errorf("del %s failed", key)
}

/* }}} */

/* {{{ func OmqTask(msg ...string) error
 *
 */
func OmqTask(msg ...string) error {
	if requester, e := OmqPool().Get(); e == nil && len(msg) > 1 {
		defer requester.Close()
		key := msg[0]
		values := msg[1:]
		if reply, e := requester.Do(10000*time.Millisecond, "TASK", key, values); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] == "OK" {
				return nil
			}
		} else {
			Info("task %s error: %s", key, e)
			return e
		}
	} else {
		return e
	}
	return fmt.Errorf("push task failed")
}

/* }}} */
<commit_msg>update<commit_after>// Ogo

package ogo

import (
	"fmt"
	"time"
)

/* {{{ func OmqSet(key, value string, expire int) error
 *
 */
func OmqSet(key, value string, expire int) error {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "SET", "redis", key, value, expire); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] == "OK" {
				return nil
			}
		} else {
			Info("set %s error: %s", key, e)
			return e
		}
	} else {
		return e
	}
	return fmt.Errorf("set failed")
}

/* }}} */

/* {{{ func OmqGet(key string) (string, error)
 *
 */
func OmqGet(key string) (string, error) {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "GET", "redis", key); e == nil {
			Debug("Received: %s", reply)
			if reply[0] == "OK" {
				return reply[1], nil
			}
		} else {
			Info("get %s error: %s", key, e)
			return "", e
		}
	} else {
		return "", e
	}
	return "", fmt.Errorf("get %s failed", key)
}

/* }}} */

/* {{{ func OmqDel(key string) (string, error)
 *
 */
func OmqDel(key string) (string, error) {
	if requester, e := OmqPool().Get(); e == nil {
		defer requester.Close()
		if reply, e := requester.Do(10000*time.Millisecond, "DEL", "redis", key); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] != "" {
				return reply[0], nil
			}
		} else {
			Info("del %s error: %s", key, e)
			return "", e
		}
	} else {
		return "", e
	}
	return "", fmt.Errorf("del %s failed", key)
}

/* }}} */

/* {{{ func OmqTask(msg ...string) error
 *
 */
func OmqTask(msg ...string) error {
	if requester, e := OmqPool().Get(); e == nil && len(msg) > 1 {
		defer requester.Close()
		key := msg[0]
		values := msg[1:]
		if reply, e := requester.Do(10000*time.Millisecond, "TASK", key, values); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] == "OK" {
				return nil
			}
		} else {
			Info("task %s error: %s", key, e)
			return e
		}
	} else {
		return e
	}
	return fmt.Errorf("push task failed")
}

/* }}} */

/* {{{ func OmqPop(msg ...string) error
 *
 */
func OmqPop(msg ...string) error {
	if requester, e := OmqPool().Get(); e == nil && len(msg) > 1 {
		defer requester.Close()
		key := msg[0]
		values := msg[1:]
		if reply, e := requester.Do(10000*time.Millisecond, "POP", key, values); e == nil {
			Debug("Received: %s", reply[0])
			if reply[0] == "OK" {
				return nil
			}
		} else {
			Info("task %s error: %s", key, e)
			return e
		}
	} else {
		return e
	}
	return fmt.Errorf("pop task failed")
}

/* }}} */
<|endoftext|>


<commit_before>package fineline

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"strings"
)

const (
	// this is intentionally the zero value
	// it saves us a lot of typing
	opPutc = iota
	opPuts
	opBackspace
	opDelete
	opDeleteToBeginning
	opDeleteToEnd
	opClear
	opHome
	opEnd
	opRight
	opLeft
	opUp
	opDown
	opComplete
	opCancel
	opEof
	opSubmit
	opTranspose
	opEscape
	noop
)

var keyMap = [...]int{
	1:    opHome,      // ctrl-a
	2:    opLeft,      // ctrl-b
	3:    opCancel,    // ctrl-c
	4:    opEof,       // ctrl-d
	5:    opEnd,       // ctrl-e
	6:    opRight,     // ctrl-f
	7:    noop,        // ctrl-g
	8:    opBackspace, // ctrl-h
	'\t': opComplete,
	'\r': opSubmit,
	11:   opDeleteToEnd, // ctrl-k
	12:   opClear,       // ctrl-l
	'\n': opSubmit,
	14:   noop,                // ctrl-n
	15:   opSubmit,            // ctrl-o
	16:   opUp,                // ctrl-p
	17:   noop,                // ctrl-q; should be quoted insert
	18:   noop,                // ctrl-r; should be reverse history search
	19:   noop,                // ctrl-s; should be forward history search?
	20:   opTranspose,         // ctrl-t
	21:   opDeleteToBeginning, // ctrl-u
	22:   noop,                // ctrl-v; should be quoted insert
	23:   noop,                // ctrl-w; should be kill word
	24:   noop,                // ctrl-x; should be something?
	25:   noop,                // ctrl-y; should be yank
	26:   noop,                // ctrl-z
	27:   opEscape,
	127:  opBackspace,
}

var cancelled = errors.New("line cancelled")

func (l *LineReader) exec(r *bufio.Reader, op int, c rune) (bool, error) {
	switch op {
	case opPutc:
		l.putc(c)
	case opHome:
		l.home()
	case opLeft:
		l.left()
	case opCancel:
		return false, cancelled
	case opEof:
		if l.pos < l.buf.len()-1 {
			l.delete()
			break
		}
		return false, io.EOF
	case opEnd:
		l.end()
	case opRight:
		l.right()
	case opBackspace:
		l.backspace()
	case opComplete:
		if l.c != nil {
			l.complete()
		} else {
			l.putc(c)
		}
	case opDeleteToEnd:
		l.deleteToEnd()
	case opClear:
		l.clearScreen()
	case opSubmit:
		l.putc('\n')
		l.setCursor(0, -l.y)
		return false, nil
	case opTranspose:
		l.transpose()
	case opDeleteToBeginning:
		l.deleteToBeginning()
	case opEscape:
		var seq [2]byte
		r.Read(seq[:])
		if seq[0] == 91 {
			if 48 < seq[1] && seq[1] < 55 {
				// extended escape
				var seq2 [2]byte
				_, err := r.Read(seq2[:])
				if err != nil {
					break
				}
				if seq[1] == 51 && seq2[0] == 126 {
					// delete
					l.delete()
				}
				break
			}
			switch seq[1] {
			case 65:
				// down arrow
				l.currentEntry++
				if l.currentEntry > len(l.history) {
					l.currentEntry = 0
				}
			case 66:
				// up arrow
				l.currentEntry--
				if l.currentEntry < 0 {
					l.currentEntry = len(l.history) - 1
				}
			case 67:
				// right arrow
				l.right()
			case 68:
				// left arrow
				l.left()
			case 70:
				// end
				l.end()
			case 72:
				// home
				l.home()
			}
		} else if seq[0] == 79 {
			switch seq[1] {
			case 70:
				// end
				l.end()
			case 72:
				// home
				l.home()
			}
		}
	}
	return true, nil
}

func (l *LineReader) putc(c rune) {
	l.buf.WriteRune(c, l.pos)
	l.pos++
	l.refreshLine()
}

func (l *LineReader) puts(s string) {
	l.buf.WriteString(s, l.pos)
	l.pos += len(s)
	l.refreshLine()
}

// finds the longest common string between the end of the first string and the
// beginning of the second
func findIntersect(head, tail string) string {
	n, m := len(head), len(tail)
	if n == 0 || m == 0 {
		return ""
	}
	for len(tail) > 0 {
		pos := strings.LastIndex(head, tail)
		if pos >= 0 {
			return head[pos:]
		}
		tail = tail[:len(tail)-1]
	}
	return ""
}

// finds the longest common prefix of two strings
func commonPrefix(x, y string) string {
	n, m := len(x), len(y)
	if m < n {
		n = m
	}
	i := 0
	for i < n && x[i] == y[i] {
		i++
	}
	return x[:i]
}

func (l *LineReader) complete() {
	if l.display {
		l.printCandidates()
		return
	}
	str := l.buf.String()
	candidates := l.c.Complete(str[:l.pos])
	n := len(candidates)
	if n == 0 {
		return
	}
	var complete string
	if n == 1 {
		complete = candidates[0]
	} else {
		// look for a common prefix to see if we can fill in anything
		prefix := commonPrefix(candidates[0], candidates[1])
		for i := 2; i < n && prefix != str; i++ {
			prefix = commonPrefix(prefix, candidates[i])
		}
		complete = prefix
		l.display = true
		l.candidates = candidates
	}
	inter := findIntersect(str, complete)
	if inter == complete {
		return
	}
	l.puts(complete[len(inter):])
	l.refreshLine()
}

func (l *LineReader) backspace() {
	if l.pos > 0 {
		l.pos--
		l.buf.remove(l.pos)
		l.refreshLine()
	}
}

// delete the character in front of the cursor, like the delete key
func (l *LineReader) delete() {
	l.buf.remove(l.pos)
	l.refreshLine()
}

func (l *LineReader) deleteToBeginning() {
	l.buf.pretruncate(l.pos)
	l.pos = 0
	l.refreshLine()
}

func (l *LineReader) deleteToEnd() {
	l.buf.truncate(l.pos)
	l.refreshLine()
}

func (l *LineReader) home() {
	l.pos = 0
	l.refreshLine()
}

func (l *LineReader) end() {
	l.pos = l.buf.len()
	l.refreshLine()
}

func (l *LineReader) transpose() {
	l.buf.transpose(l.pos)
	l.refreshLine()
}

// move the cursor left
func (l *LineReader) left() {
	if l.pos > 0 {
		l.pos--
		l.refreshLine()
	}
}

// move the cursor right
func (l *LineReader) right() {
	if l.pos < l.buf.len() {
		l.pos++
		l.refreshLine()
	}
}

func (l *LineReader) refreshLine() {
	// move to origin of the current line
	l.setCursor(0, -l.y)
	// assuming the prompt won't wrap
	fmt.Print(l.Prompt)
	bufStr := l.buf.String()
	n := len(bufStr)
	pl := len(l.Prompt)
	if n > l.cols-pl {
		n = l.cols - pl
	}
	fmt.Print(bufStr[:n])
	bufStr = bufStr[n:]
	l.lines = 0
	wrapCursor := n == l.cols-pl
	n = len(bufStr)
	for n > 0 {
		if n > l.cols {
			n = l.cols
		}
		fmt.Print(bufStr[:n])
		bufStr = bufStr[n:]
		wrapCursor = n == l.cols
		n = len(bufStr)
		l.lines++
	}
	l.eraseToEnd()
	if wrapCursor {
		l.lines++
		// move to next line
		fmt.Print("\n")
	}
	x := (pl + l.pos) % l.cols
	l.y = (pl + l.pos) / l.cols
	l.setCursor(x, l.y-l.lines)
}
<commit_msg>Get tab completion working better<commit_after>package fineline

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"strings"
)

const (
	// this is intentionally the zero value
	// it saves us a lot of typing
	opPutc = iota
	opPuts
	opBackspace
	opDelete
	opDeleteToBeginning
	opDeleteToEnd
	opClear
	opHome
	opEnd
	opRight
	opLeft
	opUp
	opDown
	opComplete
	opCancel
	opEof
	opSubmit
	opTranspose
	opEscape
	noop
)

var keyMap = [...]int{
	1:    opHome,      // ctrl-a
	2:    opLeft,      // ctrl-b
	3:    opCancel,    // ctrl-c
	4:    opEof,       // ctrl-d
	5:    opEnd,       // ctrl-e
	6:    opRight,     // ctrl-f
	7:    noop,        // ctrl-g
	8:    opBackspace, // ctrl-h
	'\t': opComplete,
	'\r': opSubmit,
	11:   opDeleteToEnd, // ctrl-k
	12:   opClear,       // ctrl-l
	'\n': opSubmit,
	14:   noop,                // ctrl-n
	15:   opSubmit,            // ctrl-o
	16:   opUp,                // ctrl-p
	17:   noop,                // ctrl-q; should be quoted insert
	18:   noop,                // ctrl-r; should be reverse history search
	19:   noop,                // ctrl-s; should be forward history search?
	20:   opTranspose,         // ctrl-t
	21:   opDeleteToBeginning, // ctrl-u
	22:   noop,                // ctrl-v; should be quoted insert
	23:   noop,                // ctrl-w; should be kill word
	24:   noop,                // ctrl-x; should be something?
	25:   noop,                // ctrl-y; should be yank
	26:   noop,                // ctrl-z
	27:   opEscape,
	127:  opBackspace,
}

var cancelled = errors.New("line cancelled")

func (l *LineReader) exec(r *bufio.Reader, op int, c rune) (bool, error) {
	if op != opComplete {
		l.display = false
	}
	switch op {
	case opPutc:
		l.putc(c)
	case opHome:
		l.home()
	case opLeft:
		l.left()
	case opCancel:
		return false, cancelled
	case opEof:
		if l.pos < l.buf.len()-1 {
			l.delete()
			break
		}
		return false, io.EOF
	case opEnd:
		l.end()
	case opRight:
		l.right()
	case opBackspace:
		l.backspace()
	case opComplete:
		if l.c != nil {
			l.complete()
		} else {
			l.putc(c)
		}
	case opDeleteToEnd:
		l.deleteToEnd()
	case opClear:
		l.clearScreen()
	case opSubmit:
		l.putc('\n')
		l.setCursor(0, -l.y)
		return false, nil
	case opTranspose:
		l.transpose()
	case opDeleteToBeginning:
		l.deleteToBeginning()
	case opEscape:
		var seq [2]byte
		r.Read(seq[:])
		if seq[0] == 91 {
			if 48 < seq[1] && seq[1] < 55 {
				// extended escape
				var seq2 [2]byte
				_, err := r.Read(seq2[:])
				if err != nil {
					break
				}
				if seq[1] == 51 && seq2[0] == 126 {
					// delete
					l.delete()
				}
				break
			}
			switch seq[1] {
			case 65:
				// down arrow
				l.currentEntry++
				if l.currentEntry > len(l.history) {
					l.currentEntry = 0
				}
			case 66:
				// up arrow
				l.currentEntry--
				if l.currentEntry < 0 {
					l.currentEntry = len(l.history) - 1
				}
			case 67:
				// right arrow
				l.right()
			case 68:
				// left arrow
				l.left()
			case 70:
				// end
				l.end()
			case 72:
				// home
				l.home()
			}
		} else if seq[0] == 79 {
			switch seq[1] {
			case 70:
				// end
				l.end()
			case 72:
				// home
				l.home()
			}
		}
	}
	return true, nil
}

func (l *LineReader) putc(c rune) {
	l.buf.WriteRune(c, l.pos)
	l.pos++
	l.refreshLine()
}

func (l *LineReader) puts(s string) {
	l.buf.WriteString(s, l.pos)
	l.pos += len(s)
	l.refreshLine()
}

// finds the longest common string between the end of the first string and the
// beginning of the second
func findIntersect(head, tail string) string {
	n, m := len(head), len(tail)
	if n == 0 || m == 0 {
		return ""
	}
	for len(tail) > 0 {
		pos := strings.LastIndex(head, tail)
		if pos >= 0 {
			return head[pos:]
		}
		tail = tail[:len(tail)-1]
	}
	return ""
}

// finds the longest common prefix of two strings
func commonPrefix(x, y string) string {
	n, m := len(x), len(y)
	if m < n {
		n = m
	}
	i := 0
	for i < n && x[i] == y[i] {
		i++
	}
	return x[:i]
}

func (l *LineReader) complete() {
	if l.display {
		l.printCandidates()
		return
	}
	str := l.buf.String()
	candidates := l.c.Complete(str[:l.pos])
	n := len(candidates)
	if n == 0 {
		return
	}
	var complete string
	if n == 1 {
		complete = candidates[0]
	} else {
		// look for a common prefix to see if we can fill in anything
		prefix := commonPrefix(candidates[0], candidates[1])
		for i := 2; i < n && prefix != str; i++ {
			prefix = commonPrefix(prefix, candidates[i])
		}
		complete = prefix
		l.display = true
		l.candidates = candidates
	}
	inter := findIntersect(str, complete)
	if inter == complete {
		return
	}
	l.puts(complete[len(inter):])
	l.refreshLine()
}

func (l *LineReader) backspace() {
	if l.pos > 0 {
		l.pos--
		l.buf.remove(l.pos)
		l.refreshLine()
	}
}

// delete the character in front of the cursor, like the delete key
func (l *LineReader) delete() {
	l.buf.remove(l.pos)
	l.refreshLine()
}

func (l *LineReader) deleteToBeginning() {
	l.buf.pretruncate(l.pos)
	l.pos = 0
	l.refreshLine()
}

func (l *LineReader) deleteToEnd() {
	l.buf.truncate(l.pos)
	l.refreshLine()
}

func (l *LineReader) home() {
	l.pos = 0
	l.refreshLine()
}

func (l *LineReader) end() {
	l.pos = l.buf.len()
	l.refreshLine()
}

func (l *LineReader) transpose() {
	l.buf.transpose(l.pos)
	l.refreshLine()
}

// move the cursor left
func (l *LineReader) left() {
	if l.pos > 0 {
		l.pos--
		l.refreshLine()
	}
}

// move the cursor right
func (l *LineReader) right() {
	if l.pos < l.buf.len() {
		l.pos++
		l.refreshLine()
	}
}

func (l *LineReader) refreshLine() {
	// move to origin of the current line
	l.setCursor(0, -l.y)
	// assuming the prompt won't wrap
	fmt.Print(l.Prompt)
	bufStr := l.buf.String()
	n := len(bufStr)
	pl := len(l.Prompt)
	if n > l.cols-pl {
		n = l.cols - pl
	}
	fmt.Print(bufStr[:n])
	bufStr = bufStr[n:]
	l.lines = 0
	wrapCursor := n == l.cols-pl
	n = len(bufStr)
	for n > 0 {
		if n > l.cols {
			n = l.cols
		}
		fmt.Print(bufStr[:n])
		bufStr = bufStr[n:]
		wrapCursor = n == l.cols
		n = len(bufStr)
		l.lines++
	}
	l.eraseToEnd()
	if wrapCursor {
		l.lines++
		// move to next line
		fmt.Print("\n")
	}
	x := (pl + l.pos) % l.cols
	l.y = (pl + l.pos) / l.cols
	l.setCursor(x, l.y-l.lines)
}
<|endoftext|>


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

package w32

/*
#include <windows.h>

typedef struct _PDH_FMT_COUNTERVALUE {
  DWORD  CStatus;
  double doubleValue;
} PDH_FMT_COUNTERVALUE, *PPDH_FMT_COUNTERVALUE;

typedef struct _PDH_FMT_COUNTERVALUE_ITEM {
  LPTSTR               szName;
  PDH_FMT_COUNTERVALUE FmtValue;
} PDH_FMT_COUNTERVALUE_ITEM, *PPDH_FMT_COUNTERVALUE_ITEM;

PDH_FMT_COUNTERVALUE_ITEM* getCounterValueBuffer(DWORD buffersize){
	 PDH_FMT_COUNTERVALUE_ITEM *pItems = NULL;
	 pItems = (PDH_FMT_COUNTERVALUE_ITEM *) malloc(buffersize);

	return pItems;
}
*/
import "C"

import (
	"fmt"
	"reflect"
	"syscall"
	"unsafe"
)

const (
	PDH_SUCCESS = 0x00
	// Status codes returned from PDH functions
	PDH_MORE_DATA                 = 0x800007D2
	PDH_CSTATUS_BAD_COUNTERNAME   = 0xC0000BC0
	PDH_CSTATUS_NO_COUNTER        = 0xC0000BB9
	PDH_CSTATUS_NO_COUNTERNAME    = 0xC0000BBF
	PDH_CSTATUS_NO_INSTANCE       = 0x800007D1
	PDH_CSTATUS_NO_MACHINE        = 0x800007D0
	PDH_CSTATUS_NO_OBJECT         = 0xC0000BB8
	PDH_FUNCTION_NOT_FOUND        = 0xC0000BBE
	PDH_INVALID_ARGUMENT          = 0xC0000BBD
	PDH_INVALID_HANDLE            = 0xC0000BBC
	PDH_MEMORY_ALLOCATION_FAILURE = 0xC0000BBB

	// Argument options
	PERF_DETAIL_NOVICE   = 100
	PERF_DETAIL_ADVANCED = 200
	PERF_DETAIL_EXPERT   = 300
	PERF_DETAIL_WIZARD   = 400

	// Counter Formats options
	PDH_FMT_LONG   = 0x00000100 // long integer
	PDH_FMT_DOUBLE = 0x00000200 // double precision float
	PDH_FMT_LARGE  = 0x00000400 // 64-bit integer

	PDH_FMT_NOSCALE = 0x00001000 // Do not apply the counter's default scaling factor.
	PDH_FMT_1000    = 0x00002000 // Multiply the actual value by 1,000.

	/* Counter values greater than 100 (for example,
	   counter values measuring the processor load on multiprocessor computers)
	   will not be reset to 100. The default behavior is that
	   counter values are capped at a value of 100. */
	PDH_FMT_NOCAP100 = 0x00008000
)

var (
	modpdh = syscall.NewLazyDLL("pdh.dll")

	procPdhEnumObjects              = modpdh.NewProc("PdhEnumObjectsA")
	procPdhEnumObjectItems          = modpdh.NewProc("PdhEnumObjectItemsA")
	procPdhOpenQuery                = modpdh.NewProc("PdhOpenQueryA")
	procPdhAddCounter               = modpdh.NewProc("PdhAddCounterA")
	procPdhCollectQueryData         = modpdh.NewProc("PdhCollectQueryData")
	procPdhValidatePath             = modpdh.NewProc("PdhValidatePathA")
	procPdhGetFormattedCounterValue = modpdh.NewProc("PdhGetFormattedCounterValue")
	procPdhGetFormattedCounterArray = modpdh.NewProc("PdhGetFormattedCounterArrayA")
)

type PDH_ERROR struct {
	Code        uint32
	Description string
}

func (err *PDH_ERROR) Error() string {
	return err.Description
}

func (err *PDH_ERROR) String() string {
	return err.Description
}

func PdhErrorCode(sCode uint32) *PDH_ERROR {
	switch sCode {
	case 0:
		return nil
	case PDH_MORE_DATA:
		return &PDH_ERROR{Code: sCode,
			Description: "There is more data to return than would fit in the supplied buffer. Allocate a larger buffer and call the function again."}
	case PDH_CSTATUS_BAD_COUNTERNAME:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to parse the counter path. Check the format and syntax of the specified path."}
	case PDH_CSTATUS_NO_COUNTER:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified counter could not be found."}
	case PDH_CSTATUS_NO_COUNTERNAME:
		return &PDH_ERROR{Code: sCode,
			Description: "No counter was specified."}
	case PDH_CSTATUS_NO_INSTANCE:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified instance is not present."}
	case PDH_CSTATUS_NO_MACHINE:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to connect to the specified computer, or the computer is offline."}
	case PDH_CSTATUS_NO_OBJECT:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified object is not found on the system."}
	case PDH_FUNCTION_NOT_FOUND:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to find the specified function."}
	case PDH_INVALID_ARGUMENT:
		return &PDH_ERROR{Code: sCode,
			Description: "A required argument is missing or incorrect."}
	case PDH_INVALID_HANDLE:
		return &PDH_ERROR{Code: sCode,
			Description: "The handle is not a valid PDH object."}
	case PDH_MEMORY_ALLOCATION_FAILURE:
		return &PDH_ERROR{Code: sCode,
			Description: "A PDH function could not allocate enough temporary memory to complete the operation. Close some applications or extend the page file and retry the function."}
	default:
		return &PDH_ERROR{Code: sCode, Description: fmt.Sprintf("Unimplemented error, code 0x%X", sCode)}
	}
}

func PdhEnumObjects(detailLevel uint32, refreshCache bool, bufferlen uint32) (uint32, []byte, uint32) {
	var buf []byte
	var buflen uint32
	var bufPtr *byte
	if bufferlen > 0 {
		buf = make([]byte, bufferlen+1)
		bufPtr = &buf[0]
		buf[bufferlen] = 0
		buflen = uint32(bufferlen)
	}

	refreshCacheWin := TRUE
	if !refreshCache {
		refreshCacheWin = FALSE
	}

	ret, _, _ := procPdhEnumObjects.Call(
		uintptr(0),
		uintptr(0),
		uintptr(unsafe.Pointer(bufPtr)),
		uintptr(unsafe.Pointer(&buflen)),
		uintptr(detailLevel),
		uintptr(refreshCacheWin),
	)

	if buflen > 0 && uint32(ret) == 0 {
		return uint32(ret), buf[:buflen-1], buflen
	}
	return uint32(ret), []byte{}, buflen
}

func PdhEnumObjectItems(object string, counterBufferLen, instanceBufferLen uint32,
	detailLevel uint32) (status uint32, outCountBuf []byte, outCountLen uint32,
	outInstBuf []byte, outInstLen uint32) {
	var counterBuf []byte
	var counterBufLen uint32
	var counterBufPtr *byte
	if counterBufferLen > 0 {
		counterBuf = make([]byte, counterBufferLen+1)
		counterBufPtr = &counterBuf[0]
		counterBuf[counterBufferLen] = 0
		counterBufLen = uint32(counterBufferLen)
	}

	var instanceBuf []byte
	var instanceBufLen uint32
	var instanceBufPtr *byte
	if instanceBufferLen > 0 {
		instanceBuf = make([]byte, instanceBufferLen+1)
		instanceBufPtr = &instanceBuf[0]
		instanceBuf[instanceBufferLen] = 0
		instanceBufLen = uint32(instanceBufferLen)
	}

	objName, _ := syscall.BytePtrFromString(object)
	ret, _, _ := procPdhEnumObjectItems.Call(
		uintptr(0),
		uintptr(0),
		uintptr(unsafe.Pointer(objName)),
		uintptr(unsafe.Pointer(counterBufPtr)),
		uintptr(unsafe.Pointer(&counterBufLen)),
		uintptr(unsafe.Pointer(instanceBufPtr)),
		uintptr(unsafe.Pointer(&instanceBufLen)),
		uintptr(detailLevel),
		uintptr(uint32(0)),
	)

	if uint32(ret) != 0 {
		return uint32(ret), []byte{}, counterBufLen, []byte{}, instanceBufLen
	}
	if counterBufLen > 0 {
		outCountBuf = counterBuf[:counterBufLen-1]
		outCountLen = counterBufLen
	}
	if instanceBufLen > 0 {
		outInstBuf = instanceBuf[:instanceBufLen-1]
		outInstLen = instanceBufLen
	}

	return uint32(ret), outCountBuf, outCountLen, outInstBuf, outInstLen
}

func PdhOpenQuery() (uint32, HANDLE) {
	var pdhQuery HANDLE
	var userData uintptr
	ret, _, _ := procPdhOpenQuery.Call(
		uintptr(0),
		uintptr(userData),
		uintptr(unsafe.Pointer(&pdhQuery)),
	)

	return uint32(ret), pdhQuery
}

func PdhAddCounter(pdhQuery HANDLE, counterPath string) (uint32, HANDLE) {
	var counter HANDLE
	var userData uintptr

	path, _ := syscall.BytePtrFromString(counterPath)

	ret, _, _ := procPdhAddCounter.Call(
		uintptr(pdhQuery),
		uintptr(unsafe.Pointer(path)),
		uintptr(userData),
		uintptr(unsafe.Pointer(&counter)),
	)

	return uint32(ret), counter
}

func PdhCollectQueryData(pdhQuery HANDLE) (uint32, HANDLE) {
	ret, _, _ := procPdhCollectQueryData.Call(
		uintptr(pdhQuery),
	)

	return uint32(ret), pdhQuery
}

func PdhGetFormattedCounterValue(hCounter HANDLE) (status uint32, outBuffer PDH_FMT_COUNTERVALUE) {
	counterType := uint32(0)
	var buffer C.PDH_FMT_COUNTERVALUE
	ret, _, _ := procPdhGetFormattedCounterValue.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&counterType)),
		uintptr(unsafe.Pointer(&buffer)),
	)

	if status = uint32(ret); status != 0 {
		return
	}

	outBuffer = PDH_FMT_COUNTERVALUE{
		Status: uint32(buffer.CStatus),
		Value:  float64(buffer.doubleValue),
	}

	return
}

func PdhGetFormattedCounterArray(hCounter HANDLE) (status uint32,
	respLen uint32, numItems uint32, outBuffer []PDH_FMT_COUNTERVALUE_ITEM) {

	// Call once to get buffer size
	ret, _, _ := procPdhGetFormattedCounterArray.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&respLen)),
		uintptr(unsafe.Pointer(&numItems)),
		uintptr(0),
	)

	if status = uint32(ret); status != PDH_MORE_DATA {
		return
	}

	itemBuffer := C.getCounterValueBuffer(C.DWORD(respLen))
	defer C.free(unsafe.Pointer(itemBuffer))

	ret, _, _ = procPdhGetFormattedCounterArray.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&respLen)),
		uintptr(unsafe.Pointer(&numItems)),
		uintptr(unsafe.Pointer(itemBuffer)),
	)

	if status = uint32(ret); status != 0 {
		return
	}

	bufferArray := *(*[]C.PDH_FMT_COUNTERVALUE_ITEM)(unsafe.Pointer(&reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(itemBuffer)),
		Len:  int(numItems),
		Cap:  int(numItems),
	}))

	outBuffer = make([]PDH_FMT_COUNTERVALUE_ITEM, numItems)
	// Loop through the array and print the instance name and counter value.
	for i := uint32(0); i < numItems; i++ {
		outBuffer[i].Name = C.GoString((*C.char)(unsafe.Pointer(bufferArray[i].szName)))
		outBuffer[i].Item = PDH_FMT_COUNTERVALUE{
			Status: uint32(bufferArray[i].FmtValue.CStatus),
			Value:  float64(bufferArray[i].FmtValue.doubleValue),
		}
	}

	return
}

func PdhValidatePath(processCounterPath string) uint32 {
	path, _ := syscall.BytePtrFromString(processCounterPath)

	ret, _, _ := procPdhValidatePath.Call(
		uintptr(unsafe.Pointer(path)),
	)

	return uint32(ret)
}
<commit_msg>Add closeQuery function<commit_after>// Copyright 2010-2012 The W32 Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package w32

/*
#include <windows.h>

typedef struct _PDH_FMT_COUNTERVALUE {
  DWORD  CStatus;
  double doubleValue;
} PDH_FMT_COUNTERVALUE, *PPDH_FMT_COUNTERVALUE;

typedef struct _PDH_FMT_COUNTERVALUE_ITEM {
  LPTSTR               szName;
  PDH_FMT_COUNTERVALUE FmtValue;
} PDH_FMT_COUNTERVALUE_ITEM, *PPDH_FMT_COUNTERVALUE_ITEM;

PDH_FMT_COUNTERVALUE_ITEM* getCounterValueBuffer(DWORD buffersize){
	 PDH_FMT_COUNTERVALUE_ITEM *pItems = NULL;
	 pItems = (PDH_FMT_COUNTERVALUE_ITEM *) malloc(buffersize);

	return pItems;
}
*/
import "C"

import (
	"fmt"
	"reflect"
	"syscall"
	"unsafe"
)

const (
	PDH_SUCCESS = 0x00
	// Status codes returned from PDH functions
	PDH_MORE_DATA                 = 0x800007D2
	PDH_CSTATUS_BAD_COUNTERNAME   = 0xC0000BC0
	PDH_CSTATUS_NO_COUNTER        = 0xC0000BB9
	PDH_CSTATUS_NO_COUNTERNAME    = 0xC0000BBF
	PDH_CSTATUS_NO_INSTANCE       = 0x800007D1
	PDH_CSTATUS_NO_MACHINE        = 0x800007D0
	PDH_CSTATUS_NO_OBJECT         = 0xC0000BB8
	PDH_FUNCTION_NOT_FOUND        = 0xC0000BBE
	PDH_INVALID_ARGUMENT          = 0xC0000BBD
	PDH_INVALID_HANDLE            = 0xC0000BBC
	PDH_MEMORY_ALLOCATION_FAILURE = 0xC0000BBB

	// Argument options
	PERF_DETAIL_NOVICE   = 100
	PERF_DETAIL_ADVANCED = 200
	PERF_DETAIL_EXPERT   = 300
	PERF_DETAIL_WIZARD   = 400

	// Counter Formats options
	PDH_FMT_LONG   = 0x00000100 // long integer
	PDH_FMT_DOUBLE = 0x00000200 // double precision float
	PDH_FMT_LARGE  = 0x00000400 // 64-bit integer

	PDH_FMT_NOSCALE = 0x00001000 // Do not apply the counter's default scaling factor.
	PDH_FMT_1000    = 0x00002000 // Multiply the actual value by 1,000.

	/* Counter values greater than 100 (for example,
	   counter values measuring the processor load on multiprocessor computers)
	   will not be reset to 100. The default behavior is that
	   counter values are capped at a value of 100. */
	PDH_FMT_NOCAP100 = 0x00008000
)

var (
	modpdh = syscall.NewLazyDLL("pdh.dll")

	procPdhEnumObjects              = modpdh.NewProc("PdhEnumObjectsA")
	procPdhEnumObjectItems          = modpdh.NewProc("PdhEnumObjectItemsA")
	procPdhOpenQuery                = modpdh.NewProc("PdhOpenQueryA")
	procPdhCloseQuery                = modpdh.NewProc("PdhCloseQuery")
	procPdhAddCounter               = modpdh.NewProc("PdhAddCounterA")
	procPdhCollectQueryData         = modpdh.NewProc("PdhCollectQueryData")
	procPdhValidatePath             = modpdh.NewProc("PdhValidatePathA")
	procPdhGetFormattedCounterValue = modpdh.NewProc("PdhGetFormattedCounterValue")
	procPdhGetFormattedCounterArray = modpdh.NewProc("PdhGetFormattedCounterArrayA")
)

type PDH_ERROR struct {
	Code        uint32
	Description string
}

func (err *PDH_ERROR) Error() string {
	return err.Description
}

func (err *PDH_ERROR) String() string {
	return err.Description
}

func PdhErrorCode(sCode uint32) *PDH_ERROR {
	switch sCode {
	case 0:
		return nil
	case PDH_MORE_DATA:
		return &PDH_ERROR{Code: sCode,
			Description: "There is more data to return than would fit in the supplied buffer. Allocate a larger buffer and call the function again."}
	case PDH_CSTATUS_BAD_COUNTERNAME:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to parse the counter path. Check the format and syntax of the specified path."}
	case PDH_CSTATUS_NO_COUNTER:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified counter could not be found."}
	case PDH_CSTATUS_NO_COUNTERNAME:
		return &PDH_ERROR{Code: sCode,
			Description: "No counter was specified."}
	case PDH_CSTATUS_NO_INSTANCE:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified instance is not present."}
	case PDH_CSTATUS_NO_MACHINE:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to connect to the specified computer, or the computer is offline."}
	case PDH_CSTATUS_NO_OBJECT:
		return &PDH_ERROR{Code: sCode,
			Description: "The specified object is not found on the system."}
	case PDH_FUNCTION_NOT_FOUND:
		return &PDH_ERROR{Code: sCode,
			Description: "Unable to find the specified function."}
	case PDH_INVALID_ARGUMENT:
		return &PDH_ERROR{Code: sCode,
			Description: "A required argument is missing or incorrect."}
	case PDH_INVALID_HANDLE:
		return &PDH_ERROR{Code: sCode,
			Description: "The handle is not a valid PDH object."}
	case PDH_MEMORY_ALLOCATION_FAILURE:
		return &PDH_ERROR{Code: sCode,
			Description: "A PDH function could not allocate enough temporary memory to complete the operation. Close some applications or extend the page file and retry the function."}
	default:
		return &PDH_ERROR{Code: sCode, Description: fmt.Sprintf("Unimplemented error, code 0x%X", sCode)}
	}
}

func PdhEnumObjects(detailLevel uint32, refreshCache bool, bufferlen uint32) (uint32, []byte, uint32) {
	var buf []byte
	var buflen uint32
	var bufPtr *byte
	if bufferlen > 0 {
		buf = make([]byte, bufferlen+1)
		bufPtr = &buf[0]
		buf[bufferlen] = 0
		buflen = uint32(bufferlen)
	}

	refreshCacheWin := TRUE
	if !refreshCache {
		refreshCacheWin = FALSE
	}

	ret, _, _ := procPdhEnumObjects.Call(
		uintptr(0),
		uintptr(0),
		uintptr(unsafe.Pointer(bufPtr)),
		uintptr(unsafe.Pointer(&buflen)),
		uintptr(detailLevel),
		uintptr(refreshCacheWin),
	)

	if buflen > 0 && uint32(ret) == 0 {
		return uint32(ret), buf[:buflen-1], buflen
	}
	return uint32(ret), []byte{}, buflen
}

func PdhEnumObjectItems(object string, counterBufferLen, instanceBufferLen uint32,
	detailLevel uint32) (status uint32, outCountBuf []byte, outCountLen uint32,
	outInstBuf []byte, outInstLen uint32) {
	var counterBuf []byte
	var counterBufLen uint32
	var counterBufPtr *byte
	if counterBufferLen > 0 {
		counterBuf = make([]byte, counterBufferLen+1)
		counterBufPtr = &counterBuf[0]
		counterBuf[counterBufferLen] = 0
		counterBufLen = uint32(counterBufferLen)
	}

	var instanceBuf []byte
	var instanceBufLen uint32
	var instanceBufPtr *byte
	if instanceBufferLen > 0 {
		instanceBuf = make([]byte, instanceBufferLen+1)
		instanceBufPtr = &instanceBuf[0]
		instanceBuf[instanceBufferLen] = 0
		instanceBufLen = uint32(instanceBufferLen)
	}

	objName, _ := syscall.BytePtrFromString(object)
	ret, _, _ := procPdhEnumObjectItems.Call(
		uintptr(0),
		uintptr(0),
		uintptr(unsafe.Pointer(objName)),
		uintptr(unsafe.Pointer(counterBufPtr)),
		uintptr(unsafe.Pointer(&counterBufLen)),
		uintptr(unsafe.Pointer(instanceBufPtr)),
		uintptr(unsafe.Pointer(&instanceBufLen)),
		uintptr(detailLevel),
		uintptr(uint32(0)),
	)

	if uint32(ret) != 0 {
		return uint32(ret), []byte{}, counterBufLen, []byte{}, instanceBufLen
	}
	if counterBufLen > 0 {
		outCountBuf = counterBuf[:counterBufLen-1]
		outCountLen = counterBufLen
	}
	if instanceBufLen > 0 {
		outInstBuf = instanceBuf[:instanceBufLen-1]
		outInstLen = instanceBufLen
	}

	return uint32(ret), outCountBuf, outCountLen, outInstBuf, outInstLen
}

func PdhOpenQuery() (uint32, HANDLE) {
	var pdhQuery HANDLE
	var userData uintptr
	ret, _, _ := procPdhOpenQuery.Call(
		uintptr(0),
		uintptr(userData),
		uintptr(unsafe.Pointer(&pdhQuery)),
	)

	return uint32(ret), pdhQuery
}

func PdhCloseQuery(pdhQuery HANDLE) uint32 {
	ret, _, _ := procPdhCloseQuery.Call(
		uintptr(pdhQuery),
	)

	return uint32(ret)
}

func PdhAddCounter(pdhQuery HANDLE, counterPath string) (uint32, HANDLE) {
	var counter HANDLE
	var userData uintptr

	path, _ := syscall.BytePtrFromString(counterPath)

	ret, _, _ := procPdhAddCounter.Call(
		uintptr(pdhQuery),
		uintptr(unsafe.Pointer(path)),
		uintptr(userData),
		uintptr(unsafe.Pointer(&counter)),
	)

	return uint32(ret), counter
}

func PdhCollectQueryData(pdhQuery HANDLE) (uint32, HANDLE) {
	ret, _, _ := procPdhCollectQueryData.Call(
		uintptr(pdhQuery),
	)

	return uint32(ret), pdhQuery
}

func PdhGetFormattedCounterValue(hCounter HANDLE) (status uint32, outBuffer PDH_FMT_COUNTERVALUE) {
	counterType := uint32(0)
	var buffer C.PDH_FMT_COUNTERVALUE
	ret, _, _ := procPdhGetFormattedCounterValue.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&counterType)),
		uintptr(unsafe.Pointer(&buffer)),
	)

	if status = uint32(ret); status != 0 {
		return
	}

	outBuffer = PDH_FMT_COUNTERVALUE{
		Status: uint32(buffer.CStatus),
		Value:  float64(buffer.doubleValue),
	}

	return
}

func PdhGetFormattedCounterArray(hCounter HANDLE) (status uint32,
	respLen uint32, numItems uint32, outBuffer []PDH_FMT_COUNTERVALUE_ITEM) {

	// Call once to get buffer size
	ret, _, _ := procPdhGetFormattedCounterArray.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&respLen)),
		uintptr(unsafe.Pointer(&numItems)),
		uintptr(0),
	)

	if status = uint32(ret); status != PDH_MORE_DATA {
		return
	}

	itemBuffer := C.getCounterValueBuffer(C.DWORD(respLen))
	defer C.free(unsafe.Pointer(itemBuffer))

	ret, _, _ = procPdhGetFormattedCounterArray.Call(
		uintptr(hCounter),
		uintptr(uint32(PDH_FMT_DOUBLE|PDH_FMT_NOSCALE)),
		uintptr(unsafe.Pointer(&respLen)),
		uintptr(unsafe.Pointer(&numItems)),
		uintptr(unsafe.Pointer(itemBuffer)),
	)

	if status = uint32(ret); status != 0 {
		return
	}

	bufferArray := *(*[]C.PDH_FMT_COUNTERVALUE_ITEM)(unsafe.Pointer(&reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(itemBuffer)),
		Len:  int(numItems),
		Cap:  int(numItems),
	}))

	outBuffer = make([]PDH_FMT_COUNTERVALUE_ITEM, numItems)
	// Loop through the array and print the instance name and counter value.
	for i := uint32(0); i < numItems; i++ {
		outBuffer[i].Name = C.GoString((*C.char)(unsafe.Pointer(bufferArray[i].szName)))
		outBuffer[i].Item = PDH_FMT_COUNTERVALUE{
			Status: uint32(bufferArray[i].FmtValue.CStatus),
			Value:  float64(bufferArray[i].FmtValue.doubleValue),
		}
	}

	return
}

func PdhValidatePath(processCounterPath string) uint32 {
	path, _ := syscall.BytePtrFromString(processCounterPath)

	ret, _, _ := procPdhValidatePath.Call(
		uintptr(unsafe.Pointer(path)),
	)

	return uint32(ret)
}
<|endoftext|>