monad.nix

with import <nixpkgs> {};

let

# type monad = {
#   type m a
#   map : (a -> b) -> m a -> m b
#   apply : m (a -> b) -> m a -> m b
#   pure : a -> m a
#   bind : m a -> (a -> m b) -> m b
# }

# type monad-fail = {
#   inherit monad
#   fail : m a
# }

# type monad-log = {
#   inherit monad
#   log : a -> m null
# }

# type monad-lift = {
#   inherit monad
#   context : monad
#   lift : context.m a -> m a
# }

# monad combinators

# bind with arguments reversed
lift-monad = m: amb: ma: m.bind ma amb;

replicate-monad = m: n: a: if n == 0 then m.pure null else m.bind a (_: replicate-monad m (n - 1) a);

# data Maybe
just = a: { just = a; };
nothing = { nothing = null; };
maybe = m: on-just: on-nothing:
  if m ? just then on-just m.just else on-nothing;
maybe-monad = rec {
  map = f: ma: maybe ma (a: just (f a)) nothing;
  apply = mf: ma: maybe mf (f: map f ma) nothing;
  pure = just;
  bind = ma: amb: maybe ma amb nothing;
  fail = nothing;
};

# data Writer
make-writer = l: r: { log = l; result = r; };
writer-result = w: w.result;
writer-log = w: w.log;
map-log = f: w: make-writer (f (writer-log w)) (writer-result w);
writer = w: f: f (writer-log w) (writer-result w);
writer-monad = {
  map = f: wa: writer wa (l: r: make-writer l (f r));
  apply = wf: wa: writer wf (l1: f: writer wa (l2: a: make-writer (l1 ++ l2) (f a)));
  pure = a: make-writer [] a;
  bind = wa: awb: writer wa (l1: a: writer (awb a) (l2: b: make-writer (l1 ++ l2) b));
  log = msg: make-writer [ msg ] null;
};

# data MaybeT m a = MkMaybeT { runMaybeT :: m (Maybe a) }
maybe-t-monad = m: lib.fix (self: ({
  map = f: m.map (maybe-monad.map f);
  apply = mf: ma: m.apply (m.map maybe-monad.apply mf) ma;
  pure = a: m.pure (just a);
  bind = ma: amb: m.bind ma (maybe-a: maybe maybe-a amb (m.pure nothing));
  fail = m.pure nothing;
  context = m;
  lift = ma: m.map just ma;
} // (if m ? log then { log = a: self.lift (m.log a); } else {})));

# type effects a = {
#   pure: a;
# } | {
#   cont: b -> effects a;
#   effect-xxx: [ args ];
# }

map-effect = f: e: assert e ? cont; e // { cont = b: f (e.cont b); };

effects-monad = rec {
  map = f: e: if e ? pure then e // { pure = f e.pure; } else map-effect (map f) e;
  apply = ef: ea: if ef ? pure then map ef.pure ea else map-effect (ef': apply ef' ea) ef;
  pure = a: { pure = a; };
  bind = ea: aeb: if ea ? pure then aeb ea.pure else map-effect (ea': bind ea' aeb) ea;
};

yield = a: { cont = b: effects-monad.pure null; effect-yield = a; };
handle-yield = e:
  if e ? pure then effects-monad.pure (writer-monad.pure e.pure)
  else if e ? effect-yield then
    let rest = handle-yield (e.cont null); in
    effects-monad.map (map-log (l: [ e.effect-yield ] ++ l)) rest
  else
    map-effect handle-yield e;

fail = { cont = b: effects-monad.pure null; effect-fail = null; };
handle-fail = e:
  if e ? pure then effects-monad.pure (maybe-monad.pure e.pure)
  else if e ? effect-fail then
    effects-monad.pure nothing
  else
    map-effect handle-fail e;

schedule = { cont = b: effects-monad.pure null; effect-schedule = null; };
handle-schedule = ready-list: e:
  if e ? pure then
    if builtins.length ready-list == 0 then
      effects-monad.pure null
    else
      handle-schedule (builtins.tail ready-list) (builtins.head ready-list)
  else if e ? effect-schedule then
    if builtins.length ready-list == 0 then
      handle-schedule [] (e.cont null)
    else
      handle-schedule (builtins.tail ready-list ++ [ (e.cont null) ]) (builtins.head ready-list)
  else
    map-effect (handle-schedule ready-list) e;


in

let

test1 = let m = maybe-monad; in
  m.bind (m.pure 1) (one:
    m.bind (m.pure (one + 1)) (two:
      m.bind m.fail (wat:
        m.pure 42)));

test2 = let m = writer-monad; in
  m.bind (make-writer ["making 1"] 1) (one:
    m.bind (make-writer ["making 2"] (one + 1)) (two:
      m.bind (m.log "and...") (wat:
        m.pure 42)));

test3 = let m = writer-monad; mm = maybe-t-monad m; in
  m.bind (make-writer ["making 1"] 1) (one:
    m.bind
      (mm.bind (mm.pure 2) (two:
        mm.bind (mm.log "haha") (wat:
          mm.bind mm.fail (wat2:
            mm.pure 42))))
      (maybe-n: m.log (maybe maybe-n (n: "got ${toString n}") "got nothing")));

test3'5 = let m = maybe-t-monad writer-monad; in
  m.bind (m.log 1) (wat:
    m.bind (m.log 2) (wat2:
      m.bind m.fail (wat3:
        m.pure 42)));

test4 = let m = effects-monad; in
  m.bind (yield 1) (wat:
    m.bind (yield 2) (wat2:
      m.bind fail (wat3:
        m.pure 42)));

test5 = [
  (replicate-monad effects-monad 2 (effects-monad.bind (yield 2) (_: schedule)))
  (replicate-monad effects-monad 3 (effects-monad.bind (yield 3) (_: schedule)))
  (replicate-monad effects-monad 4 (effects-monad.bind (yield 4) (_: schedule)))
];

test-apply = let test = action: m: m.apply (m.pure (a: a + 1)) (m.bind action (_: m.pure 1)); in [
  (test maybe-monad.fail maybe-monad)
  (test (maybe-monad.pure 0) maybe-monad)
  (test (writer-monad.log "hi") writer-monad)
  (test (maybe-t-monad writer-monad).fail (maybe-t-monad writer-monad))
  (test ((maybe-t-monad writer-monad).log "hi") (maybe-t-monad writer-monad))
  (handle-yield (test (yield "hi") effects-monad))
];

in

[
  test1
  test2
  test3
  test3'5
  (handle-fail (handle-yield test4))
  (handle-yield (handle-fail test4))
  (handle-yield (handle-schedule test5 (effects-monad.pure null)))
  test-apply
]