fix to allow at-pure optimization for generated thunks that have the pure field set

Co-authored-by: Valentin Churavy <[email protected]>

Author: jrevels

base/compiler/abstractinterpretation.jl

@@ -605,9 +605,17 @@ function pure_eval_call(@nospecialize(f), argtypes::Vector{Any}, @nospecialize(a
         return false
     end
     meth = meth[1]::SimpleVector
+    sig = meth[1]::DataType
+    sparams = meth[2]::SimpleVector
     method = meth[3]::Method
-    # TODO: check pure on the inferred thunk
-    if isdefined(method, :generator) || !method.pure
+
+    if isdefined(method, :generator)
+        method.generator.expand_early || return false
+        mi = specialize_method(method, sig, sparams, false)
+        isa(mi, MethodInstance) || return false
+        staged = get_staged(mi)
+        (staged isa CodeInfo && (staged::CodeInfo).pure) || return false
+    elseif !method.pure
         return false
     end
 

test/compiler/contextual.jl

@@ -108,3 +108,9 @@ f() = 2
 
 # Test that MiniCassette is at least somewhat capable by overdubbing gcd
 @test overdub(Ctx(), gcd, 10, 20) === gcd(10, 20)
+
+# Test that pure propagates for Cassette
+Base.@pure isbitstype(T) = T.isbitstype
+f31012(T) = Val(isbitstype(T))
+@test @inferred overdub(Ctx(), f31012, Int64) == Val(true)
+