Here's the solution:
let asker : char -> int Reader_custom.m =
fun (a : char) -> fun (env : char -> int) -> env a;;
let seed t = TR.monadize asker t;;
let v3 = TreeCont.monadize (fun a k ->
fun e -> k a (update_env e a)
) tree seed (fun a -> 0);;
What's going on here? Our distributed function takes a leaf element `a`,
a continuation `k` (which takes leaf elements and environments), and an
environment `e`, and returns the result of passing the leaf element and
an updated environment to `k`.
As always, the seed function operates on trees of leaf elements where
the distributed function operated on leaf elements. So the seed function
takes a tree and a environment, and returns whatever you want.
What we want is a tree-reader, that is, a function from environments to
trees. Once this gets distributed over the tree using the
`TreeCont.monadize` function, we'll have a tree-reader that operates on
the updated environment instead of the initial environment (which we
still have to supply---it's the final `fun a -> 0`).
How can we build such a tree-reader? The same way we formerly turned a tree
of `int`s and an int-to-b-reader function into a b-tree-reader: using the `TR.monadize` function.
And `v3` is just what we were looking for:
Node (Leaf 2, Node (Leaf 1, Node (Leaf 1, Leaf 2))).
Now all of this should be implementable in the "monads.ml" library too. But as
I said earlier, the encapsulation enforced by that library may make it somewhat harder to work with.
module T = Tree_monad;;
module C = Continuation_monad;;
module S = State_monad(struct type store = char -> int end);;
module R = Reader_monad(struct type env = char -> int end);;
module TC = T.T(C);;
module TS = T.T(S);;
module TR = T.T(R);;
let tree = Some (T.Node(T.Leaf '1', T.Node(T.Leaf '2', T.Node(T.Leaf '3', T.Leaf '1'))));;
# let v0 = TC.(run_exn (distribute (fun a ->
C.(shift (fun k -> k a >>= fun v -> unit (1+v)))
) tree)) (fun a -> 0);;
- : int = 4
# let v1 = TC.(run_exn (distribute (fun a ->
C.(shift (fun k -> k a >>= fun ka ->
unit (fun e -> let e_prev = ka e in (update_env e_prev a))))
) tree)) (fun t e -> e) (fun a -> 0);;
- : char -> int =
(* now
v1 '0';; ~~> 0
v1 '1';; ~~> 2
v1 '2';; ~~> 1
*)
# let annotater : char -> ('x, char) S.m =
fun a -> S.(puts (fun s -> (update_env s a)) >> unit a);;
# let v2 = TS.(run (distribute annotater tree)) (fun a -> 0);;
# let (t, env) = v2 in TR.(run (distribute (fun a -> R.asks (fun e -> e a)) t)) env;;
(* returns tree with leafs replaced with their numbers of occurrences *)
# let v3 = TC.(run_exn (distribute (fun a ->
C.(shift (fun k -> k a >>= fun ka ->
unit (fun e -> ka (update_env e a))))
) tree)) (fun t ->
TR.(run(distribute (fun a -> R.asks (fun e -> e a)) (Some t)))
) (fun a -> 0);;
(* also returns tree with leafs replaced with their numbers of occurrences *)