From ed0bd51c8b38dd4bd007f8f4f112701d1cf4df9a Mon Sep 17 00:00:00 2001 From: Noel Welsh Date: Mon, 4 Aug 2008 18:38:47 +0000 Subject: [PATCH] Add support for start/stop/step parameters to in-vector comprehension forms. Vector like comprehensions (strings etc) actually support this as well, when used in the body of the for comprehension. However this will not be documented as support do not extend to sequences defined outside a comprehension. svn: r11068 --- collects/scheme/private/for.ss | 97 ++++++++++++++++--- .../mzscheme/benchmarks/mz/comprehensions.ss | 44 +++++++++ collects/tests/mzscheme/for.ss | 6 ++ 3 files changed, 132 insertions(+), 15 deletions(-) create mode 100644 collects/tests/mzscheme/benchmarks/mz/comprehensions.ss diff --git a/collects/scheme/private/for.ss b/collects/scheme/private/for.ss index 2ab6639789..b9b0c80743 100644 --- a/collects/scheme/private/for.ss +++ b/collects/scheme/private/for.ss @@ -95,7 +95,7 @@ (define cert-key (gensym 'for-cert)) (define (certify-clause src-stx clause certifier introducer) - ;; This is slightly painful. The painsion into `:do-in' involves a lot of pieces + ;; This is slightly painful. The expansion into `:do-in' involves a lot of pieces ;; that are no treated as sub-expressions. We have to push the certificates ;; down to all the relevant identifiers and expressions: (define (recert s) (syntax-recertify s src-stx (current-inspector) cert-key)) @@ -385,19 +385,39 @@ (define (:list-gen l) (values car cdr l pair? (lambda (x) #t) (lambda (p x) #t))) - (define (in-vector l) - (unless (vector? l) (raise-type-error 'in-vector "vector" l)) - (make-do-sequence (lambda () (:vector-gen l)))) + (define in-vector + (case-lambda + [(v) (in-vector v 0 (vector-length v) 1)] + [(v start) (in-vector v start (vector-length v) 1)] + [(v start stop) (in-vector v start stop 1)] + [(v start stop step) + (unless (vector? v) (raise-type-error 'in-vector "vector" v)) + (when (and (< start stop) (< step 0)) + (raise-mismatch-error 'in-vector "start is less than stop but step is negative" (list start stop step))) + (when (and (< stop start) (> step 0)) + (raise-mismatch-error 'in-vector "stop is less than start but step is positive" (list start stop step))) + (when (zero? step) + (raise-mismatch-error 'in-vector "step is zero" step)) + (make-do-sequence (lambda () (:vector-gen v start stop step)))])) - (define (:vector-gen v) - (let ([len (vector-length v)]) - (values (lambda (i) - (vector-ref v i)) - add1 - 0 - (lambda (i) (< i len)) - (lambda (x) #t) - (lambda (x y) #t)))) + (define :vector-gen + (case-lambda + [(v) (:vector-gen v 0 (vector-length v) 1)] + [(v start stop step) + (values + ;; pos->element + (lambda (i) (vector-ref v i)) + ;; next-pos + ;; Minor optimisation. I assume add1 is faster than \x.x+1 + (if (= step 1) add1 (lambda (i) (+ i step))) + ;; initial pos + start + ;; continue? + (if (> step 0) + (lambda (i) (< i stop)) + (lambda (i) (> i stop))) + (lambda (x) #t) + (lambda (x y) #t))])) (define (in-string l) (unless (string? l) (raise-type-error 'in-string "string" l)) @@ -970,12 +990,13 @@ vector-length-id in-vector-id vector-ref-id) - (lambda (stx) + (define (in-vector-like stx) (with-syntax ([vector? vector?-id] [in-vector in-vector-id] [vector-length vector-length-id] [vector-ref vector-ref-id]) (syntax-case stx () + ;; Fast case [((id) (_ vec-expr)) #'[(id) (:do-in @@ -998,7 +1019,53 @@ #t ;; loop args ((add1 pos)))]] - [_ #f])))) + ;; General case + [((id) (_ vec-expr start)) + (in-vector-like (syntax ((id) (_ vec-expr start #f 1))))] + [((id) (_ vec-expr start stop)) + (in-vector-like (syntax ((id) (_ vec-expr start stop 1))))] + [((id) (_ vec-expr start stop step)) + #`[(id) + (:do-in + ;; Outer bindings + ;; Prevent multiple evaluation + ([(v* stop*) (let ([vec vec-expr] + [stop* stop]) + (if stop* + (values vec stop*) + (values vec (vector-length vec))))] + [(start*) start] + [(step*) step]) + ;; Outer check + (when (or (zero? step*) + (and (< start* stop*) (< step* 0)) + (and (< start* stop*) (< step* 0))) + (if (vector? v*) + ;; Let in-vector report the error + (in-vector v* start* stop* step*) + (raise-type-error in-vector + "start, stop, and step incompatible" + (list start* stop* step*)))) + ;; Loop bindings + ([idx start*]) + ;; Pos guard + #,(cond + [(not (number? (syntax-e #'step))) + #`(if (step* . >= . 0) (< idx stop*) (> idx stop*))] + [((syntax-e #'step) . >= . 0) + #'(< idx stop*)] + [else + #'(> idx stop*)]) + ;; Inner bindings + ([(id) (vector-ref v* idx)]) + ;; Pre guard + #t + ;; Post guard + #t + ;; Loop args + ((+ idx step)))]] + [_ #f]))) + in-vector-like) (define-sequence-syntax *in-vector (lambda () #'in-vector) diff --git a/collects/tests/mzscheme/benchmarks/mz/comprehensions.ss b/collects/tests/mzscheme/benchmarks/mz/comprehensions.ss new file mode 100644 index 0000000000..de0ea9257c --- /dev/null +++ b/collects/tests/mzscheme/benchmarks/mz/comprehensions.ss @@ -0,0 +1,44 @@ +#lang scheme/base + +(require (planet schematics/schemeunit:3) + (planet schematics/benchmark:2)) + +;; Test that comprehensions are as fast as hand-written +;; loops + + +;; +;; Vector comprehensions +;; + +(define big-vector (make-vector 65536 1)) + +(test-case + "simplest vector comprehension" + (check-as-fast + "comprehension" + (lambda () + (for/fold ([sum 0]) ([x (in-vector big-vector)]) + (+ sum x))) + "hand-written loop" + (lambda () + (let ([end (vector-length big-vector)]) + (let loop ([i 0] [sum 0]) + (if (= i end) + sum + (loop (add1 i) (+ (vector-ref big-vector i) sum)))))))) + +(test-case + "vector comprehension with step" + (check-as-fast + "comprehension" + (lambda () + (for/fold ([sum 0]) ([x (in-vector big-vector 0 (vector-length big-vector) 2)]) + (+ sum x))) + "hand-written loop" + (lambda () + (let ([end (vector-length big-vector)]) + (let loop ([i 0] [sum 0]) + (if (= i end) + sum + (loop (+ i 2) (+ (vector-ref big-vector i) sum)))))))) \ No newline at end of file diff --git a/collects/tests/mzscheme/for.ss b/collects/tests/mzscheme/for.ss index 6ee941ab69..8bfab238bd 100644 --- a/collects/tests/mzscheme/for.ss +++ b/collects/tests/mzscheme/for.ss @@ -105,6 +105,12 @@ (test-generator [(a b c)] (in-list '(a b c))) (test-generator [(a b c)] #(a b c)) (test-generator [(a b c)] (in-vector #(a b c))) +(test-generator [(b c d)] (in-vector #(a b c d) 1)) +(test-generator [(b c d)] (in-vector #(a b c d e) 1 4)) +(test-generator [(b d f)] (in-vector #(a b c d e f g h) 1 7 2)) +(test-generator [(h f d)] (in-vector #(a b c d e f g h) 7 1 -2)) +(test-generator [(b d f)] (in-vector #(a b c d e f g h) 1 6 2)) +(test-generator [(h f d)] (in-vector #(a b c d e f g h) 7 2 -2)) (test-generator [(#\a #\b #\c)] "abc") (test-generator [(#\a #\b #\c)] (in-string "abc")) (test-generator [(65 66 67)] #"ABC")