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 ints 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 = <fun>

(* 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 *)