From 62b82e2f0a2496903f15b874d45673c0c5119489 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Wed, 1 Sep 2021 15:02:57 +0900
Subject: [PATCH] update to https://github.com/JuliaLang/julia/pull/41328

---
 docs/src/internals.md |  2 +-
 src/analyzer.jl       | 27 ++++++++++++++++++++++++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/docs/src/internals.md b/docs/src/internals.md
index 234eb1519..fd06923a5 100644
--- 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
index 2638920fe..80119e4f3 100644
--- 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
 # ----------------