(define-module (control monad state)
  #:use-module (oop goops)
  #:use-module (ice-9 match)
  #:use-module (control monad)
  #:export (return-state run-state
                         get put
                         ;; TODO
                         ;; should these even be exported?
                         define-stateful
                         define-stateful-method)
  #:re-export (>>= >>))


;; Alternative implementation of get.
;; See https://hackage.haskell.org/package/mtl-2.2.1/docs/src/Control.Monad.State.Class.html#get

;;; newtype State = st-list -> st-list

;;; state procedure <=> st-list -> st-list
;;; state list <=> (list ret st)

;;; Wrapper around a procedure with signature:
;;; (st-list -> st-list). Wrapped to allow goops
;;; multiple dispatch to do its thing.
(define-class <state> ()
  (proc #:init-keyword #:proc
        #:getter proc))

;; (st-list -> st-list) -> State
(define (make-state proc)
  "Creates a state object from a State procedure"
  (make <state> #:proc proc))

;;; Define a procedure which is in the state monad. This means that it takes a
;;; state list as a curried argument, and it's return is wrappen in a <state>
;;; object.
;;; It's fully possible to create stateful objects without these macros, but it's
;;; ill adviced since that would just be boilerplate.

(define-syntax-rule (define-stateful ((proc args ...) st) body ...)
  (define (proc args ...)
    (make-state
     (lambda (st) body ...))))

(define-syntax-rule (define-stateful-method ((proc args ...) st) body ...)
  (define-method (proc args ...)
    (make-state
     (lambda (st) body ...))))

(define-stateful-method ((>>= (st <state>) f) st-list)
  (let ((new-st-list ((proc st) st-list)))
    (match new-st-list
      ((v _)
       ((proc (f v)) new-st-list)))))

(define-stateful-method ((>> (a <state>) (b <state>)) st-list-a)
  (let ((st-list-b ((proc a) st-list-a)))
    ((proc b) st-list-b)))

(define-stateful ((return-state v) st-list)
  "Sets the return value to v"
  (cons v (cdr st-list)))

(define-stateful ((get) st-list)
  "Sets the return value of state to st."
  (match st-list
    ((_ st)
     (list st st))))

(define-stateful ((put v) st-list)
  "Sets st to v."
  (list '() v))

(define-method (run-state (st <state>) init)
  "Exec state with init as starting state value and st."
  ((proc st) (list init init)))