1
struct EnumerableMapMany{T,SO,CS<:Function,RS<:Function} <: Enumerable
2 23
    source::SO
3
    collectionSelector::CS
4
    resultSelector::RS
5
end
6

7 0
Base.eltype(::Type{EnumerableMapMany{T,SO,CS,RS}}) where {T,SO,CS,RS} = T
8

9
# TODO Make sure this is actually correct. We might have to be more selective,
10
# i.e. only scan arguments for certain types of expression etc.
11
function expr_contains_ref_to(expr::Expr, var_name::Symbol)
12 23
    for sub_expr in expr.args
13 23
        if isa(sub_expr, Symbol)
14 23
            if sub_expr==var_name
15 23
                return true
16
            end
17
        else
18 23
            test_sub = expr_contains_ref_to(sub_expr, var_name)
19 23
            if test_sub
20 23
                return true
21
            end
22
        end
23
    end
24 0
    return false
25
end
26

27
function expr_contains_ref_to(expr::Symbol, var_name::Symbol)
28 23
    return expr==var_name
29
end
30

31
function expr_contains_ref_to(expr::QuoteNode, var_name::Symbol)
32 0
    return expr==var_name
33
end
34

35
function expr_contains_ref_to(expr::Function, var_name::Symbol)
36 0
    return false
37
end
38

39 0
expr_contains_ref_to(::Number, ::Symbol) = false
40

41
function mapmany(source::Enumerable, f_collectionSelector::Function, collectionSelector::Expr, f_resultSelector::Function, resultSelector::Expr)
42 23
    TS = eltype(source)
43
    # First detect whether the collectionSelector return value depends at all
44
    # on the value of the anonymous function argument
45 23
    anon_var = collectionSelector.head==:escape ? collectionSelector.args[1].args[1] : collectionSelector.args[1]
46 23
    body = collectionSelector.head==:escape ? collectionSelector.args[1].args[2].args[2] : collectionSelector.args[2].args[2]
47 23
    crossJoin = !expr_contains_ref_to(body, anon_var)
48

49 23
    if crossJoin
50 0
        inner_collection = f_collectionSelector(nothing)
51
        input_type_collection_selector = typeof(inner_collection)
52 23
        TCE = input_type_collection_selector.parameters[1]
53
    else
54
        input_type_collection_selector = Base._return_type(f_collectionSelector, Tuple{TS,})
55
        TCE = typeof(input_type_collection_selector)==Union || input_type_collection_selector==Any ? Any : eltype(input_type_collection_selector)
56
    end
57

58 23
    T = Base._return_type(f_resultSelector, Tuple{TS,TCE})
59
    SO = typeof(source)
60

61
    CS = typeof(f_collectionSelector)
62 23
    RS = typeof(f_resultSelector)
63

64 23
    return EnumerableMapMany{T,SO,CS,RS}(source,f_collectionSelector,f_resultSelector)
65
end
66

67
# TODO This should be changed to a lazy implementation
68
function Base.iterate(iter::EnumerableMapMany{T,SO,CS,RS}) where {T,SO,CS,RS}
69 23
    results = Array{T}(undef, 0)
70 23
    for i in iter.source
71 23
        for j in iter.collectionSelector(i)
72 23
            push!(results,iter.resultSelector(i,j))
73
        end
74
    end
75

76 23
    if length(results)==0
77 0
        return nothing
78
    end
79

80 23
    return results[1], (results,2)
81
end
82

83
function Base.iterate(iter::EnumerableMapMany{T,SO,CS,RS}, state) where {T,SO,CS,RS}
84 10
    if state[2]>length(state[1])
85 10
        return nothing
86
    else
87 10
        return state[1][state[2]], (state[1], state[2]+1)
88
    end
89
end

Read our documentation on viewing source code .

Loading