<commit_before>package main

import (
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"regexp"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/buger/jsonparser"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	uuid "github.com/satori/go.uuid"
	"github.com/syhlion/greq"
	"github.com/syhlion/redisocket.v2"
)

var DefaultSubHandler = func(channel string, p *redisocket.Payload) (err error) {
	return nil
}

type commandResponse struct {
	cmdType string
	handler func(string, *redisocket.Payload) (err error)
	msg     []byte
	data    string
}

func Ping() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("pong"))
	}
}
func WsAuth(sc SlaveConfig, pool *redis.Pool, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		jwt := r.FormValue("jwt")

		v := url.Values{}
		v.Add("data", jwt)

		b, _, err := reqClient.Post(sc.DecodeServiceAddr, v)
		if err != nil {
			logger.WithError(err).Warn("post decode service error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		a := &JwtPack{}
		err = json.Unmarshal(b, a)
		if err != nil {
			logger.WithError(err).Warn("json marshal error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		b, err = json.Marshal(a.Gusher)
		if err != nil {
			logger.WithError(err).Warn("json marshl error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		uid := uuid.NewV1()
		conn := pool.Get()
		defer conn.Close()
		conn.Send("SET", uid.String(), string(b))
		conn.Send("EXPIRE", uid.String(), 60)
		conn.Flush()
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(struct {
			Token string `json:"token"`
		}{
			Token: uid.String(),
		})
		return
	}

}
func WtfConnect(sc SlaveConfig, pool *redis.Pool, rHub *redisocket.Hub, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		appKey := params["app_key"]
		token := r.FormValue("token")
		if appKey == "" || token == "" {
			logger.Warn("app_key or token is nil")
			http.Error(w, "app_key is nil", http.StatusUnauthorized)
			return
		}
		conn := pool.Get()
		reply, err := redis.Bytes(conn.Do("GET", token))
		if err != nil {
			conn.Close()
			logger.WithError(err).Warn("redis get error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		conn.Close()
		auth := Auth{}
		err = json.Unmarshal(reply, &auth)
		if err != nil {
			logger.WithError(err).Warn("json unmarshal error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		if appKey != auth.AppKey {
			http.Error(w, "appkey error", http.StatusUnauthorized)
			return
		}

		s, err := rHub.Upgrade(w, r, nil, auth.UserId, appKey)
		if err != nil {
			logger.WithError(err).Warnf("upgrade ws connection error")
			return
		}
		defer s.Close()

		s.Listen(func(data []byte) (b []byte, err error) {
			logger.WithFields(logrus.Fields{
				"data": string(data),
			}).Info("receive start")
			h, err := CommanRouter(data, pool)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data": string(data),
				}).WithError(err).Warn("command router error")
				return
			}

			d, _, _, err := jsonparser.Get(data, "data")
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data": string(data),
				}).WithError(err).Warn("jsonparser data error")
				return
			}
			logger.WithFields(logrus.Fields{
				"data":  string(data),
				"pdata": string(d),
			}).Info("receive to sub")
			res, err := h(appKey, auth, d)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data":  string(data),
					"pdata": string(d),
					"res":   res,
				}).WithError(err).Warn("sub error")
				return
			}
			switch res.cmdType {
			case "SUB":
				s.On(res.data, res.handler)
			case "UNSUB":
				s.Off(res.data)

			}
			logger.WithFields(logrus.Fields{
				"data":  string(data),
				"pdata": string(d),
				"res":   res,
			}).Info("receive to sub")
			return res.msg, nil

		})
		return
	}

}

func WsConnect(sc SlaveConfig, pool *redis.Pool, rHub *redisocket.Hub, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		appKey := params["app_key"]
		token := r.FormValue("token")
		if appKey == "" || token == "" {
			logger.Warn("app_key or token is nil")
			http.Error(w, "app_key is nil", http.StatusUnauthorized)
			return
		}
		conn := pool.Get()
		reply, err := redis.Bytes(conn.Do("GET", token))
		if err != nil {
			conn.Close()
			logger.WithError(err).Warn("redis get error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		conn.Close()
		auth := Auth{}
		err = json.Unmarshal(reply, &auth)
		if err != nil {
			logger.WithError(err).Warn("json unmarshal error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		if appKey != auth.AppKey {
			http.Error(w, "appkey error", http.StatusUnauthorized)
			return
		}

		s, err := rHub.Upgrade(w, r, nil, auth.UserId, appKey)
		if err != nil {
			logger.WithError(err).Warnf("upgrade ws connection error")
			return
		}
		defer s.Close()

		s.Listen(func(data []byte) (b []byte, err error) {
			h, err := CommanRouter(data, pool)
			if err != nil {
				return
			}

			d, _, _, err := jsonparser.Get(data, "data")
			if err != nil {
				return
			}
			res, err := h(appKey, auth, d)
			if err != nil {
				return
			}
			switch res.cmdType {
			case "SUB":
				s.On(res.data, res.handler)
			case "UNSUB":
				s.Off(res.data)

			}
			return res.msg, nil
		})
		return
	}

}

func SubscribeCommand(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {

	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	msg = &commandResponse{
		handler: DefaultSubHandler,
		cmdType: "SUB",
	}
	command := &ChannelCommand{}
	exist := false
	for _, ch := range auth.Channels {
		//新增萬用字元  如果找到這個 任何頻道皆可訂閱
		if ch == "*" {
			exist = true
			break
		}
		ech := regexp.QuoteMeta(ch)
		rch := strings.Replace(ech, `\*`, "*", -1)
		r, err := regexp.Compile(rch)
		if err != nil {
			break
		}
		if r.MatchString(channel) {
			exist = true
			break
		}
	}
	var reply []byte
	if exist {
		msg.data = channel
		command.Event = SubscribeReplySucceeded
		command.Data.Channel = channel
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	} else {
		command.Event = SubscribeReplyError
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply

	}

	return
}
func Remote(pool *redis.Pool) func(string, Auth, []byte) (msg *commandResponse, err error) {
	return func(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {

		remote, err := jsonparser.GetString(data, "remote")
		if err != nil {
			return
		}
		uid, err := jsonparser.GetString(data, "uid")
		if err != nil {
			return
		}
		payload, err := jsonparser.GetString(data, "payload")
		if err != nil {
			return
		}
		p := JsonCheck(payload)
		msg = &commandResponse{
			cmdType: "REMOTE",
		}
		var reply []byte
		command := &RemoteCommand{}
		command.Data.Remote = remote
		b, ok := auth.Remotes[remote]

		//沒有這個remote 返回錯誤訊息不斷線
		if !ok || !b {
			command.Event = RemoteReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				return
			}
			msg.msg = reply
			return
		}
		wp := WorkerPayload{
			UserId: auth.UserId,
			Data:   p,
			Uid:    uid,
			AppKey: auth.AppKey,
		}
		d, err := json.Marshal(wp)
		conn := pool.Get()
		_, err = conn.Do("RPUSH", auth.AppKey+"@"+remote, d)
		if err != nil {
			return
		}
		msg.msg = reply
		return
	}

}
func UnSubscribeCommand(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {
	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	exist := false
	for _, ch := range auth.Channels {
		//新增萬用字元  如果找到這個 任何頻道皆可訂閱
		if ch == "*" {
			exist = true
			break
		}
		ech := regexp.QuoteMeta(ch)
		rch := strings.Replace(ech, `\*`, "*", -1)
		r, err := regexp.Compile(rch)
		if err != nil {
			break
		}
		if r.MatchString(channel) {
			exist = true
			break
		}
	}
	msg = &commandResponse{
		cmdType: "UNSUB",
	}
	command := &ChannelCommand{}
	var reply []byte
	//反訂閱處理
	if exist {
		msg.data = channel
		command.Event = UnSubscribeReplySucceeded
		command.Data.Channel = channel
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	} else {
		command.Event = UnSubscribeReplyError
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	}
	return
}

func CommanRouter(data []byte, pool *redis.Pool) (fn func(appkey string, auth Auth, data []byte) (msg *commandResponse, err error), err error) {

	val, err := jsonparser.GetString(data, "event")
	if err != nil {
		return
	}
	switch val {
	case RemoteEvent:
		return Remote(pool), nil
	case SubscribeEvent:
		return SubscribeCommand, nil
	case UnSubscribeEvent:
		return UnSubscribeCommand, nil
	default:
		err = errors.New("event errors")
		break
	}
	return
}
<commit_msg>fix json bug<commit_after>package main

import (
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"regexp"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/buger/jsonparser"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	uuid "github.com/satori/go.uuid"
	"github.com/syhlion/greq"
	"github.com/syhlion/redisocket.v2"
)

var DefaultSubHandler = func(channel string, p *redisocket.Payload) (err error) {
	return nil
}

type commandResponse struct {
	cmdType string
	handler func(string, *redisocket.Payload) (err error)
	msg     []byte
	data    string
}

func Ping() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("pong"))
	}
}
func WsAuth(sc SlaveConfig, pool *redis.Pool, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		jwt := r.FormValue("jwt")

		v := url.Values{}
		v.Add("data", jwt)

		b, _, err := reqClient.Post(sc.DecodeServiceAddr, v)
		if err != nil {
			logger.WithError(err).Warn("post decode service error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		a := &JwtPack{}
		err = json.Unmarshal(b, a)
		if err != nil {
			logger.WithError(err).Warn("json marshal error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		b, err = json.Marshal(a.Gusher)
		if err != nil {
			logger.WithError(err).Warn("json marshl error")
			http.Error(w, "jwt decode fail", http.StatusUnauthorized)
			return
		}
		uid := uuid.NewV1()
		conn := pool.Get()
		defer conn.Close()
		conn.Send("SET", uid.String(), string(b))
		conn.Send("EXPIRE", uid.String(), 60)
		conn.Flush()
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(struct {
			Token string `json:"token"`
		}{
			Token: uid.String(),
		})
		return
	}

}
func WtfConnect(sc SlaveConfig, pool *redis.Pool, rHub *redisocket.Hub, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		appKey := params["app_key"]
		token := r.FormValue("token")
		if appKey == "" || token == "" {
			logger.Warn("app_key or token is nil")
			http.Error(w, "app_key is nil", http.StatusUnauthorized)
			return
		}
		conn := pool.Get()
		reply, err := redis.Bytes(conn.Do("GET", token))
		if err != nil {
			conn.Close()
			logger.WithError(err).Warn("redis get error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		conn.Close()
		auth := Auth{}
		err = json.Unmarshal(reply, &auth)
		if err != nil {
			logger.WithError(err).Warn("json unmarshal error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		if appKey != auth.AppKey {
			http.Error(w, "appkey error", http.StatusUnauthorized)
			return
		}

		s, err := rHub.Upgrade(w, r, nil, auth.UserId, appKey)
		if err != nil {
			logger.WithError(err).Warnf("upgrade ws connection error")
			return
		}
		defer s.Close()

		s.Listen(func(data []byte) (b []byte, err error) {
			logger.WithFields(logrus.Fields{
				"data": string(data),
			}).Info("receive start")
			h, err := CommanRouter(data, pool)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data": string(data),
				}).WithError(err).Warn("command router error")
				return
			}

			d, _, _, err := jsonparser.Get(data, "data")
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data": string(data),
				}).WithError(err).Warn("jsonparser data error")
				return
			}
			logger.WithFields(logrus.Fields{
				"data":  string(data),
				"pdata": string(d),
			}).Info("receive to sub")
			res, err := h(appKey, auth, d)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"data":  string(data),
					"pdata": string(d),
					"res":   res,
				}).WithError(err).Warn("sub error")
				return
			}
			switch res.cmdType {
			case "SUB":
				s.On(res.data, res.handler)
			case "UNSUB":
				s.Off(res.data)

			}
			logger.WithFields(logrus.Fields{
				"data":  string(data),
				"pdata": string(d),
				"res":   res,
			}).Info("receive to sub")
			return res.msg, nil

		})
		return
	}

}

func WsConnect(sc SlaveConfig, pool *redis.Pool, rHub *redisocket.Hub, reqClient *greq.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		appKey := params["app_key"]
		token := r.FormValue("token")
		if appKey == "" || token == "" {
			logger.Warn("app_key or token is nil")
			http.Error(w, "app_key is nil", http.StatusUnauthorized)
			return
		}
		conn := pool.Get()
		reply, err := redis.Bytes(conn.Do("GET", token))
		if err != nil {
			conn.Close()
			logger.WithError(err).Warn("redis get error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		conn.Close()
		auth := Auth{}
		err = json.Unmarshal(reply, &auth)
		if err != nil {
			logger.WithError(err).Warn("json unmarshal error")
			http.Error(w, "token error", http.StatusUnauthorized)
			return
		}
		if appKey != auth.AppKey {
			http.Error(w, "appkey error", http.StatusUnauthorized)
			return
		}

		s, err := rHub.Upgrade(w, r, nil, auth.UserId, appKey)
		if err != nil {
			logger.WithError(err).Warnf("upgrade ws connection error")
			return
		}
		defer s.Close()

		s.Listen(func(data []byte) (b []byte, err error) {
			h, err := CommanRouter(data, pool)
			if err != nil {
				return
			}

			d, _, _, err := jsonparser.Get(data, "data")
			if err != nil {
				return
			}
			res, err := h(appKey, auth, d)
			if err != nil {
				return
			}
			switch res.cmdType {
			case "SUB":
				s.On(res.data, res.handler)
			case "UNSUB":
				s.Off(res.data)

			}
			return res.msg, nil
		})
		return
	}

}

func SubscribeCommand(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {

	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	msg = &commandResponse{
		handler: DefaultSubHandler,
		cmdType: "SUB",
	}
	command := &ChannelCommand{}
	exist := false
	for _, ch := range auth.Channels {
		//新增萬用字元  如果找到這個 任何頻道皆可訂閱
		if ch == "*" {
			exist = true
			break
		}
		ech := regexp.QuoteMeta(ch)
		rch := strings.Replace(ech, `\*`, "*", -1)
		r, err := regexp.Compile(rch)
		if err != nil {
			break
		}
		if r.MatchString(channel) {
			exist = true
			break
		}
	}
	var reply []byte
	if exist {
		msg.data = channel
		command.Event = SubscribeReplySucceeded
		command.Data.Channel = channel
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	} else {
		command.Event = SubscribeReplyError
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply

	}

	return
}
func Remote(pool *redis.Pool) func(string, Auth, []byte) (msg *commandResponse, err error) {
	return func(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {

		remote, err := jsonparser.GetString(data, "remote")
		if err != nil {
			return
		}
		uid, err := jsonparser.GetString(data, "uid")
		if err != nil {
			return
		}
		payload, _, _, err := jsonparser.Get(data, "payload")
		if err != nil {
			return
		}
		p := JsonCheck(string(payload))
		msg = &commandResponse{
			cmdType: "REMOTE",
		}
		var reply []byte
		command := &RemoteCommand{}
		command.Data.Remote = remote
		b, ok := auth.Remotes[remote]

		//沒有這個remote 返回錯誤訊息不斷線
		if !ok || !b {
			command.Event = RemoteReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				return
			}
			msg.msg = reply
			return
		}
		wp := WorkerPayload{
			UserId: auth.UserId,
			Data:   p,
			Uid:    uid,
			AppKey: auth.AppKey,
		}
		d, err := json.Marshal(wp)
		conn := pool.Get()
		_, err = conn.Do("RPUSH", auth.AppKey+"@"+remote, d)
		if err != nil {
			return
		}
		msg.msg = reply
		return
	}

}
func UnSubscribeCommand(appkey string, auth Auth, data []byte) (msg *commandResponse, err error) {
	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	exist := false
	for _, ch := range auth.Channels {
		//新增萬用字元  如果找到這個 任何頻道皆可訂閱
		if ch == "*" {
			exist = true
			break
		}
		ech := regexp.QuoteMeta(ch)
		rch := strings.Replace(ech, `\*`, "*", -1)
		r, err := regexp.Compile(rch)
		if err != nil {
			break
		}
		if r.MatchString(channel) {
			exist = true
			break
		}
	}
	msg = &commandResponse{
		cmdType: "UNSUB",
	}
	command := &ChannelCommand{}
	var reply []byte
	//反訂閱處理
	if exist {
		msg.data = channel
		command.Event = UnSubscribeReplySucceeded
		command.Data.Channel = channel
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	} else {
		command.Event = UnSubscribeReplyError
		reply, err = json.Marshal(command)
		if err != nil {
			return
		}
		msg.msg = reply
	}
	return
}

func CommanRouter(data []byte, pool *redis.Pool) (fn func(appkey string, auth Auth, data []byte) (msg *commandResponse, err error), err error) {

	val, err := jsonparser.GetString(data, "event")
	if err != nil {
		return
	}
	switch val {
	case RemoteEvent:
		return Remote(pool), nil
	case SubscribeEvent:
		return SubscribeCommand, nil
	case UnSubscribeEvent:
		return UnSubscribeCommand, nil
	default:
		err = errors.New("event errors")
		break
	}
	return
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package message

var _ InboundMessage = &inboundMessage{}
var _ OutboundMessage = &outboundMessage{}

// InboundMessage represents a set of fields for an inbound message that can be serialized into a byte stream
type InboundMessage interface {
	BytesSavedCompression() int
	Op() Op
	Bytes() []byte
	Get(Field) interface{}
}

type inboundMessage struct {
	op                    Op
	bytes                 []byte
	bytesSavedCompression int
	fields                map[Field]interface{}
}

// OutboundMessage represents a set of fields for an outbound message that can be serialized into a byte stream
type OutboundMessage interface {
	BytesSavedCompression() int
	Bytes() []byte
	Op() Op
}
type outboundMessage struct {
	bytes                 []byte
	bytesSavedCompression int
	op                    Op
}

// Op returns the value of the specified operation in this message
func (inMsg *inboundMessage) Op() Op { return inMsg.op }

// Bytes returns this message in bytes
func (inMsg *inboundMessage) Bytes() []byte { return inMsg.bytes }

// BytesSavedCompression returns the number of bytes this message saved due to
// compression. That is, the number of bytes we did not receive over the
// network due to the message being compressed. Empty slice for messages that were not
// compressed.
func (inMsg *inboundMessage) BytesSavedCompression() int { return inMsg.bytesSavedCompression }

// Field returns the value of the specified field in this message
func (inMsg *inboundMessage) Get(field Field) interface{} { return inMsg.fields[field] }

// Op returns the value of the specified operation in this message
func (outMsg *outboundMessage) Op() Op { return outMsg.op }

// Bytes returns this message in bytes
func (outMsg *outboundMessage) Bytes() []byte { return outMsg.bytes }

// BytesSavedCompression returns the number of bytes this message saved due to
// compression. That is, the number of bytes we did not send over the
// network due to the message being compressed. Empty slice for messages that were not
// compressed.
func (outMsg *outboundMessage) BytesSavedCompression() int { return outMsg.bytesSavedCompression }
<commit_msg>chore: cleanup interfaces<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package message

var (
	_ InboundMessage  = &inboundMessage{}
	_ OutboundMessage = &outboundMessage{}
)

// InboundMessage represents a set of fields for an inbound message that can be serialized into a byte stream
type InboundMessage interface {
	BytesSavedCompression() int
	Op() Op
	Bytes() []byte
	Get(Field) interface{}
}

type inboundMessage struct {
	op                    Op
	bytes                 []byte
	bytesSavedCompression int
	fields                map[Field]interface{}
}

// OutboundMessage represents a set of fields for an outbound message that can be serialized into a byte stream
type OutboundMessage interface {
	BytesSavedCompression() int
	Bytes() []byte
	Op() Op
}

type outboundMessage struct {
	bytes                 []byte
	bytesSavedCompression int
	op                    Op
}

// Op returns the value of the specified operation in this message
func (inMsg *inboundMessage) Op() Op { return inMsg.op }

// Bytes returns this message in bytes
func (inMsg *inboundMessage) Bytes() []byte { return inMsg.bytes }

// BytesSavedCompression returns the number of bytes this message saved due to
// compression. That is, the number of bytes we did not receive over the
// network due to the message being compressed. Empty slice for messages that were not
// compressed.
func (inMsg *inboundMessage) BytesSavedCompression() int { return inMsg.bytesSavedCompression }

// Field returns the value of the specified field in this message
func (inMsg *inboundMessage) Get(field Field) interface{} { return inMsg.fields[field] }

// Op returns the value of the specified operation in this message
func (outMsg *outboundMessage) Op() Op { return outMsg.op }

// Bytes returns this message in bytes
func (outMsg *outboundMessage) Bytes() []byte { return outMsg.bytes }

// BytesSavedCompression returns the number of bytes this message saved due to
// compression. That is, the number of bytes we did not send over the
// network due to the message being compressed. Empty slice for messages that were not
// compressed.
func (outMsg *outboundMessage) BytesSavedCompression() int { return outMsg.bytesSavedCompression }
<|endoftext|>


<commit_before>package compress

import (
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/mitchellh/packer/builder/file"
	env "github.com/mitchellh/packer/helper/builder/testing"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template"
)

func TestDetectFilename(t *testing.T) {
	// Test default / fallback with no file extension
	nakedFilename := Config{OutputPath: "test"}
	nakedFilename.detectFromFilename()
	if nakedFilename.Archive != "tar" {
		t.Error("Expected to find tar archive setting")
	}
	if nakedFilename.Algorithm != "pgzip" {
		t.Error("Expected to find pgzip algorithm setting")
	}

	// Test .archive
	zipFilename := Config{OutputPath: "test.zip"}
	zipFilename.detectFromFilename()
	if zipFilename.Archive != "zip" {
		t.Error("Expected to find zip archive setting")
	}
	if zipFilename.Algorithm != "" {
		t.Error("Expected to find empty algorithm setting")
	}

	// Test .compress
	lz4Filename := Config{OutputPath: "test.lz4"}
	lz4Filename.detectFromFilename()
	if lz4Filename.Archive != "" {
		t.Error("Expected to find empty archive setting")
	}
	if lz4Filename.Algorithm != "lz4" {
		t.Error("Expected to find lz4 algorithm setting")
	}

	// Test .archive.compress with some.extra.dots...
	lotsOfDots := Config{OutputPath: "test.blah.bloo.blee.tar.lz4"}
	lotsOfDots.detectFromFilename()
	if lotsOfDots.Archive != "tar" {
		t.Error("Expected to find tar archive setting")
	}
	if lotsOfDots.Algorithm != "lz4" {
		t.Error("Expected to find lz4 algorithm setting")
	}
}

func TestSimpleCompress(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.tar.gz"
	        }
	    ]
	}
	`
	artifact := testArchive(t, config)
	defer artifact.Destroy()

	fi, err := os.Stat("package.tar.gz")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
	if fi.IsDir() {
		t.Error("Archive should not be a directory")
	}
}

func TestZipArchive(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.zip"
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	// Verify things look good
	_, err := os.Stat("package.zip")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
}

func TestTarArchive(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.tar"
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	// Verify things look good
	_, err := os.Stat("package.tar")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
}

func TestCompressOptions(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.gz",
	            "compression_level": 9
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	// Verify things look good
	_, err := os.Stat("package.gz")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
}

// Test Helpers

func setup(t *testing.T) (packer.Ui, packer.Artifact, error) {
	// Create fake UI and Cache
	ui := packer.TestUi(t)
	cache := &packer.FileCache{CacheDir: os.TempDir()}

	// Create config for file builder
	const fileConfig = `{"builders":[{"type":"file","target":"package.txt","content":"Hello world!"}]}`
	tpl, err := template.Parse(strings.NewReader(fileConfig))
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to parse setup configuration: %s", err)
	}

	// Prepare the file builder
	builder := file.Builder{}
	warnings, err := builder.Prepare(tpl.Builders["file"].Config)
	if len(warnings) > 0 {
		for _, warn := range warnings {
			return nil, nil, fmt.Errorf("Configuration warning: %s", warn)
		}
	}
	if err != nil {
		return nil, nil, fmt.Errorf("Invalid configuration: %s", err)
	}

	// Run the file builder
	artifact, err := builder.Run(ui, nil, cache)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to build artifact: %s", err)
	}

	return ui, artifact, err
}

func testArchive(t *testing.T, config string) packer.Artifact {
	if os.Getenv(env.TestEnvVar) == "" {
		t.Skip(fmt.Sprintf(
			"Acceptance tests skipped unless env '%s' set", env.TestEnvVar))
	}

	ui, artifact, err := setup(t)
	if err != nil {
		t.Fatalf("Error bootstrapping test: %s", err)
	}
	if artifact != nil {
		defer artifact.Destroy()
	}

	tpl, err := template.Parse(strings.NewReader(config))
	if err != nil {
		t.Fatalf("Unable to parse test config: %s", err)
	}

	compressor := PostProcessor{}
	compressor.Configure(tpl.PostProcessors[0][0].Config)
	artifactOut, _, err := compressor.PostProcess(ui, artifact)
	if err != nil {
		t.Fatalf("Failed to compress artifact: %s", err)
	}

	return artifactOut
}
<commit_msg>Added test case for gzip that expands the data and compares to what we put in<commit_after>package compress

import (
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"

	"github.com/mitchellh/packer/builder/file"
	env "github.com/mitchellh/packer/helper/builder/testing"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template"
)

func TestDetectFilename(t *testing.T) {
	// Test default / fallback with no file extension
	nakedFilename := Config{OutputPath: "test"}
	nakedFilename.detectFromFilename()
	if nakedFilename.Archive != "tar" {
		t.Error("Expected to find tar archive setting")
	}
	if nakedFilename.Algorithm != "pgzip" {
		t.Error("Expected to find pgzip algorithm setting")
	}

	// Test .archive
	zipFilename := Config{OutputPath: "test.zip"}
	zipFilename.detectFromFilename()
	if zipFilename.Archive != "zip" {
		t.Error("Expected to find zip archive setting")
	}
	if zipFilename.Algorithm != "" {
		t.Error("Expected to find empty algorithm setting")
	}

	// Test .compress
	lz4Filename := Config{OutputPath: "test.lz4"}
	lz4Filename.detectFromFilename()
	if lz4Filename.Archive != "" {
		t.Error("Expected to find empty archive setting")
	}
	if lz4Filename.Algorithm != "lz4" {
		t.Error("Expected to find lz4 algorithm setting")
	}

	// Test .archive.compress with some.extra.dots...
	lotsOfDots := Config{OutputPath: "test.blah.bloo.blee.tar.lz4"}
	lotsOfDots.detectFromFilename()
	if lotsOfDots.Archive != "tar" {
		t.Error("Expected to find tar archive setting")
	}
	if lotsOfDots.Algorithm != "lz4" {
		t.Error("Expected to find lz4 algorithm setting")
	}
}

const expectedFileContents = "Hello world!"

func TestSimpleCompress(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.tar.gz"
	        }
	    ]
	}
	`
	artifact := testArchive(t, config)
	defer artifact.Destroy()

	fi, err := os.Stat("package.tar.gz")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
	if fi.IsDir() {
		t.Error("Archive should not be a directory")
	}
}

func TestZipArchive(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.zip"
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	// Verify things look good
	_, err := os.Stat("package.zip")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
}

func TestTarArchive(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.tar"
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	// Verify things look good
	_, err := os.Stat("package.tar")
	if err != nil {
		t.Errorf("Unable to read archive: %s", err)
	}
}

func TestCompressOptions(t *testing.T) {
	const config = `
	{
	    "post-processors": [
	        {
	            "type": "compress",
	            "output": "package.gz",
	            "compression_level": 9
	        }
	    ]
	}
	`

	artifact := testArchive(t, config)
	defer artifact.Destroy()

	filename := "package.gz"
	archive, _ := os.Open(filename)
	gzipReader, _ := gzip.NewReader(archive)
	data, _ := ioutil.ReadAll(gzipReader)

	if string(data) != expectedFileContents {
		t.Errorf("Expected:\n%s\nFound:\n%s\n", expectedFileContents, data)
	}
}

// Test Helpers

func setup(t *testing.T) (packer.Ui, packer.Artifact, error) {
	// Create fake UI and Cache
	ui := packer.TestUi(t)
	cache := &packer.FileCache{CacheDir: os.TempDir()}

	// Create config for file builder
	const fileConfig = `{"builders":[{"type":"file","target":"package.txt","content":"Hello world!"}]}`
	tpl, err := template.Parse(strings.NewReader(fileConfig))
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to parse setup configuration: %s", err)
	}

	// Prepare the file builder
	builder := file.Builder{}
	warnings, err := builder.Prepare(tpl.Builders["file"].Config)
	if len(warnings) > 0 {
		for _, warn := range warnings {
			return nil, nil, fmt.Errorf("Configuration warning: %s", warn)
		}
	}
	if err != nil {
		return nil, nil, fmt.Errorf("Invalid configuration: %s", err)
	}

	// Run the file builder
	artifact, err := builder.Run(ui, nil, cache)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to build artifact: %s", err)
	}

	return ui, artifact, err
}

func testArchive(t *testing.T, config string) packer.Artifact {
	if os.Getenv(env.TestEnvVar) == "" {
		t.Skip(fmt.Sprintf(
			"Acceptance tests skipped unless env '%s' set", env.TestEnvVar))
	}

	ui, artifact, err := setup(t)
	if err != nil {
		t.Fatalf("Error bootstrapping test: %s", err)
	}
	if artifact != nil {
		defer artifact.Destroy()
	}

	tpl, err := template.Parse(strings.NewReader(config))
	if err != nil {
		t.Fatalf("Unable to parse test config: %s", err)
	}

	compressor := PostProcessor{}
	compressor.Configure(tpl.PostProcessors[0][0].Config)
	artifactOut, _, err := compressor.PostProcess(ui, artifact)
	if err != nil {
		t.Fatalf("Failed to compress artifact: %s", err)
	}

	return artifactOut
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/pkg/errors"

	"github.com/spf13/pflag"

	"github.com/golang/protobuf/proto"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/cmd/profile"
	"github.com/havoc-io/mutagen/pkg/filesystem/behavior"
	"github.com/havoc-io/mutagen/pkg/sync"
)

const (
	snapshotFile = "snapshot_test"
	cacheFile    = "cache_test"
)

var usage = `scan_bench [-h|--help] [-p|--profile] [-i|--ignore=<pattern>] <path>
`

func main() {
	// Parse command line arguments.
	flagSet := pflag.NewFlagSet("scan_bench", pflag.ContinueOnError)
	flagSet.SetOutput(ioutil.Discard)
	var ignores []string
	var enableProfile bool
	flagSet.StringSliceVarP(&ignores, "ignore", "i", nil, "specify ignore paths")
	flagSet.BoolVarP(&enableProfile, "profile", "p", false, "enable profiling")
	if err := flagSet.Parse(os.Args[1:]); err != nil {
		if err == pflag.ErrHelp {
			fmt.Fprint(os.Stdout, usage)
			return
		} else {
			cmd.Fatal(errors.Wrap(err, "unable to parse command line"))
		}
	}
	arguments := flagSet.Args()
	if len(arguments) != 1 {
		cmd.Fatal(errors.New("invalid number of paths specified"))
	}
	path := arguments[0]

	// Print information.
	fmt.Println("Analyzing", path)

	// Create a snapshot without any cache. If requested, enable CPU and memory
	// profiling.
	var profiler *profile.Profile
	var err error
	if enableProfile {
		if profiler, err = profile.New("scan_cold"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start := time.Now()
	snapshot, preservesExecutability, recomposeUnicode, cache, ignoreCache, err := sync.Scan(
		path,
		nil,
		nil,
		sha1.New(),
		nil,
		ignores,
		nil,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target doesn't exist"))
	}
	stop := time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cold scan took", stop.Sub(start))
	fmt.Println("Root preserves executability:", preservesExecutability)
	fmt.Println("Root requires Unicode recomposition:", recomposeUnicode)

	// Create a snapshot with a cache. If requested, enable CPU and memory
	// profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_warm"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	snapshot, preservesExecutability, recomposeUnicode, _, _, err = sync.Scan(
		path,
		nil,
		nil,
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Warm scan took", stop.Sub(start))
	fmt.Println("Root preserves executability:", preservesExecutability)
	fmt.Println("Root requires Unicode recomposition:", recomposeUnicode)

	// Create a snapshot with a baseline and one re-check path. If requested,
	// enable CPU and memory profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_fast"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	snapshot, preservesExecutability, recomposeUnicode, _, _, err = sync.Scan(
		path,
		snapshot,
		map[string]bool{"fake path": true},
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Fast scan took", stop.Sub(start))
	fmt.Println("Root preserves executability:", preservesExecutability)
	fmt.Println("Root requires Unicode recomposition:", recomposeUnicode)

	// Create a snapshot with a baseline and no re-check paths. If requested,
	// enable CPU and memory profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_ultra_fast"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	snapshot, preservesExecutability, recomposeUnicode, _, _, err = sync.Scan(
		path,
		snapshot,
		nil,
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Ultra fast scan took", stop.Sub(start))
	fmt.Println("Root preserves executability:", preservesExecutability)
	fmt.Println("Root requires Unicode recomposition:", recomposeUnicode)

	// Serialize it.
	if enableProfile {
		if profiler, err = profile.New("serialize_snapshot"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	buffer := proto.NewBuffer(nil)
	buffer.SetDeterministic(true)
	if err := buffer.Marshal(snapshot); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to serialize snapshot"))
	}
	serializedSnapshot := buffer.Bytes()
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Snapshot serialization took", stop.Sub(start))

	// Deserialize it.
	if enableProfile {
		if profiler, err = profile.New("deserialize_snapshot"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	deserializedSnapshot := &sync.Entry{}
	if err = proto.Unmarshal(serializedSnapshot, deserializedSnapshot); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to deserialize snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Snapshot deserialization took", stop.Sub(start))

	// Validate the deserialized snapshot.
	start = time.Now()
	if err = deserializedSnapshot.EnsureValid(); err != nil {
		cmd.Fatal(errors.Wrap(err, "deserialized snapshot invalid"))
	}
	stop = time.Now()
	fmt.Println("Snapshot validation took", stop.Sub(start))

	// Write the serialized snapshot to disk.
	start = time.Now()
	if err = ioutil.WriteFile(snapshotFile, serializedSnapshot, 0600); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to write snapshot"))
	}
	stop = time.Now()
	fmt.Println("Snapshot write took", stop.Sub(start))

	// Read the serialized snapshot from disk.
	start = time.Now()
	if _, err = ioutil.ReadFile(snapshotFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to read snapshot"))
	}
	stop = time.Now()
	fmt.Println("Snapshot read took", stop.Sub(start))

	// Wipe the temporary file.
	if err = os.Remove(snapshotFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to remove snapshot"))
	}

	// TODO: I'd like to add a stable serialization benchmark since that's what
	// we really care about (especially since it has to copy the entire entry
	// tree), but I also don't want to expose that machinery publicly.

	// Print serialized snapshot size.
	fmt.Println("Serialized snapshot size is", len(serializedSnapshot), "bytes")

	// Print whether or not snapshots are equivalent.
	fmt.Println(
		"Original/deserialized snapshots equivalent?",
		deserializedSnapshot.Equal(snapshot),
	)

	// Checksum it.
	start = time.Now()
	sha1.Sum(serializedSnapshot)
	stop = time.Now()
	fmt.Println("SHA-1 snapshot digest took", stop.Sub(start))

	// Count entries.
	start = time.Now()
	snapshot.Count()
	stop = time.Now()
	fmt.Println("Snapshot entry counting took", stop.Sub(start))

	// Perform a deep copy of the snapshot.
	start = time.Now()
	snapshot.Copy()
	stop = time.Now()
	fmt.Println("Snapshot copying took", stop.Sub(start))

	// Serialize the cache.
	if enableProfile {
		if profiler, err = profile.New("serialize_cache"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	serializedCache, err := proto.Marshal(cache)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to serialize cache"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cache serialization took", stop.Sub(start))

	// Deserialize the cache.
	if enableProfile {
		if profiler, err = profile.New("deserialize_cache"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	deserializedCache := &sync.Cache{}
	if err = proto.Unmarshal(serializedCache, deserializedCache); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to deserialize cache"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cache deserialization took", stop.Sub(start))

	// Write the serialized cache to disk.
	start = time.Now()
	if err = ioutil.WriteFile(cacheFile, serializedCache, 0600); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to write cache"))
	}
	stop = time.Now()
	fmt.Println("Cache write took", stop.Sub(start))

	// Read the serialized cache from disk.
	start = time.Now()
	if _, err = ioutil.ReadFile(cacheFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to read cache"))
	}
	stop = time.Now()
	fmt.Println("Cache read took", stop.Sub(start))

	// Wipe the temporary file.
	if err = os.Remove(cacheFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to remove cache"))
	}

	// Print serialized cache size.
	fmt.Println("Serialized cache size is", len(serializedCache), "bytes")

	// Generate a reverse lookup map for the cache.
	start = time.Now()
	if _, err = cache.GenerateReverseLookupMap(); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to generate reverse lookup map"))
	}
	stop = time.Now()
	fmt.Println("Reverse lookup map generation took", stop.Sub(start))
}
<commit_msg>Added additional verifications to scan_bench.<commit_after>package main

import (
	"bytes"
	"crypto/sha1"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/pkg/errors"

	"github.com/spf13/pflag"

	"github.com/golang/protobuf/proto"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/cmd/profile"
	"github.com/havoc-io/mutagen/pkg/filesystem/behavior"
	"github.com/havoc-io/mutagen/pkg/sync"
)

const (
	snapshotFile = "snapshot_test"
	cacheFile    = "cache_test"
)

var usage = `scan_bench [-h|--help] [-p|--profile] [-i|--ignore=<pattern>] <path>
`

// cachesEqual checks that two scan caches have equivalent contents.
func cachesEqual(first, second *sync.Cache) bool {
	// Handle fast paths and verify non-nilness.
	if first == second {
		return true
	} else if first == nil || second == nil {
		return false
	}

	// Check lengths.
	if len(first.Entries) != len(second.Entries) {
		return false
	}

	// Check contents.
	for path, firstEntry := range first.Entries {
		// Extract corresponding content.
		secondEntry, ok := second.Entries[path]
		if !ok {
			return false
		}

		// Watch for nil values as a sanity check.
		if firstEntry == nil || secondEntry == nil {
			panic("cache has nil entries")
		} else if firstEntry.ModificationTime == nil || secondEntry.ModificationTime == nil {
			panic("nil modification time")
		}

		// Verify equivalence
		equivalent := secondEntry.Mode == firstEntry.Mode &&
			secondEntry.ModificationTime.Seconds == firstEntry.ModificationTime.Seconds &&
			secondEntry.ModificationTime.Nanos == firstEntry.ModificationTime.Nanos &&
			secondEntry.Size == firstEntry.Size &&
			secondEntry.FileID == firstEntry.FileID &&
			bytes.Equal(secondEntry.Digest, firstEntry.Digest)
		if !equivalent {
			return false
		}
	}

	// Success.
	return true
}

// ignoreCachesIntersectionEqual compares two ignore caches, ensuring that keys
// which are present in both caches have the same value.
func ignoreCachesIntersectionEqual(first, second sync.IgnoreCache) bool {
	// Check matches from first in second.
	for key, firstValue := range first {
		if secondValue, ok := second[key]; ok && secondValue != firstValue {
			return false
		}
	}

	// Check matches from second in first.
	for key, secondValue := range second {
		if firstValue, ok := first[key]; ok && firstValue != secondValue {
			return false
		}
	}

	// Success.
	return true
}

func main() {
	// Parse command line arguments.
	flagSet := pflag.NewFlagSet("scan_bench", pflag.ContinueOnError)
	flagSet.SetOutput(ioutil.Discard)
	var ignores []string
	var enableProfile bool
	flagSet.StringSliceVarP(&ignores, "ignore", "i", nil, "specify ignore paths")
	flagSet.BoolVarP(&enableProfile, "profile", "p", false, "enable profiling")
	if err := flagSet.Parse(os.Args[1:]); err != nil {
		if err == pflag.ErrHelp {
			fmt.Fprint(os.Stdout, usage)
			return
		} else {
			cmd.Fatal(errors.Wrap(err, "unable to parse command line"))
		}
	}
	arguments := flagSet.Args()
	if len(arguments) != 1 {
		cmd.Fatal(errors.New("invalid number of paths specified"))
	}
	path := arguments[0]

	// Print information.
	fmt.Println("Analyzing", path)

	// Perform a full (cold) scan. If requested, enable CPU and memory
	// profiling.
	var profiler *profile.Profile
	var err error
	if enableProfile {
		if profiler, err = profile.New("scan_full_cold"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start := time.Now()
	snapshot, preservesExecutability, decomposesUnicode, cache, ignoreCache, err := sync.Scan(
		path,
		nil,
		nil,
		sha1.New(),
		nil,
		ignores,
		nil,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target doesn't exist"))
	}
	stop := time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cold scan took", stop.Sub(start))
	fmt.Println("Root preserves executability:", preservesExecutability)
	fmt.Println("Root requires Unicode recomposition:", decomposesUnicode)

	// Perform a full (warm) scan. If requested, enable CPU and memory
	// profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_full_warm"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	newSnapshot, newPreservesExecutability, newDecomposesUnicode, newCache, newIgnoreCache, err := sync.Scan(
		path,
		nil,
		nil,
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Warm scan took", stop.Sub(start))

	// Compare the warm scan results with the baseline results.
	if !newSnapshot.Equal(snapshot) {
		cmd.Fatal(errors.New("snapshot mismatch"))
	} else if newPreservesExecutability != preservesExecutability {
		cmd.Fatal(errors.Errorf(
			"preserves executability mismatch: %t != %t",
			newPreservesExecutability,
			preservesExecutability,
		))
	} else if newDecomposesUnicode != decomposesUnicode {
		cmd.Fatal(errors.Errorf(
			"decomposes Unicode mismatch: %t != %t",
			newDecomposesUnicode,
			decomposesUnicode,
		))
	} else if !cachesEqual(newCache, cache) {
		cmd.Fatal(errors.New("cache mismatch"))
	} else if len(newIgnoreCache) != len(ignoreCache) {
		cmd.Fatal(errors.New("ignore cache length mismatch"))
	} else if !ignoreCachesIntersectionEqual(newIgnoreCache, ignoreCache) {
		cmd.Fatal(errors.New("ignore cache mismatch"))
	}

	// Perform an accelerated scan (with a re-check path). If requested, enable
	// CPU and memory profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_accelerated_recheck"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	newSnapshot, newPreservesExecutability, newDecomposesUnicode, newCache, newIgnoreCache, err = sync.Scan(
		path,
		snapshot,
		map[string]bool{"fake path": true},
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Accelerated scan (with re-check paths) took", stop.Sub(start))

	// Compare the accelerated scan results with the baseline results.
	if !newSnapshot.Equal(snapshot) {
		cmd.Fatal(errors.New("snapshot mismatch"))
	} else if newPreservesExecutability != preservesExecutability {
		cmd.Fatal(errors.Errorf(
			"preserves executability mismatch: %t != %t",
			newPreservesExecutability,
			preservesExecutability,
		))
	} else if newDecomposesUnicode != decomposesUnicode {
		cmd.Fatal(errors.Errorf(
			"decomposes Unicode mismatch: %t != %t",
			newDecomposesUnicode,
			decomposesUnicode,
		))
	} else if !cachesEqual(newCache, cache) {
		cmd.Fatal(errors.New("cache mismatch"))
	} else if !ignoreCachesIntersectionEqual(newIgnoreCache, ignoreCache) {
		cmd.Fatal(errors.New("ignore cache mismatch"))
	}

	// Perform an accelerated scan (without any re-check paths). If requested,
	// enable CPU and memory profiling.
	if enableProfile {
		if profiler, err = profile.New("scan_accelerated_no_recheck"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	newSnapshot, newPreservesExecutability, newDecomposesUnicode, newCache, newIgnoreCache, err = sync.Scan(
		path,
		snapshot,
		nil,
		sha1.New(),
		cache,
		ignores,
		ignoreCache,
		behavior.ProbeMode_ProbeModeProbe,
		sync.SymlinkMode_SymlinkModePortable,
	)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to create snapshot"))
	} else if snapshot == nil {
		cmd.Fatal(errors.New("target has been deleted since original snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Accelerated scan (without re-check paths) took", stop.Sub(start))

	// Compare the accelerated scan results with the baseline results.
	if !newSnapshot.Equal(snapshot) {
		cmd.Fatal(errors.New("snapshot mismatch"))
	} else if newPreservesExecutability != preservesExecutability {
		cmd.Fatal(errors.Errorf(
			"preserves executability mismatch: %t != %t",
			newPreservesExecutability,
			preservesExecutability,
		))
	} else if newDecomposesUnicode != decomposesUnicode {
		cmd.Fatal(errors.Errorf(
			"decomposes Unicode mismatch: %t != %t",
			newDecomposesUnicode,
			decomposesUnicode,
		))
	} else if !cachesEqual(newCache, cache) {
		cmd.Fatal(errors.New("cache mismatch"))
	} else if !ignoreCachesIntersectionEqual(newIgnoreCache, ignoreCache) {
		cmd.Fatal(errors.New("ignore cache mismatch"))
	}

	// Serialize it.
	if enableProfile {
		if profiler, err = profile.New("serialize_snapshot"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	buffer := proto.NewBuffer(nil)
	buffer.SetDeterministic(true)
	if err := buffer.Marshal(snapshot); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to serialize snapshot"))
	}
	serializedSnapshot := buffer.Bytes()
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Snapshot serialization took", stop.Sub(start))

	// Deserialize it.
	if enableProfile {
		if profiler, err = profile.New("deserialize_snapshot"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	deserializedSnapshot := &sync.Entry{}
	if err = proto.Unmarshal(serializedSnapshot, deserializedSnapshot); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to deserialize snapshot"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Snapshot deserialization took", stop.Sub(start))

	// Validate the deserialized snapshot.
	start = time.Now()
	if err = deserializedSnapshot.EnsureValid(); err != nil {
		cmd.Fatal(errors.Wrap(err, "deserialized snapshot invalid"))
	}
	stop = time.Now()
	fmt.Println("Snapshot validation took", stop.Sub(start))

	// Write the serialized snapshot to disk.
	start = time.Now()
	if err = ioutil.WriteFile(snapshotFile, serializedSnapshot, 0600); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to write snapshot"))
	}
	stop = time.Now()
	fmt.Println("Snapshot write took", stop.Sub(start))

	// Read the serialized snapshot from disk.
	start = time.Now()
	if _, err = ioutil.ReadFile(snapshotFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to read snapshot"))
	}
	stop = time.Now()
	fmt.Println("Snapshot read took", stop.Sub(start))

	// Wipe the temporary file.
	if err = os.Remove(snapshotFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to remove snapshot"))
	}

	// TODO: I'd like to add a stable serialization benchmark since that's what
	// we really care about (especially since it has to copy the entire entry
	// tree), but I also don't want to expose that machinery publicly.

	// Print serialized snapshot size.
	fmt.Println("Serialized snapshot size is", len(serializedSnapshot), "bytes")

	// Print whether or not snapshots are equivalent.
	fmt.Println(
		"Original/deserialized snapshots equivalent?",
		deserializedSnapshot.Equal(snapshot),
	)

	// Checksum it.
	start = time.Now()
	sha1.Sum(serializedSnapshot)
	stop = time.Now()
	fmt.Println("SHA-1 snapshot digest took", stop.Sub(start))

	// Count entries.
	start = time.Now()
	snapshot.Count()
	stop = time.Now()
	fmt.Println("Snapshot entry counting took", stop.Sub(start))

	// Perform a deep copy of the snapshot.
	start = time.Now()
	snapshot.Copy()
	stop = time.Now()
	fmt.Println("Snapshot copying took", stop.Sub(start))

	// Serialize the cache.
	if enableProfile {
		if profiler, err = profile.New("serialize_cache"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	serializedCache, err := proto.Marshal(cache)
	if err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to serialize cache"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cache serialization took", stop.Sub(start))

	// Deserialize the cache.
	if enableProfile {
		if profiler, err = profile.New("deserialize_cache"); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to create profiler"))
		}
	}
	start = time.Now()
	deserializedCache := &sync.Cache{}
	if err = proto.Unmarshal(serializedCache, deserializedCache); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to deserialize cache"))
	}
	stop = time.Now()
	if enableProfile {
		if err = profiler.Finalize(); err != nil {
			cmd.Fatal(errors.Wrap(err, "unable to finalize profiler"))
		}
		profiler = nil
	}
	fmt.Println("Cache deserialization took", stop.Sub(start))

	// Write the serialized cache to disk.
	start = time.Now()
	if err = ioutil.WriteFile(cacheFile, serializedCache, 0600); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to write cache"))
	}
	stop = time.Now()
	fmt.Println("Cache write took", stop.Sub(start))

	// Read the serialized cache from disk.
	start = time.Now()
	if _, err = ioutil.ReadFile(cacheFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to read cache"))
	}
	stop = time.Now()
	fmt.Println("Cache read took", stop.Sub(start))

	// Wipe the temporary file.
	if err = os.Remove(cacheFile); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to remove cache"))
	}

	// Print serialized cache size.
	fmt.Println("Serialized cache size is", len(serializedCache), "bytes")

	// Generate a reverse lookup map for the cache.
	start = time.Now()
	if _, err = cache.GenerateReverseLookupMap(); err != nil {
		cmd.Fatal(errors.Wrap(err, "unable to generate reverse lookup map"))
	}
	stop = time.Now()
	fmt.Println("Reverse lookup map generation took", stop.Sub(start))
}
<|endoftext|>


<commit_before>package pongo2

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
)

type Value struct {
	v reflect.Value
}

func AsValue(i interface{}) *Value {
	return &Value{
		v: reflect.ValueOf(i),
	}
}

func (v *Value) getResolvedValue() reflect.Value {
	if v.v.IsValid() && v.v.Kind() == reflect.Ptr {
		return v.v.Elem()
	}
	return v.v
}

func (v *Value) IsString() bool {
	return v.getResolvedValue().Kind() == reflect.String
}

func (v *Value) IsFloat() bool {
	return v.getResolvedValue().Kind() == reflect.Float32 ||
		v.getResolvedValue().Kind() == reflect.Float64
}

func (v *Value) IsInteger() bool {
	return v.getResolvedValue().Kind() == reflect.Int ||
		v.getResolvedValue().Kind() == reflect.Int8 ||
		v.getResolvedValue().Kind() == reflect.Int16 ||
		v.getResolvedValue().Kind() == reflect.Int32 ||
		v.getResolvedValue().Kind() == reflect.Int64 ||
		v.getResolvedValue().Kind() == reflect.Uint ||
		v.getResolvedValue().Kind() == reflect.Uint8 ||
		v.getResolvedValue().Kind() == reflect.Uint16 ||
		v.getResolvedValue().Kind() == reflect.Uint32 ||
		v.getResolvedValue().Kind() == reflect.Uint64
}

func (v *Value) IsNumber() bool {
	return v.IsInteger() || v.IsFloat()
}

func (v *Value) IsNil() bool {
	//fmt.Printf("%+v\n", v.getResolvedValue().Type().String())
	return !v.getResolvedValue().IsValid()
}

func (v *Value) String() string {
	switch v.getResolvedValue().Kind() {
	case reflect.String:
		return v.getResolvedValue().String()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return strconv.FormatInt(v.getResolvedValue().Int(), 10)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return strconv.FormatUint(v.getResolvedValue().Uint(), 10)
	case reflect.Float32, reflect.Float64:
		return fmt.Sprintf("%f", v.getResolvedValue().Float())
	case reflect.Bool:
		if v.Bool() {
			return "True"
		} else {
			return "False"
		}
	default:
		if v.IsNil() {
			return ""
		}
		logf("Value.String() not implemented for type: %s\n", v.getResolvedValue().Kind().String())
		return v.getResolvedValue().String()
	}
}

func (v *Value) Integer() int {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return int(v.getResolvedValue().Int())
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return int(v.getResolvedValue().Uint())
	case reflect.Float32, reflect.Float64:
		return int(v.getResolvedValue().Float())
	case reflect.String:
		// Try to convert from string to int (base 10)
		i, err := strconv.Atoi(v.getResolvedValue().String())
		if err != nil {
			return 0
		}
		return i
	default:
		logf("Value.Integer() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0
	}
}

func (v *Value) Float() float64 {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return float64(v.getResolvedValue().Int())
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return float64(v.getResolvedValue().Uint())
	case reflect.Float32, reflect.Float64:
		return v.getResolvedValue().Float()
	case reflect.String:
		// Try to convert from string to float64 (base 10)
		f, err := strconv.ParseFloat(v.getResolvedValue().String(), 64)
		if err != nil {
			return 0.0
		}
		return f
	default:
		logf("Value.Float() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0.0
	}
}

func (v *Value) Bool() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Bool:
		return v.getResolvedValue().Bool()
	default:
		logf("Value.Bool() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) IsTrue() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v.getResolvedValue().Int() != 0
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v.getResolvedValue().Uint() != 0
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return v.getResolvedValue().Len() > 0
	case reflect.Bool:
		return v.getResolvedValue().Bool()
	default:
		logf("Value.IsTrue() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) Negate() *Value {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		if v.Integer() != 0 {
			return AsValue(0)
		} else {
			return AsValue(1)
		}
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return AsValue(v.getResolvedValue().Len() == 0)
	case reflect.Bool:
		return AsValue(!v.getResolvedValue().Bool())
	default:
		logf("Value.IsTrue() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return AsValue(true)
	}
}

func (v *Value) Len() int {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return v.getResolvedValue().Len()
	default:
		logf("Value.Len() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0
	}
}

func (v *Value) Slice(i, j int) *Value {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Slice, reflect.String:
		return AsValue(v.getResolvedValue().Slice(i, j).Interface())
	default:
		logf("Value.Slice() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return AsValue([]int{})
	}
}

func (v *Value) Contains(other *Value) bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Struct:
		field_value := v.getResolvedValue().FieldByName(other.String())
		return field_value.IsValid()
	case reflect.Map:
		var map_value reflect.Value
		switch other.Interface().(type) {
		case int:
			map_value = v.getResolvedValue().MapIndex(other.getResolvedValue())
		case string:
			map_value = v.getResolvedValue().MapIndex(other.getResolvedValue())
		default:
			fmt.Printf("Value.Contains() does not support lookup type '%s'\n", other.getResolvedValue().Kind().String())
			return false
		}

		return map_value.IsValid()
	case reflect.String:
		return strings.Contains(v.getResolvedValue().String(), other.String())

	// TODO: reflect.Array, reflect.Slice

	default:
		logf("Value.Contains() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) CanSlice() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Slice, reflect.String:
		return true
	}
	return false
}

func (v *Value) Iterate(fn func(idx, count int, key, value *Value) bool, empty func()) {
	switch v.getResolvedValue().Kind() {
	case reflect.Map:
		keys := v.getResolvedValue().MapKeys()
		keyLen := len(keys)
		for idx, key := range keys {
			value := v.getResolvedValue().MapIndex(key)
			if !fn(idx, keyLen, &Value{key}, &Value{value}) {
				return
			}
		}
		if keyLen == 0 {
			empty()
		}
		return // done
	case reflect.Array, reflect.Slice, reflect.String:
		itemCount := v.getResolvedValue().Len()
		if itemCount > 0 {
			for i := 0; i < itemCount; i++ {
				if !fn(i, itemCount, &Value{v.getResolvedValue().Index(i)}, nil) {
					return
				}
			}
		} else {
			empty()
		}
		return // done
	default:
		logf("Value.Iterate() not available for type: %s\n", v.getResolvedValue().Kind().String())
	}
	empty()
}

func (v *Value) Interface() interface{} {
	if v.v.IsValid() {
		return v.v.Interface()
	}
	return nil
}

func (v *Value) EqualValueTo(other *Value) bool {
	return v.Interface() == other.Interface()
}
<commit_msg>time.Time will now be pretty-printed by default by calling pongo2.Value.String()<commit_after>package pongo2

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

type Value struct {
	v reflect.Value
}

func AsValue(i interface{}) *Value {
	return &Value{
		v: reflect.ValueOf(i),
	}
}

func (v *Value) getResolvedValue() reflect.Value {
	if v.v.IsValid() && v.v.Kind() == reflect.Ptr {
		return v.v.Elem()
	}
	return v.v
}

func (v *Value) IsString() bool {
	return v.getResolvedValue().Kind() == reflect.String
}

func (v *Value) IsFloat() bool {
	return v.getResolvedValue().Kind() == reflect.Float32 ||
		v.getResolvedValue().Kind() == reflect.Float64
}

func (v *Value) IsInteger() bool {
	return v.getResolvedValue().Kind() == reflect.Int ||
		v.getResolvedValue().Kind() == reflect.Int8 ||
		v.getResolvedValue().Kind() == reflect.Int16 ||
		v.getResolvedValue().Kind() == reflect.Int32 ||
		v.getResolvedValue().Kind() == reflect.Int64 ||
		v.getResolvedValue().Kind() == reflect.Uint ||
		v.getResolvedValue().Kind() == reflect.Uint8 ||
		v.getResolvedValue().Kind() == reflect.Uint16 ||
		v.getResolvedValue().Kind() == reflect.Uint32 ||
		v.getResolvedValue().Kind() == reflect.Uint64
}

func (v *Value) IsNumber() bool {
	return v.IsInteger() || v.IsFloat()
}

func (v *Value) IsNil() bool {
	//fmt.Printf("%+v\n", v.getResolvedValue().Type().String())
	return !v.getResolvedValue().IsValid()
}

func (v *Value) String() string {
	switch v.getResolvedValue().Kind() {
	case reflect.String:
		return v.getResolvedValue().String()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return strconv.FormatInt(v.getResolvedValue().Int(), 10)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return strconv.FormatUint(v.getResolvedValue().Uint(), 10)
	case reflect.Float32, reflect.Float64:
		return fmt.Sprintf("%f", v.getResolvedValue().Float())
	case reflect.Bool:
		if v.Bool() {
			return "True"
		} else {
			return "False"
		}
	case reflect.Struct:
		if v.getResolvedValue().Type() == reflect.TypeOf(time.Time{}) {
			return v.getResolvedValue().Interface().(time.Time).String()
		}
	}

	if v.IsNil() {
		return ""
	}
	logf("Value.String() not implemented for type: %s\n", v.getResolvedValue().Kind().String())
	return v.getResolvedValue().String()
}

func (v *Value) Integer() int {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return int(v.getResolvedValue().Int())
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return int(v.getResolvedValue().Uint())
	case reflect.Float32, reflect.Float64:
		return int(v.getResolvedValue().Float())
	case reflect.String:
		// Try to convert from string to int (base 10)
		i, err := strconv.Atoi(v.getResolvedValue().String())
		if err != nil {
			return 0
		}
		return i
	default:
		logf("Value.Integer() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0
	}
}

func (v *Value) Float() float64 {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return float64(v.getResolvedValue().Int())
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return float64(v.getResolvedValue().Uint())
	case reflect.Float32, reflect.Float64:
		return v.getResolvedValue().Float()
	case reflect.String:
		// Try to convert from string to float64 (base 10)
		f, err := strconv.ParseFloat(v.getResolvedValue().String(), 64)
		if err != nil {
			return 0.0
		}
		return f
	default:
		logf("Value.Float() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0.0
	}
}

func (v *Value) Bool() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Bool:
		return v.getResolvedValue().Bool()
	default:
		logf("Value.Bool() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) IsTrue() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v.getResolvedValue().Int() != 0
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v.getResolvedValue().Uint() != 0
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return v.getResolvedValue().Len() > 0
	case reflect.Bool:
		return v.getResolvedValue().Bool()
	default:
		logf("Value.IsTrue() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) Negate() *Value {
	switch v.getResolvedValue().Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		if v.Integer() != 0 {
			return AsValue(0)
		} else {
			return AsValue(1)
		}
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return AsValue(v.getResolvedValue().Len() == 0)
	case reflect.Bool:
		return AsValue(!v.getResolvedValue().Bool())
	default:
		logf("Value.IsTrue() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return AsValue(true)
	}
}

func (v *Value) Len() int {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
		return v.getResolvedValue().Len()
	default:
		logf("Value.Len() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return 0
	}
}

func (v *Value) Slice(i, j int) *Value {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Slice, reflect.String:
		return AsValue(v.getResolvedValue().Slice(i, j).Interface())
	default:
		logf("Value.Slice() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return AsValue([]int{})
	}
}

func (v *Value) Contains(other *Value) bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Struct:
		field_value := v.getResolvedValue().FieldByName(other.String())
		return field_value.IsValid()
	case reflect.Map:
		var map_value reflect.Value
		switch other.Interface().(type) {
		case int:
			map_value = v.getResolvedValue().MapIndex(other.getResolvedValue())
		case string:
			map_value = v.getResolvedValue().MapIndex(other.getResolvedValue())
		default:
			fmt.Printf("Value.Contains() does not support lookup type '%s'\n", other.getResolvedValue().Kind().String())
			return false
		}

		return map_value.IsValid()
	case reflect.String:
		return strings.Contains(v.getResolvedValue().String(), other.String())

	// TODO: reflect.Array, reflect.Slice

	default:
		logf("Value.Contains() not available for type: %s\n", v.getResolvedValue().Kind().String())
		return false
	}
}

func (v *Value) CanSlice() bool {
	switch v.getResolvedValue().Kind() {
	case reflect.Array, reflect.Slice, reflect.String:
		return true
	}
	return false
}

func (v *Value) Iterate(fn func(idx, count int, key, value *Value) bool, empty func()) {
	switch v.getResolvedValue().Kind() {
	case reflect.Map:
		keys := v.getResolvedValue().MapKeys()
		keyLen := len(keys)
		for idx, key := range keys {
			value := v.getResolvedValue().MapIndex(key)
			if !fn(idx, keyLen, &Value{key}, &Value{value}) {
				return
			}
		}
		if keyLen == 0 {
			empty()
		}
		return // done
	case reflect.Array, reflect.Slice, reflect.String:
		itemCount := v.getResolvedValue().Len()
		if itemCount > 0 {
			for i := 0; i < itemCount; i++ {
				if !fn(i, itemCount, &Value{v.getResolvedValue().Index(i)}, nil) {
					return
				}
			}
		} else {
			empty()
		}
		return // done
	default:
		logf("Value.Iterate() not available for type: %s\n", v.getResolvedValue().Kind().String())
	}
	empty()
}

func (v *Value) Interface() interface{} {
	if v.v.IsValid() {
		return v.v.Interface()
	}
	return nil
}

func (v *Value) EqualValueTo(other *Value) bool {
	return v.Interface() == other.Interface()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/miekg/pkcs11"
	"log"
	"os"
)

func main() {
	flag.Parse()
	wd, _ := os.Getwd()
	os.Setenv("SOFTHSM_CONF", wd+"/softhsm.conf")
	p := pkcs11.New("/usr/lib/softhsm/libsofthsm.so")
	if len(flag.Args()) > 0 {
		p = pkcs11.New(flag.Arg(0))
	}
	if p == nil {
		log.Fatalf("new error\n")
	}
	if e := p.Initialize(); e != nil {
		log.Fatalf("init error %s\n", e.Error())
	}

	defer p.Destroy()
	defer p.Finalize()
	slots, e := p.GetSlotList(true)
	log.Printf("slots %v\n", slots)
	if e != nil {
		log.Fatalf("slots %s\n", e.Error())
		return
	}
	session, e := p.OpenSession(slots[0], pkcs11.CKF_SERIAL_SESSION|pkcs11.CKF_RW_SESSION)
	if e != nil {
		log.Fatalf("session %s\n", e.Error())
	}
	defer p.CloseSession(session)
	log.Printf("%v %v\n", slots, session)

	if e := p.Login(session, pkcs11.CKU_USER, "1234"); e != nil {
		log.Fatal("user pin %s\n", e.Error())
	}
	defer p.Logout(session)
	publicKeyTemplate := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_KEY_TYPE, pkcs11.CKO_PUBLIC_KEY),
		pkcs11.NewAttribute(pkcs11.CKA_TOKEN, true),
		pkcs11.NewAttribute(pkcs11.CKA_ENCRYPT, true),
		pkcs11.NewAttribute(pkcs11.CKA_PUBLIC_EXPONENT, []byte{3}),
		pkcs11.NewAttribute(pkcs11.CKA_MODULUS_BITS, 1024),
		pkcs11.NewAttribute(pkcs11.CKA_LABEL, "MyFirstKey"),
	}
	privateKeyTemplate := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_KEY_TYPE, pkcs11.CKO_PRIVATE_KEY),
		pkcs11.NewAttribute(pkcs11.CKA_TOKEN, true),
		pkcs11.NewAttribute(pkcs11.CKA_PRIVATE, true),
		pkcs11.NewAttribute(pkcs11.CKA_SIGN, true),
		pkcs11.NewAttribute(pkcs11.CKA_LABEL, "MyFirstKey"),
	}
	pub, priv, e := p.GenerateKeyPair(session,
		[]*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_RSA_PKCS_KEY_PAIR_GEN, nil)},
		publicKeyTemplate, privateKeyTemplate)
	if e != nil {
		log.Fatalf("%s\n", e.Error())
	}
	log.Printf("pub id %d, priv id %d", pub, priv)
	e = p.SignInit(session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_SHA1_RSA_PKCS, nil)}, priv)
	if e != nil {
		log.Fatalf("SignInit: %s\n", e.Error())
	}

	// Sign something with priv
	data := []byte{1, 2, 3, 4}

	sig, err := p.Sign(session, data)
	if err != nil {
		log.Fatalf("sig: %s\n", err.Error())
	}
	log.Printf("%v\n", sig)
}
<commit_msg>remove sign directory<commit_after><|endoftext|>


<commit_before>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

import (
	"fmt"
	"reflect"
	"testing"
)

func TestLitArrayToSlice(t *testing.T) {
	// null pointers
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 0) })
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 1) })
	// not zero terminated
	badPtr := &cArray123[0]
	assertPanics(t, "not zero terminated - 1", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 1) })
	assertPanics(t, "not zero terminated - 2", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 2) })
	// zero length
	for maxLen := 0; maxLen <= 2; maxLen++ {
		if ls := litArrayToSlice(&cZero, 0); len(ls) != 0 {
			t.Errorf("Test 0-length litArrayToSlice, maxLen==%d: return value has length %d",
				maxLen, len(ls))
		}
	}
	// works correctly
	ptr := &cArray1230[0]
	expected := []Literal{1, 2, 3}
	for maxLen := 3; maxLen <= 10; maxLen++ {
		if ls := litArrayToSlice(ptr, maxLen); !reflect.DeepEqual(ls, expected) {
			t.Errorf("Test litArrayToSlice correct, maxLen==%d: expected %v but got %v",
				maxLen, expected, ls)
		}
	}
}

func TestAssumptionsSucceeding(t *testing.T) {
	successTests := []struct {
		assumpts  []Literal // The literals which should be assumed true/false
		solutions int       // The number of solutions that we expect to produce
	}{
		{[]Literal{1}, 10}, // We use formulaTests[0].formula below
		{[]Literal{2}, 9},
		{[]Literal{3}, 6},
		{[]Literal{4, 5}, 4},
	}

	for i, at := range successTests {
		p, _ := New(nil)
		p.AddClauses(formulaTests[0].formula)

		count := 0
		for ; ; count++ {
			for _, lit := range at.assumpts {
				p.Assume(lit)
			}
			status, sol := p.Solve()
			if status != Satisfiable {
				break
			}

			// All the UNSAT methods should give zero answers.
			if p.FailedAssumption(at.assumpts[0]) {
				t.Errorf("Test %d: FailedAssumption: expected %v not to be failed", i, at.assumpts[0])
			}
			if r := p.FailedAssumptions(); !reflect.DeepEqual(r, []Literal{}) {
				t.Errorf("Test %d: FailedAssumptions: expected [], got %v", i, r)
			}

			p.BlockSolution(sol)
		}
		if count != at.solutions {
			t.Errorf("Test %d: Expected %d solution(s) for assumptions %v; got %d",
				i, at.solutions, at.assumpts, count)
		}
	}
}

func TestAssumptionsFailing(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	failed := []Literal{3, 4}
	if actual := p.FailedAssumptions(); !reflect.DeepEqual(failed, actual) {
		t.Errorf("Expected failed assumptions %v != %v actual", failed, actual)
	}
	for _, f := range failed {
		if !p.FailedAssumption(f) {
			t.Errorf("Expected literal %v to be a failed assumption", f)
		}
	}

	i, wanted := 0, [][]Literal{{3, 5}, {3, 5}, {5, 4}, {}}
	for a := p.MaxSatisfiableAssumptions(); len(a) > 0; i++ {
		if !reflect.DeepEqual(a, wanted[i]) {
			t.Errorf("(Next)MaxSat'Assumpt's: Got %v wanted %v", a, wanted[i])
		}
		a = p.NextMaxSatisfiableAssumptions()
	}
	if a := p.MaxSatisfiableAssumptions(); !reflect.DeepEqual(a, []Literal{}) {
		t.Errorf("MaxSatisfiableAssumptions: CNF inconsistent, so wanted []Literal{}, but got %v", a)
	}
}

// TestAssumptionsNoCrash tests that if you call Assume with a satisfiable
// assumption after Solve returns UNSAT then FailedAssumption(s) does not crash.
func TestCrashOnAssumeSatAfterUnsatThenCallFailedAssumptions(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}
	if !p.FailedAssumption(3) {
		t.Fatalf("Expected literal 3 to be a failed assumption")
	}
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	p.Assume(3)
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	// Either the next two assertions work or they crash with this message:
	//   *** picosat: API usage: expected to be in UNSAT state
	//   SIGABRT: abort
	if p.FailedAssumption(3) {
		t.Errorf("FailedAssumption: Expected false, got true")
	}
	if r := p.FailedAssumptions(); len(r) != 0 {
		t.Errorf("FailedAssumptions: Expected []Literal{}, got %v", r)
	}
}

// ExamplePigosat_Assume demonstrates how to use Assume and related methods.
func ExamplePigosat_Assume() {
	var formula = Formula{{1, 2, 3}, {1, 2}, {2, 3}}
	p, _ := New(nil)
	p.AddClauses(formula)
	fmt.Println("Formula:", formula)
	status, solution := p.Solve()
	fmt.Println("No assumptions:", status, "solution ==", solution)

	// Satisfiable assumptions
	fmt.Println()
	fmt.Println("**** SATISFIABLE ASSUMPTIONS ****")
	p.Assume(1)
	p.Assume(-2)
	// Assumptions do not change the number of clauses.
	fmt.Println("Assume  1, -2 : Number of clauses:", p.AddedOriginalClauses())
	status, solution = p.Solve()
	fmt.Println("               ", status, "solution ==", solution)

	// Calls to p.AddClauses or p.Assume cancel assumptions 1 and -2
	// immediately, or a second call to p.Solve also cancels the assumptions.
	p.Assume(-3)
	status, solution = p.Solve()
	fmt.Println("Assume      -3:", status, "solution ==", solution)

	// Unsatisfiable assumptions
	fmt.Println()
	fmt.Println("**** UNSATISFIABLE ASSUMPTIONS ****")
	p.Assume(-1)                 // assume unit clause Clause{-1, 0}
	p.Assume(-2)                 // assume unit clause Clause{-2, 0}
	status, solution = p.Solve() // assumes -1 and -2 hold
	fmt.Println("Assume -1, -2 :", status, "solution ==", solution)
	// Assumptions -1 and -2 caused unsatisfiability.
	fmt.Println("                Failed assumptions:", p.FailedAssumptions())
	fmt.Println("                Assumption -1 failed:", p.FailedAssumption(-1))
	fmt.Println("                Assumption -2 failed:", p.FailedAssumption(-2))
	// Not every subset of the assumptions causes unsatisfiability.
	// p.MaxSatisfiableAssumptions would return the same as
	// p.NextMaxSatisfiableAssumptions, but the return value wouldn't change
	// with each call like it does with p.NextMaxSatisfiableAssumptions.
	fmt.Println("                Maximal satisfiable subset of assumptions 1:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 2:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 3:", p.NextMaxSatisfiableAssumptions())
	// p.NextMaxSatisfiableAssumptions does add clauses.
	fmt.Println("                Number of clauses:", p.AddedOriginalClauses())

	// Unknown status
	// Assumptions are valid but p.Solve returns no Solution assignment. The
	// solver knowns the status is Unknown until a call to p.Assume,
	// p.AddClauses, or p.Solve resets the assumptions.

	// Output:
	// Formula: [[1 2 3] [1 2] [2 3]]
	// No assumptions: Satisfiable solution == {1:true , 2:true , 3:true}
	//
	// **** SATISFIABLE ASSUMPTIONS ****
	// Assume  1, -2 : Number of clauses: 3
	//                 Satisfiable solution == {1:true , 2:false, 3:true}
	// Assume      -3: Satisfiable solution == {1:true , 2:true , 3:false}
	//
	// **** UNSATISFIABLE ASSUMPTIONS ****
	// Assume -1, -2 : Unsatisfiable solution == {}
	//                 Failed assumptions: [-1 -2]
	//                 Assumption -1 failed: true
	//                 Assumption -2 failed: true
	//                 Maximal satisfiable subset of assumptions 1: [-1]
	//                 Maximal satisfiable subset of assumptions 2: [-2]
	//                 Maximal satisfiable subset of assumptions 3: []
	//                 Number of clauses: 5
}
<commit_msg>Add test to demonstrate issue #22: calling BlockSolution from inside Solve<commit_after>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

import (
	"fmt"
	"reflect"
	"testing"
)

func TestLitArrayToSlice(t *testing.T) {
	// null pointers
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 0) })
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 1) })
	// not zero terminated
	badPtr := &cArray123[0]
	assertPanics(t, "not zero terminated - 1", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 1) })
	assertPanics(t, "not zero terminated - 2", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 2) })
	// zero length
	for maxLen := 0; maxLen <= 2; maxLen++ {
		if ls := litArrayToSlice(&cZero, 0); len(ls) != 0 {
			t.Errorf("Test 0-length litArrayToSlice, maxLen==%d: return value has length %d",
				maxLen, len(ls))
		}
	}
	// works correctly
	ptr := &cArray1230[0]
	expected := []Literal{1, 2, 3}
	for maxLen := 3; maxLen <= 10; maxLen++ {
		if ls := litArrayToSlice(ptr, maxLen); !reflect.DeepEqual(ls, expected) {
			t.Errorf("Test litArrayToSlice correct, maxLen==%d: expected %v but got %v",
				maxLen, expected, ls)
		}
	}
}

func TestAssumptionsSucceeding(t *testing.T) {
	successTests := []struct {
		assumpts  []Literal // The literals which should be assumed true/false
		solutions int       // The number of solutions that we expect to produce
	}{
		{[]Literal{1}, 10}, // We use formulaTests[0].formula below
		{[]Literal{2}, 9},
		{[]Literal{3}, 6},
		{[]Literal{4, 5}, 4},
	}

	for i, at := range successTests {
		p, _ := New(nil)
		p.AddClauses(formulaTests[0].formula)

		count := 0
		for ; ; count++ {
			for _, lit := range at.assumpts {
				p.Assume(lit)
			}
			status, sol := p.Solve()
			if status != Satisfiable {
				break
			}

			// All the UNSAT methods should give zero answers.
			if p.FailedAssumption(at.assumpts[0]) {
				t.Errorf("Test %d: FailedAssumption: expected %v not to be failed", i, at.assumpts[0])
			}
			if r := p.FailedAssumptions(); !reflect.DeepEqual(r, []Literal{}) {
				t.Errorf("Test %d: FailedAssumptions: expected [], got %v", i, r)
			}

			p.BlockSolution(sol)
		}
		if count != at.solutions {
			t.Errorf("Test %d: Expected %d solution(s) for assumptions %v; got %d",
				i, at.solutions, at.assumpts, count)
		}
	}
}

func TestAssumptionsFailing(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	failed := []Literal{3, 4}
	if actual := p.FailedAssumptions(); !reflect.DeepEqual(failed, actual) {
		t.Errorf("Expected failed assumptions %v != %v actual", failed, actual)
	}
	for _, f := range failed {
		if !p.FailedAssumption(f) {
			t.Errorf("Expected literal %v to be a failed assumption", f)
		}
	}

	i, wanted := 0, [][]Literal{{3, 5}, {3, 5}, {5, 4}, {}}
	for a := p.MaxSatisfiableAssumptions(); len(a) > 0; i++ {
		if !reflect.DeepEqual(a, wanted[i]) {
			t.Errorf("(Next)MaxSat'Assumpt's: Got %v wanted %v", a, wanted[i])
		}
		a = p.NextMaxSatisfiableAssumptions()
	}
	if a := p.MaxSatisfiableAssumptions(); !reflect.DeepEqual(a, []Literal{}) {
		t.Errorf("MaxSatisfiableAssumptions: CNF inconsistent, so wanted []Literal{}, but got %v", a)
	}
}

// TestAssumptionsNoCrash tests that if you call Assume with a satisfiable
// assumption after Solve returns UNSAT then FailedAssumption(s) does not crash.
func TestCrashOnAssumeSatAfterUnsatThenCallFailedAssumptions(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}
	if !p.FailedAssumption(3) {
		t.Fatalf("Expected literal 3 to be a failed assumption")
	}
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	p.Assume(3)
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	// Either the next two assertions work or they crash with this message:
	//   *** picosat: API usage: expected to be in UNSAT state
	//   SIGABRT: abort
	if p.FailedAssumption(3) {
		t.Errorf("FailedAssumption: Expected false, got true")
	}
	if r := p.FailedAssumptions(); len(r) != 0 {
		t.Errorf("FailedAssumptions: Expected []Literal{}, got %v", r)
	}
}

// TestNextMaxSatisfiableAssumptionsAsIterator tests that NextMaxSatisfiableAssumptions
// can be used as an iterator. In particular, this test panics if Solve calls
// BlockSolution when Solve returns Satisfiable.
func TestNextMaxSatisfiableAssumptionsAsIterator(t *testing.T) {
	var formula = Formula{{1, 2, 3}, {1, 2}, {2, 3}}
	p, _ := New(nil)
	p.AddClauses(formula)
	p.Assume(1)
	p.Assume(-2)
	p.Solve()
	p.Assume(-1)
	p.Assume(-2)
	ms := make([][]Literal, 0)
	for m := p.NextMaxSatisfiableAssumptions(); len(m) > 0; m = p.NextMaxSatisfiableAssumptions() {
		ms = append(ms, m)
	}
	expected := [][]Literal{{-1}, {-2}}
	if !reflect.DeepEqual(ms, expected) {
		t.Errorf("Expected %v. Got %v.", expected, ms)
	}
}

// ExamplePigosat_Assume demonstrates how to use Assume and related methods.
func ExamplePigosat_Assume() {
	var formula = Formula{{1, 2, 3}, {1, 2}, {2, 3}}
	p, _ := New(nil)
	p.AddClauses(formula)
	fmt.Println("Formula:", formula)
	status, solution := p.Solve()
	fmt.Println("No assumptions:", status, "solution ==", solution)

	// Satisfiable assumptions
	fmt.Println()
	fmt.Println("**** SATISFIABLE ASSUMPTIONS ****")
	p.Assume(1)
	p.Assume(-2)
	// Assumptions do not change the number of clauses.
	fmt.Println("Assume  1, -2 : Number of clauses:", p.AddedOriginalClauses())
	status, solution = p.Solve()
	fmt.Println("               ", status, "solution ==", solution)

	// Calls to p.AddClauses or p.Assume cancel assumptions 1 and -2
	// immediately, or a second call to p.Solve also cancels the assumptions.
	p.Assume(-3)
	status, solution = p.Solve()
	fmt.Println("Assume      -3:", status, "solution ==", solution)

	// Unsatisfiable assumptions
	fmt.Println()
	fmt.Println("**** UNSATISFIABLE ASSUMPTIONS ****")
	p.Assume(-1)                 // assume unit clause Clause{-1, 0}
	p.Assume(-2)                 // assume unit clause Clause{-2, 0}
	status, solution = p.Solve() // assumes -1 and -2 hold
	fmt.Println("Assume -1, -2 :", status, "solution ==", solution)
	// Assumptions -1 and -2 caused unsatisfiability.
	fmt.Println("                Failed assumptions:", p.FailedAssumptions())
	fmt.Println("                Assumption -1 failed:", p.FailedAssumption(-1))
	fmt.Println("                Assumption -2 failed:", p.FailedAssumption(-2))
	// Not every subset of the assumptions causes unsatisfiability.
	// p.MaxSatisfiableAssumptions would return the same as
	// p.NextMaxSatisfiableAssumptions, but the return value wouldn't change
	// with each call like it does with p.NextMaxSatisfiableAssumptions.
	fmt.Println("                Maximal satisfiable subset of assumptions 1:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 2:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 3:", p.NextMaxSatisfiableAssumptions())
	// p.NextMaxSatisfiableAssumptions does add clauses.
	fmt.Println("                Number of clauses:", p.AddedOriginalClauses())

	// Unknown status
	// Assumptions are valid but p.Solve returns no Solution assignment. The
	// solver knowns the status is Unknown until a call to p.Assume,
	// p.AddClauses, or p.Solve resets the assumptions.

	// Output:
	// Formula: [[1 2 3] [1 2] [2 3]]
	// No assumptions: Satisfiable solution == {1:true , 2:true , 3:true}
	//
	// **** SATISFIABLE ASSUMPTIONS ****
	// Assume  1, -2 : Number of clauses: 3
	//                 Satisfiable solution == {1:true , 2:false, 3:true}
	// Assume      -3: Satisfiable solution == {1:true , 2:true , 3:false}
	//
	// **** UNSATISFIABLE ASSUMPTIONS ****
	// Assume -1, -2 : Unsatisfiable solution == {}
	//                 Failed assumptions: [-1 -2]
	//                 Assumption -1 failed: true
	//                 Assumption -2 failed: true
	//                 Maximal satisfiable subset of assumptions 1: [-1]
	//                 Maximal satisfiable subset of assumptions 2: [-2]
	//                 Maximal satisfiable subset of assumptions 3: []
	//                 Number of clauses: 5
}
<|endoftext|>