Type-stable codegen for specialized in(v, ::Tuple)

base/operators.jl

@@ -1332,24 +1332,43 @@ used to implement specialized methods.
 """
 in(x) = Fix2(in, x)
 
-for ItrT = (Tuple,Any)
- # define a generic method and a specialized version for `Tuple`,
- # whose method bodies are identical, while giving better effects to the later
- @eval function in(x, itr::$ItrT)
- $(ItrT === Tuple ? :(@_terminates_locally_meta) : :nothing)
- anymissing = false
- for y in itr
- v = (y == x)
- if ismissing(v)
- anymissing = true
- elseif v
- return true
- end
+function in(x, itr::Any)
+ anymissing = false
+ for y in itr
+ v = (y == x)
+ if ismissing(v)
+ anymissing = true
+ elseif v
+ return true
  end
+ end
+ return anymissing ? missing : false
+end
+
+# Specialized variant of in for Tuple, which can generate typed comparisons for each element
+# of the tuple, skipping values that are statically known to be != at compile time.
+in(x, itr::Tuple) = _in_tuple(x, itr, false)
+# This recursive function will be unrolled at compiletime, and will not generate separate
+# llvm-compiled specializations for each step of the recursion.
+function _in_tuple(x, @nospecialize(itr::Tuple), anymissing::Bool)
+ @inline
+ # Base case
+ if isempty(itr)
  return anymissing ? missing : false
  end
+ # Recursive case
+ v = (itr[1] == x)
+ if ismissing(v)
+ anymissing = true
+ elseif v
+ return true
+ end
+ return _in_tuple(x, tail(itr), anymissing)
 end
 
+# fallback to the loop implementation after some number of arguments to avoid inference blowup
+in(x, itr::Any32) = invoke(in, Tuple{Any,Any}, x, itr)
+
 const ∈ = in
 ∉(x, itr) = !∈(x, itr)
 ∉(itr) = Fix2(∉, itr)

test/operators.jl

@@ -328,9 +328,21 @@ end
  @test lt5(4) && !lt5(5)
 end
 
+@testset "in tuples" begin
+ @test ∈(5, (1,5,10,11))
+ @test ∉(0, (1,5,10,11))
+ @test ∈(5, (1,"hi","hey",5.0))
+ @test ∉(0, (1,"hi","hey",5.0))
+ @test ∈(5, (5,))
+ @test ∉(0, (5,))
+ @test ∉(5, ())
+end
+
 @testset "ni" begin
  @test ∋([1,5,10,11], 5)
  @test !∋([1,10,11], 5)
+ @test ∋((1,5,10,11), 5)
+ @test ∌((1,10,11), 5)
  @test ∋(5)([5,1])
  @test !∋(42)([0,1,100])
  @test ∌(0)(1:10)
@@ -367,7 +379,8 @@ end
  @test B46327() <= B46327()
 end
 
-@testset "concrete eval `x in itr::Tuple`" begin
+@testset "inference for `x in itr::Tuple`" begin
+ # concrete evaluation
  @test Core.Compiler.is_foldable(Base.infer_effects(in, (Int,Tuple{Int,Int,Int})))
  @test Core.Compiler.is_foldable(Base.infer_effects(in, (Char,Tuple{Char,Char,Char})))
  for i = (1,2,3)
@@ -377,10 +390,23 @@ end
  end |> only == Val{true}
  end
  end
- @test Base.return_types() do
+ @test Base.infer_return_type() do
  Val(4 in (1,2,3))
- end |> only == Val{false}
- @test Base.return_types() do
+ end == Val{false}
+ @test Base.infer_return_type() do
  Val('1' in ('1','2','3'))
- end |> only == Val{true}
+ end == Val{true}
+
+ # constant propagation
+ @test Base.infer_return_type((Int,Int)) do x, y
+ Val(1 in (x,2,y))
+ end >: Val{true}
+ @test Base.infer_return_type((Int,Int)) do x, y
+ Val(2 in (x,2,y))
+ end == Val{true}
+
+ # should use the loop implementation given large tuples to avoid inference blowup
+ let t = ntuple(x->'A', 10000);
+ @test Base.infer_return_type(in, (Char,typeof(t))) == Bool
+ end
 end