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

…pure field set

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

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)
+