6 files changed, 59 insertions, 27 deletions
diff --git a/server/src/core/diagnostics.lua b/server/src/core/diagnostics.lua
index 7ebe297a..241d70ba 100644
--- a/server/src/core/diagnostics.lua
+++ b/server/src/core/diagnostics.lua
@@ -245,7 +245,11 @@ function mt:doDiagnostics(func, code, callback)
         self.datas[#self.datas+1] = data
     end)
     if coroutine.isyieldable() then
-        coroutine.yield()
+        if self.vm:isRemoved() then
+            coroutine.yield('stop')
+        else
+            coroutine.yield()
+        end
     end
 end
 
diff --git a/server/src/service.lua b/server/src/service.lua
index 586aa937..4fdd7f9c 100644
--- a/server/src/service.lua
+++ b/server/src/service.lua
@@ -432,7 +432,7 @@ function mt:doDiagnostics(uri)
     end
     local name = 'textDocument/publishDiagnostics'
     local obj = self._file[uri]
-    if not obj or not obj.vm then
+    if not obj or not obj.vm or obj.vm:isRemoved() then
         self._needDiagnostics[uri] = nil
         self:clearDiagnostics(uri)
         return
@@ -611,10 +611,12 @@ end
 function mt:_testMemory()
     local cachedVM = 0
     local cachedSource = 0
+    local cachedFunction = 0
     for _, obj in pairs(self._file) do
-        if obj.vm then
+        if obj.vm and not obj.vm:isRemoved() then
             cachedVM = cachedVM + 1
             cachedSource = cachedSource + #obj.vm.sources
+            cachedFunction = cachedFunction + #obj.vm.funcs
         end
     end
     local aliveVM = 0
@@ -648,33 +650,51 @@ function mt:_testMemory()
     end
 
     local totalFunction = 0
-    for _ in pairs(functionMgr.watch) do
-        totalFunction = totalFunction + 1
+    local alivedFunction = 0
+    local deadFunction = 0
+    for _, id in pairs(functionMgr.watch) do
+        if functionMgr.list[id] then
+            alivedFunction = alivedFunction + 1
+        else
+            deadFunction = deadFunction + 1
+        end
     end
 
     local mem = collectgarbage 'count'
     log.debug(('\n\z
     State\n\z
-    Mem:      [%.3f]kb\n\z
-    CachedVM: [%d]\n\z
-    AlivedVM: [%d]\n\z
-    DeadVM:   [%d]\n\z
-    CachedSrc:[%d]\n\z
-    AlivedSrc:[%d]\n\z
-    DeadSrc:  [%d]\n\z
-    TotalLoc: [%d]\n\z
-    TotalVal: [%d]\n\z
-    TotalFunc:[%d]\n\z'):format(
+    Mem:       [%.3f]kb\n\z
+                   \n\z
+    CachedVM:  [%d]\n\z
+    AlivedVM:  [%d]\n\z
+    DeadVM:    [%d]\n\z
+                   \n\z
+    CachedSrc: [%d]\n\z
+    AlivedSrc: [%d]\n\z
+    DeadSrc:   [%d]\n\z
+                   \n\z
+    CachedFunc:[%d]\n\z
+    AlivedFunc:[%d]\n\z
+    DeadFunc:  [%d]\n\z
+                   \n\z
+    TotalLoc:  [%d]\n\z
+    TotalVal:  [%d]\n\z'):format(
         mem,
+
         cachedVM,
         aliveVM,
         deadVM,
+
         cachedSource,
         alivedSource,
         deadSource,
+
+        cachedFunction,
+        alivedFunction,
+        deadFunction,
+
         totalLocal,
-        totalValue,
-        totalFunction
+        totalValue
     ))
 end
 
diff --git a/server/src/vm/function.lua b/server/src/vm/function.lua
index 2fe1ee41..c6915f25 100644
--- a/server/src/vm/function.lua
+++ b/server/src/vm/function.lua
@@ -3,6 +3,8 @@ local valueMgr = require 'vm.value'
 local localMgr = require 'vm.local'
 local sourceMgr = require 'vm.source'
 
+local Id = 0
+local List = {}
 local Watch = setmetatable({}, {__mode = 'kv'})
 
 local mt = {}
@@ -296,7 +298,11 @@ function mt:getSource()
 end
 
 function mt:kill()
+    if self._removed then
+        return
+    end
     self._removed = true
+    List[self.id] = nil
 end
 
 local function create(source)
@@ -307,19 +313,23 @@ local function create(source)
     if not id then
         error('Not instanted source')
     end
+    Id = Id + 1
     local self = setmetatable({
         source = id,
+        id = id,
         locals = {},
         finishs = {},
         args = {},
         argValues = {},
     }, mt)
     self:push(source)
-    Watch[self] = true
+    Watch[self] = id
+    List[id] = self
     return self
 end
 
 return {
     create = create,
     watch = Watch,
+    list = List,
 }
diff --git a/server/src/vm/source.lua b/server/src/vm/source.lua
index b74d10d2..af8b9ab3 100644
--- a/server/src/vm/source.lua
+++ b/server/src/vm/source.lua
@@ -103,7 +103,7 @@ local function instant(source)
     Id = Id + 1
     source.id = Id
     List[Id] = source
-    Watch[source] = Id 
+    Watch[source] = Id
     setmetatable(source, mt)
     return true
 end
diff --git a/server/src/vm/value.lua b/server/src/vm/value.lua
index a7170659..ad1deccf 100644
--- a/server/src/vm/value.lua
+++ b/server/src/vm/value.lua
@@ -389,6 +389,9 @@ function mt:getFunction()
     if not func then
         return nil
     end
+    if func._removed then
+        return nil
+    end
     if not func:getSource() then
         self._func = nil
     end
diff --git a/server/src/vm/vm.lua b/server/src/vm/vm.lua
index 4b044916..77f428ae 100644
--- a/server/src/vm/vm.lua
+++ b/server/src/vm/vm.lua
@@ -637,14 +637,6 @@ function mt:getUnary(exp)
 end
 
 function mt:getExp(exp)
-    if coroutine.isyieldable() then
-        if self.lsp:isNeedCompile(self.uri) then
-            coroutine.yield()
-        else
-            self:remove()
-            coroutine.yield('stop')
-        end
-    end
     self:instantSource(exp)
     local tp = exp.type
     if     tp == 'nil' then
@@ -1129,6 +1121,9 @@ function mt:createEnvironment(ast)
 end
 
 function mt:eachSource(callback)
+    if self._removed then
+        return
+    end
     local sources = self.sources
     for i = 1, #sources do
         callback(sources[i])