fix(table): make hash traversal linear

doc_test.go

@@ -1,22 +1,37 @@
 package lua_test
 
 import (
+	"fmt"
+
 	"github.com/Shopify/go-lua"
 )
 
 // This example receives a variable number of numerical arguments and returns their average and sum.
-func ExampleFunction(l *lua.State) int {
-	n := l.Top() // Number of arguments.
-	var sum float64
-	for i := 1; i <= n; i++ {
-		f, ok := l.ToNumber(i)
-		if !ok {
-			l.PushString("incorrect argument")
-			l.Error()
+func ExampleFunction() {
+	averageAndSum := func(l *lua.State) int {
+		n := l.Top() // Number of arguments.
+		var sum float64
+		for i := 1; i <= n; i++ {
+			f, ok := l.ToNumber(i)
+			if !ok {
+				l.PushString("incorrect argument")
+				l.Error()
+			}
+			sum += f
 		}
-		sum += f
+		l.PushNumber(sum / float64(n)) // First result.
+		l.PushNumber(sum)              // Second result.
+		return 2                       // Result count.
 	}
-	l.PushNumber(sum / float64(n)) // First result.
-	l.PushNumber(sum)              // Second result.
-	return 2                       // Result count.
+
+	l := lua.NewState()
+	l.PushGoFunction(averageAndSum)
+	l.PushNumber(1)
+	l.PushNumber(2)
+	l.PushNumber(3)
+	l.Call(3, 2)
+	average, _ := l.ToNumber(-2)
+	sum, _ := l.ToNumber(-1)
+	fmt.Println(average, sum)
+	// Output: 2 6
 }

pow_test.go

@@ -0,0 +1,101 @@
+package lua
+
+import (
+	"fmt"
+	"math"
+	"testing"
+)
+
+func TestPowerConsistentAcrossEvaluationPaths(t *testing.T) {
+	engines := []struct {
+		name string
+		run  func(*State)
+	}{
+		{"function table", (*State).executeFunctionTable},
+		{"switch", (*State).executeSwitch},
+	}
+	cases := []struct {
+		base, exponent             string
+		baseNumber, exponentNumber float64
+	}{
+		{"10", "-33", 10, -33},
+		{"10", "33", 10, 33},
+		{"10", "100", 10, 100},
+		{"10", "308", 10, 308},
+		{"10", "309", 10, 309},
+		{"10", "-323", 10, -323},
+		{"10", "-324", 10, -324},
+		{"-2", "3", -2, 3},
+		{"-2", "-3", -2, -3},
+		{"-2", "0.5", -2, 0.5},
+		{"0", "0", 0, 0},
+		{"0", "-3", 0, -3},
+		{"-0.0", "3", math.Copysign(0, -1), 3},
+		{"-0.0", "-3", math.Copysign(0, -1), -3},
+		{"-0.0", "2", math.Copysign(0, -1), 2},
+		{"-0.0", "-2", math.Copysign(0, -1), -2},
+		{"2", "0.5", 2, 0.5},
+	}
+	for _, engine := range engines {
+		for _, tc := range cases {
+			t.Run(engine.name+"/"+tc.base+"^"+tc.exponent, func(t *testing.T) {
+				l := NewState()
+				OpenLibraries(l)
+				source := fmt.Sprintf(`
+local base, exponent = (%s), (%s)
+local folded = (%s)^(%s)
+local runtime = base^exponent
+local coerced = %q^%q
+local library = math.pow(base, exponent)
+local keyMatches = true
+if folded == folded then
+  local keys = {[folded] = "found"}
+  keyMatches = keys[runtime] == "found" and keys[coerced] == "found" and keys[library] == "found"
+end
+return folded, runtime, coerced, library, keyMatches
+`, tc.base, tc.exponent, tc.base, tc.exponent, tc.base, tc.exponent)
+				if err := LoadString(l, source); err != nil {
+					t.Fatal(err)
+				}
+				if err := l.protectedCall(func() {
+					if l.preCall(l.top-1, MultipleReturns) {
+						t.Fatal("expected a Lua function")
+					}
+					engine.run(l)
+				}, l.top-1, 0); err != nil {
+					t.Fatal(err)
+				}
+				// Use the runtime's numeric primitive, not a decimal approximation
+				// or the result from a different host math library.
+				want := math.Pow(tc.baseNumber, tc.exponentNumber)
+				for index, name := range []string{"folded", "runtime", "coerced", "math.pow"} {
+					got, ok := l.ToNumber(index + 1)
+					if !ok || !samePowerNumber(got, want) {
+						t.Errorf("%s = %g (%016x), want %g (%016x)", name, got, math.Float64bits(got), want, math.Float64bits(want))
+					}
+				}
+				if !l.ToBoolean(5) {
+					t.Error("folded and runtime powers did not address the same table key")
+				}
+				for _, coerce := range []bool{false, true} {
+					api := NewState()
+					if coerce {
+						api.PushString(tc.base)
+						api.PushString(tc.exponent)
+					} else {
+						api.PushNumber(tc.baseNumber)
+						api.PushNumber(tc.exponentNumber)
+					}
+					api.Arith(OpPow)
+					if got, ok := api.ToNumber(-1); !ok || !samePowerNumber(got, want) {
+						t.Errorf("Arith(OpPow), coerce=%v: %g, want %g", coerce, got, want)
+					}
+				}
+			})
+		}
+	}
+}
+
+func samePowerNumber(a, b float64) bool {
+	return math.Float64bits(a) == math.Float64bits(b) || math.IsNaN(a) && math.IsNaN(b)
+}

tables.go

@@ -5,11 +5,12 @@ import (
 )
 
 type table struct {
-	array         []value
-	hash          map[value]value
-	metaTable     *table
-	flags         byte
-	iterationKeys []value
+	array          []value
+	hash           map[value]value
+	metaTable      *table
+	flags          byte
+	iterationKeys  []value
+	iterationIndex map[value]int
 }
 
 func newTable() *table                     { return &table{hash: make(map[value]value)} }
@@ -82,6 +83,7 @@ func (t *table) maybeResizeArray(key int) bool {
 func (t *table) addOrInsertHash(k, v value) {
 	if _, ok := t.hash[k]; !ok {
 		t.iterationKeys = nil // invalidate iterations when adding an entry
+		t.iterationIndex = nil
 	}
 	t.hash[k] = v
 }
@@ -240,23 +242,32 @@ func (l *State) next(t *table, key int) bool {
 	}
 	if t.iterationKeys == nil {
 		j, keys := 0, make([]value, len(t.hash))
+		var indices map[value]int
+		if len(t.hash) > 0 {
+			indices = make(map[value]int, len(t.hash))
+		}
 		for hk := range t.hash {
 			keys[j] = hk
+			indices[hk] = j
 			j++
 		}
 		t.iterationKeys = keys
+		t.iterationIndex = indices
+	}
+	start := 0
+	if k != nil {
+		start = t.iterationIndex[k] + 1
 	}
-	found := k == nil
-	for i, hk := range t.iterationKeys {
+	for i := start; i < len(t.iterationKeys); i++ {
+		hk := t.iterationKeys[i]
 		if hk == nil { // skip deleted key
 		} else if _, present := t.hash[hk]; !present {
 			t.iterationKeys[i] = nil // mark key as deleted
-		} else if found {
+			delete(t.iterationIndex, hk)
+		} else {
 			l.stack[key] = hk
 			l.stack[key+1] = t.hash[hk]
 			return true
-		} else if l.equalObjects(hk, k) {
-			found = true
 		}
 	}
 	return false // no more elements

tables_test.go

@@ -0,0 +1,150 @@
+package lua
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestNextUsesRawKeyEquality(t *testing.T) {
+	for _, kind := range []string{"table", "userdata"} {
+		t.Run(kind, func(t *testing.T) {
+			l := NewState()
+			OpenLibraries(l)
+			if err := DoString(l, `return function(a, b, c)
+				local mt = {__eq = function() error("next must not invoke __eq") end}
+				debug.setmetatable(a, mt)
+				debug.setmetatable(b, mt)
+				debug.setmetatable(c, mt)
+				local t = {[a] = 1, [b] = 2, [c] = 3}
+				local seen, count = {}, 0
+				for k, v in pairs(t) do
+					assert(not seen[k])
+					seen[k] = true
+					assert(t[k] == v)
+					count = count + 1
+				end
+				assert(count == 3)
+			end`); err != nil {
+				t.Fatal(err)
+			}
+			for i := 0; i < 3; i++ {
+				if kind == "table" {
+					l.CreateTable(0, 0)
+				} else {
+					l.PushUserData(i)
+				}
+			}
+			if err := l.ProtectedCall(3, 0, 0); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+func TestNextHashTableChanges(t *testing.T) {
+	tests := []struct {
+		name, program string
+	}{
+		{"interleaved", `
+			local t = {a = 1, b = 2, c = 3, [false] = 4, [100] = 5, 6, 7}
+			local outer, count = {}, 0
+			for k, v in pairs(t) do
+				assert(not outer[k])
+				outer[k] = true
+				local inner, n = {}, 0
+				for x, y in pairs(t) do
+					assert(not inner[x])
+					inner[x] = true
+					assert(t[x] == y)
+					n = n + 1
+				end
+				assert(n == 7)
+				count = count + 1
+			end
+			assert(count == 7)
+		`},
+		{"update_and_delete", `
+			local t = {a = 1, b = 2, c = 3, d = 4}
+			local first = next(t)
+			local second = next(t, first)
+			t[second] = nil
+			local seen, count = {[first] = true}, 1
+			local k = first
+			while true do
+				local v
+				k, v = next(t, k)
+				if k == nil then break end
+				assert(not seen[k])
+				seen[k] = true
+				t[k] = v * 10
+				count = count + 1
+			end
+			assert(count == 3 and not seen[second])
+			count = 0
+			for k, v in pairs(t) do
+				assert(seen[k] and v == t[k])
+				count = count + 1
+			end
+			assert(count == 3)
+		`},
+		{"insert_after_iteration", `
+			local t = {a = 1, b = 2, c = 3}
+			for _ in pairs(t) do end
+			t.d = 4
+			local seen, count = {}, 0
+			for k, v in pairs(t) do
+				assert(not seen[k] and t[k] == v)
+				seen[k] = true
+				count = count + 1
+			end
+			assert(count == 4 and seen.d)
+		`},
+		{"promote_numeric_keys", `
+			local t = {[4] = 4, [8] = 8, [16] = 16, a = 17}
+			for _ in pairs(t) do end
+			for i = 1, 16 do t[i] = i end
+			local seen, count = {}, 0
+			for k, v in pairs(t) do
+				assert(not seen[k] and t[k] == v)
+				seen[k] = true
+				count = count + 1
+			end
+			assert(count == 17 and seen.a)
+		`},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := NewState()
+			OpenLibraries(l)
+			if err := DoString(l, tt.program); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+func BenchmarkNextHash(b *testing.B) {
+	for _, size := range []int{1000, 10000, 50000} {
+		b.Run(fmt.Sprint(size), func(b *testing.B) {
+			l := NewState()
+			l.CreateTable(0, size)
+			for i := 0; i < size; i++ {
+				l.PushString(fmt.Sprintf("%d_MySuperKey", 100000000+i))
+				l.PushInteger(i)
+				l.SetTable(-3)
+			}
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				count := 0
+				for l.PushNil(); l.Next(-2); {
+					count++
+					l.Pop(1)
+				}
+				if count != size {
+					b.Fatalf("visited %d entries, want %d", count, size)
+				}
+			}
+		})
+	}
+}

types.go

@@ -238,10 +238,6 @@ func arith(op Operator, v1, v2 float64) float64 {
 	case OpMod:
 		return v1 - math.Floor(v1/v2)*v2
 	case OpPow:
-		// Golang bug: math.Pow(10.0, 33.0) is incorrect by 1 bit.
-		if v1 == 10.0 && float64(int(v2)) == v2 {
-			return math.Pow10(int(v2))
-		}
 		return math.Pow(v1, v2)
 	case OpUnaryMinus:
 		return -v1