implement generator expressions (#4470)

This introduces the types `Generator`, which maps a function over
an iterator, and `IteratorND`, which wraps an iterator with a
shape tuple.

base/boot.jl

@@ -341,6 +341,7 @@ unsafe_convert{T}(::Type{T}, x::T) = x
 (::Type{Array{T}}){T}(m::Int, n::Int, o::Int) = Array{T,3}(m, n, o)
 
 # TODO: possibly turn these into deprecations
+Array{T,N}(::Type{T}, d::NTuple{N,Int}) = Array{T}(d)
 Array{T}(::Type{T}, d::Int...) = Array{T}(d)
 Array{T}(::Type{T}, m::Int)               = Array{T,1}(m)
 Array{T}(::Type{T}, m::Int,n::Int)        = Array{T,2}(m,n)

base/essentials.jl

@@ -193,3 +193,17 @@ const (:) = Colon()
 # For passing constants through type inference
 immutable Val{T}
 end
+
+immutable Generator{I,F}
+    f::F
+    iter::I
+end
+
+start(g::Generator) = start(g.iter)
+done(g::Generator, s) = done(g.iter, s)
+function next(g::Generator, s)
+    v, s2 = next(g.iter, s)
+    g.f(v), s2
+end
+
+collect(g::Generator) = map(g.f, g.iter)

base/iterator.jl

@@ -292,3 +292,65 @@ eltype{I1,I2}(::Type{Prod{I1,I2}}) = tuple_type_cons(eltype(I1), eltype(I2))
     x = prod_next(p, st)
     ((x[1][1],x[1][2]...), x[2])
 end
+
+_size(p::Prod2) = (length(p.a), length(p.b))
+_size(p::Prod) = (length(p.a), _size(p.b)...)
+
+immutable IteratorND{I,N}
+    iter::I
+    dims::NTuple{N,Int}
+
+    function (::Type{IteratorND}){I,N}(iter::I, shape::NTuple{N,Integer})
+        if length(iter) != prod(shape)
+            throw(DimensionMismatch("dimensions $shape must be consistent with iterator length $(iter(a))"))
+        end
+        new{I,N}(iter, shape)
+    end
+    (::Type{IteratorND}){I<:AbstractProdIterator}(p::I) = IteratorND(p, _size(p))
+end
+
+start(i::IteratorND) = start(i.iter)
+done(i::IteratorND, s) = done(i.iter, s)
+next(i::IteratorND, s) = next(i.iter, s)
+
+size(i::IteratorND) = i.dims
+length(i::IteratorND) = length(i.iter)
+ndims{I,N}(::IteratorND{I,N}) = N
+
+eltype{I}(::IteratorND{I}) = eltype(I)
+
+collect(i::IteratorND) = copy!(Array(eltype(i),size(i)), i)
+
+function collect{I<:IteratorND}(g::Generator{I})
+    sz = size(g.iter)
+    if length(g.iter) == 0
+        return Array(Union{}, sz)
+    end
+    st = start(g)
+    first, st = next(g, st)
+    dest = Array(typeof(first), sz)
+    dest[1] = first
+    return map_to!(g.f, 2, st, dest, g.iter)
+end
+
+# special case for 2d
+function collect{F,I<:Prod2}(g::Generator{IteratorND{I,2},F})
+    f = g.f
+    a = g.iter.iter.a
+    b = g.iter.iter.b
+    sz = size(g.iter)
+    if prod(sz) == 0
+        return Array(Union{}, sz)
+    end
+    fst = f((first(a), first(b)))  # TODO: don't recompute this in the loop
+    dest = Array(typeof(fst), sz)
+    k = 1
+    for j in b
+        for i in a
+            val = f((i, j))  # TODO: handle type changes
+            @inbounds dest[k] = val
+            k += 1
+        end
+    end
+    return dest
+end

src/julia-parser.scm

@@ -1400,21 +1400,22 @@
   (parse-comma-separated s parse-eq*))
 
 ;; as above, but allows both "i=r" and "i in r"
+(define (parse-iteration-spec s)
+  (let ((r (parse-eq* s)))
+    (cond ((and (pair? r) (eq? (car r) '=))  r)
+          ((eq? r ':)  r)
+          ((and (length= r 4) (eq? (car r) 'comparison)
+                (or (eq? (caddr r) 'in) (eq? (caddr r) '∈)))
+           `(= ,(cadr r) ,(cadddr r)))
+          (else
+           (error "invalid iteration specification")))))
+
 (define (parse-comma-separated-iters s)
   (let loop ((ranges '()))
-    (let ((r (parse-eq* s)))
-      (let ((r (cond ((and (pair? r) (eq? (car r) '=))
-                      r)
-                     ((eq? r ':)
-                      r)
-                     ((and (length= r 4) (eq? (car r) 'comparison)
-                           (or (eq? (caddr r) 'in) (eq? (caddr r) '∈)))
-                      `(= ,(cadr r) ,(cadddr r)))
-                     (else
-                      (error "invalid iteration specification")))))
-        (case (peek-token s)
-          ((#\,)  (take-token s) (loop (cons r ranges)))
-          (else   (reverse! (cons r ranges))))))))
+    (let ((r (parse-iteration-spec s)))
+      (case (peek-token s)
+        ((#\,)  (take-token s) (loop (cons r ranges)))
+        (else   (reverse! (cons r ranges)))))))
 
 (define (parse-space-separated-exprs s)
   (with-space-sensitive
@@ -1471,6 +1472,12 @@
                        (begin (take-token s) (loop (cons nxt lst))))
                       ((eqv? c #\;)          (loop (cons nxt lst)))
                       ((equal? c closer)     (loop (cons nxt lst)))
+                      ((eq? c 'for)
+                       (take-token s)
+                       (let ((gen (parse-generator s nxt #f)))
+                         (if (eqv? (require-token s) #\,)
+                             (take-token s))
+                         (loop (cons gen lst))))
                       ((eq? c 'for)          (take-token s) (parse-generator s t closer))
                       ;; newline character isn't detectable here
                       #;((eqv? c #\newline)
@@ -1516,7 +1523,7 @@
 (define (parse-comprehension s first closer)
   (let ((r (parse-comma-separated-iters s)))
     (if (not (eqv? (require-token s) closer))
-        (error (string "expected " closer))
+        (error (string "expected \"" closer "\""))
         (take-token s))
     `(comprehension ,first ,@r)))
 
@@ -1526,12 +1533,11 @@
         `(dict_comprehension ,@(cdr c))
         (error "invalid dict comprehension"))))
 
-(define (parse-generator s first closer)
-  (let ((r (parse-comma-separated-iters s)))
-    (if (not (eqv? (require-token s) closer))
-	(error (string "expected " closer))
-        (take-token s))
-    `(macrocall @generator ,first ,@r)))
+(define (parse-generator s first allow-comma)
+  (let ((r (if allow-comma
+               (parse-comma-separated-iters s)
+               (list (parse-iteration-spec s)))))
+    `(generator ,first ,@r)))
 
 (define (parse-matrix s first closer gotnewline)
   (define (fix head v) (cons head (reverse v)))
@@ -1961,6 +1967,13 @@
                             `(tuple ,ex)
                             ;; value in parentheses (x)
                             ex))
+                       ((eq? t 'for)
+                        (take-token s)
+                        (let ((gen (parse-generator s ex #t)))
+                          (if (eqv? (require-token s) #\) )
+                              (take-token s)
+                              (error "expected \")\""))
+                          gen))
                        (else
                         ;; tuple (x,) (x,y) (x...) etc.
                         (if (eqv? t #\, )

src/julia-syntax.scm

@@ -1846,6 +1846,24 @@
               (lower-ccall name RT (cdr argtypes) args))))
          e))
 
+   'generator
+   (lambda (e)
+     (let ((expr   (cadr e))
+           (vars   (map cadr  (cddr e)))
+           (ranges (map caddr (cddr e))))
+       (let* ((argname (if (and (length= vars 1) (symbol? (car vars)))
+                           (car vars)
+                           (gensy)))
+              (splat (if (eq? argname (car vars))
+                         '()
+                         `((= (tuple ,@vars) ,argname)))))
+         (expand-forms
+          (expand-binding-forms
+           `(call (top Generator) (-> ,argname (block ,@splat ,expr))
+                  ,(if (length= ranges 1)
+                       (car ranges)
+                       `(call (top IteratorND) (call (top product) ,@ranges)))))))))
+
    'comprehension
    (lambda (e)
      (expand-forms (lower-comprehension #f       (cadr e) (cddr e))))

test/functional.jl

@@ -178,3 +178,22 @@ let
     foreach((args...)->push!(a,args), [2,4,6], [10,20,30])
     @test a == [(2,10),(4,20),(6,30)]
 end
+
+# generators (#4470, #14848)
+
+@test sum(i/2 for i=1:2) == 1.5
+@test collect(2i for i=2:5) == [4,6,8,10]
+@test collect((i+10j for i=1:2,j=3:4)) == [31 41; 32 42]
+
+let I = Base.IteratorND(1:27,(3,3,3))
+    @test collect(I) == reshape(1:27,(3,3,3))
+    @test size(I) == (3,3,3)
+    @test length(I) == 27
+    @test eltype(I) === Int
+    @test ndims(I) == 3
+end
+
+let A = collect(Base.Generator(x->2x, Real[1.5,2.5]))
+    @test A == [3,5]
+    @test isa(A,Vector{Float64})
+end