From 9da0d0ffbdaf9612319b2f074d61ef99cff2ef08 Mon Sep 17 00:00:00 2001
From: CPunch <sethtstubbs@gmail.com>
Date: Sun, 14 Aug 2022 01:36:05 -0500
Subject: [PATCH] lp: support FORPREP && FORLOOP

---
 README.md  | 58 +++++++++++++++++++++++++++++-------------------------
 lparser.py | 23 +++++++++++++++++++---
 lundump.py |  2 +-
 3 files changed, 52 insertions(+), 31 deletions(-)

diff --git a/README.md b/README.md
index 432abf6..d461be5 100644
--- a/README.md
+++ b/README.md
@@ -12,47 +12,51 @@ Lua has a relatively small instruction set (only 38 different opcodes!). This ma
 
 ```sh
 > cat example.lua && luac5.1 -o example.luac example.lua
-local i, x = 10, 2
+local total = 0
 
-repeat
-    print(i + x)
-    i = i - 1
-until i < 0
+for i = 0, 9, 1 do
+    total = total + i
+    print(total)
+end
 > python main.py example.luac
 example.luac
 
 ==== [[example.lua's constants]] ====
 
-0: [NUMBER] 10.0
-1: [NUMBER] 2.0
-2: [STRING] print
-3: [NUMBER] 1.0
-4: [NUMBER] 0.0
+0: [NUMBER] 0.0
+1: [NUMBER] 9.0
+2: [NUMBER] 1.0
+3: [STRING] print
 
 ==== [[example.lua's locals]] ====
 
-R[0]: i
-R[1]: x
+R[0]: total
+R[1]: (for index)
+R[2]: (for limit)
+R[3]: (for step)
+R[4]: i
 
 ==== [[example.lua's dissassembly]] ====
 
-[  0]      LOADK :   R[0]   K[0]               ; load 10.0 into R[0]
-[  1]      LOADK :   R[1]   K[1]               ; load 2.0 into R[1]
-[  2]  GETGLOBAL :   R[2]   K[2]               ; move _G["print"] into R[2]
-[  3]        ADD :   R[3]   R[0]   R[1]        ; add R[1] to R[0], place into R[3]
-[  4]       CALL :      2      2      1        ; 
-[  5]        SUB :   R[0]   R[0]   K[3]        ; sub K[3] from R[0], place into R[0]
-[  6]         LT :   R[0]   R[0]   K[4]        ; 
-[  7]        JMP :   R[0]     -6               ; 
-[  8]     RETURN :      0      1      0        ; 
+[  0]      LOADK :   R[0]   K[0]               ; load 0.0 into R[0]
+[  1]      LOADK :   R[1]   K[0]               ; load 0.0 into R[1]
+[  2]      LOADK :   R[2]   K[1]               ; load 9.0 into R[2]
+[  3]      LOADK :   R[3]   K[2]               ; load 1.0 into R[3]
+[  4]    FORPREP :   R[1]      4               ; 
+[  5]        ADD :   R[0]   R[0]   R[4]        ; add R[4] to R[0], place into R[0]
+[  6]  GETGLOBAL :   R[5]   K[3]               ; move _G["print"] into R[5]
+[  7]       MOVE :      6      0      0        ; move R[0] into R[6]
+[  8]       CALL :      5      2      1        ; 
+[  9]    FORLOOP :   R[1]     -5               ; 
+[ 10]     RETURN :      0      1      0        ; 
 
 ==== [[example.lua's decompiled source]] ====
 
-local i = 10.0
-local x = 2.0
-repeat
-    print((i + x))
-    i = (i - 1.0)
-until i < 0.0 
+local total = 0.0
+for i = 0.0, 9.0, 1.0 do
+    total = (total + i)
+    print(total)
+end
+
 
 ```
\ No newline at end of file
diff --git a/lparser.py b/lparser.py
index b85e84d..b4ecb4d 100644
--- a/lparser.py
+++ b/lparser.py
@@ -29,6 +29,13 @@ class _Line:
         self.src = src
         self.scope = scope
 
+def isValidLocal(ident: str) -> bool:
+    for c in ident:
+        if c not in "abcdefghijklmnopqrstuvwxyz1234567890_":
+            return False
+
+    return True
+
 class LuaDecomp:
     def __init__(self, chunk: Chunk):
         self.chunk = chunk
@@ -124,11 +131,16 @@ class LuaDecomp:
 
     def __loadLocals(self):
         for i in range(len(self.chunk.locals)):
-            if not self.chunk.locals[i].name == "":
-                self.locals[i] = self.chunk.locals[i].name 
-            else:
+            name = self.chunk.locals[i].name
+            if isValidLocal(name):
+                self.locals[i] = name
+            elif "(for " not in name: # if it's a for loop register, ignore
                 self.__makeLocalIdentifier(i)
 
+    # when you *know* the register *has* to be a local (for loops, etc.)
+    def __getLocal(self, indx: int) -> str:
+        return self.locals[indx] if indx in self.locals else self.__makeLocalIdentifier(indx)
+
     def __getReg(self, indx: int) -> str:
         self.__addUseTraceback(indx)
 
@@ -336,5 +348,10 @@ class LuaDecomp:
         elif instr.opcode == Opcodes.RETURN:
             self.__endStatement()
             pass # no-op for now
+        elif instr.opcode == Opcodes.FORLOOP:
+            pass # no-op for now
+        elif instr.opcode == Opcodes.FORPREP:
+            self.__addExpr("for %s = %s, %s, %s " % (self.__getLocal(instr.A+3), self.__getReg(instr.A), self.__getReg(instr.A + 1), self.__getReg(instr.A + 2)))
+            self.__startScope("do", self.pc, instr.B)
         else:
             raise Exception("unsupported instruction: %s" % instr.toString())
\ No newline at end of file
diff --git a/lundump.py b/lundump.py
index 7fef536..15b8c28 100644
--- a/lundump.py
+++ b/lundump.py
@@ -400,7 +400,7 @@ class LuaUndump:
         # locals
         num = self.get_int()
         for i in range(num):
-            name = self.get_string(None) # local name
+            name = self.get_string(None)[:-1] # local name ([:-1] to remove the NULL terminator)
             start = self.get_int() # local start PC
             end = self.get_int() # local end PC
             chunk.appendLocal(Local(name, start, end))