Phased AST

PhasedAST.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

module Main where

data Exp (p :: Phase)
  = (HasInfix p ~ 'True) => Infix (Anno p) !Op (Exp p) (Exp p)
  | (HasStack p ~ 'False) => Call (Anno p) (Exp p) (Exp p)
  | Lit (Anno p) !Integer
  | Var (Anno p) !Name
  | (HasStack p ~ 'True) => CallStack (Anno p) !Name
  | (HasStack p ~ 'True) => Seq (Anno p) (Exp p) (Exp p)

data Op = AddOp | MulOp

type Name = String

data Phase
  = Parsed
  | Checked
  | Desugared
  | Generated

type family Anno (p :: Phase) :: * where
  Anno 'Parsed = ()
  Anno 'Checked = Type
  Anno 'Desugared = Type
  Anno 'Generated = Type

type family HasInfix (p :: Phase) :: Bool where
  HasInfix 'Parsed = 'True
  HasInfix 'Checked = 'True
  HasInfix 'Desugared = 'False
  HasInfix 'Generated = 'False

type family HasStack (p :: Phase) :: Bool where
  HasStack 'Parsed = 'False
  HasStack 'Checked = 'False
  HasStack 'Desugared = 'False
  HasStack 'Generated = 'True

data Type
  = IntType
  | FunType !Type !Type

typecheck :: Exp 'Parsed -> Exp 'Checked
typecheck = \ case
  Infix _anno op lhs rhs -> Infix IntType op (typecheck lhs) (typecheck rhs)
  Call _anno fun arg -> Call IntType (typecheck fun) (typecheck arg)
  Lit _anno n -> Lit IntType n
  Var _anno name -> Var IntType name

desugar :: Exp 'Checked -> Exp 'Desugared
desugar = \ case
  Infix type_ op lhs rhs -> Call
    type_
    (Call
      (FunType IntType type_)
      (Var
        (FunType IntType (FunType IntType type_))
        (opName op))
      (desugar lhs))
    (desugar rhs)
  Call type_ fun arg -> Call type_ (desugar fun) (desugar arg)
  Lit type_ n -> Lit type_ n
  Var type_ name -> Var type_ name

generate :: Exp 'Desugared -> Exp 'Generated
generate = \ case
  Call type_ (Var varType name) arg -> Seq type_ (generate arg) (CallStack varType name)
  Call type_ fun arg -> Seq type_ (generate arg) (generate fun)
  Lit type_ n -> Lit type_ n
  Var type_ name -> Var type_ name

opName :: Op -> Name
opName AddOp = "__add__"
opName MulOp = "__mul__"

codegen :: Exp 'Generated -> String
codegen = \ case
  Lit _type n -> "push " ++ show n
  Var _type name -> "push " ++ name
  CallStack _type name -> "call " ++ name
  Seq _type a b -> codegen a ++ "\n" ++ codegen b

main :: IO ()
main = putStrLn $ codegen $ generate $ desugar $ typecheck 
  $ Infix () AddOp
    (Infix () MulOp (Lit () 2) (Lit () 3))
    (Infix () MulOp (Lit () 5) (Var () "x"))