stash

3 files changed, 79 insertions, 29 deletions

diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index e301ee25..060a2173 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -1252,22 +1252,29 @@ local compilerSwitch = util.switch()
                 end
             end
         else
-            ---@cast key string
-            vm.compileByParentNode(source.node, key, function (src)
-                if src.value then
-                    if bindDocs(src) then
-                        vm.setNode(source, vm.compileNode(src))
-                    elseif src.value.type ~= 'nil' then
-                        vm.setNode(source, vm.compileNode(src.value))
-                        local node = vm.getNode(src)
-                        if node then
-                            vm.setNode(source, node)
+            if guide.isAssign(source) then
+                ---@cast key string
+                vm.compileByParentNode(source.node, key, function (src)
+                    if src.value then
+                        if bindDocs(src) then
+                            vm.setNode(source, vm.compileNode(src))
+                        elseif src.value.type ~= 'nil' then
+                            vm.setNode(source, vm.compileNode(src.value))
+                            local node = vm.getNode(src)
+                            if node then
+                                vm.setNode(source, node)
+                            end
                         end
+                    else
+                        vm.setNode(source, vm.compileNode(src))
                     end
-                else
-                    vm.setNode(source, vm.compileNode(src))
+                end)
+            else
+                local node = vm.traceNode(source)
+                if node then
+                    vm.setNode(source, node)
                 end
-            end)
+            end
         end
     end)
     : case 'setglobal'
diff --git a/script/vm/tracer.lua b/script/vm/tracer.lua
index 3f2380bd..33af9d0e 100644
--- a/script/vm/tracer.lua
+++ b/script/vm/tracer.lua
@@ -781,8 +781,9 @@ end
 ---@field package _tracer vm.tracer
 
 ---@param source parser.object
+---@param name string
 ---@return vm.tracer?
-local function createTracer(source)
+local function createTracer(source, name)
     local node = vm.compileNode(source)
     local tracer = node._tracer
     if tracer then
@@ -794,6 +795,7 @@ local function createTracer(source)
     end
     tracer = setmetatable({
         source    = source,
+        name      = name,
         assigns   = {},
         assignMap = {},
         getMap    = {},
@@ -808,10 +810,8 @@ local function createTracer(source)
 
     if source.type == 'local'
     or source.type == 'self' then
-        tracer.name = source[1]
         tracer:collectLocal()
     else
-        tracer.name = source.global:getName()
         tracer:collectGlobal()
     end
 
@@ -821,15 +821,29 @@ end
 ---@param source parser.object
 ---@return vm.node?
 function vm.traceNode(source)
-    local base
+    local base, name
     if source.type == 'getlocal'
     or source.type == 'setlocal' then
         base = source.node
-    else
+        ---@type string
+        name = source[1]
+    elseif vm.getGlobalNode(source) then
         base = vm.getGlobalBase(source)
+        if not base then
+            return nil
+        end
+        name = base.global:getCodeName()
+    else
+        base = vm.getVariableHead(source)
+        if not base then
+            return nil
+        end
+        name = vm.getVariableName(source)
+        if not name then
+            return nil
+        end
     end
-    assert(base)
-    local tracer = createTracer(base)
+    local tracer = createTracer(base, name)
     if not tracer then
         return nil
     end
diff --git a/script/vm/variable-id.lua b/script/vm/variable-id.lua
index cbc11409..a2aece65 100644
--- a/script/vm/variable-id.lua
+++ b/script/vm/variable-id.lua
@@ -6,8 +6,6 @@ local vm    = require 'vm.vm'
 ---@class vm.variable
 ---@field sets parser.object[]
 ---@field gets parser.object[]
----@field node? vm.node
----@field type 'global'|'local'
 
 ---@class parser.object
 ---@field package _variableID string|false
@@ -19,7 +17,7 @@ local compileSwitch = util.switch()
     : case 'local'
     : case 'self'
     : call(function (source)
-        source._variableID = ('l|%d'):format(source.start)
+        source._variableID = ('%d'):format(source.start)
         if not source.ref then
             return
         end
@@ -30,7 +28,7 @@ local compileSwitch = util.switch()
     : case 'setlocal'
     : case 'getlocal'
     : call(function (source)
-        source._variableID = ('l|%d'):format(source.node.start)
+        source._variableID = ('%d'):format(source.node.start)
         compileVariableID(source.next)
     end)
     : case 'getfield'
@@ -123,7 +121,6 @@ function vm.insertVariableID(id, source)
     if not root._variableIDs then
         root._variableIDs = util.multiTable(2, function (head)
             return {
-                type = head:sub(1, 1) == 'l' and 'local' or 'global',
                 sets = {},
                 gets = {},
             }
@@ -154,18 +151,33 @@ function compileVariableID(source)
 end
 
 ---@param source parser.object
----@return string|false
+---@return string?
 function vm.getVariableID(source)
     if source._variableID ~= nil then
-        return source._variableID
+        return source._variableID or nil
     end
     source._variableID = false
     local loc = getVariable(source)
     if not loc then
-        return source._variableID
+        return source._variableID or nil
     end
     compileVariableID(loc)
-    return source._variableID
+    return source._variableID or nil
+end
+
+---@param source parser.object
+---@return string?
+function vm.getVariableName(source)
+    local id = vm.getVariableID(source)
+    if not id then
+        return nil
+    end
+    local head = vm.getVariableHead(source)
+    if not head then
+        return nil
+    end
+    local name = id:gsub('%d+', head[1]):gsub(vm.ID_SPLITE, '.')
+    return name
 end
 
 ---@param source parser.object
@@ -248,3 +260,20 @@ function vm.getVariableFields(source, includeGets)
     end
     return fields
 end
+
+---@param source parser.object
+---@return parser.object?
+function vm.getVariableHead(source)
+    local id = vm.getVariableID(source)
+    if not id then
+        return nil
+    end
+    for _ = 1, 1000 do
+        if source.type == 'local'
+        or source.type == 'self' then
+            return source
+        end
+        source = source.node
+    end
+    return nil
+end