Merge branch 'refactor-map'

src/FixedSizeArrays.jl

@@ -18,13 +18,15 @@ if VERSION <= v"0.5.0"
 end
 
 # put them here due to #JuliaLang/julia#12814
-# needs to be befor indexing and ops, but after constructors
+# needs to be before indexing and ops, but after constructors
 immutable Mat{Row, Column, T} <: FixedMatrix{Row, Column, T}
     _::NTuple{Column, NTuple{Row, T}}
 end
 function similar{FSA <: Mat, T}(::Type{FSA}, ::Type{T}, SZ::NTuple{2, Int})
     Mat{SZ[1], SZ[2], T}
 end
+similar{FSA <: Mat}(::Type{FSA}, SZ::NTuple{2,Int}) = similar(FSA, eltype(FSA), SZ)
+similar{N,M,S, T}(::Type{Mat{N,M,S}}, ::Type{T}) = Mat{N,M,T}
 
 # most common FSA types
 immutable Vec{N, T} <: FixedVector{N, T}
@@ -68,6 +70,7 @@ export MutableFixedVector
 export MutableFixedMatrix
 export Mat, Vec, Point
 export @fsa
+export construct_similar
 
 export unit
 export normalize

src/core.jl

@@ -102,3 +102,65 @@ end
 function get_tuple(f::FixedArray)
     f._
 end
+
+
+# Infrastructure for construct_similar
+"Compute promoted element type of the potentially nested Tuple type `ty`"
+function promote_type_nested(ty)
+    if ty.name.primary === Tuple
+        promote_type([promote_type_nested(t) for t in ty.parameters]...)
+    else
+        ty
+    end
+end
+
+"""
+Construct tuple expression converting inner elements of the nested Tuple type
+`ty` with name `varexpr` to the scalar `outtype`
+"""
+function convert_nested_tuple_expr(outtype, varexpr, ty)
+    if ty.name.primary === Tuple
+        Expr(:tuple, [convert_nested_tuple_expr(outtype, :($varexpr[$i]), t)
+                      for (i,t) in enumerate(ty.parameters)]...)
+    else
+        :(convert($outtype, $varexpr))
+    end
+end
+
+"""
+Compute the N-dimensional array shape of a nested Tuple if it were used as
+column-major storage for a FixedArray.
+"""
+function nested_Tuple_shape(ty)
+    if ty.name.primary !== Tuple
+        return 0
+    end
+    subshapes = [nested_Tuple_shape(t) for t in ty.parameters]
+    if isempty(subshapes)
+        return ()
+    end
+    if any(subshapes .!= subshapes[1])
+        throw(DimensionMismatch("Nested tuples must have equal length to form a FixedSizeArray"))
+    end
+    if subshapes[1] == 0
+        return (length(subshapes),) # Scalar elements
+    end
+    return (subshapes[1]..., length(subshapes))
+end
+
+"""
+    construct_similar(::Type{FSA}, elements::Tuple)
+
+Construct FixedArray as similar as possible to `FSA`, but with shape given by
+the shape of the nested tuple `elements` and with `eltype` equal to the
+promoted element type of the nested tuple `elements`.
+"""
+@generated function construct_similar{FSA <: FixedArray}(::Type{FSA}, elements::Tuple)
+    etype = promote_type_nested(elements)
+    shape = nested_Tuple_shape(elements)
+    outtype = similar(FSA, etype, shape)
+    converted_elements = convert_nested_tuple_expr(etype, :elements, elements)
+    constructor_expr(outtype, converted_elements)
+end
+
+

src/mapreduce.jl

@@ -1,3 +1,4 @@
+# reduce() versions
 @inline function reduce{FSA <: Union{FixedArray, Tuple}}(f, a::FSA)
     length(a) == 1 && return a[1]
     @inbounds begin
@@ -20,73 +21,123 @@ end
     red
 end
 
-index_expr{T <: Number}(::Type{T}, i::Int, inds::Int...) = :($(Symbol("arg$i")))
-index_expr{T <: FixedArray}(::Type{T}, i::Int, inds::Int...) = :($(Symbol("arg$i"))[$(inds...)])
-inner_expr{N}(args::NTuple{N, DataType}, inds::Int...) = :( F($(ntuple(i -> index_expr(args[i], i, inds...), N)...)) )
 
+#------------------------------------------------------------------------------
+# map() machinery
 
-# This solves the combinational explosion from FixedVectorNoTuple while staying fast.
-function constructor_expr{T <: FixedVector}(::Type{T}, tuple_expr::Expr)
+# Get an expression indexing the collection `name` of type `T` for use in map()
+index_expr{T <: Number}(::Type{T},     name, inds::Int...)    = :($name)
+index_expr{T <: FixedArray}(::Type{T}, name, inds::Int...)    = :($name[$(inds...)])
+index_expr{T <: AbstractArray}(::Type{T}, name, inds::Int...) = :($name[$(inds...)])
+
+# Get expression checking size of collection `name` against `SIZE`
+function sizecheck_expr{T <: Number}(::Type{T}, name, SIZE)
+    :nothing
+end
+function sizecheck_expr{FSA<:FixedArray}(::Type{FSA}, name, SIZE)
+    if size(FSA) == SIZE
+        :nothing
+    else
+        :(throw(DimensionMismatch(string($FSA)*" is wrong size")))
+    end
+end
+function sizecheck_expr{A<:AbstractArray}(::Type{A}, name, SIZE)
+    quote
+        # Note - should be marked with @boundscheck in 0.5
+        size($name) == $SIZE || throw(DimensionMismatch(string($A)*" is wrong size"))
+    end
+end
+
+# Type wrapper to signify that a similar FSA should be constructed as the map()
+# output type via construct_similar()
+immutable SimilarTo{FSA}; end
+
+# Get expression to construct FSA from a nested tuple of storage generated by map()
+constructor_expr{FSA <: FixedArray}(::Type{FSA}, tuple_expr::Expr) = :($FSA($tuple_expr))
+constructor_expr{FSA <: FixedVectorNoTuple}(::Type{FSA}, tuple_expr::Expr) = :($FSA($tuple_expr...))
+constructor_expr{FSA <: FixedArray}(::Type{SimilarTo{FSA}}, tuple_expr::Expr) = :(construct_similar($FSA, $tuple_expr))
+
+# Generate an unrolled nested tuple of calls mapping funcname across the input,
+# constructing an `OutFSA` to store the result.
+#
+# julia> FixedSizeArrays.unrolled_map_expr(:f, Vec{2,Bool}, (2,), (Vec{2,Int},Int), (:A,:b))
+#
+# generates, after cleaning up:
+#
+#    (FixedSizeArrays.Vec{2,Bool})(
+#        tuple(tuple(f(A[1,1],b), f(A[2,1],b)),
+#              tuple(f(A[1,2],b), f(A[2,2],b)))
+#    )
+function unrolled_map_expr(funcname, OutFSA, SIZE, argtypes, argnames)
+    sizecheck = [sizecheck_expr(T,n,SIZE) for (T,n) in zip(argtypes,argnames)]
+    tuple_expr = fill_tuples_expr(SIZE) do inds...
+        Expr(:call, funcname,
+            [index_expr(argtypes[i], argnames[i], inds...) for i=1:length(argtypes)]...
+        )
+    end
     quote
-        $(Expr(:boundscheck, false))
         $(Expr(:meta, :inline))
-        rvalue = FSA($(tuple_expr))
+        $(sizecheck...)
+        $(Expr(:boundscheck, false))
+        rvalue = $(constructor_expr(OutFSA, tuple_expr))
         $(Expr(:boundscheck,:pop))
         rvalue
     end
 end
-constructor_expr{T <: Mat}(::Type{T}, tuple_expr::Expr) = quote
-    $(Expr(:boundscheck, false))
-    $(Expr(:meta, :inline))
-    rvalue = Mat($(tuple_expr))
-    $(Expr(:boundscheck,:pop))
-    rvalue
+
+
+# map() comes in two flavours:
+#
+# 1) You can specify the output type with the first argument
+#   map(func, ::Type{OutFSA}, arg1, arg2, ...)
+
+# General N-ary version with explicit output.  (Unary and binary versions are
+# written out below since this generates better code in julia-0.4.)
+@generated function map{OutFSA<:FixedArray}(func, ::Type{OutFSA}, args...)
+    argexprs = ntuple(i->:(args[$i]), length(args))
+    unrolled_map_expr(:func, OutFSA, size(OutFSA), args, argexprs)
 end
-constructor_expr{T <: FixedVectorNoTuple}(::Type{T}, tuple_expr::Expr) = quote
-    $(Expr(:boundscheck, false))
-    $(Expr(:meta, :inline))
-    rvalue = FSA($(tuple_expr)...)
-    $(Expr(:boundscheck,:pop))
-    rvalue
+# Binary
+@generated function map{OutFSA<:FixedArray}(func, ::Type{OutFSA}, arg1, arg2)
+    unrolled_map_expr(:func, OutFSA, size(OutFSA), (arg1,arg2), (:arg1,:arg2))
 end
-@generated function map{FSA <: FixedArray}(F, arg1::FSA, arg2::FSA)
-    inner = fill_tuples_expr((inds...) -> inner_expr((arg1, arg2), inds...), size(FSA))
-    constructor_expr(FSA, inner)
+# Unary
+@generated function map{OutFSA<:FixedArray}(func, ::Type{OutFSA}, arg1)
+    unrolled_map_expr(:func, OutFSA, size(OutFSA), (arg1,), (:arg1,))
 end
-@generated function map{FSA <: FixedArray}(F, arg1::FSA, arg2::Number)
-    inner = fill_tuples_expr((inds...) -> inner_expr((arg1, arg2), inds...), size(FSA))
-    constructor_expr(FSA, inner)
+
+
+# 2) You can let map() infer the output FixedArray type from `func`
+#   map(func, fixed_arrays...)
+
+# Binary, inferred output
+@generated function map{F1<:FixedArray, F2<:FixedArray}(func, arg1::F1, arg2::F2)
+    unrolled_map_expr(:func, SimilarTo{arg1}, size(F1), (arg1,arg2), (:arg1,:arg2))
 end
-@generated function map{FSA <: FixedArray}(F, arg1::Number, arg2::FSA)
-    inner = fill_tuples_expr((inds...) -> inner_expr((arg1, arg2), inds...), size(FSA))
-    constructor_expr(FSA, inner)
+@generated function map{F<:FixedArray}(func, arg1::F, arg2::Union{Number,AbstractArray})
+    unrolled_map_expr(:func, SimilarTo{arg1}, size(F), (arg1,arg2), (:arg1,:arg2))
 end
-@generated function map{FSA <: FixedArray}(F, ::Type{FSA})
-    inner = fill_tuples_expr((inds...) -> :(F($(inds...))), size(FSA))
-    :( FSA($inner) )
+@generated function map{F<:FixedArray}(func, arg1::Union{Number,AbstractArray}, arg2::F)
+    unrolled_map_expr(:func, SimilarTo{arg2}, size(F), (arg1,arg2), (:arg1,:arg2))
 end
-
-@generated function map{FSA <: FixedArray}(F, arg1::FSA)
-    inner = fill_tuples_expr((inds...) -> :( F(arg1[$(inds...)]) ), size(FSA))
-    constructor_expr(FSA, inner)
+# Unary, inferred output
+@generated function map{F<:FixedArray}(func, arg1::F)
+    unrolled_map_expr(:func, SimilarTo{arg1}, size(F), (arg1,), (:arg1,))
 end
+# Unary versions for type conversion.  Need to override these explicitly to
+# prevent conflicts with Base.
+@inline map{T,N,S}(::Type{T}, arg1::FixedArray{T,N,S}) = arg1 # nop version
+immutable ConstructTypeFun{T}; end
+call{T}(::ConstructTypeFun{T}, x) = T(x)
+@inline map{T,FSA<:FixedArray}(::Type{T}, arg1::FSA) = map(ConstructTypeFun{T}(), similar(FSA, T), arg1)
 
 
-@generated function map{T}(::Type{T}, arg1::FixedArray)
-    eltype(arg1) == T && return :(arg1)
-    FSA = similar(arg1, T)
-    inner = fill_tuples_expr((inds...) -> :( T(arg1[$(inds...)]) ), size(FSA))
-    :( $FSA($(inner)) )
-end
-
-@inline map{R,C,T}(F::Type{T}, arg1::Mat{R,C,T}) = arg1
-@generated function map{R,C,T}(F::DataType, arg1::Mat{R,C,T})
-    inner = fill_tuples_expr((inds...) -> :( F(arg1[$(inds...)]) ), (R, C))
-    :( Mat{R, C, F}($(inner)) )
+# Nullary special case version.
+#
+# TODO: This is inconsistent, since it maps *indices* through the functor
+# rather than using a functor with no arguments.
+@generated function map{FSA <: FixedArray}(F, ::Type{FSA})
+    tuple_expr = fill_tuples_expr((inds...) -> :(F($(inds...))), size(FSA))
+    constructor_expr(FSA, tuple_expr)
 end
 
-@generated function map{FSA <: FixedVectorNoTuple}(F::DataType, arg1::FSA)
-    eltype(FSA) == F && return :(arg1)
-    inner = ntuple(i-> :(F(arg1[$i])), length(FSA))
-    :( similar(FSA, F)($(inner...)) )
-end

src/ops.jl

@@ -36,7 +36,7 @@ const unaryOps = (:-, :~, :conj, :abs,
                   :gamma, :lfact, :frexp, :modf, :airy, :airyai,
                   :airyprime, :airyaiprime, :airybi, :airybiprime,
                   :besselj0, :besselj1, :bessely0, :bessely1,
-                  :eta, :zeta, :digamma)
+                  :eta, :zeta, :digamma, :real, :imag)
 
 # vec-vec and vec-scalar
 const binaryOps = (:.+, :.-,:.*, :./, :.\, :.^,
@@ -80,12 +80,11 @@ for op in binaryOps
     functor_name, functor_expr = gen_functor(op, 2)
     eval(quote
         $functor_expr
-        @inline $op{T <: FixedArray}(x::T, y::T) = map($functor_name(), x, y)
-        @inline $op{T, T2, NDIM, SIZE}(x::FixedArray{T, NDIM, SIZE}, y::FixedArray{T2, NDIM, SIZE}) = $op(promote(x, y)...)
-        @inline $op{T <: Number}(x::T, y::FixedArray{T}) = map($functor_name(), x, y)
-        @inline $op{T1 <: Number, T2}(x::T1, y::FixedArray{T2}) = $op(promote(x, y)...)
-        @inline $op{T <: Number}(x::FixedArray{T}, y::T) = map($functor_name(), x, y)
-        @inline $op{T1, T2 <: Number}(x::FixedArray{T1}, y::T2) = $op(promote(x, y)...)
+        @inline $op{F1<:FixedArray, F2<:FixedArray}(x::F1, y::F2) = map($functor_name(), x, y)
+        @inline $op{F<:FixedArray}(x::F, y::Number) = map($functor_name(), x, y)
+        @inline $op{F<:FixedArray}(x::Number, y::F) = map($functor_name(), x, y)
+        @inline $op{T,N}(x::FixedArray{T,N}, y::AbstractArray{T,N}) = map($functor_name(), x, y)
+        @inline $op{T,N}(x::AbstractArray{T,N}, y::FixedArray{T,N}) = map($functor_name(), x, y)
     end)
 end