From f8254f1fa5074d0102ba258ac459dd1186afefd8 Mon Sep 17 00:00:00 2001 From: Jeff Bezanson Date: Mon, 27 Jun 2016 15:46:09 -0400 Subject: [PATCH] add comprehension filters (fixes #550) and nested loops (fixes #4867) this also uses the new lowering for typed comprehensions, allowing all comprehensions on unknown-length iterables (fixes #1457) --- src/julia-parser.scm | 20 +++++++--- src/julia-syntax.scm | 90 +++++++++++++++++++++++++++++--------------- 2 files changed, 73 insertions(+), 37 deletions(-) diff --git a/src/julia-parser.scm b/src/julia-parser.scm index 3cd3896707f19..d29f7c04ebeb2 100644 --- a/src/julia-parser.scm +++ b/src/julia-parser.scm @@ -1511,22 +1511,30 @@ (else (error "missing separator in array expression"))))))) +(define (parse-generator s first) + (let ((iters (parse-comma-separated-iters s))) + (let ((iters (if (eq? (peek-token s) 'if) + (begin (take-token s) + (list `(filter ,(parse-cond s) ,@iters))) + iters))) + (if (eq? (peek-token s) 'for) + (begin (take-token s) + `(flatten (generator ,(parse-generator s first) ,@iters))) + `(generator ,first ,@iters))))) + (define (parse-comprehension s first closer) - (let ((r (parse-comma-separated-iters s))) + (let ((gen (parse-generator s first))) (if (not (eqv? (require-token s) closer)) (error (string "expected \"" closer "\"")) (take-token s)) - `(comprehension ,first ,@r))) + `(comprehension ,gen))) (define (parse-dict-comprehension s first closer) (let ((c (parse-comprehension s first closer))) - (if (dict-literal? (cadr c)) + (if (dict-literal? (cadr (cadr c))) `(dict_comprehension ,@(cdr c)) (error "invalid dict comprehension")))) -(define (parse-generator s first) - `(generator ,first ,@(parse-comma-separated-iters s))) - (define (parse-matrix s first closer gotnewline) (define (fix head v) (cons head (reverse v))) (define (update-outer v outer) diff --git a/src/julia-syntax.scm b/src/julia-syntax.scm index ac6cf8cb9ac6b..24442a7a0a55d 100644 --- a/src/julia-syntax.scm +++ b/src/julia-syntax.scm @@ -1526,6 +1526,26 @@ (define (syntactic-op-to-call e) `(call ,(car e) ,(expand-forms (cadr e)) ,(expand-forms (caddr e)))) +;; wrap `expr` in a function appropriate for consuming values from given ranges +(define (func-for-generator-ranges expr range-exprs) + (let* ((vars (map cadr range-exprs)) + (argname (if (and (length= vars 1) (symbol? (car vars))) + (car vars) + (gensy))) + (splat (cond ((eq? argname (car vars)) '()) + ((length= vars 1) + `(,@(map (lambda (v) `(local ,v)) (lhs-vars (car vars))) + (= ,(car vars) ,argname))) + (else + `(,@(map (lambda (v) `(local ,v)) (lhs-vars `(tuple ,@vars))) + (= (tuple ,@vars) ,argname)))))) + (if (and (null? splat) + (length= expr 3) (eq? (car expr) 'call) + (eq? (caddr expr) argname) + (not (expr-contains-eq argname (cadr expr)))) + (cadr expr) ;; eta reduce `x->f(x)` => `f` + `(-> ,argname (block ,@splat ,expr))))) + ;; table mapping expression head to a function expanding that form (define expand-table (table @@ -1960,52 +1980,60 @@ '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 (cond ((eq? argname (car vars)) '()) - ((length= vars 1) - `(,@(map (lambda (v) `(local ,v)) (lhs-vars (car vars))) - (= ,(car vars) ,argname))) - (else - `(,@(map (lambda (v) `(local ,v)) (lhs-vars `(tuple ,@vars))) - (= (tuple ,@vars) ,argname)))))) - (expand-forms - `(call (top Generator) - ,(if (and (null? splat) - (length= expr 3) (eq? (car expr) 'call) - (eq? (caddr expr) argname) - (not (expr-contains-eq argname (cadr expr)))) - (cadr expr) ;; eta reduce `x->f(x)` => `f` - `(-> ,argname (block ,@splat ,expr))) - ,(if (length= ranges 1) + (let* ((expr (cadr e)) + (filt? (eq? (car (caddr e)) 'filter)) + (range-exprs (if filt? (cddr (caddr e)) (cddr e))) + (ranges (map caddr range-exprs)) + (iter (if (length= ranges 1) (car ranges) - `(call (top product) ,@ranges))))))) + `(call (top product) ,@ranges))) + (iter (if filt? + `(call (top Filter) + ,(func-for-generator-ranges (cadr (caddr e)) range-exprs) + ,iter) + iter))) + (expand-forms + `(call (top Generator) + ,(func-for-generator-ranges expr range-exprs) + ,iter)))) + + 'flatten + (lambda (e) `(call (top Flatten) ,(expand-forms (cadr e)))) 'comprehension (lambda (e) - (if (any (lambda (x) (eq? x ':)) (cddr e)) - (error "comprehension syntax with `:` ranges has been removed")) - (expand-forms `(call (top collect) (generator ,(cadr e) ,@(cddr e))))) + (if (length> e 2) + ;; backwards compat for macros that generate :comprehension exprs + (expand-forms `(comprehension (generator ,@(cdr e)))) + (begin (if (and (eq? (caadr e) 'generator) + (any (lambda (x) (eq? x ':)) (cddr (cadr e)))) + (error "comprehension syntax with `:` ranges has been removed")) + (expand-forms `(call (top collect) ,(cadr e)))))) 'typed_comprehension (lambda (e) - (if (any (lambda (x) (eq? x ':)) (cdddr e)) - (error "comprehension syntax with `:` ranges has been removed")) - (expand-forms (lower-comprehension (cadr e) (caddr e) (cdddr e)))) + (expand-forms + (or (and (eq? (caaddr e) 'generator) + (let ((ranges (cddr (caddr e)))) + (if (any (lambda (x) (eq? x ':)) ranges) + (error "comprehension syntax with `:` ranges has been removed")) + (and (every (lambda (x) (and (pair? x) (eq? (car x) '=) + (pair? (caddr x)) (eq? (car (caddr x)) ':))) + ranges) + ;; TODO: this is a hack to lower simple comprehensions to loops very + ;; early, to greatly reduce the # of functions and load on the compiler + (lower-comprehension (cadr e) (cadr (caddr e)) ranges)))) + `(call (top collect) ,(cadr e) ,(caddr e))))) 'dict_comprehension (lambda (e) (syntax-deprecation #f "[a=>b for (a,b) in c]" "Dict(a=>b for (a,b) in c)") - (expand-forms `(call (top Dict) (generator ,(cadr e) ,@(cddr e))))) + (expand-forms `(call (top Dict) ,(cadr e)))) 'typed_dict_comprehension (lambda (e) (expand-forms `(call (call (core apply_type) (top Dict) ,@(cdr (cadr e))) - (generator ,(caddr e) ,@(cdddr e))))))) + ,(caddr e)))))) (define (lower-comprehension atype expr ranges) (let ((result (make-ssavalue))