monad.clj

(ns user
  (:use clojure.set))
  
;; core

(defstruct monad :unit :bind :run :zero :plus)

(defmacro defmonad [name & fdecl]
  (let [[doc defs]
        (if (string? (first fdecl))
          [(first fdecl) (rest fdecl)]
          [nil fdecl])]
  `(def ~(with-meta name (into (or (meta name) {}) (if doc {:doc doc} {})))
        (struct-map monad ~@defs))))

;; generic dispatchers

(declare run-seq)

(deftype monad-expr [forall]
;; When Seqable becomes a protocol instead of an interface, we can pull this into
;; the seq monad implementation, where it belongs.
  :as this
  clojure.lang.Seqable
  (seq [] (run-seq this)))

(defmacro forall [[t] x]
  `(monad-expr (fn [~t] ~x)))

(defn at [m x]
  (if (monad-expr? x)
    ((:forall x) m)
    x))

(defn unit [x]
  (forall [t] ((:unit t) x)))

(defn bind [m f]
  (forall [t] ((:bind t) (at t m) (fn [x] (at t (f x))))))

(defn run [t & xs]
  (let [m (last xs)
        xs (butlast xs)]
    (apply (:run t) (concat xs [(at t m)]))))

(defn zero []
  (forall [t] (:zero t)))

(defn plus [m1 m2]
  (forall [t] ((:plus t) (at t m1) (at t m2))))

;; helpers

(defmacro >>= [& body]
  (reduce (fn [x y]
            (let [[v m] (if (vector? y) y ['% y])]
              `(bind ~m (fn [~v] ~x))))
          (last body)
          (reverse (butlast body))))

(defmacro >> [& body]
  `(>>= ~@(butlast body) (unit ~(last body))))

(defn either [& ms]
  (reduce #(plus %2 %1) (reverse ms)))

(defn mfilter [mpred coll]
  (if (seq coll)
      (>> [b (mpred (first coll))]
          [xs (mfilter mpred (rest coll))]
          (if b
            (lazy-seq (cons (first coll) xs))
            xs))
      (>> ())))

;; monad implementations

(defmonad identity-monad
  "The identity monad"
  :unit identity
  :bind (fn [m f] (f m))
  :run identity)

(def run-identity (partial run identity-monad))

(defmonad maybe-monad
  "The maybe/option monad."
  :unit identity
  :bind (fn [m f] (when-not (nil? m) (f m)))
  :zero nil
  :plus (fn [m1 m2] (if (nil? m1) m2 m1))
  :run identity)

(def run-maybe (partial run maybe-monad))

(defn flip [f]
  (fn [x y] (f y x)))

(defmonad seq-monad
  "The sequence monad."
  :unit list
  :bind (flip mapcat)
  :zero ()
  :plus concat
  :run identity)

(def run-seq (partial run seq-monad))

(defmonad set-monad
  "The set monad."
  :unit (fn [x] #{x})
  :bind (fn [m f] (apply union (map f m)))
  :zero #{}
  :plus union
  :run identity)

(def run-set (partial run set-monad))

(defmonad state-monad
  "The state monad."
  :unit (fn [x] (fn [s] [x s]))
  :bind (fn [m f] (fn [s]
                    (let [[x s-new] (m s)]
                      ((f x) s-new))))
  :run (fn [s m] (m s)))

(def run-state (partial run state-monad))

(def get-state (fn [s] [s s]))

(defn set-state [x] (fn [s] [nil x]))

(defn update-state [f] (fn [s] [s (f s)]))

;; examples

(run-seq
  (>> [x (range 100)]
       [y (range 100)]
       (+ x y)))

(run-seq
  (>> [x (>> [y (range 100)] (* y y))]
      (+ x x)))

(run-seq
  (>> [x (either (unit 1) (unit 2) (unit 3))]
      (* x x)))

(run-seq (mfilter (constantly [false true]) [1 2 3 4]))

(run-set
  (>> [x (range 100)]
      [y (range 100)]
      (+ x y)))

(run-state 100
  (>> [x get-state]
      (set-state 42)
      [y (update-state #(+ % 8))]
      [z get-state]
      [x y z]))

(run-state 0 (mfilter (fn [x] (>> (update-state inc) (= (mod % 5) 0))) (range 100)))

(let [xs [1 2]]
  (run-maybe
    (>> [x (first xs)]
        [y (second xs)]
        (+ x y))))

(let [xs [1]]
  (run-maybe
    (>> [x (first xs)]
        [y (second xs)]
        (+ x y))))

(run-identity (>>= 42 (- % 2) (* % %)))