cleanup implementation

Author: vchuravy

src/overdub.jl

@@ -319,20 +319,59 @@ function overdub_pass!(reflection::Reflection,
     end
 
     #=== mark all `llvmcall`s as nooverdub ===#
-    # TODO: this only works for: `Intrinsics.llvmcall` and not `Core.Intrinsics.llvmcall`
-    # since there is a getproperty call in the way.
+    function unravel_intrinsics(x)
+        stmt = Base.Meta.isexpr(x, :(=)) ? x.args[2] : x
+        if Base.Meta.isexpr(stmt, :call)
+            applycall = is_ir_element(stmt.args[1], GlobalRef(Core, :_apply), overdubbed_code)
+            f = applycall ? stmt.args[2] : stmt.args[1]
+            f = ir_element(f, overdubbed_code)
+            if f isa Expr && Base.Meta.isexpr(f, :call) &&
+                 is_ir_element(f.args[1], GlobalRef(Base, :getproperty), overdubbed_code)
+
+                # resolve getproperty here
+                mod = ir_element(f.args[2], overdubbed_code)
+                if mod isa GlobalRef
+                    mod = resolve_early(mod) # returns nothing if fails
+                    mod === nothing && return nothing
+                end
+                fname = ir_element(f.args[3], overdubbed_code)
+                if fname isa QuoteNode
+                    fname = fname.value
+                end
+                f = GlobalRef(mod, fname)
+            end
+            if f isa GlobalRef
+                f = resolve_early(f)
+            end
+            return f 
+        end
+        return nothing
+    end
+
     # TODO: Need to fix for `istaggingenabled == true`
+    # TODO: add user-facing flag to do this for all intrinsics
     if !iskwfunc && !istaggingenabled
         insert_statements!(overdubbed_code, overdubbed_codelocs,
                             (x, i) -> begin
-                                if Base.Meta.isexpr(x, :call) &&
-                                  is_ir_element(x.args[1], GlobalRef(Core.Intrinsics, :llvmcall), overdubbed_code)
+                                intrinsic = unravel_intrinsics(x)
+                                if intrinsic === nothing
+                                    return nothing
+                                end
+                                if intrinsic === Core.Intrinsics.llvmcall
                                     return 1
                                 end
-                                return nothing
                             end,
                             (x, i) -> begin
-                                [Expr(:call, Expr(:nooverdub, GlobalRef(Core.Intrinsics, :llvmcall)), x.args[2:end]...)]
+                                stmt = Base.Meta.isexpr(x, :(=)) ? x.args[2] : x
+                                applycall = is_ir_element(stmt.args[1], GlobalRef(Core, :_apply), overdubbed_code)
+                                intrinsic = unravel_intrinsics(x)
+                                if applycall
+                                    # using stmt.args[2] instead of `intrinsic` leads to a bug
+                                    stmt.args[2] = Expr(:nooverdub, intrinsic)
+                                else
+                                    stmt.args[1] = Expr(:nooverdub, intrinsic)
+                                end
+                                [x]
                             end)
     end
 

src/pass.jl

@@ -197,3 +197,27 @@ function is_ir_element(x, y, code::Vector)
     end
     return result
 end
+
+"""
+    ir_element(x, code::Vector)
+
+Follows the series of `SSAValue` that define `x`.
+
+See also: [`is_ir_element`](@ref)
+"""
+function ir_element(x, code::Vector)
+    while isa(x, Core.SSAValue)
+        x = code[x.id]
+    end
+    return x
+end
+
+function resolve_early(ref::GlobalRef)
+    mod = ref.mod
+    name = ref.name
+    if Base.isbindingresolved(mod, name) && Base.isdefined(mod, name)
+        return getfield(mod, name)
+    else
+        return nothing
+    end
+end

test/misctests.jl

@@ -686,10 +686,8 @@ Cassette.@context LLVMCallCtx
     end
 end
 
-import Core.Intrinsics
 function llvm_sin(x::Float64)
-    # Needs fix for Core.Intrinsics.llvmcall
-    Intrinsics.llvmcall(
+    Core.Intrinsics.llvmcall(
         (
             """declare double @llvm.sin.f64(double)""",
             """%2 = call double @llvm.sin.f64(double %0)