Haskell Recursion Schemes: Label the tree with intermediate results

Question

Using cata I can fold an AST to a result. With Cofree I can store additional annotations on the AST. How can I take an AST and return an annotated AST with the results at each step of the way?

alg :: Term Result -> Result
alg = undefined

run :: Fix Term -> Result
run ast = cata alg ast

run' :: Fix Term -> Cofree Term Result
run' = ???

Answer 1

Does this modified algebra work?

alg' :: Term (Cofree Term Result) -> Cofree Term Result
alg' t = alg (fmap extract t) :< t  

run' :: Fix Term -> Cofree Term Result
run' ast = cata alg' ast

extract is from Control.Comonad. We are using it here with type Cofree Term Result -> Result. It simply returns the annotation at the root.

fmap extract :: Term (Cofree Term Result) -> Term Result lets us reuse our previous alg definition.

Answer 2

If you just need a simple catamorphism, you could use a function like cataM. This allows you to fold monadic values such that the actions are properly ordered. Plus, you don't need to write any boilerplate to wrap your F-algebra.

Then

alg :: Term Result -> Cofree Term Result
alg = undefined

run' :: Fix Term -> Cofree Term Result
run' = cataM alg