ana.hs

#!/usr/bin/env cabal
{- cabal:
build-depends:
  , base
  , recursion-schemes > 5
  , free
  , comonad
-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wall #-}

module Main where

    import Control.Comonad
    import Control.Comonad.Cofree
    import Control.Comonad.Trans.Cofree qualified as CofreeF
    import Data.Functor.Foldable
    import Data.Functor.Classes
    import Data.Functor.Foldable.TH

    -- A simple arithmetic language

    data Expr = Expr :+ Expr | I Int
      deriving (Eq, Show)

    infixl 6 :+

    makeBaseFunctor ''Expr

    deriving instance Show a => Show (ExprF a)

    instance Show1 ExprF where
      liftShowsPrec p l n = \case
        IF i -> showsPrec n i
        a :+$ b -> showParen True (p n a . showString " :$+ " . p n b)

    example :: Expr
    example = I 6 :+ I 5 :+ I 4

    type Annotated t = Cofree (Base t)

    annotateIn ::
      (Recursive t, Comonad (Cofree (Base t))) =>
      (Base t b -> b) ->
      t ->
      Cofree (Base t) b
    annotateIn f = cata $ \t -> (f $ fmap extract t) :< t

    anno :: ExprF Int -> Int
    anno (IF i) = i
    anno (a :+$ b) = a + b

    depth :: (Int, Expr) -> Base (Cofree ExprF Int) (Int, Expr)
    depth (n, e) = n CofreeF.:< x
      where x = case e of
              I i -> IF i
              a :+ b -> (succ n, a) :+$ (succ n, b)

    gyah :: Annotated Expr Int
    gyah = annWithDepth example

    strip :: Annotated Expr Int -> Expr
    strip a = ana unwrap a

    annWithDepth :: Expr -> Annotated Expr Int
    annWithDepth e = ana depth (0, e)

    main :: IO ()
    main = do
      print example
      print (annotateIn anno example)
      print gyah
      print (strip gyah)