gyehoek-hs/app/Gyehoek/QBE.hs

{-# LANGUAGE RequiredTypeArguments #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo #-}
module Gyehoek.QBE
  ( Unique
  , runUnique
  , unique
  , module QBE
  , render
  )
  where

import Language.QBE as QBE
import Effectful.State.Dynamic
import Effectful.Dispatch.Dynamic
import Effectful
import Numeric.Natural
import Data.String (IsString(fromString))
import Prettyprinter (Pretty(pretty), layoutPretty, defaultLayoutOptions)
import Data.Text (Text)
import Prettyprinter.Render.Text (renderStrict)
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L
import Data.Void (Void)
import Data.Char (isAlpha, isAlphaNum)
import Control.Lens.Wrapped
import Data.Functor.Contravariant (Predicate(Predicate))
import qualified Data.Text as T
import Data.Functor
import Data.List (List)
import Data.Foldable (fold)
import Data.Maybe (isJust, fromMaybe)
import Control.Monad.Fix (MonadFix(mfix))


render :: Pretty a => a -> Text
render = renderStrict . layoutPretty defaultLayoutOptions . pretty

data Unique :: Effect where
  Unique :: Unique m (Ident s)

type instance DispatchOf Unique = Dynamic

unique :: forall s es. Unique :> es => Eff es (Ident s)
unique = send Unique

runUnique :: Eff (Unique : es) a -> Eff es a
runUnique = reinterpret (evalStateLocal (0 :: Natural)) \_ Unique ->
  state \n -> (Ident . fromString $ '.' : show n, succ n)

-- class SigilChar (s :: Sigil) where
--   sigilChar' :: Proxy s -> Char
-- instance SigilChar AggregateTy where sigilChar' _ = ':'
-- instance SigilChar Global      where sigilChar' _ = '$'
-- instance SigilChar Temporary   where sigilChar' _ = '%'
-- instance SigilChar Label       where sigilChar' _ = '@'

-- sigilChar :: forall (s :: Sigil) -> SigilChar s => Char
-- sigilChar s = sigilChar' (Proxy @s)



type P = Parsec Void Text

sc :: P ()
sc = L.space hspace1 (L.skipLineComment "#") empty

lexeme :: P a -> P a
lexeme = L.lexeme sc

symbol :: Text -> P Text
symbol = L.symbol sc

infixr 8 .:
(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
(.:) f g x y = f (g x y)

rawIdent :: P QBE.RawIdent
rawIdent = (fromString .: (:) <$> lead <*> trail) <?> "ident"
  where
    lead = satisfy \x -> isAlpha x || (x=='.') || (x=='_')
    trail = fmap T.unpack . takeWhileP Nothing $ \x ->
      isAlphaNum x || (x=='.') || (x=='_')

class ParseIdent (s :: Sigil) where
  ident :: P (QBE.Ident s)

rawIdentWithSigil :: Char -> P (Ident t)
rawIdentWithSigil c = Ident <$> lexeme (char c *> rawIdent)

instance ParseIdent AggregateTy where ident = rawIdentWithSigil ':'
instance ParseIdent Global      where ident = rawIdentWithSigil '$'
instance ParseIdent Temporary   where ident = rawIdentWithSigil '%'
instance ParseIdent QBE.Label   where ident = rawIdentWithSigil '@'

const :: P QBE.Const
const = cint <|> csingle <|> cdouble <|> cglobal <?> "const"
  where
    cint = CInt <$> lexeme (L.signed empty L.decimal) <?> "integer"
    csingle = empty <?> "single-precision float"
    cdouble = empty <?> "double-precision float"
    cglobal = CGlobal <$> ident <?> "global symbol"

val :: P QBE.Val
val = vconst <|> vtemp <?> "val"
  where
    vconst = ValConst <$> Gyehoek.QBE.const
    vtemp = ValTemporary <$> ident <?> "temporary symbol"

assignment :: P QBE.Assignment
assignment =
  Assignment <$> ident <*> (char '=' *> basety)

basety :: P QBE.BaseTy
basety = lexeme $ char 'w' $> Word
              <|> char 'l' $> Long
              <|> char 's' $> Single
              <|> char 'd' $> Double

abity :: P AbiTy
abity = AbiBaseTy <$> basety
    <|> AbiAggregateTy <$> ident

binaryOp :: P QBE.BinaryOp
binaryOp = lexeme $ "add" $> Add
                <|> "sub" $> Sub
                <|> "mul" $> Mul
                <|> "div" $> Div Signed

comma :: P a -> P a
comma p = symbol "," *> p

inst :: P QBE.Inst
inst = try binaryOpInst <|> negInst <?> "inst"
  where
    binaryOpInst =
      BinaryOp
        <$> assignment
        <*> binaryOp
        <*> val <*> comma val
    negInst = Neg <$> assignment <*> (symbol "neg" *> val)

jump :: P QBE.Jump
jump = jmp <|> jnz <|> ret <|> hlt <?> "jump"
  where
    jmp = symbol "jmp" *> (Jmp <$> ident)
    jnz = symbol "jnz" *> (Jnz <$> val <*> ident <*> comma ident)
    ret = symbol "ret" *> (Ret <$> optional val)
    hlt = empty

nl :: P ()
nl = void (some (newline *> sc)) <?> "newline"

phi :: P QBE.Phi
phi = empty

block :: P QBE.Block
block = Block
  <$> (ident <* nl)
  <*> sepBy phi nl
  <*> sepBy inst nl
  <*> jump

sepByTry :: MonadParsec e s m => m a -> m sep -> m (List a)
sepByTry p sep = do
  x <- p
  xs <- many (try $ sep *> p)
  pure (x:xs)

params :: P (Maybe (Ident Temporary), List Param, Variadic)
params = between (symbol "(") (symbol ")") do
  e <- optional env
  ps <- optional . try $ do
    commaIf (isJust e)
    sepByTry reg (symbol ",")
  v <- optional do
    commaIf (isJust e || isJust ps)
    variadic
  pure (e, fromMaybe [] ps, fromMaybe NoVariadic v)
  where
    continue d m = fromMaybe d <$> optional (symbol "," *> m)
    commaIf True = void $ symbol ","
    commaIf False = pure ()
    env = symbol "env" *> ident @Temporary <?> "environment parameter"
    reg = Param <$> abity <*> ident <?> "regular parameter"
    variadic = symbol "..." $> Variadic <?> "variadic parameter"

funcdef :: P QBE.FuncDef
funcdef = do
  linkages <- _
  symbol "function"
  returnTy <- optional abity
  name <- ident @Global
  (env,params,variadic) <- params
  code <- _
  pure $ FuncDef linkages returnTy name env params variadic code