Improve inference for unique with abstract eltypes

#20317 improved inference of unique, but problematic cases still arise
for containers with known but abstract eltypes. Here, we short-circuit
the `typejoin` when the return type is determined by the element type
of the input container.

For `unique(f, itr)`, this commit also allows the caller to supply
`seen::Set` to circumvent the inference challenges.

Author: timholy

NEWS.md

@@ -48,6 +48,8 @@ Standard library changes
 * `view`, `@view`, and `@views` now work on `AbstractString`s, returning a `SubString` when appropriate ([#35879]).
 * All `AbstractUnitRange{<:Integer}`s now work with `SubString`, `view`, `@view` and `@views` on strings ([#35879]).
 * `sum`, `prod`, `maximum`, and `minimum` now support `init` keyword argument ([#36188], [#35839]).
+* `unique(f, itr; seen=Set{T}())` now allows you to declare the container type used for
+  keeping track of values returned by `f` on elements of `itr` ([#36280]).
 
 #### LinearAlgebra
 * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]).

base/set.jl

@@ -120,19 +120,25 @@ julia> unique(Real[1, 1.0, 2])
 ```
 """
 function unique(itr)
+    if isa(IteratorEltype(itr), HasEltype)
+        T = eltype(itr)
+        out = Vector{T}()
+        seen = Set{T}()
+        for x in itr
+            if !in(x, seen)
+                push!(seen, x)
+                push!(out, x)
+            end
+        end
+        return out
+    end
     T = @default_eltype(itr)
-    out = Vector{T}()
-    seen = Set{T}()
     y = iterate(itr)
-    y === nothing && return out
+    y === nothing && return T[]
     x, i = y
-    if !isconcretetype(T) && IteratorEltype(itr) == EltypeUnknown()
-        S = typeof(x)
-        return _unique_from(itr, S[x], Set{S}((x,)), i)
-    end
-    push!(seen, x)
-    push!(out, x)
-    return unique_from(itr, out, seen, i)
+    S = typeof(x)
+    R = isconcretetype(T) ? T : S
+    return _unique_from(itr, R[x], Set{R}((x,)), i)
 end
 
 _unique_from(itr, out, seen, i) = unique_from(itr, out, seen, i)
@@ -175,8 +181,18 @@ julia> unique(x -> x^2, [1, -1, 3, -3, 4])
  4
 ```
 """
-function unique(f, C)
+function unique(f, C; seen::Union{Nothing,Set}=nothing)
     out = Vector{eltype(C)}()
+    if seen !== nothing
+        for x in C
+            y = f(x)
+            if y ∉ seen
+                push!(out, x)
+                push!(seen, y)
+            end
+        end
+        return out
+    end
 
     s = iterate(C)
     if s === nothing
@@ -241,15 +257,17 @@ julia> unique!(iseven, [2, 3, 5, 7, 9])
  3
 ```
 """
-function unique!(f, A::AbstractVector)
+function unique!(f, A::AbstractVector; seen::Union{Nothing,Set}=nothing)
     if length(A) <= 1
         return A
     end
 
     i = firstindex(A)
     x = @inbounds A[i]
     y = f(x)
-    seen = Set{typeof(y)}()
+    if seen === nothing
+        seen = Set{typeof(y)}()
+    end
     push!(seen, y)
     return _unique!(f, A, seen, i, i+1)
 end

test/sets.jl

@@ -384,27 +384,30 @@ end
 end
 
 @testset "unique" begin
-    u = unique([1, 1, 2])
+    u = @inferred(unique([1, 1, 2]))
     @test in(1, u)
     @test in(2, u)
     @test length(u) == 2
-    @test unique(iseven, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 8]
-    @test unique(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 1, 9]
+    @test @inferred(unique(iseven, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6])) == [5, 8]
+    @test @inferred(unique(x->x^2, Integer[3, -4, 5, 4])) == Integer[3, -4, 5]
+    @test @inferred(unique(iseven, Integer[3, -4, 5, 4]; seen=Set{Bool}())) == Integer[3, -4]
+    @test @inferred(unique(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6])) == [5, 1, 9]
 end
 
 @testset "issue 20105" begin
     @test @inferred(unique(x for x in 1:1)) == [1]
     @test unique(x for x in Any[1, 1.0])::Vector{Real} == [1]
     @test unique(x for x in Real[1, 1.0])::Vector{Real} == [1]
-    @test unique(Integer[1, 1, 2])::Vector{Integer} == [1, 2]
+    @test @inferred(unique(Integer[1, 1, 2]))::Vector{Integer} == [1, 2]
+    @test unique(x for x in []) isa Vector{Any}
 end
 
 @testset "unique!" begin
     u = [1,1,3,2,1]
-    unique!(u)
+    @inferred(unique!(u))
     @test u == [1,3,2]
-    @test unique!([]) == []
-    @test unique!(Float64[]) == Float64[]
+    @test @inferred(unique!([])) == []
+    @test @inferred(unique!(Float64[])) == Float64[]
     u = [1,2,2,3,5,5]
     @test unique!(u) === u
     @test u == [1,2,3,5]
@@ -434,8 +437,9 @@ end
     u = [1,2,5,1,3,2]
     @test unique!(x -> x ^ 2, [1, -1, 3, -3, 5, -5]) == [1, 3, 5]
     @test unique!(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 1, 9]
-    @test unique!(iseven, [2, 3, 5, 7, 9]) == [2, 3]
-    @test unique!(x -> x % 2 == 0 ? :even : :odd, [1, 2, 3, 4, 2, 2, 1]) == [1, 2]
+    @test @inferred(unique!(iseven, [2, 3, 5, 7, 9])) == [2, 3]
+    @test @inferred(unique!(x -> x % 2 == 0 ? :even : :odd, [1, 2, 3, 4, 2, 2, 1])) == [1, 2]
+    @test @inferred(unique!(x -> x % 2 == 0 ? :even : "odd", [1, 2, 3, 4, 2, 2, 1]; seen=Set{Union{Symbol,String}}())) == [1, 2]
 end
 
 @testset "allunique" begin