Fix Module tracking across ownership transfer. (#494)

Since we use simple structs, the same handle results
in the same object object, triggering user-after-free warnings.

Author: maleadt

src/debug.jl

@@ -9,12 +9,12 @@ const memcheck_enabled = parse(Bool, @load_preference("memcheck", "false"))
 
 const tracked_objects = Dict{Any,Any}()
 
-function mark_alloc(obj::Any)
+function mark_alloc(obj::Any; allow_overwrite::Bool=false)
     @static if memcheck_enabled
         io = Core.stdout
         new_alloc_bt = backtrace()[2:end]
 
-        if haskey(tracked_objects, obj)
+        if haskey(tracked_objects, obj) && !allow_overwrite
             old_alloc_bt, dispose_bt = tracked_objects[obj]
             if dispose_bt == nothing
                 print("\nWARNING: An instance of $(typeof(obj)) was not properly disposed of, and a new allocation will overwrite it.")

src/executionengine/ts_module.jl

@@ -89,13 +89,12 @@ function ThreadSafeModule(mod::Module)
         #      Module and a regular Context, only a method to create one from a Module
         #      and a pre-existing TSContext, which isn't useful...
         # TODO: expose the other convenience method?
-        # XXX: work around this by serializing/deserializing the module in the correct contex
+        # XXX: work around this by serializing/deserializing in the correct context
         bitcode = convert(MemoryBuffer, mod)
-        new_mod = context!(context(ts_context())) do
+        dispose(mod)
+        mod = context!(context(ts_context())) do
             parse(Module, bitcode)
         end
-        dispose(mod)
-        mod = new_mod
     end
     @assert context(mod) == context(ts_context())
 
@@ -140,7 +139,7 @@ end
 
 function tsm_callback(data::Ptr{Cvoid}, ref::API.LLVMModuleRef)
     cb = Base.unsafe_pointer_to_objref(data)::ThreadSafeModuleCallback
-    mod = Module(ref)
+    mod = mark_alloc(Module(ref); allow_overwrite=true)
     ctx = context(mod)
     activate(ctx)
     try
@@ -149,6 +148,7 @@ function tsm_callback(data::Ptr{Cvoid}, ref::API.LLVMModuleRef)
         msg = sprint(Base.display_error, err, Base.catch_backtrace())
         return API.LLVMCreateStringError(msg)
     finally
+        mark_dispose(mod)
         deactivate(ctx)
     end
     return convert(API.LLVMErrorRef, C_NULL)