; Haskell style monads for scheme
; Copyright 2005 Brian Alliet

(define create-monad-dict vector)

; Primitive functions to deal with monad dictionaries
; These functions extract functions from the monad dictionary
(define (return-value dict) (vector-ref dict 0))
(define (bind-value   dict) (vector-ref dict 1))
(define (fail-value   dict) (vector-ref dict 2))
(define (mzero-value  dict) (vector-ref dict 3))
(define (mplus-value  dict) (vector-ref dict 4))

; And these call the above functions with arguments
(define (return dict x)  ((return-value dict) x  ))
(define (bind dict m f)  ((bind-value   dict) m f))
(define (fail dict s)    ((fail-value   dict) s  ))
(define mzero mzero-value)
(define (mplus dict a b) ((mplus-value  dict) a b))

; Generic monad function defined in terms of the above interface
(define (bind_ d p q) (bind d p (lambda (x) q)))
(define (msum d: xs) (foldr (mplus-value d:) (mzero-value d:) xs))
(define (lift-m d: f) (lambda (x) (bind d: x (lambda (x2) (return d: (f x2))))))
(define (sequence d: xs)
    (let*
        ((b (bind-value d:))
         (r (return-value d:))
         (mcons (lambda (p q) (b p (lambda (x) (b q (lambda (y) (r (cons x y)))))))))
        (foldr mcons (return d: '()) xs)))
(define (sequence_ d: xs) (foldr (curry bind_ d:) (return d: '()) xs))
(define (map-m d f xs)  (sequence d (map f xs)))
(define (map-m_ d f xs) (sequence_ d (map f xs)))

(define-syntax do
    (syntax-rules (<- let)
        ((do d: (p <- e) rest ...) (bind d: e
            (lambda-match
                (p (do d: rest ...))
                (_ (fail d: "pattern match failure in do expression")))))
        ((do d: (let (var expr) ...) rest ...) (let ((var expr) ...) (do d: rest ...)))
        ((do d: e) e)
        ((do d: e rest ...) (bind d: e (lambda (ignore) (do d: rest ...))))
    ))

; Monad instances

; List Monad
(define (list:return x) (list x))
(define (list:bind x f) (concat-map f x))
(define (list:fail s) '())
(define list:mzero '())
(define list:mplus append)
(define list: (create-monad-dict list:return list:bind list:fail list:mzero list:mplus))

; Maybe Monad (which isn't lazy, so doesn't work as well as in Haskell)
(define (maybe:return x) (cons 'just x))
(define (maybe:bind mx f)
    (match mx
        ((cons 'just x) (f x))
        ('nothing 'nothing)
        (_ (error "maybe:bind thats not a maybe"))))
(define (maybe:fail s) 'nothing)
(define maybe:mzero 'nothing)
(define (maybe:mplus l r)
    (match l
        ((cons 'just _) l)
        ('nothing r)
        (_ (error "mplus: thats not a maybe"))))
(define maybe: (create-monad-dict  maybe:return maybe:bind maybe:fail maybe:mzero maybe:mplus))

; Identity Monad
(define (id:return x) x)
(define (id:bind x f) (f x))
(define (id:fail s) (error s))
(define id: (create-monad-dict id:return id:bind id:fail))

; State Monad
(define (state:return x) (lambda (state) (cons state x)))
(define (state:bind r f) (lambda (state) (match (r state) ((cons state2 x) ((f x) state2)))))
(define (state:fail s) (error s))
(define state: (create-monad-dict state:return state:bind state:fail))

(define get-state (lambda (state) (cons state state)))
(define (put-state x) (lambda (state) (cons x '())))
(define (modify-state f) (state:bind get-state (lambda (s) (put-state (f s)))))