<commit_before>package slirc

import (
	"crypto/tls"
	"fmt"
	"log"
	"strings"
	"time"

	ircc "github.com/fluffle/goirc/client"

	"github.com/simonkern/slirc/slack"
)

// Bridge links an irc and a slack channel
type Bridge struct {
	SlackChan string
	IRCChan   string
	slack     *slack.Client
	irc       *ircc.Conn
}

type messager interface {
	Usernick() string
	Msg() string
	Chan() string
}

// IRCAuth stores authentification target and the message that needs to be send in order to auth
// e.g. "NickServ" and "IDENTIFY fooBarPassword"
type IRCAuth struct {
	Target string
	Msg    string
}

// NewBridge instantiates a Bridge object and sets up the required irc and slack clients
func NewBridge(slackBotToken, slackUserToken, slackChannel, ircServer, ircChannel, ircNick string, ircSSL bool, tlsConfig *tls.Config, ircAuth *IRCAuth) (bridge *Bridge) {
	sc := slack.NewClient(slackBotToken)

	sc.UserToken = slackUserToken

	ircCfg := ircc.NewConfig(ircNick, "slirc", "Powered by Slirc")
	ircCfg.QuitMessage = "Slack <-> IRC Bridge shutting down"
	ircCfg.Server = ircServer
	ircCfg.NewNick = func(n string) string {
		if n != ircNick && len(n) > len(ircNick)+2 {
			return ircNick
		}
		return n + "_"
	}
	if ircSSL {
		ircCfg.SSL = true
		if tlsConfig != nil {
			ircCfg.SSLConfig = tlsConfig
		}
	}
	c := ircc.Client(ircCfg)

	bridge = &Bridge{SlackChan: slackChannel, IRCChan: ircChannel, slack: sc, irc: c}

	// IRC Handlers
	c.HandleFunc(ircc.CONNECTED,
		func(conn *ircc.Conn, line *ircc.Line) {
			if ircAuth != nil {
				log.Println("IRC Authentication")
				<-time.After(5 * time.Second)
				conn.Privmsg(ircAuth.Target, ircAuth.Msg)
				<-time.After(3 * time.Second)
			}
			conn.Join(ircChannel)
			bridge.slack.Send(bridge.SlackChan, "Connected to IRC.")
			log.Println("Connected to IRC.")
		})

	c.HandleFunc(ircc.DISCONNECTED,
		func(conn *ircc.Conn, line *ircc.Line) {
			bridge.slack.Send(bridge.SlackChan, "Disconnected from IRC. Reconnecting...")
			log.Println("Disconnected from IRC. Reconnecting...")
			for {
				if err := conn.Connect(); err != nil {
					log.Println("IRC reconnect failed: ", err)
					log.Println("Trying again in 30 seconds...")
					time.Sleep(30 * time.Second)
					continue
				}
				// success
				break
			}
		})

	c.HandleFunc(ircc.PRIVMSG,
		func(conn *ircc.Conn, line *ircc.Line) {
			if line.Target() == bridge.IRCChan {
				msg := fmt.Sprintf("[%s]: %s", line.Nick, line.Text())
				bridge.slack.Send(bridge.SlackChan, msg)
			}
		})

	// thanks jn__
	c.HandleFunc(ircc.ACTION,
		func(conn *ircc.Conn, line *ircc.Line) {
			if line.Target() == bridge.IRCChan {
				msg := fmt.Sprintf(" * %s %s", line.Nick, line.Text())
				bridge.slack.Send(bridge.SlackChan, msg)
			}
		})

	// Slack Handlers
	sc.HandleFunc("shutdown",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Shutting down slack client")
			log.Println("Shutting down slack client")

		})

	sc.HandleFunc("disconnected",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Disconnected from Slack. Reconnecting...")
			log.Println("Disconnected from Slack. Reconnecting...")
			sc.Connect()

		})

	sc.HandleFunc("connected",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Connected to Slack.")
			log.Println("Connected to Slack.")
		})

	sc.HandleFunc("message",
		func(sc *slack.Client, e *slack.Event) {
			if e.Chan() == bridge.SlackChan && !sc.IsSelfMsg(e) && e.Text != "" && e.Usernick() != "rwthirc" {
				msg := fmt.Sprintf("[%s]: %s", e.Usernick(), e.Msg())
				// IRC has problems with newlines, therefore we split the message
				for _, line := range strings.SplitAfter(msg, "\n") {
					// we do not want to send empty lines...
					if strings.TrimSpace(line) != "" {
						bridge.irc.Privmsg(bridge.IRCChan, line)
					}
				}
			}

		})

	go func() {
		if err := c.Connect(); err != nil {
			log.Fatal("Could not connect to IRC: ", err)
		}
	}()
	go func() {
		if err := sc.Connect(); err != nil {
			log.Fatal("Could not connect to Slack: ", err)
		}
	}()
	return bridge
}
<commit_msg>delete condition that was used for debugging purposes<commit_after>package slirc

import (
	"crypto/tls"
	"fmt"
	"log"
	"strings"
	"time"

	ircc "github.com/fluffle/goirc/client"

	"github.com/simonkern/slirc/slack"
)

// Bridge links an irc and a slack channel
type Bridge struct {
	SlackChan string
	IRCChan   string
	slack     *slack.Client
	irc       *ircc.Conn
}

type messager interface {
	Usernick() string
	Msg() string
	Chan() string
}

// IRCAuth stores authentification target and the message that needs to be send in order to auth
// e.g. "NickServ" and "IDENTIFY fooBarPassword"
type IRCAuth struct {
	Target string
	Msg    string
}

// NewBridge instantiates a Bridge object and sets up the required irc and slack clients
func NewBridge(slackBotToken, slackUserToken, slackChannel, ircServer, ircChannel, ircNick string, ircSSL bool, tlsConfig *tls.Config, ircAuth *IRCAuth) (bridge *Bridge) {
	sc := slack.NewClient(slackBotToken)

	sc.UserToken = slackUserToken

	ircCfg := ircc.NewConfig(ircNick, "slirc", "Powered by Slirc")
	ircCfg.QuitMessage = "Slack <-> IRC Bridge shutting down"
	ircCfg.Server = ircServer
	ircCfg.NewNick = func(n string) string {
		if n != ircNick && len(n) > len(ircNick)+2 {
			return ircNick
		}
		return n + "_"
	}
	if ircSSL {
		ircCfg.SSL = true
		if tlsConfig != nil {
			ircCfg.SSLConfig = tlsConfig
		}
	}
	c := ircc.Client(ircCfg)

	bridge = &Bridge{SlackChan: slackChannel, IRCChan: ircChannel, slack: sc, irc: c}

	// IRC Handlers
	c.HandleFunc(ircc.CONNECTED,
		func(conn *ircc.Conn, line *ircc.Line) {
			if ircAuth != nil {
				log.Println("IRC Authentication")
				<-time.After(5 * time.Second)
				conn.Privmsg(ircAuth.Target, ircAuth.Msg)
				<-time.After(3 * time.Second)
			}
			conn.Join(ircChannel)
			bridge.slack.Send(bridge.SlackChan, "Connected to IRC.")
			log.Println("Connected to IRC.")
		})

	c.HandleFunc(ircc.DISCONNECTED,
		func(conn *ircc.Conn, line *ircc.Line) {
			bridge.slack.Send(bridge.SlackChan, "Disconnected from IRC. Reconnecting...")
			log.Println("Disconnected from IRC. Reconnecting...")
			for {
				if err := conn.Connect(); err != nil {
					log.Println("IRC reconnect failed: ", err)
					log.Println("Trying again in 30 seconds...")
					time.Sleep(30 * time.Second)
					continue
				}
				// success
				break
			}
		})

	c.HandleFunc(ircc.PRIVMSG,
		func(conn *ircc.Conn, line *ircc.Line) {
			if line.Target() == bridge.IRCChan {
				msg := fmt.Sprintf("[%s]: %s", line.Nick, line.Text())
				bridge.slack.Send(bridge.SlackChan, msg)
			}
		})

	// thanks jn__
	c.HandleFunc(ircc.ACTION,
		func(conn *ircc.Conn, line *ircc.Line) {
			if line.Target() == bridge.IRCChan {
				msg := fmt.Sprintf(" * %s %s", line.Nick, line.Text())
				bridge.slack.Send(bridge.SlackChan, msg)
			}
		})

	// Slack Handlers
	sc.HandleFunc("shutdown",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Shutting down slack client")
			log.Println("Shutting down slack client")

		})

	sc.HandleFunc("disconnected",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Disconnected from Slack. Reconnecting...")
			log.Println("Disconnected from Slack. Reconnecting...")
			sc.Connect()

		})

	sc.HandleFunc("connected",
		func(sc *slack.Client, e *slack.Event) {
			bridge.irc.Privmsg(bridge.IRCChan, "Connected to Slack.")
			log.Println("Connected to Slack.")
		})

	sc.HandleFunc("message",
		func(sc *slack.Client, e *slack.Event) {
			if e.Chan() == bridge.SlackChan && !sc.IsSelfMsg(e) && e.Text != "" {
				msg := fmt.Sprintf("[%s]: %s", e.Usernick(), e.Msg())
				// IRC has problems with newlines, therefore we split the message
				for _, line := range strings.SplitAfter(msg, "\n") {
					// we do not want to send empty lines...
					if strings.TrimSpace(line) != "" {
						bridge.irc.Privmsg(bridge.IRCChan, line)
					}
				}
			}

		})

	go func() {
		if err := c.Connect(); err != nil {
			log.Fatal("Could not connect to IRC: ", err)
		}
	}()
	go func() {
		if err := sc.Connect(); err != nil {
			log.Fatal("Could not connect to Slack: ", err)
		}
	}()
	return bridge
}
<|endoftext|>


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

package agent

import (
	"github.com/romana/core/common"
//	"github.com/romana/core/topology"

	"log"
	"net"
)

// NetworkConfig holds the agent's current configuration.
// This consists of data parsed from the config file as well as
// runtime or discovered configuration, such as the network
// config of the current host.
type NetworkConfig struct {
	// Current host network configuration
	hostIP net.IP
	romanaIP net.IP
	otherHosts []common.HostMessage
	dc               common.Datacenter
}

// EndpointNetmaskSize returns integer value (aka size) of endpoint netmask.
func (c *NetworkConfig) EndpointNetmaskSize() uint64 {
	// TODO make this depend on the IP version
	return 32 - uint64(c.dc.EndpointSpaceBits)
}

// PNetCIDR returns pseudo net cidr in net.IPNet format.
func (c *NetworkConfig) PNetCIDR() (cidr *net.IPNet, err error) {
	_, cidr, err = net.ParseCIDR(c.dc.Cidr)
	return
}

// TenantBits returns tenant bits value from POC config.
func (c *NetworkConfig) TenantBits() uint {
	return c.dc.TenantBits
}

// SegmentBits returns segment bits value from POC config.
func (c *NetworkConfig) SegmentBits() uint {
	return c.dc.SegmentBits
}

// EndpointBits returns endpoint bits value from POC config.
func (c *NetworkConfig) EndpointBits() uint {
	return c.dc.EndpointBits
}

// identifyCurrentHost discovers network configuration
// of the host we are running on.
// We need to know public IP and Romana gateway IP of the current host.
// This is done by matching current host IP addresses against what topology
// service thinks the host address is.
// If no match is found we assume we are running on host which is not
// part of the Romana setup and spit error out.
func (a Agent) identifyCurrentHost() error {
	topologyURL, err := common.GetServiceUrl(a.config.Common.Api.RootServiceUrl, "topology")
	if err != nil {
		return agentError(err)
	}

	client, err := common.NewRestClient(topologyURL)
	if err != nil {
		return agentError(err)
	}
	index := common.IndexResponse{}
	err = client.Get(topologyURL, &index)
	if err != nil {
		return agentError(err)
	}
	dcURL := index.Links.FindByRel("datacenter")
	a.networkConfig.dc = common.Datacenter{}
	err = client.Get(dcURL, &a.networkConfig.dc)
	if err != nil {
		return agentError(err)
	}

	hostURL := index.Links.FindByRel("host-list")
	hosts := []common.HostMessage{}
	err = client.Get(hostURL, &hosts)
	if err != nil {
		return agentError(err)
	}
	log.Println("Retrieived hosts list, found", len(hosts))

	// find our host in the list of hosts
	// checking for a matching IP address
	// to get the romana CIDR
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return err
	}
	log.Println("Searching", len(addrs), "for a matching host configuration")

	matched := false
	romanaCIDR := &net.IPNet{}
	FindHost:
	for i, host := range hosts {
		hostIP := net.ParseIP(host.Ip)
		if hostIP == nil {
			continue
		}
		for _, addr := range addrs {
			ipnet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}
			if ! ipnet.IP.Equal(hostIP) {
				continue
			}
			log.Println("Found matching host and IP:", hostIP, "==", ipnet)
			// match has been found
			_, romanaCIDR, err = net.ParseCIDR(host.RomanaIp)
			if err != nil {
				// this would mean a problem with the data itself
				continue
			}
			matched = true
			// update otherHosts list
			a.networkConfig.otherHosts = append(a.networkConfig.otherHosts, hosts[0:i]...)
			a.networkConfig.otherHosts = append(a.networkConfig.otherHosts, hosts[1+1:]...)

			break FindHost
		}
	}
	if ! matched {
		log.Println("Host configuration not found")
		return agentErrorString("Unable to find host configuration")
	}

	log.Println("Found host configuration using address")

	// now find the interface address that matches it
	for _, addr := range addrs {
		ipnet, ok := addr.(*net.IPNet)
		if !ok {
			continue
		}
		if romanaCIDR.Contains(ipnet.IP) {
			// check that it's the same subnet size
			s1, _ := romanaCIDR.Mask.Size()
			s2, _ := ipnet.Mask.Size()
			if s1 != s2 {
				continue
			}
			// OK, we're happy with this result
			a.networkConfig.romanaIP = ipnet.IP
			return nil
		}
	}
	log.Println("Unable to find interface matching the Romana CIDR", romanaCIDR)
	return agentErrorString("Unable to find interface matching Romana CIDR")
}
<commit_msg>Better log message<commit_after>// Copyright (c) 2016 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package agent

import (
	"github.com/romana/core/common"
//	"github.com/romana/core/topology"

	"log"
	"net"
)

// NetworkConfig holds the agent's current configuration.
// This consists of data parsed from the config file as well as
// runtime or discovered configuration, such as the network
// config of the current host.
type NetworkConfig struct {
	// Current host network configuration
	hostIP net.IP
	romanaIP net.IP
	otherHosts []common.HostMessage
	dc               common.Datacenter
}

// EndpointNetmaskSize returns integer value (aka size) of endpoint netmask.
func (c *NetworkConfig) EndpointNetmaskSize() uint64 {
	// TODO make this depend on the IP version
	return 32 - uint64(c.dc.EndpointSpaceBits)
}

// PNetCIDR returns pseudo net cidr in net.IPNet format.
func (c *NetworkConfig) PNetCIDR() (cidr *net.IPNet, err error) {
	_, cidr, err = net.ParseCIDR(c.dc.Cidr)
	return
}

// TenantBits returns tenant bits value from POC config.
func (c *NetworkConfig) TenantBits() uint {
	return c.dc.TenantBits
}

// SegmentBits returns segment bits value from POC config.
func (c *NetworkConfig) SegmentBits() uint {
	return c.dc.SegmentBits
}

// EndpointBits returns endpoint bits value from POC config.
func (c *NetworkConfig) EndpointBits() uint {
	return c.dc.EndpointBits
}

// identifyCurrentHost discovers network configuration
// of the host we are running on.
// We need to know public IP and Romana gateway IP of the current host.
// This is done by matching current host IP addresses against what topology
// service thinks the host address is.
// If no match is found we assume we are running on host which is not
// part of the Romana setup and spit error out.
func (a Agent) identifyCurrentHost() error {
	topologyURL, err := common.GetServiceUrl(a.config.Common.Api.RootServiceUrl, "topology")
	if err != nil {
		return agentError(err)
	}

	client, err := common.NewRestClient(topologyURL)
	if err != nil {
		return agentError(err)
	}
	index := common.IndexResponse{}
	err = client.Get(topologyURL, &index)
	if err != nil {
		return agentError(err)
	}
	dcURL := index.Links.FindByRel("datacenter")
	a.networkConfig.dc = common.Datacenter{}
	err = client.Get(dcURL, &a.networkConfig.dc)
	if err != nil {
		return agentError(err)
	}

	hostURL := index.Links.FindByRel("host-list")
	hosts := []common.HostMessage{}
	err = client.Get(hostURL, &hosts)
	if err != nil {
		return agentError(err)
	}
	log.Println("Retrieived hosts list, found", len(hosts), "hosts")

	// find our host in the list of hosts
	// checking for a matching IP address
	// to get the romana CIDR
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return err
	}
	log.Println("Searching", len(addrs), "interfaces for a matching host configuration")

	matched := false
	romanaCIDR := &net.IPNet{}
	FindHost:
	for i, host := range hosts {
		hostIP := net.ParseIP(host.Ip)
		if hostIP == nil {
			continue
		}
		for _, addr := range addrs {
			ipnet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}
			if ! ipnet.IP.Equal(hostIP) {
				continue
			}
			log.Println("Found matching host and IP:", hostIP, "==", ipnet)
			// match has been found
			_, romanaCIDR, err = net.ParseCIDR(host.RomanaIp)
			if err != nil {
				// this would mean a problem with the data itself
				continue
			}
			matched = true
			// update otherHosts list
			a.networkConfig.otherHosts = append(a.networkConfig.otherHosts, hosts[0:i]...)
			a.networkConfig.otherHosts = append(a.networkConfig.otherHosts, hosts[1+1:]...)

			break FindHost
		}
	}
	if ! matched {
		log.Println("Host configuration not found")
		return agentErrorString("Unable to find host configuration")
	}

	log.Println("Found host configuration using address")

	// now find the interface address that matches it
	for _, addr := range addrs {
		ipnet, ok := addr.(*net.IPNet)
		if !ok {
			continue
		}
		if romanaCIDR.Contains(ipnet.IP) {
			// check that it's the same subnet size
			s1, _ := romanaCIDR.Mask.Size()
			s2, _ := ipnet.Mask.Size()
			if s1 != s2 {
				continue
			}
			// OK, we're happy with this result
			a.networkConfig.romanaIP = ipnet.IP
			return nil
		}
	}
	log.Println("Unable to find interface matching the Romana CIDR", romanaCIDR)
	return agentErrorString("Unable to find interface matching Romana CIDR")
}
<|endoftext|>


<commit_before>// Copyright 2010 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 net

import (
	"time"
)

func parseDialNetwork(net string) (afnet string, proto int, err error) {
	i := last(net, ':')
	if i < 0 { // no colon
		switch net {
		case "tcp", "tcp4", "tcp6":
		case "udp", "udp4", "udp6":
		case "unix", "unixgram", "unixpacket":
		default:
			return "", 0, UnknownNetworkError(net)
		}
		return net, 0, nil
	}
	afnet = net[:i]
	switch afnet {
	case "ip", "ip4", "ip6":
		protostr := net[i+1:]
		proto, i, ok := dtoi(protostr, 0)
		if !ok || i != len(protostr) {
			proto, err = lookupProtocol(protostr)
			if err != nil {
				return "", 0, err
			}
		}
		return afnet, proto, nil
	}
	return "", 0, UnknownNetworkError(net)
}

func resolveNetAddr(op, net, addr string) (afnet string, a Addr, err error) {
	afnet, _, err = parseDialNetwork(net)
	if err != nil {
		return "", nil, &OpError{op, net, nil, err}
	}
	if op == "dial" && addr == "" {
		return "", nil, &OpError{op, net, nil, errMissingAddress}
	}
	switch afnet {
	case "tcp", "tcp4", "tcp6":
		if addr != "" {
			a, err = ResolveTCPAddr(afnet, addr)
		}
	case "udp", "udp4", "udp6":
		if addr != "" {
			a, err = ResolveUDPAddr(afnet, addr)
		}
	case "ip", "ip4", "ip6":
		if addr != "" {
			a, err = ResolveIPAddr(afnet, addr)
		}
	case "unix", "unixgram", "unixpacket":
		if addr != "" {
			a, err = ResolveUnixAddr(afnet, addr)
		}
	}
	return
}

// Dial connects to the address addr on the network net.
//
// Known networks are "tcp", "tcp4" (IPv4-only), "tcp6" (IPv6-only),
// "udp", "udp4" (IPv4-only), "udp6" (IPv6-only), "ip", "ip4"
// (IPv4-only), "ip6" (IPv6-only), "unix", "unixgram" and "unixpacket".
//
// For TCP and UDP networks, addresses have the form host:port.
// If host is a literal IPv6 address, it must be enclosed
// in square brackets.  The functions JoinHostPort and SplitHostPort
// manipulate addresses in this form.
//
// Examples:
//	Dial("tcp", "12.34.56.78:80")
//	Dial("tcp", "google.com:80")
//	Dial("tcp", "[de:ad:be:ef::ca:fe]:80")
//
// For IP networks, addr must be "ip", "ip4" or "ip6" followed
// by a colon and a protocol number or name.
//
// Examples:
//	Dial("ip4:1", "127.0.0.1")
//	Dial("ip6:ospf", "::1")
//
func Dial(net, addr string) (Conn, error) {
	_, addri, err := resolveNetAddr("dial", net, addr)
	if err != nil {
		return nil, err
	}
	return dialAddr(net, addr, addri)
}

func dialAddr(net, addr string, addri Addr) (c Conn, err error) {
	switch ra := addri.(type) {
	case *TCPAddr:
		c, err = DialTCP(net, nil, ra)
	case *UDPAddr:
		c, err = DialUDP(net, nil, ra)
	case *IPAddr:
		c, err = DialIP(net, nil, ra)
	case *UnixAddr:
		c, err = DialUnix(net, nil, ra)
	default:
		err = &OpError{"dial", net + " " + addr, nil, UnknownNetworkError(net)}
	}
	if err != nil {
		return nil, err
	}
	return
}

// DialTimeout acts like Dial but takes a timeout.
// The timeout includes name resolution, if required.
func DialTimeout(net, addr string, timeout time.Duration) (Conn, error) {
	// TODO(bradfitz): the timeout should be pushed down into the
	// net package's event loop, so on timeout to dead hosts we
	// don't have a goroutine sticking around for the default of
	// ~3 minutes.
	t := time.NewTimer(timeout)
	defer t.Stop()
	type pair struct {
		Conn
		error
	}
	ch := make(chan pair, 1)
	resolvedAddr := make(chan Addr, 1)
	go func() {
		_, addri, err := resolveNetAddr("dial", net, addr)
		if err != nil {
			ch <- pair{nil, err}
			return
		}
		resolvedAddr <- addri // in case we need it for OpError
		c, err := dialAddr(net, addr, addri)
		ch <- pair{c, err}
	}()
	select {
	case <-t.C:
		// Try to use the real Addr in our OpError, if we resolved it
		// before the timeout. Otherwise we just use stringAddr.
		var addri Addr
		select {
		case a := <-resolvedAddr:
			addri = a
		default:
			addri = &stringAddr{net, addr}
		}
		err := &OpError{
			Op:   "dial",
			Net:  net,
			Addr: addri,
			Err:  &timeoutError{},
		}
		return nil, err
	case p := <-ch:
		return p.Conn, p.error
	}
	panic("unreachable")
}

type stringAddr struct {
	net, addr string
}

func (a stringAddr) Network() string { return a.net }
func (a stringAddr) String() string  { return a.addr }

// Listen announces on the local network address laddr.
// The network string net must be a stream-oriented network:
// "tcp", "tcp4", "tcp6", or "unix", or "unixpacket".
func Listen(net, laddr string) (Listener, error) {
	afnet, a, err := resolveNetAddr("listen", net, laddr)
	if err != nil {
		return nil, err
	}
	switch afnet {
	case "tcp", "tcp4", "tcp6":
		var la *TCPAddr
		if a != nil {
			la = a.(*TCPAddr)
		}
		return ListenTCP(net, la)
	case "unix", "unixpacket":
		var la *UnixAddr
		if a != nil {
			la = a.(*UnixAddr)
		}
		return ListenUnix(net, la)
	}
	return nil, UnknownNetworkError(net)
}

// ListenPacket announces on the local network address laddr.
// The network string net must be a packet-oriented network:
// "udp", "udp4", "udp6", "ip", "ip4", "ip6" or "unixgram".
func ListenPacket(net, addr string) (PacketConn, error) {
	afnet, a, err := resolveNetAddr("listen", net, addr)
	if err != nil {
		return nil, err
	}
	switch afnet {
	case "udp", "udp4", "udp6":
		var la *UDPAddr
		if a != nil {
			la = a.(*UDPAddr)
		}
		return ListenUDP(net, la)
	case "ip", "ip4", "ip6":
		var la *IPAddr
		if a != nil {
			la = a.(*IPAddr)
		}
		return ListenIP(net, la)
	case "unixgram":
		var la *UnixAddr
		if a != nil {
			la = a.(*UnixAddr)
		}
		return DialUnix(net, la, nil)
	}
	return nil, UnknownNetworkError(net)
}
<commit_msg>net: fix comment on Dial with unixgram<commit_after>// Copyright 2010 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 net

import (
	"time"
)

func parseDialNetwork(net string) (afnet string, proto int, err error) {
	i := last(net, ':')
	if i < 0 { // no colon
		switch net {
		case "tcp", "tcp4", "tcp6":
		case "udp", "udp4", "udp6":
		case "unix", "unixgram", "unixpacket":
		default:
			return "", 0, UnknownNetworkError(net)
		}
		return net, 0, nil
	}
	afnet = net[:i]
	switch afnet {
	case "ip", "ip4", "ip6":
		protostr := net[i+1:]
		proto, i, ok := dtoi(protostr, 0)
		if !ok || i != len(protostr) {
			proto, err = lookupProtocol(protostr)
			if err != nil {
				return "", 0, err
			}
		}
		return afnet, proto, nil
	}
	return "", 0, UnknownNetworkError(net)
}

func resolveNetAddr(op, net, addr string) (afnet string, a Addr, err error) {
	afnet, _, err = parseDialNetwork(net)
	if err != nil {
		return "", nil, &OpError{op, net, nil, err}
	}
	if op == "dial" && addr == "" {
		return "", nil, &OpError{op, net, nil, errMissingAddress}
	}
	switch afnet {
	case "tcp", "tcp4", "tcp6":
		if addr != "" {
			a, err = ResolveTCPAddr(afnet, addr)
		}
	case "udp", "udp4", "udp6":
		if addr != "" {
			a, err = ResolveUDPAddr(afnet, addr)
		}
	case "ip", "ip4", "ip6":
		if addr != "" {
			a, err = ResolveIPAddr(afnet, addr)
		}
	case "unix", "unixgram", "unixpacket":
		if addr != "" {
			a, err = ResolveUnixAddr(afnet, addr)
		}
	}
	return
}

// Dial connects to the address addr on the network net.
//
// Known networks are "tcp", "tcp4" (IPv4-only), "tcp6" (IPv6-only),
// "udp", "udp4" (IPv4-only), "udp6" (IPv6-only), "ip", "ip4"
// (IPv4-only), "ip6" (IPv6-only), "unix" and "unixpacket".
//
// For TCP and UDP networks, addresses have the form host:port.
// If host is a literal IPv6 address, it must be enclosed
// in square brackets.  The functions JoinHostPort and SplitHostPort
// manipulate addresses in this form.
//
// Examples:
//	Dial("tcp", "12.34.56.78:80")
//	Dial("tcp", "google.com:80")
//	Dial("tcp", "[de:ad:be:ef::ca:fe]:80")
//
// For IP networks, addr must be "ip", "ip4" or "ip6" followed
// by a colon and a protocol number or name.
//
// Examples:
//	Dial("ip4:1", "127.0.0.1")
//	Dial("ip6:ospf", "::1")
//
func Dial(net, addr string) (Conn, error) {
	_, addri, err := resolveNetAddr("dial", net, addr)
	if err != nil {
		return nil, err
	}
	return dialAddr(net, addr, addri)
}

func dialAddr(net, addr string, addri Addr) (c Conn, err error) {
	switch ra := addri.(type) {
	case *TCPAddr:
		c, err = DialTCP(net, nil, ra)
	case *UDPAddr:
		c, err = DialUDP(net, nil, ra)
	case *IPAddr:
		c, err = DialIP(net, nil, ra)
	case *UnixAddr:
		c, err = DialUnix(net, nil, ra)
	default:
		err = &OpError{"dial", net + " " + addr, nil, UnknownNetworkError(net)}
	}
	if err != nil {
		return nil, err
	}
	return
}

// DialTimeout acts like Dial but takes a timeout.
// The timeout includes name resolution, if required.
func DialTimeout(net, addr string, timeout time.Duration) (Conn, error) {
	// TODO(bradfitz): the timeout should be pushed down into the
	// net package's event loop, so on timeout to dead hosts we
	// don't have a goroutine sticking around for the default of
	// ~3 minutes.
	t := time.NewTimer(timeout)
	defer t.Stop()
	type pair struct {
		Conn
		error
	}
	ch := make(chan pair, 1)
	resolvedAddr := make(chan Addr, 1)
	go func() {
		_, addri, err := resolveNetAddr("dial", net, addr)
		if err != nil {
			ch <- pair{nil, err}
			return
		}
		resolvedAddr <- addri // in case we need it for OpError
		c, err := dialAddr(net, addr, addri)
		ch <- pair{c, err}
	}()
	select {
	case <-t.C:
		// Try to use the real Addr in our OpError, if we resolved it
		// before the timeout. Otherwise we just use stringAddr.
		var addri Addr
		select {
		case a := <-resolvedAddr:
			addri = a
		default:
			addri = &stringAddr{net, addr}
		}
		err := &OpError{
			Op:   "dial",
			Net:  net,
			Addr: addri,
			Err:  &timeoutError{},
		}
		return nil, err
	case p := <-ch:
		return p.Conn, p.error
	}
	panic("unreachable")
}

type stringAddr struct {
	net, addr string
}

func (a stringAddr) Network() string { return a.net }
func (a stringAddr) String() string  { return a.addr }

// Listen announces on the local network address laddr.
// The network string net must be a stream-oriented network:
// "tcp", "tcp4", "tcp6", or "unix", or "unixpacket".
func Listen(net, laddr string) (Listener, error) {
	afnet, a, err := resolveNetAddr("listen", net, laddr)
	if err != nil {
		return nil, err
	}
	switch afnet {
	case "tcp", "tcp4", "tcp6":
		var la *TCPAddr
		if a != nil {
			la = a.(*TCPAddr)
		}
		return ListenTCP(net, la)
	case "unix", "unixpacket":
		var la *UnixAddr
		if a != nil {
			la = a.(*UnixAddr)
		}
		return ListenUnix(net, la)
	}
	return nil, UnknownNetworkError(net)
}

// ListenPacket announces on the local network address laddr.
// The network string net must be a packet-oriented network:
// "udp", "udp4", "udp6", "ip", "ip4", "ip6" or "unixgram".
func ListenPacket(net, addr string) (PacketConn, error) {
	afnet, a, err := resolveNetAddr("listen", net, addr)
	if err != nil {
		return nil, err
	}
	switch afnet {
	case "udp", "udp4", "udp6":
		var la *UDPAddr
		if a != nil {
			la = a.(*UDPAddr)
		}
		return ListenUDP(net, la)
	case "ip", "ip4", "ip6":
		var la *IPAddr
		if a != nil {
			la = a.(*IPAddr)
		}
		return ListenIP(net, la)
	case "unixgram":
		var la *UnixAddr
		if a != nil {
			la = a.(*UnixAddr)
		}
		return DialUnix(net, la, nil)
	}
	return nil, UnknownNetworkError(net)
}
<|endoftext|>


<commit_before>package mcts

import (
	"log"
	"math"
	"math/rand"
	"sort"
	"strings"
	"time"

	"context"

	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type MCTSConfig struct {
	Debug int
	Limit time.Duration
	C     float64
	Seed  int64

	InitialVisits int
	MMDepth       int
	MaxRollout    int
	EvalThreshold int64
	Policy        string

	Size int

	DumpTree string
}

type Policy interface {
	Select(ctx context.Context, m *MonteCarloAI, p *tak.Position) *tak.Position
}

type MonteCarloAI struct {
	c bitboard.Constants

	cfg  MCTSConfig
	mm   *ai.MinimaxAI
	eval ai.EvaluationFunc

	policy Policy

	r *rand.Rand
}

type tree struct {
	position    *tak.Position
	move        tak.Move
	simulations int

	proven int
	value  int

	parent   *tree
	children []*tree
}

func (t *tree) ucb(C float64, N int) float64 {
	if t.proven > 0 {
		return -100
	} else if t.proven < 0 {
		return 100
	} else if t.simulations == 0 {
		return 10
	} else {
		return -float64(t.value)/float64(t.simulations) +
			C*math.Sqrt(math.Log(float64(N))/float64(t.simulations))
	}
}

type bySims []*tree

func (b bySims) Len() int           { return len(b) }
func (b bySims) Less(i, j int) bool { return b[j].simulations < b[i].simulations }
func (b bySims) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}

func (ai *MonteCarloAI) GetMove(ctx context.Context, p *tak.Position) tak.Move {
	tree := &tree{
		position: p,
	}
	start := time.Now()
	deadline, limited := ctx.Deadline()
	if !limited || deadline.Sub(start) > ai.cfg.Limit {
		deadline = time.Now().Add(ai.cfg.Limit)
	}

	var tick <-chan time.Time
	if ai.cfg.Debug > 2 {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		tick = ticker.C
	}
	for time.Now().Before(deadline) {
		node := ai.descend(tree)
		ai.populate(ctx, node)
		if tree.proven != 0 {
			break
		}
		var val int
		if node.proven == 0 {
			val = ai.rollout(ctx, node)
		}
		if ai.cfg.Debug > 4 {
			var s []string
			t := node
			for t.parent != nil {
				s = append(s, ptn.FormatMove(t.move))
				t = t.parent
			}
			log.Printf("evaluate: [%s] = %d p=%d",
				strings.Join(s, "<-"), val, node.proven)
		}
		ai.update(node, val)
		if tick != nil {
			select {
			case <-tick:
				ai.printdbg(tree)
			default:
			}
		}
	}
	if tree.proven != 0 {
		return ai.mm.GetMove(ctx, p)
	}
	best := tree.children[0]
	i := 0
	sort.Sort(bySims(tree.children))
	if ai.cfg.Debug > 2 {
		log.Printf("=== mcts done ===")
	}
	for _, c := range tree.children {
		if ai.cfg.Debug > 2 {
			log.Printf("[mcts][%s]: n=%d v=%d:%d(%0.3f) ucb=%f",
				ptn.FormatMove(c.move), c.simulations, c.proven, c.value,
				float64(c.value)/float64(c.simulations),
				c.ucb(ai.cfg.C, tree.simulations),
			)
		}
		if c.simulations > best.simulations {
			best = c
			i = 1
		} else if c.simulations == best.simulations {
			i++
			if ai.r.Intn(i) == 0 {
				best = c
				i = 1
			}
		}
	}
	if ai.cfg.Debug > 1 {
		log.Printf("[mcts] evaluated simulations=%d value=%d", tree.simulations, tree.value)
	}
	if ai.cfg.DumpTree != "" {
		ai.dumpTree(tree)
	}
	return best.move
}

func (mc *MonteCarloAI) printdbg(t *tree) {
	log.Printf("===")
	for _, c := range t.children {
		if c.simulations*20 > t.simulations {
			log.Printf("[mcts][%s]: n=%d v=%d:%d(%0.3f) ucb=%f",
				ptn.FormatMove(c.move), c.simulations, c.proven, c.value,
				float64(c.value)/float64(c.simulations),
				c.ucb(mc.cfg.C, t.simulations),
			)
		}
	}
}

func (mc *MonteCarloAI) populate(ctx context.Context, t *tree) {
	_, v, _ := mc.mm.Analyze(ctx, t.position)
	if v > ai.WinThreshold {
		t.proven = 1
		return
	} else if v < -ai.WinThreshold {
		t.proven = -1
		return
	}

	moves := t.position.AllMoves(nil)
	t.children = make([]*tree, 0, len(moves))
	for _, m := range moves {
		child, e := t.position.Move(m)
		if e != nil {
			continue
		}
		t.children = append(t.children, &tree{
			position: child,
			move:     m,
			parent:   t,
		})
	}
}

func (mc *MonteCarloAI) descendPolicy(t *tree) *tree {
	var best *tree
	val := ai.MinEval
	i := 0
	for _, c := range t.children {
		v := mc.eval(&mc.c, c.position)
		if v > val {
			best = c
			val = v
			i = 1
		} else if v == val {
			i++
			if mc.r.Intn(i) == 0 {
				best = c
			}
		}
	}
	return best
}

func (ai *MonteCarloAI) descend(t *tree) *tree {
	if t.children == nil {
		return t
	}
	var best *tree
	var val float64
	i := 0
	for _, c := range t.children {
		var s float64
		if c.proven > 0 {
			s = 0.01
		} else if c.proven < 0 {
			s = 100
		} else if c.simulations == 0 {
			s = 10
		} else {
			s = -float64(c.value)/float64(c.simulations) +
				ai.cfg.C*math.Sqrt(math.Log(float64(t.simulations))/float64(c.simulations))
		}
		if s > val {
			best = c
			val = s
			i = 1
		} else if s == val {
			i++
			if ai.r.Intn(i) == 0 {
				best = c
			}
		}
	}
	if best == nil {
		return t.children[0]
	}
	return ai.descend(best)
}

func (ai *MonteCarloAI) rollout(ctx context.Context, t *tree) int {
	p := t.position.Clone()

	for i := 0; i < ai.cfg.MaxRollout; i++ {
		if ok, c := p.GameOver(); ok {
			switch c {
			case tak.NoColor:
				return 0
			case t.position.ToMove():
				return 1
			default:
				return -1
			}
		}
		next := ai.policy.Select(ctx, ai, p)
		if next == nil {
			return 0
		}
		p = next
	}
	v := ai.eval(&ai.c, p)
	if v > ai.cfg.EvalThreshold {
		return 1
	} else if v < -ai.cfg.EvalThreshold {
		return -1
	}
	return 0
}

func (mc *MonteCarloAI) update(t *tree, value int) {
	for t != nil {
		t.simulations++
		if t.proven != 0 {
			if t.parent == nil {
				return
			}
			// Minimax backup
			if t.proven < 0 {
				// My best move is a loss; therefore
				// my parent should choose this branch
				// and win
				t.parent.proven = 1
				value = -1
			} else {
				// This move is a win for me; My
				// parent is a loss only if *all* of
				// its children are wins
				all := true
				for _, ch := range t.parent.children {
					if ch.proven <= 0 {
						all = false
						break
					}
				}
				if all {
					t.parent.proven = -1
				}
				value = 1
			}
		} else {
			t.value += value
		}
		value = -value
		t = t.parent
	}
}

func NewMonteCarlo(cfg MCTSConfig) *MonteCarloAI {
	mc := &MonteCarloAI{
		cfg: cfg,
		c:   bitboard.Precompute(uint(cfg.Size)),
	}
	if mc.cfg.C == 0 {
		mc.cfg.C = 0.7
	}
	if mc.cfg.Seed == 0 {
		mc.cfg.Seed = time.Now().Unix()
	}
	if mc.cfg.InitialVisits == 0 {
		mc.cfg.InitialVisits = 3
	}
	if mc.cfg.MMDepth == 0 {
		mc.cfg.MMDepth = 3
	}
	if mc.cfg.MaxRollout == 0 {
		mc.cfg.MaxRollout = 50
	}
	if mc.cfg.EvalThreshold == 0 {
		mc.cfg.EvalThreshold = 2000
	}
	mc.policy = mc.buildPolicy()
	mc.r = rand.New(rand.NewSource(mc.cfg.Seed))
	mc.mm = ai.NewMinimax(ai.MinimaxConfig{
		Size:     cfg.Size,
		Evaluate: ai.EvaluateWinner,
		Depth:    mc.cfg.MMDepth,
		Seed:     mc.cfg.Seed,
	})
	mc.eval = ai.MakeEvaluator(mc.cfg.Size, nil)
	return mc
}
<commit_msg>defunct<commit_after>package mcts

import (
	"log"
	"math"
	"math/rand"
	"sort"
	"strings"
	"time"

	"context"

	"github.com/nelhage/taktician/ai"
	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type MCTSConfig struct {
	Debug int
	Limit time.Duration
	C     float64
	Seed  int64

	InitialVisits int
	MMDepth       int
	MaxRollout    int
	EvalThreshold int64
	Policy        string

	Size int

	DumpTree string
}

type Policy interface {
	Select(ctx context.Context, m *MonteCarloAI, p *tak.Position) *tak.Position
}

type MonteCarloAI struct {
	c bitboard.Constants

	cfg  MCTSConfig
	mm   *ai.MinimaxAI
	eval ai.EvaluationFunc

	policy Policy

	r *rand.Rand
}

type tree struct {
	position    *tak.Position
	move        tak.Move
	simulations int

	proven int
	value  int

	parent   *tree
	children []*tree
}

func (t *tree) ucb(C float64, N int) float64 {
	if t.proven > 0 {
		return -100
	} else if t.proven < 0 {
		return 100
	} else if t.simulations == 0 {
		return 10
	} else {
		return -float64(t.value)/float64(t.simulations) +
			C*math.Sqrt(math.Log(float64(N))/float64(t.simulations))
	}
}

type bySims []*tree

func (b bySims) Len() int           { return len(b) }
func (b bySims) Less(i, j int) bool { return b[j].simulations < b[i].simulations }
func (b bySims) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}

func (ai *MonteCarloAI) GetMove(ctx context.Context, p *tak.Position) tak.Move {
	tree := &tree{
		position: p,
	}
	start := time.Now()
	deadline, limited := ctx.Deadline()
	if !limited || deadline.Sub(start) > ai.cfg.Limit {
		deadline = time.Now().Add(ai.cfg.Limit)
	}

	var tick <-chan time.Time
	if ai.cfg.Debug > 2 {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		tick = ticker.C
	}
	for time.Now().Before(deadline) {
		node := ai.descend(tree)
		ai.populate(ctx, node)
		if tree.proven != 0 {
			break
		}
		var val int
		if node.proven == 0 {
			val = ai.rollout(ctx, node)
		}
		if ai.cfg.Debug > 4 {
			var s []string
			t := node
			for t.parent != nil {
				s = append(s, ptn.FormatMove(t.move))
				t = t.parent
			}
			log.Printf("evaluate: [%s] = %d p=%d",
				strings.Join(s, "<-"), val, node.proven)
		}
		ai.update(node, val)
		if tick != nil {
			select {
			case <-tick:
				ai.printdbg(tree)
			default:
			}
		}
	}
	if tree.proven != 0 {
		return ai.mm.GetMove(ctx, p)
	}
	best := tree.children[0]
	i := 0
	sort.Sort(bySims(tree.children))
	if ai.cfg.Debug > 2 {
		log.Printf("=== mcts done ===")
	}
	for _, c := range tree.children {
		if ai.cfg.Debug > 2 {
			log.Printf("[mcts][%s]: n=%d v=%d:%d(%0.3f) ucb=%f",
				ptn.FormatMove(c.move), c.simulations, c.proven, c.value,
				float64(c.value)/float64(c.simulations),
				c.ucb(ai.cfg.C, tree.simulations),
			)
		}
		if c.simulations > best.simulations {
			best = c
			i = 1
		} else if c.simulations == best.simulations {
			i++
			if ai.r.Intn(i) == 0 {
				best = c
				i = 1
			}
		}
	}
	if ai.cfg.Debug > 1 {
		log.Printf("[mcts] evaluated simulations=%d value=%d", tree.simulations, tree.value)
	}
	if ai.cfg.DumpTree != "" {
		ai.dumpTree(tree)
	}
	return best.move
}

func (mc *MonteCarloAI) printdbg(t *tree) {
	log.Printf("===")
	for _, c := range t.children {
		if c.simulations*20 > t.simulations {
			log.Printf("[mcts][%s]: n=%d v=%d:%d(%0.3f) ucb=%f",
				ptn.FormatMove(c.move), c.simulations, c.proven, c.value,
				float64(c.value)/float64(c.simulations),
				c.ucb(mc.cfg.C, t.simulations),
			)
		}
	}
}

func (mc *MonteCarloAI) populate(ctx context.Context, t *tree) {
	_, v, _ := mc.mm.Analyze(ctx, t.position)
	if v > ai.WinThreshold {
		t.proven = 1
		return
	} else if v < -ai.WinThreshold {
		t.proven = -1
		return
	}

	moves := t.position.AllMoves(nil)
	t.children = make([]*tree, 0, len(moves))
	for _, m := range moves {
		child, e := t.position.Move(m)
		if e != nil {
			continue
		}
		t.children = append(t.children, &tree{
			position: child,
			move:     m,
			parent:   t,
		})
	}
}

func (ai *MonteCarloAI) descend(t *tree) *tree {
	if t.children == nil {
		return t
	}
	var best *tree
	var val float64
	i := 0
	for _, c := range t.children {
		var s float64
		if c.proven > 0 {
			s = 0.01
		} else if c.proven < 0 {
			s = 100
		} else if c.simulations == 0 {
			s = 10
		} else {
			s = -float64(c.value)/float64(c.simulations) +
				ai.cfg.C*math.Sqrt(math.Log(float64(t.simulations))/float64(c.simulations))
		}
		if s > val {
			best = c
			val = s
			i = 1
		} else if s == val {
			i++
			if ai.r.Intn(i) == 0 {
				best = c
			}
		}
	}
	if best == nil {
		return t.children[0]
	}
	return ai.descend(best)
}

func (ai *MonteCarloAI) rollout(ctx context.Context, t *tree) int {
	p := t.position.Clone()

	for i := 0; i < ai.cfg.MaxRollout; i++ {
		if ok, c := p.GameOver(); ok {
			switch c {
			case tak.NoColor:
				return 0
			case t.position.ToMove():
				return 1
			default:
				return -1
			}
		}
		next := ai.policy.Select(ctx, ai, p)
		if next == nil {
			return 0
		}
		p = next
	}
	v := ai.eval(&ai.c, p)
	if v > ai.cfg.EvalThreshold {
		return 1
	} else if v < -ai.cfg.EvalThreshold {
		return -1
	}
	return 0
}

func (mc *MonteCarloAI) update(t *tree, value int) {
	for t != nil {
		t.simulations++
		if t.proven != 0 {
			if t.parent == nil {
				return
			}
			// Minimax backup
			if t.proven < 0 {
				// My best move is a loss; therefore
				// my parent should choose this branch
				// and win
				t.parent.proven = 1
				value = -1
			} else {
				// This move is a win for me; My
				// parent is a loss only if *all* of
				// its children are wins
				all := true
				for _, ch := range t.parent.children {
					if ch.proven <= 0 {
						all = false
						break
					}
				}
				if all {
					t.parent.proven = -1
				}
				value = 1
			}
		} else {
			t.value += value
		}
		value = -value
		t = t.parent
	}
}

func NewMonteCarlo(cfg MCTSConfig) *MonteCarloAI {
	mc := &MonteCarloAI{
		cfg: cfg,
		c:   bitboard.Precompute(uint(cfg.Size)),
	}
	if mc.cfg.C == 0 {
		mc.cfg.C = 0.7
	}
	if mc.cfg.Seed == 0 {
		mc.cfg.Seed = time.Now().Unix()
	}
	if mc.cfg.InitialVisits == 0 {
		mc.cfg.InitialVisits = 3
	}
	if mc.cfg.MMDepth == 0 {
		mc.cfg.MMDepth = 3
	}
	if mc.cfg.MaxRollout == 0 {
		mc.cfg.MaxRollout = 50
	}
	if mc.cfg.EvalThreshold == 0 {
		mc.cfg.EvalThreshold = 2000
	}
	mc.policy = mc.buildPolicy()
	mc.r = rand.New(rand.NewSource(mc.cfg.Seed))
	mc.mm = ai.NewMinimax(ai.MinimaxConfig{
		Size:     cfg.Size,
		Evaluate: ai.EvaluateWinner,
		Depth:    mc.cfg.MMDepth,
		Seed:     mc.cfg.Seed,
	})
	mc.eval = ai.MakeEvaluator(mc.cfg.Size, nil)
	return mc
}
<|endoftext|>


<commit_before>package api

import (
	"os"
	"testing"
)

func try(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func assert(t *testing.T, expr bool, desc string) {
	if !expr {
		t.Fatal("Failed:", desc)
	}
}

func TestParseThread(t *testing.T) {
	file, err := os.Open("example.json")
	try(t, err)
	defer file.Close()

	thread, err := ParseThread(file, "ck")
	try(t, err)

	assert(t, thread.OP.Name == "Anonymous", "OP's name should be Anonymous")
	assert(t, thread.Id() == 3856791, "Thread id should be 3856791")
	assert(t, thread.OP.File != nil, "OP post should have a file")
	assert(t, len(thread.Posts) == 38, "Thread should have 38 posts")
	imageURL := thread.OP.ImageURL()
	assert(t, imageURL == "http://i.4cdn.org/ck/1346968817055.jpg", "Image URL should be 'http://i.4cdn.org/ck/1346968817055.jpg' (got '"+imageURL+"')")
	thumbURL := thread.OP.ThumbURL()
	assert(t, thumbURL == "http://t.4cdn.org/ck/1346968817055s.jpg", "Thumb URL should be 'http://t.4cdn.org/ck/1346968817055s.jpg' (got '"+thumbURL+"')")
}

func TestGetIndex(t *testing.T) {
	threads, err := GetIndex("a", 0)
	try(t, err)
	assert(t, len(threads) > 0, "Threads should exist")
}

func TestGetThreads(t *testing.T) {
	n, err := GetThreads("a")
	try(t, err)
	for _, q := range n {
		for _, p := range q {
			if p == 0 {
				t.Fatal("There are #0 posts")
			}
		}
	}
}
<commit_msg>Fixed tests a little<commit_after>package api

import (
	"os"
	"testing"
)

func try(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func assert(t *testing.T, expr bool, desc string) {
	if !expr {
		t.Fatal("Failed:", desc)
	}
}

func TestParseThread(t *testing.T) {
	file, err := os.Open("example.json")
	try(t, err)
	defer file.Close()

	thread, err := ParseThread(file, "ck")
	try(t, err)

	assert(t, thread.OP.Name == "Anonymous", "OP's name should be Anonymous")
	assert(t, thread.Id() == 3856791, "Thread id should be 3856791")
	assert(t, thread.OP.File != nil, "OP post should have a file")
	assert(t, len(thread.Posts) == 38, "Thread should have 38 posts")
	imageURL := thread.OP.ImageURL()
	assert(t, imageURL == "http://i.4cdn.org/ck/1346968817055.jpg", "Image URL should be 'http://i.4cdn.org/ck/1346968817055.jpg' (got '"+imageURL+"')")
	thumbURL := thread.OP.ThumbURL()
	assert(t, thumbURL == "http://i.4cdn.org/ck/1346968817055s.jpg", "Thumb URL should be 'http://i.4cdn.org/ck/1346968817055s.jpg' (got '"+thumbURL+"')")
}

func TestGetIndex(t *testing.T) {
	threads, err := GetIndex("a", 0)
	try(t, err)
	assert(t, len(threads) > 0, "Threads should exist")
}

func TestGetThreads(t *testing.T) {
	n, err := GetThreads("a")
	try(t, err)
	for _, q := range n {
		for _, p := range q {
			if p == 0 {
				t.Fatal("There are #0 posts")
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"regexp"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/sqs"
)

const (
	invalidContent = "InvalidMessageContents"
)

var (
	emailRegexp = regexp.MustCompile(`^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,4}$`)
)

type ResponseError struct {
	Errors map[string]string `json:"errors"`
}

func NewResponseError(errors map[string]string) *ResponseError {
	return &ResponseError{errors}
}

func NewBaseResponseError(errorMsg string) *ResponseError {
	return &ResponseError{map[string]string{"base": errorMsg}}
}

type SendEmailRequest struct {
	Email Email `json:"email"`
}

type Email struct {
	FromEmail string `json:"fromEmail"`
	FromName  string `json:"fromName"`
	ToEmail   string `json:"toEmail"`
	ToName    string `json:"toName"`
	Subject   string `json:"subject"`
	Body      string `json:"body"`
}

type SendEmailResponse struct {
	MessageId string `json:"messageId"`
}

func (e Email) Validate() (bool, *ResponseError) {
	errors := make(map[string]string)

	if valid, errorMsg := validateEmail("From email", e.FromEmail); !valid {
		errors["fromEmail"] = errorMsg
	}
	if valid, errorMsg := validateEmail("To email", e.ToEmail); !valid {
		errors["toEmail"] = errorMsg
	}

	if len(errors) > 0 {
		return false, NewResponseError(errors)
	}
	return true, nil
}

func validateEmail(fieldName, email string) (bool, string) {
	if email == "" {
		return false, fieldName + " is required"
	}
	if !emailRegexp.MatchString(email) {
		return false, fieldName + " is not a valid email"
	}
	return true, ""
}

func respondWithError(w http.ResponseWriter, respErr *ResponseError, httpStatus int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(httpStatus)
	errorBytes, err := json.Marshal(respErr)
	if err != nil {
		// This should never happen
		panic("Could not marshal error response: " + err.Error())
	}

	w.Write(errorBytes)
}

func SendEmailHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, config.MaxBodySizeBytes))
	if err != nil {
		respondWithError(w, NewBaseResponseError("Could not read body"), http.StatusBadRequest)
		return
	}
	defer r.Body.Close()

	var request SendEmailRequest
	if err := json.Unmarshal(body, &request); err != nil {
		respondWithError(w, NewBaseResponseError("Could not decode JSON body"), http.StatusBadRequest)
		return
	}
	email := request.Email
	if valid, respErr := email.Validate(); !valid {
		respondWithError(w, respErr, http.StatusUnprocessableEntity)
		return
	}

	// get a random queue url from config
	queueUrl := config.QueueUrls[rand.Intn(len(config.QueueUrls))]

	bodyStr := string(body)
	resp, err := sqsClient.SendMessage(&sqs.SendMessageInput{
		QueueUrl:    &queueUrl,
		MessageBody: &bodyStr,
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == invalidContent {
			respondWithError(
				w,
				NewBaseResponseError("Body contains characters outside the allowed set"),
				http.StatusBadRequest,
			)
			return
		} else {
			respondWithError(w, NewBaseResponseError("Service unavailable"), http.StatusServiceUnavailable)
			return
		}
	}

	response := &SendEmailResponse{MessageId: *resp.MessageId}
	respBytes, err := json.Marshal(response)
	if err != nil {
		// This should never happen
		panic("Could not marshal response: " + err.Error())
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(respBytes)
}
<commit_msg>Added more logging to API<commit_after>package main

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"regexp"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/sqs"
)

const (
	invalidContent = "InvalidMessageContents"
)

var (
	emailRegexp = regexp.MustCompile(`^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,4}$`)
)

type ResponseError struct {
	Errors map[string]string `json:"errors"`
}

func NewResponseError(errors map[string]string) *ResponseError {
	return &ResponseError{errors}
}

func NewBaseResponseError(errorMsg string) *ResponseError {
	return &ResponseError{map[string]string{"base": errorMsg}}
}

type SendEmailRequest struct {
	Email Email `json:"email"`
}

type Email struct {
	FromEmail string `json:"fromEmail"`
	FromName  string `json:"fromName"`
	ToEmail   string `json:"toEmail"`
	ToName    string `json:"toName"`
	Subject   string `json:"subject"`
	Body      string `json:"body"`
}

type SendEmailResponse struct {
	MessageId string `json:"messageId"`
}

func (e Email) Validate() (bool, *ResponseError) {
	errors := make(map[string]string)

	if valid, errorMsg := validateEmail("From email", e.FromEmail); !valid {
		errors["fromEmail"] = errorMsg
	}
	if valid, errorMsg := validateEmail("To email", e.ToEmail); !valid {
		errors["toEmail"] = errorMsg
	}

	if len(errors) > 0 {
		return false, NewResponseError(errors)
	}
	return true, nil
}

func validateEmail(fieldName, email string) (bool, string) {
	if email == "" {
		return false, fieldName + " is required"
	}
	if !emailRegexp.MatchString(email) {
		return false, fieldName + " is not a valid email"
	}
	return true, ""
}

func respondWithError(w http.ResponseWriter, respErr *ResponseError, httpStatus int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(httpStatus)
	errorBytes, err := json.Marshal(respErr)
	if err != nil {
		// This should never happen
		panic("Could not marshal error response: " + err.Error())
	}

	w.Write(errorBytes)
}

func SendEmailHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(io.LimitReader(r.Body, config.MaxBodySizeBytes))
	if err != nil {
		log.Print("[REQUEST ERROR] Could not read body: %v", err.Error())
		respondWithError(w, NewBaseResponseError("Could not read body"), http.StatusBadRequest)
		return
	}
	defer r.Body.Close()

	var request SendEmailRequest
	if err := json.Unmarshal(body, &request); err != nil {
		log.Print("[REQUEST ERROR] Could not decode JSON body: %v", err.Error())
		respondWithError(w, NewBaseResponseError("Could not decode JSON body"), http.StatusBadRequest)
		return
	}
	email := request.Email
	if valid, respErr := email.Validate(); !valid {
		log.Print("[REQUEST ERROR] Email is invalid: %v", email)
		respondWithError(w, respErr, http.StatusUnprocessableEntity)
		return
	}

	// get a random queue url from config
	queueUrl := config.QueueUrls[rand.Intn(len(config.QueueUrls))]

	bodyStr := string(body)
	resp, err := sqsClient.SendMessage(&sqs.SendMessageInput{
		QueueUrl:    &queueUrl,
		MessageBody: &bodyStr,
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == invalidContent {
			log.Print("[REQUEST ERROR] Invalid content in body: %v", err.Error())
			respondWithError(
				w,
				NewBaseResponseError("Body contains characters outside the allowed set"),
				http.StatusBadRequest,
			)
			return
		} else {
			log.Print("[REQUEST ERROR] SQS service returned an error: %v", err.Error())
			respondWithError(w, NewBaseResponseError("Service unavailable"), http.StatusServiceUnavailable)
			return
		}
	}

	response := &SendEmailResponse{MessageId: *resp.MessageId}
	respBytes, err := json.Marshal(response)
	if err != nil {
		// This should never happen
		panic("Could not marshal response: " + err.Error())
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(respBytes)
}
<|endoftext|>


<commit_before>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/consensus"
	"github.com/bytom/consensus/segwit"
	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/net/http/reqid"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

var defaultTxTTL = 5 * time.Minute

func (a *API) actionDecoder(action string) (func([]byte) (txbuilder.Action, error), bool) {
	var decoder func([]byte) (txbuilder.Action, error)
	switch action {
	case "control_address":
		decoder = txbuilder.DecodeControlAddressAction
	case "control_program":
		decoder = txbuilder.DecodeControlProgramAction
	case "control_receiver":
		decoder = txbuilder.DecodeControlReceiverAction
	case "issue":
		decoder = a.wallet.AssetReg.DecodeIssueAction
	case "retire":
		decoder = txbuilder.DecodeRetireAction
	case "spend_account":
		decoder = a.wallet.AccountMgr.DecodeSpendAction
	case "spend_account_unspent_output":
		decoder = a.wallet.AccountMgr.DecodeSpendUTXOAction
	default:
		return nil, false
	}
	return decoder, true
}

// TODO modify mergeActions to loadSpendAction
func mergeActions(req *BuildRequest) ([]map[string]interface{}, error) {
	var actions []map[string]interface{}
	actionMap := make(map[string]map[string]interface{})

	for _, m := range req.Actions {
		if actionType := m["type"].(string); actionType != "spend_account" {
			actions = append(actions, m)
			continue
		}

		if m["amount"] == nil {
			return nil, errEmptyAmount
		}

		amountNumber := m["amount"].(json.Number)
		amount, err := amountNumber.Int64()
		if err != nil || amount == 0 {
			return nil, errBadAmount
		}

		actionKey := m["asset_id"].(string) + m["account_id"].(string)
		if tmpM, ok := actionMap[actionKey]; ok {
			if tmpM["amount"] == nil {
				return nil, errEmptyAmount
			}

			tmpNumber := tmpM["amount"].(json.Number)
			tmpAmount, err := tmpNumber.Int64()
			if err != nil || tmpAmount == 0 {
				return nil, errBadAmount
			}

			tmpM["amount"] = json.Number(fmt.Sprintf("%v", tmpAmount+amount))
		} else {
			actionMap[actionKey] = m
			actions = append(actions, m)
		}
	}

	return actions, nil
}

func onlyHaveSpendActions(req *BuildRequest) bool {
	count := 0
	for _, m := range req.Actions {
		if actionType := m["type"].(string); strings.HasPrefix(actionType, "spend") {
			count++
		}
	}

	return count == len(req.Actions)
}

func (a *API) buildSingle(ctx context.Context, req *BuildRequest) (*txbuilder.Template, error) {
	err := a.filterAliases(ctx, req)
	if err != nil {
		return nil, err
	}

	if onlyHaveSpendActions(req) {
		return nil, errors.New("transaction only contain spend actions, didn't have output actions")
	}

	reqActions, err := mergeActions(req)
	if err != nil {
		return nil, errors.WithDetail(err, "unmarshal json amount error in mergeActions")
	}

	actions := make([]txbuilder.Action, 0, len(reqActions))
	for i, act := range reqActions {
		typ, ok := act["type"].(string)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "no action type provided on action %d", i)
		}
		decoder, ok := a.actionDecoder(typ)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "unknown action type %q on action %d", typ, i)
		}

		// Remarshal to JSON, the action may have been modified when we
		// filtered aliases.
		b, err := json.Marshal(act)
		if err != nil {
			return nil, err
		}
		action, err := decoder(b)
		if err != nil {
			return nil, errors.WithDetailf(errBadAction, "%s on action %d", err.Error(), i)
		}
		actions = append(actions, action)
	}

	ttl := req.TTL.Duration
	if ttl == 0 {
		ttl = defaultTxTTL
	}
	maxTime := time.Now().Add(ttl)

	tpl, err := txbuilder.Build(ctx, req.Tx, actions, maxTime, req.TimeRange)
	if errors.Root(err) == txbuilder.ErrAction {
		// append each of the inner errors contained in the data.
		var Errs string
		for _, innerErr := range errors.Data(err)["actions"].([]error) {
			Errs = Errs + "<" + innerErr.Error() + ">"
		}
		err = errors.New(err.Error() + "-" + Errs)
	}
	if err != nil {
		return nil, err
	}

	// ensure null is never returned for signing instructions
	if tpl.SigningInstructions == nil {
		tpl.SigningInstructions = []*txbuilder.SigningInstruction{}
	}
	return tpl, nil
}

// POST /build-transaction
func (a *API) build(ctx context.Context, buildReqs *BuildRequest) Response {
	subctx := reqid.NewSubContext(ctx, reqid.New())

	tmpl, err := a.buildSingle(subctx, buildReqs)
	if err != nil {
		return NewErrorResponse(err)
	}

	return NewSuccessResponse(tmpl)
}

type submitTxResp struct {
	TxID *bc.Hash `json:"tx_id"`
}

// POST /submit-transaction
func (a *API) submit(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	if err := txbuilder.FinalizeTx(ctx, a.chain, &ins.Tx); err != nil {
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", ins.Tx.ID).Info("submit single tx")
	return NewSuccessResponse(&submitTxResp{TxID: &ins.Tx.ID})
}

// EstimateTxGasResp estimate transaction consumed gas
type EstimateTxGasResp struct {
	TotalNeu   int64 `json:"total_neu"`
	StorageNeu int64 `json:"storage_neu"`
	VMNeu      int64 `json:"vm_neu"`
}

// EstimateTxGas estimate consumed neu for transaction
func EstimateTxGas(template txbuilder.Template) (*EstimateTxGasResp, error) {
	// base tx size and not include sign
	data, err := template.Transaction.TxData.MarshalText()
	if err != nil {
		return nil, err
	}
	baseTxSize := int64(len(data))

	// extra tx size for sign witness parts
	baseWitnessSize := int64(300)
	lenSignInst := int64(len(template.SigningInstructions))
	signSize := baseWitnessSize * lenSignInst

	// total gas for tx storage
	totalTxSizeGas, ok := checked.MulInt64(baseTxSize+signSize, consensus.StorageGasRate)
	if !ok {
		return nil, errors.New("calculate txsize gas got a math error")
	}

	// consume gas for run VM
	totalP2WPKHGas := int64(0)
	totalP2WSHGas := int64(0)
	baseP2WPKHGas := int64(1419)
	baseP2WSHGas := int64(2499)

	for _, inpID := range template.Transaction.Tx.InputIDs {
		sp, err := template.Transaction.Spend(inpID)
		if err != nil {
			continue
		}

		resOut, err := template.Transaction.Output(*sp.SpentOutputId)
		if err != nil {
			continue
		}

		if segwit.IsP2WPKHScript(resOut.ControlProgram.Code) {
			totalP2WPKHGas += baseP2WPKHGas
		} else if segwit.IsP2WSHScript(resOut.ControlProgram.Code) {
			totalP2WSHGas += baseP2WSHGas
		}
	}

	// total estimate gas
	totalGas := totalTxSizeGas + totalP2WPKHGas + totalP2WSHGas

	// rounding totalNeu with base 100000
	totalNeu := float64(totalGas*consensus.VMGasRate) / float64(100000)
	roundingNeu := math.Ceil(totalNeu)
	estimateNeu := int64(roundingNeu) * int64(100000)

	return &EstimateTxGasResp{
		TotalNeu:   estimateNeu,
		StorageNeu: totalTxSizeGas * consensus.VMGasRate,
		VMNeu:      (totalP2WPKHGas + totalP2WSHGas) * consensus.VMGasRate,
	}, nil
}

// POST /estimate-transaction-gas
func (a *API) estimateTxGas(ctx context.Context, in struct {
	TxTemplate txbuilder.Template `json:"transaction_template"`
}) Response {
	txGasResp, err := EstimateTxGas(in.TxTemplate)
	if err != nil {
		return NewErrorResponse(err)
	}
	return NewSuccessResponse(txGasResp)
}
<commit_msg>optimise<commit_after>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/consensus"
	"github.com/bytom/consensus/segwit"
	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/net/http/reqid"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

var (
	defaultTxTTL    = 5 * time.Minute
	defaultBaseRate = float64(100000)
)

func (a *API) actionDecoder(action string) (func([]byte) (txbuilder.Action, error), bool) {
	var decoder func([]byte) (txbuilder.Action, error)
	switch action {
	case "control_address":
		decoder = txbuilder.DecodeControlAddressAction
	case "control_program":
		decoder = txbuilder.DecodeControlProgramAction
	case "control_receiver":
		decoder = txbuilder.DecodeControlReceiverAction
	case "issue":
		decoder = a.wallet.AssetReg.DecodeIssueAction
	case "retire":
		decoder = txbuilder.DecodeRetireAction
	case "spend_account":
		decoder = a.wallet.AccountMgr.DecodeSpendAction
	case "spend_account_unspent_output":
		decoder = a.wallet.AccountMgr.DecodeSpendUTXOAction
	default:
		return nil, false
	}
	return decoder, true
}

// TODO modify mergeActions to loadSpendAction
func mergeActions(req *BuildRequest) ([]map[string]interface{}, error) {
	var actions []map[string]interface{}
	actionMap := make(map[string]map[string]interface{})

	for _, m := range req.Actions {
		if actionType := m["type"].(string); actionType != "spend_account" {
			actions = append(actions, m)
			continue
		}

		if m["amount"] == nil {
			return nil, errEmptyAmount
		}

		amountNumber := m["amount"].(json.Number)
		amount, err := amountNumber.Int64()
		if err != nil || amount == 0 {
			return nil, errBadAmount
		}

		actionKey := m["asset_id"].(string) + m["account_id"].(string)
		if tmpM, ok := actionMap[actionKey]; ok {
			if tmpM["amount"] == nil {
				return nil, errEmptyAmount
			}

			tmpNumber := tmpM["amount"].(json.Number)
			tmpAmount, err := tmpNumber.Int64()
			if err != nil || tmpAmount == 0 {
				return nil, errBadAmount
			}

			tmpM["amount"] = json.Number(fmt.Sprintf("%v", tmpAmount+amount))
		} else {
			actionMap[actionKey] = m
			actions = append(actions, m)
		}
	}

	return actions, nil
}

func onlyHaveSpendActions(req *BuildRequest) bool {
	count := 0
	for _, m := range req.Actions {
		if actionType := m["type"].(string); strings.HasPrefix(actionType, "spend") {
			count++
		}
	}

	return count == len(req.Actions)
}

func (a *API) buildSingle(ctx context.Context, req *BuildRequest) (*txbuilder.Template, error) {
	err := a.filterAliases(ctx, req)
	if err != nil {
		return nil, err
	}

	if onlyHaveSpendActions(req) {
		return nil, errors.New("transaction only contain spend actions, didn't have output actions")
	}

	reqActions, err := mergeActions(req)
	if err != nil {
		return nil, errors.WithDetail(err, "unmarshal json amount error in mergeActions")
	}

	actions := make([]txbuilder.Action, 0, len(reqActions))
	for i, act := range reqActions {
		typ, ok := act["type"].(string)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "no action type provided on action %d", i)
		}
		decoder, ok := a.actionDecoder(typ)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "unknown action type %q on action %d", typ, i)
		}

		// Remarshal to JSON, the action may have been modified when we
		// filtered aliases.
		b, err := json.Marshal(act)
		if err != nil {
			return nil, err
		}
		action, err := decoder(b)
		if err != nil {
			return nil, errors.WithDetailf(errBadAction, "%s on action %d", err.Error(), i)
		}
		actions = append(actions, action)
	}

	ttl := req.TTL.Duration
	if ttl == 0 {
		ttl = defaultTxTTL
	}
	maxTime := time.Now().Add(ttl)

	tpl, err := txbuilder.Build(ctx, req.Tx, actions, maxTime, req.TimeRange)
	if errors.Root(err) == txbuilder.ErrAction {
		// append each of the inner errors contained in the data.
		var Errs string
		for _, innerErr := range errors.Data(err)["actions"].([]error) {
			Errs = Errs + "<" + innerErr.Error() + ">"
		}
		err = errors.New(err.Error() + "-" + Errs)
	}
	if err != nil {
		return nil, err
	}

	// ensure null is never returned for signing instructions
	if tpl.SigningInstructions == nil {
		tpl.SigningInstructions = []*txbuilder.SigningInstruction{}
	}
	return tpl, nil
}

// POST /build-transaction
func (a *API) build(ctx context.Context, buildReqs *BuildRequest) Response {
	subctx := reqid.NewSubContext(ctx, reqid.New())

	tmpl, err := a.buildSingle(subctx, buildReqs)
	if err != nil {
		return NewErrorResponse(err)
	}

	return NewSuccessResponse(tmpl)
}

type submitTxResp struct {
	TxID *bc.Hash `json:"tx_id"`
}

// POST /submit-transaction
func (a *API) submit(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	if err := txbuilder.FinalizeTx(ctx, a.chain, &ins.Tx); err != nil {
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", ins.Tx.ID).Info("submit single tx")
	return NewSuccessResponse(&submitTxResp{TxID: &ins.Tx.ID})
}

// EstimateTxGasResp estimate transaction consumed gas
type EstimateTxGasResp struct {
	TotalNeu   int64 `json:"total_neu"`
	StorageNeu int64 `json:"storage_neu"`
	VMNeu      int64 `json:"vm_neu"`
}

// EstimateTxGas estimate consumed neu for transaction
func EstimateTxGas(template txbuilder.Template) (*EstimateTxGasResp, error) {
	// base tx size and not include sign
	data, err := template.Transaction.TxData.MarshalText()
	if err != nil {
		return nil, err
	}
	baseTxSize := int64(len(data))

	// extra tx size for sign witness parts
	baseWitnessSize := int64(300)
	lenSignInst := int64(len(template.SigningInstructions))
	signSize := baseWitnessSize * lenSignInst

	// total gas for tx storage
	totalTxSizeGas, ok := checked.MulInt64(baseTxSize+signSize, consensus.StorageGasRate)
	if !ok {
		return nil, errors.New("calculate txsize gas got a math error")
	}

	// consume gas for run VM
	totalP2WPKHGas := int64(0)
	totalP2WSHGas := int64(0)
	baseP2WPKHGas := int64(1419)
	baseP2WSHGas := int64(2499)

	for _, inpID := range template.Transaction.Tx.InputIDs {
		sp, err := template.Transaction.Spend(inpID)
		if err != nil {
			continue
		}

		resOut, err := template.Transaction.Output(*sp.SpentOutputId)
		if err != nil {
			continue
		}

		if segwit.IsP2WPKHScript(resOut.ControlProgram.Code) {
			totalP2WPKHGas += baseP2WPKHGas
		} else if segwit.IsP2WSHScript(resOut.ControlProgram.Code) {
			totalP2WSHGas += baseP2WSHGas
		}
	}

	// total estimate gas
	totalGas := totalTxSizeGas + totalP2WPKHGas + totalP2WSHGas

	// rounding totalNeu with base rate 100000
	totalNeu := float64(totalGas*consensus.VMGasRate) / defaultBaseRate
	roundingNeu := math.Ceil(totalNeu)
	estimateNeu := int64(roundingNeu) * int64(defaultBaseRate)

	return &EstimateTxGasResp{
		TotalNeu:   estimateNeu,
		StorageNeu: totalTxSizeGas * consensus.VMGasRate,
		VMNeu:      (totalP2WPKHGas + totalP2WSHGas) * consensus.VMGasRate,
	}, nil
}

// POST /estimate-transaction-gas
func (a *API) estimateTxGas(ctx context.Context, in struct {
	TxTemplate txbuilder.Template `json:"transaction_template"`
}) Response {
	txGasResp, err := EstimateTxGas(in.TxTemplate)
	if err != nil {
		return NewErrorResponse(err)
	}
	return NewSuccessResponse(txGasResp)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"html/template"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/nlopes/slack"
)

const (
	MessageGroupDisplayTimestampFormat = "3:04pm"
	MessageTextBlockquotePrefix        = ">"
	MessageTextControlRegexp           = "<(.*?)>"
)

type Message struct {
	*slack.Message
	timezoneLocation *time.Location
	slackClient      *slack.Client
}

func (m *Message) TimestampTime() time.Time {
	floatTimestamp, err := strconv.ParseFloat(m.Timestamp, 64)
	if err != nil {
		log.Println("Could not parse timestamp \"%s\".", m.Timestamp, err)
		return time.Time{}
	}
	return time.Unix(int64(floatTimestamp), 0).In(m.timezoneLocation)
}

func (m *Message) TextHtml() template.HTML {
	lines := strings.Split(m.Text, "\n")
	htmlPieces := []string{}
	controlRegexp := regexp.MustCompile(MessageTextControlRegexp)
	for _, line := range lines {
		linePrefix := ""
		lineSuffix := ""
		if strings.HasPrefix(line, MessageTextBlockquotePrefix) {
			line = strings.TrimPrefix(line, MessageTextBlockquotePrefix)
			if line == "" {
				// Ensure that even empty blockquote lines get rendered.
				line = "\u200b"
			}
			linePrefix = fmt.Sprintf("
",
				Style("message.blockquote"))
			lineSuffix = "
"
		} else {
			lineSuffix = "
"
		}
		line = controlRegexp.ReplaceAllStringFunc(line, func(control string) string {
			control = control[1 : len(control)-1]
			anchorText := ""
			pipeIndex := strings.LastIndex(control, "|")
			if pipeIndex != -1 {
				anchorText = control[pipeIndex+1:]
				control = control[:pipeIndex]
			}
			if strings.HasPrefix(control, "@U") {
				userId := strings.TrimPrefix(control, "@")
				userLookup, err := newUserLookup(m.slackClient)
				if err == nil {
					user, err := userLookup.GetUser(userId)
					if err == nil {
						anchorText = fmt.Sprintf("@%s", user.Name)
						authTest, err := m.slackClient.AuthTest()
						if err == nil {
							control = fmt.Sprintf("%s/team/%s", authTest.URL, user.Name)
						} else {
							log.Printf("Could get team URL: %s", err)
						}
					} else {
						log.Printf("Could not render user mention: %s", err)
					}
				} else {
					log.Printf("Could not render user mention: %s", err)
				}
			}
			if anchorText == "" {
				anchorText = control
			}
			return fmt.Sprintf("%s", control, anchorText)
		})

		htmlPieces = append(htmlPieces, linePrefix)
		// Slack's API claims that all HTML is already escaped
		htmlPieces = append(htmlPieces, line)
		htmlPieces = append(htmlPieces, lineSuffix)
	}
	return template.HTML(strings.Join(htmlPieces, ""))
}

type MessageGroup struct {
	Messages []*Message
	Author   *slack.User
}

func safeFormattedDate(date string) string {
	// Insert zero-width spaces every few characters so that Apple Data
	// Detectors and Gmail's calendar event dection don't pick up on these
	// dates.
	var buffer bytes.Buffer
	dateLength := len(date)
	for i := 0; i < dateLength; i += 2 {
		if i == dateLength-1 {
			buffer.WriteString(date[i : i+1])
		} else {
			buffer.WriteString(date[i : i+2])
			if date[i] != ' ' && date[i+1] != ' ' && i < dateLength-2 {
				buffer.WriteString("\u200b")
			}
		}
	}
	return buffer.String()
}

func (mg *MessageGroup) shouldContainMessage(message *Message, messageAuthor *slack.User) bool {
	if messageAuthor.ID != mg.Author.ID {
		return false
	}
	lastMessage := mg.Messages[len(mg.Messages)-1]
	timestampDelta := message.TimestampTime().Sub(lastMessage.TimestampTime())
	if timestampDelta > time.Minute*10 {
		return false
	}
	return true
}

func (mg *MessageGroup) DisplayTimestamp() string {
	return safeFormattedDate(mg.Messages[0].TimestampTime().Format(
		MessageGroupDisplayTimestampFormat))
}

func groupMessages(messages []*slack.Message, slackClient *slack.Client, timezoneLocation *time.Location) ([]*MessageGroup, error) {
	var currentGroup *MessageGroup = nil
	groups := make([]*MessageGroup, 0)
	userLookup, err := newUserLookup(slackClient)
	if err != nil {
		return nil, err
	}
	for i := range messages {
		message := &Message{messages[i], timezoneLocation, slackClient}
		if message.Hidden {
			continue
		}
		messageAuthor, _ := userLookup.GetUserForMessage(messages[i])
		if messageAuthor == nil {
			log.Printf("Could not determine author for message type %s "+
				"(subtype %s), skipping", message.Type, message.SubType)
			continue
		}
		if currentGroup == nil || !currentGroup.shouldContainMessage(message, messageAuthor) {
			currentGroup = &MessageGroup{
				Messages: make([]*Message, 0),
				Author:   messageAuthor,
			}
			groups = append(groups, currentGroup)
		}
		currentGroup.Messages = append(currentGroup.Messages, message)
	}
	return groups, nil
}
<commit_msg>Render channel mentions.<commit_after>package main

import (
	"bytes"
	"fmt"
	"html/template"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/nlopes/slack"
)

const (
	MessageGroupDisplayTimestampFormat = "3:04pm"
	MessageTextBlockquotePrefix        = ">"
	MessageTextControlRegexp           = "<(.*?)>"
)

type Message struct {
	*slack.Message
	timezoneLocation *time.Location
	slackClient      *slack.Client
}

func (m *Message) TimestampTime() time.Time {
	floatTimestamp, err := strconv.ParseFloat(m.Timestamp, 64)
	if err != nil {
		log.Println("Could not parse timestamp \"%s\".", m.Timestamp, err)
		return time.Time{}
	}
	return time.Unix(int64(floatTimestamp), 0).In(m.timezoneLocation)
}

func (m *Message) TextHtml() template.HTML {
	lines := strings.Split(m.Text, "\n")
	htmlPieces := []string{}
	controlRegexp := regexp.MustCompile(MessageTextControlRegexp)
	for _, line := range lines {
		linePrefix := ""
		lineSuffix := ""
		if strings.HasPrefix(line, MessageTextBlockquotePrefix) {
			line = strings.TrimPrefix(line, MessageTextBlockquotePrefix)
			if line == "" {
				// Ensure that even empty blockquote lines get rendered.
				line = "\u200b"
			}
			linePrefix = fmt.Sprintf("
",
				Style("message.blockquote"))
			lineSuffix = "
"
		} else {
			lineSuffix = "
"
		}
		line = controlRegexp.ReplaceAllStringFunc(line, func(control string) string {
			control = control[1 : len(control)-1]
			anchorText := ""
			pipeIndex := strings.LastIndex(control, "|")
			if pipeIndex != -1 {
				anchorText = control[pipeIndex+1:]
				control = control[:pipeIndex]
			}
			if strings.HasPrefix(control, "@U") {
				userId := strings.TrimPrefix(control, "@")
				userLookup, err := newUserLookup(m.slackClient)
				if err == nil {
					user, err := userLookup.GetUser(userId)
					if err == nil {
						anchorText = fmt.Sprintf("@%s", user.Name)
						authTest, err := m.slackClient.AuthTest()
						if err == nil {
							control = fmt.Sprintf("%s/team/%s", authTest.URL, user.Name)
						} else {
							log.Printf("Could get team URL: %s", err)
						}
					} else {
						log.Printf("Could not render user mention: %s", err)
					}
				} else {
					log.Printf("Could not render user mention: %s", err)
				}
			} else if strings.HasPrefix(control, "#C") {
				channelId := strings.TrimPrefix(control, "#")
				channel, err := m.slackClient.GetChannelInfo(channelId)
				if err == nil {
					anchorText = fmt.Sprintf("#%s", channel.Name)
					authTest, err := m.slackClient.AuthTest()
					if err == nil {
						control = fmt.Sprintf("%s/team/%s", authTest.URL, channel.Name)
					} else {
						log.Printf("Could get team URL: %s", err)
					}
				} else {
					log.Printf("Could not render channel mention: %s", err)
				}
			}
			if anchorText == "" {
				anchorText = control
			}
			return fmt.Sprintf("%s", control, anchorText)
		})

		htmlPieces = append(htmlPieces, linePrefix)
		// Slack's API claims that all HTML is already escaped
		htmlPieces = append(htmlPieces, line)
		htmlPieces = append(htmlPieces, lineSuffix)
	}
	return template.HTML(strings.Join(htmlPieces, ""))
}

type MessageGroup struct {
	Messages []*Message
	Author   *slack.User
}

func safeFormattedDate(date string) string {
	// Insert zero-width spaces every few characters so that Apple Data
	// Detectors and Gmail's calendar event dection don't pick up on these
	// dates.
	var buffer bytes.Buffer
	dateLength := len(date)
	for i := 0; i < dateLength; i += 2 {
		if i == dateLength-1 {
			buffer.WriteString(date[i : i+1])
		} else {
			buffer.WriteString(date[i : i+2])
			if date[i] != ' ' && date[i+1] != ' ' && i < dateLength-2 {
				buffer.WriteString("\u200b")
			}
		}
	}
	return buffer.String()
}

func (mg *MessageGroup) shouldContainMessage(message *Message, messageAuthor *slack.User) bool {
	if messageAuthor.ID != mg.Author.ID {
		return false
	}
	lastMessage := mg.Messages[len(mg.Messages)-1]
	timestampDelta := message.TimestampTime().Sub(lastMessage.TimestampTime())
	if timestampDelta > time.Minute*10 {
		return false
	}
	return true
}

func (mg *MessageGroup) DisplayTimestamp() string {
	return safeFormattedDate(mg.Messages[0].TimestampTime().Format(
		MessageGroupDisplayTimestampFormat))
}

func groupMessages(messages []*slack.Message, slackClient *slack.Client, timezoneLocation *time.Location) ([]*MessageGroup, error) {
	var currentGroup *MessageGroup = nil
	groups := make([]*MessageGroup, 0)
	userLookup, err := newUserLookup(slackClient)
	if err != nil {
		return nil, err
	}
	for i := range messages {
		message := &Message{messages[i], timezoneLocation, slackClient}
		if message.Hidden {
			continue
		}
		messageAuthor, _ := userLookup.GetUserForMessage(messages[i])
		if messageAuthor == nil {
			log.Printf("Could not determine author for message type %s "+
				"(subtype %s), skipping", message.Type, message.SubType)
			continue
		}
		if currentGroup == nil || !currentGroup.shouldContainMessage(message, messageAuthor) {
			currentGroup = &MessageGroup{
				Messages: make([]*Message, 0),
				Author:   messageAuthor,
			}
			groups = append(groups, currentGroup)
		}
		currentGroup.Messages = append(currentGroup.Messages, message)
	}
	return groups, nil
}
<|endoftext|>


<commit_before>package torrent

import (
	"container/heap"
	"fmt"
	"io"
	"math/rand"

	pp "bitbucket.org/anacrolix/go.torrent/peer_protocol"
)

type DownloadStrategy interface {
	// Tops up the outgoing pending requests. Returns the indices of pieces
	// that would be requested. This is used to determine if pieces require
	// hashing so the completed state is known.
	FillRequests(*torrent, *connection) (pieces []int)
	TorrentStarted(*torrent)
	TorrentStopped(*torrent)
	DeleteRequest(*torrent, request)
	TorrentPrioritize(t *torrent, off, _len int64)
	TorrentGotChunk(*torrent, request)
	TorrentGotPiece(t *torrent, piece int)
	WriteStatus(w io.Writer)
	AssertNotRequested(*torrent, request)
}

type DefaultDownloadStrategy struct {
	heat map[*torrent]map[request]int
}

func (me *DefaultDownloadStrategy) AssertNotRequested(t *torrent, r request) {
	if me.heat[t][r] != 0 {
		panic("outstanding requests break invariant")
	}
}

func (me *DefaultDownloadStrategy) WriteStatus(w io.Writer) {}

func (s *DefaultDownloadStrategy) FillRequests(t *torrent, c *connection) (pieces []int) {
	if c.Interested {
		if c.PeerChoked {
			return
		}
		if len(c.Requests) >= (c.PeerMaxRequests+1)/2 {
			return
		}
	}
	th := s.heat[t]
	addRequest := func(req request) (again bool) {
		piece := t.Pieces[req.Index]
		if piece.Hashing || piece.QueuedForHash {
			// We can't be sure we want this.
			return true
		}
		if piece.Complete() {
			// We already have this.
			return true
		}
		if c.RequestPending(req) {
			return true
		}
		again = c.Request(req)
		if c.RequestPending(req) {
			th[req]++
		}
		return
	}
	// Then finish off incomplete pieces in order of bytes remaining.
	for _, heatThreshold := range []int{1, 4, 15, 60} {
		for e := t.IncompletePiecesByBytesLeft.Front(); e != nil; e = e.Next() {
			pieceIndex := pp.Integer(e.Value.(int))
			piece := t.Pieces[pieceIndex]
			if !piece.EverHashed {
				pieces = append(pieces, int(pieceIndex))
				return
			}
			for chunkSpec := range t.Pieces[pieceIndex].PendingChunkSpecs {
				r := request{pieceIndex, chunkSpec}
				if th[r] >= heatThreshold {
					continue
				}
				if !addRequest(r) {
					return
				}
			}
		}
	}
	return
}

func (s *DefaultDownloadStrategy) TorrentStarted(t *torrent) {
	if s.heat[t] != nil {
		panic("torrent already started")
	}
	if s.heat == nil {
		s.heat = make(map[*torrent]map[request]int, 10)
	}
	s.heat[t] = make(map[request]int, t.ChunkCount())
}

func (s *DefaultDownloadStrategy) TorrentStopped(t *torrent) {
	if _, ok := s.heat[t]; !ok {
		panic("torrent not yet started")
	}
	delete(s.heat, t)
}

func (s *DefaultDownloadStrategy) DeleteRequest(t *torrent, r request) {
	m := s.heat[t]
	if m[r] <= 0 {
		panic("no pending requests")
	}
	m[r]--
}

func (me *DefaultDownloadStrategy) TorrentGotChunk(t *torrent, c request)      {}
func (me *DefaultDownloadStrategy) TorrentGotPiece(t *torrent, piece int)      {}
func (*DefaultDownloadStrategy) TorrentPrioritize(t *torrent, off, _len int64) {}

func NewResponsiveDownloadStrategy(readahead int64) *responsiveDownloadStrategy {
	return &responsiveDownloadStrategy{
		Readahead:      readahead,
		lastReadOffset: make(map[*torrent]int64),
		priorities:     make(map[*torrent]map[request]struct{}),
		requestHeat:    make(map[*torrent]map[request]int),
		rand:           rand.New(rand.NewSource(1337)),
	}
}

type responsiveDownloadStrategy struct {
	// How many bytes to preemptively download starting at the beginning of
	// the last piece read for a given torrent.
	Readahead      int64
	lastReadOffset map[*torrent]int64
	priorities     map[*torrent]map[request]struct{}
	requestHeat    map[*torrent]map[request]int
	rand           *rand.Rand // Avoid global lock
}

func (me *responsiveDownloadStrategy) WriteStatus(w io.Writer) {
	fmt.Fprintf(w, "Priorities:\n")
	for t, pp := range me.priorities {
		fmt.Fprintf(w, "\t%s:", t.Name())
		for r := range pp {
			fmt.Fprintf(w, " %v", r)
		}
		fmt.Fprintln(w)
	}
}

func (me *responsiveDownloadStrategy) TorrentStarted(t *torrent) {
	me.priorities[t] = make(map[request]struct{})
	me.requestHeat[t] = make(map[request]int)
}

func (me *responsiveDownloadStrategy) TorrentStopped(t *torrent) {
	delete(me.lastReadOffset, t)
	delete(me.priorities, t)
}
func (me *responsiveDownloadStrategy) DeleteRequest(t *torrent, r request) {
	rh := me.requestHeat[t]
	if rh[r] <= 0 {
		panic("request heat invariant broken")
	}
	rh[r]--
}

type requestFiller struct {
	c *connection
	t *torrent
	s *responsiveDownloadStrategy

	// The set of pieces that were considered for requesting.
	pieces map[int]struct{}
}

// Wrapper around connection.request that tracks request heat.
func (me *requestFiller) request(req request) bool {
	if me.pieces == nil {
		me.pieces = make(map[int]struct{})
	}
	me.pieces[int(req.Index)] = struct{}{}
	if me.c.RequestPending(req) {
		return true
	}
	if !me.t.wantChunk(req) {
		return true
	}
	again := me.c.Request(req)
	if me.c.RequestPending(req) {
		me.s.requestHeat[me.t][req]++
	}
	return again
}

// Adds additional constraints around the request heat wrapper.
func (me *requestFiller) conservativelyRequest(req request) bool {
	again := me.request(req)
	if len(me.c.Requests) >= 50 {
		return false
	}
	return again
}

// Fill priority requests.
func (me *requestFiller) priorities() bool {
	for req := range me.s.priorities[me.t] {
		// TODO: Perhaps this filter should be applied to every request?
		if _, ok := me.t.Pieces[req.Index].PendingChunkSpecs[req.chunkSpec]; !ok {
			panic(req)
		}
		if !me.request(req) {
			return false
		}
	}
	return true
}

// Fill requests, with all contextual information available in the receiver.
func (me requestFiller) Run() {
	if !me.priorities() {
		return
	}
	if len(me.c.Requests) > 25 {
		return
	}
	if !me.readahead() {
		return
	}
	if len(me.c.Requests) > 0 {
		return
	}
	me.completePartial()
}

// Request partial pieces that aren't in the readahead zone.
func (me *requestFiller) completePartial() bool {
	t := me.t
	th := me.s.requestHeat[t]
	lro, lroOk := me.s.lastReadOffset[t]
	for e := t.IncompletePiecesByBytesLeft.Front(); e != nil; e = e.Next() {
		p := e.Value.(int)
		// Stop when we reach pieces that aren't partial and aren't smaller
		// than usual.
		if !t.PiecePartiallyDownloaded(p) && int(t.PieceLength(pp.Integer(p))) == t.UsualPieceSize() {
			break
		}
		// Skip pieces that are entirely inside the readahead zone.
		if lroOk {
			pieceOff := int64(p) * int64(t.UsualPieceSize())
			pieceEndOff := pieceOff + int64(t.PieceLength(pp.Integer(p)))
			if pieceOff >= lro && pieceEndOff < lro+me.s.Readahead {
				continue
			}
		}
		for chunkSpec := range t.Pieces[p].PendingChunkSpecs {
			r := request{pp.Integer(p), chunkSpec}
			if th[r] >= 1 {
				continue
			}
			if lroOk {
				off := me.t.requestOffset(r)
				if off >= lro && off < lro+me.s.Readahead {
					continue
				}
			}
			if !me.conservativelyRequest(r) {
				return false
			}
		}
	}
	return true
}

// Returns all wanted chunk specs in the readahead zone.
func (me *requestFiller) pendingReadaheadChunks() (ret []request) {
	t := me.t
	lastReadOffset, ok := me.s.lastReadOffset[t]
	if !ok {
		return
	}
	ret = make([]request, 0, (me.s.Readahead+chunkSize-1)/chunkSize)
	for pi := int(lastReadOffset / int64(t.UsualPieceSize())); pi < t.NumPieces() && int64(pi)*int64(t.UsualPieceSize()) < lastReadOffset+me.s.Readahead; pi++ {
		if t.havePiece(pi) || !me.c.PeerHasPiece(pp.Integer(pi)) {
			continue
		}
		for cs := range t.Pieces[pi].PendingChunkSpecs {
			r := request{pp.Integer(pi), cs}
			if _, ok := me.c.Requests[r]; ok {
				continue
			}
			if off := t.requestOffset(r); off < lastReadOffset || off >= lastReadOffset+me.s.Readahead {
				continue
			}
			ret = append(ret, r)
		}
	}
	return
}

// Min-heap of int.
type intHeap []int

func (h intHeap) Len() int            { return len(h) }
func (h intHeap) Less(i, j int) bool  { return h[i] < h[j] }
func (h intHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *intHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *intHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func (me *requestFiller) readahead() bool {
	rr := me.pendingReadaheadChunks()
	if len(rr) == 0 {
		return true
	}
	// Produce a partially sorted random permutation into the readahead chunks
	// to somewhat preserve order but reducing wasted chunks due to overlap
	// with other peers.
	ii := new(intHeap)
	*ii = me.s.rand.Perm(len(rr))
	heap.Init(ii)
	for _, i := range *ii {
		if !me.conservativelyRequest(rr[i]) {
			return false
		}
	}
	return true
}

func (me *responsiveDownloadStrategy) FillRequests(t *torrent, c *connection) (pieces []int) {
	rf := requestFiller{c: c, t: t, s: me}
	rf.Run()
	for p := range rf.pieces {
		pieces = append(pieces, p)
	}
	return
}

func (me *responsiveDownloadStrategy) TorrentGotChunk(t *torrent, req request) {
	delete(me.priorities[t], req)
}

func (me *responsiveDownloadStrategy) TorrentGotPiece(t *torrent, piece int) {
	for _, cs := range t.pieceChunks(piece) {
		delete(me.priorities[t], request{pp.Integer(piece), cs})
	}
}

func (s *responsiveDownloadStrategy) TorrentPrioritize(t *torrent, off, _len int64) {
	s.lastReadOffset[t] = off
	for _len > 0 {
		req, ok := t.offsetRequest(off)
		if !ok {
			panic("bad offset")
		}
		reqOff := t.requestOffset(req)
		// Gain the alignment adjustment.
		_len += off - reqOff
		// Lose the length of this block.
		_len -= int64(req.Length)
		off = reqOff + int64(req.Length)
		if !t.haveChunk(req) {
			s.priorities[t][req] = struct{}{}
		}
	}
}

func (s *responsiveDownloadStrategy) AssertNotRequested(t *torrent, r request) {
	if s.requestHeat[t][r] != 0 {
		panic("outstanding requests invariant broken")
	}
}
<commit_msg>Fix accidental pass by value<commit_after>package torrent

import (
	"container/heap"
	"fmt"
	"io"
	"math/rand"

	pp "bitbucket.org/anacrolix/go.torrent/peer_protocol"
)

type DownloadStrategy interface {
	// Tops up the outgoing pending requests. Returns the indices of pieces
	// that would be requested. This is used to determine if pieces require
	// hashing so the completed state is known.
	FillRequests(*torrent, *connection) (pieces []int)
	TorrentStarted(*torrent)
	TorrentStopped(*torrent)
	DeleteRequest(*torrent, request)
	TorrentPrioritize(t *torrent, off, _len int64)
	TorrentGotChunk(*torrent, request)
	TorrentGotPiece(t *torrent, piece int)
	WriteStatus(w io.Writer)
	AssertNotRequested(*torrent, request)
}

type DefaultDownloadStrategy struct {
	heat map[*torrent]map[request]int
}

func (me *DefaultDownloadStrategy) AssertNotRequested(t *torrent, r request) {
	if me.heat[t][r] != 0 {
		panic("outstanding requests break invariant")
	}
}

func (me *DefaultDownloadStrategy) WriteStatus(w io.Writer) {}

func (s *DefaultDownloadStrategy) FillRequests(t *torrent, c *connection) (pieces []int) {
	if c.Interested {
		if c.PeerChoked {
			return
		}
		if len(c.Requests) >= (c.PeerMaxRequests+1)/2 {
			return
		}
	}
	th := s.heat[t]
	addRequest := func(req request) (again bool) {
		piece := t.Pieces[req.Index]
		if piece.Hashing || piece.QueuedForHash {
			// We can't be sure we want this.
			return true
		}
		if piece.Complete() {
			// We already have this.
			return true
		}
		if c.RequestPending(req) {
			return true
		}
		again = c.Request(req)
		if c.RequestPending(req) {
			th[req]++
		}
		return
	}
	// Then finish off incomplete pieces in order of bytes remaining.
	for _, heatThreshold := range []int{1, 4, 15, 60} {
		for e := t.IncompletePiecesByBytesLeft.Front(); e != nil; e = e.Next() {
			pieceIndex := pp.Integer(e.Value.(int))
			piece := t.Pieces[pieceIndex]
			if !piece.EverHashed {
				pieces = append(pieces, int(pieceIndex))
				return
			}
			for chunkSpec := range t.Pieces[pieceIndex].PendingChunkSpecs {
				r := request{pieceIndex, chunkSpec}
				if th[r] >= heatThreshold {
					continue
				}
				if !addRequest(r) {
					return
				}
			}
		}
	}
	return
}

func (s *DefaultDownloadStrategy) TorrentStarted(t *torrent) {
	if s.heat[t] != nil {
		panic("torrent already started")
	}
	if s.heat == nil {
		s.heat = make(map[*torrent]map[request]int, 10)
	}
	s.heat[t] = make(map[request]int, t.ChunkCount())
}

func (s *DefaultDownloadStrategy) TorrentStopped(t *torrent) {
	if _, ok := s.heat[t]; !ok {
		panic("torrent not yet started")
	}
	delete(s.heat, t)
}

func (s *DefaultDownloadStrategy) DeleteRequest(t *torrent, r request) {
	m := s.heat[t]
	if m[r] <= 0 {
		panic("no pending requests")
	}
	m[r]--
}

func (me *DefaultDownloadStrategy) TorrentGotChunk(t *torrent, c request)      {}
func (me *DefaultDownloadStrategy) TorrentGotPiece(t *torrent, piece int)      {}
func (*DefaultDownloadStrategy) TorrentPrioritize(t *torrent, off, _len int64) {}

func NewResponsiveDownloadStrategy(readahead int64) *responsiveDownloadStrategy {
	return &responsiveDownloadStrategy{
		Readahead:      readahead,
		lastReadOffset: make(map[*torrent]int64),
		priorities:     make(map[*torrent]map[request]struct{}),
		requestHeat:    make(map[*torrent]map[request]int),
		rand:           rand.New(rand.NewSource(1337)),
	}
}

type responsiveDownloadStrategy struct {
	// How many bytes to preemptively download starting at the beginning of
	// the last piece read for a given torrent.
	Readahead      int64
	lastReadOffset map[*torrent]int64
	priorities     map[*torrent]map[request]struct{}
	requestHeat    map[*torrent]map[request]int
	rand           *rand.Rand // Avoid global lock
}

func (me *responsiveDownloadStrategy) WriteStatus(w io.Writer) {
	fmt.Fprintf(w, "Priorities:\n")
	for t, pp := range me.priorities {
		fmt.Fprintf(w, "\t%s:", t.Name())
		for r := range pp {
			fmt.Fprintf(w, " %v", r)
		}
		fmt.Fprintln(w)
	}
}

func (me *responsiveDownloadStrategy) TorrentStarted(t *torrent) {
	me.priorities[t] = make(map[request]struct{})
	me.requestHeat[t] = make(map[request]int)
}

func (me *responsiveDownloadStrategy) TorrentStopped(t *torrent) {
	delete(me.lastReadOffset, t)
	delete(me.priorities, t)
}
func (me *responsiveDownloadStrategy) DeleteRequest(t *torrent, r request) {
	rh := me.requestHeat[t]
	if rh[r] <= 0 {
		panic("request heat invariant broken")
	}
	rh[r]--
}

type requestFiller struct {
	c *connection
	t *torrent
	s *responsiveDownloadStrategy

	// The set of pieces that were considered for requesting.
	pieces map[int]struct{}
}

// Wrapper around connection.request that tracks request heat.
func (me *requestFiller) request(req request) bool {
	if me.pieces == nil {
		me.pieces = make(map[int]struct{})
	}
	me.pieces[int(req.Index)] = struct{}{}
	if me.c.RequestPending(req) {
		return true
	}
	if !me.t.wantChunk(req) {
		return true
	}
	again := me.c.Request(req)
	if me.c.RequestPending(req) {
		me.s.requestHeat[me.t][req]++
	}
	return again
}

// Adds additional constraints around the request heat wrapper.
func (me *requestFiller) conservativelyRequest(req request) bool {
	again := me.request(req)
	if len(me.c.Requests) >= 50 {
		return false
	}
	return again
}

// Fill priority requests.
func (me *requestFiller) priorities() bool {
	for req := range me.s.priorities[me.t] {
		// TODO: Perhaps this filter should be applied to every request?
		if _, ok := me.t.Pieces[req.Index].PendingChunkSpecs[req.chunkSpec]; !ok {
			panic(req)
		}
		if !me.request(req) {
			return false
		}
	}
	return true
}

// Fill requests, with all contextual information available in the receiver.
func (me *requestFiller) Run() {
	if !me.priorities() {
		return
	}
	if len(me.c.Requests) > 25 {
		return
	}
	if !me.readahead() {
		return
	}
	if len(me.c.Requests) > 0 {
		return
	}
	me.completePartial()
}

// Request partial pieces that aren't in the readahead zone.
func (me *requestFiller) completePartial() bool {
	t := me.t
	th := me.s.requestHeat[t]
	lro, lroOk := me.s.lastReadOffset[t]
	for e := t.IncompletePiecesByBytesLeft.Front(); e != nil; e = e.Next() {
		p := e.Value.(int)
		// Stop when we reach pieces that aren't partial and aren't smaller
		// than usual.
		if !t.PiecePartiallyDownloaded(p) && int(t.PieceLength(pp.Integer(p))) == t.UsualPieceSize() {
			break
		}
		// Skip pieces that are entirely inside the readahead zone.
		if lroOk {
			pieceOff := int64(p) * int64(t.UsualPieceSize())
			pieceEndOff := pieceOff + int64(t.PieceLength(pp.Integer(p)))
			if pieceOff >= lro && pieceEndOff < lro+me.s.Readahead {
				continue
			}
		}
		for chunkSpec := range t.Pieces[p].PendingChunkSpecs {
			r := request{pp.Integer(p), chunkSpec}
			if th[r] >= 1 {
				continue
			}
			if lroOk {
				off := me.t.requestOffset(r)
				if off >= lro && off < lro+me.s.Readahead {
					continue
				}
			}
			if !me.conservativelyRequest(r) {
				return false
			}
		}
	}
	return true
}

// Returns all wanted chunk specs in the readahead zone.
func (me *requestFiller) pendingReadaheadChunks() (ret []request) {
	t := me.t
	lastReadOffset, ok := me.s.lastReadOffset[t]
	if !ok {
		return
	}
	ret = make([]request, 0, (me.s.Readahead+chunkSize-1)/chunkSize)
	for pi := int(lastReadOffset / int64(t.UsualPieceSize())); pi < t.NumPieces() && int64(pi)*int64(t.UsualPieceSize()) < lastReadOffset+me.s.Readahead; pi++ {
		if t.havePiece(pi) || !me.c.PeerHasPiece(pp.Integer(pi)) {
			continue
		}
		for cs := range t.Pieces[pi].PendingChunkSpecs {
			r := request{pp.Integer(pi), cs}
			if _, ok := me.c.Requests[r]; ok {
				continue
			}
			if off := t.requestOffset(r); off < lastReadOffset || off >= lastReadOffset+me.s.Readahead {
				continue
			}
			ret = append(ret, r)
		}
	}
	return
}

// Min-heap of int.
type intHeap []int

func (h intHeap) Len() int            { return len(h) }
func (h intHeap) Less(i, j int) bool  { return h[i] < h[j] }
func (h intHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *intHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *intHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func (me *requestFiller) readahead() bool {
	rr := me.pendingReadaheadChunks()
	if len(rr) == 0 {
		return true
	}
	// Produce a partially sorted random permutation into the readahead chunks
	// to somewhat preserve order but reducing wasted chunks due to overlap
	// with other peers.
	ii := new(intHeap)
	*ii = me.s.rand.Perm(len(rr))
	heap.Init(ii)
	for _, i := range *ii {
		if !me.conservativelyRequest(rr[i]) {
			return false
		}
	}
	return true
}

func (me *responsiveDownloadStrategy) FillRequests(t *torrent, c *connection) (pieces []int) {
	rf := requestFiller{c: c, t: t, s: me}
	rf.Run()
	for p := range rf.pieces {
		pieces = append(pieces, p)
	}
	return
}

func (me *responsiveDownloadStrategy) TorrentGotChunk(t *torrent, req request) {
	delete(me.priorities[t], req)
}

func (me *responsiveDownloadStrategy) TorrentGotPiece(t *torrent, piece int) {
	for _, cs := range t.pieceChunks(piece) {
		delete(me.priorities[t], request{pp.Integer(piece), cs})
	}
}

func (s *responsiveDownloadStrategy) TorrentPrioritize(t *torrent, off, _len int64) {
	s.lastReadOffset[t] = off
	for _len > 0 {
		req, ok := t.offsetRequest(off)
		if !ok {
			panic("bad offset")
		}
		reqOff := t.requestOffset(req)
		// Gain the alignment adjustment.
		_len += off - reqOff
		// Lose the length of this block.
		_len -= int64(req.Length)
		off = reqOff + int64(req.Length)
		if !t.haveChunk(req) {
			s.priorities[t][req] = struct{}{}
		}
	}
}

func (s *responsiveDownloadStrategy) AssertNotRequested(t *torrent, r request) {
	if s.requestHeat[t][r] != 0 {
		panic("outstanding requests invariant broken")
	}
}
<|endoftext|>