;; Monads in scheme. Dynamic implementation with polymorphic 'bind'.

(module monad mzscheme

  ;; (require (lib "match.ss"))

  (provide unit
           bind
           define-monad
           monad->list   ;; for debugging
           Mlist)

  ;; Monads are characterized by

  ;; - a type constructor M
  ;; - unit :: a -> M a
  ;; - bind :: M a -> (a -> M b) -> M b

  ;; In words: something that creates the type (a struct in scheme),
  ;; something that puts a value into a monad (unit) and something
  ;; that takes values out of a monad, applies them to generate
  ;; several instances of the monad, and combines them into one.


  ;; The 'class' and 'instance' structures.
  
  (define-struct monad (name make unit bind))
  (define-struct monad-instance (monad contents))

  ;; Definition of a new monad in terms of its unit and bind
  ;; functions. This provides all the structure wrapping necessary for
  ;; polymorhpy, so the implementation of 'unit' and 'bind' can focus
  ;; on the representation.

  ;; The representative of the monad is a single object which can be
  ;; used to convert a representation into a wrapped one.

  (define ((monad-rep monad-class) . args)
    (if (null? args) monad-class
        (apply wrap monad-class args)))
  
  (define-syntax define-monad
    (syntax-rules ()
      ((_ name unit bind)
       (begin
         (define monad-class
           (make-monad
            'name
            (lambda (contents)
              (make-monad-instance monad-class contents))
            unit bind))
         (define name
           (monad-rep monad-class))))))
  
  
  ;; Creation of a monad structure from its internal representation
  ;; and back.
  
  (define (wrap M x)  ((monad-make M) x))
  (define (unwrap Mx) (monad-instance-contents Mx))

  ;; For debugging
  (define (monad->list Ma . port)
    (let ((M (monad-instance-monad Ma)))
      (list
       (monad-name M)
       (monad-instance-contents Ma))))
  
  
  ;; The generic interface functions. We can't infer the type of unit,
  ;; so it needs to be provided explicitly.

  (define (_unit M a) (wrap M ((monad-unit M) a)))
  (define (unit M a)  (_unit (M) a))

              

  ;; For bind, the type can be inferred from the data. Assume the
  ;; function a->Mb has correct type, which means it operates on base
  ;; type and produces a wrapped monad.

  (define (bind Ma a->Mb)
    (let ((M      (monad-instance-monad Ma))
          (_Ma    (monad-instance-contents Ma))
          (_a->Mb (lambda (a) (unwrap (a->Mb a)))))

      (wrap M ((monad-bind M) _Ma _a->Mb))))
             

  ;; Map and Join are defined in terms of bind and unit.
  (define (join MMa)
    (bind MMa (lambda (Ma) Ma)))

  (define (fmap a->b Ma)
    (bind Ma (lambda (a)
               (_unit (monad-instance-monad Ma)
                      (a->b a)))))

  ;; 'do' notation for bind, schemefied in the form of a let* 
  (define-syntax letM*
    (syntax-rules ()
      ((_ () expr) expr)
      ((_ ((n Mv) bindings ...) expr)
       (bind Mv
             (lambda (n)
               (letM* (bindings ...) expr))))))
  

  ;; Some examples. The functions 'unit' and 'bind' operate on
  ;; internal representation only, and are independent of the way
  ;; monads are represented.
  
  (define-monad Mlist
    (lambda (a)        (list a))
    (lambda (Ma a->Mb) (apply append (map a->Mb Ma))))

  ;; TEST: (bind (Mlist '(1 2 3)) (lambda (x) (unit Mlist (+ 1 x))))
  
  ;; The maybe monad
  (define-monad Mfail
    (lambda (a)        a)
    (lambda (Ma a->Mb) (if (eq? Ma 'fail) 'fail
                           (a->Mb Ma))))

  (define fail (Mfail 'fail))
  
  ;; TEST: (bind (Mfail 0) (lambda (x) (if (zero? x) fail (unit Mfail (/ 1 x)))))


  ;; The writer monad.
  (define-monad Mwriter
    (lambda (a)         (list a))
    (lambda (Ma a->Mb)  (cons (a->Mb 'dummy) Ma)))


  ;; The state monad. Values in the state monad are state
  ;; transformers: transforming a state into a value and a new state.

  (define-monad Mstate

    ;; Converting any value into a state monad creates a function that
    ;; returns the state it got passed.
    
    (lambda (a)
      (lambda (s) (values a s)))
    
    (lambda (Ma a->Mb)

      ;; Bind creates a new state transformer that,
      (lambda (s0)

        ;; when applied to a state 's0', first transforms this state
        ;; to obtain a value 'a' and an intermediate state 's1', and
        (let-values (((a s1) (Ma s0)))

          ;; applies the function we're binding to obtain a new state
          ;; transformer, which is applied to the intermediate state.
          ((a->Mb a) s1))))
    
    )



  ;; To test, we create a simple generator written as a state
  ;; transformer. The representation is a function which accepts a
  ;; state and returns a value and a new state.
  
  
  )