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")