Haskell: Labeling an AST with type information using Algorithm W

Question

We have an AST definition:

data Term a 
  = Lam String a
  | App a a
  | Var String
  deriving(Read,Show,Eq,Functor,Foldable,Traversable)

And an F-Algebra for the type inference:

type Wrapped m a = Enviroment -> m a
type Result m = Wrapped (Type,Substitution)
w ::
  (MonadState Int, MonadError TypeError) 
  => Term (Result m)
  -> Result m
w term env = ...

We can get the result of running the inference using cata:

infer :: 
  (MonadState Int, MonadError TypeError) 
  => Fix Term 
  -> Result m
infer ast = cata hm ast

But now I want the result to be the original AST annotated with type information for each expression, So now infer' :: Fix Term -> Wrapped (Labeled Term Type).

1. What data structure should I use to annotate the tree (Cofree,Product,Fix with a custom Label)?
2. How can I implement this using recursion schemes, without modifying the original w function?

Answer 1

This answer does modify the function w, but it still aims to keep the "workhorse" functions disentangled from the recursion-schemes machinery.

Let's keep the Term type as is, and assume we have a type E for the environment that is calculated downwards, and a type R for the final annotation that is calculated upwards from the leaves.

Let's also assume that we have these two functions:

calcEnv :: E -> Term x -> E -- calculates the environment which will be passed downwards

calcResult :: E -> Term R -> IO R -- effectfully calculates the result flowing upwards

I'm using IO as the monad for simplicity.

(Notice that I'm assuming that "calculating the environment" can't have effects. I this isn't the case then this solution won't do.)

We work in two phases. First we construct a tree in which the nodes have been annotated with their environments. We use an anamorphism, instead of the "trick" of returning a function from a catamorphism.

import qualified Control.Comonad.Trans.Cofree as COFREEF

annotate :: E -> Fix Term -> Cofree Term E
annotate = curry (ana coalg)
    where
    coalg :: (E, Fix Term) -> COFREEF.CofreeF Term E (E, Fix Term)
    coalg (environment, Fix term) =
        let environment' = calcEnv environment term
        in  environment COFREEF.:< fmap ((,) environment') term

(Keep in mind that type instance Base (Cofree f a) = CofreeF f a. That's where the COFREEF.:< comes from. It's basically a pair of a pure value and another value wrapped in the functor.)

And in the next phase we effectfully consume the annotated tree from the leaves to produce the final result —a tree with R annotations:

calculate :: Cofree Term E -> IO (Cofree Term R)
calculate = cata alg
    where
    alg :: COFREEF.CofreeF Term E (IO (Cofree Term R)) -> IO (Cofree Term R)
    alg (environment COFREEF.:< termio) = do
        term :: Term (Cofree Term R) <- sequenceA termio
        result :: R <- calcResult environment (fmap extract term)
        return $ result :< term

I did in two phases because I was having trouble combining the "returning a function" trick with returning an annotated tree.

---

An anamorphism followed by a catamorphism is known as hylomorphism. We could define the composed function using hylo, like this:

together :: E -> Fix Term -> IO (Cofree Term R)
together = curry (hylo alg coalg)
    where
    coalg (environment, Fix term) = ...
    alg (environment COFREEF.:< termio) = ...

---

You can put together calcEnv and calcResult in the form of the original algebra like this:

w :: Term (E -> IO R) -> E -> IO R
w term environment = do
    let downwards = calcEnv environment term
    tr :: Term R <- sequenceA $ fmap ($ downwards) term
    calcResult environment tr