Merge pull request #304 from maleadt/tb/metadata

Add support for replacing values in metadata.

Author: maleadt

Manifest.toml

@@ -34,9 +34,9 @@ version = "1.4.1"
 
 [[LLVMExtra_jll]]
 deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
-git-tree-sha1 = "00d23b26d194507028b9a1c2728a691ab9914262"
+git-tree-sha1 = "771bfe376249626d3ca12bcd58ba243d3f961576"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.15+0"
+version = "0.0.16+0"
 
 [[LazyArtifacts]]
 deps = ["Artifacts", "Pkg"]

Project.toml

@@ -11,5 +11,5 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [compat]
 CEnum = "0.2, 0.3, 0.4"
-LLVMExtra_jll = "=0.0.15"
+LLVMExtra_jll = "=0.0.16"
 julia = "1.6"

deps/LLVMExtra/include/LLVMExtra.h

@@ -154,6 +154,8 @@ LLVMValueRef LLVMBuildCallWithOpBundle(LLVMBuilderRef B, LLVMValueRef Fn,
                                        LLVMOperandBundleDefRef *Bundles, unsigned NumBundles,
                                        const char *Name);
 LLVMValueRef LLVMMetadataAsValue2(LLVMContextRef C, LLVMMetadataRef Metadata);
+void LLVMReplaceAllMetadataUsesWith(LLVMValueRef Old, LLVMValueRef New);
+void LLVMReplaceMDNodeOperandWith(LLVMMetadataRef MD, unsigned I, LLVMMetadataRef New);
 
 LLVM_C_EXTERN_C_END
 #endif

deps/LLVMExtra/lib/llvm-api.cpp

@@ -526,3 +526,11 @@ LLVMValueRef LLVMMetadataAsValue2(LLVMContextRef C, LLVMMetadataRef Metadata) {
   else
     return wrap(MetadataAsValue::get(*unwrap(C), MD));
 }
+
+void LLVMReplaceAllMetadataUsesWith(LLVMValueRef Old, LLVMValueRef New) {
+  ValueAsMetadata::handleRAUW(unwrap<Value>(Old), unwrap<Value>(New));
+}
+
+void LLVMReplaceMDNodeOperandWith(LLVMMetadataRef MD, unsigned I, LLVMMetadataRef New) {
+    unwrap<MDNode>(MD)->replaceOperandWith(I, unwrap(New));
+}

lib/libLLVM_extra.jl

@@ -398,3 +398,11 @@ end
 function LLVMMetadataAsValue2(C, MD)
     ccall((:LLVMMetadataAsValue2, libLLVMExtra), LLVMValueRef, (LLVMContextRef, LLVMMetadataRef), C, MD)
 end
+
+function LLVMReplaceAllMetadataUsesWith(OldVal, NewVal)
+    ccall((:LLVMReplaceAllMetadataUsesWith, libLLVMExtra), Cvoid, (LLVMValueRef, LLVMValueRef), OldVal, NewVal)
+end
+
+function LLVMReplaceMDNodeOperandWith(MD, I, New)
+    ccall((:LLVMReplaceMDNodeOperandWith, libLLVMExtra), Cvoid, (LLVMMetadataRef, Cuint, LLVMMetadataRef), MD, I, New)
+end

src/core/metadata.jl

@@ -111,6 +111,11 @@ function operands(md::MDNode)
     return [Metadata(op) for op in ops]
 end
 
+# TODO: setindex?
+function replace_operand(md::MDNode, i, new::Metadata)
+    API.LLVMReplaceMDNodeOperandWith(md, i-1, new)
+end
+
 
 ## tuples
 

src/core/value.jl

@@ -40,7 +40,7 @@ end
 
 ## general APIs
 
-export llvmtype, llvmeltype, name, name!, replace_uses!, isconstant, isundef, ispoison, context
+export llvmtype, llvmeltype, name, name!, replace_uses!, replace_metadata_uses!, isconstant, isundef, ispoison, context
 
 llvmtype(val::Value) = LLVMType(API.LLVMTypeOf(val))
 llvmeltype(val::Value) = eltype(llvmtype(val))
@@ -55,6 +55,39 @@ end
 
 replace_uses!(old::Value, new::Value) = API.LLVMReplaceAllUsesWith(old, new)
 
+function replace_metadata_uses!(old::Value, new::Value)
+    if llvmtype(old) == llvmtype(new)
+        API.LLVMReplaceAllMetadataUsesWith(old, new)
+    else
+        # NOTE: LLVM does not support replacing values of different types, either using
+        #       regular RAUW or only on metadata. The latter should probably be supported.
+        #       Instead, we replace by a bitcast to the old type.
+        compat_new = const_bitcast(new, llvmtype(old))
+        replace_metadata_uses!(old, compat_new)
+
+        # the above is often invalid, e.g. for module-level metadata identifying functions.
+        # so we peek into such metadata and try to get rid of the bitcast. see also
+        # https://discourse.llvm.org/t/replacing-module-metadata-uses-of-function/62431/4
+        mod = LLVM.parent(new)
+        while !isa(mod, LLVM.Module)
+            mod = LLVM.parent(new)
+        end
+        ctx = context(mod)
+        function recurse(md)
+            for (i, op) in enumerate(operands(md))
+                if op isa ValueAsMetadata && Value(op; ctx) == compat_new
+                    LLVM.replace_operand(md, i, Metadata(new))
+                elseif isa(op, MDTuple)
+                    recurse(op)
+                end
+            end
+        end
+        for (key, md) in metadata(mod)
+            recurse(md)
+        end
+    end
+end
+
 isconstant(val::Value) = convert(Core.Bool, API.LLVMIsConstant(val))
 
 isundef(val::Value) = convert(Core.Bool, API.LLVMIsUndef(val))

test/core.jl

@@ -815,6 +815,34 @@ Context() do ctx
     @test convert(ConstantInt, val) == int
 end
 
+Context() do ctx
+    mod = LLVM.Module("SomeModule"; ctx)
+    ft = LLVM.FunctionType(LLVM.VoidType(ctx))
+    f1 = LLVM.Function(mod, "f1", ft)
+
+    push!(metadata(mod)["function"], MDNode([f1]; ctx))
+    @test Value(operands(operands(metadata(mod)["function"])[1])[1]; ctx) == f1
+
+    f2 = LLVM.Function(mod, "f2", ft)
+    replace_metadata_uses!(f1, f2)
+    @test Value(operands(operands(metadata(mod)["function"])[1])[1]; ctx) == f2
+end
+
+# different type; requires a hack
+Context() do ctx
+    mod = LLVM.Module("SomeModule"; ctx)
+    ft1 = LLVM.FunctionType(LLVM.VoidType(ctx))
+    f1 = LLVM.Function(mod, "f1", ft1)
+
+    push!(metadata(mod)["function"], MDNode([f1]; ctx))
+    @test Value(operands(operands(metadata(mod)["function"])[1])[1]; ctx) == f1
+
+    ft2 = LLVM.FunctionType(LLVM.Int32Type(ctx))
+    f2 = LLVM.Function(mod, "f2", ft2)
+    replace_metadata_uses!(f1, f2)
+    @test Value(operands(operands(metadata(mod)["function"])[1])[1]; ctx) == f2
+end
+
 Context() do ctx
     str = MDString("foo"; ctx)
     node = MDNode([str]; ctx)