With newtype Cook m a = Cook (m a)

cook.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE NoDeriveAnyClass #-}

import Prelude
import Data.Kind
import Data.Coerce
import Control.Monad.Identity
import Control.Monad.Writer
import Control.Monad.State

program
  :: forall m . 
  ( DSL m
  , Fun (m ()) m
  , Fun (m (Expr Int)) m)
  -- ^ ISSUE: needs a constraint for every return type, e.g if any
  -- of the functions in body returns Expr Bool, then Fun (m (Expr Bool)) m
  -- needs to be added. And so on for every other return type.
  => m ()
program = do
  stm $ Stm (Var "a")
  -- this way:
  f <- fun $ \(a :: Expr a) (b :: Expr b) -> do
    -- function body is of the same monad `m` as the one in top level
    stm @m $ Stm (Var "b")
    return ()

  stm $ Stm $ Apply f (Var "c")
  pure ()

-- | An AST
data Expr a
  = Var String
  | Apply String (Expr a)
  | Function String [String] [Stm]
  | AnyExpr
  deriving Show

data Stm = Stm (Expr ())
  deriving Show

-- | The "effect"
class Monad m => DSL m where
  freshName :: m String -- generate fresh variable name
  stm :: Stm -> m () -- emit statement
  toAST :: m a -> m [Stm] -- turn code `m a` into [Stm] (without emitting it)
  fun :: Fun f m => f -> m String -- emit function f, return its name

-- | Helper class to convert literal haskell functions to the dsl
class Fun f m where
  mkFun :: DSL m => f -> [String] -> m ([String], [Stm])
instance Fun f m => Fun (Expr a -> f) m where -- COMMENT: kept the base case as is
  mkFun f acc = do
    name <- freshName -- [2]
    mkFun (f $ Var name) (name : acc)

-- * SampleImplRaw

type SampleImplRaw = StateT Int (Writer [Stm])

-- | A sample implementation
instance DSL SampleImplRaw where
  freshName = do
    n <- get
    put $ n + 1
    return $ "var" <> show n
  stm stm' = tell [stm']
  toAST m = do
    state0 <- get
    let ((_, state1), w) = runSampleImplRaw m state0
    put state1
    return w
  fun f = do
    (args, body) <- mkFun f []
    name <- freshName
    stm $ Stm $ Function name args body -- [1]
    return name

runSampleImplRaw :: SampleImplRaw a -> Int -> ((a, Int), [Stm])
runSampleImplRaw m s = runWriter $ runStateT m s

-- * SampleImpl

newtype Cook m a = Cook (m a)
  deriving (Functor, Applicative, Monad, Show)
instance MonadTrans Cook where -- COMMENT: needed lift in `DSL SampleImpl` below
  lift = Cook

type SampleImpl = Cook SampleImplRaw

instance DSL SampleImpl where
  freshName = lift freshName
  stm = lift . stm
  toAST (Cook m) = lift $ toAST m
  fun f = do
    (args, body) <- mkFun f []
    name <- freshName
    stm $ Stm $ Function name args body
    return name

instance (m0 ~ m) => Fun (Cook m0 a) (Cook m) where -- COMMENT: the new base case
  mkFun m args = do
    fname <- freshName
    body <- toAST m
    return (args, body)

runCook :: Cook SampleImplRaw a -> Int -> ((a, Int), [Stm])
runCook m = runSampleImplRaw (coerce m)

main :: IO ()
main = print $ runCook program 0