Add a memory checker (#420)

Detects leaks, double-frees, and use-after-frees
of manually memory managed LLVM objects

Author: maleadt

.github/workflows/ci.yml

@@ -88,7 +88,12 @@ jobs:
         if: runner.os == 'Windows'
 
       - name: Run tests
-        run: julia --project -e 'using Pkg; Pkg.test(; coverage=true, julia_args=`-g2`)'
+        run: |
+          julia -e 'open("LocalPreferences.toml", "a") do io
+                      println(io, "typecheck = \"true\"")
+                      println(io, "memcheck = \"true\"")
+                    end'
+          julia --project -e 'using Pkg; Pkg.test(; coverage=true)'
         env:
           JULIA_LLVM_ARGS: ${{ matrix.llvm_args }}
       - uses: julia-actions/julia-processcoverage@v1
@@ -151,7 +156,12 @@ jobs:
         run: julia --project=deps deps/build_ci.jl
 
       - name: Run tests
-        run: julia --project -e 'using Pkg; Pkg.test(; coverage=true, julia_args=`-g2`)'
+        run: |
+          julia -e 'open("LocalPreferences.toml", "a") do io
+                      println(io, "typecheck = \"true\"")
+                      println(io, "memcheck = \"true\"")
+                    end'
+          julia --project -e 'using Pkg; Pkg.test(; coverage=true)'
         env:
           JULIA_LLVM_ARGS: ${{ matrix.llvm_args }}
       - uses: julia-actions/julia-processcoverage@v1

LocalPreferences.toml

@@ -3,3 +3,10 @@
 # but if you are using a custom version of LLVM you will need to provide your own,
 # e.g., by running `deps/build_local.jl`.
 #libLLVMExtra = "/path/to/libLLVMExtra.so"
+
+# whether to enable additional object type checking
+#typecheck = "false"
+
+# whether to enable object memory checking. these are expensive checks that keep
+# track of allocated objects, whether they are freed correctly, etc.
+#memcheck = "false"

examples/constrained.jl

@@ -44,6 +44,7 @@ meta(::Type{FPExceptStrict}) = "fpexcept.strict"
         intrinsic = Intrinsic("llvm.experimental.constrained.$(func(F))")
         intrinsic_fun = LLVM.Function(mod, intrinsic, [typ])
         ftype = LLVM.FunctionType(intrinsic,[typ])
+
         # generate IR
         @dispose builder=IRBuilder() begin
             entry = BasicBlock(llvm_f, "entry")

src/LLVM.jl

@@ -1,5 +1,6 @@
 module LLVM
 
+using Preferences
 using Unicode
 using Printf
 using Libdl
@@ -59,6 +60,9 @@ has_oldpm() = LLVM.version() < v"17"
 has_newpm() = LLVM.version() >= v"15"
 has_julia_ojit() = VERSION >= v"1.10.0-DEV.1395"
 
+# helpers
+include("debug.jl")
+
 # LLVM API wrappers
 include("support.jl")
 if LLVM.version() < v"17"
@@ -133,6 +137,7 @@ function __init__()
 
     _install_handlers()
     _install_handlers(GlobalContext())
+    atexit(report_leaks)
 end
 
 end

src/analysis.jl

@@ -29,10 +29,11 @@ export DomTree, dominates
     ref::API.LLVMDominatorTreeRef
 end
 
-Base.unsafe_convert(::Type{API.LLVMDominatorTreeRef}, domtree::DomTree) = domtree.ref
+Base.unsafe_convert(::Type{API.LLVMDominatorTreeRef}, domtree::DomTree) =
+    mark_use(domtree).ref
 
-DomTree(f::Function) = DomTree(API.LLVMCreateDominatorTree(f))
-dispose(domtree::DomTree) = API.LLVMDisposeDominatorTree(domtree)
+DomTree(f::Function) = mark_alloc(DomTree(API.LLVMCreateDominatorTree(f)))
+dispose(domtree::DomTree) = mark_dispose(API.LLVMDisposeDominatorTree, domtree)
 
 function dominates(domtree::DomTree, A::Instruction, B::Instruction)
     API.LLVMDominatorTreeInstructionDominates(domtree, A, B) |> Bool
@@ -48,10 +49,11 @@ export PostDomTree, dominates
 end
 
 Base.unsafe_convert(::Type{API.LLVMPostDominatorTreeRef}, postdomtree::PostDomTree) =
-    postdomtree.ref
+    mark_use(postdomtree).ref
 
-PostDomTree(f::Function) = PostDomTree(API.LLVMCreatePostDominatorTree(f))
-dispose(postdomtree::PostDomTree) = API.LLVMDisposePostDominatorTree(postdomtree)
+PostDomTree(f::Function) = mark_alloc(PostDomTree(API.LLVMCreatePostDominatorTree(f)))
+dispose(postdomtree::PostDomTree) =
+    mark_dispose(API.LLVMDisposePostDominatorTree, postdomtree)
 
 function dominates(postdomtree::PostDomTree, A::Instruction, B::Instruction)
     API.LLVMPostDominatorTreeInstructionDominates(postdomtree, A, B) |> Bool

src/bitcode.jl

@@ -9,7 +9,11 @@ function Base.parse(::Type{Module}, membuf::MemoryBuffer)
     Module(out_ref[])
 end
 
-Base.parse(::Type{Module}, data::Vector) = parse(Module, MemoryBuffer(data, "", false))
+function Base.parse(::Type{Module}, data::Vector)
+    @dispose membuf = MemoryBuffer(data, "", false) begin
+        parse(Module, membuf)
+    end
+end
 
 
 ## writer

src/buffer.jl

@@ -4,16 +4,18 @@ export MemoryBuffer, MemoryBufferFile, dispose
     ref::API.LLVMMemoryBufferRef
 end
 
-Base.unsafe_convert(::Type{API.LLVMMemoryBufferRef}, membuf::MemoryBuffer) = membuf.ref
+Base.unsafe_convert(::Type{API.LLVMMemoryBufferRef}, membuf::MemoryBuffer) =
+    mark_use(membuf).ref
 
 function MemoryBuffer(data::Vector{T}, name::String="", copy::Bool=true) where T<:Union{UInt8,Int8}
     ptr = pointer(data)
     len = Csize_t(length(data))
-    if copy
-        return MemoryBuffer(API.LLVMCreateMemoryBufferWithMemoryRangeCopy(ptr, len, name))
+    membuf = if copy
+        MemoryBuffer(API.LLVMCreateMemoryBufferWithMemoryRangeCopy(ptr, len, name))
     else
-        return MemoryBuffer(API.LLVMCreateMemoryBufferWithMemoryRange(ptr, len, name, false))
+        MemoryBuffer(API.LLVMCreateMemoryBufferWithMemoryRange(ptr, len, name, false))
     end
+    mark_alloc(membuf)
 end
 
 function MemoryBuffer(f::Core.Function, args...; kwargs...)
@@ -48,7 +50,7 @@ function MemoryBufferFile(f::Core.Function, args...; kwargs...)
     end
 end
 
-dispose(membuf::MemoryBuffer) = API.LLVMDisposeMemoryBuffer(membuf)
+dispose(membuf::MemoryBuffer) = mark_dispose(API.LLVMDisposeMemoryBuffer, membuf)
 
 Base.length(membuf::MemoryBuffer) = API.LLVMGetBufferSize(membuf)
 

src/core/context.jl

@@ -6,10 +6,10 @@ export Context, dispose, GlobalContext
     ref::API.LLVMContextRef
 end
 
-Base.unsafe_convert(::Type{API.LLVMContextRef}, ctx::Context) = ctx.ref
+Base.unsafe_convert(::Type{API.LLVMContextRef}, ctx::Context) = mark_use(ctx).ref
 
 function Context(; opaque_pointers=nothing)
-    ctx = Context(API.LLVMContextCreate())
+    ctx = mark_alloc(Context(API.LLVMContextCreate()))
     if opaque_pointers !== nothing
         opaque_pointers!(ctx, opaque_pointers)
     end
@@ -20,7 +20,7 @@ end
 
 function dispose(ctx::Context)
     deactivate(ctx)
-    API.LLVMContextDispose(ctx)
+    mark_dispose(API.LLVMContextDispose, ctx)
 end
 
 function Context(f::Core.Function; kwargs...)

src/core/instructions.jl

@@ -18,7 +18,7 @@ end
 
 function refcheck(::Type{T}, ref::API.LLVMValueRef) where T<:Instruction
     ref==C_NULL && throw(UndefRefError())
-    if Base.JLOptions().debug_level >= 2
+    if typecheck_enabled
         T′ = identify(Instruction, ref)
         if T != T′
             error("invalid conversion of $T′ instruction reference to $T")

src/core/metadata.jl

@@ -21,7 +21,7 @@ end
 
 function refcheck(::Type{T}, ref::API.LLVMMetadataRef) where T<:Metadata
     ref==C_NULL && throw(UndefRefError())
-    if Base.JLOptions().debug_level >= 2
+    if typecheck_enabled
         T′ = identify(Metadata, ref)
         if T != T′
             error("invalid conversion of $T′ metadata reference to $T")

src/core/module.jl

@@ -9,11 +9,7 @@ export dispose, context,
 
 # forward definition of Module in src/core/value/constant.jl
 
-function Base.unsafe_convert(::Type{API.LLVMModuleRef}, mod::Module)
-    # modules can get destroyed, so be sure to check for validity
-    mod.ref == C_NULL && throw(UndefRefError())
-    mod.ref
-end
+Base.unsafe_convert(::Type{API.LLVMModuleRef}, mod::Module) = mark_use(mod).ref
 
 Base.:(==)(x::Module, y::Module) = (x.ref === y.ref)
 
@@ -23,12 +19,12 @@ Base.:(==)(x::Module, y::Module) = (x.ref === y.ref)
 end
 
 Module(name::String) =
-    Module(API.LLVMModuleCreateWithNameInContext(name, context()))
+    mark_alloc(Module(API.LLVMModuleCreateWithNameInContext(name, context())))
 
-Module(mod::Module) = Module(API.LLVMCloneModule(mod))
+Module(mod::Module) = mark_alloc(Module(API.LLVMCloneModule(mod)))
 Base.copy(mod::Module) = Module(mod)
 
-dispose(mod::Module) = API.LLVMDisposeModule(mod)
+dispose(mod::Module) = mark_dispose(API.LLVMDisposeModule, mod)
 
 function Module(f::Core.Function, args...; kwargs...)
     mod = Module(args...; kwargs...)

src/core/type.jl

@@ -23,7 +23,7 @@ end
 
 function refcheck(::Type{T}, ref::API.LLVMTypeRef) where T<:LLVMType
     ref==C_NULL && throw(UndefRefError())
-    if Base.JLOptions().debug_level >= 2
+    if typecheck_enabled
         T′ = identify(LLVMType, ref)
         if T != T′
             error("invalid conversion of $T′ type reference to $T")

src/core/value.jl

@@ -21,7 +21,7 @@ end
 
 function refcheck(::Type{T}, ref::API.LLVMValueRef) where T<:Value
     ref==C_NULL && throw(UndefRefError())
-    if Base.JLOptions().debug_level >= 2
+    if typecheck_enabled
         T′ = identify(Value, ref)
         if T != T′
             error("invalid conversion of $T′ value reference to $T")

src/core/value/constant.jl

@@ -12,7 +12,7 @@ abstract type Constant <: User end
 unsafe_destroy!(constant::Constant) = API.LLVMDestroyConstant(constant)
 
 # forward declarations
-@checked mutable struct Module
+@checked struct Module
     ref::API.LLVMModuleRef
 end
 abstract type Instruction <: User end

src/datalayout.jl

@@ -8,11 +8,11 @@ export DataLayout, dispose,
 
 # forward definition of DataLayout in src/module.jl
 
-Base.unsafe_convert(::Type{API.LLVMTargetDataRef}, dl::DataLayout) = dl.ref
+Base.unsafe_convert(::Type{API.LLVMTargetDataRef}, dl::DataLayout) = mark_use(dl).ref
 
-DataLayout(rep::String) = DataLayout(API.LLVMCreateTargetData(rep))
+DataLayout(rep::String) = mark_alloc(DataLayout(API.LLVMCreateTargetData(rep)))
 
-DataLayout(tm::TargetMachine) = DataLayout(API.LLVMCreateTargetDataLayout(tm))
+DataLayout(tm::TargetMachine) = mark_alloc(DataLayout(API.LLVMCreateTargetDataLayout(tm)))
 
 function DataLayout(f::Core.Function, args...; kwargs...)
     data = DataLayout(args...; kwargs...)
@@ -23,7 +23,7 @@ function DataLayout(f::Core.Function, args...; kwargs...)
     end
 end
 
-dispose(data::DataLayout) = API.LLVMDisposeTargetData(data)
+dispose(data::DataLayout) = mark_dispose(API.LLVMDisposeTargetData, data)
 
 Base.string(data::DataLayout) =
     unsafe_message(API.LLVMCopyStringRepOfTargetData(data))

src/debug.jl

@@ -0,0 +1,111 @@
+## typecheck: ensuring that the types of objects is as expected
+
+const typecheck_enabled = parse(Bool, @load_preference("typecheck", "false"))
+
+
+## memcheck: keeping track when objects are valid
+
+const memcheck_enabled = parse(Bool, @load_preference("memcheck", "false"))
+
+const tracked_objects = Dict{Any,Any}()
+
+function mark_alloc(obj::Any)
+    @static if memcheck_enabled
+        io = Core.stdout
+        new_alloc_bt = backtrace()[2:end]
+
+        if haskey(tracked_objects, obj)
+            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.")
+                print("\nThe original allocation was at:")
+                Base.show_backtrace(io, old_alloc_bt)
+                print("\nThe new allocation is at:")
+                Base.show_backtrace(io, new_alloc_bt)
+                println(io)
+            end
+        end
+
+        tracked_objects[obj] = (new_alloc_bt, nothing)
+    end
+    return obj
+end
+
+function mark_use(obj::Any)
+    @static if memcheck_enabled
+        io = Core.stdout
+
+        if !haskey(tracked_objects, obj)
+            # we have to ignore unknown objects, as they may originate externally.
+            # for example, a Julia-created Type we call `context` on.
+            return obj
+        end
+
+        alloc_bt, dispose_bt = tracked_objects[obj]
+        if dispose_bt !== nothing
+            print("\nWARNING: An instance of $(typeof(obj)) is being used after it was disposed.")
+            print("\nThe object was allocated at:")
+            Base.show_backtrace(io, alloc_bt)
+            print("\nThe object was disposed at:")
+            Base.show_backtrace(io, dispose_bt)
+            print("\nThe object is being used at:")
+            Base.show_backtrace(io, backtrace()[2:end])
+            println(io)
+        end
+    end
+    return obj
+end
+
+function mark_dispose(obj)
+    @static if memcheck_enabled
+        io = Core.stdout
+        new_dispose_bt = backtrace()[2:end]
+
+        if !haskey(tracked_objects, obj)
+            print(io, "\nWARNING: An unknown instance of $(typeof(obj)) is being disposed of.")
+            Base.show_backtrace(io, new_dispose_bt)
+            return
+        end
+
+        alloc_bt, old_dispose_bt = tracked_objects[obj]
+        if old_dispose_bt !== nothing
+            print("\nWARNING: An instance of $(typeof(obj)) is being disposed twice.")
+            print("\nThe object was allocated at:")
+            Base.show_backtrace(io, alloc_bt)
+            print("\nThe object was already disposed at:")
+            Base.show_backtrace(io, old_dispose_bt)
+            print("\nThe object is being disposed again at:")
+            Base.show_backtrace(io, new_dispose_bt)
+            println(io)
+        end
+
+        tracked_objects[obj] = (alloc_bt, new_dispose_bt)
+    end
+    return
+end
+
+# helper for single-line disposal without a use-after-free warning
+function mark_dispose(f, obj)
+    ret = f(obj)
+    mark_dispose(obj)
+    return ret
+end
+
+function report_leaks(code=0)
+    # if we errored, we can't trust the memory state
+    if code != 0
+        return
+    end
+
+    @static if memcheck_enabled
+        io = Core.stdout
+        for (obj, (alloc_bt, dispose_bt)) in tracked_objects
+            if dispose_bt === nothing
+                print(io, "\nWARNING: An instance of $(typeof(obj)) was not properly disposed of.")
+                print("\nThe object was allocated at:")
+                Base.show_backtrace(io, alloc_bt)
+                println(io)
+            end
+        end
+    end
+end