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.

NEWS.md

@@ -48,6 +48,7 @@ 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 storing observed values of `f` applied to element of `itr`.
 
 #### LinearAlgebra
 * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]).

base/set.jl

@@ -121,18 +121,22 @@ julia> unique(Real[1, 1.0, 2])
 """
 function unique(itr)
  T = @default_eltype(itr)
- out = Vector{T}()
- seen = Set{T}()
+ if isconcretetype(T) || IteratorEltype(itr) == HasEltype()
+ 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
  y = iterate(itr)
  y === nothing && return out
  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)
+ return _unique_from(itr, S[x], Set{S}((x,)), i)
 end
 
 _unique_from(itr, out, seen, i) = unique_from(itr, out, seen, i)
@@ -175,8 +179,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 +255,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

@@ -389,6 +389,8 @@ end
  @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(x->x^2, Integer[3, -4, 5, 4]) == Integer[3, -4, 5]
+ @test unique(iseven, Integer[3, -4, 5, 4]; seen=Set{Bool}()) == Integer[3, -4]
  @test unique(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 1, 9]
 end
 
@@ -436,6 +438,7 @@ end
  @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 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