RosettaCodeData/Task/Stream-merge/Racket/stream-merge.rkt
2023-07-01 13:44:08 -04:00

73 lines
2.7 KiB
Racket
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

;; This module produces a sequence that merges streams in order (by <)
#lang racket/base
(require racket/stream)
(define-values (tl-first tl-rest tl-empty?)
(values stream-first stream-rest stream-empty?))
(define-struct merged-stream (< ss v ss)
#:mutable ; sadly, so we don't have to redo potentially expensive <
#:methods gen:stream
[(define (stream-empty? S)
;; andmap defined to be true when ss is null
(andmap tl-empty? (merged-stream-ss S)))
(define (cache-next-head S)
(unless (box? (merged-stream-v S))
(define < (merged-stream-< S))
(define ss (merged-stream-ss S))
(define-values (best-f best-i)
(for/fold ((F #f) (I 0)) ((s (in-list ss)) (i (in-naturals)))
(if (tl-empty? s) (values F I)
(let ((f (tl-first s)))
(if (or (not F) (< f (unbox F))) (values (box f) i) (values F I))))))
(set-merged-stream-v! S best-f)
(define ss (for/list ((s ss) (i (in-naturals)) #:unless (tl-empty? s))
(if (= i best-i) (tl-rest s) s)))
(set-merged-stream-ss! S ss))
S)
(define (stream-first S)
(cache-next-head S)
(unbox (merged-stream-v S)))
(define (stream-rest S)
(cache-next-head S)
(struct-copy merged-stream S [ss (merged-stream-ss S)] [v #f]))])
(define ((merge-sequences <) . sqs)
(let ((strms (map sequence->stream sqs)))
(merged-stream < strms #f #f)))
;; ---------------------------------------------------------------------------------------------------
(module+ main
(require racket/string)
;; there are file streams and all sorts of other streams -- we can even read lines from strings
(for ((l ((merge-sequences string<?)
(in-lines (open-input-string "aardvark
dog
fox"))
(in-list (string-split "cat donkey elephant"))
(in-port read (open-input-string #<<<
"boy"
"emu"
"monkey"
<
)))))
(displayln l)))
;; ---------------------------------------------------------------------------------------------------
(module+ test
(require rackunit)
(define merge-sequences/< (merge-sequences <))
(check-equal?
(for/list ((i (in-stream (merge-sequences/< (in-list '(1 3 5)))))) i)
'(1 3 5))
;; in-stream (and in-list) is optional (but may increase performance)
(check-equal? (for/list ((i (merge-sequences/<))) i) null)
(check-equal? (for/list ((i (merge-sequences/< '(1 3 5) '(2 4 6)))) i) '(1 2 3 4 5 6))
(check-equal? (for/list ((i (merge-sequences/< '(1 3 5) '(2 4 6 7 8 9 10)))) i)
'(1 2 3 4 5 6 7 8 9 10))
(check-equal? (for/list ((i (merge-sequences/< '(2 4 6 7 8 9 10) '(1 3 5)))) i)
'(1 2 3 4 5 6 7 8 9 10)))