preserves/implementations/racket/preserves/preserves/iolist.rkt

70 lines
2.5 KiB
Racket

#lang racket/base
(provide counted-iolist?
counted-iolist-value
counted-iolist-length
count-iolist
iolist-length
iolist->bytes
write-iolist)
(struct counted-iolist (value length) #:prefab)
(define (bad-iolist who v)
(error who "Invalid iolist: ~v" v))
(define (iolist-length i [acc 0])
(cond [(byte? i) (+ acc 1)]
[(bytes? i) (+ acc (bytes-length i))]
[(pair? i) (iolist-length (cdr i) (iolist-length (car i) acc))]
[(null? i) acc]
[(void? i) acc]
[(counted-iolist? i) (+ acc (counted-iolist-length i))]
[else (bad-iolist 'iolist-length i)]))
(define (count-iolist i)
(if (counted-iolist? i)
i
(counted-iolist i (iolist-length i))))
(define (iolist->bytes i)
(if (bytes? i)
i
(let ((buffer (make-bytes (iolist-length i))))
(let walk ((i i) (offset 0))
(cond [(byte? i) (bytes-set! buffer offset i) (+ offset 1)]
[(bytes? i) (bytes-copy! buffer offset i) (+ offset (bytes-length i))]
[(pair? i) (walk (cdr i) (walk (car i) offset))]
[(null? i) offset]
[(void? i) offset]
[(counted-iolist? i) (walk (counted-iolist-value i) offset)]))
buffer)))
(define (write-iolist i [out-port (current-output-port)])
(cond [(byte? i) (write-byte i out-port)]
[(bytes? i) (write-bytes i out-port)]
[(pair? i) (write-iolist (car i) out-port) (write-iolist (cdr i) out-port)]
[(null? i) (void)]
[(void? i) (void)]
[(counted-iolist? i) (write-iolist (counted-iolist-value i) out-port)]
[else (bad-iolist 'write-iolist i)]))
(module+ test
(require rackunit)
(require (only-in racket/port with-output-to-bytes))
(check-equal? (iolist-length '((1 . 1) (1 . #"xxx") . 1)) 7)
(check-equal? (iolist-length '((1 . 1) (1 . #"xxx"))) 6)
(check-equal? (iolist-length 123) 1)
(check-equal? (iolist-length #"123") 3)
(check-equal? (iolist-length (cons 123 #"123")) 4)
(check-equal? (iolist->bytes '((1 . 1) (1 . #"xxx") . 1)) (bytes 1 1 1 120 120 120 1))
(check-equal? (iolist->bytes '((1 . 1) (1 . #"xxx"))) (bytes 1 1 1 120 120 120))
(check-equal? (iolist->bytes 123) (bytes 123))
(check-equal? (iolist->bytes #"123") (bytes 49 50 51))
(check-equal? (iolist->bytes (cons 123 #"123")) (bytes 123 49 50 51))
(check-equal? (with-output-to-bytes (lambda () (write-iolist '((1 . 1) (1 . #"xxx") . 1))))
(bytes 1 1 1 120 120 120 1)))