From e9e7350447642fb855786f749768e36d183c7f1a Mon Sep 17 00:00:00 2001
From: Roxy Light <roxy@256lights.llc>
Date: Sat, 14 Dec 2024 16:42:36 -0800
Subject: [PATCH] Add `*Prototype.LocalName` method

---
 internal/luacode/prototype.go      | 25 ++++++++++-
 internal/luacode/prototype_test.go | 68 ++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/internal/luacode/prototype.go b/internal/luacode/prototype.go
index 919524a..e44a06b 100644
--- a/internal/luacode/prototype.go
+++ b/internal/luacode/prototype.go
@@ -41,7 +41,9 @@ type Prototype struct {
 
 	// Debug information:
 
-	Source          Source
+	Source Source
+	// LocalVariables is a list of the function's local variables in declaration order.
+	// It is guaranteed that LocalVariables[i].StartPC <= LocalVariables[i+1].StartPC.
 	LocalVariables  []LocalVariable
 	LineInfo        LineInfo
 	LineDefined     int
@@ -74,6 +76,27 @@ func (f *Prototype) StripDebug() *Prototype {
 	return f2
 }
 
+// LocalName returns the name of the local variable the given register represents
+// during the execution of the given instruction,
+// or the empty string if the register does not represent a local variable
+// (or the debug information has been stripped).
+func (f *Prototype) LocalName(register uint8, pc int) string {
+	for _, v := range f.LocalVariables {
+		if v.StartPC > pc {
+			// Local variables are ordered by StartPC,
+			// so this variable and any subsequent ones will be out of scope.
+			break
+		}
+		if pc < v.EndPC {
+			if register == 0 {
+				return v.Name
+			}
+			register--
+		}
+	}
+	return ""
+}
+
 func (f *Prototype) hasUpvalueNames() bool {
 	for _, upval := range f.Upvalues {
 		if upval.Name != "" {
diff --git a/internal/luacode/prototype_test.go b/internal/luacode/prototype_test.go
index 21b94b0..db21b57 100644
--- a/internal/luacode/prototype_test.go
+++ b/internal/luacode/prototype_test.go
@@ -5,6 +5,7 @@ package luacode
 
 import (
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -46,3 +47,70 @@ func FuzzPrototypeMarshalBinary(f *testing.F) {
 		}
 	})
 }
+
+func TestLocalName(t *testing.T) {
+	tests := []struct {
+		register uint8
+		pc       int
+		want     string
+	}{
+		{pc: 0, register: 0, want: ""},
+		{pc: 0, register: 1, want: ""},
+		{pc: 0, register: 2, want: ""},
+
+		{pc: 1, register: 0, want: ""},
+		{pc: 1, register: 1, want: ""},
+		{pc: 1, register: 2, want: ""},
+
+		{pc: 2, register: 0, want: "a"},
+		{pc: 2, register: 1, want: ""},
+		{pc: 2, register: 2, want: ""},
+
+		{pc: 3, register: 0, want: "a"},
+		{pc: 3, register: 1, want: "c"},
+		{pc: 3, register: 2, want: ""},
+
+		{pc: 5, register: 0, want: "a"},
+		{pc: 5, register: 1, want: "c"},
+		{pc: 5, register: 2, want: ""},
+
+		{pc: 6, register: 0, want: "a"},
+		{pc: 6, register: 1, want: "c"},
+		{pc: 6, register: 2, want: "b"},
+
+		{pc: 7, register: 0, want: "a"},
+		{pc: 7, register: 1, want: "c"},
+		{pc: 7, register: 2, want: "b"},
+
+		{pc: 8, register: 0, want: "a"},
+		{pc: 8, register: 1, want: "c"},
+		{pc: 8, register: 2, want: ""},
+
+		{pc: 9, register: 0, want: "a"},
+		{pc: 9, register: 1, want: "c"},
+		{pc: 9, register: 2, want: ""},
+
+		{pc: 10, register: 0, want: "a"},
+		{pc: 10, register: 1, want: "c"},
+		{pc: 10, register: 2, want: "d"},
+	}
+
+	chunk, err := os.ReadFile(filepath.Join("testdata", "Scoping", "luac.out"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	p := new(Prototype)
+	if err := p.UnmarshalBinary(chunk); err != nil {
+		t.Fatal(err)
+	}
+	t.Log("Locals:")
+	for _, v := range p.LocalVariables {
+		t.Logf("%s\tscope:[%d, %d)", v.Name, v.StartPC, v.EndPC)
+	}
+
+	for _, test := range tests {
+		if got := p.LocalName(test.register, test.pc); got != test.want {
+			t.Errorf("p.LocalName(%d, %d) = %q; want %q", test.register, test.pc, got, test.want)
+		}
+	}
+}