func.go

package xpath

import (
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"
	"sync"
	"unicode"
)

// Defined an interface of stringBuilder that compatible with
// strings.Builder(go 1.10) and bytes.Buffer(< go 1.10)
type stringBuilder interface {
	WriteRune(r rune) (n int, err error)
	WriteString(s string) (int, error)
	Reset()
	Grow(n int)
	String() string
}

var builderPool = sync.Pool{New: func() interface{} {
	return newStringBuilder()
}}

// The XPath function list.

func predicate(q query) func(NodeNavigator) bool {
	type Predicater interface {
		Test(NodeNavigator) bool
	}
	if p, ok := q.(Predicater); ok {
		return p.Test
	}
	return func(NodeNavigator) bool { return true }
}

// positionFunc is a XPath Node Set functions position().
func positionFunc() func(query, iterator) interface{} {
	return func(q query, t iterator) interface{} {
		var (
			count = 1
			node  = t.Current().Copy()
		)
		test := predicate(q)
		for node.MoveToPrevious() {
			if test(node) {
				count++
			}
		}
		return float64(count)
	}
}

// lastFunc is a XPath Node Set functions last().
func lastFunc() func(query, iterator) interface{} {
	return func(q query, t iterator) interface{} {
		var (
			count = 0
			node  = t.Current().Copy()
		)
		test := predicate(q)
		node.MoveToFirst()
		for {
			if test(node) {
				count++
			}
			if !node.MoveToNext() {
				break
			}
		}
		return float64(count)
	}
}

// countFunc is a XPath Node Set functions count(node-set).
func countFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var count = 0
		q := functionArgs(arg)
		test := predicate(q)
		switch typ := q.Evaluate(t).(type) {
		case query:
			for node := typ.Select(t); node != nil; node = typ.Select(t) {
				if test(node) {
					count++
				}
			}
		}
		return float64(count)
	}
}

// sumFunc is a XPath Node Set functions sum(node-set).
func sumFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var sum float64
		switch typ := functionArgs(arg).Evaluate(t).(type) {
		case query:
			for node := typ.Select(t); node != nil; node = typ.Select(t) {
				if v, err := strconv.ParseFloat(node.Value(), 64); err == nil {
					sum += v
				}
			}
		case float64:
			sum = typ
		case string:
			v, err := strconv.ParseFloat(typ, 64)
			if err != nil {
				panic(errors.New("sum() function argument type must be a node-set or number"))
			}
			sum = v
		}
		return sum
	}
}

func asNumber(t iterator, o interface{}) float64 {
	switch typ := o.(type) {
	case query:
		node := typ.Select(t)
		if node == nil {
			return math.NaN()
		}
		if v, err := strconv.ParseFloat(node.Value(), 64); err == nil {
			return v
		}
	case float64:
		return typ
	case string:
		v, err := strconv.ParseFloat(typ, 64)
		if err == nil {
			return v
		}
	}
	return math.NaN()
}

// ceilingFunc is a XPath Node Set functions ceiling(node-set).
func ceilingFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		val := asNumber(t, functionArgs(arg).Evaluate(t))
		// if math.IsNaN(val) {
		// 	panic(errors.New("ceiling() function argument type must be a valid number"))
		// }
		return math.Ceil(val)
	}
}

// floorFunc is a XPath Node Set functions floor(node-set).
func floorFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		val := asNumber(t, functionArgs(arg).Evaluate(t))
		return math.Floor(val)
	}
}

// roundFunc is a XPath Node Set functions round(node-set).
func roundFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		val := asNumber(t, functionArgs(arg).Evaluate(t))
		//return math.Round(val)
		return round(val)
	}
}

// nameFunc is a XPath functions name([node-set]).
func nameFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var v NodeNavigator
		if arg == nil {
			v = t.Current()
		} else {
			v = arg.Clone().Select(t)
			if v == nil {
				return ""
			}
		}
		ns := v.Prefix()
		if ns == "" {
			return v.LocalName()
		}
		return ns + ":" + v.LocalName()
	}
}

// localNameFunc is a XPath functions local-name([node-set]).
func localNameFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var v NodeNavigator
		if arg == nil {
			v = t.Current()
		} else {
			v = arg.Clone().Select(t)
			if v == nil {
				return ""
			}
		}
		return v.LocalName()
	}
}

// namespaceFunc is a XPath functions namespace-uri([node-set]).
func namespaceFunc(arg query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var v NodeNavigator
		if arg == nil {
			v = t.Current()
		} else {
			// Get the first node in the node-set if specified.
			v = arg.Clone().Select(t)
			if v == nil {
				return ""
			}
		}
		// fix about namespace-uri() bug: https://github.com/antchfx/xmlquery/issues/22
		// TODO: In the next version, add NamespaceURL() to the NodeNavigator interface.
		type namespaceURL interface {
			NamespaceURL() string
		}
		if f, ok := v.(namespaceURL); ok {
			return f.NamespaceURL()
		}
		return v.Prefix()
	}
}

func asBool(t iterator, v interface{}) bool {
	switch v := v.(type) {
	case nil:
		return false
	case *NodeIterator:
		return v.MoveNext()
	case bool:
		return v
	case float64:
		return v != 0
	case string:
		return v != ""
	case query:
		return v.Select(t) != nil
	default:
		panic(fmt.Errorf("unexpected type: %T", v))
	}
}

func asString(t iterator, v interface{}) string {
	switch v := v.(type) {
	case nil:
		return ""
	case bool:
		if v {
			return "true"
		}
		return "false"
	case float64:
		return strconv.FormatFloat(v, 'g', -1, 64)
	case string:
		return v
	case query:
		node := v.Select(t)
		if node == nil {
			return ""
		}
		return node.Value()
	default:
		panic(fmt.Errorf("unexpected type: %T", v))
	}
}

// booleanFunc is a XPath functions boolean([node-set]).
func booleanFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		v := functionArgs(arg1).Evaluate(t)
		return asBool(t, v)
	}
}

// numberFunc is a XPath functions number([node-set]).
func numberFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		v := functionArgs(arg1).Evaluate(t)
		return asNumber(t, v)
	}
}

// stringFunc is a XPath functions string([node-set]).
func stringFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		v := functionArgs(arg1).Evaluate(t)
		return asString(t, v)
	}
}

// startwithFunc is a XPath functions starts-with(string, string).
func startwithFunc(arg1, arg2 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var (
			m, n string
			ok   bool
		)
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			m = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return false
			}
			m = node.Value()
		default:
			panic(errors.New("starts-with() function argument type must be string"))
		}
		n, ok = functionArgs(arg2).Evaluate(t).(string)
		if !ok {
			panic(errors.New("starts-with() function argument type must be string"))
		}
		return strings.HasPrefix(m, n)
	}
}

// endwithFunc is a XPath functions ends-with(string, string).
func endwithFunc(arg1, arg2 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var (
			m, n string
			ok   bool
		)
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			m = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return false
			}
			m = node.Value()
		default:
			panic(errors.New("ends-with() function argument type must be string"))
		}
		n, ok = functionArgs(arg2).Evaluate(t).(string)
		if !ok {
			panic(errors.New("ends-with() function argument type must be string"))
		}
		return strings.HasSuffix(m, n)
	}
}

// containsFunc is a XPath functions contains(string or @attr, string).
func containsFunc(arg1, arg2 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var (
			m, n string
			ok   bool
		)
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			m = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return false
			}
			m = node.Value()
		default:
			panic(errors.New("contains() function argument type must be string"))
		}

		n, ok = functionArgs(arg2).Evaluate(t).(string)
		if !ok {
			panic(errors.New("contains() function argument type must be string"))
		}

		return strings.Contains(m, n)
	}
}

// matchesFunc is an XPath function that tests a given string against a regexp pattern.
// Note: does not support https://www.w3.org/TR/xpath-functions-31/#func-matches 3rd optional `flags` argument; if
// needed, directly put flags in the regexp pattern, such as `(?i)^pattern$` for `i` flag.
func matchesFunc(arg1, arg2 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var s string
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			s = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return ""
			}
			s = node.Value()
		}
		var pattern string
		var ok bool
		if pattern, ok = functionArgs(arg2).Evaluate(t).(string); !ok {
			panic(errors.New("matches() function second argument type must be string"))
		}
		re, err := getRegexp(pattern)
		if err != nil {
			panic(fmt.Errorf("matches() function second argument is not a valid regexp pattern, err: %s", err.Error()))
		}
		return re.MatchString(s)
	}
}

// normalizespaceFunc is XPath functions normalize-space(string?)
func normalizespaceFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var m string
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			m = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return ""
			}
			m = node.Value()
		}
		var b = builderPool.Get().(stringBuilder)
		b.Grow(len(m))

		runeStr := []rune(strings.TrimSpace(m))
		l := len(runeStr)
		for i := range runeStr {
			r := runeStr[i]
			isSpace := unicode.IsSpace(r)
			if !(isSpace && (i+1 < l && unicode.IsSpace(runeStr[i+1]))) {
				if isSpace {
					r = ' '
				}
				b.WriteRune(r)
			}
		}
		result := b.String()
		b.Reset()
		builderPool.Put(b)

		return result
	}
}

// substringFunc is XPath functions substring function returns a part of a given string.
func substringFunc(arg1, arg2, arg3 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var m string
		switch typ := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			m = typ
		case query:
			node := typ.Select(t)
			if node == nil {
				return ""
			}
			m = node.Value()
		}

		var start, length float64
		var ok bool
		if start, ok = functionArgs(arg2).Evaluate(t).(float64); !ok {
			panic(errors.New("substring() function first argument type must be number"))
		}
		// fix https://github.com/antchfx/xpath/issues/109
		start = math.Round(start)
		if start > float64(len(m)) {
			return ""
		}
		if arg3 == nil {
			if start <= 0 {
				return m
			}
			return m[int(start)-1:]
		}

		if length, ok = functionArgs(arg3).Evaluate(t).(float64); !ok {
			panic(errors.New("substring() function second argument type must be number"))
		}
		length = math.Round(length)
		if length <= 0 {
			return ""
		}
		if length > float64(len(m)) {
			length = float64(len(m))
		}
		if start < 0 {
			length = length - math.Abs(start)
			if length <= 1 {
				return ""
			}
			return m[:int(length-1)]
		}
		if start == 0 {
			return m[:int(length-1)]
		}
		return m[int(start-1):int(length+start-1)]
	}
}

// substringIndFunc is XPath functions substring-before/substring-after function returns a part of a given string.
func substringIndFunc(arg1, arg2 query, after bool) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var str string
		switch v := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			str = v
		case query:
			node := v.Select(t)
			if node == nil {
				return ""
			}
			str = node.Value()
		}
		var word string
		switch v := functionArgs(arg2).Evaluate(t).(type) {
		case string:
			word = v
		case query:
			node := v.Select(t)
			if node == nil {
				return ""
			}
			word = node.Value()
		}
		if word == "" {
			return ""
		}

		i := strings.Index(str, word)
		if i < 0 {
			return ""
		}
		if after {
			return str[i+len(word):]
		}
		return str[:i]
	}
}

// stringLengthFunc is XPATH string-length( [string] ) function that returns a number
// equal to the number of characters in a given string.
func stringLengthFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		switch v := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			return float64(len(v))
		case query:
			node := v.Select(t)
			if node == nil {
				break
			}
			return float64(len(node.Value()))
		}
		return float64(0)
	}
}

// translateFunc is XPath functions translate() function returns a replaced string.
func translateFunc(arg1, arg2, arg3 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		str := asString(t, functionArgs(arg1).Evaluate(t))
		src := asString(t, functionArgs(arg2).Evaluate(t))
		dst := asString(t, functionArgs(arg3).Evaluate(t))

		replace := make([]string, 0, len(src))
		for i, s := range src {
			d := ""
			if i < len(dst) {
				d = string(dst[i])
			}
			replace = append(replace, string(s), d)
		}
		return strings.NewReplacer(replace...).Replace(str)
	}
}

// replaceFunc is XPath functions replace() function returns a replaced string.
func replaceFunc(arg1, arg2, arg3 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		str := asString(t, functionArgs(arg1).Evaluate(t))
		src := asString(t, functionArgs(arg2).Evaluate(t))
		dst := asString(t, functionArgs(arg3).Evaluate(t))
		e, err := getRegexp(src)
		if err != nil {
			panic(fmt.Errorf("replace() function second argument is not a valid regexp pattern, err: %s", err.Error()))
		}

		// replace all $i to ${i} for golang regexp.Expand
		for idx := e.NumSubexp(); idx > 0; idx-- {
			dst = strings.ReplaceAll(dst, fmt.Sprintf("$%d", idx), fmt.Sprintf("${%d}", idx))
		}

		return e.ReplaceAllString(str, dst)
	}
}

// notFunc is XPATH functions not(expression) function operation.
func notFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		switch v := functionArgs(arg1).Evaluate(t).(type) {
		case bool:
			return !v
		case query:
			node := v.Select(t)
			return node == nil
		default:
			return false
		}
	}
}

// concatFunc is the concat function concatenates two or more
// strings and returns the resulting string.
// concat( string1 , string2 [, stringn]* )
func concatFunc(args ...query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		b := builderPool.Get().(stringBuilder)
		for _, v := range args {
			v = functionArgs(v)

			switch v := v.Evaluate(t).(type) {
			case string:
				b.WriteString(v)
			case query:
				node := v.Select(t)
				if node != nil {
					b.WriteString(node.Value())
				}
			}
		}
		result := b.String()
		b.Reset()
		builderPool.Put(b)

		return result
	}
}

// https://github.com/antchfx/xpath/issues/43
func functionArgs(q query) query {
	if _, ok := q.(*functionQuery); ok {
		return q
	}
	return q.Clone()
}

func reverseFunc(q query, t iterator) func() NodeNavigator {
	var list []NodeNavigator
	for {
		node := q.Select(t)
		if node == nil {
			break
		}
		list = append(list, node.Copy())
	}
	i := len(list)
	return func() NodeNavigator {
		if i <= 0 {
			return nil
		}
		i--
		node := list[i]
		return node
	}
}

// string-join is a XPath Node Set functions string-join(node-set, separator).
func stringJoinFunc(q, arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		var separator string
		switch v := functionArgs(arg1).Evaluate(t).(type) {
		case string:
			separator = v
		case query:
			node := v.Select(t)
			if node != nil {
				separator = node.Value()
			}
		}

		q = functionArgs(q)
		test := predicate(q)
		var parts []string
		switch v := q.Evaluate(t).(type) {
		case string:
			return v
		case query:
			for node := v.Select(t); node != nil; node = v.Select(t) {
				if test(node) {
					parts = append(parts, node.Value())
				}
			}
		}
		return strings.Join(parts, separator)
	}
}

// lower-case is XPATH function that converts a string to lower case.
func lowerCaseFunc(arg1 query) func(query, iterator) interface{} {
	return func(_ query, t iterator) interface{} {
		v := functionArgs(arg1).Evaluate(t)
		return strings.ToLower(asString(t, v))
	}
}