eachMeta

3 files changed, 33 insertions, 27 deletions

diff --git a/script-beta/vm/eachField.lua b/script-beta/vm/eachField.lua
index 656d6fe7..dce5eec4 100644
--- a/script-beta/vm/eachField.lua
+++ b/script-beta/vm/eachField.lua
@@ -35,7 +35,7 @@ local function ofENV(source, callback)
             callback(ref)
             if guide.getName(ref) == '_G' then
                 checkNext(ref, callback)
-                vm.checkMetaValue(ref, callback)
+                vm.eachMetaValue(ref, callback)
                 local call, index = vm.getArgInfo(ref)
                 local special = vm.getSpecial(call)
                 if (special == 'rawset' or special == 'rawget')
@@ -45,7 +45,7 @@ local function ofENV(source, callback)
             end
         elseif ref.type == 'getlocal' then
             checkNext(ref, callback)
-            vm.checkMetaValue(ref, callback)
+            vm.eachMetaValue(ref, callback)
         end
         vm.eachFieldInTable(ref.value, callback)
     end
@@ -57,7 +57,7 @@ local function ofLocal(source, callback)
     else
         vm.eachRef(source, function (src)
             checkNext(src, callback)
-            vm.checkMetaValue(src, callback)
+            vm.eachMetaValue(src, callback)
             vm.eachFieldInTable(src.value, callback)
         end)
     end
@@ -66,7 +66,7 @@ end
 local function ofGlobal(source, callback)
     vm.eachRef(source, function (src)
         checkNext(src, callback)
-        vm.checkMetaValue(src, callback)
+        vm.eachMetaValue(src, callback)
         vm.eachFieldInTable(src.value, callback)
     end)
 end
@@ -74,7 +74,7 @@ end
 local function ofGetField(source, callback)
     vm.eachRef(source, function (src)
         checkNext(src, callback)
-        vm.checkMetaValue(src, callback)
+        vm.eachMetaValue(src, callback)
         vm.eachFieldInTable(src.value, callback)
     end)
 end
@@ -91,7 +91,7 @@ end
 local function ofTableField(source, callback)
     vm.eachRef(source, function (src)
         checkNext(src, callback)
-        vm.checkMetaValue(src, callback)
+        vm.eachMetaValue(src, callback)
     end)
 end
 
@@ -107,7 +107,7 @@ function vm.eachFieldInTable(value, callback)
         if field.type == 'tablefield'
         or field.type == 'tableindex' then
             callback(field)
-            vm.checkMetaValue(field, callback)
+            vm.eachMetaValue(field, callback)
         end
     end
 end
diff --git a/script-beta/vm/getClass.lua b/script-beta/vm/getClass.lua
index f1a50568..295ffae3 100644
--- a/script-beta/vm/getClass.lua
+++ b/script-beta/vm/getClass.lua
@@ -1,10 +1,9 @@
 local vm = require 'vm.vm'
 
-local function getClass(source, deep)
+local function getClass(source, classes, deep)
     if deep > 3 then
-        return nil
+        return
     end
-    local classes = {}
     vm.eachField(source, function (src)
         local key = vm.getKeyName(src)
         local lkey = key:lower()
@@ -17,12 +16,19 @@ local function getClass(source, deep)
             end
         end
     end)
-    if #classes == 0 then
-        return nil
+    if #classes ~= 0 then
+        return
     end
-    return vm.mergeTypeViews(table.unpack(classes))
+    vm.eachMeta(source, function (mt)
+        getClass(mt, classes, deep + 1)
+    end)
 end
 
 function vm.getClass(source)
-    return getClass(source, 1)
+    local classes = {}
+    getClass(source, classes, 1)
+    if #classes == 0 then
+        return nil
+    end
+    return vm.mergeTypeViews(table.unpack(classes))
 end
diff --git a/script-beta/vm/getMeta.lua b/script-beta/vm/getMeta.lua
index 5bb341b2..4915f616 100644
--- a/script-beta/vm/getMeta.lua
+++ b/script-beta/vm/getMeta.lua
@@ -1,23 +1,17 @@
 local vm = require 'vm.vm'
 
-local function checkMetaArg1(source, callback)
+local function eachMetaOfArg1(source, callback)
     local node, index = vm.getArgInfo(source)
     local special = vm.getSpecial(node)
     if special == 'setmetatable' and index == 1 then
         local mt = node.next.args[2]
         if mt then
-            vm.eachField(mt, function (src)
-                if vm.getKeyName(src) == 's|__index' then
-                    if src.value then
-                        vm.eachField(src.value, callback)
-                    end
-                end
-            end)
+            callback(mt)
         end
     end
 end
 
-local function checkMetaRecv(source, callback)
+local function eachMetaOfRecv(source, callback)
     if not source or source.type ~= 'select' then
         return
     end
@@ -35,6 +29,12 @@ local function checkMetaRecv(source, callback)
     vm.eachFieldInTable(call.args[1])
     local mt = call.args[2]
     if mt then
+        callback(mt)
+    end
+end
+
+function vm.eachMetaValue(source, callback)
+    vm.eachMeta(source, function (mt)
         vm.eachField(mt, function (src)
             if vm.getKeyName(src) == 's|__index' then
                 if src.value then
@@ -42,10 +42,10 @@ local function checkMetaRecv(source, callback)
                 end
             end
         end)
-    end
+    end)
 end
 
-function vm.checkMetaValue(source, callback)
-    checkMetaArg1(source, callback)
-    checkMetaRecv(source.value, callback)
+function vm.eachMeta(source, callback)
+    eachMetaOfArg1(source, callback)
+    eachMetaOfRecv(source.value, callback)
 end