pipes in racket (typed no-check)

#lang typed/racket/no-check

(require racket/match)

;; Types ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(struct: (I O R) Yield ([out : O] [next : (Pipe I O R)]))
(struct: (I O R) Await ([next/f : (I -> (Pipe I O R))]))
(struct: (I O R) Effect ([thunk : (-> (Pipe I O R))]))
(struct: (R) Done ([result : R]))

(define-type (Pipe I O R)
  (U (Done R)
     (Effect I O R)
     (Yield I O R)
     (Await I O R)))
(define-type (Conduit I O)
  (All (R) (Pipe I O R)))
(define-type (Source O)
  (All (I) (Pipe I O Void)))
(define-type (Sink I)
  (All (O) (Pipe I O Void)))

;; Pipe monad ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(: return/pipe (All (R) (R -> (All (I O) (Pipe I O R)))))
(define (return/pipe x) (Done x))

(: bind/pipe
   (All (I O A R) ((Pipe I O A) (A -> (Pipe I O R)) -> (Pipe I O R))))
(define (bind/pipe m f)
  (define (go x) (bind/pipe x f))
  (match m
    [(Done result) (f result)]
    [(Effect thunk) (Effect (λ () (go (thunk))))]
    [(Yield out next) (Yield out (go next))]
    [(Await next/f) (Await (compose go next/f))]))

(define-syntax begin/pipe
  (syntax-rules ()
    [(begin/pipe) done]
    [(begin/pipe p) p]
    [(begin/pipe p1 pn ...)
     (bind/pipe p1 (λ (i) (begin/pipe pn ...)))]))

(: forever/pipe (All (I O A) ((Pipe I O A) -> (All (R) (Pipe I O R)))))
(define (forever/pipe m)
  (begin/pipe
    m
    (forever/pipe m)))

;; Pipe monad transformer ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define-syntax-rule (lift/pipe e)
  (Effect (λ ()
            (define r e)
            (return/pipe r))))

(: lift./pipe (All (A R) ((A -> R) -> (A -> (All (I O) (Pipe I O R))))))
(define (lift./pipe proc)
  (λ (x)
    (Effect (λ ()
              (define r (proc x))
              (return/pipe r)))))

;; Pipe primitives ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(: done (All (I O) (Pipe I O Void)))
(define done (Done (void)))

(: yield (All (O) (O -> (All (I) (Pipe I O Void)))))
(define (yield out) (Yield out done))

(: await (All (I O) (Pipe I O I)))
(define await (Await Done))

;; A helper function

(: absurd (Nothing -> (All (A) A)))
(define (absurd n)
  (error "I can't believe you made it this far"))

;; Running pipes ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(: simulate/pipe
   (All (I O R A) ((Pipe I O R) (-> I) (O -> Void) (R -> A) -> A)))
(define (simulate/pipe p on-await on-yield on-done)
  (define (go n)
        (simulate/pipe n on-await on-yield on-done))
  (match p
    [(Done result)    (on-done result)]
    [(Effect thunk)   (go (thunk))]
    [(Yield out next) (begin (on-yield out) (go next))]
    [(Await next/f)   (go (next/f (on-await)))]))

(: eval/pipe (All (O R) ((Pipe Void O R) -> (Values (Listof O) R))))
(define (eval/pipe p)
  (define outs (box empty))
  (define (on-yield out)
    (set-box! outs (cons out (unbox outs))))
  (define (on-done result)
    (values (unbox outs) result))
  (simulate/pipe p void on-yield on-done))

(: run/pipe ((Pipe Void Nothing Void) -> Void))
(define (run/pipe p)
  (simulate/pipe p void absurd void))

;; Composing pipes ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(: compose/pipe (All (I X O R) ((Pipe X O R) (Pipe I X R) -> (Pipe I O R))))
(define (compose/pipe p1 p2)
  (match p1
    [(Done result)    (Done result)]
    [(Effect thunk)   (Effect (λ () (compose/pipe (thunk) p2)))]
    [(Yield out next) (Yield out (compose/pipe next p2))]
    [(Await next/f1)
     (match p2
       [(Done result)    (Done result)]
       [(Effect thunk)   (Effect (λ () (compose/pipe p1 (thunk))))]
       [(Yield out next) (compose/pipe (next/f1 out) next)]
       [(Await next/f2)  (Await (λ (x) (compose/pipe p1 (next/f2 x))))]
       )]
    ))

(define (compose*/pipe . ps)
  (match ps
    [(list) done]
    [(list p) p]
    [(list-rest p1 p2 p-rest)
     (apply compose*/pipe (cons (compose/pipe p1 p2) p-rest))]))

;; A few simple pipes ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(: id/pipe (All (I) (Conduit I I)))
(define id/pipe (forever/pipe (bind/pipe await yield)))

(: list->pipe (All (O) ((Listof O) -> (Source O))))
(define (list->pipe xs)
  (match xs
    [(list) done]
    [(cons y ys)
     (begin/pipe
       (yield y)
       (list->pipe ys))]))

(: printer (All (I) (Sink I)))
(define printer (forever/pipe (bind/pipe await (lift./pipe print))))

(: func->pipe (All (I O) ((I -> O) -> (Conduit I O))))
(define (func->pipe f)
  (forever/pipe (bind/pipe await (compose yield f))))

(: add1/pipe (Conduit Integer Integer))
(define add1/pipe (func->pipe add1))

;; A few simple test runs ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; compose/pipe is associative
(run/pipe
 (compose/pipe printer (compose/pipe add1/pipe (list->pipe '(1 2 3)))))
(printf "\n")
(run/pipe
 (compose/pipe (compose/pipe printer add1/pipe) (list->pipe '(1 2 3))))
(printf "\n")


;; id/pipe is the identity of composition
(run/pipe
 (compose*/pipe printer add1/pipe (list->pipe '(1 2 3))))
(printf "\n")
(run/pipe
 (compose*/pipe id/pipe printer
                id/pipe add1/pipe
                id/pipe (list->pipe '(1 2 3))
                id/pipe))
(printf "\n")