From b72019000c0c0487ce6ce6868a17dc3eb90eb8ca Mon Sep 17 00:00:00 2001
From: Mark Tully <mark.tully@lockwood-publishing.com>
Date: Fri, 26 Jul 2024 22:16:53 +0100
Subject: [PATCH] fix compiler producing incorrect LOADNIL byte code

---
 compile.go     |  2 +-
 script_test.go | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/compile.go b/compile.go
index c3736777..e46e0d75 100644
--- a/compile.go
+++ b/compile.go
@@ -237,7 +237,7 @@ func (cd *codeStore) PropagateMV(top int, save *int, reg *int, inc int) {
 
 func (cd *codeStore) AddLoadNil(a, b, line int) {
 	last := cd.Last()
-	if opGetOpCode(last) == OP_LOADNIL && (opGetArgA(last)+opGetArgB(last)) == a {
+	if opGetOpCode(last) == OP_LOADNIL && (opGetArgB(last)+1) == a {
 		cd.SetB(cd.LastPC(), b)
 	} else {
 		cd.AddABC(OP_LOADNIL, a, b, 0, line)
diff --git a/script_test.go b/script_test.go
index 3c0675d6..48f48b3c 100644
--- a/script_test.go
+++ b/script_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -146,3 +147,86 @@ func TestGlua(t *testing.T) {
 func TestLua(t *testing.T) {
 	testScriptDir(t, luaTests, "_lua5.1-tests")
 }
+
+func TestMergingLoadNilBug(t *testing.T) {
+	// there was a bug where a multiple load nils were being incorrectly merged, and the following code exposed it
+	s := `
+    function test()
+        local a = 0
+        local b = 1
+        local c = 2
+        local d = 3
+        local e = 4		-- reg 4
+        local f = 5
+        local g = 6
+        local h = 7
+
+        if e == 4 then
+            e = nil		-- should clear reg 4, but clears regs 4-8 by mistake
+        end
+        if f == nil then
+            error("bad f")
+        end
+        if g == nil then
+            error("bad g")
+        end
+        if h == nil then
+            error("bad h")
+        end
+    end
+
+    test()
+`
+
+	L := NewState()
+	defer L.Close()
+	if err := L.DoString(s); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestMergingLoadNil(t *testing.T) {
+	// multiple nil assignments to consecutive registers should be merged
+	s := `
+		function test()
+			local a = 0
+			local b = 1
+			local c = 2
+
+			-- this should generate just one LOADNIL byte code instruction
+			a = nil
+			b = nil
+			c = nil
+
+			print(a,b,c)
+		end
+
+		test()`
+
+	chunk, err := parse.Parse(strings.NewReader(s), "test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	compiled, err := Compile(chunk, "test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(compiled.FunctionPrototypes) != 1 {
+		t.Fatal("expected 1 function prototype")
+	}
+
+	// there should be exactly 1 LOADNIL instruction in the byte code generated for the above
+	// anymore, and the LOADNIL merging is not working correctly
+	count := 0
+	for _, instr := range compiled.FunctionPrototypes[0].Code {
+		if opGetOpCode(instr) == OP_LOADNIL {
+			count++
+		}
+	}
+
+	if count != 1 {
+		t.Fatalf("expected 1 LOADNIL instruction, found %d", count)
+	}
+}