gyehoek-hs/app/Gyehoek/ANF.hs

{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-orphans -Wno-unused-matches -Wno-missing-signatures #-}
{- HLINT ignore "Avoid lambda using `infix`" -}
module Gyehoek.ANF
  (toANF, lower)
  where

import Data.Text (Text)
import Effectful
import Gyehoek.QBE qualified as QBE
import Data.List (List)
import Data.Text.IO qualified as TIO
import Control.Lens
import Data.Generics.Labels
import Data.Vector.Strict (Vector)
import Data.Function (fix)
import Effectful.Writer.Static.Local
import Gyehoek.Syntax qualified as Lam
import Gyehoek.Syntax (Name, Prim(..), Val(..))
import Gyehoek.GenSym
import Control.Monad.Cont
import Data.Foldable
import Data.List.NonEmpty (NonEmpty((:|)), toList)
import Gyehoek.QBE (FuncDef(FuncDef))
import Data.Foldable1
import qualified Data.Text as T
import Data.String (fromString)


-- data Val
--   = ValInt Int
--   | ValNil
--   | ValPrim Prim
--   | ValLambda (List Name) Exp
--   | ValVar Name
--   deriving (Show)

data Exp
  = ExpLetApply Name Val (List Val) Exp
  | ExpProgn (List Exp)
  | ExpVal Val
  deriving (Show)



blah' :: List ((a -> r) -> r) -> (List a -> r) -> r
blah' = go [] where
  go acc []     k = k (reverse acc)
  go acc (f:fs) k = f \a -> go (a:acc) fs k

-- 뻘짓이어라
blah :: forall t a r. Foldable t => t ((a -> r) -> r) -> (List a -> r) -> r
blah xs k =
  foldr (\i l acc -> i \x -> l (x:acc)) (k . reverse) xs []

toANF'
  :: forall es. GenSym :> es
  => Lam.Exp
  -> (Val -> Eff es Exp)
  -> Eff es Exp

toANF' (Lam.ExpVal v) k = k v

toANF' (Lam.ExpApply f xs) k =
  blah (toANF' <$> (f:|xs)) \(f':xs') -> do
    r <- gensym
    ExpLetApply r f' xs' <$> k (ValVar r)
  where
    allToANF' es k = traverse (\e -> toANF' e k) es
    -- blah = traverse g (f :| xs)
    -- g :: Lam.Exp -> Eff es Val
    -- g = ContT . toANF'

toANF' e k = _

toANF e = toANF' e (pure . ExpVal)

expr =
  Lam.ExpApply (Lam.ExpVal (ValPrim PrimAdd))
    [ Lam.ExpVal (ValInt 1)
    , Lam.ExpApply
      (Lam.ExpVal (ValPrim PrimMul))
      [ Lam.ExpVal (ValInt 2)
      , Lam.ExpVal (ValInt 4)
      ]
    ]

expr2 =
  Lam.ExpApply (Lam.ExpVal (ValPrim PrimAdd))
    [ Lam.ExpApply (Lam.ExpVal (ValPrim PrimMul)) [Lam.ExpVal (ValInt 1)]
    , Lam.ExpApply (Lam.ExpVal (ValPrim PrimMul)) [Lam.ExpVal (ValInt 2)]
    , Lam.ExpApply (Lam.ExpVal (ValPrim PrimMul)) [Lam.ExpVal (ValInt 3)]
    ]



type CodeGen = Writer (Vector QBE.Inst)

emit :: CodeGen :> es => QBE.Inst -> Eff es ()
emit = tell . pure

instance Semigroup QBE.Program where
  QBE.Program ts ds fs <> QBE.Program ts' ds' fs' =
    QBE.Program (ts <> ts') (ds <> ds') (fs <> fs')

instance Monoid QBE.Program where
  mempty :: QBE.Program
  mempty = QBE.Program mempty mempty mempty

funcdef :: QBE.Ident QBE.Global -> [QBE.Param] -> NonEmpty QBE.Block -> FuncDef
funcdef name ps = QBE.FuncDef mempty Nothing name Nothing ps QBE.NoVariadic

prims :: QBE.Program
prims = QBE.Program mempty mempty primfns where
  primfns = [ mkArith "plus" QBE.Add
            , mkArith "star" QBE.Mul
            , mkArith "_" QBE.Sub
            , mkArith "slash" (QBE.Div QBE.Signed)
            ]
  mkArith name bop =
    funcdef name
    [ QBE.Param (QBE.AbiBaseTy QBE.Long) "x"
    , QBE.Param (QBE.AbiBaseTy QBE.Long) "y"
    ]
    [ QBE.Block "start" []
      [ QBE.BinaryOp ("r" QBE.:= QBE.Long) bop
        (QBE.ValTemporary "x") (QBE.ValTemporary "y")
      ]
      (QBE.Ret (Just (QBE.ValTemporary "r")))
    ]

data BlockBuilder
  = Emit (Vector QBE.Inst) !BlockBuilder
  | Exit QBE.Jump
  deriving (Show)

instance Each BlockBuilder BlockBuilder QBE.Inst QBE.Inst where
  each k (Emit is bb) = Emit <$> traverse k is <*> each k bb
  each k (Exit j) = pure (Exit j)

evalBlockBuilder :: BlockBuilder -> (Vector QBE.Inst, QBE.Jump)
evalBlockBuilder (Emit is bb) = evalBlockBuilder bb & _1 <>:~ is
evalBlockBuilder (Exit j)     = ([],j)

buildBlock :: QBE.Ident QBE.Label -> BlockBuilder -> QBE.Block
buildBlock n bb = QBE.Block n [] (is ^.. each) j
  where (is,j) = evalBlockBuilder bb

lowerName :: Name -> QBE.Ident t
lowerName = fromString . T.unpack

lowerPrim :: Prim -> _ -> _
lowerPrim PrimAdd k = k $ QBE.ValGlobal "plus"
lowerPrim PrimMul k = k $ QBE.ValGlobal "star"
lowerPrim PrimSub k = k $ QBE.ValGlobal "_"
lowerPrim PrimDiv k = k $ QBE.ValGlobal "slash"
lowerPrim p k = _

lowerVal
  :: forall es. (GenSym :> es)
  => Val
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder

lowerVal (ValInt n) k = k . QBE.ValConst . QBE.CInt . fromIntegral $ n
lowerVal (ValVar x) k = k . QBE.ValTemporary . lowerName $ x
lowerVal (ValPrim p) k = lowerPrim p k
lowerVal _ k = _

lower'
  :: forall es. (GenSym :> es)
  => Exp
  -> (QBE.Val -> Eff es BlockBuilder)
  -> Eff es BlockBuilder

lower' (ExpVal v) k = lowerVal v k

lower' (ExpLetApply r f xs e) k =
  blah (lowerVal @es <$> (f:|xs)) \(f':xs') -> do
    Emit [ QBE.Call
            (Just (lowerName r, QBE.AbiBaseTy QBE.Long))
            f'
            Nothing
            (QBE.Arg (QBE.AbiBaseTy QBE.Long) <$> xs')
            []
          ]
      <$> lower' e k

lower' _ k = _

lower :: GenSym :> es => QBE.Ident QBE.Label -> Exp -> Eff es QBE.Block
lower n e = buildBlock n <$> lower' e (pure . Exit . QBE.Ret . Just)

wrapProgram :: Foldable1 t => t QBE.Block -> QBE.Program
wrapProgram bs = QBE.Program [] [] [main] where
  main = QBE.FuncDef [QBE.Export]
    (Just (QBE.AbiBaseTy QBE.Word))
    "main" Nothing [] QBE.NoVariadic (toNonEmpty bs)