Polykinded by adding the (x :: ()) kind

Polykinded.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
  , forall a . Fun (m a) m)
  => m ()
program = do
  stm $ Stm (Var "a")
  f <- fun $ \(a :: Expr a) (b :: Expr b) -> do
    stm @m $ Stm (Var "b")
    -- ^ ISSUE: commenting above line out (= not asking the monad in the
    -- function body to be the same as in outer scope) stops it from compiling
    stm $ 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)

-- * Base

type Base = StateT Int (Writer [Stm])

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

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

-- * Polykinded

newtype PolykindedRaw (x :: ()) a = PolykindedRaw (Base a)
  deriving (Functor, Applicative, Monad)

type Polykinded = PolykindedRaw '() :: Type -> Type

instance (m ~ m', p ~ p', p ~ '(), DSL (m '())) => Fun (m p a) (m' p') where
  mkFun m args = do
    fname <- freshName
    body <- toAST @(m' p') m
    return (args, body)

instance DSL Polykinded where
  freshName = PolykindedRaw $ do
    n <- get
    put $ n + 1
    return $ "var" <> show n
  stm = PolykindedRaw . stm
  toAST (PolykindedRaw m) = PolykindedRaw $ toAST m
  fun f = do
    (args, body) <- mkFun f []
    name <- freshName
    stm $ Stm $ Function name args body
    return name

runPolykinded :: Polykinded a -> Int -> ((a, Int), [Stm])
runPolykinded m s = runBase (coerce m) s

main = print $ runPolykinded (program) 0