#
# fast types
#

"""
    $(@__MODULE__).isfasttype(::Type{T}) where T -> Bool

Return `true` if elements of type `T` permit fast (for instance, vectorized) operations.

By default, this holds for all primitive types except `String` and `Symbol`
and for all `struct` and `mutable struct` types whose fields have bits type
and recursively satisfy `isfasttype`.
"""
@assume_effects :total function isfasttype(::Type{T}) where T
# all non-struct types should be covered by the methods below
    isstructtype(T) && all(isfastfield, fieldtypes(T))
end

isfastfield(::Type{T}) where T = isbitstype(T) && isfasttype(T)

isfasttype(::Type{<:HWType}) = true
isfasttype(::Type{<:Union{Symbol,String}}) = false

#
# MapStyle definitions
#

using Base.FastMath: abs2_fast, abs_fast, add_fast, cmp_fast, conj_fast,
    eq_fast, ge_fast, gt_fast, inv_fast, isfinite_fast, isinf_fast,
    isnan_fast, issubnormal_fast, le_fast, lt_fast, max_fast, min_fast,
    minmax_fast, mul_fast, ne_fast, sign_fast, sqrt_fast, sub_fast

"""
    $(@__MODULE__).MapStyle

    MapStyle(f, types::Type...) -> MapStyle


A `MapStyle` determines how `$(@__MODULE__)` evaluates certain functions like
`map` or `findfirst` that take a function `f` as argument. The available subtypes
of `MapStyle` are as follows, from the least to the most efficient:

- `LazyStyle`: the function `f` is only evaluated when strictly necessary,
  and it is not assumed to be defined for default values. This is the default
  `MapStyle` for unknown functions.

- `EagerStyle`: `f` may be evaluated more often than strictly necessary,
  and it is assumed to be defined for default values. However, it need not
  map default values to default values.

- `RigidStyle`: `f` may be evaluated more often than strictly necessary. It is
  assumed to be defined for default values and to return a default value if all
  arguments are default values.

- `StrictStyle`: `f` may be evaluated more often than strictly necessary. It is
  assumed to be defined for default values and to return a default value if at
  least one argument is a default value. (This is the same as `RigidStyle` for
  functions with a single argument.)

The `MapStyle` is predefined for many functions from `Base` as well as for operations
that produce new functions out of old. Unnamed functions are not recognized. However,
several functions from `$(@__MODULE__)` allow to specify a `MapStyle` as keyword argument.
In addition, users can define a `MapStyle` for their own types.

See also [`$(@__MODULE__).default`](@ref), [`$(@__MODULE__).isfasttype`](@ref).

# Examples
```jldoctest
julia> using $(@__MODULE__): MapStyle

julia> MapStyle(iszero, Int)   # not RigidStyle: iszero(0) is true, not false
$(@__MODULE__).EagerStyle()

julia> MapStyle(+, Int, Int)
$(@__MODULE__).RigidStyle()

julia> MapStyle(*, Int, Float64)   # not StrictStyle: 0 * Inf is NaN, not 0.0
$(@__MODULE__).RigidStyle()

julia> MapStyle(*, Int, Int)
$(@__MODULE__).StrictStyle()

julia> MapStyle(-, Int)
$(@__MODULE__).StrictStyle()

julia> MapStyle(x -> -x, Int)   # not StrictStyle: anonymous function not recognized
$(@__MODULE__).LazyStyle()

julia> MapStyle(-, Int128)   # Int128 is not a fast type, so better be lazy
$(@__MODULE__).LazyStyle()

julia> MapStyle(isfinite∘inv, Float64)   # function composition is recognized
$(@__MODULE__).EagerStyle()

julia> MapStyle(!isodd, Int)  # function composition again
$(@__MODULE__).EagerStyle()

julia> MapStyle(!iszero, Int)   # separately defined to override EagerStyle
$(@__MODULE__).StrictStyle()

julia> MapStyle(>=(1), Int)   # >=(1) is Base.Fix2(>=, 1), which is recognized
$(@__MODULE__).EagerStyle()
```
"""
abstract type MapStyle end,
struct LazyStyle <: MapStyle end,
struct EagerStyle <: MapStyle end,
struct RigidStyle <: MapStyle end,
struct StrictStyle <: MapStyle end

iffasttypes(style::MapStyle) = style

function iffasttypes(style::MapStyle, ::Type{T}, types::Type...) where T
    isfasttype(T) ? iffasttypes(style, types...) : LazyStyle()
end

MapStyle(::Any, ::Type...) = LazyStyle()
# MapStyle(f, args...) = MapStyle(f, map(typeof, args)...)

MapStyle(::Union{typeof.(
        (-, identity, signbit, isodd, isone, isinf, isinf_fast, isnan, isnan_fast,
            issubnormal, issubnormal_fast, zero, round, floor, ceil, trunc,
            abs, abs_fast, sign, sign_fast, sqrt, sqrt_fast, conj, conj_fast,
            length, first, last, capacity)
    )...}, ::Type{T}) where T = iffasttypes(StrictStyle(), T)
MapStyle(::Union{typeof.(
        (&,)
    )...}, types::Type...) = iffasttypes(StrictStyle(), types...)
MapStyle(::Fix2{<:Union{typeof.(
        (<<, >>, >>>)
    )...}, S}, ::Type{T}) where {S <: Integer, T} = iffasttypes(StrictStyle(), S, T)

MapStyle(::Union{typeof.(
        (!==, !=, <, >, ne_fast, lt_fast, gt_fast, cmp_fast, -, abs2, abs2_fast)
    )...}, ::Type{T1}, ::Type{T2}) where {T1,T2} = iffasttypes(RigidStyle(), T1, T2)
MapStyle(::Union{typeof.(
        (|, xor, +, add_fast, min, min_fast, max, max_fast, minmax, minmax_fast)
    )...}, types::Type...) = iffasttypes(RigidStyle(), types...)
MapStyle(::typeof(ifelse), ::Type{Bool}, ::Type{T1}, ::Type{T2}) where {T1,T2} = iffasttypes(RigidStyle(), T1, T2)

MapStyle(::Union{typeof.(
        (!, ~, iseven, iszero, isfinite, isfinite_fast, one, inv, inv_fast)
    )...}, ::Type{T}) where T = iffasttypes(EagerStyle(), T)
MapStyle(::Union{typeof.(
        # note: 1/0 = Inf
        (===, ==, <=, >=, isequal, eq_fast, le_fast, ge_fast, /)
    )...}, ::Type{T1}, ::Type{T2}) where {T1,T2} = iffasttypes(EagerStyle(), T1, T2)
MapStyle(::Union{typeof.(
        (nand, nor)
    )...}, types::Type...) = iffasttypes(EagerStyle(), types...)

# definitions depending on specific types

MapStyle(::Fix1{typeof(<), U}, ::Type{T}) where {T, U <: Unsigned} = iffasttypes(StrictStyle(), T, U)
MapStyle(::Fix2{typeof(>), U}, ::Type{T}) where {T, U <: Unsigned} = iffasttypes(StrictStyle(), T, U)

MapStyle(::typeof(rem), ::Type{T}, ::Type{Type{U}}) where {T,U} = iffasttypes(StrictStyle(), T, U)
MapStyle(::Fix2{typeof(rem),U}, ::Type{T}) where {T,U} = iffasttypes(StrictStyle(), T, U)

MapStyle(::typeof(min), types::Type{<:Unsigned}...) = iffasttypes(StrictStyle(), types...)

hasfloats(::Type) = false
hasfloats(::Type{<:AbstractFloat}) = true
hasfloats(::Type{<:Complex{T}}) where T = hasfloats(T)
@generated hasfloats(::Type{T}) where T <: Tuple = any(hasfloats, fieldtypes(T))
hasfloats(::Type{<:AbstractArray{T}}) where T = hasfloats(T)
hasfloats(::Type{<:Ref{T}}) where T = hasfloats(T)

# -(0.0) === -0.0, not 0.0
MapStyle(::typeof(-), ::Type{T}) where T = iffasttypes(hasfloats(T) ? EagerStyle() : StrictStyle(), T)

MapStyle(::Union{typeof.(
        (cmp, isless)
    )...}, ::Type{T1}, ::Type{T2}) where {T1,T2} = hasfloats(Tuple{T1,T2}) ? LazyStyle() : iffasttypes(RigidStyle(), T1, T2)

hasfixedlength(::Type) = false
hasfixedlength(::Type{<:Number}) = true

# (-1) * 0.0 === -0.0, not 0.0
# also, 0 * [1] is not a vector of length 0
@assume_effects :total function MapStyle(::Union{typeof.(
        (*, mul_fast, dot, dot_fast)
    )...}, types::Type...)
    style = if !all(hasfixedlength, types)
        EagerStyle()
    elseif hasfloats(Tuple{types...})
        RigidStyle()
    else
        StrictStyle()
    end
    iffasttypes(style, types...)
end

function MapStyle(::typeof(Base.literal_pow), ::Type{typeof(^)}, ::Type{T}, ::Type{Val{N}}) where {T,N}
    if 1 <= N <= 3
        iffasttypes(StrictStyle(), T)
    elseif -2 <= N <= 0
        iffasttypes(EagerStyle(), T)
    else
        LazyStyle()
    end
end

MapStyle(::typeof(in), ::Type{S}, ::Type{T}) where {S, T <: AbstractSet} = iffasttypes(StrictStyle(), T, S)

MapStyle(::typeof(intersect), ::Type{T}, types::Type...) where T <: AbstractSet = iffasttypes(StrictStyle(), T, types...)

MapStyle(::Union{typeof.(
        (union, setdiff, symdiff)
    )...}, ::Type{T}, types::Type...) where T <: AbstractSet = iffasttypes(RigidStyle(), T, types...)

MapStyle(::Union{typeof.(
        (<, >, ⊊, ⊋)
    )...}, ::Type{T1}, ::Type{T2}) where {T1 <: AbstractSet, T2 <: AbstractSet} = iffasttypes(RigidStyle(), T1, T2)

MapStyle(::Union{
        Fix2{typeof(setdiff), U}, Fix1{typeof(<), U}, Fix2{typeof(>), U}
    }, ::Type{T}) where {T <: AbstractSet, U <: AbstractSet} = iffasttypes(StrictStyle(), T, U)

MapStyle(::Union{typeof.(
        (issubset, ⊇)
    )...}, ::Type{T1}, ::Type{T2}) where {T1 <: AbstractSet, T2 <: AbstractSet} = iffasttypes(EagerStyle(), T1, T2)

# definitions for constructors

MapStyle(::Type{Ref}, ::Type{T}) where T = iffasttypes(StrictStyle(), S, T)
MapStyle(::Type{Ref{S}}, ::Type{T}) where {S,T} = iffasttypes(StrictStyle(), S, T)
MapStyle(::Type{<:AbstractOneTo}, ::Type{T}) where T = iffasttypes(StrictStyle(), T)
MapStyle(::Type{<:AbstractOneTo{S}}, ::Type{T}) where {S,T} = iffasttypes(StrictStyle(), S, T)

# definitions for constructors of new functions

MapStyle(f::Returns{T}, types::Type...) where T = iffasttypes(EagerStyle(), T)

function MapStyle(f::ComposedFunction, types::Type...)
    istyle = MapStyle(f.inner, types...)
    istyle isa LazyStyle && return LazyStyle()
    T = Core.Compiler.return_type(f.inner, Tuple{types...})
    isconcretetype(T) || return LazyStyle()
    ostyle = MapStyle(f.outer, T)
    if ostyle isa Union{LazyStyle, EagerStyle}
        ostyle
    elseif istyle isa Union{EagerStyle, RigidStyle}
        istyle
    else
        StrictStyle()
    end
end

MapStyle(::ComposedFunction{typeof(!), <:Union{typeof.(
        (~, iseven, iszero, isfinite, isfinite_fast)
    )...}}, ::Type{T}) where T = iffasttypes(StrictStyle(), T)

MapStyle(f::Base.Splat, ::Type{T}) where T <: Tuple = MapStyle(f.f, fieldtypes(T)...)

# TODO: use StrictStyle
if VERSION > v"1.12-alpha"
    function MapStyle(g::Base.Fix{N,F,T}, types::Type...) where {N,F,T}
        style = MapStyle(g.f, types[1:N-1]..., T, types[N:end]...)
        style isa RigidStyle ? EagerStyle() : style
    end
else
    function MapStyle(g::Fix1{F,T}, types::Type...) where {F,T}
        style = MapStyle(g.f, T, types...)
        style isa RigidStyle ? EagerStyle() : style
    end
    function MapStyle(g::Fix2{F,T}, type1::Type, types::Type...) where {F,T}
        style = MapStyle(g.f, type1, T, types...)
        style isa RigidStyle ? EagerStyle() : style
    end
end

# Inbounds

"""
    $(@__MODULE__).Inbounds(f)

Returns a callable object for calling `f` in an `@inbounds` context.
This is useful for `map` and broadcasting, for instance.

See also `Base.@inbounds`.

# Example
```jldoctest
julia> using $(@__MODULE__): Inbounds

julia> @inline f(x) = (@boundscheck @assert x < 3; -x);

julia> v = FixedVector{4}(1:4);

julia> f.(v)
ERROR: AssertionError: x < 3
[...]

julia> Inbounds(f).(v)
4-element FixedVector{4, Int64}:
 -1
 -2
 -3
 -4
```
"""
struct Inbounds{F}
    f::F
end

Inbounds(::Type{T}) where T = Inbounds{Type{T}}(T)

(inb::Inbounds)(xs...; kw...) = @inbounds inb.f(xs...; kw...)

MapStyle(inb::Inbounds, types::Type...) = MapStyle(inb.f, types...)
