<commit_before>package cmdhandler

import (
	"errors"
	"fmt"
	"strings"
	"sync"

	"regexp"

	"github.com/lrstanley/girc"
)

// This may eventually get merged directly into girc in the future.

type Input struct {
	Origin *girc.Event
	Args   []string

	// TODO: utilize Event.Source and PRIVMSG's Params to lookup these to
	// make it easier on the end user.
	User    *girc.User
	Channel *girc.Channel
}

type Command struct {
	Help    string
	MinArgs int
	Fn      func(*girc.Client, *Input)
}

type CmdHandler struct {
	prefix string
	re     *regexp.Regexp

	mu   sync.Mutex
	cmds map[string]*Command
}

var cmdMatch = `^%s([a-z0-9-_]{1,20})(?: (.*))?$`

func New(prefix string) (*CmdHandler, error) {
	re, err := regexp.Compile(fmt.Sprintf(cmdMatch, regexp.QuoteMeta(prefix)))
	if err != nil {
		return nil, err
	}

	return &CmdHandler{prefix: prefix, re: re, cmds: make(map[string]*Command)}, nil
}

var validName = regexp.MustCompile(`^[a-zA-Z0-9-_]{1,20}$`)

func (ch *CmdHandler) Add(name string, cmd *Command) error {
	if cmd == nil {
		return errors.New("nil command provided to CmdHandler")
	}

	name = strings.ToLower(name)

	if !validName.MatchString(name) {
		return fmt.Errorf("invalid command name: %q (req: %q)", name, validName.String())
	}

	if cmd.MinArgs < 0 {
		cmd.MinArgs = 0
	}

	ch.mu.Lock()
	defer ch.mu.Unlock()

	if _, ok := ch.cmds[name]; ok {
		return fmt.Errorf("command already registered: %s", name)
	}

	ch.cmds[name] = cmd

	return nil
}

// Execute satisfies the girc.Handler interface.
func (ch *CmdHandler) Execute(client *girc.Client, event girc.Event) {
	if event.Source == nil || event.Command != girc.PRIVMSG {
		return
	}

	parsed := ch.re.FindStringSubmatch(event.Trailing)
	if len(parsed) != 3 {
		return
	}

	invCmd := strings.ToLower(parsed[1])
	args := strings.Split(parsed[2], " ")
	if len(args) == 1 && args[0] == "" {
		args = []string{}
	}

	ch.mu.Lock()
	defer ch.mu.Unlock()

	if invCmd == "help" {
		if len(args) == 0 {
			client.Commands.ReplyTo(event, girc.Fmt("type '{b}!help {blue}<command>{c}{b}' to optionally get more info about a specific command."))
			return
		}

		args[0] = strings.ToLower(args[0])

		if _, ok := ch.cmds[args[0]]; !ok {
			client.Commands.ReplyTof(event, girc.Fmt("unknown command {b}%q{b}."), args[0])
			return
		}

		if ch.cmds[args[0]].Help == "" {
			client.Commands.ReplyTof(event, girc.Fmt("there is no help documentaiton for {b}%q{b}"), args[0])
			return
		}

		client.Commands.ReplyTof(event, girc.Fmt("{b}%s%s{b} :: "+ch.cmds[args[0]].Help), ch.prefix, args[0])
		return
	}

	cmd, ok := ch.cmds[invCmd]
	if !ok {
		// TODO: return "no command found" if config option set to do so?
		return
	}

	if len(args) < cmd.MinArgs {
		client.Commands.ReplyTof(event, girc.Fmt("not enough arguments supplied for {b}%q{b}. try '{b}%shelp %s{b}'?"), invCmd, ch.prefix, invCmd)
		return
	}

	in := &Input{
		Origin: &event,
		Args:   args,
	}

	go cmd.Fn(client, in)
}
<commit_msg>clean up cmdhandler a bit<commit_after>package cmdhandler

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"

	"github.com/lrstanley/girc"
)

// Input is a wrapper for events, based around private messages.
type Input struct {
	Origin *girc.Event
	Args   []string
}

// Command is an IRC command, supporting aliases, help documentation and easy
// wrapping for message inputs.
type Command struct {
	// Name of command, e.g. "search" or "ping".
	Name string
	// Aliases for the above command, e.g. "s" for search, or "p" for "ping".
	Aliases []string
	// Help documentation. Should be in the format "<arg> <arg> [arg] --
	// something useful here"
	Help string
	// MinArgs is the minimum required arguments for the command. Defaults to
	// 0, which means multiple, or no arguments can be supplied. If set
	// above 0, this means that the command handler will throw an error asking
	// the person to check "<prefix>help <command>" for more info.
	MinArgs int
	// Fn is the function which is executed when the command is ran from a
	// private message, or channel.
	Fn func(*girc.Client, *Input)
}

func (c *Command) genHelp(prefix string) string {
	out := "{b}" + prefix + c.Name + "{b}"

	if c.Aliases != nil && len(c.Aliases) > 0 {
		out += " ({b}" + prefix + strings.Join(c.Aliases, "{b}, {b}"+prefix) + "{b})"
	}

	out += " :: " + c.Help

	return out
}

// CmdHandler is an irc command parser and execution format which you could
// use as an example for building your own version/bot.
//
// An example of how you would register this with girc:
//
//   ch, err := cmdhandler.New("!")
//   if err != nil {
//   	panic(err)
//   }

//   ch.Add(&cmdhandler.Command{
//   	Name:    "ping",
//   	Help:    "Sends a pong reply back to the original user.",
//   	Fn: func(c *girc.Client, input *cmdhandler.Input) {
//   		c.Commands.ReplyTo(*input.Origin, "pong!")
//   	},
//   })
//
//   client.Handlers.AddHandler(girc.PRIVMSG, ch)
type CmdHandler struct {
	prefix string
	re     *regexp.Regexp

	mu   sync.Mutex
	cmds map[string]*Command
}

var cmdMatch = `^%s([a-z0-9-_]{1,20})(?: (.*))?$`

// New returns a new CmdHandler based on the specified command prefix. A good
// prefix is a single character, and easy to remember/use. E.g. "!", or ".".
func New(prefix string) (*CmdHandler, error) {
	re, err := regexp.Compile(fmt.Sprintf(cmdMatch, regexp.QuoteMeta(prefix)))
	if err != nil {
		return nil, err
	}

	return &CmdHandler{prefix: prefix, re: re, cmds: make(map[string]*Command)}, nil
}

var validName = regexp.MustCompile(`^[a-z0-9-_]{1,20}$`)

// Add registers a new command to the handler. Note that you cannot remove
// commands once added, unless you add another CmdHandler to the client.
func (ch *CmdHandler) Add(cmd *Command) error {
	if cmd == nil {
		return errors.New("nil command provided to CmdHandler")
	}

	cmd.Name = strings.ToLower(cmd.Name)
	if !validName.MatchString(cmd.Name) {
		return fmt.Errorf("invalid command name: %q (req: %q)", cmd.Name, validName.String())
	}

	if cmd.Aliases != nil {
		for i := 0; i < len(cmd.Aliases); i++ {
			cmd.Aliases[i] = strings.ToLower(cmd.Aliases[i])
			if !validName.MatchString(cmd.Aliases[i]) {
				return fmt.Errorf("invalid command name: %q (req: %q)", cmd.Aliases[i], validName.String())
			}
		}
	}

	if cmd.MinArgs < 0 {
		cmd.MinArgs = 0
	}

	ch.mu.Lock()
	defer ch.mu.Unlock()

	if _, ok := ch.cmds[cmd.Name]; ok {
		return fmt.Errorf("command already registered: %s", cmd.Name)
	}

	ch.cmds[cmd.Name] = cmd

	// Since we'd be storing pointers, duplicates do not matter.
	for i := 0; i < len(cmd.Aliases); i++ {
		if _, ok := ch.cmds[cmd.Aliases[i]]; ok {
			return fmt.Errorf("alias already registered: %s", cmd.Aliases[i])
		}

		ch.cmds[cmd.Aliases[i]] = cmd
	}

	return nil
}

// Execute satisfies the girc.Handler interface.
func (ch *CmdHandler) Execute(client *girc.Client, event girc.Event) {
	if event.Source == nil || event.Command != girc.PRIVMSG {
		return
	}

	parsed := ch.re.FindStringSubmatch(event.Trailing)
	if len(parsed) != 3 {
		return
	}

	invCmd := strings.ToLower(parsed[1])
	args := strings.Split(parsed[2], " ")
	if len(args) == 1 && args[0] == "" {
		args = []string{}
	}

	ch.mu.Lock()
	defer ch.mu.Unlock()

	if invCmd == "help" {
		if len(args) == 0 {
			client.Commands.ReplyTo(event, girc.Fmt("type '{b}!help {blue}<command>{c}{b}' to optionally get more info about a specific command."))
			return
		}

		args[0] = strings.ToLower(args[0])

		if _, ok := ch.cmds[args[0]]; !ok {
			client.Commands.ReplyTof(event, girc.Fmt("unknown command {b}%q{b}."), args[0])
			return
		}

		if ch.cmds[args[0]].Help == "" {
			client.Commands.ReplyTof(event, girc.Fmt("there is no help documentaiton for {b}%q{b}"), args[0])
			return
		}

		client.Commands.ReplyTo(event, girc.Fmt(ch.cmds[args[0]].genHelp(ch.prefix)))
		return
	}

	cmd, ok := ch.cmds[invCmd]
	if !ok {
		return
	}

	if len(args) < cmd.MinArgs {
		client.Commands.ReplyTof(event, girc.Fmt("not enough arguments supplied for {b}%q{b}. try '{b}%shelp %s{b}'?"), invCmd, ch.prefix, invCmd)
		return
	}

	in := &Input{
		Origin: &event,
		Args:   args,
	}

	go cmd.Fn(client, in)
}
<|endoftext|>


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

package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/klauspost/pgzip"
	"github.com/spf13/pflag"
	"github.com/u-root/u-root/pkg/null"
)

const version string = "0.0.1"

type options struct {
	blocksize  int
	level      int
	processes  int
	decompress bool
	force      bool
	help       bool
	keep       bool
	license    bool
	quiet      bool
	stdout     bool
	test       bool
	verbose    bool
	version    bool
	suffix     string
}

func (o *options) parseArgs() error {
	var levels [10]bool

	pflag.IntVarP(&o.blocksize, "blocksize", "b", 128, "Set compression block size in KiB")
	pflag.BoolVarP(&o.decompress, "decompress", "d", false, "Decompress the compressed input")
	pflag.BoolVarP(&o.force, "force", "f", false, "Force overwrite of output file and compress links")
	pflag.BoolVarP(&o.help, "help", "h", false, "Display a help screen and quit")
	pflag.BoolVarP(&o.keep, "keep", "k", false, "Do not delete original file after processing")
	pflag.BoolVarP(&o.license, "license", "L", false, "Display license")
	// TODO: implement list option here
	pflag.IntVarP(&o.processes, "processes", "p", runtime.NumCPU(), "Allow up to n compression threads")
	pflag.BoolVarP(&o.quiet, "quiet", "q", false, "Print no messages, even on error")
	// TODO: implement recursive option here
	pflag.BoolVarP(&o.stdout, "stdout", "c", false, "Write all processed output to stdout (won't delete)")
	pflag.StringVarP(&o.suffix, "suffix", "S", ".gz", "Specify suffix for compression")
	pflag.BoolVarP(&o.test, "test", "t", false, "Test the integrity of the compressed input")
	pflag.BoolVarP(&o.verbose, "verbose", "v", false, "Produce more verbose output")
	pflag.BoolVarP(&o.version, "version", "V", false, "Print version")
	pflag.BoolVarP(&levels[1], "fast", "1", false, "Compression level 1")
	pflag.BoolVarP(&levels[2], "two", "2", false, "Compression level 2")
	pflag.BoolVarP(&levels[3], "three", "3", false, "Compression level 3")
	pflag.BoolVarP(&levels[4], "four", "4", false, "Compression level 4")
	pflag.BoolVarP(&levels[5], "five", "5", false, "Compression level 5")
	pflag.BoolVarP(&levels[6], "six", "6", false, "Compression level 6")
	pflag.BoolVarP(&levels[7], "seven", "7", false, "Compression level 7")
	pflag.BoolVarP(&levels[8], "eight", "8", false, "Compression level 8")
	pflag.BoolVarP(&levels[9], "best", "9", false, "Compression level 9")

	// Hide the compression levels 2 - 8 from usage.
	_ = pflag.CommandLine.MarkHidden("two")
	_ = pflag.CommandLine.MarkHidden("three")
	_ = pflag.CommandLine.MarkHidden("four")
	_ = pflag.CommandLine.MarkHidden("five")
	_ = pflag.CommandLine.MarkHidden("six")
	_ = pflag.CommandLine.MarkHidden("seven")
	_ = pflag.CommandLine.MarkHidden("eight")

	pflag.Parse()

	var err error
	o.level, err = parseLevels(levels)
	if err != nil {
		return err
	}

	return o.validate()
}

func (o *options) validate() error {

	if o.help {
		pflag.Usage()
		os.Exit(0)
	}

	if o.version {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if o.license {
		fmt.Printf("%s\n%s\n%s\n",
			version,
			"Copyright (c) 2012-2017 The u-root Authors. All rights reserved.",
			"Subject to the terms of the BSD 3-Clause license.")
		os.Exit(0)
	}

	if o.test {
		o.decompress = true
	}

	// Support gunzip and gzcat symlinks
	if filepath.Base(os.Args[0]) == "gunzip" {
		o.decompress = true
	} else if filepath.Base(os.Args[0]) == "gzcat" {
		o.decompress = true
		o.stdout = true
	}

	return nil
}

// parseLevels loops through a [10]bool and returns the index of the element
// thats true. If more than one element is true return an error. If no
// element is true, return the constant pgzip.DefaultCompression (-1).
func parseLevels(levels [10]bool) (int, error) {
	var level int

	for i, l := range levels {
		if l && level != 0 {
			return 0, errors.New("Multiple compression levels specified")
		} else if l {
			level = i
		}
	}

	if level == 0 {
		level = pgzip.DefaultCompression
	}
	return level, nil
}

// file is a file path to be compressed or decompressed.
type file struct {
	path    string
	options *options
}

// outputPath removes the path suffix on decompress and adds it on compress.
func (f *file) outputPath() string {
	if f.options.decompress {
		return f.path[:len(f.path)-len(f.options.suffix)]
	}
	return f.path + f.options.suffix
}

// checkPath validates the input file path. Checking on compression
// if the path has the correct suffix, and on decompression checks
// that it doesn't have the suffix. Allows override by force option.
func (f *file) checkPath() error {
	if !f.options.force {
		if f.options.decompress {
			if !strings.HasSuffix(f.path, f.options.suffix) {
				return fmt.Errorf("%s does not have %s suffix", f.path, f.options.suffix)
			}
		} else {
			if strings.HasSuffix(f.path, f.options.suffix) {
				return fmt.Errorf("%s already has %s suffix", f.path, f.options.suffix)
			}
		}
	}
	_, err := os.Stat(f.path)
	if os.IsNotExist(err) {
		return fmt.Errorf("%s does not exist", f.path)
	} else if os.IsPermission(err) {
		return fmt.Errorf("%s permission denied", f.path)
	}
	return f.checkOutPath()
}

// checkOutPath checks if output is attempting to write binary to the stdout.
// Check if output path already exists. Allow override via force option.
func (f *file) checkOutPath() error {
	if f.options.stdout {
		if !f.options.decompress && !f.options.force {
			return errors.New("trying to write compressed data to a terminal (use -f to force)")
		}
		return nil
	}
	_, err := os.Stat(f.outputPath())
	if !os.IsNotExist(err) && !f.options.force {
		return fmt.Errorf("%s already exist", f.outputPath())
	} else if os.IsPermission(err) {
		return fmt.Errorf("%s permission denied", f.outputPath())
	}
	return nil
}

// Cleanup removes input file. Overrided with keep option. Skipped if
// stdout or test option is true
func (f *file) cleanup() error {
	if !f.options.keep && !f.options.stdout && !f.options.test {
		return os.Remove(f.path)
	}
	return nil
}

func (f *file) process() error {
	i, err := os.Open(f.path)
	if err != nil {
		return err
	}
	defer i.Close()

	var o null.WriteNameCloser
	if f.options.stdout {
		o = os.Stdout
	} else if f.options.test {
		o = null.WriteNameClose
	} else {
		if o, err = os.Create(f.outputPath()); err != nil {
			return err
		}
	}

	if f.options.verbose && !f.options.quiet {
		fmt.Fprintf(os.Stderr, "%s to %s\n", i.Name(), o.Name())
	}

	if f.options.decompress {
		if err := decompressFile(i, o, f.options.blocksize, f.options.processes); err != nil {
			if !f.options.stdout {
				o.Close()
			}
			return err
		}
	} else {
		if err := compressFile(i, o, f.options.level, f.options.blocksize, f.options.processes); err != nil {
			if !f.options.stdout {
				o.Close()
			}
			return err
		}
	}

	if f.options.stdout {
		return nil
	}

	return o.Close()
}

func compressFile(r io.Reader, w io.Writer, level int, blocksize int, processes int) error {
	zw, err := pgzip.NewWriterLevel(w, level)
	if err != nil {
		return err
	}

	if err := zw.SetConcurrency(blocksize*1024, processes); err != nil {
		zw.Close()
		return err
	}

	if _, err := io.Copy(zw, r); err != nil {
		zw.Close()
		return err
	}

	return zw.Close()
}

func decompressFile(r io.Reader, w io.Writer, blocksize int, processes int) error {
	zr, err := pgzip.NewReaderN(r, blocksize*1024, processes)
	if err != nil {
		return err
	}

	if _, err := io.Copy(w, zr); err != nil {
		zr.Close()
		return err
	}

	return zr.Close()
}

func main() {

	var opts options

	if err := opts.parseArgs(); err != nil {
		fmt.Fprintf(os.Stderr, "Argument error: %s\n", err)
		os.Exit(1)
	}

	for _, path := range pflag.Args() {

		f := file{path: path, options: &opts}
		if err := f.checkPath(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "skipping, %s\n", err)
			}
			continue
		}

		if err := f.process(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "error, %s %s\n", f.path, err)
			}
			os.Exit(1)
		}

		if err := f.cleanup(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "warning, %s %s\n", f.path, err)
			}
			continue
		}
	}
}
<commit_msg>cmds/gzip: Add stdin pipe support, detect stdout redirect<commit_after>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/klauspost/pgzip"
	"github.com/spf13/pflag"
	"github.com/u-root/u-root/pkg/null"
)

const version string = "0.0.1"

type options struct {
	blocksize  int
	level      int
	processes  int
	decompress bool
	force      bool
	help       bool
	keep       bool
	license    bool
	quiet      bool
	stdin      bool
	stdout     bool
	test       bool
	verbose    bool
	version    bool
	suffix     string
}

func (o *options) parseArgs() error {
	var levels [10]bool

	pflag.IntVarP(&o.blocksize, "blocksize", "b", 128, "Set compression block size in KiB")
	pflag.BoolVarP(&o.decompress, "decompress", "d", false, "Decompress the compressed input")
	pflag.BoolVarP(&o.force, "force", "f", false, "Force overwrite of output file and compress links")
	pflag.BoolVarP(&o.help, "help", "h", false, "Display a help screen and quit")
	pflag.BoolVarP(&o.keep, "keep", "k", false, "Do not delete original file after processing")
	pflag.BoolVarP(&o.license, "license", "L", false, "Display license")
	// TODO: implement list option here
	pflag.IntVarP(&o.processes, "processes", "p", runtime.NumCPU(), "Allow up to n compression threads")
	pflag.BoolVarP(&o.quiet, "quiet", "q", false, "Print no messages, even on error")
	// TODO: implement recursive option here
	pflag.BoolVarP(&o.stdout, "stdout", "c", false, "Write all processed output to stdout (won't delete)")
	pflag.StringVarP(&o.suffix, "suffix", "S", ".gz", "Specify suffix for compression")
	pflag.BoolVarP(&o.test, "test", "t", false, "Test the integrity of the compressed input")
	pflag.BoolVarP(&o.verbose, "verbose", "v", false, "Produce more verbose output")
	pflag.BoolVarP(&o.version, "version", "V", false, "Print version")
	pflag.BoolVarP(&levels[1], "fast", "1", false, "Compression level 1")
	pflag.BoolVarP(&levels[2], "two", "2", false, "Compression level 2")
	pflag.BoolVarP(&levels[3], "three", "3", false, "Compression level 3")
	pflag.BoolVarP(&levels[4], "four", "4", false, "Compression level 4")
	pflag.BoolVarP(&levels[5], "five", "5", false, "Compression level 5")
	pflag.BoolVarP(&levels[6], "six", "6", false, "Compression level 6")
	pflag.BoolVarP(&levels[7], "seven", "7", false, "Compression level 7")
	pflag.BoolVarP(&levels[8], "eight", "8", false, "Compression level 8")
	pflag.BoolVarP(&levels[9], "best", "9", false, "Compression level 9")

	// Hide the compression levels 2 - 8 from usage.
	_ = pflag.CommandLine.MarkHidden("two")
	_ = pflag.CommandLine.MarkHidden("three")
	_ = pflag.CommandLine.MarkHidden("four")
	_ = pflag.CommandLine.MarkHidden("five")
	_ = pflag.CommandLine.MarkHidden("six")
	_ = pflag.CommandLine.MarkHidden("seven")
	_ = pflag.CommandLine.MarkHidden("eight")

	pflag.Parse()

	var err error
	o.level, err = parseLevels(levels)
	if err != nil {
		return err
	}

	return o.validate()
}

// Validate checks options and handles help, version, and license
// output to the user. If forces decompression to be enabled when
// test mode is enabled. It further modifies options if the running
// binary is named gunzip or gzcat to allow for expected behavor for
// those binaries. Finally it checks if there is piped stdin data.
func (o *options) validate() error {

	if o.help {
		pflag.Usage()
		os.Exit(0)
	}

	if o.version {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if o.license {
		fmt.Printf("%s\n%s\n%s\n",
			version,
			"Copyright (c) 2012-2017 The u-root Authors. All rights reserved.",
			"Subject to the terms of the BSD 3-Clause license.")
		os.Exit(0)
	}

	if o.test {
		o.decompress = true
	}

	// Support gunzip and gzcat symlinks
	if filepath.Base(os.Args[0]) == "gunzip" {
		o.decompress = true
	} else if filepath.Base(os.Args[0]) == "gzcat" {
		o.decompress = true
		o.stdout = true
	}

	// Stat os.Stdin and ignore errors. stat will be nil FileInfo if there is an
	// error.
	stat, _ := os.Stdin.Stat()

	// No files passed and arguments and Stdin piped data found.
	// Stdin piped data is ignored if arguments are found.
	if len(pflag.Args()) == 0 && (stat.Mode()&os.ModeNamedPipe) != 0 {
		o.stdin = true
		// Enable force to ignore suffix checks
		o.force = true
		// Since there's no filename to derive the output path from, only support
		// outputting to stdout when data is piped from stdin
		o.stdout = true
	} else if len(pflag.Args()) == 0 {
		// No stdin piped data found and no files passed as arguments
		pflag.Usage()
		os.Exit(0)
	}

	return nil
}

// parseLevels loops through a [10]bool and returns the index of the element
// thats true. If more than one element is true return an error. If no
// element is true, return the constant pgzip.DefaultCompression (-1).
func parseLevels(levels [10]bool) (int, error) {
	var level int

	for i, l := range levels {
		if l && level != 0 {
			return 0, errors.New("Multiple compression levels specified")
		} else if l {
			level = i
		}
	}

	if level == 0 {
		level = pgzip.DefaultCompression
	}
	return level, nil
}

// file is a file path to be compressed or decompressed.
type file struct {
	path    string
	options *options
}

// outputPath removes the path suffix on decompress and adds it on compress.
// In the case of when options stdout or test are enabled it returns the path
// as is.
func (f *file) outputPath() string {
	if f.options.stdout || f.options.test {
		return f.path
	} else if f.options.decompress {
		return f.path[:len(f.path)-len(f.options.suffix)]
	}
	return f.path + f.options.suffix
}

// checkPath validates the input file path. Checks on compression
// if the path has the correct suffix, and on decompression checks
// that it doesn't have the suffix. Allows override by force option.
func (f *file) checkPath() error {
	_, err := os.Stat(f.path)
	if os.IsNotExist(err) {
		return fmt.Errorf("%s does not exist", f.path)
	} else if os.IsPermission(err) {
		return fmt.Errorf("%s permission denied", f.path)
	}

	if !f.options.force {
		if f.options.decompress {
			if !strings.HasSuffix(f.path, f.options.suffix) {
				return fmt.Errorf("%s does not have %s suffix", f.path, f.options.suffix)
			}
		} else {
			if strings.HasSuffix(f.path, f.options.suffix) {
				return fmt.Errorf("%s already has %s suffix", f.path, f.options.suffix)
			}
		}
	}
	return f.checkOutPath()
}

// checkOutPath checks if output is attempting to write binary to stdout if
// stdout is a device. Also checks if output path already exists. Allow
// override via force option.
func (f *file) checkOutPath() error {
	if f.options.stdout {
		stat, _ := os.Stdout.Stat()
		if !f.options.decompress && !f.options.force && (stat.Mode()&os.ModeDevice) != 0 {
			return errors.New("trying to write compressed data to a terminal/device (use -f to force)")
		}
		return nil
	}
	_, err := os.Stat(f.outputPath())
	if !os.IsNotExist(err) && !f.options.stdout && !f.options.test && !f.options.force {
		return fmt.Errorf("%s already exist", f.outputPath())
	} else if os.IsPermission(err) {
		return fmt.Errorf("%s permission denied", f.outputPath())
	}
	return nil
}

// Cleanup removes input file. Overrided with keep option. Skipped if
// stdout or test option is true.
func (f *file) cleanup() error {
	if !f.options.keep && !f.options.stdout && !f.options.test {
		return os.Remove(f.path)
	}
	return nil
}

// Process either compresses or decompressed the input file based on
// the associated file.options.
func (f *file) process() error {
	i, err := os.Open(f.path)
	if err != nil {
		return err
	}
	defer i.Close()

	// Use the null.WriteNameCloser interface so both *os.File and
	// null.WriteNameClose can be assigned to var o without any type casting below.
	var o null.WriteNameCloser

	if f.options.test {
		o = null.WriteNameClose
	} else if f.options.stdout {
		o = os.Stdout
	} else {
		if o, err = os.Create(f.outputPath()); err != nil {
			return err
		}
	}

	if f.options.verbose && !f.options.quiet {
		fmt.Fprintf(os.Stderr, "%s to %s\n", i.Name(), o.Name())
	}

	if f.options.decompress {
		if err := decompressFile(i, o, f.options.blocksize, f.options.processes); err != nil {
			if !f.options.stdout {
				o.Close()
			}
			return err
		}
	} else {
		if err := compressFile(i, o, f.options.level, f.options.blocksize, f.options.processes); err != nil {
			if !f.options.stdout {
				o.Close()
			}
			return err
		}
	}

	if f.options.stdout {
		return nil
	}

	return o.Close()
}

func compressFile(r io.Reader, w io.Writer, level int, blocksize int, processes int) error {
	zw, err := pgzip.NewWriterLevel(w, level)
	if err != nil {
		return err
	}

	if err := zw.SetConcurrency(blocksize*1024, processes); err != nil {
		zw.Close()
		return err
	}

	if _, err := io.Copy(zw, r); err != nil {
		zw.Close()
		return err
	}

	return zw.Close()
}

func decompressFile(r io.Reader, w io.Writer, blocksize int, processes int) error {
	zr, err := pgzip.NewReaderN(r, blocksize*1024, processes)
	if err != nil {
		return err
	}

	if _, err := io.Copy(w, zr); err != nil {
		zr.Close()
		return err
	}

	return zr.Close()
}

func main() {

	var opts options

	if err := opts.parseArgs(); err != nil {
		fmt.Fprintf(os.Stderr, "Argument error: %s\n", err)
		os.Exit(1)
	}

	var files []string

	if opts.stdin {
		files = []string{"/dev/stdin"}
	} else {
		files = pflag.Args()
	}

	for _, path := range files {

		f := file{path: path, options: &opts}
		if err := f.checkPath(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "skipping, %s\n", err)
			}
			continue
		}

		if err := f.process(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "error, %s %s\n", f.path, err)
			}
			os.Exit(1)
		}

		if err := f.cleanup(); err != nil {
			if !opts.quiet {
				fmt.Fprintf(os.Stderr, "warning, %s %s\n", f.path, err)
			}
			continue
		}
	}
}
<|endoftext|>


<commit_before>package comments

import (
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
	"sync"

	"github.com/google/go-github/github"
	"github.com/parkr/changelog"
)

var (
	mergeCommentRegexp = regexp.MustCompile("@[a-zA-Z-_]+: (merge|:shipit:|:ship:)( \\+([a-zA-Z-_ ]+))?")

	HandlerMergeAndLabel = func(client *github.Client, event github.IssueCommentEvent) error {
		// Is this a pull request?
		if !isPullRequest(event) {
			return errors.New("not a pull request")
		}

		var changeSectionLabel string
		isReq, labelFromComment := parseMergeRequestComment(*event.Comment.Body)

		// Is It a merge request comment?
		if !isReq {
			return errors.New("not a merge request comment")
		}

		if os.Getenv("AUTO_REPLY_DEBUG") == "true" {
			log.Println("received event:", event)
		}

		var wg sync.WaitGroup

		owner := *event.Repo.Owner.Login
		repo := *event.Repo.Name
		number := *event.Issue.Number

		// Does the user have merge/label abilities?
		if !isAuthorizedCommenter(client, event) {
			log.Printf("%s isn't authenticated to merge anything on %s", *event.Comment.User.Login, *event.Repo.FullName)
			return errors.New("commenter isn't allowed to merge")
		}

		// Should it be labeled?
		if labelFromComment != "" {
			changeSectionLabel = sectionForLabel(labelFromComment)
		} else {
			changeSectionLabel = "none"
		}
		fmt.Printf("changeSectionLabel = '%s'\n", changeSectionLabel)

		// Merge
		commitMsg := fmt.Sprintf("Merge pull request %v", number)
		_, _, mergeErr := client.PullRequests.Merge(owner, repo, number, commitMsg)
		if mergeErr != nil {
			fmt.Printf("comments: error merging %v\n", mergeErr)
			return mergeErr
		}

		// Delete branch
		repoInfo, _, getRepoErr := client.PullRequests.Get(owner, repo, number)
		if getRepoErr != nil {
			fmt.Printf("comments: error fetching pull request: %v\n", getRepoErr)
			return getRepoErr
		}

		// Delete branch
		if deletableRef(repoInfo, owner) {
			wg.Add(1)
			go func() {
				ref := fmt.Sprintf("heads/%s", *repoInfo.Head.Ref)
				_, deleteBranchErr := client.Git.DeleteRef(owner, repo, ref)
				if deleteBranchErr != nil {
					fmt.Printf("comments: error deleting branch %v\n", mergeErr)
				}
				wg.Done()
			}()
		}

		wg.Add(1)
		go func() {
			err := addLabelsForSubsection(client, owner, repo, number, changeSectionLabel)
			if err != nil {
				fmt.Printf("comments: error applying labels: %v\n", err)
			}
			wg.Done()
		}()

		wg.Add(1)
		go func() {
			// Read History.markdown, add line to appropriate change section
			historyFileContents, historySHA := getHistoryContents(client, owner, repo)

			// Add to
			newHistoryFileContents := addMergeReference(historyFileContents, changeSectionLabel, *repoInfo.Title, number)

			// Commit change to History.markdown
			commitErr := commitHistoryFile(client, historySHA, owner, repo, number, newHistoryFileContents)
			if commitErr != nil {
				fmt.Printf("comments: error committing updated history %v\n", mergeErr)
			}
			wg.Done()
		}()

		wg.Wait()

		return nil
	}
)

func parseMergeRequestComment(commentBody string) (bool, string) {
	matches := mergeCommentRegexp.FindAllStringSubmatch(commentBody, -1)
	if matches == nil || matches[0] == nil {
		return false, ""
	}

	var label string
	if len(matches[0]) >= 4 {
		if labelFromComment := matches[0][3]; labelFromComment != "" {
			label = downcaseAndHyphenize(labelFromComment)
		}
	}

	return true, normalizeLabel(label)
}

func downcaseAndHyphenize(label string) string {
	return strings.Replace(strings.ToLower(label), " ", "-", -1)
}

func normalizeLabel(label string) string {
	if strings.HasPrefix(label, "major") {
		return "major-enhancements"
	}

	if strings.HasPrefix(label, "minor") {
		return "minor-enhancements"
	}

	if strings.HasPrefix(label, "bug") {
		return "bug-fixes"
	}

	if strings.HasPrefix(label, "dev") {
		return "development-fixes"
	}

	if strings.HasPrefix(label, "site") {
		return "site-enhancements"
	}

	return label
}

func sectionForLabel(label string) string {
	switch label {
	case "major-enhancements":
		return "Major Enhancements"
	case "minor-enhancements":
		return "Minor Enhancements"
	case "bug-fixes":
		return "Bug Fixes"
	case "development-fixes":
		return "Development Fixes"
	case "site-enhancements":
		return "Site Enhancements"
	default:
		return label
	}
}

func selectSectionLabel(labels []github.Label) string {
	for _, label := range labels {
		if sectionForLabel(*label.Name) != *label.Name {
			return *label.Name
		}
	}
	return ""
}

func containsChangeLabel(commentBody string) bool {
	_, labelFromComment := parseMergeRequestComment(commentBody)
	return labelFromComment != ""
}

func labelsForSubsection(changeSectionLabel string) []string {
	labels := []string{}

	switch changeSectionLabel {
	case "Major Enhancements":
		labels = append(labels, "feature")
	case "Minor Enhancements":
		labels = append(labels, "enhancement")
	case "Bug Fixes":
		labels = append(labels, "bug")
		labels = append(labels, "fix")
	case "Development Fixes":
		labels = append(labels, "internal")
		labels = append(labels, "fix")
	case "Site Enhancements":
		labels = append(labels, "documentation")
	}

	return labels
}

func addLabelsForSubsection(client *github.Client, owner, repo string, number int, changeSectionLabel string) error {
	labels := labelsForSubsection(changeSectionLabel)

	if len(labels) < 1 {
		return fmt.Errorf("no labels for changeSectionLabel='%s'", changeSectionLabel)
	}

	_, _, err := client.Issues.AddLabelsToIssue(owner, repo, number, labels)
	return err
}

func getHistoryContents(client *github.Client, owner, repo string) (content, sha string) {
	contents, _, _, err := client.Repositories.GetContents(
		owner,
		repo,
		"History.markdown",
		&github.RepositoryContentGetOptions{Ref: "heads/master"},
	)
	if err != nil {
		fmt.Printf("comments: error getting History.markdown %v\n", err)
		return "", ""
	}
	return base64Decode(*contents.Content), *contents.SHA
}

func base64Decode(encoded string) string {
	decoded, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		fmt.Printf("comments: error decoding string: %v\n", err)
		return ""
	}
	return string(decoded)
}

func addMergeReference(historyFileContents, changeSectionLabel, prTitle string, number int) string {
	changes, err := changelog.NewChangelogFromReader(strings.NewReader(historyFileContents))
	if historyFileContents == "" {
		err = nil
		changes = changelog.NewChangelog()
	}
	if err != nil {
		fmt.Printf("comments: error %v\n", err)
		return historyFileContents
	}

	changeLine := &changelog.ChangeLine{
		Summary:   prTitle,
		Reference: fmt.Sprintf("#%d", number),
	}

	// Put either directly in the version history or in a subsection.
	if changeSectionLabel == "none" {
		changes.AddLineToVersion("HEAD", changeLine)
	} else {
		changes.AddLineToSubsection("HEAD", changeSectionLabel, changeLine)
	}

	return changes.String()
}

func deletableRef(pr *github.PullRequest, owner string) bool {
	return *pr.Head.Repo.Owner.Login == owner && *pr.Head.Ref != "master" && *pr.Head.Ref != "gh-pages"
}

func commitHistoryFile(client *github.Client, historySHA, owner, repo string, number int, newHistoryFileContents string) error {
	repositoryContentsOptions := &github.RepositoryContentFileOptions{
		Message: github.String(fmt.Sprintf("Update history to reflect merge of #%d [ci skip]", number)),
		Content: []byte(newHistoryFileContents),
		SHA:     github.String(historySHA),
		Committer: &github.CommitAuthor{
			Name:  github.String("jekyllbot"),
			Email: github.String("jekyllbot@jekyllrb.com"),
		},
	}
	updateResponse, _, err := client.Repositories.UpdateFile(owner, repo, "History.markdown", repositoryContentsOptions)
	if err != nil {
		fmt.Printf("comments: error committing History.markdown: %v\n", err)
		return err
	}
	fmt.Printf("comments: updateResponse: %s\n", updateResponse)
	return nil
}
<commit_msg>for merge_and_label too<commit_after>package comments

import (
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
	"sync"

	"github.com/google/go-github/github"
	"github.com/parkr/changelog"
)

var (
	mergeCommentRegexp = regexp.MustCompile("@[a-zA-Z-_]+: (merge|:shipit:|:ship:)( \\+([a-zA-Z-_ ]+))?")

	HandlerMergeAndLabel = func(client *github.Client, event github.IssueCommentEvent) error {
		// Is this a pull request?
		if !isPullRequest(event) {
			return errors.New("not a pull request")
		}

		var changeSectionLabel string
		isReq, labelFromComment := parseMergeRequestComment(*event.Comment.Body)

		// Is It a merge request comment?
		if !isReq {
			return errors.New("not a merge request comment")
		}

		if os.Getenv("AUTO_REPLY_DEBUG") == "true" {
			log.Println("[merge_and_label]: received event:", event)
		}

		var wg sync.WaitGroup

		owner := *event.Repo.Owner.Login
		repo := *event.Repo.Name
		number := *event.Issue.Number

		// Does the user have merge/label abilities?
		if !isAuthorizedCommenter(client, event) {
			log.Printf("%s isn't authenticated to merge anything on %s", *event.Comment.User.Login, *event.Repo.FullName)
			return errors.New("commenter isn't allowed to merge")
		}

		// Should it be labeled?
		if labelFromComment != "" {
			changeSectionLabel = sectionForLabel(labelFromComment)
		} else {
			changeSectionLabel = "none"
		}
		fmt.Printf("changeSectionLabel = '%s'\n", changeSectionLabel)

		// Merge
		commitMsg := fmt.Sprintf("Merge pull request %v", number)
		_, _, mergeErr := client.PullRequests.Merge(owner, repo, number, commitMsg)
		if mergeErr != nil {
			fmt.Printf("comments: error merging %v\n", mergeErr)
			return mergeErr
		}

		// Delete branch
		repoInfo, _, getRepoErr := client.PullRequests.Get(owner, repo, number)
		if getRepoErr != nil {
			fmt.Printf("comments: error fetching pull request: %v\n", getRepoErr)
			return getRepoErr
		}

		// Delete branch
		if deletableRef(repoInfo, owner) {
			wg.Add(1)
			go func() {
				ref := fmt.Sprintf("heads/%s", *repoInfo.Head.Ref)
				_, deleteBranchErr := client.Git.DeleteRef(owner, repo, ref)
				if deleteBranchErr != nil {
					fmt.Printf("comments: error deleting branch %v\n", mergeErr)
				}
				wg.Done()
			}()
		}

		wg.Add(1)
		go func() {
			err := addLabelsForSubsection(client, owner, repo, number, changeSectionLabel)
			if err != nil {
				fmt.Printf("comments: error applying labels: %v\n", err)
			}
			wg.Done()
		}()

		wg.Add(1)
		go func() {
			// Read History.markdown, add line to appropriate change section
			historyFileContents, historySHA := getHistoryContents(client, owner, repo)

			// Add to
			newHistoryFileContents := addMergeReference(historyFileContents, changeSectionLabel, *repoInfo.Title, number)

			// Commit change to History.markdown
			commitErr := commitHistoryFile(client, historySHA, owner, repo, number, newHistoryFileContents)
			if commitErr != nil {
				fmt.Printf("comments: error committing updated history %v\n", mergeErr)
			}
			wg.Done()
		}()

		wg.Wait()

		return nil
	}
)

func parseMergeRequestComment(commentBody string) (bool, string) {
	matches := mergeCommentRegexp.FindAllStringSubmatch(commentBody, -1)
	if matches == nil || matches[0] == nil {
		return false, ""
	}

	var label string
	if len(matches[0]) >= 4 {
		if labelFromComment := matches[0][3]; labelFromComment != "" {
			label = downcaseAndHyphenize(labelFromComment)
		}
	}

	return true, normalizeLabel(label)
}

func downcaseAndHyphenize(label string) string {
	return strings.Replace(strings.ToLower(label), " ", "-", -1)
}

func normalizeLabel(label string) string {
	if strings.HasPrefix(label, "major") {
		return "major-enhancements"
	}

	if strings.HasPrefix(label, "minor") {
		return "minor-enhancements"
	}

	if strings.HasPrefix(label, "bug") {
		return "bug-fixes"
	}

	if strings.HasPrefix(label, "dev") {
		return "development-fixes"
	}

	if strings.HasPrefix(label, "site") {
		return "site-enhancements"
	}

	return label
}

func sectionForLabel(label string) string {
	switch label {
	case "major-enhancements":
		return "Major Enhancements"
	case "minor-enhancements":
		return "Minor Enhancements"
	case "bug-fixes":
		return "Bug Fixes"
	case "development-fixes":
		return "Development Fixes"
	case "site-enhancements":
		return "Site Enhancements"
	default:
		return label
	}
}

func selectSectionLabel(labels []github.Label) string {
	for _, label := range labels {
		if sectionForLabel(*label.Name) != *label.Name {
			return *label.Name
		}
	}
	return ""
}

func containsChangeLabel(commentBody string) bool {
	_, labelFromComment := parseMergeRequestComment(commentBody)
	return labelFromComment != ""
}

func labelsForSubsection(changeSectionLabel string) []string {
	labels := []string{}

	switch changeSectionLabel {
	case "Major Enhancements":
		labels = append(labels, "feature")
	case "Minor Enhancements":
		labels = append(labels, "enhancement")
	case "Bug Fixes":
		labels = append(labels, "bug")
		labels = append(labels, "fix")
	case "Development Fixes":
		labels = append(labels, "internal")
		labels = append(labels, "fix")
	case "Site Enhancements":
		labels = append(labels, "documentation")
	}

	return labels
}

func addLabelsForSubsection(client *github.Client, owner, repo string, number int, changeSectionLabel string) error {
	labels := labelsForSubsection(changeSectionLabel)

	if len(labels) < 1 {
		return fmt.Errorf("no labels for changeSectionLabel='%s'", changeSectionLabel)
	}

	_, _, err := client.Issues.AddLabelsToIssue(owner, repo, number, labels)
	return err
}

func getHistoryContents(client *github.Client, owner, repo string) (content, sha string) {
	contents, _, _, err := client.Repositories.GetContents(
		owner,
		repo,
		"History.markdown",
		&github.RepositoryContentGetOptions{Ref: "heads/master"},
	)
	if err != nil {
		fmt.Printf("comments: error getting History.markdown %v\n", err)
		return "", ""
	}
	return base64Decode(*contents.Content), *contents.SHA
}

func base64Decode(encoded string) string {
	decoded, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		fmt.Printf("comments: error decoding string: %v\n", err)
		return ""
	}
	return string(decoded)
}

func addMergeReference(historyFileContents, changeSectionLabel, prTitle string, number int) string {
	changes, err := changelog.NewChangelogFromReader(strings.NewReader(historyFileContents))
	if historyFileContents == "" {
		err = nil
		changes = changelog.NewChangelog()
	}
	if err != nil {
		fmt.Printf("comments: error %v\n", err)
		return historyFileContents
	}

	changeLine := &changelog.ChangeLine{
		Summary:   prTitle,
		Reference: fmt.Sprintf("#%d", number),
	}

	// Put either directly in the version history or in a subsection.
	if changeSectionLabel == "none" {
		changes.AddLineToVersion("HEAD", changeLine)
	} else {
		changes.AddLineToSubsection("HEAD", changeSectionLabel, changeLine)
	}

	return changes.String()
}

func deletableRef(pr *github.PullRequest, owner string) bool {
	return *pr.Head.Repo.Owner.Login == owner && *pr.Head.Ref != "master" && *pr.Head.Ref != "gh-pages"
}

func commitHistoryFile(client *github.Client, historySHA, owner, repo string, number int, newHistoryFileContents string) error {
	repositoryContentsOptions := &github.RepositoryContentFileOptions{
		Message: github.String(fmt.Sprintf("Update history to reflect merge of #%d [ci skip]", number)),
		Content: []byte(newHistoryFileContents),
		SHA:     github.String(historySHA),
		Committer: &github.CommitAuthor{
			Name:  github.String("jekyllbot"),
			Email: github.String("jekyllbot@jekyllrb.com"),
		},
	}
	updateResponse, _, err := client.Repositories.UpdateFile(owner, repo, "History.markdown", repositoryContentsOptions)
	if err != nil {
		fmt.Printf("comments: error committing History.markdown: %v\n", err)
		return err
	}
	fmt.Printf("comments: updateResponse: %s\n", updateResponse)
	return nil
}
<|endoftext|>


<commit_before>package commitlog

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"testing"

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

var (
	msgs = []Message{
		NewMessage([]byte("one")),
		NewMessage([]byte("two")),
		NewMessage([]byte("three")),
		NewMessage([]byte("four")),
	}
	msgSets = []MessageSet{
		NewMessageSet(0, msgs...),
		NewMessageSet(1, msgs...),
	}
	maxBytes = msgSets[0].Size()
	path     = filepath.Join(os.TempDir(), fmt.Sprintf("commitlogtest%d", rand.Int63()))
)

func TestNewCommitLog(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	maxBytes := msgSets[0].Size()
	r, err := l.NewReader(0, maxBytes)
	assert.NoError(t, err)

	for i := range msgSets {
		p := make([]byte, maxBytes)
		_, err = r.Read(p)
		assert.NoError(t, err)

		ms := MessageSet(p)
		assert.Equal(t, int64(i), ms.Offset())

		payload := ms.Payload()
		var offset int
		for _, msg := range msgs {
			assert.Equal(t, []byte(msg), payload[offset:offset+len(msg)])
			offset += len(msg)
		}
	}
}

func TestTruncateTo(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	assert.Equal(t, int64(2), l.NewestOffset())
	assert.Equal(t, 2, len(l.Segments()))

	err = l.TruncateTo(int64(1))
	assert.NoError(t, err)
	assert.Equal(t, 1, len(l.Segments()))

	maxBytes := msgSets[0].Size()
	_, err = l.NewReader(0, maxBytes)
	assert.Error(t, err)

	r, err := l.NewReader(1, maxBytes)
	assert.NoError(t, err)

	for i := range msgSets[1:] {
		p := make([]byte, maxBytes)
		_, err = r.Read(p)
		assert.NoError(t, err)

		ms := MessageSet(p)
		assert.Equal(t, int64(i+1), ms.Offset())

		payload := ms.Payload()
		var offset int
		for _, msg := range msgs {
			assert.Equal(t, []byte(msg), payload[offset:offset+len(msg)])
			offset += len(msg)
		}
	}
}

func TestCleaner(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	segments := l.Segments()
	assert.Equal(t, 2, len(segments))

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	assert.Equal(t, 2, len(l.Segments()))
	for i, s := range l.Segments() {
		assert.NotEqual(t, s, segments[i])
	}
}

func check(t *testing.T, got, want []byte) {
	if !bytes.Equal(got, want) {
		t.Errorf("got = %s, want %s", string(got), string(want))
	}
}

func setup(t *testing.T) *CommitLog {
	opts := Options{
		Path:            path,
		MaxSegmentBytes: 6,
		MaxLogBytes:     30,
	}
	l, err := New(opts)
	assert.NoError(t, err)

	assert.NoError(t, l.Init())
	assert.NoError(t, l.Open())

	return l
}

func cleanup(t *testing.T) {
	os.RemoveAll(path)
	os.MkdirAll(path, 0755)
}
<commit_msg>commitlog: add benchmark test<commit_after>package commitlog

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"testing"

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

var (
	msgs = []Message{
		NewMessage([]byte("one")),
		NewMessage([]byte("two")),
		NewMessage([]byte("three")),
		NewMessage([]byte("four")),
	}
	msgSets = []MessageSet{
		NewMessageSet(0, msgs...),
		NewMessageSet(1, msgs...),
	}
	maxBytes = msgSets[0].Size()
	path     = filepath.Join(os.TempDir(), fmt.Sprintf("commitlogtest%d", rand.Int63()))
)

func TestNewCommitLog(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	maxBytes := msgSets[0].Size()
	r, err := l.NewReader(0, maxBytes)
	assert.NoError(t, err)

	for i := range msgSets {
		p := make([]byte, maxBytes)
		_, err = r.Read(p)
		assert.NoError(t, err)

		ms := MessageSet(p)
		assert.Equal(t, int64(i), ms.Offset())

		payload := ms.Payload()
		var offset int
		for _, msg := range msgs {
			assert.Equal(t, []byte(msg), payload[offset:offset+len(msg)])
			offset += len(msg)
		}
	}
}

func BenchmarkCommitLog(b *testing.B) {
	var err error
	l := setup(b)
	defer cleanup(b)

	msgSet := msgSets[0]

	for i := 0; i < b.N; i++ {
		_, err = l.Append(msgSet)
		assert.NoError(b, err)
	}
}

func TestTruncateTo(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	assert.Equal(t, int64(2), l.NewestOffset())
	assert.Equal(t, 2, len(l.Segments()))

	err = l.TruncateTo(int64(1))
	assert.NoError(t, err)
	assert.Equal(t, 1, len(l.Segments()))

	maxBytes := msgSets[0].Size()
	_, err = l.NewReader(0, maxBytes)
	assert.Error(t, err)

	r, err := l.NewReader(1, maxBytes)
	assert.NoError(t, err)

	for i := range msgSets[1:] {
		p := make([]byte, maxBytes)
		_, err = r.Read(p)
		assert.NoError(t, err)

		ms := MessageSet(p)
		assert.Equal(t, int64(i+1), ms.Offset())

		payload := ms.Payload()
		var offset int
		for _, msg := range msgs {
			assert.Equal(t, []byte(msg), payload[offset:offset+len(msg)])
			offset += len(msg)
		}
	}
}

func TestCleaner(t *testing.T) {
	var err error
	l := setup(t)
	defer cleanup(t)

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	segments := l.Segments()
	assert.Equal(t, 2, len(segments))

	for _, msgSet := range msgSets {
		_, err = l.Append(msgSet)
		assert.NoError(t, err)
	}
	assert.Equal(t, 2, len(l.Segments()))
	for i, s := range l.Segments() {
		assert.NotEqual(t, s, segments[i])
	}
}

func check(t assert.TestingT, got, want []byte) {
	if !bytes.Equal(got, want) {
		t.Errorf("got = %s, want %s", string(got), string(want))
	}
}

func setup(t assert.TestingT) *CommitLog {
	opts := Options{
		Path:            path,
		MaxSegmentBytes: 6,
		MaxLogBytes:     30,
	}
	l, err := New(opts)
	assert.NoError(t, err)

	assert.NoError(t, l.Init())
	assert.NoError(t, l.Open())

	return l
}

func cleanup(t assert.TestingT) {
	os.RemoveAll(path)
	os.MkdirAll(path, 0755)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2019 by Farsight Security, Inc.
 *
 * 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 dnstap

import (
	"io"
	"net"

	"github.com/golang/protobuf/proto"
)

// An Encoder serializes and writes Dnstap messages to an underlying
// io.Writer.
type Encoder struct {
	w *Writer
}

// NewEncoder creates an Encoder using the given io.Writer and options.
func NewEncoder(w io.Writer, opt *WriterOptions) (*Encoder, error) {
	ew, err := NewWriter(w, opt)
	if err != nil {
		return nil, err
	}
	return &Encoder{w: ew}, nil
}

// Encode serializes and writes the Dnstap message m to the encoder's
// Writer.
func (e *Encoder) Encode(m *Dnstap) error {
	b, err := proto.Marshal(m)
	if err != nil {
		return err
	}

	_, err = e.w.Write(b)
	return err
}
<commit_msg>Add SocketEncoder to Encode to a SocketWriter<commit_after>/*
 * Copyright (c) 2019 by Farsight Security, Inc.
 *
 * 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 dnstap

import (
	"io"
	"net"

	"github.com/golang/protobuf/proto"
)

// An Encoder serializes and writes Dnstap messages to an underlying
// io.Writer.
type Encoder struct {
	w *Writer
}

// NewEncoder creates an Encoder using the given io.Writer and options.
func NewEncoder(w io.Writer, opt *WriterOptions) (*Encoder, error) {
	ew, err := NewWriter(w, opt)
	if err != nil {
		return nil, err
	}
	return &Encoder{w: ew}, nil
}

// Encode serializes and writes the Dnstap message m to the encoder's
// Writer.
func (e *Encoder) Encode(m *Dnstap) error {
	b, err := proto.Marshal(m)
	if err != nil {
		return err
	}

	_, err = e.w.Write(b)
	return err
}

type SocketEncoder struct {
	w *SocketWriter
}

func NewSocketEncoder(addr net.Addr, opt *SocketWriterOptions) *SocketEncoder {
	return &SocketEncoder{
		w: NewSocketWriter(addr, opt),
	}
}

func (se *SocketEncoder) Encode(m *Dnstap) error {
	b, err := proto.Marshal(m)
	if err != nil {
		return err
	}

	_, err = se.w.Write(b)
	return err
}
<|endoftext|>


<commit_before>package bolt

import (
	"context"

	"github.com/boltdb/bolt"
	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt/internal"
)

// Ensure SourcesStore implements chronograf.SourcesStore.
var _ chronograf.SourcesStore = &SourcesStore{}

var SourcesBucket = []byte("Sources")

type SourcesStore struct {
	client *Client
}

// All returns all known sources
func (s *SourcesStore) All(ctx context.Context) ([]chronograf.Source, error) {
	var srcs []chronograf.Source
	if err := s.client.db.View(func(tx *bolt.Tx) error {
		if err := tx.Bucket(SourcesBucket).ForEach(func(k, v []byte) error {
			var src chronograf.Source
			if err := internal.UnmarshalSource(v, &src); err != nil {
				return err
			}
			srcs = append(srcs, src)
			return nil
		}); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return srcs, nil

}

// Add creates a new Source in the SourceStore.
func (s *SourcesStore) Add(ctx context.Context, src chronograf.Source) (chronograf.Source, error) {

	// force first source added to be default
	if srcs, err := s.All(ctx); err != nil {
		return chronograf.Source{}, err
	} else if len(srcs) == 0 {
		src.Default = true
	}

	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(SourcesBucket)
		seq, err := b.NextSequence()
		if err != nil {
			return err
		}
		src.ID = int(seq)

		if src.Default {
			if err := s.resetDefaultSource(b, ctx); err != nil {
				return err
			}
		}

		if v, err := internal.MarshalSource(src); err != nil {
			return err
		} else if err := b.Put(itob(src.ID), v); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return chronograf.Source{}, err
	}

	return src, nil
}

// Delete removes the Source from the SourcesStore
func (s *SourcesStore) Delete(ctx context.Context, src chronograf.Source) error {

	if err := s.setRandomDefault(ctx, src); err != nil {
		return err
	}

	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		if err := tx.Bucket(SourcesBucket).Delete(itob(src.ID)); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}

	return nil
}

// Get returns a Source if the id exists.
func (s *SourcesStore) Get(ctx context.Context, id int) (chronograf.Source, error) {
	var src chronograf.Source
	if err := s.client.db.View(func(tx *bolt.Tx) error {
		if v := tx.Bucket(SourcesBucket).Get(itob(id)); v == nil {
			return chronograf.ErrSourceNotFound
		} else if err := internal.UnmarshalSource(v, &src); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return chronograf.Source{}, err
	}

	return src, nil
}

// Update a Source
func (s *SourcesStore) Update(ctx context.Context, src chronograf.Source) error {
	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		// Get an existing soource with the same ID.
		b := tx.Bucket(SourcesBucket)
		if v := b.Get(itob(src.ID)); v == nil {
			return chronograf.ErrSourceNotFound
		}

		if src.Default {
			if err := s.resetDefaultSource(b, ctx); err != nil {
				return err
			}
		}

		if v, err := internal.MarshalSource(src); err != nil {
			return err
		} else if err := b.Put(itob(src.ID), v); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}

	return nil
}

// resetDefaultSource unsets the Default flag on all sources
func (s *SourcesStore) resetDefaultSource(b *bolt.Bucket, ctx context.Context) error {
	srcs, err := s.All(ctx)
	if err != nil {
		return err
	}

	for _, other := range srcs {
		if other.Default {
			other.Default = false
			if v, err := internal.MarshalSource(other); err != nil {
				return err
			} else if err := b.Put(itob(other.ID), v); err != nil {
				return err
			}
		}
	}
	return nil
}

// setRandomDefault will locate a source other than the provided
// chronograf.Source and set it as the default source. If no other sources are
// available, the provided source will be set to the default source if is not
// already. It assumes that the provided chronograf.Source has been persisted.
func (s *SourcesStore) setRandomDefault(ctx context.Context, src chronograf.Source) error {
	// Check if requested source is the current default
	if target, err := s.Get(ctx, src.ID); err != nil {
		return err
	} else if target.Default {
		// Locate another source to be the new default
		if srcs, err := s.All(ctx); err != nil {
			return err
		} else {
			var other *chronograf.Source
			for idx, _ := range srcs {
				other = &srcs[idx]
				// avoid selecting the source we're about to delete as the new default
				if other.ID != target.ID {
					break
				}
			}

			// set the other to be the default
			other.Default = true
			if err := s.Update(ctx, *other); err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>Prevent deadlocks with source default enforcement<commit_after>package bolt

import (
	"context"

	"github.com/boltdb/bolt"
	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt/internal"
)

// Ensure SourcesStore implements chronograf.SourcesStore.
var _ chronograf.SourcesStore = &SourcesStore{}

var SourcesBucket = []byte("Sources")

type SourcesStore struct {
	client *Client
}

// All returns all known sources
func (s *SourcesStore) All(ctx context.Context) ([]chronograf.Source, error) {
	var srcs []chronograf.Source
	if err := s.client.db.View(func(tx *bolt.Tx) error {
		var err error
		srcs, err = s.all(ctx, tx)
		if err != nil {
			return err
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return srcs, nil

}

// Add creates a new Source in the SourceStore.
func (s *SourcesStore) Add(ctx context.Context, src chronograf.Source) (chronograf.Source, error) {

	// force first source added to be default
	if srcs, err := s.All(ctx); err != nil {
		return chronograf.Source{}, err
	} else if len(srcs) == 0 {
		src.Default = true
	}

	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		return s.add(ctx, &src, tx)
	}); err != nil {
		return chronograf.Source{}, err
	}

	return src, nil
}

// Delete removes the Source from the SourcesStore
func (s *SourcesStore) Delete(ctx context.Context, src chronograf.Source) error {
	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		if err := s.setRandomDefault(ctx, tx, src); err != nil {
			return err
		}
		return s.delete(ctx, src, tx)
	}); err != nil {
		return err
	}

	return nil
}

// Get returns a Source if the id exists.
func (s *SourcesStore) Get(ctx context.Context, id int) (chronograf.Source, error) {
	var src chronograf.Source
	if err := s.client.db.View(func(tx *bolt.Tx) error {
		var err error
		src, err = s.get(ctx, id, tx)
		if err != nil {
			return err
		}
		return nil
	}); err != nil {
		return chronograf.Source{}, err
	}

	return src, nil
}

// Update a Source
func (s *SourcesStore) Update(ctx context.Context, src chronograf.Source) error {
	if err := s.client.db.Update(func(tx *bolt.Tx) error {
		return s.update(ctx, src, tx)
	}); err != nil {
		return err
	}

	return nil
}

func (s *SourcesStore) all(ctx context.Context, tx *bolt.Tx) ([]chronograf.Source, error) {
	var srcs []chronograf.Source
	if err := tx.Bucket(SourcesBucket).ForEach(func(k, v []byte) error {
		var src chronograf.Source
		if err := internal.UnmarshalSource(v, &src); err != nil {
			return err
		}
		srcs = append(srcs, src)
		return nil
	}); err != nil {
		return srcs, err
	}
	return srcs, nil
}

func (s *SourcesStore) add(ctx context.Context, src *chronograf.Source, tx *bolt.Tx) error {
	b := tx.Bucket(SourcesBucket)
	seq, err := b.NextSequence()
	if err != nil {
		return err
	}
	src.ID = int(seq)

	if src.Default {
		if err := s.resetDefaultSource(tx, ctx); err != nil {
			return err
		}
	}

	if v, err := internal.MarshalSource(*src); err != nil {
		return err
	} else if err := b.Put(itob(src.ID), v); err != nil {
		return err
	}
	return nil
}

func (s *SourcesStore) delete(ctx context.Context, src chronograf.Source, tx *bolt.Tx) error {
	if err := tx.Bucket(SourcesBucket).Delete(itob(src.ID)); err != nil {
		return err
	}
	return nil
}

func (s *SourcesStore) get(ctx context.Context, id int, tx *bolt.Tx) (chronograf.Source, error) {
	var src chronograf.Source
	if v := tx.Bucket(SourcesBucket).Get(itob(id)); v == nil {
		return src, chronograf.ErrSourceNotFound
	} else if err := internal.UnmarshalSource(v, &src); err != nil {
		return src, err
	}
	return src, nil
}

func (s *SourcesStore) update(ctx context.Context, src chronograf.Source, tx *bolt.Tx) error {
	// Get an existing soource with the same ID.
	b := tx.Bucket(SourcesBucket)
	if v := b.Get(itob(src.ID)); v == nil {
		return chronograf.ErrSourceNotFound
	}

	if src.Default {
		if err := s.resetDefaultSource(tx, ctx); err != nil {
			return err
		}
	}

	if v, err := internal.MarshalSource(src); err != nil {
		return err
	} else if err := b.Put(itob(src.ID), v); err != nil {
		return err
	}
	return nil
}

// resetDefaultSource unsets the Default flag on all sources
func (s *SourcesStore) resetDefaultSource(tx *bolt.Tx, ctx context.Context) error {
	b := tx.Bucket(SourcesBucket)
	srcs, err := s.all(ctx, tx)
	if err != nil {
		return err
	}

	for _, other := range srcs {
		if other.Default {
			other.Default = false
			if v, err := internal.MarshalSource(other); err != nil {
				return err
			} else if err := b.Put(itob(other.ID), v); err != nil {
				return err
			}
		}
	}
	return nil
}

// setRandomDefault will locate a source other than the provided
// chronograf.Source and set it as the default source. If no other sources are
// available, the provided source will be set to the default source if is not
// already. It assumes that the provided chronograf.Source has been persisted.
func (s *SourcesStore) setRandomDefault(ctx context.Context, tx *bolt.Tx, src chronograf.Source) error {
	// Check if requested source is the current default
	if target, err := s.get(ctx, src.ID, tx); err != nil {
		return err
	} else if target.Default {
		// Locate another source to be the new default
		if srcs, err := s.all(ctx, tx); err != nil {
			return err
		} else {
			var other *chronograf.Source
			for idx, _ := range srcs {
				other = &srcs[idx]
				// avoid selecting the source we're about to delete as the new default
				if other.ID != target.ID {
					break
				}
			}

			// set the other to be the default
			other.Default = true
			if err := s.update(ctx, *other, tx); err != nil {
				return err
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package emitter

import (
	"errors"
	"strconv"
	"strings"

	"github.com/PuerkitoBio/agora/bytecode"
	"github.com/PuerkitoBio/agora/compiler/parser"
)

var (
	binSym2op = map[string]bytecode.Opcode{
		"+":  bytecode.OP_ADD,
		"-":  bytecode.OP_SUB,
		"*":  bytecode.OP_MUL,
		"/":  bytecode.OP_DIV,
		"%":  bytecode.OP_MOD,
		"<":  bytecode.OP_LT,
		"<=": bytecode.OP_LTE,
		">":  bytecode.OP_GT,
		">=": bytecode.OP_GTE,
		"==": bytecode.OP_EQ,
	}
)

type Emitter struct {
	err     error
	kMap    map[*bytecode.Fn]map[string]int
	stackSz map[*bytecode.Fn]int64
}

func (e *Emitter) Emit(id string, syms []*parser.Symbol, scps *parser.Scope) (*bytecode.File, error) {
	// Reset the internal fields
	e.err = nil
	e.kMap = make(map[*bytecode.Fn]map[string]int)
	e.stackSz = make(map[*bytecode.Fn]int64)

	// Create the bytecode representation structure
	f := bytecode.NewFile(id)
	fn := new(bytecode.Fn)
	fn.Header.Name = f.Name // Expected args is always 0 for top-level func
	// TODO : Line start and end, ExpVars
	f.Fns = append(f.Fns, fn)
	for _, sym := range syms {
		e.emitSymbol(f, fn, sym, false)
	}
	return f, e.err
}

func (e *Emitter) emitFn(f *bytecode.File, sym *parser.Symbol) {
	if e.err != nil {
		return
	}
	e.assert(sym.Ar == parser.ArFunction, errors.New("expected `"+sym.Id+"` to have function arity"))
	fn := new(bytecode.Fn)
	fn.Header.Name = sym.Name
	args := sym.First.([]*parser.Symbol)
	fn.Header.ExpArgs = int64(len(args))
	// TODO : ExpVars, Line Start, Line End
	f.Fns = append(f.Fns, fn)
	// Define the expected args in the K table - *MUST* be defined in spots 0..ExpArgs - 1
	for _, arg := range args {
		e.assert(arg.Ar == parser.ArName, errors.New("expected argument to have name arity"))
		e.registerK(fn, arg.Val, true)
	}
	stmts := sym.Second.([]*parser.Symbol)
	for _, stmt := range stmts {
		e.emitSymbol(f, fn, stmt, false)
	}
}

func (e *Emitter) emitSymbol(f *bytecode.File, fn *bytecode.Fn, sym *parser.Symbol, asg bool) {
	if e.err != nil {
		return
	}
	switch sym.Id {
	case "nil":
		e.assert(!asg, errors.New("invalid assignment to nil"))
		e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_N, 0)
	case "(name)":
		// TODO : For expected vars, the correct scope is required
		// Register the symbol
		kix := e.registerK(fn, sym.Val, true)
		if asg {
			e.addInstr(fn, bytecode.OP_POP, bytecode.FLG_V, kix)
		} else {
			e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_V, kix)
		}
	case "(literal)":
		// Register the symbol
		e.assert(!asg, errors.New("invalid assignment to a literal"))
		kix := e.registerK(fn, sym.Val, false)
		e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_K, kix)
	case ":=":
		e.assert(sym.Ar == parser.ArBinary, errors.New("expected `:=` to have binary arity"))
		e.emitSymbol(f, fn, sym.Second.(*parser.Symbol), false)
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), true)
	case "!":
		e.assert(sym.Ar == parser.ArUnary, errors.New("expected `!` to have unary arity"))
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.addInstr(fn, bytecode.OP_NOT, bytecode.FLG__, 0)
	case "-":
		if sym.Ar == parser.ArUnary {
			e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
			e.addInstr(fn, bytecode.OP_UNM, bytecode.FLG__, 0)
			break
		}
		fallthrough
	case "+", "*", "/", "%", "<", ">", "<=", ">=", "==":
		e.assert(sym.Ar == parser.ArBinary, errors.New("expected `"+sym.Id+"` to have binary arity"))
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.emitSymbol(f, fn, sym.Second.(*parser.Symbol), false)
		e.addInstr(fn, binSym2op[sym.Id], bytecode.FLG__, 0)
	case "func":
		if sym.Name != "" {
			// Function defined as a statement, register the name as a K,
			// and push the function's value into this variable.
			kix := e.registerK(fn, sym.Name, true)
			e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_F, uint64(len(f.Fns))) // New Fn will be added at this index
			e.addInstr(fn, bytecode.OP_POP, bytecode.FLG_V, kix)
		}
		e.emitFn(f, sym)
	case "(":
		if sym.Ar == parser.ArBinary {
			// Push parameters
			parms := sym.Second.([]*parser.Symbol)
			for _, parm := range parms {
				e.emitSymbol(f, fn, parm, false)
			}
			// Push function name
			e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
			// Call
			e.addInstr(fn, bytecode.OP_CALL, bytecode.FLG_nA, uint64(len(parms)))
		} else if sym.Ar == parser.ArTernary {

		} else {
			e.assert(false, errors.New("expected `(` to have binary or ternary arity"))
		}
	case "if":
		e.assert(sym.Ar == parser.ArStatement, errors.New("expected `if` to have statement arity"))
		// First is the condition, always a *Symbol
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		// Next comes the TEST, but we don't know yet how many instructions to jump
		// insert a placeholder (invalid op) so that it fails explicitly should it ever make it to
		// the VM.
		e.addInstr(fn, bytecode.OP_INVL, bytecode.FLG_INVL, 0)
		// Then comes the body
		// TODO : Emit block ([]*Symbol)? Reused by emitRoot and emitFn?
		// TODO : Keep count of instrs in map, or simply emit, then count - index of INVL?
		// Then comes the ELSE/ELSE IF, maybe
	case "return":
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.addInstr(fn, bytecode.OP_RET, bytecode.FLG__, 0)
	default:
		e.err = errors.New("unexpected symbol id: " + sym.Id)
	}
}

func (e *Emitter) addInstr(fn *bytecode.Fn, op bytecode.Opcode, flg bytecode.Flag, ix uint64) {
	if e.err != nil {
		return
	}
	switch op {
	case bytecode.OP_PUSH:
		e.stackSz[fn] += 1
	case bytecode.OP_POP, bytecode.OP_RET, bytecode.OP_UNM, bytecode.OP_NOT:
		e.stackSz[fn] -= 1
	case bytecode.OP_ADD, bytecode.OP_SUB, bytecode.OP_MUL, bytecode.OP_DIV, bytecode.OP_MOD:
		e.stackSz[fn] -= 2
	}
	if e.stackSz[fn] > fn.Header.StackSz {
		fn.Header.StackSz = e.stackSz[fn]
	}
	fn.Is = append(fn.Is, bytecode.NewInstr(op, flg, ix))
}

func (e *Emitter) registerK(fn *bytecode.Fn, val interface{}, isName bool) uint64 {
	var kt bytecode.KType
	s, ok := val.(string)
	if ok {
		if isName {
			val = s
			kt = bytecode.KtString
		} else if s[0] == '"' || s[0] == '`' {
			// Strip the quotes
			s = s[1 : len(s)-1]
			val = s
			kt = bytecode.KtString
		} else if strings.Index(s, ".") >= 0 {
			val, e.err = strconv.ParseFloat(s, 64)
			kt = bytecode.KtFloat
		} else {
			val, e.err = strconv.ParseInt(s, 10, 64)
			kt = bytecode.KtInteger
		}
	} else {
		kt = bytecode.KtBoolean
		if v := val.(bool); v {
			s = "true"
			val = 1
		} else {
			s = "false"
			val = 0
		}
	}
	m, ok := e.kMap[fn]
	if !ok {
		m = make(map[string]int)
		e.kMap[fn] = m
	}
	i, ok := m[s]
	if !ok {
		i = len(m)
		m[s] = i
		fn.Ks = append(fn.Ks, &bytecode.K{Type: kt, Val: val})
	}
	return uint64(i)
}

func (e *Emitter) assert(cond bool, err error) {
	if !cond {
		e.err = err
	}
}
<commit_msg>emit new opcodes, run test file up to 05<commit_after>package emitter

import (
	"errors"
	"strconv"
	"strings"

	"github.com/PuerkitoBio/agora/bytecode"
	"github.com/PuerkitoBio/agora/compiler/parser"
)

var (
	binSym2op = map[string]bytecode.Opcode{
		"+":  bytecode.OP_ADD,
		"-":  bytecode.OP_SUB,
		"*":  bytecode.OP_MUL,
		"/":  bytecode.OP_DIV,
		"%":  bytecode.OP_MOD,
		"<":  bytecode.OP_LT,
		"<=": bytecode.OP_LTE,
		">":  bytecode.OP_GT,
		">=": bytecode.OP_GTE,
		"==": bytecode.OP_EQ,
	}
)

type Emitter struct {
	err     error
	kMap    map[*bytecode.Fn]map[string]int
	stackSz map[*bytecode.Fn]int64
}

func (e *Emitter) Emit(id string, syms []*parser.Symbol, scps *parser.Scope) (*bytecode.File, error) {
	// Reset the internal fields
	e.err = nil
	e.kMap = make(map[*bytecode.Fn]map[string]int)
	e.stackSz = make(map[*bytecode.Fn]int64)

	// Create the bytecode representation structure
	f := bytecode.NewFile(id)
	fn := new(bytecode.Fn)
	fn.Header.Name = f.Name // Expected args is always 0 for top-level func
	// TODO : Line start and end, ExpVars
	f.Fns = append(f.Fns, fn)
	e.emitBlock(f, fn, syms)
	return f, e.err
}

func (e *Emitter) emitFn(f *bytecode.File, sym *parser.Symbol) {
	if e.err != nil {
		return
	}
	e.assert(sym.Ar == parser.ArFunction, errors.New("expected `"+sym.Id+"` to have function arity"))
	fn := new(bytecode.Fn)
	fn.Header.Name = sym.Name
	args := sym.First.([]*parser.Symbol)
	fn.Header.ExpArgs = int64(len(args))
	// TODO : ExpVars, Line Start, Line End
	f.Fns = append(f.Fns, fn)
	// Define the expected args in the K table - *MUST* be defined in spots 0..ExpArgs - 1
	for _, arg := range args {
		e.assert(arg.Ar == parser.ArName, errors.New("expected argument to have name arity"))
		e.registerK(fn, arg.Val, true)
	}
	stmts := sym.Second.([]*parser.Symbol)
	e.emitBlock(f, fn, stmts)
}

func (e *Emitter) emitBlock(f *bytecode.File, fn *bytecode.Fn, syms []*parser.Symbol) {
	for _, sym := range syms {
		e.emitSymbol(f, fn, sym, false)
	}
}

func (e *Emitter) emitSymbol(f *bytecode.File, fn *bytecode.Fn, sym *parser.Symbol, asg bool) {
	if e.err != nil {
		return
	}
	switch sym.Id {
	case "nil":
		e.assert(!asg, errors.New("invalid assignment to nil"))
		e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_N, 0)
	case "(name)":
		// TODO : For expected vars, the correct scope is required
		// Register the symbol
		e.assert(sym.Ar == parser.ArName || sym.Ar == parser.ArLiteral, errors.New("expected `(name)` to have name or literal arity"))
		kix := e.registerK(fn, sym.Val, true)
		if asg {
			e.addInstr(fn, bytecode.OP_POP, bytecode.FLG_V, kix)
		} else if sym.Ar == parser.ArLiteral {
			e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_K, kix)
		} else {
			e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_V, kix)
		}
	case "(literal)":
		// Register the symbol
		e.assert(!asg, errors.New("invalid assignment to a literal"))
		kix := e.registerK(fn, sym.Val, false)
		e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_K, kix)
	case ":=":
		e.assert(sym.Ar == parser.ArBinary, errors.New("expected `:=` to have binary arity"))
		e.emitSymbol(f, fn, sym.Second.(*parser.Symbol), false)
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), true)
	case "!":
		e.assert(sym.Ar == parser.ArUnary, errors.New("expected `!` to have unary arity"))
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.addInstr(fn, bytecode.OP_NOT, bytecode.FLG__, 0)
	case "-":
		if sym.Ar == parser.ArUnary {
			e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
			e.addInstr(fn, bytecode.OP_UNM, bytecode.FLG__, 0)
			break
		}
		fallthrough
	case "+", "*", "/", "%", "<", ">", "<=", ">=", "==":
		e.assert(sym.Ar == parser.ArBinary, errors.New("expected `"+sym.Id+"` to have binary arity"))
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.emitSymbol(f, fn, sym.Second.(*parser.Symbol), false)
		e.addInstr(fn, binSym2op[sym.Id], bytecode.FLG__, 0)
	case "func":
		if sym.Name != "" {
			// Function defined as a statement, register the name as a K,
			// and push the function's value into this variable.
			kix := e.registerK(fn, sym.Name, true)
			e.addInstr(fn, bytecode.OP_PUSH, bytecode.FLG_F, uint64(len(f.Fns))) // New Fn will be added at this index
			e.addInstr(fn, bytecode.OP_POP, bytecode.FLG_V, kix)
		}
		e.emitFn(f, sym)
	case "(":
		e.assert(sym.Ar == parser.ArBinary || sym.Ar == parser.ArTernary, errors.New("expected `(` to have binary or ternary arity"))
		// Push parameters
		var parms []*parser.Symbol
		var op bytecode.Opcode
		if sym.Ar == parser.ArBinary {
			parms = sym.Second.([]*parser.Symbol)
			op = bytecode.OP_CALL
		} else {
			parms = sym.Third.([]*parser.Symbol)
			op = bytecode.OP_CFLD
		}
		for _, parm := range parms {
			e.emitSymbol(f, fn, parm, false)
		}
		// If ternary, push field (Second)
		if sym.Ar == parser.ArTernary {
			e.emitSymbol(f, fn, sym.Second.(*parser.Symbol), false)
		}
		// Push function name (or parent object of the field if ternary)
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		// Call
		e.addInstr(fn, op, bytecode.FLG_nA, uint64(len(parms)))
	case "if":
		e.assert(sym.Ar == parser.ArStatement, errors.New("expected `if` to have statement arity"))
		// First is the condition, always a *Symbol
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		// Next comes the TEST, but we don't know yet how many instructions to jump
		// insert a placeholder (invalid op) so that it fails explicitly should it ever make it to
		// the VM.
		e.addInstr(fn, bytecode.OP_INVL, bytecode.FLG_INVL, 0)
		tstIx := len(fn.Is) - 1
		// Then comes the body
		e.emitBlock(f, fn, sym.Second.([]*parser.Symbol))
		// Update the test instruction, now that we know where to jump to
		fn.Is[tstIx] = bytecode.NewInstr(bytecode.OP_TEST, bytecode.FLG_J, uint64(len(fn.Is)-tstIx-1))
		// Then comes the ELSE/ELSE IF, maybe
		if sym.Third != nil {
			switch t := sym.Third.(type) {
			case *parser.Symbol:
				e.emitSymbol(f, fn, t, false)
			case []*parser.Symbol:
				e.emitBlock(f, fn, t)
			default:
				e.assert(false, errors.New("expected third branch of `if` to be a symbol or a slice of symbols"))
			}
		}
	case "return":
		e.emitSymbol(f, fn, sym.First.(*parser.Symbol), false)
		e.addInstr(fn, bytecode.OP_RET, bytecode.FLG__, 0)
	case "import":
		// TODO : Import should be changed to a built-in function:
		//   i.e. `fmt := import("fmt")`
		// For now, it's First is a slice of pairs of names-literals
		var ixV, ixI uint64
		for i, s := range sym.First.([]*parser.Symbol) {
			if i%2 == 0 {
				ixV = e.registerK(fn, s.Val, true)
			} else {
				ixI = e.registerK(fn, s.Val, false)
				e.addInstr(fn, bytecode.OP_LOAD, bytecode.FLG_K, ixI)
				e.addInstr(fn, bytecode.OP_POP, bytecode.FLG_V, ixV)
			}
		}
	default:
		e.err = errors.New("unexpected symbol id: " + sym.Id)
	}
}

func (e *Emitter) addInstr(fn *bytecode.Fn, op bytecode.Opcode, flg bytecode.Flag, ix uint64) {
	if e.err != nil {
		return
	}
	switch op {
	case bytecode.OP_PUSH:
		e.stackSz[fn] += 1
	case bytecode.OP_POP, bytecode.OP_RET, bytecode.OP_UNM, bytecode.OP_NOT, bytecode.OP_TEST,
		bytecode.OP_LT, bytecode.OP_LTE, bytecode.OP_GT, bytecode.OP_GTE, bytecode.OP_EQ,
		bytecode.OP_AND, bytecode.OP_OR:
		e.stackSz[fn] -= 1
	case bytecode.OP_ADD, bytecode.OP_SUB, bytecode.OP_MUL, bytecode.OP_DIV, bytecode.OP_MOD:
		e.stackSz[fn] -= 2
	case bytecode.OP_CALL:
		e.stackSz[fn] -= (int64(ix) + 1)
	}
	if e.stackSz[fn] > fn.Header.StackSz {
		fn.Header.StackSz = e.stackSz[fn]
	}
	fn.Is = append(fn.Is, bytecode.NewInstr(op, flg, ix))
}

func (e *Emitter) registerK(fn *bytecode.Fn, val interface{}, isName bool) uint64 {
	var kt bytecode.KType
	s, ok := val.(string)
	if ok {
		if isName {
			val = s
			kt = bytecode.KtString
		} else if s[0] == '"' || s[0] == '`' {
			// Strip the quotes
			s = s[1 : len(s)-1]
			val = s
			kt = bytecode.KtString
		} else if strings.Index(s, ".") >= 0 {
			val, e.err = strconv.ParseFloat(s, 64)
			kt = bytecode.KtFloat
		} else {
			val, e.err = strconv.ParseInt(s, 10, 64)
			kt = bytecode.KtInteger
		}
	} else {
		kt = bytecode.KtBoolean
		if v := val.(bool); v {
			s = "true"
			val = 1
		} else {
			s = "false"
			val = 0
		}
	}
	m, ok := e.kMap[fn]
	if !ok {
		m = make(map[string]int)
		e.kMap[fn] = m
	}
	i, ok := m[s]
	if !ok {
		i = len(m)
		m[s] = i
		fn.Ks = append(fn.Ks, &bytecode.K{Type: kt, Val: val})
	}
	return uint64(i)
}

func (e *Emitter) assert(cond bool, err error) {
	if !cond {
		e.err = err
	}
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
)

const POINTING_PAIR_COL_GRID = `3|.|6|.|.|.|.|.|.
.|.|.|.|7|.|.|.|.
4|.|5|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.`

func TestSolveOnlyLegalNumber(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)
	cell := grid.Cell(3, 3)
	num := cell.Number()
	cell.SetNumber(0)

	//Now that cell should be filled by this technique.

	solver := &nakedSingleTechnique{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The only legal number technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The only legal number technique identified the wrong cell.")
		t.Fail()
	}
	numFromStep := step.Nums[0]

	if numFromStep != num {
		t.Log("The only legal number technique identified the wrong number.")
		t.Fail()
	}
	if grid.Solved() {
		t.Log("The only legal number technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInRow(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Row(3) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInRow{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in row technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in row technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in row technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in row technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInCol(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Col(3) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInCol{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in col technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in col technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in col technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in col technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInBlock(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Block(4) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInBlock{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in block technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in block technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in block technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in block technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestPointingPairCol(t *testing.T) {
	grid := NewGrid()
	grid.Load(POINTING_PAIR_COL_GRID)
	solver := &pointingPairCol{}
	step := solver.Find(grid)
	if step == nil {
		t.Log("The pointing pair col didn't find a cell it should have")
		t.Fail()
	}
	if len(step.TargetCells) != BLOCK_DIM*2 {
		t.Log("The pointing pair col gave back the wrong number of target cells")
		t.Fail()
	}
	if len(step.PointerCells) != BLOCK_DIM-1 {
		t.Log("The pointing pair col gave back the wrong number of pointer cells")
		t.Fail()
	}
	if !step.TargetCells.SameCol() || step.TargetCells.Col() != 1 {
		t.Log("The target cells in the pointing pair col technique were wrong col")
		t.Fail()
	}
	if len(step.Nums) != 1 || step.Nums[0] != 7 {
		t.Log("Pointing pair col technique gave the wrong number")
		t.Fail()
	}
}

func TestHumanSolve(t *testing.T) {
	grid := NewGrid()
	grid.Load(TEST_GRID)
	steps := grid.HumanSolve()
	//TODO: test to make sure that we use a wealth of different techniques. This will require a cooked random for testing.
	if steps == nil {
		t.Log("Human solve returned 0 techniques")
		t.Fail()
	}
	if !grid.Solved() {
		t.Log("Human solve failed to solve the simple grid.")
		t.Fail()
	}
}
<commit_msg>Added another test to make sure we applied the result of pointing pair col correctly.<commit_after>package sudoku

import (
	"testing"
)

const POINTING_PAIR_COL_GRID = `3|.|6|.|.|.|.|.|.
.|.|.|.|7|.|.|.|.
4|.|5|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.
.|.|.|.|.|.|.|.|.`

func TestSolveOnlyLegalNumber(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)
	cell := grid.Cell(3, 3)
	num := cell.Number()
	cell.SetNumber(0)

	//Now that cell should be filled by this technique.

	solver := &nakedSingleTechnique{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The only legal number technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The only legal number technique identified the wrong cell.")
		t.Fail()
	}
	numFromStep := step.Nums[0]

	if numFromStep != num {
		t.Log("The only legal number technique identified the wrong number.")
		t.Fail()
	}
	if grid.Solved() {
		t.Log("The only legal number technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInRow(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Row(3) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInRow{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in row technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in row technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in row technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in row technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInCol(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Col(3) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInCol{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in col technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in col technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in col technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in col technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestNecessaryInBlock(t *testing.T) {
	grid := NewGrid()
	//Load up a solved grid
	grid.Load(SOLVED_TEST_GRID)

	//We're going to cheat an set up an unrealistic grid.

	impossibles := make([]int, DIM)

	for i := 0; i < DIM-1; i++ {
		impossibles[i] = 0
	}
	impossibles[DIM-1] = 1

	//SetNumber will affect the other cells in row, so do it first.
	for _, cell := range grid.Block(4) {
		cell.number = 0
		copy(cell.impossibles[:], impossibles)
	}

	cell := grid.Cell(3, 3)
	//This is the only cell where DIM will be allowed.
	cell.impossibles[DIM-1] = 0

	//Now that cell should be filled by this technique.

	solver := &hiddenSingleInBlock{}

	step := solver.Find(grid)

	if step == nil {
		t.Log("The necessary in block technique did not solve a puzzle it should have.")
		t.FailNow()
	}

	cellFromStep := step.TargetCells[0]

	if cellFromStep.Col != 3 || cellFromStep.Row != 3 {
		t.Log("The necessary in block technique identified the wrong cell.")
		t.Fail()
	}

	numFromStep := step.Nums[0]

	if numFromStep != DIM {
		t.Log("The necessary in block technique identified the wrong number.")
		t.Fail()
	}
	//Can't check if grid is solved because we un-set all the other cells in the row.
	if cell.Number() != 0 {
		t.Log("The necessary in block technique did actually mutate the grid.")
		t.Fail()
	}
}

func TestPointingPairCol(t *testing.T) {
	grid := NewGrid()
	grid.Load(POINTING_PAIR_COL_GRID)
	solver := &pointingPairCol{}
	step := solver.Find(grid)
	if step == nil {
		t.Log("The pointing pair col didn't find a cell it should have")
		t.Fail()
	}
	if len(step.TargetCells) != BLOCK_DIM*2 {
		t.Log("The pointing pair col gave back the wrong number of target cells")
		t.Fail()
	}
	if len(step.PointerCells) != BLOCK_DIM-1 {
		t.Log("The pointing pair col gave back the wrong number of pointer cells")
		t.Fail()
	}
	if !step.TargetCells.SameCol() || step.TargetCells.Col() != 1 {
		t.Log("The target cells in the pointing pair col technique were wrong col")
		t.Fail()
	}
	if len(step.Nums) != 1 || step.Nums[0] != 7 {
		t.Log("Pointing pair col technique gave the wrong number")
		t.Fail()
	}
	step.Apply(grid)
	num := step.Nums[0]
	for _, cell := range step.TargetCells {
		if cell.Possible(num) {
			t.Log("The pointing pairs col technique was not applied correclty")
			t.Fail()
		}
	}
}

func TestHumanSolve(t *testing.T) {
	grid := NewGrid()
	grid.Load(TEST_GRID)
	steps := grid.HumanSolve()
	//TODO: test to make sure that we use a wealth of different techniques. This will require a cooked random for testing.
	if steps == nil {
		t.Log("Human solve returned 0 techniques")
		t.Fail()
	}
	if !grid.Solved() {
		t.Log("Human solve failed to solve the simple grid.")
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package command

import (
	"github.com/codegangsta/cli"
	"github.com/farmer-project/farmer-cli/api"
	"github.com/farmer-project/farmer-cli/hub"
	"github.com/farmer-project/farmer/api/request"
)

func DeployCmd() cli.Command {
	return cli.Command{
		Name:        "deploy",
		Usage:       "<boxname> --pathspec=BRANCH",
		Description: "Updates a box's code from provided Git branch specifier. Note code will be pulled from repository Url you've provided when creating the box.",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "pathspec, p",
				Value: "master",
				Usage: "Branch specifier used as the Git path when cloning the code, e.g. master, tags/v2.3",
			},
		},
		Action: deployAction,
	}
}

func deployAction(context *cli.Context) {
	if !context.Args().Present() {
		println("You must specify a 'name' for the box you want to create.\nSee 'farmer create --help' for more info.")
		return
	}

	if context.String("pathspec") == "" {
		println("You must specify a 'pathspec' (Git branch specifier) to pull the code from.\nSee 'farmer create --help' for more info.")
		return
	}

	stream := hub.Stream{}
	request := request.DeployRequest{
		Pathspec: context.String("pathspec"),
	}

	if err := api.Put("/boxes/"+context.Args().First(), &request, &stream); err != nil {
		println(err.Error())
		return
	}

	if err := stream.Consume(); err != nil {
		println("Could not consume the stream from Farmer server.")
		return
	}
}
<commit_msg>feature(deploy): Optional pathspec on deploy #10<commit_after>package command

import (
	"github.com/codegangsta/cli"
	"github.com/farmer-project/farmer-cli/api"
	"github.com/farmer-project/farmer-cli/hub"
	"github.com/farmer-project/farmer/api/request"
)

func DeployCmd() cli.Command {
	return cli.Command{
		Name:        "deploy",
		Usage:       "<boxname> --pathspec=BRANCH",
		Description: "Updates a box's code from provided Git branch specifier. Note code will be pulled from repository Url you've provided when creating the box.",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "pathspec, p",
				Usage: "Branch specifier used as the Git path when cloning the code, e.g. master, tags/v2.3",
			},
		},
		Action: deployAction,
	}
}

func deployAction(context *cli.Context) {
	if !context.Args().Present() {
		println("You must specify a 'name' for the box you want to create.\nSee 'farmer create --help' for more info.")
		return
	}

	stream := hub.Stream{}
	request := request.DeployRequest{
		Pathspec: context.String("pathspec"),
	}

	if err := api.Put("/boxes/"+context.Args().First(), &request, &stream); err != nil {
		println(err.Error())
		return
	}

	if err := stream.Consume(); err != nil {
		println("Could not consume the stream from Farmer server.")
		return
	}
}
<|endoftext|>


<commit_before>package command

import (
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"raindfs/raftlayer"
	"raindfs/server"
	"raindfs/util"

	"github.com/gorilla/mux"
	"github.com/neverlee/glog"
	"github.com/soheilhy/cmux"
)

var cmdMaster = &Command{
	UsageLine: "master -port=10000",
	Short:     "start a master server",
	Long: `start a master server to provide volume=>location mapping service
  and sequence number of file ids

  `,
}

type masterServerOption struct {
	addr       *string
	bindAll    *bool
	metaDir    *string
	clusters   *string
	pulse      *int
	timeout    *int
	maxCpu     *int
	cpuProfile *string
}

var msopt masterServerOption

func init() {
	cmdMaster.Run = runMaster // break init cycle
	msopt = masterServerOption{
		addr:       cmdMaster.Flag.String("addr", "0.0.0.0:10000", "address to bind to"),
		bindAll:    cmdMaster.Flag.Bool("bindall", true, "bind address to 0.0.0.0 default except one ip"),
		metaDir:    cmdMaster.Flag.String("mdir", "./meta", "data directory to store meta data"),
		clusters:   cmdMaster.Flag.String("clusters", "", "master nodes in comma separated ip:port list, example: 127.0.0.1:9093,127.0.0.1:9094"),
		pulse:      cmdMaster.Flag.Int("pulseseconds", 5, "number of seconds between heartbeats"),
		timeout:    cmdMaster.Flag.Int("idletimeout", 10, "connection idle seconds"),
		maxCpu:     cmdMaster.Flag.Int("maxcpu", 0, "maximum number of CPUs. 0 means all available CPUs"),
		cpuProfile: cmdMaster.Flag.String("cpuprofile", "", "cpu profile output file"),
	}
}

func runMaster(cmd *Command, args []string) bool {
	if *msopt.maxCpu > 0 {
		runtime.GOMAXPROCS(*msopt.maxCpu)
	}
	if *msopt.cpuProfile != "" {
		f, err := os.Create(*msopt.cpuProfile)
		if err != nil {
			glog.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
		OnInterrupt(func() {
			pprof.StopCPUProfile()
		})
	}

	addr := *msopt.addr
	bindall := *msopt.bindAll
	clusters := util.StrSplit(*msopt.clusters, ",")
	metaDir := *msopt.metaDir
	pulse := *msopt.pulse
	timeout := time.Duration(*msopt.timeout) * time.Second
	if err := util.MkdirOrExist(metaDir); err != nil {
		glog.Fatalf("Check Meta Folder (-mdir) Writable %s : %s", metaDir, err)
	}

	baddr := addr
	if bindall {
		ipport := strings.Split(addr, ":")
		baddr = "0.0.0.0:" + ipport[1]
	}

	listener, err := net.Listen("tcp", baddr)
	if err != nil {
		glog.Fatalf("Master startup error: %v", err)
		return false
	}

	lmux := cmux.New(listener)
	HTTPListener := &util.Listener{
		Listener:     lmux.Match(cmux.HTTP1Fast()),
		ReadTimeout:  timeout,
		WriteTimeout: timeout,
	}
	RaftListener := lmux.Match(cmux.Any())

	raftserver := raftlayer.NewRaftServer(RaftListener, addr, clusters, metaDir)
	if raftserver == nil {
		glog.Fatalf("Fail to create raft serve")
	}
	go lmux.Serve()

	router := mux.NewRouter()
	ms := server.NewMasterServer(raftserver, pulse)
	if ms == nil {
		glog.Fatalf("Fail to serve")
	}

	ms.SetMasterServer(router)

	glog.V(0).Infoln("Start Seaweed Master", util.VERSION, "at", addr)

	if e := http.Serve(HTTPListener, router); e != nil {
		glog.Fatalf("Fail to serve: %v", e)
	}
	return true
}
<commit_msg>fix master close<commit_after>package command

import (
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"raindfs/raftlayer"
	"raindfs/server"
	"raindfs/util"

	"github.com/gorilla/mux"
	"github.com/neverlee/glog"
	"github.com/soheilhy/cmux"
)

var cmdMaster = &Command{
	UsageLine: "master -port=10000",
	Short:     "start a master server",
	Long: `start a master server to provide volume=>location mapping service
  and sequence number of file ids

  `,
}

type masterServerOption struct {
	addr       *string
	bindAll    *bool
	metaDir    *string
	clusters   *string
	pulse      *int
	timeout    *int
	maxCpu     *int
	cpuProfile *string
}

var msopt masterServerOption

func init() {
	cmdMaster.Run = runMaster // break init cycle
	msopt = masterServerOption{
		addr:       cmdMaster.Flag.String("addr", "0.0.0.0:10000", "address to bind to"),
		bindAll:    cmdMaster.Flag.Bool("bindall", true, "bind address to 0.0.0.0 default except one ip"),
		metaDir:    cmdMaster.Flag.String("mdir", "./meta", "data directory to store meta data"),
		clusters:   cmdMaster.Flag.String("clusters", "", "master nodes in comma separated ip:port list, example: 127.0.0.1:9093,127.0.0.1:9094"),
		pulse:      cmdMaster.Flag.Int("pulseseconds", 5, "number of seconds between heartbeats"),
		timeout:    cmdMaster.Flag.Int("idletimeout", 10, "connection idle seconds"),
		maxCpu:     cmdMaster.Flag.Int("maxcpu", 0, "maximum number of CPUs. 0 means all available CPUs"),
		cpuProfile: cmdMaster.Flag.String("cpuprofile", "", "cpu profile output file"),
	}
}

func runMaster(cmd *Command, args []string) bool {
	if *msopt.maxCpu > 0 {
		runtime.GOMAXPROCS(*msopt.maxCpu)
	}
	if *msopt.cpuProfile != "" {
		f, err := os.Create(*msopt.cpuProfile)
		if err != nil {
			glog.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
		OnInterrupt(func() {
			pprof.StopCPUProfile()
		})
	}

	addr := *msopt.addr
	bindall := *msopt.bindAll
	clusters := util.StrSplit(*msopt.clusters, ",")
	metaDir := *msopt.metaDir
	pulse := *msopt.pulse
	timeout := time.Duration(*msopt.timeout) * time.Second
	if err := util.MkdirOrExist(metaDir); err != nil {
		glog.Fatalf("Check Meta Folder (-mdir) Writable %s : %s", metaDir, err)
	}

	baddr := addr
	if bindall {
		ipport := strings.Split(addr, ":")
		baddr = "0.0.0.0:" + ipport[1]
	}

	listener, err := net.Listen("tcp", baddr)
	if err != nil {
		glog.Fatalf("Master startup error: %v", err)
		return false
	}

	lmux := cmux.New(listener)
	HTTPListener := &util.Listener{
		Listener:     lmux.Match(cmux.HTTP1Fast()),
		ReadTimeout:  timeout,
		WriteTimeout: timeout,
	}
	RaftListener := lmux.Match(cmux.Any())

	raftserver := raftlayer.NewRaftServer(RaftListener, addr, clusters, metaDir)
	if raftserver == nil {
		glog.Fatalf("Fail to create raft serve")
	}
	go lmux.Serve()

	router := mux.NewRouter()
	ms := server.NewMasterServer(raftserver, pulse)
	if ms == nil {
		glog.Fatalf("Fail to serve")
	}

	ms.SetMasterServer(router)

	OnInterrupt(func() {
		ms.Close()
	})

	glog.V(0).Infoln("Start Seaweed Master", util.VERSION, "at", addr)

	if e := http.Serve(HTTPListener, router); e != nil {
		glog.Fatalf("Fail to serve: %v", e)
	}
	return true
}
<|endoftext|>