<commit_before>package interp

import (
	"fmt"
	"go/ast"
	"go/token"
	"log"
	"reflect"

	"code.google.com/p/go.tools/go/exact"
)

func (env *environ) runStmt(stmt ast.Stmt, topLevel bool) {
	switch stmt := stmt.(type) {
	case *ast.ReturnStmt:
		if topLevel {
			log.Fatal("Return from top-level not allowed")
		}
		log.Fatal("Return statements not implemented yet")
	case *ast.AssignStmt:
		switch stmt.Tok {
		case token.DEFINE:
			// Short variable declaration
			// Collect LHS identifiers
			idents := make([]*ast.Ident, len(stmt.Lhs))
			for i, lhs := range stmt.Lhs {
				if ident, ok := lhs.(*ast.Ident); ok {
					// LHS is identifier
					idents[i] = ident
				} else {
					log.Fatal("short declaration LHS not identifier!")
				}
			}
			// Evaluate RHS expressions
			var rhs []Object
			if len(stmt.Rhs) == 1 {
				// Single rhs expression, potentially multi-valued
				rhs = env.Eval(stmt.Rhs[0])
			} else {
				// Multiple rhs expressions, each single-valued
				rhs = make([]Object, len(stmt.Lhs))
				for i, expr := range stmt.Rhs {
					rhs[i] = env.Eval(expr)[0]
				}
			}
			// Create LHS variables and set them to RHS
			for i, ident := range idents {
				obj := rhs[i]
				rval := obj.Value.(reflect.Value)
				identDef := env.info.Defs[ident]
				if identDef == nil || identDef.Pos() != ident.Pos() {
					// Redeclaration: variable already exists in current scope. Just assign the new value.
					lhs, _ := env.lookup(ident.Name)
					lval := lhs.Value.(reflect.Value)
					lval.Set(rval)
				} else {
					// New variable declaration. Create new variable with the right value.
					typ := rval.Type()
					sv := reflect.New(typ).Elem()
					sv.Set(rval)
					obj.Value = sv
					// Add the name we're declaring to env.names if it's a new name
					if _, ok := env.lookup(ident.Name); !ok {
						env.names = append(env.names, ident.Name)
					}
					// Add the object to env.objs
					env.objs[ident.Name] = obj
				}
			}
		case token.ASSIGN:
			// assignment
			// TODO: this only covers very simple assignment. There are more complicated rules
			//   not yet implemented (see http://golang.org/ref/spec#Assignments), such as:
			//   1) nil
			//   2) Blank identifier
			lhs := make([]Object, len(stmt.Lhs))
			rhs := make([]reflect.Value, len(stmt.Lhs))
			for i, expr := range stmt.Lhs {
				lhs[i] = env.Eval(expr)[0]
			}
			if len(stmt.Rhs) == 1 {
				// Single rhs expression, potentially multi-valued
				robjs := env.Eval(stmt.Rhs[0])
				for i, robj := range robjs {
					rv, ok := robj.Value.(reflect.Value)
					if !ok {
						// This means that rval is untyped.
						// In assignment context, this should only happen if it's untyped nil
						rhs[i] = reflect.ValueOf(nil)
					}
					rhs[i] = rv
				}
			} else {
				// Multiple rhs expressions, each single-valued
				for i, expr := range stmt.Rhs {
					robj := env.Eval(expr)[0]
					rv, ok := robj.Value.(reflect.Value)
					if !ok {
						rhs[i] = reflect.ValueOf(nil)
					}
					rhs[i] = rv
				}
			}

			for i, obj := range lhs {
				v := obj.Value.(reflect.Value)
				if rhs[i].IsValid() {
					v.Set(rhs[i])
				} else {
					v.Set(reflect.Zero(v.Type()))
				}
			}
		}
	case *ast.ExprStmt:
		// If we're not at top level, then only call expressions and receive operations are valid statements
		if !topLevel {
			if _, ok := stmt.X.(*ast.CallExpr); !ok {
				// TODO: handle this error better
				log.Fatal("Expression used as statement inappropriately")
			}
			// TODO: what about receive operations?
		}
		objs := env.Eval(stmt.X)
		if topLevel {
			for _, obj := range objs {
				// TODO: do something better than print the results to stdout
				switch v := obj.Value.(type) {
				case reflect.Value:
					fmt.Printf("=> %s: %v\n", TypeString(obj.Typ), v.Interface())
				case exact.Value, nil:
					fmt.Printf("=> %s: %v\n", TypeString(obj.Typ), v)
				}
			}
		}
	}
}
<commit_msg>Implemented the "go" statement (with 2 lines of code!)<commit_after>package interp

import (
	"fmt"
	"go/ast"
	"go/token"
	"log"
	"reflect"

	"code.google.com/p/go.tools/go/exact"
)

func (env *environ) runStmt(stmt ast.Stmt, topLevel bool) {
	switch stmt := stmt.(type) {
	case *ast.ReturnStmt:
		if topLevel {
			log.Fatal("Return from top-level not allowed")
		}
		log.Fatal("Return statements not implemented yet")
	case *ast.AssignStmt:
		switch stmt.Tok {
		case token.DEFINE:
			// Short variable declaration
			// Collect LHS identifiers
			idents := make([]*ast.Ident, len(stmt.Lhs))
			for i, lhs := range stmt.Lhs {
				if ident, ok := lhs.(*ast.Ident); ok {
					// LHS is identifier
					idents[i] = ident
				} else {
					log.Fatal("short declaration LHS not identifier!")
				}
			}
			// Evaluate RHS expressions
			var rhs []Object
			if len(stmt.Rhs) == 1 {
				// Single rhs expression, potentially multi-valued
				rhs = env.Eval(stmt.Rhs[0])
			} else {
				// Multiple rhs expressions, each single-valued
				rhs = make([]Object, len(stmt.Lhs))
				for i, expr := range stmt.Rhs {
					rhs[i] = env.Eval(expr)[0]
				}
			}
			// Create LHS variables and set them to RHS
			for i, ident := range idents {
				obj := rhs[i]
				rval := obj.Value.(reflect.Value)
				identDef := env.info.Defs[ident]
				if identDef == nil || identDef.Pos() != ident.Pos() {
					// Redeclaration: variable already exists in current scope. Just assign the new value.
					lhs, _ := env.lookup(ident.Name)
					lval := lhs.Value.(reflect.Value)
					lval.Set(rval)
				} else {
					// New variable declaration. Create new variable with the right value.
					typ := rval.Type()
					sv := reflect.New(typ).Elem()
					sv.Set(rval)
					obj.Value = sv
					// Add the name we're declaring to env.names if it's a new name
					if _, ok := env.lookup(ident.Name); !ok {
						env.names = append(env.names, ident.Name)
					}
					// Add the object to env.objs
					env.objs[ident.Name] = obj
				}
			}
		case token.ASSIGN:
			// assignment
			// TODO: this only covers very simple assignment. There are more complicated rules
			//   not yet implemented (see http://golang.org/ref/spec#Assignments), such as:
			//   1) nil
			//   2) Blank identifier
			lhs := make([]Object, len(stmt.Lhs))
			rhs := make([]reflect.Value, len(stmt.Lhs))
			for i, expr := range stmt.Lhs {
				lhs[i] = env.Eval(expr)[0]
			}
			if len(stmt.Rhs) == 1 {
				// Single rhs expression, potentially multi-valued
				robjs := env.Eval(stmt.Rhs[0])
				for i, robj := range robjs {
					rv, ok := robj.Value.(reflect.Value)
					if !ok {
						// This means that rval is untyped.
						// In assignment context, this should only happen if it's untyped nil
						rhs[i] = reflect.ValueOf(nil)
					}
					rhs[i] = rv
				}
			} else {
				// Multiple rhs expressions, each single-valued
				for i, expr := range stmt.Rhs {
					robj := env.Eval(expr)[0]
					rv, ok := robj.Value.(reflect.Value)
					if !ok {
						rhs[i] = reflect.ValueOf(nil)
					}
					rhs[i] = rv
				}
			}

			for i, obj := range lhs {
				v := obj.Value.(reflect.Value)
				if rhs[i].IsValid() {
					v.Set(rhs[i])
				} else {
					v.Set(reflect.Zero(v.Type()))
				}
			}
		}
	case *ast.ExprStmt:
		// If we're not at top level, then only call expressions and receive operations are valid statements
		if !topLevel {
			if _, ok := stmt.X.(*ast.CallExpr); !ok {
				// TODO: handle this error better
				log.Fatal("Expression used as statement inappropriately")
			}
			// TODO: what about receive operations?
		}
		objs := env.Eval(stmt.X)
		if topLevel {
			for _, obj := range objs {
				// TODO: do something better than print the results to stdout
				switch v := obj.Value.(type) {
				case reflect.Value:
					fmt.Printf("=> %s: %v\n", TypeString(obj.Typ), v.Interface())
				case exact.Value, nil:
					fmt.Printf("=> %s: %v\n", TypeString(obj.Typ), v)
				}
			}
		}
	case *ast.GoStmt:
		env.evalFuncCall(stmt.Call, true)
	}
}
<|endoftext|>


<commit_before>package interactive

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"
	"go/parser"
)

func readline(prompt string, in *bufio.Reader) (string, error) {
	fmt.Printf(prompt)
	line, err := in.ReadString('\n')
	if err == nil {
		line = strings.TrimRight(line, "r\n")
	}
	return line, err
}

func Run(env *Env, results *([]interface{})) {

	var err error
	exprs := 0
	in := bufio.NewReader(os.Stdin)
	line, err := readline("go> ", in)
	for line != "quit" {
		if err != nil {
			if err == io.EOF { break }
			panic(err)
		}
		//line = "func() {" + line + "}"
		if expr, err := parser.ParseExpr(line); err != nil {
			fmt.Printf("parse error: %s\n", err)
		} else if vals, _, err := EvalExpr(expr, env); err != nil {
			fmt.Printf("eval error: %s\n", err)
		} else if vals == nil {
			fmt.Printf("nil\n")
		} else if len(*vals) == 0 {
			fmt.Printf("void\n")
		} else if len(*vals) == 1 {
			value := (*vals)[0]
			kind := value.Kind().String()
			fmt.Printf("Kind = %v\n", kind)
			if kind == "string" {
				fmt.Printf("Results[%d] = \"%v\"\n", exprs, (value.Interface()))
			} else {
				fmt.Printf("Results[%d] = %v\n", exprs, (value.Interface()))
			}
			exprs  += 1
			*results = append(*results, (*vals)[0].Interface())
		} else {
			sep := "("
			for _, v := range *vals {
				fmt.Printf("%s%v", sep, v.Interface())
			}
			fmt.Printf(")\n")
		}

		line, err = readline("go> ", in)
	}
	//WriteHistory("data/deleteme.history")
	// rl.Rl_reset_terminal(term)
}
<commit_msg>r should have been \r in TrimRight<commit_after>package interactive

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"
	"go/parser"
)

func readline(prompt string, in *bufio.Reader) (string, error) {
	fmt.Printf(prompt)
	line, err := in.ReadString('\n')
	if err == nil {
		line = strings.TrimRight(line, "\r\n")
	}
	return line, err
}

func Run(env *Env, results *([]interface{})) {

	var err error
	exprs := 0
	in := bufio.NewReader(os.Stdin)
	line, err := readline("go> ", in)
	for line != "quit" {
		if err != nil {
			if err == io.EOF { break }
			panic(err)
		}
		//line = "func() {" + line + "}"
		if expr, err := parser.ParseExpr(line); err != nil {
			fmt.Printf("parse error: %s\n", err)
		} else if vals, _, err := EvalExpr(expr, env); err != nil {
			fmt.Printf("eval error: %s\n", err)
		} else if vals == nil {
			fmt.Printf("nil\n")
		} else if len(*vals) == 0 {
			fmt.Printf("void\n")
		} else if len(*vals) == 1 {
			value := (*vals)[0]
			kind := value.Kind().String()
			fmt.Printf("Kind = %v\n", kind)
			if kind == "string" {
				fmt.Printf("Results[%d] = \"%v\"\n", exprs, (value.Interface()))
			} else {
				fmt.Printf("Results[%d] = %v\n", exprs, (value.Interface()))
			}
			exprs  += 1
			*results = append(*results, (*vals)[0].Interface())
		} else {
			sep := "("
			for _, v := range *vals {
				fmt.Printf("%s%v", sep, v.Interface())
			}
			fmt.Printf(")\n")
		}

		line, err = readline("go> ", in)
	}
	//WriteHistory("data/deleteme.history")
	// rl.Rl_reset_terminal(term)
}
<|endoftext|>


<commit_before><commit_msg>Syncing repos<commit_after><|endoftext|>


<commit_before>package memory

import (
	"errors"
	"math"
	"regexp"
	"strings"

	"github.com/grafana/metrictank/idx"
)

var (
	expressionRe    = regexp.MustCompile("^([^!=~]+)([!=~]{1,3})(.*)$")
	errUnknownId    = errors.New("id does not exist")
	errInvalidTag   = errors.New("found an invalid tag in the index")
	errInvalidQuery = errors.New("invalid query")
)

type kv struct {
	key   string
	value string
}

const (
	PARSING_ERROR = iota
	EQUAL
	NOT_EQUAL
	MATCH
	NOT_MATCH
)

type expression struct {
	kv
	operator int
}

type TagQuery struct {
	equal    []kv
	match    []kv
	notEqual []kv
	notMatch []kv
}

// returns expression,
// in case of error the operator will be PARSING_ERROR
func parseExpression(expr string) expression {
	var key []byte
	regex, not := false, false
	var pos int

	// get key
	for ; pos < len(expr); pos++ {
		// ! || =
		if expr[pos] == 33 || expr[pos] == 61 {
			// key must not be empty
			if len(key) == 0 {
				return expression{operator: PARSING_ERROR}
			}
			break
		}

		// a-z A-Z 0-9
		if !(expr[pos] >= 97 && expr[pos] <= 122) && !(expr[pos] >= 65 && expr[pos] <= 90) && !(expr[pos] >= 48 && expr[pos] <= 57) {
			return expression{operator: PARSING_ERROR}
		}

		key = append(key, expr[pos])
	}

	// if !
	if expr[pos] == 33 {
		not = true
		pos++
	}

	// expecting a =
	if expr[pos] != 61 {
		return expression{operator: PARSING_ERROR}
	}
	pos++

	// if ~
	if len(expr) > pos && expr[pos] == 126 {
		regex = true
		pos++
	}

	var value []byte
	for ; pos < len(expr); pos++ {
		// enforce a-z A-Z 0-9 -_ if query is not regex
		if !regex && !(expr[pos] >= 97 && expr[pos] <= 122) && !(expr[pos] >= 65 && expr[pos] <= 90) && !(expr[pos] >= 48 && expr[pos] <= 57) && expr[pos] != 45 && expr[pos] != 95 {
			return expression{operator: PARSING_ERROR}
		}

		value = append(value, expr[pos])
	}

	if not {
		if regex {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: NOT_MATCH}
		} else {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: NOT_EQUAL}
		}
	} else {
		if regex {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: MATCH}
		} else {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: EQUAL}
		}
	}
}

func NewTagQuery(expressions []string) (TagQuery, error) {
	q := TagQuery{}

	if len(expressions) == 0 {
		return q, errInvalidQuery
	}

	for _, expr := range expressions {
		e := parseExpression(expr)
		if e.operator == PARSING_ERROR {
			return q, errInvalidQuery
		}

		// special case of empty value
		if len(e.value) == 0 {
			if e.operator == EQUAL || e.operator == MATCH {
				q.notMatch = append(q.notMatch, kv{
					key:   e.key,
					value: "^.+",
				})
			} else {
				q.match = append(q.match, kv{
					key:   e.key,
					value: "^.+",
				})
			}
		} else {
			// always anchor all regular expressions at the beginning
			if (e.operator == MATCH || e.operator == NOT_MATCH) && len(e.value) > 0 && e.value[0] != byte('^') {
				e.value = "^" + e.value
			}

			switch e.operator {
			case EQUAL:
				q.equal = append(q.equal, e.kv)
			case NOT_EQUAL:
				q.notEqual = append(q.notEqual, e.kv)
			case MATCH:
				q.match = append(q.match, e.kv)
			case NOT_MATCH:
				q.notMatch = append(q.notMatch, e.kv)
			}
		}
	}

	return q, nil
}

func (q *TagQuery) getInitialByMatch(index TagIndex) (int, map[string]struct{}, error) {
	resultSet := make(map[string]struct{})
	lowestCount := math.MaxUint32
	startId := 0

	// choose key that has the smallest number of values
	for i := range q.match {
		l := len(index[q.match[i].key])
		if l < lowestCount {
			lowestCount = l
			startId = i
		}
	}

	re, err := regexp.Compile(q.match[startId].value)
	if err != nil {
		return 0, nil, errInvalidQuery
	}

	for v, ids := range index[q.match[startId].key] {
		if !re.MatchString(v) {
			continue
		}

		for id := range ids {
			resultSet[id] = struct{}{}
		}
	}
	return startId, resultSet, nil
}

func (q *TagQuery) getInitialByEqual(index TagIndex) (int, map[string]struct{}) {
	resultSet := make(map[string]struct{})
	lowestCount := math.MaxUint32
	startId := 0

	// choose key-value combo that has the smallest number of ids
	for i := range q.equal {
		l := len(index[q.equal[i].key][q.equal[i].value])
		if l < lowestCount {
			lowestCount = l
			startId = i
		}
	}

	// copy the map, because we'll later delete items from it
	for k, v := range index[q.equal[startId].key][q.equal[startId].value] {
		resultSet[k] = v
	}

	return startId, resultSet
}

func (q *TagQuery) filterByEqual(skipEqual int, resultSet map[string]struct{}, index TagIndex) {
	for i, e := range q.equal {
		if i == skipEqual {
			continue
		}

		for id := range resultSet {
			if _, ok := index[e.key][e.value][id]; !ok {
				delete(resultSet, id)
			}
		}
	}
}

func (q *TagQuery) filterByNotEqual(resultSet map[string]struct{}, index TagIndex, byId map[string]*idx.Archive) {
	for _, e := range q.notEqual {
		fullTag := e.key + "=" + e.value
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}
			for _, tag := range def.Tags {
				if tag == fullTag {
					delete(resultSet, id)
					continue IDS
				}
			}
		}
	}
}

func (q *TagQuery) filterByMatch(skipMatch int, resultSet map[string]struct{}, byId map[string]*idx.Archive) error {
	for i, e := range q.match {
		if i == skipMatch {
			continue
		}

		re, err := regexp.Compile(e.value)
		if err != nil {
			return errInvalidQuery
		}

		matchingTags := make(map[string]struct{})
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}
			for _, tag := range def.Tags {
				// optimization to reduce regex matching
				if _, ok := matchingTags[tag]; ok {
					continue IDS
				}
			}

			for _, tag := range def.Tags {
				tagSplits := strings.SplitN(tag, "=", 2)
				if len(tagSplits) != 2 {
					// corrupt index
					delete(resultSet, id)
					continue IDS
				}

				if e.key == tagSplits[0] && re.MatchString(tagSplits[1]) {
					matchingTags[tag] = struct{}{}
					continue IDS
				}
			}
			delete(resultSet, id)
		}
	}
	return nil
}

func (q *TagQuery) filterByNotMatch(skipMatch int, resultSet map[string]struct{}, byId map[string]*idx.Archive) error {
	for _, e := range q.notMatch {
		re, err := regexp.Compile(e.value)
		if err != nil {
			return errInvalidQuery
		}

		matchingTags := make(map[string]struct{})
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}

			// optimization to reduce regex matching
			for _, tag := range def.Tags {
				if _, ok := matchingTags[tag]; ok {
					delete(resultSet, id)
					continue IDS
				}
			}

			for _, tag := range def.Tags {
				tagSplits := strings.SplitN(tag, "=", 2)
				if len(tagSplits) != 2 {
					// corrupt index
					delete(resultSet, id)
					continue IDS
				}

				if e.key == tagSplits[0] && re.MatchString(tagSplits[1]) {
					matchingTags[tag] = struct{}{}
					delete(resultSet, id)
					continue IDS
				}
			}
		}
	}
	return nil
}

func (q *TagQuery) Run(index TagIndex, byId map[string]*idx.Archive) (map[string]struct{}, error) {
	var skipMatch, skipEqual = -1, -1
	var resultSet map[string]struct{}
	var err error

	// find the best expression to start with
	if len(q.equal) == 0 {
		if len(q.match) == 0 {
			return nil, errInvalidQuery
		}

		skipMatch, resultSet, err = q.getInitialByMatch(index)
		if err != nil {
			return nil, errInvalidQuery
		}
	} else {
		skipEqual, resultSet = q.getInitialByEqual(index)
	}

	q.filterByEqual(skipEqual, resultSet, index)
	q.filterByNotEqual(resultSet, index, byId)
	err = q.filterByMatch(skipMatch, resultSet, byId)
	if err != nil {
		return nil, err
	}
	err = q.filterByNotMatch(skipMatch, resultSet, byId)
	if err != nil {
		return nil, err
	}

	return resultSet, nil
}
<commit_msg>add shortcuts if regex is not necessary<commit_after>package memory

import (
	"errors"
	"math"
	"regexp"
	"strings"

	"github.com/grafana/metrictank/idx"
)

var (
	expressionRe    = regexp.MustCompile("^([^!=~]+)([!=~]{1,3})(.*)$")
	errUnknownId    = errors.New("id does not exist")
	errInvalidTag   = errors.New("found an invalid tag in the index")
	errInvalidQuery = errors.New("invalid query")
)

type kv struct {
	key   string
	value string
}

const (
	PARSING_ERROR = iota
	EQUAL
	NOT_EQUAL
	MATCH
	NOT_MATCH
)

type expression struct {
	kv
	operator int
}

type TagQuery struct {
	equal    []kv
	match    []kv
	notEqual []kv
	notMatch []kv
}

// returns expression,
// in case of error the operator will be PARSING_ERROR
func parseExpression(expr string) expression {
	var key []byte
	regex, not := false, false
	var pos int

	// get key
	for ; pos < len(expr); pos++ {
		// ! || =
		if expr[pos] == 33 || expr[pos] == 61 {
			// key must not be empty
			if len(key) == 0 {
				return expression{operator: PARSING_ERROR}
			}
			break
		}

		// a-z A-Z 0-9
		if !(expr[pos] >= 97 && expr[pos] <= 122) && !(expr[pos] >= 65 && expr[pos] <= 90) && !(expr[pos] >= 48 && expr[pos] <= 57) {
			return expression{operator: PARSING_ERROR}
		}

		key = append(key, expr[pos])
	}

	// if !
	if expr[pos] == 33 {
		not = true
		pos++
	}

	// expecting a =
	if expr[pos] != 61 {
		return expression{operator: PARSING_ERROR}
	}
	pos++

	// if ~
	if len(expr) > pos && expr[pos] == 126 {
		regex = true
		pos++
	}

	var value []byte
	for ; pos < len(expr); pos++ {
		// enforce a-z A-Z 0-9 -_ if query is not regex
		if !regex && !(expr[pos] >= 97 && expr[pos] <= 122) && !(expr[pos] >= 65 && expr[pos] <= 90) && !(expr[pos] >= 48 && expr[pos] <= 57) && expr[pos] != 45 && expr[pos] != 95 {
			return expression{operator: PARSING_ERROR}
		}

		value = append(value, expr[pos])
	}

	if not {
		if regex {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: NOT_MATCH}
		} else {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: NOT_EQUAL}
		}
	} else {
		if regex {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: MATCH}
		} else {
			return expression{kv: kv{key: string(key), value: string(value)}, operator: EQUAL}
		}
	}
}

func NewTagQuery(expressions []string) (TagQuery, error) {
	q := TagQuery{}

	if len(expressions) == 0 {
		return q, errInvalidQuery
	}

	for _, expr := range expressions {
		e := parseExpression(expr)
		if e.operator == PARSING_ERROR {
			return q, errInvalidQuery
		}

		// special case of empty value
		if len(e.value) == 0 {
			if e.operator == EQUAL || e.operator == MATCH {
				q.notMatch = append(q.notMatch, kv{
					key:   e.key,
					value: "^.+",
				})
			} else {
				q.match = append(q.match, kv{
					key:   e.key,
					value: "^.+",
				})
			}
		} else {
			// always anchor all regular expressions at the beginning
			if (e.operator == MATCH || e.operator == NOT_MATCH) && len(e.value) > 0 && e.value[0] != byte('^') {
				e.value = "^" + e.value
			}

			switch e.operator {
			case EQUAL:
				q.equal = append(q.equal, e.kv)
			case NOT_EQUAL:
				q.notEqual = append(q.notEqual, e.kv)
			case MATCH:
				q.match = append(q.match, e.kv)
			case NOT_MATCH:
				q.notMatch = append(q.notMatch, e.kv)
			}
		}
	}

	return q, nil
}

func (q *TagQuery) getInitialByMatch(index TagIndex) (int, map[string]struct{}, error) {
	resultSet := make(map[string]struct{})
	lowestCount := math.MaxUint32
	startId := 0

	// choose key that has the smallest number of values
	for i := range q.match {
		l := len(index[q.match[i].key])
		if l < lowestCount {
			lowestCount = l
			startId = i
		}
	}

	// shortcut if pattern is ^.+ (f.e. expression "key!=~" will be translated to "key=~^.+")
	if q.match[startId].value == "^.+" {
		for _, ids := range index[q.match[startId].key] {
			for id := range ids {
				resultSet[id] = struct{}{}
			}
		}
		return startId, resultSet, nil
	}

	re, err := regexp.Compile(q.match[startId].value)
	if err != nil {
		return 0, nil, errInvalidQuery
	}

	for v, ids := range index[q.match[startId].key] {
		if !re.MatchString(v) {
			continue
		}

		for id := range ids {
			resultSet[id] = struct{}{}
		}
	}
	return startId, resultSet, nil
}

func (q *TagQuery) getInitialByEqual(index TagIndex) (int, map[string]struct{}) {
	resultSet := make(map[string]struct{})
	lowestCount := math.MaxUint32
	startId := 0

	// choose key-value combo that has the smallest number of ids
	for i := range q.equal {
		l := len(index[q.equal[i].key][q.equal[i].value])
		if l < lowestCount {
			lowestCount = l
			startId = i
		}
	}

	// copy the map, because we'll later delete items from it
	for k, v := range index[q.equal[startId].key][q.equal[startId].value] {
		resultSet[k] = v
	}

	return startId, resultSet
}

func (q *TagQuery) filterByEqual(skipEqual int, resultSet map[string]struct{}, index TagIndex) {
	for i, e := range q.equal {
		if i == skipEqual {
			continue
		}

		for id := range resultSet {
			if _, ok := index[e.key][e.value][id]; !ok {
				delete(resultSet, id)
			}
		}
	}
}

func (q *TagQuery) filterByNotEqual(resultSet map[string]struct{}, index TagIndex, byId map[string]*idx.Archive) {
	for _, e := range q.notEqual {
		fullTag := e.key + "=" + e.value
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}
			for _, tag := range def.Tags {
				if tag == fullTag {
					delete(resultSet, id)
					continue IDS
				}
			}
		}
	}
}

func (q *TagQuery) filterByMatch(skipMatch int, resultSet map[string]struct{}, byId map[string]*idx.Archive) error {
	for i, e := range q.match {
		if i == skipMatch {
			continue
		}

		var shortCut bool
		var re *regexp.Regexp
		var err error

		// shortcut if pattern is ^.+ (f.e. expression "key!=" will be translated to "key=~^.+")
		if e.value == "^.+" {
			shortCut = true
		} else {
			re, err = regexp.Compile(e.value)
			if err != nil {
				return errInvalidQuery
			}
		}

		matchingTags := make(map[string]struct{})
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}

			for _, tag := range def.Tags {
				// optimization to reduce regex matching
				if _, ok := matchingTags[tag]; ok {
					continue IDS
				}
			}

			for _, tag := range def.Tags {
				tagSplits := strings.SplitN(tag, "=", 2)
				if len(tagSplits) != 2 {
					// corrupt index
					delete(resultSet, id)
					continue IDS
				}

				if e.key == tagSplits[0] && (shortCut || re.MatchString(tagSplits[1])) {
					matchingTags[tag] = struct{}{}
					continue IDS
				}
			}
			delete(resultSet, id)
		}
	}
	return nil
}

func (q *TagQuery) filterByNotMatch(skipMatch int, resultSet map[string]struct{}, byId map[string]*idx.Archive) error {
	for _, e := range q.notMatch {
		re, err := regexp.Compile(e.value)
		if err != nil {
			return errInvalidQuery
		}

		matchingTags := make(map[string]struct{})
	IDS:
		for id := range resultSet {
			var def *idx.Archive
			var ok bool
			if def, ok = byId[id]; !ok {
				// corrupt index
				delete(resultSet, id)
				continue IDS
			}

			// optimization to reduce regex matching
			for _, tag := range def.Tags {
				if _, ok := matchingTags[tag]; ok {
					delete(resultSet, id)
					continue IDS
				}
			}

			for _, tag := range def.Tags {
				tagSplits := strings.SplitN(tag, "=", 2)
				if len(tagSplits) != 2 {
					// corrupt index
					delete(resultSet, id)
					continue IDS
				}

				if e.key == tagSplits[0] && re.MatchString(tagSplits[1]) {
					matchingTags[tag] = struct{}{}
					delete(resultSet, id)
					continue IDS
				}
			}
		}
	}
	return nil
}

func (q *TagQuery) Run(index TagIndex, byId map[string]*idx.Archive) (map[string]struct{}, error) {
	var skipMatch, skipEqual = -1, -1
	var resultSet map[string]struct{}
	var err error

	// find the best expression to start with
	if len(q.equal) == 0 {
		if len(q.match) == 0 {
			return nil, errInvalidQuery
		}

		skipMatch, resultSet, err = q.getInitialByMatch(index)
		if err != nil {
			return nil, errInvalidQuery
		}
	} else {
		skipEqual, resultSet = q.getInitialByEqual(index)
	}

	q.filterByEqual(skipEqual, resultSet, index)
	q.filterByNotEqual(resultSet, index, byId)
	err = q.filterByMatch(skipMatch, resultSet, byId)
	if err != nil {
		return nil, err
	}
	err = q.filterByNotMatch(skipMatch, resultSet, byId)
	if err != nil {
		return nil, err
	}

	return resultSet, nil
}
<|endoftext|>


<commit_before>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	return sendSample(span)
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err, ok := r.(error)
			if !ok {
				err = fmt.Errorf("encountered panic while sending sample %#v", err)
			}
			log.WithError(err).Error("Panicked while dialing UDP connection - traces will not be sent")
		}
	}()

	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)

	// at this point, the connection should already be initialized

	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<commit_msg>Fix syntax error<commit_after>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	return sendSample(span)
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err, ok := r.(error)
			if !ok {
				err = fmt.Errorf("encountered panic while sending sample %#v", err)
			}
			logrus.WithError(err).Error("Panicked while dialing UDP connection - traces will not be sent")
		}
	}()

	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)

	// at this point, the connection should already be initialized

	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"bytes"
	"database/sql"
	"log"
	"regexp"
	"testing"

	"github.com/mattn/go-sqlite3"
	"github.com/shogo82148/txmanager"
)

func TestTraceProxy(t *testing.T) {
	buf := &bytes.Buffer{}
	logger := log.New(buf, "", log.Lshortfile)
	sql.Register("sqlite3-trace-proxy", NewTraceProxy(&sqlite3.SQLiteDriver{}, logger))

	db, err := sql.Open("sqlite3-trace-proxy", ":memory:")
	if err != nil {
		t.Fatalf("Open filed: %v", err)
	}

	_, err = db.Exec(
		"CREATE TABLE t1 (id INTEGER PRIMARY KEY)",
	)
	if err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	dbm := txmanager.NewDB(db)
	txmanager.Do(dbm, func(tx txmanager.Tx) error {
		_, err := tx.Exec("INSERT INTO t1 (id) VALUES(?)", 1)
		return err
	})
	if err != nil {
		t.Fatalf("do failed: %v", err)
	}

	row := dbm.QueryRow("SELECT id FROM t1 WHERE id = ?", 1)
	var id int
	if err = row.Scan(&id); err != nil {
		t.Fatalf("selecting row failed: %v", err)
	}
	if id != 1 {
		t.Errorf("got %d\nwant 1", id)
	}

	expected := []*regexp.Regexp{
		// Fake time compinent with (\d+\.\d+[^\)]+)
		regexp.MustCompile(`tracer_test.go:27: Open \(\d+\.\d+[^\)]+\)`),
		regexp.MustCompile(`tracer_test.go:27: Exec: CREATE TABLE t1 \(id INTEGER PRIMARY KEY\); args = \[\] \(\d+\.\d+[^\)]+\)`),
		// The line numbers on these two lines may change depending on the version
		// of txmanager that you have. For now, we cross our fingers that the
		// filename doesn't change ever, but the line numbers may change.
		regexp.MustCompile(`txmanager.go:\d+: Begin`),
		regexp.MustCompile(`txmanager.go:\d+: Exec: INSERT INTO t1 \(id\) VALUES\(\?\); args = \[1\] \(\d+\.\d+[^\)]+\)`),
		regexp.MustCompile(`tracer_test.go:36: Commit`),
		regexp.MustCompile(`tracer_test.go:41: Query: SELECT id FROM t1 WHERE id = \?; args = \[1\] \(\d+\.\d+[^\)]+\)`),
	}

	scanner := bufio.NewScanner(buf)
	i := 0
	for scanner.Scan() {
		line := scanner.Text()
		if i >= len(expected) {
			t.Errorf("Got more lines than expected (%s)", line)
			break
		}

		if !expected[i].MatchString(line) {
			t.Errorf("\ngot: %s\nwant: %s", line, expected[i])
		}
		i++
	}
}
<commit_msg>elapsed time might be integer<commit_after>package proxy

import (
	"bufio"
	"bytes"
	"database/sql"
	"log"
	"regexp"
	"testing"

	"github.com/mattn/go-sqlite3"
	"github.com/shogo82148/txmanager"
)

func TestTraceProxy(t *testing.T) {
	buf := &bytes.Buffer{}
	logger := log.New(buf, "", log.Lshortfile)
	sql.Register("sqlite3-trace-proxy", NewTraceProxy(&sqlite3.SQLiteDriver{}, logger))

	db, err := sql.Open("sqlite3-trace-proxy", ":memory:")
	if err != nil {
		t.Fatalf("Open filed: %v", err)
	}

	_, err = db.Exec(
		"CREATE TABLE t1 (id INTEGER PRIMARY KEY)",
	)
	if err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	dbm := txmanager.NewDB(db)
	txmanager.Do(dbm, func(tx txmanager.Tx) error {
		_, err := tx.Exec("INSERT INTO t1 (id) VALUES(?)", 1)
		return err
	})
	if err != nil {
		t.Fatalf("do failed: %v", err)
	}

	row := dbm.QueryRow("SELECT id FROM t1 WHERE id = ?", 1)
	var id int
	if err = row.Scan(&id); err != nil {
		t.Fatalf("selecting row failed: %v", err)
	}
	if id != 1 {
		t.Errorf("got %d\nwant 1", id)
	}

	expected := []*regexp.Regexp{
		// Fake time compinent with (\d+\.\d+[^\)]+)
		regexp.MustCompile(`tracer_test.go:27: Open \(\d+(?:\.\d+)?[^\)]+\)`),
		regexp.MustCompile(`tracer_test.go:27: Exec: CREATE TABLE t1 \(id INTEGER PRIMARY KEY\); args = \[\] \(\d+(?:\.\d+)?[^\)]+\)`),
		// The line numbers on these two lines may change depending on the version
		// of txmanager that you have. For now, we cross our fingers that the
		// filename doesn't change ever, but the line numbers may change.
		regexp.MustCompile(`txmanager.go:\d+: Begin`),
		regexp.MustCompile(`txmanager.go:\d+: Exec: INSERT INTO t1 \(id\) VALUES\(\?\); args = \[1\] \(\d+(?:\.\d+)?[^\)]+\)`),
		regexp.MustCompile(`tracer_test.go:36: Commit`),
		regexp.MustCompile(`tracer_test.go:41: Query: SELECT id FROM t1 WHERE id = \?; args = \[1\] \(\d+(?:\.\d+)?[^\)]+\)`),
	}

	scanner := bufio.NewScanner(buf)
	i := 0
	for scanner.Scan() {
		line := scanner.Text()
		if i >= len(expected) {
			t.Errorf("Got more lines than expected (%s)", line)
			break
		}

		if !expected[i].MatchString(line) {
			t.Errorf("\ngot: %s\nwant: %s", line, expected[i])
		}
		i++
	}
}
<|endoftext|>


<commit_before>// +build !confonly

package core

import (
	"context"
	"reflect"
	"sync"

	"v2ray.com/core/common"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/features"
	"v2ray.com/core/features/dns"
	"v2ray.com/core/features/dns/localdns"
	"v2ray.com/core/features/inbound"
	"v2ray.com/core/features/outbound"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/features/stats"
)

// Server is an instance of V2Ray. At any time, there must be at most one Server instance running.
type Server interface {
	common.Runnable
}

// ServerType returns the type of the server.
func ServerType() interface{} {
	return (*Instance)(nil)
}

type resolution struct {
	deps     []reflect.Type
	callback interface{}
}

func getFeature(allFeatures []features.Feature, t reflect.Type) features.Feature {
	for _, f := range allFeatures {
		if reflect.TypeOf(f.Type()) == t {
			return f
		}
	}
	return nil
}

func (r *resolution) resolve(allFeatures []features.Feature) (bool, error) {
	var fs []features.Feature
	for _, d := range r.deps {
		f := getFeature(allFeatures, d)
		if f == nil {
			return false, nil
		}
		fs = append(fs, f)
	}

	callback := reflect.ValueOf(r.callback)
	var input []reflect.Value
	callbackType := callback.Type()
	for i := 0; i < callbackType.NumIn(); i++ {
		pt := callbackType.In(i)
		for _, f := range fs {
			if reflect.TypeOf(f).AssignableTo(pt) {
				input = append(input, reflect.ValueOf(f))
				break
			}
		}
	}

	if len(input) != callbackType.NumIn() {
		panic("Can't get all input parameters")
	}

	var err error
	ret := callback.Call(input)
	errInterface := reflect.TypeOf((*error)(nil)).Elem()
	for i := len(ret) - 1; i >= 0; i-- {
		if ret[i].Type() == errInterface {
			v := ret[i].Interface()
			if v != nil {
				err = v.(error)
			}
			break
		}
	}

	return true, err
}

// Instance combines all functionalities in V2Ray.
type Instance struct {
	access             sync.Mutex
	features           []features.Feature
	featureResolutions []resolution
	running            bool
}

func AddInboundHandler(server *Instance, config *InboundHandlerConfig) error {
	inboundManager := server.GetFeature(inbound.ManagerType()).(inbound.Manager)
	rawHandler, err := CreateObject(server, config)
	if err != nil {
		return err
	}
	handler, ok := rawHandler.(inbound.Handler)
	if !ok {
		return newError("not an InboundHandler")
	}
	if err := inboundManager.AddHandler(context.Background(), handler); err != nil {
		return err
	}
	return nil
}

func addInboundHandlers(server *Instance, configs []*InboundHandlerConfig) error {
	for _, inboundConfig := range configs {
		if err := AddInboundHandler(server, inboundConfig); err != nil {
			return err
		}
	}

	return nil
}

func AddOutboundHandler(server *Instance, config *OutboundHandlerConfig) error {
	outboundManager := server.GetFeature(outbound.ManagerType()).(outbound.Manager)
	rawHandler, err := CreateObject(server, config)
	if err != nil {
		return err
	}
	handler, ok := rawHandler.(outbound.Handler)
	if !ok {
		return newError("not an OutboundHandler")
	}
	if err := outboundManager.AddHandler(context.Background(), handler); err != nil {
		return err
	}
	return nil
}

func addOutboundHandlers(server *Instance, configs []*OutboundHandlerConfig) error {
	for _, outboundConfig := range configs {
		if err := AddOutboundHandler(server, outboundConfig); err != nil {
			return err
		}
	}

	return nil
}

// RequireFeatures is a helper function to require features from Instance in context.
// See Instance.RequireFeatures for more information.
func RequireFeatures(ctx context.Context, callback interface{}) error {
	v := MustFromContext(ctx)
	return v.RequireFeatures(callback)
}

// New returns a new V2Ray instance based on given configuration.
// The instance is not started at this point.
// To ensure V2Ray instance works properly, the config must contain one Dispatcher, one InboundHandlerManager and one OutboundHandlerManager. Other features are optional.
func New(config *Config) (*Instance, error) {
	var server = &Instance{}

	if config.Transport != nil {
		features.PrintDeprecatedFeatureWarning("global transport settings")
	}
	if err := config.Transport.Apply(); err != nil {
		return nil, err
	}

	for _, appSettings := range config.App {
		settings, err := appSettings.GetInstance()
		if err != nil {
			return nil, err
		}
		obj, err := CreateObject(server, settings)
		if err != nil {
			return nil, err
		}
		if feature, ok := obj.(features.Feature); ok {
			if err := server.AddFeature(feature); err != nil {
				return nil, err
			}
		}
	}

	essentialFeatures := []struct {
		Type     interface{}
		Instance features.Feature
	}{
		{dns.ClientType(), localdns.New()},
		{policy.ManagerType(), policy.DefaultManager{}},
		{routing.RouterType(), routing.DefaultRouter{}},
		{stats.ManagerType(), stats.NoopManager{}},
	}

	for _, f := range essentialFeatures {
		if server.GetFeature(f.Type) == nil {
			if err := server.AddFeature(f.Instance); err != nil {
				return nil, err
			}
		}
	}

	if server.featureResolutions != nil {
		return nil, newError("not all dependency are resolved.")
	}

	if err := addInboundHandlers(server, config.Inbound); err != nil {
		return nil, err
	}

	if err := addOutboundHandlers(server, config.Outbound); err != nil {
		return nil, err
	}

	return server, nil
}

// Type implements common.HasType.
func (s *Instance) Type() interface{} {
	return ServerType()
}

// Close shutdown the V2Ray instance.
func (s *Instance) Close() error {
	s.access.Lock()
	defer s.access.Unlock()

	s.running = false

	var errors []interface{}
	for _, f := range s.features {
		if err := f.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		return newError("failed to close all features").Base(newError(serial.Concat(errors...)))
	}

	return nil
}

// RequireFeatures registers a callback, which will be called when all dependent features are registered.
// The callback must be a func(). All its parameters must be features.Feature.
func (s *Instance) RequireFeatures(callback interface{}) error {
	callbackType := reflect.TypeOf(callback)
	if callbackType.Kind() != reflect.Func {
		panic("not a function")
	}

	var featureTypes []reflect.Type
	for i := 0; i < callbackType.NumIn(); i++ {
		featureTypes = append(featureTypes, reflect.PtrTo(callbackType.In(i)))
	}

	r := resolution{
		deps:     featureTypes,
		callback: callback,
	}
	if finished, err := r.resolve(s.features); finished {
		return err
	}
	s.featureResolutions = append(s.featureResolutions, r)
	return nil
}

// AddFeature registers a feature into current Instance.
func (s *Instance) AddFeature(feature features.Feature) error {
	s.features = append(s.features, feature)

	if s.running {
		if err := feature.Start(); err != nil {
			newError("failed to start feature").Base(err).WriteToLog()
		}
		return nil
	}

	if s.featureResolutions == nil {
		return nil
	}

	var pendingResolutions []resolution
	for _, r := range s.featureResolutions {
		finished, err := r.resolve(s.features)
		if finished && err != nil {
			return err
		}
		if !finished {
			pendingResolutions = append(pendingResolutions, r)
		}
	}
	if len(pendingResolutions) == 0 {
		s.featureResolutions = nil
	} else if len(pendingResolutions) < len(s.featureResolutions) {
		s.featureResolutions = pendingResolutions
	}

	return nil
}

// GetFeature returns a feature of the given type, or nil if such feature is not registered.
func (s *Instance) GetFeature(featureType interface{}) features.Feature {
	return getFeature(s.features, reflect.TypeOf(featureType))
}

// Start starts the V2Ray instance, including all registered features. When Start returns error, the state of the instance is unknown.
// A V2Ray instance can be started only once. Upon closing, the instance is not guaranteed to start again.
//
// v2ray:api:stable
func (s *Instance) Start() error {
	s.access.Lock()
	defer s.access.Unlock()

	s.running = true
	for _, f := range s.features {
		if err := f.Start(); err != nil {
			return err
		}
	}

	newError("V2Ray ", Version(), " started").AtWarning().WriteToLog()

	return nil
}
<commit_msg>Allow client to submit a context in v2ray.New<commit_after>// +build !confonly

package core

import (
	"context"
	"reflect"
	"sync"

	"v2ray.com/core/common"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/features"
	"v2ray.com/core/features/dns"
	"v2ray.com/core/features/dns/localdns"
	"v2ray.com/core/features/inbound"
	"v2ray.com/core/features/outbound"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/features/stats"
)

// Server is an instance of V2Ray. At any time, there must be at most one Server instance running.
type Server interface {
	common.Runnable
}

// ServerType returns the type of the server.
func ServerType() interface{} {
	return (*Instance)(nil)
}

type resolution struct {
	deps     []reflect.Type
	callback interface{}
}

func getFeature(allFeatures []features.Feature, t reflect.Type) features.Feature {
	for _, f := range allFeatures {
		if reflect.TypeOf(f.Type()) == t {
			return f
		}
	}
	return nil
}

func (r *resolution) resolve(allFeatures []features.Feature) (bool, error) {
	var fs []features.Feature
	for _, d := range r.deps {
		f := getFeature(allFeatures, d)
		if f == nil {
			return false, nil
		}
		fs = append(fs, f)
	}

	callback := reflect.ValueOf(r.callback)
	var input []reflect.Value
	callbackType := callback.Type()
	for i := 0; i < callbackType.NumIn(); i++ {
		pt := callbackType.In(i)
		for _, f := range fs {
			if reflect.TypeOf(f).AssignableTo(pt) {
				input = append(input, reflect.ValueOf(f))
				break
			}
		}
	}

	if len(input) != callbackType.NumIn() {
		panic("Can't get all input parameters")
	}

	var err error
	ret := callback.Call(input)
	errInterface := reflect.TypeOf((*error)(nil)).Elem()
	for i := len(ret) - 1; i >= 0; i-- {
		if ret[i].Type() == errInterface {
			v := ret[i].Interface()
			if v != nil {
				err = v.(error)
			}
			break
		}
	}

	return true, err
}

// Instance combines all functionalities in V2Ray.
type Instance struct {
	access             sync.Mutex
	features           []features.Feature
	featureResolutions []resolution
	running            bool

	ctx context.Context
}

func AddInboundHandler(server *Instance, config *InboundHandlerConfig) error {
	inboundManager := server.GetFeature(inbound.ManagerType()).(inbound.Manager)
	rawHandler, err := CreateObject(server, config)
	if err != nil {
		return err
	}
	handler, ok := rawHandler.(inbound.Handler)
	if !ok {
		return newError("not an InboundHandler")
	}
	if err := inboundManager.AddHandler(server.ctx, handler); err != nil {
		return err
	}
	return nil
}

func addInboundHandlers(server *Instance, configs []*InboundHandlerConfig) error {
	for _, inboundConfig := range configs {
		if err := AddInboundHandler(server, inboundConfig); err != nil {
			return err
		}
	}

	return nil
}

func AddOutboundHandler(server *Instance, config *OutboundHandlerConfig) error {
	outboundManager := server.GetFeature(outbound.ManagerType()).(outbound.Manager)
	rawHandler, err := CreateObject(server, config)
	if err != nil {
		return err
	}
	handler, ok := rawHandler.(outbound.Handler)
	if !ok {
		return newError("not an OutboundHandler")
	}
	if err := outboundManager.AddHandler(server.ctx, handler); err != nil {
		return err
	}
	return nil
}

func addOutboundHandlers(server *Instance, configs []*OutboundHandlerConfig) error {
	for _, outboundConfig := range configs {
		if err := AddOutboundHandler(server, outboundConfig); err != nil {
			return err
		}
	}

	return nil
}

// RequireFeatures is a helper function to require features from Instance in context.
// See Instance.RequireFeatures for more information.
func RequireFeatures(ctx context.Context, callback interface{}) error {
	v := MustFromContext(ctx)
	return v.RequireFeatures(callback)
}

// New returns a new V2Ray instance based on given configuration.
// The instance is not started at this point.
// To ensure V2Ray instance works properly, the config must contain one Dispatcher, one InboundHandlerManager and one OutboundHandlerManager. Other features are optional.
func New(config *Config) (*Instance, error) {
	var server = &Instance{ctx: context.Background()}

	err, done := initInstanceWithConfig(config, server)
	if done {
		return nil, err
	}

	return server, nil
}

func NewWithContext(config *Config, ctx context.Context) (*Instance, error) {
	var server = &Instance{ctx: ctx}

	err, done := initInstanceWithConfig(config, server)
	if done {
		return nil, err
	}

	return server, nil
}

func initInstanceWithConfig(config *Config, server *Instance) (error, bool) {
	if config.Transport != nil {
		features.PrintDeprecatedFeatureWarning("global transport settings")
	}
	if err := config.Transport.Apply(); err != nil {
		return err, true
	}

	for _, appSettings := range config.App {
		settings, err := appSettings.GetInstance()
		if err != nil {
			return err, true
		}
		obj, err := CreateObject(server, settings)
		if err != nil {
			return err, true
		}
		if feature, ok := obj.(features.Feature); ok {
			if err := server.AddFeature(feature); err != nil {
				return err, true
			}
		}
	}

	essentialFeatures := []struct {
		Type     interface{}
		Instance features.Feature
	}{
		{dns.ClientType(), localdns.New()},
		{policy.ManagerType(), policy.DefaultManager{}},
		{routing.RouterType(), routing.DefaultRouter{}},
		{stats.ManagerType(), stats.NoopManager{}},
	}

	for _, f := range essentialFeatures {
		if server.GetFeature(f.Type) == nil {
			if err := server.AddFeature(f.Instance); err != nil {
				return err, true
			}
		}
	}

	if server.featureResolutions != nil {
		return newError("not all dependency are resolved."), true
	}

	if err := addInboundHandlers(server, config.Inbound); err != nil {
		return err, true
	}

	if err := addOutboundHandlers(server, config.Outbound); err != nil {
		return err, true
	}
	return nil, false
}

// Type implements common.HasType.
func (s *Instance) Type() interface{} {
	return ServerType()
}

// Close shutdown the V2Ray instance.
func (s *Instance) Close() error {
	s.access.Lock()
	defer s.access.Unlock()

	s.running = false

	var errors []interface{}
	for _, f := range s.features {
		if err := f.Close(); err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		return newError("failed to close all features").Base(newError(serial.Concat(errors...)))
	}

	return nil
}

// RequireFeatures registers a callback, which will be called when all dependent features are registered.
// The callback must be a func(). All its parameters must be features.Feature.
func (s *Instance) RequireFeatures(callback interface{}) error {
	callbackType := reflect.TypeOf(callback)
	if callbackType.Kind() != reflect.Func {
		panic("not a function")
	}

	var featureTypes []reflect.Type
	for i := 0; i < callbackType.NumIn(); i++ {
		featureTypes = append(featureTypes, reflect.PtrTo(callbackType.In(i)))
	}

	r := resolution{
		deps:     featureTypes,
		callback: callback,
	}
	if finished, err := r.resolve(s.features); finished {
		return err
	}
	s.featureResolutions = append(s.featureResolutions, r)
	return nil
}

// AddFeature registers a feature into current Instance.
func (s *Instance) AddFeature(feature features.Feature) error {
	s.features = append(s.features, feature)

	if s.running {
		if err := feature.Start(); err != nil {
			newError("failed to start feature").Base(err).WriteToLog()
		}
		return nil
	}

	if s.featureResolutions == nil {
		return nil
	}

	var pendingResolutions []resolution
	for _, r := range s.featureResolutions {
		finished, err := r.resolve(s.features)
		if finished && err != nil {
			return err
		}
		if !finished {
			pendingResolutions = append(pendingResolutions, r)
		}
	}
	if len(pendingResolutions) == 0 {
		s.featureResolutions = nil
	} else if len(pendingResolutions) < len(s.featureResolutions) {
		s.featureResolutions = pendingResolutions
	}

	return nil
}

// GetFeature returns a feature of the given type, or nil if such feature is not registered.
func (s *Instance) GetFeature(featureType interface{}) features.Feature {
	return getFeature(s.features, reflect.TypeOf(featureType))
}

// Start starts the V2Ray instance, including all registered features. When Start returns error, the state of the instance is unknown.
// A V2Ray instance can be started only once. Upon closing, the instance is not guaranteed to start again.
//
// v2ray:api:stable
func (s *Instance) Start() error {
	s.access.Lock()
	defer s.access.Unlock()

	s.running = true
	for _, f := range s.features {
		if err := f.Start(); err != nil {
			return err
		}
	}

	newError("V2Ray ", Version(), " started").AtWarning().WriteToLog()

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"hu"
	"fmt"
	"os"
	"bufio"
	"flag"
	"io"
	"log"
)

func main() {
	filenameFlag := flag.String("filename", "-", "filename from which to read and execute program")
	flag.Parse()
	filename := *filenameFlag

	var reader io.RuneScanner
	if filename == "-" {
		reader = bufio.NewReader(os.Stdin)
	} else {
		f, err := os.Open(filename)
		if err != nil {
			log.Fatalln(err)
		}
		reader = bufio.NewReader(f)
	}

	environment := hu.NewEnvironment()
	hu.AddDefaultBindings(environment)

	for {
		fmt.Printf("hu> ")
		expression := hu.Read(reader)
		if expression != nil {
			result := environment.Evaluate(expression)
			if result != nil {
				fmt.Fprintf(os.Stdout, "%v\n", result)
			}
		} else {
			fmt.Fprintf(os.Stdout, "Goodbye!\n")
			break
		}
	}
}
<commit_msg>changed path to its remote import path syntax<commit_after>package main

import (
	"github.com/eikeon/hu"
	"fmt"
	"os"
	"bufio"
	"flag"
	"io"
	"log"
)

func main() {
	filenameFlag := flag.String("filename", "-", "filename from which to read and execute program")
	flag.Parse()
	filename := *filenameFlag

	var reader io.RuneScanner
	if filename == "-" {
		reader = bufio.NewReader(os.Stdin)
	} else {
		f, err := os.Open(filename)
		if err != nil {
			log.Fatalln(err)
		}
		reader = bufio.NewReader(f)
	}

	environment := hu.NewEnvironment()
	hu.AddDefaultBindings(environment)

	for {
		fmt.Printf("hu> ")
		expression := hu.Read(reader)
		if expression != nil {
			result := environment.Evaluate(expression)
			if result != nil {
				fmt.Fprintf(os.Stdout, "%v\n", result)
			}
		} else {
			fmt.Fprintf(os.Stdout, "Goodbye!\n")
			break
		}
	}
}
<|endoftext|>


<commit_before>package imagick

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

	"github.com/gographics/imagick/imagick"
	"github.com/pressly/imgry"
)

var (
	ErrEngineReleased = errors.New("imagick: engine has been released.")
	ErrEngineFailure  = errors.New("imagick: unable to request a MagickWand")
)

type Engine struct {
	tmpDir string
	// TODO: perhaps we have counter of wands here..
	// it ensures we'll always release them..
	// AvailableWands .. .. set to MaxSizers ..
	// then, NewMagickWand() will limit, or error..
	// perhaps we make a channel of these things..
	// then block.. + ctx will have a timeout if it wants to stop waiting..
	// we can log these timeouts etc..
}

func (ng Engine) Version() string {
	v, _ := imagick.GetVersion()
	return fmt.Sprintf("%s", v)
}

func (ng Engine) Initialize(tmpDir string) error {
	if tmpDir != "" {
		if err := os.MkdirAll(tmpDir, 0755); err != nil {
			return err
		}
		ng.tmpDir = tmpDir
		os.Setenv("MAGICK_TMPDIR", tmpDir)
		ng.SweepTmpDir()
	}
	imagick.Initialize()
	return nil
}

func (ng Engine) Terminate() {
	imagick.Terminate()
	ng.SweepTmpDir()
}

func (ng Engine) SweepTmpDir() error {
	if ng.tmpDir == "" {
		return nil
	}
	err := filepath.Walk(
		ng.tmpDir,
		func(path string, info os.FileInfo, err error) error {
			if ng.tmpDir == path {
				return nil // skip the root
			}
			if strings.Index(filepath.Base(path), "magick") >= 0 {
				if err = os.Remove(path); err != nil {
					return fmt.Errorf("failed to sweet engine tmpdir %s, because: %s", path, err)
				}
			}
			return nil
		},
	)
	if err != nil {
		return err
	}
	return nil
}

func (ng Engine) LoadFile(filename string, srcFormat ...string) (imgry.Image, error) {
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return ng.LoadBlob(b, srcFormat...)
}

func (ng Engine) LoadBlob(b []byte, srcFormat ...string) (imgry.Image, error) {
	if len(b) == 0 {
		return nil, imgry.ErrInvalidImageData
	}

	mw := imagick.NewMagickWand()
	if !mw.IsVerified() {
		return nil, ErrEngineFailure
	}

	// Offer a hint to the decoder of the file format
	if len(srcFormat) > 0 {
		f := srcFormat[0]
		if f != "" {
			mw.SetFormat(f)
		}
	}

	err := mw.ReadImageBlob(b)
	if err != nil {
		mw.Destroy()
		return nil, imgry.ErrInvalidImageData
	}

	// TODO: perhaps we pass the engine instance like Image{engine: i}
	im := &Image{mw: mw, data: b}
	if err := im.sync(); err != nil {
		mw.Destroy()
		return nil, err
	}

	return im, nil
}

func (ng Engine) GetImageInfo(b []byte, srcFormat ...string) (*imgry.ImageInfo, error) {
	if len(b) == 0 {
		return nil, imgry.ErrInvalidImageData
	}

	mw := imagick.NewMagickWand()
	defer mw.Destroy()

	if !mw.IsVerified() {
		return nil, ErrEngineFailure
	}

	err := mw.PingImageBlob(b)
	if err != nil {
		return nil, imgry.ErrInvalidImageData
	}

	w, h := int(mw.GetImageWidth()), int(mw.GetImageHeight())
	ar := float64(int(float64(w)/float64(h)*10000)) / 10000

	format := strings.ToLower(mw.GetImageFormat())
	if format == "jpeg" {
		format = "jpg"
	}

	imfo := &imgry.ImageInfo{
		Format: format, Width: w, Height: h,
		AspectRatio: ar, ContentLength: len(b),
	}

	return imfo, nil
}

type Image struct {
	mw *imagick.MagickWand

	data   []byte
	width  int
	height int
	format string
}

func (i *Image) Data() []byte {
	return i.data
}

func (i *Image) Width() int {
	return i.width
}

func (i *Image) Height() int {
	return i.height
}

func (i *Image) Format() string {
	return i.format
}

func (i *Image) SetFormat(format string) error {
	if i.Released() {
		return ErrEngineReleased
	}
	if err := i.mw.SetImageFormat(format); err != nil {
		return err
	}
	if err := i.sync(); err != nil {
		return err
	}
	return nil
}

func (i *Image) Released() bool {
	return i.mw == nil
}

func (i *Image) Release() {
	if i.mw != nil {
		i.mw.Destroy()
		i.mw = nil
	}
}

func (i *Image) Clone() imgry.Image {
	i2 := &Image{}
	i2.data = i.data
	i2.width = i.width
	i2.height = i.height
	i2.format = i.format
	if i.mw != nil && i.mw.IsVerified() {
		i2.mw = i.mw.Clone()
	}
	return i2
}

func (i *Image) SizeIt(sz *imgry.Sizing) error {
	if i.Released() {
		return ErrEngineReleased
	}

	if err := i.sizeFrames(sz); err != nil {
		return err
	}

	if sz.Format != "" {
		if err := i.mw.SetFormat(sz.Format); err != nil {
			return err
		}
	}

	if sz.Quality > 0 {
		if err := i.mw.SetImageCompressionQuality(uint(sz.Quality)); err != nil {
			return err
		}
	}

	if err := i.sync(sz.Flatten); err != nil {
		return err
	}

	return nil
}

func (i *Image) sizeFrames(sz *imgry.Sizing) error {
	// Shortcut if there is nothing to size
	if sz.Size.Equal(imgry.ZeroRect) && sz.CropBox.Equal(imgry.ZeroFloatingRect) {
		return nil
	}

	// TODO: instead of Coalesce, better to change the offsets.
	// This is required for animated image resizing
	if !sz.Flatten && i.mw.GetNumberImages() > 1 {
		i.mw = i.mw.CoalesceImages()
	}

	i.mw.SetFirstIterator()
	for n := true; n; n = i.mw.NextImage() {

		pw, ph := int(i.mw.GetImageWidth()), int(i.mw.GetImageHeight())
		srcSize := imgry.NewRect(pw, ph)

		// Initial crop of the source image
		cropBox, cropOrigin, err := sz.CalcCropBox(srcSize)
		if err != nil {
			return err
		}
		if cropBox != nil && cropOrigin != nil && !cropBox.Equal(imgry.ZeroRect) {
			err := i.mw.CropImage(uint(cropBox.Width), uint(cropBox.Height), cropOrigin.X, cropOrigin.Y)
			if err != nil {
				return err
			}
			srcSize = cropBox
			i.mw.ResetImagePage("")
		}

		// Resize the image
		resizeRect, cropBox, cropOrigin := sz.CalcResizeRect(srcSize)
		if resizeRect != nil && !resizeRect.Equal(imgry.ZeroRect) {
			err := i.mw.ResizeImage(uint(resizeRect.Width), uint(resizeRect.Height), imagick.FILTER_LANCZOS, 1.0)
			if err != nil {
				return err
			}
			i.mw.ResetImagePage("")
		}

		// Perform any final crops from an operation
		if cropBox != nil && cropOrigin != nil && !cropBox.Equal(imgry.ZeroRect) {
			err := i.mw.CropImage(uint(cropBox.Width), uint(cropBox.Height), cropOrigin.X, cropOrigin.Y)
			if err != nil {
				return err
			}
			i.mw.ResetImagePage("")
		}

		if sz.Flatten {
			break
		}
	}

	i.mw = i.mw.DeconstructImages()

	return nil
}

func (i *Image) WriteToFile(fn string) error {
	err := ioutil.WriteFile(fn, i.Data(), 0664)
	return err
}

func (i *Image) sync(optFlatten ...bool) error {
	if i.Released() {
		return ErrEngineReleased
	}

	var flatten bool
	if len(optFlatten) > 0 {
		flatten = optFlatten[0]
	}

	if flatten {
		i.data = i.mw.GetImageBlob()
	} else {
		i.data = i.mw.GetImagesBlob()
	}

	i.width = int(i.mw.GetImageWidth())
	i.height = int(i.mw.GetImageHeight())

	i.format = strings.ToLower(i.mw.GetImageFormat())
	if i.format == "jpeg" {
		i.format = "jpg"
	}

	return nil
}
<commit_msg>Executing i.mw.DeconstructImages only when i.mw.GetNumberImages() > 1.<commit_after>package imagick

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

	"github.com/gographics/imagick/imagick"
	"github.com/pressly/imgry"
)

var (
	ErrEngineReleased = errors.New("imagick: engine has been released.")
	ErrEngineFailure  = errors.New("imagick: unable to request a MagickWand")
)

type Engine struct {
	tmpDir string
	// TODO: perhaps we have counter of wands here..
	// it ensures we'll always release them..
	// AvailableWands .. .. set to MaxSizers ..
	// then, NewMagickWand() will limit, or error..
	// perhaps we make a channel of these things..
	// then block.. + ctx will have a timeout if it wants to stop waiting..
	// we can log these timeouts etc..
}

func (ng Engine) Version() string {
	v, _ := imagick.GetVersion()
	return fmt.Sprintf("%s", v)
}

func (ng Engine) Initialize(tmpDir string) error {
	if tmpDir != "" {
		if err := os.MkdirAll(tmpDir, 0755); err != nil {
			return err
		}
		ng.tmpDir = tmpDir
		os.Setenv("MAGICK_TMPDIR", tmpDir)
		ng.SweepTmpDir()
	}
	imagick.Initialize()
	return nil
}

func (ng Engine) Terminate() {
	imagick.Terminate()
	ng.SweepTmpDir()
}

func (ng Engine) SweepTmpDir() error {
	if ng.tmpDir == "" {
		return nil
	}
	err := filepath.Walk(
		ng.tmpDir,
		func(path string, info os.FileInfo, err error) error {
			if ng.tmpDir == path {
				return nil // skip the root
			}
			if strings.Index(filepath.Base(path), "magick") >= 0 {
				if err = os.Remove(path); err != nil {
					return fmt.Errorf("failed to sweet engine tmpdir %s, because: %s", path, err)
				}
			}
			return nil
		},
	)
	if err != nil {
		return err
	}
	return nil
}

func (ng Engine) LoadFile(filename string, srcFormat ...string) (imgry.Image, error) {
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return ng.LoadBlob(b, srcFormat...)
}

func (ng Engine) LoadBlob(b []byte, srcFormat ...string) (imgry.Image, error) {
	if len(b) == 0 {
		return nil, imgry.ErrInvalidImageData
	}

	mw := imagick.NewMagickWand()
	if !mw.IsVerified() {
		return nil, ErrEngineFailure
	}

	// Offer a hint to the decoder of the file format
	if len(srcFormat) > 0 {
		f := srcFormat[0]
		if f != "" {
			mw.SetFormat(f)
		}
	}

	err := mw.ReadImageBlob(b)
	if err != nil {
		mw.Destroy()
		return nil, imgry.ErrInvalidImageData
	}

	// TODO: perhaps we pass the engine instance like Image{engine: i}
	im := &Image{mw: mw, data: b}
	if err := im.sync(); err != nil {
		mw.Destroy()
		return nil, err
	}

	return im, nil
}

func (ng Engine) GetImageInfo(b []byte, srcFormat ...string) (*imgry.ImageInfo, error) {
	if len(b) == 0 {
		return nil, imgry.ErrInvalidImageData
	}

	mw := imagick.NewMagickWand()
	defer mw.Destroy()

	if !mw.IsVerified() {
		return nil, ErrEngineFailure
	}

	err := mw.PingImageBlob(b)
	if err != nil {
		return nil, imgry.ErrInvalidImageData
	}

	w, h := int(mw.GetImageWidth()), int(mw.GetImageHeight())
	ar := float64(int(float64(w)/float64(h)*10000)) / 10000

	format := strings.ToLower(mw.GetImageFormat())
	if format == "jpeg" {
		format = "jpg"
	}

	imfo := &imgry.ImageInfo{
		Format: format, Width: w, Height: h,
		AspectRatio: ar, ContentLength: len(b),
	}

	return imfo, nil
}

type Image struct {
	mw *imagick.MagickWand

	data   []byte
	width  int
	height int
	format string
}

func (i *Image) Data() []byte {
	return i.data
}

func (i *Image) Width() int {
	return i.width
}

func (i *Image) Height() int {
	return i.height
}

func (i *Image) Format() string {
	return i.format
}

func (i *Image) SetFormat(format string) error {
	if i.Released() {
		return ErrEngineReleased
	}
	if err := i.mw.SetImageFormat(format); err != nil {
		return err
	}
	if err := i.sync(); err != nil {
		return err
	}
	return nil
}

func (i *Image) Released() bool {
	return i.mw == nil
}

func (i *Image) Release() {
	if i.mw != nil {
		i.mw.Destroy()
		i.mw = nil
	}
}

func (i *Image) Clone() imgry.Image {
	i2 := &Image{}
	i2.data = i.data
	i2.width = i.width
	i2.height = i.height
	i2.format = i.format
	if i.mw != nil && i.mw.IsVerified() {
		i2.mw = i.mw.Clone()
	}
	return i2
}

func (i *Image) SizeIt(sz *imgry.Sizing) error {
	if i.Released() {
		return ErrEngineReleased
	}

	if err := i.sizeFrames(sz); err != nil {
		return err
	}

	if sz.Format != "" {
		if err := i.mw.SetFormat(sz.Format); err != nil {
			return err
		}
	}

	if sz.Quality > 0 {
		if err := i.mw.SetImageCompressionQuality(uint(sz.Quality)); err != nil {
			return err
		}
	}

	if err := i.sync(sz.Flatten); err != nil {
		return err
	}

	return nil
}

func (i *Image) sizeFrames(sz *imgry.Sizing) error {
	// Shortcut if there is nothing to size
	if sz.Size.Equal(imgry.ZeroRect) && sz.CropBox.Equal(imgry.ZeroFloatingRect) {
		return nil
	}

	numberImages := i.mw.GetNumberImages()

	// TODO: instead of Coalesce, better to change the offsets.
	// This is required for animated image resizing
	if !sz.Flatten && numberImages > 1 {
		i.mw = i.mw.CoalesceImages()
	}

	i.mw.SetFirstIterator()
	for n := true; n; n = i.mw.NextImage() {

		pw, ph := int(i.mw.GetImageWidth()), int(i.mw.GetImageHeight())
		srcSize := imgry.NewRect(pw, ph)

		// Initial crop of the source image
		cropBox, cropOrigin, err := sz.CalcCropBox(srcSize)
		if err != nil {
			return err
		}
		if cropBox != nil && cropOrigin != nil && !cropBox.Equal(imgry.ZeroRect) {
			err := i.mw.CropImage(uint(cropBox.Width), uint(cropBox.Height), cropOrigin.X, cropOrigin.Y)
			if err != nil {
				return err
			}
			srcSize = cropBox
			i.mw.ResetImagePage("")
		}

		// Resize the image
		resizeRect, cropBox, cropOrigin := sz.CalcResizeRect(srcSize)
		if resizeRect != nil && !resizeRect.Equal(imgry.ZeroRect) {
			err := i.mw.ResizeImage(uint(resizeRect.Width), uint(resizeRect.Height), imagick.FILTER_LANCZOS, 1.0)
			if err != nil {
				return err
			}
			i.mw.ResetImagePage("")
		}

		// Perform any final crops from an operation
		if cropBox != nil && cropOrigin != nil && !cropBox.Equal(imgry.ZeroRect) {
			err := i.mw.CropImage(uint(cropBox.Width), uint(cropBox.Height), cropOrigin.X, cropOrigin.Y)
			if err != nil {
				return err
			}
			i.mw.ResetImagePage("")
		}

		if sz.Flatten {
			break
		}
	}

	if numberImages > 1 {
		i.mw = i.mw.DeconstructImages()
	}

	return nil
}

func (i *Image) WriteToFile(fn string) error {
	err := ioutil.WriteFile(fn, i.Data(), 0664)
	return err
}

func (i *Image) sync(optFlatten ...bool) error {
	if i.Released() {
		return ErrEngineReleased
	}

	var flatten bool
	if len(optFlatten) > 0 {
		flatten = optFlatten[0]
	}

	if flatten {
		i.data = i.mw.GetImageBlob()
	} else {
		i.data = i.mw.GetImagesBlob()
	}

	i.width = int(i.mw.GetImageWidth())
	i.height = int(i.mw.GetImageHeight())

	i.format = strings.ToLower(i.mw.GetImageFormat())
	if i.format == "jpeg" {
		i.format = "jpg"
	}

	return nil
}
<|endoftext|>