From 0d7263928a74a8d4cf6e23b648bace6925b65dbb Mon Sep 17 00:00:00 2001
From: Alan Donovan <adonovan@google.com>
Date: Tue, 25 Jul 2023 12:14:58 -0400
Subject: [PATCH] starlarktest/assert.star: allow +/-1 ULP in float eq
 comparisons (#490)

Fixes #488
---
 starlarktest/assert.star     | 10 +++++++++-
 starlarktest/starlarktest.go | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/starlarktest/assert.star b/starlarktest/assert.star
index c6e480fc..9a699df4 100644
--- a/starlarktest/assert.star
+++ b/starlarktest/assert.star
@@ -6,12 +6,20 @@
 # matches(str, pattern): report whether str matches regular expression pattern.
 # module(**kwargs): a constructor for a module.
 # _freeze(x): freeze the value x and everything reachable from it.
+# _floateq(x, y): reports floating point equality (within 1 ULP).
 #
 # Clients may use these functions to define their own testing abstractions.
 
+_num = ("float", "int")
+
 def _eq(x, y):
     if x != y:
-        error("%r != %r" % (x, y))
+	if (type(x) == "float" and type(y) in _num or
+	    type(y) == "float" and type(x) in _num):
+	    if not _floateq(float(x), float(y)):
+		error("floats: %r != %r (delta > 1 ulp)" % (x, y))
+	else:
+            error("%r != %r" % (x, y))
 
 def _ne(x, y):
     if x == y:
diff --git a/starlarktest/starlarktest.go b/starlarktest/starlarktest.go
index 3a19f04e..88e85489 100644
--- a/starlarktest/starlarktest.go
+++ b/starlarktest/starlarktest.go
@@ -15,6 +15,7 @@ package starlarktest // import "go.starlark.net/starlarktest"
 import (
 	_ "embed"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -63,11 +64,12 @@ var (
 func LoadAssertModule() (starlark.StringDict, error) {
 	once.Do(func() {
 		predeclared := starlark.StringDict{
-			"error":   starlark.NewBuiltin("error", error_),
-			"catch":   starlark.NewBuiltin("catch", catch),
-			"matches": starlark.NewBuiltin("matches", matches),
-			"module":  starlark.NewBuiltin("module", starlarkstruct.MakeModule),
-			"_freeze": starlark.NewBuiltin("freeze", freeze),
+			"error":    starlark.NewBuiltin("error", error_),
+			"catch":    starlark.NewBuiltin("catch", catch),
+			"matches":  starlark.NewBuiltin("matches", matches),
+			"module":   starlark.NewBuiltin("module", starlarkstruct.MakeModule),
+			"_freeze":  starlark.NewBuiltin("freeze", freeze),
+			"_floateq": starlark.NewBuiltin("floateq", floateq),
 		}
 		thread := new(starlark.Thread)
 		assert, assertErr = starlark.ExecFile(thread, "assert.star", assertFileSrc, predeclared)
@@ -131,6 +133,29 @@ func freeze(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, k
 	return args[0], nil
 }
 
+// floateq(x, y) reports whether two floats are within 1 ULP of each other.
+func floateq(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var xf, yf starlark.Float
+	if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 2, &xf, &yf); err != nil {
+		return nil, err
+	}
+
+	res := false
+	switch {
+	case xf == yf:
+		res = true
+	case math.IsNaN(float64(xf)):
+		res = math.IsNaN(float64(yf))
+	case math.IsNaN(float64(yf)):
+		// false (non-NaN = Nan)
+	default:
+		x := math.Float64bits(float64(xf))
+		y := math.Float64bits(float64(yf))
+		res = x == y+1 || y == x+1
+	}
+	return starlark.Bool(res), nil
+}
+
 // DataFile returns the effective filename of the specified
 // test data resource.  The function abstracts differences between
 // 'go build', under which a test runs in its package directory,