<commit_before>/*
 Copyright 2016 Padduck, LLC

 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 sftp

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	configuration "github.com/pufferpanel/pufferd/config"
	"github.com/pufferpanel/pufferd/logging"
	"github.com/pufferpanel/pufferd/programs"
	"github.com/taruti/sftpd"
	"golang.org/x/crypto/ssh"
)

func Run() {
	e := runServer()
	if e != nil {
		logging.Error("Error starting SFTP", e)
	}
}

func runServer() error {
	config := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			return validateSSH(c.User(), string(pass))
		},
	}

	serverKeyFile := path.Join(configuration.GetOrDefault("datafolder", "data"), "server.key")

	_, e := os.Stat(serverKeyFile)

	if e != nil && os.IsNotExist(e) {
		logging.Debug("Generating new key")
		var key *rsa.PrivateKey
		key, e = rsa.GenerateKey(rand.Reader, 2048)
		if e != nil {
			return e
		}

		data := x509.MarshalPKCS1PrivateKey(key)
		block := pem.Block{
			Type:    "RSA PRIVATE KEY",
			Headers: nil,
			Bytes:   data,
		}
		ioutil.WriteFile(serverKeyFile, pem.EncodeToMemory(&block), 0700)
		if e != nil {
			return e
		}
	} else if e != nil {
		return e
	}

	logging.Debug("Loading existing key")
	var data []byte
	data, e = ioutil.ReadFile(serverKeyFile)
	if e != nil {
		return e
	}

	hkey, e := ssh.ParsePrivateKey(data)

	if e != nil {
		logging.Debug("trigger")
		return e
	}

	config.AddHostKey(hkey)

	bind := configuration.GetOrDefault("sftp", "0.0.0.0:5657")

	listener, e := net.Listen("tcp", bind)
	if e != nil {
		return e
	}
	logging.Infof("Started SFTP Server on %s", bind)

	go func() {
		for {
			conn, _ := listener.Accept()
			go HandleConn(conn, config)
		}
	}()

	return nil
}

func HandleConn(conn net.Conn, config *ssh.ServerConfig) {
	defer conn.Close()
	e := handleConn(conn, config)
	if e != nil {
		logging.Error("sftpd connection errored:", e)
	}
}
func handleConn(conn net.Conn, config *ssh.ServerConfig) error {
	sc, chans, reqs, e := ssh.NewServerConn(conn, config)
	if e != nil {
		return e
	}
	defer sc.Close()

	// The incoming Request channel must be serviced.
	go PrintDiscardRequests(reqs)

	// Service the incoming Channel channel.
	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			return err
		}

		fs := &VirtualFS{Prefix: path.Join(programs.ServerFolder, sc.Permissions.Extensions["server_id"])}

		go func(in <-chan *ssh.Request) {
			for req := range in {
				ok := false
				switch {
				case sftpd.IsSftpRequest(req):
					ok = true
					go func() {
						sftpd.ServeChannel(channel, fs)
					}()
				}
				req.Reply(ok, nil)
			}
		}(requests)

	}
	return nil
}

func PrintDiscardRequests(in <-chan *ssh.Request) {
	for req := range in {
		if req.WantReply {
			req.Reply(false, nil)
		}
	}
}

func validateSSH(username string, password string) (*ssh.Permissions, error) {
	authUrl := configuration.Get("authserver")
	client := &http.Client{}
	data := url.Values{}
	data.Set("grant_type", "password")
	data.Set("username", username)
	data.Set("password", password)
	data.Set("scope", "sftp")
	token := configuration.Get("authtoken")
	request, _ := http.NewRequest("POST", authUrl, bytes.NewBufferString(data.Encode()))
	request.Header.Add("Authorization", "Bearer "+token)
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Add("Content-Length", strconv.Itoa(len(data.Encode())))
	response, err := client.Do(request)
	if err != nil {
		logging.Error("Error talking to auth server", err)
		return nil, errors.New("Invalid response from authorization server")
	}

	//we should only get a 200 or 400 back, if we get any others, we have a problem
	if response.StatusCode != 200 && response.StatusCode != 400 {
		logging.Error("Error talking to auth server", response.StatusCode)
		return nil, errors.New("Invalid response from authorization server")
	}
	var respArr map[string]interface{}
	json.NewDecoder(response.Body).Decode(&respArr)
	if respArr["error"] != nil {
		return nil, errors.New("Incorrect username or password")
	}
	sshPerms := &ssh.Permissions{}
	scopes := strings.Split(respArr["scope"].(string), " ")
	if len(scopes) != 2 {
		return nil, errors.New("Invalid response from authorization server")
	}
	for _, v := range scopes {
		if v != "sftp" {
			sshPerms.Extensions = make(map[string]string)
			sshPerms.Extensions["server_id"] = v
			return sshPerms, nil
		}
	}
	return nil, errors.New("Incorrect username or password")
}<commit_msg>Prepare for usage of new sftp library WARNING - This exposes the entire file system currently and does not handle large files<commit_after>/*
 Copyright 2016 Padduck, LLC

 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 sftp

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	configuration "github.com/pufferpanel/pufferd/config"
	"github.com/pufferpanel/pufferd/logging"
	"golang.org/x/crypto/ssh"
	"github.com/pufferpanel/sftp"
)

func Run() {
	e := runServer()
	if e != nil {
		logging.Error("Error starting SFTP", e)
	}
}

func runServer() error {
	config := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			return validateSSH(c.User(), string(pass))
		},
	}

	serverKeyFile := path.Join(configuration.GetOrDefault("datafolder", "data"), "server.key")

	_, e := os.Stat(serverKeyFile)

	if e != nil && os.IsNotExist(e) {
		logging.Debug("Generating new key")
		var key *rsa.PrivateKey
		key, e = rsa.GenerateKey(rand.Reader, 2048)
		if e != nil {
			return e
		}

		data := x509.MarshalPKCS1PrivateKey(key)
		block := pem.Block{
			Type:    "RSA PRIVATE KEY",
			Headers: nil,
			Bytes:   data,
		}
		ioutil.WriteFile(serverKeyFile, pem.EncodeToMemory(&block), 0700)
		if e != nil {
			return e
		}
	} else if e != nil {
		return e
	}

	logging.Debug("Loading existing key")
	var data []byte
	data, e = ioutil.ReadFile(serverKeyFile)
	if e != nil {
		return e
	}

	hkey, e := ssh.ParsePrivateKey(data)

	if e != nil {
		logging.Debug("trigger")
		return e
	}

	config.AddHostKey(hkey)

	bind := configuration.GetOrDefault("sftp", "0.0.0.0:5657")

	listener, e := net.Listen("tcp", bind)
	if e != nil {
		return e
	}
	logging.Infof("Started SFTP Server on %s", bind)

	go func() {
		for {
			conn, _ := listener.Accept()
			go HandleConn(conn, config)
		}
	}()

	return nil
}

func HandleConn(conn net.Conn, config *ssh.ServerConfig) {
	defer conn.Close()
	e := handleConn(conn, config)
	if e != nil {
		logging.Error("sftpd connection errored:", e)
	}
}
func handleConn(conn net.Conn, config *ssh.ServerConfig) error {
	sc, chans, reqs, e := ssh.NewServerConn(conn, config)
	if e != nil {
		return e
	}
	defer sc.Close()

	// The incoming Request channel must be serviced.
	go PrintDiscardRequests(reqs)

	// Service the incoming Channel channel.
	for newChannel := range chans {
		// Channels have a type, depending on the application level
		// protocol intended. In the case of an SFTP session, this is "subsystem"
		// with a payload string of "<length=4>sftp"
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			return err
		}

		// Sessions have out-of-band requests such as "shell",
		// "pty-req" and "env".  Here we handle only the
		// "subsystem" request.
		go func(in <-chan *ssh.Request) {
			for req := range in {
				ok := false
				switch req.Type {
				case "subsystem":
					if string(req.Payload[4:]) == "sftp" {
						ok = true
					}
				}
				req.Reply(ok, nil)
			}
		}(requests)

		serverOptions := []sftp.ServerOption{}

		server, err := sftp.NewServer(
			channel,
			serverOptions...,
		)
		if err != nil {
			return err
		}
		if err := server.Serve(); err != nil {
			return err
		}
	}
	return nil
}

func PrintDiscardRequests(in <-chan *ssh.Request) {
	for req := range in {
		if req.WantReply {
			req.Reply(false, nil)
		}
	}
}

func validateSSH(username string, password string) (*ssh.Permissions, error) {
	authUrl := configuration.Get("authserver")
	client := &http.Client{}
	data := url.Values{}
	data.Set("grant_type", "password")
	data.Set("username", username)
	data.Set("password", password)
	data.Set("scope", "sftp")
	token := configuration.Get("authtoken")
	request, _ := http.NewRequest("POST", authUrl, bytes.NewBufferString(data.Encode()))
	request.Header.Add("Authorization", "Bearer "+token)
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Add("Content-Length", strconv.Itoa(len(data.Encode())))
	response, err := client.Do(request)
	if err != nil {
		logging.Error("Error talking to auth server", err)
		return nil, errors.New("Invalid response from authorization server")
	}

	//we should only get a 200 or 400 back, if we get any others, we have a problem
	if response.StatusCode != 200 && response.StatusCode != 400 {
		logging.Error("Error talking to auth server", response.StatusCode)
		return nil, errors.New("Invalid response from authorization server")
	}
	var respArr map[string]interface{}
	json.NewDecoder(response.Body).Decode(&respArr)
	if respArr["error"] != nil {
		return nil, errors.New("Incorrect username or password")
	}
	sshPerms := &ssh.Permissions{}
	scopes := strings.Split(respArr["scope"].(string), " ")
	if len(scopes) != 2 {
		return nil, errors.New("Invalid response from authorization server")
	}
	for _, v := range scopes {
		if v != "sftp" {
			sshPerms.Extensions = make(map[string]string)
			sshPerms.Extensions["server_id"] = v
			return sshPerms, nil
		}
	}
	return nil, errors.New("Incorrect username or password")
}
<|endoftext|>


<commit_before>// Copyright (c) 2012-2015 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

// This file sets up the variables used, including testInitFns.
// Each file should add initialization that should be performed
// after flags are parsed.
//
// init is a multi-step process:
//   - setup vars (handled by init functions in each file)
//   - parse flags
//   - setup derived vars (handled by pre-init registered functions - registered in init function)
//   - post init (handled by post-init registered functions - registered in init function)
// This way, no one has to manage carefully control the initialization
// using file names, etc.
//
// Tests which require external dependencies need the -tag=x parameter.
// They should be run as:
//    go test -tags=x -run=. <other parameters ...>
// Benchmarks should also take this parameter, to include the sereal, xdr, etc.
// To run against codecgen, etc, make sure you pass extra parameters.
// Example usage:
//    go test "-tags=x codecgen" -bench=. <other parameters ...>
//
// To fully test everything:
//    go test -tags=x -benchtime=100ms -tv -bg -bi  -brw -bu -v -run=. -bench=.

// Handling flags
// codec_test.go will define a set of global flags for testing, including:
//   - Use Reset
//   - Use IO reader/writer (vs direct bytes)
//   - Set Canonical
//   - Set InternStrings
//   - Use Symbols
//
// This way, we can test them all by running same set of tests with a different
// set of flags.
//
// Following this, all the benchmarks will utilize flags set by codec_test.go
// and will not redefine these "global" flags.

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"sync"
	"testing"
)

import . "github.com/ugorji/go/codec"

type testHED struct {
	H Handle
	E *Encoder
	D *Decoder
}

type ioReaderWrapper struct {
	r io.Reader
}

func (x ioReaderWrapper) Read(p []byte) (n int, err error) {
	return x.r.Read(p)
}

type ioWriterWrapper struct {
	w io.Writer
}

func (x ioWriterWrapper) Write(p []byte) (n int, err error) {
	return x.w.Write(p)
}

var (
	// testNoopH    = NoopHandle(8)
	testMsgpackH = &MsgpackHandle{}
	testBincH    = &BincHandle{}
	testSimpleH  = &SimpleHandle{}
	testCborH    = &CborHandle{}
	testJsonH    = &JsonHandle{}

	testHandles     []Handle
	testPreInitFns  []func()
	testPostInitFns []func()

	testOnce sync.Once

	testHEDs []testHED
)

// flag variables used by tests (and bench)
var (
	testDepth int

	testVerbose       bool
	testInitDebug     bool
	testStructToArray bool
	testCanonical     bool
	testUseReset      bool
	testSkipIntf      bool
	testInternStr     bool
	testUseMust       bool
	testCheckCircRef  bool

	testUseIoEncDec  int
	testUseIoWrapper bool

	testMaxInitLen int

	testNumRepeatString int
)

// variables that are not flags, but which can configure the handles
var (
	testEncodeOptions EncodeOptions
	testDecodeOptions DecodeOptions
)

// flag variables used by bench
var (
	benchDoInitBench      bool
	benchVerify           bool
	benchUnscientificRes  bool = false
	benchMapStringKeyOnly bool
	//depth of 0 maps to ~400bytes json-encoded string, 1 maps to ~1400 bytes, etc
	//For depth>1, we likely trigger stack growth for encoders, making benchmarking unreliable.
	benchDepth     int
	benchInitDebug bool
)

func init() {
	testHEDs = make([]testHED, 0, 32)
	testHandles = append(testHandles,
		// testNoopH,
		testMsgpackH, testBincH, testSimpleH,
		testCborH, testJsonH)
	testInitFlags()
	benchInitFlags()
}

func testInitFlags() {
	// delete(testDecOpts.ExtFuncs, timeTyp)
	flag.IntVar(&testDepth, "tsd", 0, "Test Struc Depth")
	flag.BoolVar(&testVerbose, "tv", false, "Test Verbose")
	flag.BoolVar(&testInitDebug, "tg", false, "Test Init Debug")
	flag.IntVar(&testUseIoEncDec, "ti", -1, "Use IO Reader/Writer for Marshal/Unmarshal ie >= 0")
	flag.BoolVar(&testUseIoWrapper, "tiw", false, "Wrap the IO Reader/Writer with a base pass-through reader/writer")
	flag.BoolVar(&testStructToArray, "ts", false, "Set StructToArray option")
	flag.BoolVar(&testCanonical, "tc", false, "Set Canonical option")
	flag.BoolVar(&testInternStr, "te", false, "Set InternStr option")
	flag.BoolVar(&testSkipIntf, "tf", false, "Skip Interfaces")
	flag.BoolVar(&testUseReset, "tr", false, "Use Reset")
	flag.IntVar(&testNumRepeatString, "trs", 8, "Create string variables by repeating a string N times")
	flag.IntVar(&testMaxInitLen, "tx", 0, "Max Init Len")
	flag.BoolVar(&testUseMust, "tm", true, "Use Must(En|De)code")
	flag.BoolVar(&testCheckCircRef, "tl", false, "Use Check Circular Ref")
}

func benchInitFlags() {
	flag.BoolVar(&benchMapStringKeyOnly, "bs", false, "Bench use maps with string keys only")
	flag.BoolVar(&benchInitDebug, "bg", false, "Bench Debug")
	flag.IntVar(&benchDepth, "bd", 1, "Bench Depth")
	flag.BoolVar(&benchDoInitBench, "bi", false, "Run Bench Init")
	flag.BoolVar(&benchVerify, "bv", false, "Verify Decoded Value during Benchmark")
	flag.BoolVar(&benchUnscientificRes, "bu", false, "Show Unscientific Results during Benchmark")
}

func testHEDGet(h Handle) *testHED {
	for i := range testHEDs {
		v := &testHEDs[i]
		if v.H == h {
			return v
		}
	}
	testHEDs = append(testHEDs, testHED{h, NewEncoder(nil, h), NewDecoder(nil, h)})
	return &testHEDs[len(testHEDs)-1]
}

func testReinit() {
	testOnce = sync.Once{}
	testHEDs = nil
}

func testInitAll() {
	// only parse it once.
	if !flag.Parsed() {
		flag.Parse()
	}
	for _, f := range testPreInitFns {
		f()
	}
	for _, f := range testPostInitFns {
		f()
	}
}

func sTestCodecEncode(ts interface{}, bsIn []byte, fn func([]byte) *bytes.Buffer,
	h Handle, bh *BasicHandle) (bs []byte, err error) {
	// bs = make([]byte, 0, approxSize)
	var e *Encoder
	var buf *bytes.Buffer
	if testUseReset {
		e = testHEDGet(h).E
	} else {
		e = NewEncoder(nil, h)
	}
	var oldWriteBufferSize int
	if testUseIoEncDec >= 0 {
		buf = fn(bsIn)
		// set the encode options for using a buffer
		oldWriteBufferSize = bh.WriterBufferSize
		bh.WriterBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			e.Reset(ioWriterWrapper{buf})
		} else {
			e.Reset(buf)
		}
	} else {
		bs = bsIn
		e.ResetBytes(&bs)
	}
	if testUseMust {
		e.MustEncode(ts)
	} else {
		err = e.Encode(ts)
	}
	if testUseIoEncDec >= 0 {
		bs = buf.Bytes()
		bh.WriterBufferSize = oldWriteBufferSize
	}
	return
}

func sTestCodecDecode(bs []byte, ts interface{}, h Handle, bh *BasicHandle) (err error) {
	var d *Decoder
	// var buf *bytes.Reader
	if testUseReset {
		d = testHEDGet(h).D
	} else {
		d = NewDecoder(nil, h)
	}
	var oldReadBufferSize int
	if testUseIoEncDec >= 0 {
		buf := bytes.NewReader(bs)
		oldReadBufferSize = bh.ReaderBufferSize
		bh.ReaderBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			d.Reset(ioReaderWrapper{buf})
		} else {
			d.Reset(buf)
		}
	} else {
		d.ResetBytes(bs)
	}
	if testUseMust {
		d.MustDecode(ts)
	} else {
		err = d.Decode(ts)
	}
	if testUseIoEncDec >= 0 {
		bh.ReaderBufferSize = oldReadBufferSize
	}
	return
}

// --- functions below are used by both benchmarks and tests

func logT(x interface{}, format string, args ...interface{}) {
	if t, ok := x.(*testing.T); ok && t != nil {
		t.Logf(format, args...)
	} else if b, ok := x.(*testing.B); ok && b != nil {
		b.Logf(format, args...)
	} else if testVerbose {
		if len(format) == 0 || format[len(format)-1] != '\n' {
			format = format + "\n"
		}
		fmt.Printf(format, args...)
	}
}

// --- functions below are used only by benchmarks alone

func fnBenchmarkByteBuf(bsIn []byte) (buf *bytes.Buffer) {
	// var buf bytes.Buffer
	// buf.Grow(approxSize)
	buf = bytes.NewBuffer(bsIn)
	buf.Truncate(0)
	return
}

// func benchFnCodecEncode(ts interface{}, bsIn []byte, h Handle) (bs []byte, err error) {
// 	return testCodecEncode(ts, bsIn, fnBenchmarkByteBuf, h)
// }

// func benchFnCodecDecode(bs []byte, ts interface{}, h Handle) (err error) {
// 	return testCodecDecode(bs, ts, h)
// }
<commit_msg>codec: update<commit_after>// Copyright (c) 2012-2015 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

// This file sets up the variables used, including testInitFns.
// Each file should add initialization that should be performed
// after flags are parsed.
//
// init is a multi-step process:
//   - setup vars (handled by init functions in each file)
//   - parse flags
//   - setup derived vars (handled by pre-init registered functions - registered in init function)
//   - post init (handled by post-init registered functions - registered in init function)
// This way, no one has to manage carefully control the initialization
// using file names, etc.
//
// Tests which require external dependencies need the -tag=x parameter.
// They should be run as:
//    go test -tags=x -run=. <other parameters ...>
// Benchmarks should also take this parameter, to include the sereal, xdr, etc.
// To run against codecgen, etc, make sure you pass extra parameters.
// Example usage:
//    go test "-tags=x codecgen" -bench=. <other parameters ...>
//
// To fully test everything:
//    go test -tags=x -benchtime=100ms -tv -bg -bi  -brw -bu -v -run=. -bench=.

// Handling flags
// codec_test.go will define a set of global flags for testing, including:
//   - Use Reset
//   - Use IO reader/writer (vs direct bytes)
//   - Set Canonical
//   - Set InternStrings
//   - Use Symbols
//
// This way, we can test them all by running same set of tests with a different
// set of flags.
//
// Following this, all the benchmarks will utilize flags set by codec_test.go
// and will not redefine these "global" flags.

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"sync"
	"testing"
)

import . "github.com/ugorji/go/codec"

type testHED struct {
	H Handle
	E *Encoder
	D *Decoder
}

type ioReaderWrapper struct {
	r io.Reader
}

func (x ioReaderWrapper) Read(p []byte) (n int, err error) {
	return x.r.Read(p)
}

type ioWriterWrapper struct {
	w io.Writer
}

func (x ioWriterWrapper) Write(p []byte) (n int, err error) {
	return x.w.Write(p)
}

var (
	// testNoopH    = NoopHandle(8)
	testMsgpackH = &MsgpackHandle{}
	testBincH    = &BincHandle{}
	testSimpleH  = &SimpleHandle{}
	testCborH    = &CborHandle{}
	testJsonH    = &JsonHandle{}

	testHandles     []Handle
	testPreInitFns  []func()
	testPostInitFns []func()

	testOnce sync.Once

	testHEDs []testHED
)

// flag variables used by tests (and bench)
var (
	testDepth int

	testVerbose       bool
	testInitDebug     bool
	testStructToArray bool
	testCanonical     bool
	testUseReset      bool
	testSkipIntf      bool
	testInternStr     bool
	testUseMust       bool
	testCheckCircRef  bool

	testUseIoEncDec  int
	testUseIoWrapper bool

	testMaxInitLen int

	testNumRepeatString int

	testRpcBufsize int
)

// variables that are not flags, but which can configure the handles
var (
	testEncodeOptions EncodeOptions
	testDecodeOptions DecodeOptions
)

// flag variables used by bench
var (
	benchDoInitBench      bool
	benchVerify           bool
	benchUnscientificRes  bool = false
	benchMapStringKeyOnly bool
	//depth of 0 maps to ~400bytes json-encoded string, 1 maps to ~1400 bytes, etc
	//For depth>1, we likely trigger stack growth for encoders, making benchmarking unreliable.
	benchDepth     int
	benchInitDebug bool
)

func init() {
	testHEDs = make([]testHED, 0, 32)
	testHandles = append(testHandles,
		// testNoopH,
		testMsgpackH, testBincH, testSimpleH,
		testCborH, testJsonH)
	testInitFlags()
	benchInitFlags()
}

func testInitFlags() {
	// delete(testDecOpts.ExtFuncs, timeTyp)
	flag.IntVar(&testDepth, "tsd", 0, "Test Struc Depth")
	flag.BoolVar(&testVerbose, "tv", false, "Test Verbose")
	flag.BoolVar(&testInitDebug, "tg", false, "Test Init Debug")
	flag.IntVar(&testUseIoEncDec, "ti", -1, "Use IO Reader/Writer for Marshal/Unmarshal ie >= 0")
	flag.BoolVar(&testUseIoWrapper, "tiw", false, "Wrap the IO Reader/Writer with a base pass-through reader/writer")
	flag.BoolVar(&testStructToArray, "ts", false, "Set StructToArray option")
	flag.BoolVar(&testCanonical, "tc", false, "Set Canonical option")
	flag.BoolVar(&testInternStr, "te", false, "Set InternStr option")
	flag.BoolVar(&testSkipIntf, "tf", false, "Skip Interfaces")
	flag.BoolVar(&testUseReset, "tr", false, "Use Reset")
	flag.IntVar(&testNumRepeatString, "trs", 8, "Create string variables by repeating a string N times")
	flag.IntVar(&testMaxInitLen, "tx", 0, "Max Init Len")
	flag.BoolVar(&testUseMust, "tm", true, "Use Must(En|De)code")
	flag.BoolVar(&testCheckCircRef, "tl", false, "Use Check Circular Ref")
}

func benchInitFlags() {
	flag.BoolVar(&benchMapStringKeyOnly, "bs", false, "Bench use maps with string keys only")
	flag.BoolVar(&benchInitDebug, "bg", false, "Bench Debug")
	flag.IntVar(&benchDepth, "bd", 1, "Bench Depth")
	flag.BoolVar(&benchDoInitBench, "bi", false, "Run Bench Init")
	flag.BoolVar(&benchVerify, "bv", false, "Verify Decoded Value during Benchmark")
	flag.BoolVar(&benchUnscientificRes, "bu", false, "Show Unscientific Results during Benchmark")
}

func testHEDGet(h Handle) *testHED {
	for i := range testHEDs {
		v := &testHEDs[i]
		if v.H == h {
			return v
		}
	}
	testHEDs = append(testHEDs, testHED{h, NewEncoder(nil, h), NewDecoder(nil, h)})
	return &testHEDs[len(testHEDs)-1]
}

func testReinit() {
	testOnce = sync.Once{}
	testHEDs = nil
}

func testInitAll() {
	// only parse it once.
	if !flag.Parsed() {
		flag.Parse()
	}
	for _, f := range testPreInitFns {
		f()
	}
	for _, f := range testPostInitFns {
		f()
	}
}

func sTestCodecEncode(ts interface{}, bsIn []byte, fn func([]byte) *bytes.Buffer,
	h Handle, bh *BasicHandle) (bs []byte, err error) {
	// bs = make([]byte, 0, approxSize)
	var e *Encoder
	var buf *bytes.Buffer
	if testUseReset {
		e = testHEDGet(h).E
	} else {
		e = NewEncoder(nil, h)
	}
	var oldWriteBufferSize int
	if testUseIoEncDec >= 0 {
		buf = fn(bsIn)
		// set the encode options for using a buffer
		oldWriteBufferSize = bh.WriterBufferSize
		bh.WriterBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			e.Reset(ioWriterWrapper{buf})
		} else {
			e.Reset(buf)
		}
	} else {
		bs = bsIn
		e.ResetBytes(&bs)
	}
	if testUseMust {
		e.MustEncode(ts)
	} else {
		err = e.Encode(ts)
	}
	if testUseIoEncDec >= 0 {
		bs = buf.Bytes()
		bh.WriterBufferSize = oldWriteBufferSize
	}
	return
}

func sTestCodecDecode(bs []byte, ts interface{}, h Handle, bh *BasicHandle) (err error) {
	var d *Decoder
	// var buf *bytes.Reader
	if testUseReset {
		d = testHEDGet(h).D
	} else {
		d = NewDecoder(nil, h)
	}
	var oldReadBufferSize int
	if testUseIoEncDec >= 0 {
		buf := bytes.NewReader(bs)
		oldReadBufferSize = bh.ReaderBufferSize
		bh.ReaderBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			d.Reset(ioReaderWrapper{buf})
		} else {
			d.Reset(buf)
		}
	} else {
		d.ResetBytes(bs)
	}
	if testUseMust {
		d.MustDecode(ts)
	} else {
		err = d.Decode(ts)
	}
	if testUseIoEncDec >= 0 {
		bh.ReaderBufferSize = oldReadBufferSize
	}
	return
}

// --- functions below are used by both benchmarks and tests

func logT(x interface{}, format string, args ...interface{}) {
	if t, ok := x.(*testing.T); ok && t != nil {
		t.Logf(format, args...)
	} else if b, ok := x.(*testing.B); ok && b != nil {
		b.Logf(format, args...)
	} else if testVerbose {
		if len(format) == 0 || format[len(format)-1] != '\n' {
			format = format + "\n"
		}
		fmt.Printf(format, args...)
	}
}

// --- functions below are used only by benchmarks alone

func fnBenchmarkByteBuf(bsIn []byte) (buf *bytes.Buffer) {
	// var buf bytes.Buffer
	// buf.Grow(approxSize)
	buf = bytes.NewBuffer(bsIn)
	buf.Truncate(0)
	return
}

// func benchFnCodecEncode(ts interface{}, bsIn []byte, h Handle) (bs []byte, err error) {
// 	return testCodecEncode(ts, bsIn, fnBenchmarkByteBuf, h)
// }

// func benchFnCodecDecode(bs []byte, ts interface{}, h Handle) (err error) {
// 	return testCodecDecode(bs, ts, h)
// }
<|endoftext|>


<commit_before>package vm

type Bool bool

var True = NewBool(true)
var False = NewBool(false)

func NewBool(b bool) *Thunk {
	return Normal(Bool(b))
}

func (b Bool) Equal(e Equalable) bool {
	return b == e.(Bool)
}

func And(ts ...*Thunk) *Thunk { // With short circuit
	for _, t := range ts {
		o := t.Eval()
		b, ok := o.(Bool)

		if !ok {
			return notBoolError(o)
		} else if !bool(b) {
			return False
		}
	}

	return True
}

func Or(ts ...*Thunk) *Thunk { // With short circuit
	for _, t := range ts {
		o := t.Eval()
		b, ok := o.(Bool)

		if !ok {
			return notBoolError(o)
		} else if bool(b) {
			return True
		}
	}

	return False
}

func Not(ts ...*Thunk) *Thunk {
	if len(ts) != 1 {
		return NumArgsError("not", "1")
	}

	o := ts[0].Eval()
	b, ok := o.(Bool)

	if !ok {
		return notBoolError(o)
	}

	return NewBool(!bool(b))
}

func notBoolError(o Object) *Thunk {
	return TypeError(o, "Bool")
}
<commit_msg>Refactor bool.go<commit_after>package vm

type Bool bool

var True, False = NewBool(true), NewBool(false)

func NewBool(b bool) *Thunk {
	return Normal(Bool(b))
}

func (b Bool) Equal(e Equalable) bool {
	return b == e.(Bool)
}

func And(ts ...*Thunk) *Thunk { // With short circuit
	for _, t := range ts {
		o := t.Eval()
		b, ok := o.(Bool)

		if !ok {
			return notBoolError(o)
		} else if !bool(b) {
			return False
		}
	}

	return True
}

func Or(ts ...*Thunk) *Thunk { // With short circuit
	for _, t := range ts {
		o := t.Eval()
		b, ok := o.(Bool)

		if !ok {
			return notBoolError(o)
		} else if bool(b) {
			return True
		}
	}

	return False
}

func Not(ts ...*Thunk) *Thunk {
	if len(ts) != 1 {
		return NumArgsError("not", "1")
	}

	o := ts[0].Eval()
	b, ok := o.(Bool)

	if !ok {
		return notBoolError(o)
	}

	return NewBool(!bool(b))
}

func notBoolError(o Object) *Thunk {
	return TypeError(o, "Bool")
}
<|endoftext|>


<commit_before>package match

var hasAssembler bool

// Match4 will return start indeces of all matches of a 4 byte needle
// in a haystack that is a multiple of 16 in length.
// Indeces are returned ordered from index 0 and upwards.
func Match4(needle, haystack []byte, indices []int) []int {
	if len(needle) != 4 {
		panic("length not 4")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]byte, len(haystack)/8)
	if indices == nil {
		indices = make([]int, 0, 10)
	}
	find4(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&1 == 1 {
				indices = append(indices, i*8+j)
			}
			v >>= 1
			j++
		}
	}
	return indices
}

func find4(needle, haystack, dst []byte) {
	if hasAssembler {
		find4SSE4(needle, haystack, dst)
		return
	}
	find4Go(needle, haystack, dst)
}

func find4Go(needle, haystack, dst []byte) {
	end := uint(len(haystack) - 3)
	for i := uint(0); i < end; i++ {
		if needle[0] == haystack[i] {
			if needle[1] == haystack[i+1] && needle[2] == haystack[i+2] && needle[3] == haystack[i+3] {
				dst[i>>3] |= 1 << (i & 7)
			}
		}
	}
}

// Match8 will return start indeces of all matches of a 8 byte needle
// in a haystack that is a multiple of 16 in length.
// Indeces are returned ordered from index 0 and upwards.
func Match8(needle, haystack []byte, indices []int) []int {
	if len(needle) != 8 {
		panic("length not 8")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]uint16, len(haystack)/8)
	if indices == nil {
		indices = make([]int, 0, 10)
	}
	find8(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&3 == 3 {
				indices = append(indices, i*8+j)
			}
			v >>= 2
			j++
		}
	}
	return indices
}

// Match4And will return start indeces of all matches of a 8 byte needle
// in a haystack that is a multiple of 16 in length.
// Matches for the first four bytes are returned in the first index, and 8
// byte matches are returned in the second. An index that is an 8 byte match will
// not be present in the 4-byte matches.
// Indeces are returned ordered from index 0 and upwards.
func Match8And4(needle, haystack []byte, indices8 []int, indices4 []int) ([]int, []int) {
	if len(needle) != 8 {
		panic("length not 8")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]uint16, len(haystack)/8)
	if indices8 == nil {
		indices8 = make([]int, 0, 10)
	} else {
		indices8 = indices8[:0]
	}
	if indices4 == nil {
		indices4 = make([]int, 0, 10)
	} else {
		indices4 = indices4[:0]
	}
	find8(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&3 == 3 {
				indices8 = append(indices8, i*8+j)
			} else if v&1 == 1 {
				indices4 = append(indices4, i*8+j)
			}
			v >>= 2
			j++
		}
	}
	return indices8, indices4
}

func find8(needle, haystack []byte, dst []uint16) {
	if hasAssembler {
		find8SSE4(needle, haystack, dst)
		return
	}
	find8Go(needle, haystack, dst)
}

func find8Go(needle, haystack []byte, dst []uint16) {
	end := uint(len(haystack) - 7)
	for i := uint(0); i < end; i++ {
		if needle[0] == haystack[i] && needle[1] == haystack[i+1] && needle[2] == haystack[i+2] && needle[3] == haystack[i+3] {
			dst[i>>3] |= 1 << ((i & 7) << 1)
		}
		if needle[4] == haystack[i+4] && needle[5] == haystack[i+5] && needle[6] == haystack[i+6] && needle[7] == haystack[i+7] {
			dst[i>>3] |= 2 << ((i & 7) << 1)
		}
	}
}
<commit_msg>+documentation.<commit_after>package match

var hasAssembler bool

// Match4 will return start indeces of all matches of a 4 byte needle
// in a haystack that is a multiple of 16 in length.
// Indeces are returned ordered from index 0 and upwards.
func Match4(needle, haystack []byte, indices []int) []int {
	if len(needle) != 4 {
		panic("length not 4")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]byte, len(haystack)/8)
	if indices == nil {
		indices = make([]int, 0, 10)
	}
	find4(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&1 == 1 {
				indices = append(indices, i*8+j)
			}
			v >>= 1
			j++
		}
	}
	return indices
}

func find4(needle, haystack, dst []byte) {
	if hasAssembler {
		find4SSE4(needle, haystack, dst)
		return
	}
	find4Go(needle, haystack, dst)
}

// find4Go is the reference implementation that mimmics the SSE4
// implemenation.
func find4Go(needle, haystack, dst []byte) {
	end := uint(len(haystack) - 3)
	for i := uint(0); i < end; i++ {
		if needle[0] == haystack[i] {
			if needle[1] == haystack[i+1] && needle[2] == haystack[i+2] && needle[3] == haystack[i+3] {
				dst[i>>3] |= 1 << (i & 7)
			}
		}
	}
}

// Match8 will return start indeces of all matches of a 8 byte needle
// in a haystack that is a multiple of 16 in length.
// Indeces are returned ordered from index 0 and upwards.
func Match8(needle, haystack []byte, indices []int) []int {
	if len(needle) != 8 {
		panic("length not 8")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]uint16, len(haystack)/8)
	if indices == nil {
		indices = make([]int, 0, 10)
	}
	find8(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&3 == 3 {
				indices = append(indices, i*8+j)
			}
			v >>= 2
			j++
		}
	}
	return indices
}

// Match4And will return start indeces of all matches of a 8 byte needle
// in a haystack that is a multiple of 16 in length.
// Matches for the first four bytes are returned in the first index, and 8
// byte matches are returned in the second. An index that is an 8 byte match will
// not be present in the 4-byte matches.
// Indeces are returned ordered from index 0 and upwards.
func Match8And4(needle, haystack []byte, indices8 []int, indices4 []int) ([]int, []int) {
	if len(needle) != 8 {
		panic("length not 8")
	}
	if len(haystack)&15 != 0 {
		panic("haystack must be dividable by 16")
	}
	dst := make([]uint16, len(haystack)/8)
	if indices8 == nil {
		indices8 = make([]int, 0, 10)
	} else {
		indices8 = indices8[:0]
	}
	if indices4 == nil {
		indices4 = make([]int, 0, 10)
	} else {
		indices4 = indices4[:0]
	}
	find8(needle, haystack, dst)
	for i, v := range dst {
		j := 0
		for v != 0 {
			if v&3 == 3 {
				indices8 = append(indices8, i*8+j)
			} else if v&1 == 1 {
				indices4 = append(indices4, i*8+j)
			}
			v >>= 2
			j++
		}
	}
	return indices8, indices4
}

func find8(needle, haystack []byte, dst []uint16) {
	if hasAssembler {
		find8SSE4(needle, haystack, dst)
		return
	}
	find8Go(needle, haystack, dst)
}

// find8Go is the reference implementation that mimmics the SSE4
// implemenation.
func find8Go(needle, haystack []byte, dst []uint16) {
	end := uint(len(haystack) - 7)
	for i := uint(0); i < end; i++ {
		if needle[0] == haystack[i] && needle[1] == haystack[i+1] && needle[2] == haystack[i+2] && needle[3] == haystack[i+3] {
			dst[i>>3] |= 1 << ((i & 7) << 1)
		}
		if needle[4] == haystack[i+4] && needle[5] == haystack[i+5] && needle[6] == haystack[i+6] && needle[7] == haystack[i+7] {
			dst[i>>3] |= 2 << ((i & 7) << 1)
		}
	}
}
<|endoftext|>


<commit_before>package memdb

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/t3rm1n4l/memdb/skiplist"
	"io"
	"math/rand"
	"os"
	"path"
	"sync/atomic"
)

const DiskBlockSize = 512 * 1024

var (
	ErrNotEnoughSpace = errors.New("Not enough space in the buffer")
)

type Item struct {
	bornSn, deadSn uint32
	data           []byte
}

func (itm *Item) Encode(buf []byte, w io.Writer) error {
	l := 2
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint16(buf[0:2], uint16(len(itm.data)))
	if _, err := w.Write(buf[0:2]); err != nil {
		return err
	}
	if _, err := w.Write(itm.data); err != nil {
		return err
	}

	return nil
}

func (itm *Item) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:2]); err != nil {
		return err
	}
	l := binary.BigEndian.Uint16(buf[0:2])
	itm.data = make([]byte, int(l))
	_, err := io.ReadFull(r, itm.data)

	return err
}

func NewItem(data []byte) *Item {
	return &Item{
		data: data,
	}
}

func InsertCompare(this skiplist.Item, that skiplist.Item) int {
	var v int
	if v = DataCompare(this, that); v == 0 {
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		v = int(thisItem.bornSn - thatItem.bornSn)
	}

	return v
}

func DataCompare(this skiplist.Item, that skiplist.Item) int {
	var l int
	thisItem := this.(*Item)
	thatItem := that.(*Item)

	l1 := len(thisItem.data)
	l2 := len(thatItem.data)
	if l1 < l2 {
		l = l1
	} else {
		l = l2
	}
	return bytes.Compare(thisItem.data[:l], thatItem.data[:l])
}

//
//compare item,sn
type Writer struct {
	rand *rand.Rand
	buf  *skiplist.ActionBuffer
	iter *skiplist.Iterator
	*MemDB
}

func (w *Writer) Put(x *Item) {
	x.bornSn = w.getCurrSn()
	w.store.Insert2(x, InsertCompare, w.buf, w.rand.Float32)
}

// find first item, seek until dead=0, mark dead=sn
func (w *Writer) Delete(x *Item) bool {
	gotItem := w.Get(x)
	if gotItem != nil {
		sn := w.getCurrSn()
		return atomic.CompareAndSwapUint32(&gotItem.deadSn, 0, sn)
	}

	return false
}

func (w *Writer) Get(x *Item) *Item {
	var curr *Item
	found := w.iter.Seek(x)
	if !found {
		return nil
	}

	// Seek until most recent item for key is found
	curr = w.iter.Get().(*Item)
	for {
		w.iter.Next()
		if !w.iter.Valid() {
			break
		}
		next := w.iter.Get().(*Item)
		if DataCompare(next, curr) != 0 {
			break
		}

		curr = next
	}

	if curr.deadSn != 0 {
		return nil
	}

	return curr
}

type MemDB struct {
	store       *skiplist.Skiplist
	currSn      uint32
	snapshots   *skiplist.Skiplist
	isGCRunning int32
	lastGCSn    uint32
}

func New() *MemDB {
	return &MemDB{
		store:     skiplist.New(),
		snapshots: skiplist.New(),
		currSn:    1,
	}
}

func (m *MemDB) getCurrSn() uint32 {
	return atomic.LoadUint32(&m.currSn)
}

func (m *MemDB) NewWriter() *Writer {
	buf := m.store.MakeBuf()

	return &Writer{
		rand:  rand.New(rand.NewSource(int64(rand.Int()))),
		buf:   buf,
		iter:  m.store.NewIterator(DataCompare, buf),
		MemDB: m,
	}
}

type Snapshot struct {
	sn       uint32
	refCount int32
	db       *MemDB
}

func (s *Snapshot) Encode(buf []byte, w io.Writer) error {
	l := 4
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint32(buf[0:4], s.sn)
	if _, err := w.Write(buf[0:4]); err != nil {
		return err
	}

	return nil

}

func (s *Snapshot) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:4]); err != nil {
		return err
	}
	s.sn = binary.BigEndian.Uint32(buf[0:4])
	return nil
}

func (s *Snapshot) Open() bool {
	if atomic.LoadInt32(&s.refCount) == 0 {
		return false
	}
	atomic.AddInt32(&s.refCount, 1)
	return true
}

func (s *Snapshot) Close() {
	newRefcount := atomic.AddInt32(&s.refCount, -1)
	if newRefcount == 0 {
		buf := s.db.snapshots.MakeBuf()
		defer s.db.snapshots.FreeBuf(buf)
		s.db.snapshots.Delete(s, CompareSnapshot, buf)
		if atomic.CompareAndSwapInt32(&s.db.isGCRunning, 0, 1) {
			go s.db.GC()
		}
	}
}

func CompareSnapshot(this skiplist.Item, that skiplist.Item) int {
	thisItem := this.(*Snapshot)
	thatItem := that.(*Snapshot)

	return int(thisItem.sn) - int(thatItem.sn)
}

func (m *MemDB) NewSnapshot() *Snapshot {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	snap := &Snapshot{db: m, sn: m.getCurrSn(), refCount: 1}
	m.snapshots.Insert(snap, CompareSnapshot, buf)
	atomic.AddUint32(&m.currSn, 1)
	return snap
}

type Iterator struct {
	snap *Snapshot
	iter *skiplist.Iterator
	buf  *skiplist.ActionBuffer
}

func (it *Iterator) skipUnwanted() {
loop:
	if !it.iter.Valid() {
		return
	}
	itm := it.iter.Get().(*Item)
	if itm.bornSn > it.snap.sn || (itm.deadSn > 0 && itm.deadSn <= it.snap.sn) {
		it.iter.Next()
		goto loop
	}
}

func (it *Iterator) SeekFirst() {
	it.iter.SeekFirst()
	it.skipUnwanted()
}

func (it *Iterator) Seek(itm *Item) {
	it.iter.Seek(itm)
	it.skipUnwanted()
}

func (it *Iterator) Valid() bool {
	return it.iter.Valid()
}

func (it *Iterator) Get() *Item {
	return it.iter.Get().(*Item)
}

func (it *Iterator) Next() {
	it.iter.Next()
	it.skipUnwanted()
}

func (it *Iterator) Close() {
	it.snap.Close()
	it.snap.db.store.FreeBuf(it.buf)
}

func (m *MemDB) NewIterator(snap *Snapshot) *Iterator {
	if !snap.Open() {
		return nil
	}
	buf := snap.db.store.MakeBuf()
	return &Iterator{
		snap: snap,
		iter: m.store.NewIterator(DataCompare, buf),
		buf:  buf,
	}
}

func (m *MemDB) collectDead(sn uint32) {
	buf1 := m.snapshots.MakeBuf()
	buf2 := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf1)
	defer m.snapshots.FreeBuf(buf2)
	iter := m.store.NewIterator(DataCompare, buf1)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		itm := iter.Get().(*Item)
		if itm.deadSn > 0 && itm.deadSn <= sn {
			m.store.Delete(itm, DataCompare, buf2)
		}
	}
}

func (m *MemDB) GC() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	if iter.Valid() {
		snap := iter.Get().(*Snapshot)
		if snap.sn != m.lastGCSn && snap.sn > 1 {
			m.lastGCSn = snap.sn - 1
			m.collectDead(m.lastGCSn)
		}
	}

	atomic.CompareAndSwapInt32(&m.isGCRunning, 1, 0)
}

func (m *MemDB) GetSnapshots() []*Snapshot {
	var snaps []*Snapshot
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		snaps = append(snaps, iter.Get().(*Snapshot))
	}

	return snaps
}

func (m *MemDB) StoreToDisk(dir string, snap *Snapshot) error {
	os.MkdirAll(dir, 0755)
	file, err := os.OpenFile(path.Join(dir, "records.data"), os.O_WRONLY|os.O_CREATE, 0755)
	if err != nil {
		return err
	}
	defer file.Close()

	w := bufio.NewWriterSize(file, DiskBlockSize)
	defer w.Flush()

	buf := make([]byte, 4)
	itr := m.NewIterator(snap)
	if itr == nil {
		return errors.New("Invalid snapshot")
	}
	defer itr.Close()

	if err := snap.Encode(buf, w); err != nil {
		return err
	}

	for itr.SeekFirst(); itr.Valid(); itr.Next() {
		itm := itr.Get()
		if err := itm.Encode(buf, w); err != nil {
			return err
		}
	}

	endItem := &Item{
		data: []byte(nil),
	}

	if err := endItem.Encode(buf, w); err != nil {
		return err
	}

	w.Flush()

	return nil
}

func (m *MemDB) LoadFromDisk(dir string) (*Snapshot, error) {
	file, err := os.OpenFile(path.Join(dir, "records.data"), os.O_RDONLY, 0755)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	w := m.NewWriter()

	buf := make([]byte, 4)
	r := bufio.NewReaderSize(file, DiskBlockSize)
	snap := &Snapshot{db: m, refCount: 1}
	if err := snap.Decode(buf, r); err != nil {
		return nil, err
	}
	m.currSn = snap.sn

loop:
	for {
		itm := &Item{}
		itm.Decode(buf, r)
		if len(itm.data) == 0 {
			break loop
		}

		w.Put(itm)
	}

	snbuf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(snbuf)
	m.snapshots.Insert(snap, CompareSnapshot, snbuf)
	m.currSn++

	return snap, nil
}
<commit_msg>memdb: Add worker routines for concurrently loading memdb from disk<commit_after>package memdb

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/t3rm1n4l/memdb/skiplist"
	"io"
	"math/rand"
	"os"
	"path"
	"runtime"
	"sync"
	"sync/atomic"
)

const DiskBlockSize = 512 * 1024

var (
	ErrNotEnoughSpace = errors.New("Not enough space in the buffer")
)

type Item struct {
	bornSn, deadSn uint32
	data           []byte
}

func (itm *Item) Encode(buf []byte, w io.Writer) error {
	l := 2
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint16(buf[0:2], uint16(len(itm.data)))
	if _, err := w.Write(buf[0:2]); err != nil {
		return err
	}
	if _, err := w.Write(itm.data); err != nil {
		return err
	}

	return nil
}

func (itm *Item) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:2]); err != nil {
		return err
	}
	l := binary.BigEndian.Uint16(buf[0:2])
	itm.data = make([]byte, int(l))
	_, err := io.ReadFull(r, itm.data)

	return err
}

func NewItem(data []byte) *Item {
	return &Item{
		data: data,
	}
}

func InsertCompare(this skiplist.Item, that skiplist.Item) int {
	var v int
	if v = DataCompare(this, that); v == 0 {
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		v = int(thisItem.bornSn - thatItem.bornSn)
	}

	return v
}

func DataCompare(this skiplist.Item, that skiplist.Item) int {
	var l int
	thisItem := this.(*Item)
	thatItem := that.(*Item)

	l1 := len(thisItem.data)
	l2 := len(thatItem.data)
	if l1 < l2 {
		l = l1
	} else {
		l = l2
	}
	return bytes.Compare(thisItem.data[:l], thatItem.data[:l])
}

//
//compare item,sn
type Writer struct {
	rand *rand.Rand
	buf  *skiplist.ActionBuffer
	iter *skiplist.Iterator
	*MemDB
}

func (w *Writer) Put(x *Item) {
	x.bornSn = w.getCurrSn()
	w.store.Insert2(x, InsertCompare, w.buf, w.rand.Float32)
}

// find first item, seek until dead=0, mark dead=sn
func (w *Writer) Delete(x *Item) bool {
	gotItem := w.Get(x)
	if gotItem != nil {
		sn := w.getCurrSn()
		return atomic.CompareAndSwapUint32(&gotItem.deadSn, 0, sn)
	}

	return false
}

func (w *Writer) Get(x *Item) *Item {
	var curr *Item
	found := w.iter.Seek(x)
	if !found {
		return nil
	}

	// Seek until most recent item for key is found
	curr = w.iter.Get().(*Item)
	for {
		w.iter.Next()
		if !w.iter.Valid() {
			break
		}
		next := w.iter.Get().(*Item)
		if DataCompare(next, curr) != 0 {
			break
		}

		curr = next
	}

	if curr.deadSn != 0 {
		return nil
	}

	return curr
}

type MemDB struct {
	store       *skiplist.Skiplist
	currSn      uint32
	snapshots   *skiplist.Skiplist
	isGCRunning int32
	lastGCSn    uint32
}

func New() *MemDB {
	return &MemDB{
		store:     skiplist.New(),
		snapshots: skiplist.New(),
		currSn:    1,
	}
}

func (m *MemDB) getCurrSn() uint32 {
	return atomic.LoadUint32(&m.currSn)
}

func (m *MemDB) NewWriter() *Writer {
	buf := m.store.MakeBuf()

	return &Writer{
		rand:  rand.New(rand.NewSource(int64(rand.Int()))),
		buf:   buf,
		iter:  m.store.NewIterator(DataCompare, buf),
		MemDB: m,
	}
}

type Snapshot struct {
	sn       uint32
	refCount int32
	db       *MemDB
}

func (s *Snapshot) Encode(buf []byte, w io.Writer) error {
	l := 4
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint32(buf[0:4], s.sn)
	if _, err := w.Write(buf[0:4]); err != nil {
		return err
	}

	return nil

}

func (s *Snapshot) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:4]); err != nil {
		return err
	}
	s.sn = binary.BigEndian.Uint32(buf[0:4])
	return nil
}

func (s *Snapshot) Open() bool {
	if atomic.LoadInt32(&s.refCount) == 0 {
		return false
	}
	atomic.AddInt32(&s.refCount, 1)
	return true
}

func (s *Snapshot) Close() {
	newRefcount := atomic.AddInt32(&s.refCount, -1)
	if newRefcount == 0 {
		buf := s.db.snapshots.MakeBuf()
		defer s.db.snapshots.FreeBuf(buf)
		s.db.snapshots.Delete(s, CompareSnapshot, buf)
		if atomic.CompareAndSwapInt32(&s.db.isGCRunning, 0, 1) {
			go s.db.GC()
		}
	}
}

func CompareSnapshot(this skiplist.Item, that skiplist.Item) int {
	thisItem := this.(*Snapshot)
	thatItem := that.(*Snapshot)

	return int(thisItem.sn) - int(thatItem.sn)
}

func (m *MemDB) NewSnapshot() *Snapshot {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	snap := &Snapshot{db: m, sn: m.getCurrSn(), refCount: 1}
	m.snapshots.Insert(snap, CompareSnapshot, buf)
	atomic.AddUint32(&m.currSn, 1)
	return snap
}

type Iterator struct {
	snap *Snapshot
	iter *skiplist.Iterator
	buf  *skiplist.ActionBuffer
}

func (it *Iterator) skipUnwanted() {
loop:
	if !it.iter.Valid() {
		return
	}
	itm := it.iter.Get().(*Item)
	if itm.bornSn > it.snap.sn || (itm.deadSn > 0 && itm.deadSn <= it.snap.sn) {
		it.iter.Next()
		goto loop
	}
}

func (it *Iterator) SeekFirst() {
	it.iter.SeekFirst()
	it.skipUnwanted()
}

func (it *Iterator) Seek(itm *Item) {
	it.iter.Seek(itm)
	it.skipUnwanted()
}

func (it *Iterator) Valid() bool {
	return it.iter.Valid()
}

func (it *Iterator) Get() *Item {
	return it.iter.Get().(*Item)
}

func (it *Iterator) Next() {
	it.iter.Next()
	it.skipUnwanted()
}

func (it *Iterator) Close() {
	it.snap.Close()
	it.snap.db.store.FreeBuf(it.buf)
}

func (m *MemDB) NewIterator(snap *Snapshot) *Iterator {
	if !snap.Open() {
		return nil
	}
	buf := snap.db.store.MakeBuf()
	return &Iterator{
		snap: snap,
		iter: m.store.NewIterator(DataCompare, buf),
		buf:  buf,
	}
}

func (m *MemDB) collectDead(sn uint32) {
	buf1 := m.snapshots.MakeBuf()
	buf2 := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf1)
	defer m.snapshots.FreeBuf(buf2)
	iter := m.store.NewIterator(DataCompare, buf1)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		itm := iter.Get().(*Item)
		if itm.deadSn > 0 && itm.deadSn <= sn {
			m.store.Delete(itm, DataCompare, buf2)
		}
	}
}

func (m *MemDB) GC() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	if iter.Valid() {
		snap := iter.Get().(*Snapshot)
		if snap.sn != m.lastGCSn && snap.sn > 1 {
			m.lastGCSn = snap.sn - 1
			m.collectDead(m.lastGCSn)
		}
	}

	atomic.CompareAndSwapInt32(&m.isGCRunning, 1, 0)
}

func (m *MemDB) GetSnapshots() []*Snapshot {
	var snaps []*Snapshot
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		snaps = append(snaps, iter.Get().(*Snapshot))
	}

	return snaps
}

func (m *MemDB) StoreToDisk(dir string, snap *Snapshot) error {
	os.MkdirAll(dir, 0755)
	file, err := os.OpenFile(path.Join(dir, "records.data"), os.O_WRONLY|os.O_CREATE, 0755)
	if err != nil {
		return err
	}
	defer file.Close()

	w := bufio.NewWriterSize(file, DiskBlockSize)
	defer w.Flush()

	buf := make([]byte, 4)
	itr := m.NewIterator(snap)
	if itr == nil {
		return errors.New("Invalid snapshot")
	}
	defer itr.Close()

	if err := snap.Encode(buf, w); err != nil {
		return err
	}

	for itr.SeekFirst(); itr.Valid(); itr.Next() {
		itm := itr.Get()
		if err := itm.Encode(buf, w); err != nil {
			return err
		}
	}

	endItem := &Item{
		data: []byte(nil),
	}

	if err := endItem.Encode(buf, w); err != nil {
		return err
	}

	w.Flush()

	return nil
}

func (m *MemDB) LoadFromDisk(dir string) (*Snapshot, error) {
	var wg sync.WaitGroup
	file, err := os.OpenFile(path.Join(dir, "records.data"), os.O_RDONLY, 0755)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	buf := make([]byte, 4)
	r := bufio.NewReaderSize(file, DiskBlockSize)
	snap := &Snapshot{db: m, refCount: 1}
	if err := snap.Decode(buf, r); err != nil {
		return nil, err
	}
	m.currSn = snap.sn

	ch := make(chan *Item, 100000)
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go func(wg *sync.WaitGroup) {
			defer wg.Done()
			w := m.NewWriter()
			for itm := range ch {
				w.Put(itm)
			}
		}(&wg)
	}

loop:
	for {
		itm := &Item{}
		itm.Decode(buf, r)
		if len(itm.data) == 0 {
			break loop
		}

		ch <- itm
	}

	close(ch)
	wg.Wait()

	snbuf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(snbuf)
	m.snapshots.Insert(snap, CompareSnapshot, snbuf)
	m.currSn++

	return snap, nil
}
<|endoftext|>


<commit_before>package mstree

import (
	"io"
	"regexp"
	"strings"
)

type node struct {
	Children map[string]*node
}

func newNode() *node {
	return &node{make(map[string]*node)}
}

func (n *node) insert(tokens []string, inserted *bool) {
	if len(tokens) == 0 {
		return
	}
	first, tail := tokens[0], tokens[1:]
	child, ok := n.Children[first]
	if !ok {
		*inserted = true
		child = newNode()
		n.Children[first] = child
	}
	child.insert(tail, inserted)
}

func (n *node) traverseDump(prefix string, writer io.Writer) {
	if len(n.Children) == 0 {
		io.WriteString(writer, prefix+"\n")
	} else {
		for k, node := range n.Children {
			var nPref string
			if prefix == "" {
				nPref = k
			} else {
				nPref = prefix + "." + k
			}
			node.traverseDump(nPref, writer)
		}
	}
}

func (n *node) search(pattern string) map[string]*node {
	if pattern == "*" {
		return n.Children
	}

	results := make(map[string]*node)

	wcIndex := strings.Index(pattern, "*")
	if wcIndex == -1 {
		if node, ok := n.Children[pattern]; ok {
			results[pattern] = node
		}
		return results
	}
	lwcIndex := strings.LastIndex(pattern, "*")
	if wcIndex != lwcIndex || (wcIndex != 0 && wcIndex != len(pattern)-1) {
		// more than one wildcard or one wildcard in the middle
		re := regexp.MustCompile(strings.Replace(pattern, "*", ".*", -1))
		for k, node := range n.Children {
			if re.MatchString(k) {
				results[k] = node
			}
		}
		return results
	}

	if wcIndex == len(pattern)-1 {
		// wildcard at the end
		partial := pattern[:len(pattern)-1]
		for k, node := range n.Children {
			if strings.HasPrefix(k, partial) {
				results[k] = node
			}
		}
	} else {
		// wildcard at the begining
		partial := pattern[1:]
		for k, node := range n.Children {
			if strings.HasSuffix(k, partial) {
				results[k] = node
			}
		}
	}
	return results
}
<commit_msg>question wc added<commit_after>package mstree

import (
	"io"
	"regexp"
	"strings"
)

type node struct {
	Children map[string]*node
}

func newNode() *node {
	return &node{make(map[string]*node)}
}

func (n *node) insert(tokens []string, inserted *bool) {
	if len(tokens) == 0 {
		return
	}
	first, tail := tokens[0], tokens[1:]
	child, ok := n.Children[first]
	if !ok {
		*inserted = true
		child = newNode()
		n.Children[first] = child
	}
	child.insert(tail, inserted)
}

func (n *node) traverseDump(prefix string, writer io.Writer) {
	if len(n.Children) == 0 {
		io.WriteString(writer, prefix+"\n")
	} else {
		for k, node := range n.Children {
			var nPref string
			if prefix == "" {
				nPref = k
			} else {
				nPref = prefix + "." + k
			}
			node.traverseDump(nPref, writer)
		}
	}
}

func (n *node) search(pattern string) map[string]*node {
	if pattern == "*" {
		return n.Children
	}

	results := make(map[string]*node)

	wcIndex := strings.Index(pattern, "*")
	qIndex := strings.Index(pattern, "?")

	if wcIndex == -1 && qIndex == -1 {
		if node, ok := n.Children[pattern]; ok {
			results[pattern] = node
		}
		return results
	}

	if qIndex == -1 {
		// Only *
		lwcIndex := strings.LastIndex(pattern, "*")

		if wcIndex != lwcIndex || (wcIndex != 0 && wcIndex != len(pattern)-1) {
			// more than one wildcard or one wildcard in the middle
			re := regexp.MustCompile(strings.Replace(pattern, "*", ".*", -1))
			for k, node := range n.Children {
				if re.MatchString(k) {
					results[k] = node
				}
			}
			return results
		}

		if wcIndex == len(pattern)-1 {
			// wildcard at the end
			partial := pattern[:len(pattern)-1]
			for k, node := range n.Children {
				if strings.HasPrefix(k, partial) {
					results[k] = node
				}
			}
		} else {
			// wildcard at the begining
			partial := pattern[1:]
			for k, node := range n.Children {
				if strings.HasSuffix(k, partial) {
					results[k] = node
				}
			}
		}
	} else if wcIndex == -1 {
		// Only ?
		lqIndex := strings.LastIndex(pattern, "?")
		if qIndex != lqIndex || (qIndex != 0 && qIndex != len(pattern)-1) {
			// more than one ? or one ? in the middle
			re := regexp.MustCompile(strings.Replace(pattern, "?", ".?", -1))
			for k, node := range n.Children {
				if re.MatchString(k) {
					results[k] = node
				}
			}
			return results
		}

		if qIndex == len(pattern)-1 {
			// ? at the end
			partial := pattern[:len(pattern)-1]
			for k, node := range n.Children {
				if k[:len(k)-1] == partial {
					results[k] = node
				}
			}
		} else {
			// ? at the begining
			partial := pattern[1:]
			for k, node := range n.Children {
				if k[1:] == partial {
					results[k] = node
				}
			}
		}

	} else {
		// * and ? presents
		rePattern := strings.Replace(strings.Replace(pattern, "*", ".*", -1), "?", ".?", -1)
		re := regexp.MustCompile(rePattern)
		for k, node := range n.Children {
			if re.MatchString(k) {
				results[k] = node
			}
		}
	}

	return results
}
<|endoftext|>


<commit_before>package network

import (
	proto "code.google.com/p/goprotobuf/proto"
	"fmt"
	logPkg "log"
	"net"
	"os"
)

const (
	PacketBufferSize = 10
)

var log *logPkg.Logger

func init() {
	log = logPkg.New(os.Stdout, "network:", logPkg.LstdFlags)
}

type BitXConn struct {
	conn *net.UDPConn
	on   bool
	Receive <-chan BitXPacket
}

type BitXPacket struct {
	Addr   *net.UDPAddr
	Packet *Packet
}

func ListenUDP(ip string, port int) (*BitXConn, error) {
	addr := fmt.Sprintf("%v:%v", ip, port)
	laddr, errResolve := net.ResolveUDPAddr("udp", addr)
	if errResolve != nil {
		return nil, errResolve
	}
	pudp, errListen := net.ListenUDP("udp", laddr)
	if errListen != nil {
		return nil, errListen
	}
	return &BitXConn{conn:pudp, on:false}, nil
}

func (b *BitXConn) StartServerLoop() bool {
	if b.on || b.Receive != nil {
		return false
	}
	b.on = true
	log.Printf("Server started on:%s\n", b.conn.LocalAddr())
	receive := make(chan BitXPacket, PacketBufferSize)
	b.Receive = receive
	go func() {
		buf := make([]byte, 65507)
		for b.on {
			n, addr, err := b.conn.ReadFromUDP(buf)
			if err != nil {
				if b.on {
					log.Println("error ReadFromUDP:", err)
				}
				continue
			}
			log.Println("from:", addr, "length", n)
			p := new(Packet)
			err = proto.Unmarshal(buf[:n], p)
			if err != nil {
				log.Println("error Unmarshal:", err)
				continue
			}
			receive <- BitXPacket{addr, p}
		}
		close(receive)
	}()
	return true
}
func (b *BitXConn) Close() error {
	b.on = false
	err := b.conn.Close()
	if err == nil {
		log.Printf("Server stopped on:%s\n", b.conn.LocalAddr())
	}
	return err
}

func (b *BitXConn) Send(p *Packet, addr *net.UDPAddr) (int, error) {
	data, err := proto.Marshal(p)
	if err != nil {
		return 0, err
	}
	return b.conn.WriteToUDP(data, addr)
}
<commit_msg>go fmt<commit_after>package network

import (
	proto "code.google.com/p/goprotobuf/proto"
	"fmt"
	logPkg "log"
	"net"
	"os"
)

const (
	PacketBufferSize = 10
)

var log *logPkg.Logger

func init() {
	log = logPkg.New(os.Stdout, "network:", logPkg.LstdFlags)
}

type BitXConn struct {
	conn    *net.UDPConn
	on      bool
	Receive <-chan BitXPacket
}

type BitXPacket struct {
	Addr   *net.UDPAddr
	Packet *Packet
}

func ListenUDP(ip string, port int) (*BitXConn, error) {
	addr := fmt.Sprintf("%v:%v", ip, port)
	laddr, errResolve := net.ResolveUDPAddr("udp", addr)
	if errResolve != nil {
		return nil, errResolve
	}
	pudp, errListen := net.ListenUDP("udp", laddr)
	if errListen != nil {
		return nil, errListen
	}
	return &BitXConn{conn: pudp, on: false}, nil
}

func (b *BitXConn) StartServerLoop() bool {
	if b.on || b.Receive != nil {
		return false
	}
	b.on = true
	log.Printf("Server started on:%s\n", b.conn.LocalAddr())
	receive := make(chan BitXPacket, PacketBufferSize)
	b.Receive = receive
	go func() {
		buf := make([]byte, 65507)
		for b.on {
			n, addr, err := b.conn.ReadFromUDP(buf)
			if err != nil {
				if b.on {
					log.Println("error ReadFromUDP:", err)
				}
				continue
			}
			log.Println("from:", addr, "length", n)
			p := new(Packet)
			err = proto.Unmarshal(buf[:n], p)
			if err != nil {
				log.Println("error Unmarshal:", err)
				continue
			}
			receive <- BitXPacket{addr, p}
		}
		close(receive)
	}()
	return true
}
func (b *BitXConn) Close() error {
	b.on = false
	err := b.conn.Close()
	if err == nil {
		log.Printf("Server stopped on:%s\n", b.conn.LocalAddr())
	}
	return err
}

func (b *BitXConn) Send(p *Packet, addr *net.UDPAddr) (int, error) {
	data, err := proto.Marshal(p)
	if err != nil {
		return 0, err
	}
	return b.conn.WriteToUDP(data, addr)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package newton

import (
	"github.com/szabba/md/src/vect"
)

type Body struct {
	Xs, Vs []vect.Vector
	mass   float64
	currAt int
}

// Constructs a body of specified mass suitable for working with the integrator
func NewBody(algo Integrator, mass float64) *Body {

	b := new(Body)

	b.Xs = make([]vect.Vector, algo.StateLen())
	b.Vs = make([]vect.Vector, algo.StateLen())

	b.mass = mass
	b.currAt = algo.CurrentAt()

	return b
}

// Give a body's mass
func (b *Body) Mass() float64 {

	return b.mass
}

// Shifts all the values in xs by one and puts x at the beginning.
func Shift(xs []vect.Vector, x vect.Vector) {

	for i := len(xs) - 1; i > 0; i-- {

		xs[i] = xs[i-1]
	}

	xs[0] = x
}
<commit_msg>Add newton.Body.Shift<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package newton

import (
	"github.com/szabba/md/src/vect"
)

type Body struct {
	Xs, Vs []vect.Vector
	mass   float64
	currAt int
}

// Constructs a body of specified mass suitable for working with the integrator
func NewBody(algo Integrator, mass float64) *Body {

	b := new(Body)

	b.Xs = make([]vect.Vector, algo.StateLen())
	b.Vs = make([]vect.Vector, algo.StateLen())

	b.mass = mass
	b.currAt = algo.CurrentAt()

	return b
}

// Give a body's mass
func (b *Body) Mass() float64 {

	return b.mass
}

// Put new values of x and v a the beginning of the remembered values
//
// The oldest values get discarded
func (b *Body) Shift(x, v vect.Vector) {

	Shift(b.Xs, x)
	Shift(b.Vs, v)
}

// Shifts all the values in xs by one and puts x at the beginning.
func Shift(xs []vect.Vector, x vect.Vector) {

	for i := len(xs) - 1; i > 0; i-- {

		xs[i] = xs[i-1]
	}

	xs[0] = x
}
<|endoftext|>


<commit_before>package filesize

import (
	"fmt"
	"go/parser"
	"go/token"
)

func DirLines(dir string) error {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, dir, nil, 0)
	if err != nil {
		return err
	}

	for _, pkg := range pkgs {
		for fName, file := range pkg.Files {
			fmt.Printf("%s: %d\n", fName, fset.File(file.Pos()).LineCount())
		}
	}
	return nil
}

func FileLines(fName string) error {
	return nil
}
<commit_msg>Add more stuff<commit_after>package filesize

import (
	"fmt"
	"go/parser"
	"go/token"
	"path"
)

func DirsLines(dirs []string) error {
	for _, dir := range dirs {
		if err := DirLines(dir); err != nil {
			return err
		}
	}
	return nil
}

func DirLines(dir string) error {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, dir, nil, 0)
	if err != nil {
		return err
	}

	for _, pkg := range pkgs {
		for fName, file := range pkg.Files {
			fmt.Printf(
				"%s: %d\n",
				path.Base(fName),
				fset.File(file.Pos()).LineCount())
		}
	}
	return nil
}

func FileLines(fName string) error {
	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, fName, nil, 0)
	if err != nil {
		return err
	}
	fmt.Printf("%s: %d\n", path.Base(fName), fset.File(file.Pos()).LineCount())
	return nil
}
<|endoftext|>


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

package assert

import (
	"bytes"
	"fmt"
	"reflect"
	"strings"
)

// 判断一个值是否为空(0, "", false, 空数组等)。
// []string{""}空数组里套一个空字符串，不会被判断为空。
func isEmpty(expr interface{}) bool {
	if isZero(expr) {
		return true
	}

	rv := reflect.ValueOf(expr)
	for rv.Kind() == reflect.Pointer {
		rv = rv.Elem()
	}
	switch rv.Kind() {
	case reflect.Slice, reflect.Map, reflect.Array, reflect.Chan: // 长度为 0 的数组也是 empty
		return rv.Len() == 0
	default:
		return false
	}
}

func isZero(v interface{}) bool {
	if isNil(v) || reflect.ValueOf(v).IsZero() {
		return true
	}

	rv := reflect.ValueOf(v)
	for rv.Kind() == reflect.Pointer {
		rv = rv.Elem()
	}
	return rv.IsZero()
}

// isNil 判断一个值是否为 nil。
// 当特定类型的变量，已经声明，但还未赋值时，也将返回 true
func isNil(expr interface{}) bool {
	if nil == expr {
		return true
	}

	v := reflect.ValueOf(expr)
	k := v.Kind()

	return k >= reflect.Chan && k <= reflect.Slice && v.IsNil()
}

// 判断两个值是否相等
//
// 除了通过 reflect.DeepEqual() 判断值是否相等之外，一些类似
// 可转换的数值也能正确判断，比如以下值也将会被判断为相等：
//
//	int8(5)                     == int(5)
//	[]int{1,2}                  == []int8{1,2}
//	[]int{1,2}                  == [2]int8{1,2}
//	[]int{1,2}                  == []float32{1,2}
//	map[string]int{"1":"2":2}   == map[string]int8{"1":1,"2":2}
//
//	// map 的键值不同，即使可相互转换也判断不相等。
//	map[int]int{1:1,2:2}        != map[int8]int{1:1,2:2}
func isEqual(v1, v2 interface{}) bool {
	if reflect.DeepEqual(v1, v2) {
		return true
	}

	vv1 := reflect.ValueOf(v1)
	vv2 := reflect.ValueOf(v2)

	if !vv1.IsValid() || !vv2.IsValid() {
		return vv1.IsValid() == vv2.IsValid()
	}

	if vv1 == vv2 {
		return true
	}

	vv1Type := vv1.Type()
	vv2Type := vv2.Type()

	if vv1Type.Comparable() && vv2Type.Comparable() && v1 == v2 {
		return true
	}

	// 过滤掉已经在 reflect.DeepEqual() 进行处理的类型
	switch vv1Type.Kind() {
	case reflect.Struct, reflect.Ptr, reflect.Func, reflect.Interface:
		return false
	case reflect.Slice, reflect.Array:
		// vv2.Kind() 与 vv1 的不相同
		if vv2.Kind() != reflect.Slice && vv2.Kind() != reflect.Array {
			// 虽然类型不同，但可以相互转换成 vv1 的，如：vv2 是 string，vv2 是 []byte，
			if vv2Type.ConvertibleTo(vv1Type) {
				return isEqual(vv1.Interface(), vv2.Convert(vv1Type).Interface())
			}
			return false
		}

		// reflect.DeepEqual() 未考虑类型不同但是类型可转换的情况，比如：
		// []int{8,9} == []int8{8,9}，此处重新对 slice 和 array 做比较处理。
		if vv1.Len() != vv2.Len() {
			return false
		}

		for i := 0; i < vv1.Len(); i++ {
			if !isEqual(vv1.Index(i).Interface(), vv2.Index(i).Interface()) {
				return false
			}
		}
		return true // for 中所有的值比较都相等，返回 true
	case reflect.Map:
		if vv2.Kind() != reflect.Map {
			return false
		}

		if vv1.IsNil() != vv2.IsNil() {
			return false
		}
		if vv1.Len() != vv2.Len() {
			return false
		}
		if vv1.Pointer() == vv2.Pointer() {
			return true
		}

		// 两个 map 的键名类型不同
		if vv2Type.Key().Kind() != vv1Type.Key().Kind() {
			return false
		}

		for _, index := range vv1.MapKeys() {
			vv2Index := vv2.MapIndex(index)
			if !vv2Index.IsValid() {
				return false
			}

			if !isEqual(vv1.MapIndex(index).Interface(), vv2Index.Interface()) {
				return false
			}
		}
		return true // for 中所有的值比较都相等，返回 true
	case reflect.String:
		if vv2.Kind() == reflect.String {
			return vv1.String() == vv2.String()
		}
		if vv2Type.ConvertibleTo(vv1Type) { // 考虑 v1 是 string，v2 是 []byte 的情况
			return isEqual(vv1.Interface(), vv2.Convert(vv1Type).Interface())
		}

		return false
	}

	if vv1Type.ConvertibleTo(vv2Type) {
		return vv2.Interface() == vv1.Convert(vv2Type).Interface()
	} else if vv2Type.ConvertibleTo(vv1Type) {
		return vv1.Interface() == vv2.Convert(vv1Type).Interface()
	}

	return false
}

// hasPanic 判断 fn 函数是否会发生 panic
// 若发生了 panic，将把 msg 一起返回。
func hasPanic(fn func()) (has bool, msg interface{}) {
	defer func() {
		if msg = recover(); msg != nil {
			has = true
		}
	}()
	fn()

	return
}

// isContains 判断 container 是否包含了 item 的内容。若是指针，会判断指针指向的内容
func isContains(container, item interface{}) bool {
	if container == nil { // nil不包含任何东西
		return false
	}

	cv := reflect.ValueOf(container)
	iv := reflect.ValueOf(item)

	for cv.Kind() == reflect.Ptr {
		cv = cv.Elem()
	}

	for iv.Kind() == reflect.Ptr {
		iv = iv.Elem()
	}

	if isEqual(container, item) {
		return true
	}

	// 判断是字符串的情况
	switch c := cv.Interface().(type) {
	case string:
		switch i := iv.Interface().(type) {
		case string:
			return strings.Contains(c, i)
		case []byte:
			return strings.Contains(c, string(i))
		case []rune:
			return strings.Contains(c, string(i))
		case byte:
			return bytes.IndexByte([]byte(c), i) != -1
		case rune:
			return bytes.ContainsRune([]byte(c), i)
		}
	case []byte:
		switch i := iv.Interface().(type) {
		case string:
			return bytes.Contains(c, []byte(i))
		case []byte:
			return bytes.Contains(c, i)
		case []rune:
			return strings.Contains(string(c), string(i))
		case byte:
			return bytes.IndexByte(c, i) != -1
		case rune:
			return bytes.ContainsRune(c, i)
		}
	case []rune:
		switch i := iv.Interface().(type) {
		case string:
			return strings.Contains(string(c), i)
		case []byte:
			return strings.Contains(string(c), string(i))
		case []rune:
			return strings.Contains(string(c), string(i))
		case byte:
			return strings.IndexByte(string(c), i) != -1
		case rune:
			return strings.ContainsRune(string(c), i)
		}
	}

	if (cv.Kind() == reflect.Slice) || (cv.Kind() == reflect.Array) {
		if !cv.IsValid() || cv.Len() == 0 { // 空的，就不算包含另一个，即使另一个也是空值。
			return false
		}

		if !iv.IsValid() {
			return false
		}

		// item 是 container 的一个元素
		for i := 0; i < cv.Len(); i++ {
			if isEqual(cv.Index(i).Interface(), iv.Interface()) {
				return true
			}
		}

		// 开始判断 item 的元素是否与 container 中的元素相等。

		// 若 item 的长度为 0，表示不包含
		if (iv.Kind() != reflect.Slice) || (iv.Len() == 0) {
			return false
		}

		// item 的元素比 container 的元素多
		if iv.Len() > cv.Len() {
			return false
		}

		// 依次比较 item 的各个子元素是否都存在于 container，且下标都相同
		ivIndex := 0
		for i := 0; i < cv.Len(); i++ {
			if isEqual(cv.Index(i).Interface(), iv.Index(ivIndex).Interface()) {
				if (ivIndex == 0) && (i+iv.Len() > cv.Len()) {
					return false
				}
				ivIndex++
				if ivIndex == iv.Len() { // 已经遍历完 iv
					return true
				}
			} else if ivIndex > 0 {
				return false
			}
		}
		return false
	} // end cv.Kind == reflect.Slice and reflect.Array

	if cv.Kind() == reflect.Map {
		if cv.Len() == 0 {
			return false
		}

		if (iv.Kind() != reflect.Map) || (iv.Len() == 0) {
			return false
		}

		if iv.Len() > cv.Len() {
			return false
		}

		// 判断所有 item 的项都存在于 container 中
		for _, key := range iv.MapKeys() {
			cvItem := cv.MapIndex(key)
			if !cvItem.IsValid() { // container 中不包含该值。
				return false
			}
			if !isEqual(cvItem.Interface(), iv.MapIndex(key).Interface()) {
				return false
			}
		}
		// for 中的所有判断都成立，返回 true
		return true
	}

	return false
}

func getLen(v interface{}) (l int, msg string) {
	r := reflect.ValueOf(v)
	for r.Kind() == reflect.Ptr {
		r = r.Elem()
	}

	if v == nil {
		return 0, ""
	}

	switch r.Kind() {
	case reflect.Array, reflect.String, reflect.Slice, reflect.Map:
		return r.Len(), ""
	}
	return 0, fmt.Sprintf("无法获取 %s 类型的长度信息", r.Kind())
}

func getType(ptr bool, v1, v2 interface{}) (t1, t2 reflect.Type) {
	t1 = reflect.TypeOf(v1)
	t2 = reflect.TypeOf(v2)

	if ptr {
		for t1.Kind() == reflect.Ptr {
			t1 = t1.Elem()
		}
		for t2.Kind() == reflect.Ptr {
			t2 = t2.Elem()
		}
	}

	return
}
<commit_msg>fix: 将 reflect.Pointer 改为 reflect.Ptr<commit_after>// SPDX-License-Identifier: MIT

package assert

import (
	"bytes"
	"fmt"
	"reflect"
	"strings"
)

// 判断一个值是否为空(0, "", false, 空数组等)。
// []string{""}空数组里套一个空字符串，不会被判断为空。
func isEmpty(expr interface{}) bool {
	if isZero(expr) {
		return true
	}

	rv := reflect.ValueOf(expr)
	for rv.Kind() == reflect.Ptr {
		rv = rv.Elem()
	}
	switch rv.Kind() {
	case reflect.Slice, reflect.Map, reflect.Array, reflect.Chan: // 长度为 0 的数组也是 empty
		return rv.Len() == 0
	default:
		return false
	}
}

func isZero(v interface{}) bool {
	if isNil(v) || reflect.ValueOf(v).IsZero() {
		return true
	}

	rv := reflect.ValueOf(v)
	for rv.Kind() == reflect.Ptr {
		rv = rv.Elem()
	}
	return rv.IsZero()
}

// isNil 判断一个值是否为 nil。
// 当特定类型的变量，已经声明，但还未赋值时，也将返回 true
func isNil(expr interface{}) bool {
	if nil == expr {
		return true
	}

	v := reflect.ValueOf(expr)
	k := v.Kind()

	return k >= reflect.Chan && k <= reflect.Slice && v.IsNil()
}

// 判断两个值是否相等
//
// 除了通过 reflect.DeepEqual() 判断值是否相等之外，一些类似
// 可转换的数值也能正确判断，比如以下值也将会被判断为相等：
//
//	int8(5)                     == int(5)
//	[]int{1,2}                  == []int8{1,2}
//	[]int{1,2}                  == [2]int8{1,2}
//	[]int{1,2}                  == []float32{1,2}
//	map[string]int{"1":"2":2}   == map[string]int8{"1":1,"2":2}
//
//	// map 的键值不同，即使可相互转换也判断不相等。
//	map[int]int{1:1,2:2}        != map[int8]int{1:1,2:2}
func isEqual(v1, v2 interface{}) bool {
	if reflect.DeepEqual(v1, v2) {
		return true
	}

	vv1 := reflect.ValueOf(v1)
	vv2 := reflect.ValueOf(v2)

	if !vv1.IsValid() || !vv2.IsValid() {
		return vv1.IsValid() == vv2.IsValid()
	}

	if vv1 == vv2 {
		return true
	}

	vv1Type := vv1.Type()
	vv2Type := vv2.Type()

	if vv1Type.Comparable() && vv2Type.Comparable() && v1 == v2 {
		return true
	}

	// 过滤掉已经在 reflect.DeepEqual() 进行处理的类型
	switch vv1Type.Kind() {
	case reflect.Struct, reflect.Ptr, reflect.Func, reflect.Interface:
		return false
	case reflect.Slice, reflect.Array:
		// vv2.Kind() 与 vv1 的不相同
		if vv2.Kind() != reflect.Slice && vv2.Kind() != reflect.Array {
			// 虽然类型不同，但可以相互转换成 vv1 的，如：vv2 是 string，vv2 是 []byte，
			if vv2Type.ConvertibleTo(vv1Type) {
				return isEqual(vv1.Interface(), vv2.Convert(vv1Type).Interface())
			}
			return false
		}

		// reflect.DeepEqual() 未考虑类型不同但是类型可转换的情况，比如：
		// []int{8,9} == []int8{8,9}，此处重新对 slice 和 array 做比较处理。
		if vv1.Len() != vv2.Len() {
			return false
		}

		for i := 0; i < vv1.Len(); i++ {
			if !isEqual(vv1.Index(i).Interface(), vv2.Index(i).Interface()) {
				return false
			}
		}
		return true // for 中所有的值比较都相等，返回 true
	case reflect.Map:
		if vv2.Kind() != reflect.Map {
			return false
		}

		if vv1.IsNil() != vv2.IsNil() {
			return false
		}
		if vv1.Len() != vv2.Len() {
			return false
		}
		if vv1.Pointer() == vv2.Pointer() {
			return true
		}

		// 两个 map 的键名类型不同
		if vv2Type.Key().Kind() != vv1Type.Key().Kind() {
			return false
		}

		for _, index := range vv1.MapKeys() {
			vv2Index := vv2.MapIndex(index)
			if !vv2Index.IsValid() {
				return false
			}

			if !isEqual(vv1.MapIndex(index).Interface(), vv2Index.Interface()) {
				return false
			}
		}
		return true // for 中所有的值比较都相等，返回 true
	case reflect.String:
		if vv2.Kind() == reflect.String {
			return vv1.String() == vv2.String()
		}
		if vv2Type.ConvertibleTo(vv1Type) { // 考虑 v1 是 string，v2 是 []byte 的情况
			return isEqual(vv1.Interface(), vv2.Convert(vv1Type).Interface())
		}

		return false
	}

	if vv1Type.ConvertibleTo(vv2Type) {
		return vv2.Interface() == vv1.Convert(vv2Type).Interface()
	} else if vv2Type.ConvertibleTo(vv1Type) {
		return vv1.Interface() == vv2.Convert(vv1Type).Interface()
	}

	return false
}

// hasPanic 判断 fn 函数是否会发生 panic
// 若发生了 panic，将把 msg 一起返回。
func hasPanic(fn func()) (has bool, msg interface{}) {
	defer func() {
		if msg = recover(); msg != nil {
			has = true
		}
	}()
	fn()

	return
}

// isContains 判断 container 是否包含了 item 的内容。若是指针，会判断指针指向的内容
func isContains(container, item interface{}) bool {
	if container == nil { // nil不包含任何东西
		return false
	}

	cv := reflect.ValueOf(container)
	iv := reflect.ValueOf(item)

	for cv.Kind() == reflect.Ptr {
		cv = cv.Elem()
	}

	for iv.Kind() == reflect.Ptr {
		iv = iv.Elem()
	}

	if isEqual(container, item) {
		return true
	}

	// 判断是字符串的情况
	switch c := cv.Interface().(type) {
	case string:
		switch i := iv.Interface().(type) {
		case string:
			return strings.Contains(c, i)
		case []byte:
			return strings.Contains(c, string(i))
		case []rune:
			return strings.Contains(c, string(i))
		case byte:
			return bytes.IndexByte([]byte(c), i) != -1
		case rune:
			return bytes.ContainsRune([]byte(c), i)
		}
	case []byte:
		switch i := iv.Interface().(type) {
		case string:
			return bytes.Contains(c, []byte(i))
		case []byte:
			return bytes.Contains(c, i)
		case []rune:
			return strings.Contains(string(c), string(i))
		case byte:
			return bytes.IndexByte(c, i) != -1
		case rune:
			return bytes.ContainsRune(c, i)
		}
	case []rune:
		switch i := iv.Interface().(type) {
		case string:
			return strings.Contains(string(c), i)
		case []byte:
			return strings.Contains(string(c), string(i))
		case []rune:
			return strings.Contains(string(c), string(i))
		case byte:
			return strings.IndexByte(string(c), i) != -1
		case rune:
			return strings.ContainsRune(string(c), i)
		}
	}

	if (cv.Kind() == reflect.Slice) || (cv.Kind() == reflect.Array) {
		if !cv.IsValid() || cv.Len() == 0 { // 空的，就不算包含另一个，即使另一个也是空值。
			return false
		}

		if !iv.IsValid() {
			return false
		}

		// item 是 container 的一个元素
		for i := 0; i < cv.Len(); i++ {
			if isEqual(cv.Index(i).Interface(), iv.Interface()) {
				return true
			}
		}

		// 开始判断 item 的元素是否与 container 中的元素相等。

		// 若 item 的长度为 0，表示不包含
		if (iv.Kind() != reflect.Slice) || (iv.Len() == 0) {
			return false
		}

		// item 的元素比 container 的元素多
		if iv.Len() > cv.Len() {
			return false
		}

		// 依次比较 item 的各个子元素是否都存在于 container，且下标都相同
		ivIndex := 0
		for i := 0; i < cv.Len(); i++ {
			if isEqual(cv.Index(i).Interface(), iv.Index(ivIndex).Interface()) {
				if (ivIndex == 0) && (i+iv.Len() > cv.Len()) {
					return false
				}
				ivIndex++
				if ivIndex == iv.Len() { // 已经遍历完 iv
					return true
				}
			} else if ivIndex > 0 {
				return false
			}
		}
		return false
	} // end cv.Kind == reflect.Slice and reflect.Array

	if cv.Kind() == reflect.Map {
		if cv.Len() == 0 {
			return false
		}

		if (iv.Kind() != reflect.Map) || (iv.Len() == 0) {
			return false
		}

		if iv.Len() > cv.Len() {
			return false
		}

		// 判断所有 item 的项都存在于 container 中
		for _, key := range iv.MapKeys() {
			cvItem := cv.MapIndex(key)
			if !cvItem.IsValid() { // container 中不包含该值。
				return false
			}
			if !isEqual(cvItem.Interface(), iv.MapIndex(key).Interface()) {
				return false
			}
		}
		// for 中的所有判断都成立，返回 true
		return true
	}

	return false
}

func getLen(v interface{}) (l int, msg string) {
	r := reflect.ValueOf(v)
	for r.Kind() == reflect.Ptr {
		r = r.Elem()
	}

	if v == nil {
		return 0, ""
	}

	switch r.Kind() {
	case reflect.Array, reflect.String, reflect.Slice, reflect.Map:
		return r.Len(), ""
	}
	return 0, fmt.Sprintf("无法获取 %s 类型的长度信息", r.Kind())
}

func getType(ptr bool, v1, v2 interface{}) (t1, t2 reflect.Type) {
	t1 = reflect.TypeOf(v1)
	t2 = reflect.TypeOf(v2)

	if ptr {
		for t1.Kind() == reflect.Ptr {
			t1 = t1.Elem()
		}
		for t2.Kind() == reflect.Ptr {
			t2 = t2.Elem()
		}
	}

	return
}
<|endoftext|>


<commit_before>package psn

import (
	"os"
	"strconv"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
)

func isInt(s string) bool {
	_, err := strconv.Atoi(s)
	return err == nil
}

func humanizeDurationMs(sec uint64) string {
	s := humanize.Time(time.Now().Add(-1 * time.Duration(sec) * time.Millisecond))
	if s == "now" {
		s = "0"
	}
	return strings.TrimSpace(strings.Replace(s, " ago", "", -1))
}

func humanizeDurationSecond(sec uint64) string {
	s := humanize.Time(time.Now().Add(-1 * time.Duration(sec) * time.Second))
	if s == "now" {
		s = "0"
	}
	return strings.TrimSpace(strings.Replace(s, " ago", "", -1))
}

func openToRead(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDONLY, 0444)
	if err != nil {
		return nil, err
	}
	return f, nil
}

func openToAppend(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0777)
	if err != nil {
		return nil, err
	}
	return f, nil
}

func openToOverwrite(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0777)
	if err != nil {
		return nil, err
	}
	return f, nil
}
<commit_msg>psn: update duration utils<commit_after>package psn

import (
	"os"
	"strconv"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
)

func isInt(s string) bool {
	_, err := strconv.Atoi(s)
	return err == nil
}

func humanizeDurationMs(ms uint64) string {
	s := humanize.Time(time.Now().Add(-1 * time.Duration(ms) * time.Millisecond))
	if s == "now" {
		s = "0 second"
	}
	return strings.TrimSpace(strings.Replace(s, " ago", "", -1))
}

func humanizeDurationSecond(sec uint64) string {
	s := humanize.Time(time.Now().Add(-1 * time.Duration(sec) * time.Second))
	if s == "now" {
		s = "0 second"
	}
	return strings.TrimSpace(strings.Replace(s, " ago", "", -1))
}

func openToRead(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDONLY, 0444)
	if err != nil {
		return nil, err
	}
	return f, nil
}

func openToAppend(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0777)
	if err != nil {
		return nil, err
	}
	return f, nil
}

func openToOverwrite(fpath string) (*os.File, error) {
	f, err := os.OpenFile(fpath, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0777)
	if err != nil {
		return nil, err
	}
	return f, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
)

// fmt.Printf to stdout. Convention is to outf info intended for scripting.
func outf(f string, v ...interface{}) {
	fmt.Printf(f, v...)
}

// fmt.Printf to stderr. Convention is to errf info intended for human.
func errf(f string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, f, v...)
}

// Verbose output for debugging.
func logf(fmt string, v ...interface{}) {
	log.Printf(fmt, v...)
}

// Check if the connection to tcp://addr is readable.
func read(addr string) error {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	if _, err = conn.Read(make([]byte, 1)); err != nil {
		return err
	}
	return nil
}

// Check if an addr can be successfully connected.
func ping(addr string) bool {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return false
	}
	defer conn.Close()
	return true
}

// Download the url to the dest path.
func download(dest, url string) error {
	rsp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer rsp.Body.Close()

	// Create the dest dir.
	if err := os.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}

	f, err := os.Create(fmt.Sprintf("%s.download", dest))
	if err != nil {
		return err
	}
	defer os.Remove(f.Name())

	if _, err := io.Copy(f, rsp.Body); err != nil {
		// TODO: display download progress?
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	if err := os.Rename(f.Name(), dest); err != nil {
		return err
	}
	return nil
}

// Get latest release tag name (e.g. "v0.6.0") from a repo on GitHub.
func getLatestReleaseName(url string) (string, error) {
	rsp, err := http.Get(url)
	if err != nil {
		return "", err
	}
	defer rsp.Body.Close()

	var t []struct {
		TagName string `json:"tag_name"`
	}
	if err := json.NewDecoder(rsp.Body).Decode(&t); err != nil {
		return "", err
	}
	if len(t) == 0 {
		return "", fmt.Errorf("no releases found")
	}
	return t[0].TagName, nil
}

// Convenient function to exec a command.
func cmd(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	if *verbose {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<commit_msg>Print command executed in verbose mode<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// fmt.Printf to stdout. Convention is to outf info intended for scripting.
func outf(f string, v ...interface{}) {
	fmt.Printf(f, v...)
}

// fmt.Printf to stderr. Convention is to errf info intended for human.
func errf(f string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, f, v...)
}

// Verbose output for debugging.
func logf(fmt string, v ...interface{}) {
	log.Printf(fmt, v...)
}

// Check if the connection to tcp://addr is readable.
func read(addr string) error {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	if _, err = conn.Read(make([]byte, 1)); err != nil {
		return err
	}
	return nil
}

// Check if an addr can be successfully connected.
func ping(addr string) bool {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return false
	}
	defer conn.Close()
	return true
}

// Download the url to the dest path.
func download(dest, url string) error {
	rsp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer rsp.Body.Close()

	// Create the dest dir.
	if err := os.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}

	f, err := os.Create(fmt.Sprintf("%s.download", dest))
	if err != nil {
		return err
	}
	defer os.Remove(f.Name())

	if _, err := io.Copy(f, rsp.Body); err != nil {
		// TODO: display download progress?
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	if err := os.Rename(f.Name(), dest); err != nil {
		return err
	}
	return nil
}

// Get latest release tag name (e.g. "v0.6.0") from a repo on GitHub.
func getLatestReleaseName(url string) (string, error) {
	rsp, err := http.Get(url)
	if err != nil {
		return "", err
	}
	defer rsp.Body.Close()

	var t []struct {
		TagName string `json:"tag_name"`
	}
	if err := json.NewDecoder(rsp.Body).Decode(&t); err != nil {
		return "", err
	}
	if len(t) == 0 {
		return "", fmt.Errorf("no releases found")
	}
	return t[0].TagName, nil
}

// Convenient function to exec a command.
func cmd(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	if *verbose {
		logf("executing: %v %v", name, strings.Join(args, " "))
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<|endoftext|>