3 files changed, 30 insertions, 19 deletions
diff --git a/server/src/vm/source.lua b/server/src/vm/source.lua
index 8c80262a..22132166 100644
--- a/server/src/vm/source.lua
+++ b/server/src/vm/source.lua
@@ -62,6 +62,21 @@ function mt:getUri()
     return self.uri
 end
 
+function mt:setVM()
+    self.vm = vm
+end
+
+function mt:isRemoved()
+    if not self.vm then
+        return true
+    end
+    if self.vm._removed then
+        self.vm = nil
+        return true
+    end
+    return false
+end
+
 function mt:set(name, v)
     if not self._flag then
         self._flag = {}
@@ -80,10 +95,11 @@ function mt:getName()
     return self[1]
 end
 
-return function (source)
+return function (vm, source)
     if source._hasInstant then
         return false
     end
+    source.vm = vm
     setmetatable(source, mt)
     return true
 end
diff --git a/server/src/vm/value.lua b/server/src/vm/value.lua
index fe727caf..f41d30e0 100644
--- a/server/src/vm/value.lua
+++ b/server/src/vm/value.lua
@@ -246,9 +246,12 @@ function mt:addInfo(tp, source, ...)
     if self._info[source] then
         return
     end
+    if not source then
+        return
+    end
     local info = {
         type = tp,
-        source = source or getDefaultSource(),
+        source = source,
         ...
     }
     self._info[#self._info+1] = info
diff --git a/server/src/vm/vm.lua b/server/src/vm/vm.lua
index ad25235c..d6524958 100644
--- a/server/src/vm/vm.lua
+++ b/server/src/vm/vm.lua
@@ -12,7 +12,7 @@ local mt = {}
 mt.__index = mt
 
 function mt:getDefaultSource()
-    return {
+    return self:instantSource {
         start = 0,
         finish = 0,
         uri = self:getUri(),
@@ -27,19 +27,6 @@ function mt:scopePop()
     self.currentFunction:pop()
 end
 
-function mt:eachInfo(var, callback)
-    if not self.results.infos[var] then
-        return nil
-    end
-    for _, info in ipairs(self.results.infos[var]) do
-        local res = callback(info)
-        if res ~= nil then
-            return res
-        end
-    end
-    return nil
-end
-
 function mt:buildTable(source)
     local tbl = self:createValue('table', source)
     if not source then
@@ -723,7 +710,7 @@ function mt:doReturn(action)
     values:eachValue(function (n, value)
         value.uri = self:getUri()
         func:setReturn(n, value)
-        value:addInfo('return', action[n])
+        value:addInfo('return', action[n] or value.source)
     end)
 end
 
@@ -1077,10 +1064,11 @@ function mt:getUri()
 end
 
 function mt:instantSource(source)
-    if instantSource(source) then
+    if instantSource(self, source) then
         source:setUri(self:getUri())
         self.sources[#self.sources+1] = source
     end
+    return source
 end
 
 function mt:bindLocal(source, loc, action)
@@ -1118,7 +1106,7 @@ function mt:createLocal(key, source, value)
     loc = createLocal(key, source, value)
     self:saveLocal(key, loc)
     self:bindLocal(source, loc, 'local')
-    value:addInfo('local', source)
+    value:addInfo('local', source or self:getDefaultSource())
     return loc
 end
 
@@ -1141,6 +1129,10 @@ function mt:eachSource(callback)
     end
 end
 
+function mt:remove()
+    self._removed = true
+end
+
 local function compile(ast, lsp, uri)
     local vm = setmetatable({
         funcs   = {},