Alter the behavior of equal util method along with Contains, IndexOf & LastIndexOf functions

Author: vellotis

predicate.go

@@ -20,7 +20,7 @@ func predicatesImpl(value interface{}, wantedAnswer bool, predicates interface{}
 		}
 
 		funcType := funcValue.Type()
-		if funcType.NumIn() != 1 || funcType.NumOut() != 1 || funcType.Out(0).Kind() != reflect.Bool {
+		if !IsPredicate(funcValue.Interface()) {
 			panic("Predicate function must have 1 parameter and must return boolean")
 		}
 

presence.go

@@ -118,8 +118,9 @@ func IndexOf(in interface{}, elem interface{}) int {
 	}
 
 	if inType.Kind() == reflect.Slice {
+		equalTo := equal(elem)
 		for i := 0; i < inValue.Len(); i++ {
-			if equal(inValue.Index(i).Interface(), elem) {
+			if equalTo(reflect.Value{}, inValue.Index(i)) {
 				return i
 			}
 		}
@@ -144,8 +145,9 @@ func LastIndexOf(in interface{}, elem interface{}) int {
 	if inType.Kind() == reflect.Slice {
 		length := inValue.Len()
 
+		equalTo := equal(elem)
 		for i := length - 1; i >= 0; i-- {
-			if equal(inValue.Index(i).Interface(), elem) {
+			if equalTo(reflect.Value{}, inValue.Index(i)) {
 				return i
 			}
 		}
@@ -164,14 +166,16 @@ func Contains(in interface{}, elem interface{}) bool {
 	case reflect.String:
 		return strings.Contains(inValue.String(), elemValue.String())
 	case reflect.Map:
+		equalTo := equal(elem, true)
 		for _, key := range inValue.MapKeys() {
-			if equal(key.Interface(), elem) {
+			if equalTo(key, inValue.MapIndex(key)) {
 				return true
 			}
 		}
 	case reflect.Slice, reflect.Array:
+		equalTo := equal(elem)
 		for i := 0; i < inValue.Len(); i++ {
-			if equal(inValue.Index(i).Interface(), elem) {
+			if equalTo(reflect.Value{}, inValue.Index(i)) {
 				return true
 			}
 		}

presence_test.go

@@ -54,6 +54,20 @@ func TestContains(t *testing.T) {
 
 	is.True(Contains(mapping, 1))
 	is.False(Contains(mapping, 2))
+
+	is.False(Contains(mapping, func (key int, val *Foo) bool {
+		return key == 4
+	}))
+	is.True(Contains(mapping, func (key int, val *Foo) bool {
+		return key == 1
+	}))
+
+	is.False(Contains(mapping, func (_ int, val *Foo) bool {
+		return val.FirstName == "NotPresent"
+	}))
+	is.True(Contains(mapping, func (_ int, val *Foo) bool {
+		return val.FirstName == "Harald"
+	}))
 }
 
 func TestEvery(t *testing.T) {
@@ -117,6 +131,9 @@ func TestIndexOf(t *testing.T) {
 	is := assert.New(t)
 
 	is.Equal(IndexOf([]string{"foo", "bar"}, "bar"), 1)
+	is.Equal(IndexOf([]string{"foo", "bar"}, func (value string) bool {
+		return value == "bar"
+	}), 1)
 
 	is.Equal(IndexOf(results, f), 0)
 	is.Equal(IndexOf(results, b), -1)
@@ -126,6 +143,9 @@ func TestLastIndexOf(t *testing.T) {
 	is := assert.New(t)
 
 	is.Equal(LastIndexOf([]string{"foo", "bar", "bar"}, "bar"), 2)
+	is.Equal(LastIndexOf([]string{"foo", "bar", "bar"}, func (value string) bool {
+		return value == "bar"
+	}), 2)
 	is.Equal(LastIndexOf([]int{1, 2, 2, 3}, 2), 2)
 	is.Equal(LastIndexOf([]int{1, 2, 2, 3}, 4), -1)
 }

utils.go

@@ -1,16 +1,57 @@
 package funk
 
 import (
+	"fmt"
 	"reflect"
 )
 
-func equal(expected, actual interface{}) bool {
-	if expected == nil || actual == nil {
-		return expected == actual
+func equal(expectedOrPredicate interface{}, optionalIsMap ...bool) func(keyValueIfMap, actualValue reflect.Value) bool {
+	isMap := append(optionalIsMap, false)[0]
+
+	if IsFunction(expectedOrPredicate) {
+		inTypes := []reflect.Type{nil}; if isMap {
+			inTypes = append(inTypes, nil)
+		}
+
+		if !IsPredicate(expectedOrPredicate, inTypes...) {
+			panic(fmt.Sprintf("Predicate function must have %d parameter and must return boolean", len(inTypes)))
+		}
+
+		predicateValue := reflect.ValueOf(expectedOrPredicate)
+
+		return func(keyValueIfMap, actualValue reflect.Value) bool {
+
+			if isMap && !keyValueIfMap.Type().ConvertibleTo(predicateValue.Type().In(0)) {
+				panic("Given key is not compatible with type of parameter for the predicate.")
+			}
+
+			if (isMap && !actualValue.Type().ConvertibleTo(predicateValue.Type().In(1))) ||
+				(!isMap && !actualValue.Type().ConvertibleTo(predicateValue.Type().In(0))) {
+				panic("Given value is not compatible with type of parameter for the predicate.")
+			}
+
+			args := []reflect.Value{actualValue}
+			if isMap {
+				args = append([]reflect.Value{keyValueIfMap}, args...)
+			}
+
+			return predicateValue.Call(args)[0].Bool()
+		}
 	}
 
-	return reflect.DeepEqual(expected, actual)
+	expected := expectedOrPredicate
+
+	return func(keyValueIfMap, actualValue reflect.Value) bool {
+		if isMap {
+			actualValue = keyValueIfMap
+		}
+
+		if expected == nil || actualValue.IsZero() {
+			return actualValue.Interface() == expected
+		}
 
+		return reflect.DeepEqual(actualValue.Interface(), expected)
+	}
 }
 
 func sliceElem(rtype reflect.Type) reflect.Type {