supports return

2 files changed, 110 insertions, 32 deletions

diff --git a/script/core/guide.lua b/script/core/guide.lua
index b42744ca..c7901448 100644
--- a/script/core/guide.lua
+++ b/script/core/guide.lua
@@ -37,7 +37,8 @@ function m.pushResult(status, mode, ref)
         or ref.type == 'setmethod'
         or ref.type == 'setindex'
         or ref.type == 'tableindex'
-        or ref.type == 'tablefield' then
+        or ref.type == 'tablefield'
+        or ref.type == 'function' then
             results[#results+1] = ref
         end
     elseif mode == 'ref' then
@@ -55,7 +56,8 @@ function m.pushResult(status, mode, ref)
         or ref.type == 'setindex'
         or ref.type == 'getindex'
         or ref.type == 'tableindex'
-        or ref.type == 'tablefield' then
+        or ref.type == 'tablefield'
+        or ref.type == 'function' then
             results[#results+1] = ref
         end
     elseif mode == 'field' then
@@ -106,7 +108,7 @@ function m.searchRefs(status, source, mode)
 
     local search
 
-    local function seachSource(obj, field, flag)
+    local function searchSource(obj, field, flag)
         local link = linker.getLink(obj)
         if not link then
             return
@@ -118,12 +120,60 @@ function m.searchRefs(status, source, mode)
         search(id, nil, flag)
     end
 
+    local function getReturnSetByFunc(func, index)
+        local call = func.parent
+        if call.type ~= 'call' then
+            return nil
+        end
+        if index == 0 then
+            return nil
+        end
+        if index == 1 then
+            return call.parent
+        else
+            for _, sel in ipairs(call.extParent) do
+                if sel.index == index then
+                    return sel
+                end
+            end
+        end
+        return nil
+    end
+
+    local function searchFunction(obj)
+        if obj.type ~= 'function' then
+            return
+        end
+        local link = linker.getLink(obj)
+        if not link then
+            return
+        end
+        if not link.freturn then
+            return
+        end
+        local func = guide.getParentFunction(obj)
+        if not func or func.type ~= 'function' then
+            return
+        end
+        local newStatus = m.status(status)
+        m.searchRefs(newStatus, func, 'ref')
+        for _, ref in ipairs(newStatus.results) do
+            local set = getReturnSetByFunc(ref, link.freturn)
+            if set then
+                local id = linker.getID(set)
+                if id then
+                    search(id)
+                end
+            end
+        end
+    end
+
     local function checkForward(link, field, flag)
         if not link.forward then
             return
         end
         for _, forwardSources in ipairs(link.forward) do
-            seachSource(forwardSources, field, flag)
+            searchSource(forwardSources, field, flag)
         end
     end
 
@@ -132,7 +182,18 @@ function m.searchRefs(status, source, mode)
             return
         end
         for _, backSources in ipairs(link.backward) do
-            seachSource(backSources, field, flag)
+            searchSource(backSources, field, flag)
+        end
+    end
+
+    local function checkLastID(id, field, flag)
+        local lastID = linker.getLastID(root, id)
+        if lastID then
+            local newField = id:sub(#lastID + 1)
+            if field then
+                newField = newField .. field
+            end
+            search(lastID, newField, flag)
         end
     end
 
@@ -156,26 +217,15 @@ function m.searchRefs(status, source, mode)
             if field == nil then
                 m.pushResult(status, mode, eachLink.source)
             end
-            --if flag & SEARCH_FLAG.backward == 0 then
-                checkForward(eachLink,  field, flag | SEARCH_FLAG.forward)
-            --end
-            --if flag & SEARCH_FLAG.forward  == 0 then
-                checkBackward(eachLink, field, flag | SEARCH_FLAG.backward)
-            --end
-        end
-        local lastID = linker.getLastID(root, id)
-        if lastID then
-            local newField = id:sub(#lastID + 1)
-            if field then
-                newField = newField .. field
-            end
-            search(lastID, newField, flag)
+            checkForward(eachLink,  field, flag | SEARCH_FLAG.forward)
+            checkBackward(eachLink, field, flag | SEARCH_FLAG.backward)
         end
-        --checkLastID(id, expect)
+        checkLastID(id, field, flag)
         stackCount = stackCount - 1
     end
 
-    seachSource(source)
+    searchSource(source)
+    searchFunction(source)
 end
 
 ---@class guide.status
diff --git a/script/core/linker.lua b/script/core/linker.lua
index d05f45c8..dd5e1ab7 100644
--- a/script/core/linker.lua
+++ b/script/core/linker.lua
@@ -22,6 +22,27 @@ local function getKey(source)
     elseif source.type == 'getmethod'
     or     source.type == 'setmethod' then
         return ('%q'):format(source.method and source.method[1] or ''), source.node
+    elseif source.type == 'setindex'
+    or     source.type == 'getindex' then
+        local index = source.index
+        if not index then
+            return '', source.node
+        end
+        if index.type == 'string' then
+            return ('%q'):format(index[1] or ''), source.node
+        else
+            return '', source.node
+        end
+    elseif source.type == 'tableindex' then
+        local index = source.index
+        if not index then
+            return '', source.parent
+        end
+        if index.type == 'string' then
+            return ('%q'):format(index[1] or ''), source.parent
+        else
+            return '', source.parent
+        end
     elseif source.type == 'table' then
         return source.start, nil
     elseif source.type == 'label' then
@@ -31,6 +52,10 @@ local function getKey(source)
             return source.node.start, nil
         end
         return nil, nil
+    elseif source.type == 'function' then
+        return source.start, nil
+    elseif source.type == 'select' then
+        return ('%d:%d'):format(source.start, source.index)
     end
     return nil, nil
 end
@@ -40,19 +65,13 @@ local function checkMode(source)
     or source.type == 'getglobal' then
         return 'g'
     end
-    if source.type == 'local'
-    or source.type == 'setlocal'
-    or source.type == 'getlocal' then
-        return 'l'
-    end
-    if source.type == 'label'
-    or source.type == 'goto' then
-        return 'l'
-    end
     if source.type == 'table' then
-        return 'l'
+        return 't'
     end
-    return nil
+    if source.type == 'select' then
+        return 's'
+    end
+    return 'l'
 end
 
 local function checkFunctionReturn(source)
@@ -250,6 +269,15 @@ function m.getLink(source)
     return source._link
 end
 
+---获取source的ID
+function m.getID(source)
+    local link = m.getLink(source)
+    if not link then
+        return nil
+    end
+    return link.id
+end
+
 ---编译整个文件的link
 ---@param  source parser.guide.object
 ---@return table