Printing free state monad as Haskell code

169 Views Asked by At

I've read articles about extending simple-reflect to functions and about representing the State monad as a free monad (here), so now I'd like to try to print some state monads as executable Haskell code, so that for example the monad

someComputation :: Num a => State a ()
someComputation = do
  put 42

gets printed as StateT {runStateT = \x0 -> Identity ((),42)}. For such simple cases, I'm already getting there, but for code involving binds, the output I get is far from my intuition for what should get printed, which is why I'm at a loss for how to edge closer to a solution.

My understanding from the article I linked to above is that Free corresponds to join and Pure to pure, so I thought that I could simply map the constructors to join and pure in the output, insert fmap in the right places, and generate unique variable names. Should that work?

For the monad

someHarderComputation :: Num a => State a ()
someHarderComputation = do
    i <- get
    put $ i + 1

I get the actual output

StateF {runStateF = \ x1 -> (Free \ x0 -> ("Pure (),x1)",x1 + 1)} $ 1

though my desired output (for the finished implementation) is

join ((\x1 -> StateT {runStateT = \x2 -> Identity ((),x1 + 1)}) <$> StateT {runStateT = \x0 -> Identity (x0,x0)}).

Full code (requires simple-reflect package):

{-# LANGUAGE DerivingStrategies #-}
import Control.Arrow
import Control.Monad.Free
import qualified Control.Monad.Trans.State as S
import Data.Functor.Classes
import Data.List
import Debug.SimpleReflect.Expr

newtype StateF s a = StateF { runStateF :: s -> (a, s) }
  deriving stock Functor

deriving instance (ExprArg s, Show s, Show a) => Show (StateF s a)

instance (ExprArg s, Show s) => Show1 (StateF s) where
    liftShowsPrec sp _ d (StateF m) = \str -> show ((first (($ str) . sp 0) . m))

getF :: StateF s s
getF = StateF $ \s -> (s, s)

putF :: s -> StateF s ()
putF s = StateF $ const ((), s)

type State s = Free (StateF s)

get :: State s s
get = liftF getF

put :: s -> State s ()
put = liftF . putF

someComputation :: Num a => State a ()
someComputation = do
  i <- get
  put $ i + 1

instance (Show a, ExprArg a, Show r) => Show (a -> r) where
    show f = "\\ " <> show v <> " -> " <> show (f v)
      where v = S.evalState exprArg vars
            dummy = S.evalState exprArg $ repeat "_"
            vars = supply \\ tokenize (show $ f dummy)
            supply = [ "x" <> show i | i <- [0 :: Int ..]]
            tokenize "" = []
            tokenize s = case lex s of
                (x,s') : _ -> x : tokenize s'
                _ -> error "Assertion error"

class ExprArg a where
    exprArg :: S.State [String] a

instance ExprArg Expr where
    exprArg = do
        vvs <- S.get
        S.put $ tail vvs
        pure (var $ head vvs)

instance ExprArg () where
    exprArg = pure ()

instance (ExprArg a, ExprArg b) => ExprArg (a, b) where
    exprArg = (,) <$> exprArg <*> exprArg

instance (ExprArg a, ExprArg b, ExprArg c) => ExprArg (a, b, c) where
    exprArg = (,,) <$> exprArg <*> exprArg <*> exprArg

printState :: (ExprArg s, Show s, Show a) => State s a -> s -> String
printState (Pure x) s = "pure (" <> show x <> "," <> show s <> ")"
printState (Free m) s =
    show m <> " $ " <> show s

main :: IO ()
main = do
    putStrLn $ printState someComputation (1 :: Expr)
0

There are 0 best solutions below