update to JuliaLang/julia#41328

aviatesk · Sep 1, 2021 · 62b82e2 · 62b82e2
1 parent 4cb3357
commit 62b82e2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/docs/src/internals.md b/docs/src/internals.md
@@ -31,7 +31,7 @@ JET.analyze_task_parallel_code!
 ```@docs
 JET.JET_CACHE
 JET.JETCachedResult
-JET.jet_inlining_policy
+JET.inlining_policy
 ```
 
 

diff --git a/src/analyzer.jl b/src/analyzer.jl
@@ -479,14 +479,35 @@ CC.may_compress(analyzer::AbstractAnalyzer)      = false
 CC.may_discard_trees(analyzer::AbstractAnalyzer) = false
 CC.verbose_stmt_info(analyzer::AbstractAnalyzer) = false
 
+# branch on https://github.com/JuliaLang/julia/pull/41328
+@static if isdefined(CC, :is_stmt_inline)
+@doc """
+    inlining_policy(analyzer::AbstractAnalyzer, @nospecialize(src), stmt_flag::UInt8) -> source::Any
+
+Implements inlining policy for `AbstractAnalyzer`.
+Since `AbstractAnalyzer` works on `InferenceResult` whose `src` field keeps
+[`JETResult`](@ref) or [`JETCachedResult`](@ref), this implementation just bypasses
+their wrapped source to `inlining_policy(::AbstractInterpreter, ::Any, ::UInt8)`.
 """
-    jet_inlining_policy(src)
+function CC.inlining_policy(analyzer::AbstractAnalyzer, @nospecialize(src), stmt_flag::UInt8)
+    if isa(src, JETResult)
+        src = get_source(src)
+    elseif isa(src, JETCachedResult)
+        src = get_source(src)
+    end
+    return @invoke CC.inlining_policy(analyzer::AbstractInterpreter, @nospecialize(src), stmt_flag::UInt8)
+end
+else # @static if isdefined(Compiler, :is_stmt_inline)
+@doc """
+    inlining_policy(::AbstractAnalyzer) = jet_inlining_policy
+    jet_inlining_policy(@nospecialize(src)) -> source::Any
 
-Implements `Core.Compiler.inlining_policy` for `AbstractAnalyzer`.
+`jet_inlining_policy` implements `Core.Compiler.inlining_policy` for `AbstractAnalyzer`.
 Since `AbstractAnalyzer` works on `InferenceResult` whose `src` field keeps
 [`JETResult`](@ref) or [`JETCachedResult`](@ref), `jet_inlining_policy` bypasses
 their wrapped source to `Core.Compiler.default_inlining_policy`.
 """
+CC.inlining_policy(::AbstractAnalyzer) = jet_inlining_policy
 @inline function jet_inlining_policy(@nospecialize(src))
     if isa(src, JETResult)
         src = get_source(src)
@@ -495,7 +516,7 @@ their wrapped source to `Core.Compiler.default_inlining_policy`.
     end
     return default_inlining_policy(src)
 end
-CC.inlining_policy(::AbstractAnalyzer) = jet_inlining_policy
+end # @static if isdefined(Compiler, :is_stmt_inline)
 
 # AbstractAnalyzer
 # ----------------