(define-module (control monad state)
  #:use-module (oop goops)
  #:use-module (ice-9 match)
  #:use-module (control monad)
  #:export (make-state return-state
                       get put
                       run-state)
  #: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-method (>>= (st <state>) f)
  (lambda (st-list)
    (let ((new-st-list ((proc st) st-list)))
      (match new-st-list
        ((v _)
         ((proc (f v)) new-st-list))))))

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

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

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

(define (put v)
  "Sets st to v."
  (make-state
   (lambda (st-list)
     (list '() v))))

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