;; Example: the state monad seen in class
;; translated directly into Scheme
;; MPradella, MMXIII

#lang racket

(define (return x)
  (lambda (t)
    (cons t x)))

(define (>>= f g)
  (lambda (oldstate)
    (let* ((fappl    (f oldstate))
           (newstate (car fappl))
           (val      (cdr fappl)))
      ((g val) newstate))))

(define (>> m k) ; standard definition
  (>>= m (lambda (x) k)))

(define-syntax DO
    (syntax-rules (<-)
      ((_ s)                 s)
      ((_ (x <- s0) s1 ...)  (>>= s0 (lambda (x) (DO s1 ...))))
      ((_ s0 s1 ...)         (>> s0 (DO s1 ...)))     
      ))


(define (getState)
  (lambda (state)
    (cons state state)))

(define (putState new)
  (lambda (x)
    (cons new '())))



;; some of the examples seen in class:

(define (esm) ;; this without using the DO notation
  (>>= (return 5)
       (lambda (x)
         (return (+ x 1)))))

(displayln ((esm) 333))

(define (esm2)
  (DO (x <- (getState))
      (putState (+ x 1))
      (x <- (getState))
      (return x)))

(displayln ((esm2) 333))

;; State monad and trees. 
;; For simplicity: Leaf = Atom; Branch = Cons

(define (mapTreeM f tree)
  (if (not (cons? tree))
      (DO (b <- (f tree))
          (return b))
      (let ((lhs (car tree))
            (rhs (cdr tree)))
        (DO (lhs1 <- (mapTreeM f lhs))
            (rhs1 <- (mapTreeM f rhs))
            (return (cons lhs1 rhs1))))))

(define (runStateM f st)
  (cdr (f st)))

(define (numberTree tree)
  (define (number v)
    (DO  (curr <- (getState))
         (putState (+ curr 1))
         (return (cons v curr))))
  (mapTreeM number tree))

(define testTree (cons 
                  (cons 'a 
                        (cons 'b 'c))
                  (cons 'd 'e)))

(displayln (runStateM (numberTree testTree) 1))