Handles os.exit by wrapping coroutine.(create|wrap)

Signed-off-by: Norman Clarke <norman@njclarke.com>

1 files changed, 27 insertions, 0 deletions

diff --git a/src/luacov.lua b/src/luacov.lua
index 0c97674..a2334e3 100644
--- a/src/luacov.lua
+++ b/src/luacov.lua
@@ -54,3 +54,30 @@ if not tick then
 end
 
 debug.sethook(on_line, "l")
+
+rawcoroutinecreate = coroutine.create
+
+function coroutine.create(...)
+  local co = rawcoroutinecreate(...)
+  debug.sethook(co, on_line, "l")
+  return co
+end
+
+function coroutine.wrap(...)
+  local co = rawcoroutinecreate(...)
+  debug.sethook(co, on_line, "l")
+  return function()
+    local r = { coroutine.resume(co) }
+    if not r[1] then
+      error(r[2])
+    end
+    return unpack(r, 2)
+  end
+end
+
+
+local rawexit = os.exit
+function os.exit(...)
+  on_exit()
+  rawexit(...)
+end
\ No newline at end of file