{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Frontend AST
-- | This module is split out so that the bakend/IR need not depend on
-- everything in 'AST'.
module Kempe.AST.Size ( KempeTy (..)
                      , StackType (..)
                      , MonoStackType
                      , BuiltinTy (..)
                      , ABI (..)
                      , prettyMonoStackType
                      -- * Sizing bits
                      , SizeEnv
                      , Size
                      , cSize
                      , size
                      , size'
                      , sizeStack
                      ) where

import           Control.DeepSeq (NFData)
import           Data.Int        (Int64)
import qualified Data.IntMap     as IM
import           Data.Monoid     (Sum (..))
import qualified Data.Set        as S
import           GHC.Generics    (Generic)
import           Kempe.Name
import           Kempe.Unique
import           Prettyprinter   (Doc, Pretty (pretty), parens, sep, (<+>))

data KempeTy a = TyBuiltin a BuiltinTy
               | TyNamed a (TyName a)
               | TyVar a (Name a)
               | TyApp a (KempeTy a) (KempeTy a) -- type applied to another, e.g. Just Int
               deriving ((forall x. KempeTy a -> Rep (KempeTy a) x)
-> (forall x. Rep (KempeTy a) x -> KempeTy a)
-> Generic (KempeTy a)
forall x. Rep (KempeTy a) x -> KempeTy a
forall x. KempeTy a -> Rep (KempeTy a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (KempeTy a) x -> KempeTy a
forall a x. KempeTy a -> Rep (KempeTy a) x
$cto :: forall a x. Rep (KempeTy a) x -> KempeTy a
$cfrom :: forall a x. KempeTy a -> Rep (KempeTy a) x
Generic, KempeTy a -> ()
(KempeTy a -> ()) -> NFData (KempeTy a)
forall a. NFData a => KempeTy a -> ()
forall a. (a -> ()) -> NFData a
rnf :: KempeTy a -> ()
$crnf :: forall a. NFData a => KempeTy a -> ()
NFData, a -> KempeTy b -> KempeTy a
(a -> b) -> KempeTy a -> KempeTy b
(forall a b. (a -> b) -> KempeTy a -> KempeTy b)
-> (forall a b. a -> KempeTy b -> KempeTy a) -> Functor KempeTy
forall a b. a -> KempeTy b -> KempeTy a
forall a b. (a -> b) -> KempeTy a -> KempeTy b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> KempeTy b -> KempeTy a
$c<$ :: forall a b. a -> KempeTy b -> KempeTy a
fmap :: (a -> b) -> KempeTy a -> KempeTy b
$cfmap :: forall a b. (a -> b) -> KempeTy a -> KempeTy b
Functor, KempeTy a -> KempeTy a -> Bool
(KempeTy a -> KempeTy a -> Bool)
-> (KempeTy a -> KempeTy a -> Bool) -> Eq (KempeTy a)
forall a. Eq a => KempeTy a -> KempeTy a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: KempeTy a -> KempeTy a -> Bool
$c/= :: forall a. Eq a => KempeTy a -> KempeTy a -> Bool
== :: KempeTy a -> KempeTy a -> Bool
$c== :: forall a. Eq a => KempeTy a -> KempeTy a -> Bool
Eq, Eq (KempeTy a)
Eq (KempeTy a)
-> (KempeTy a -> KempeTy a -> Ordering)
-> (KempeTy a -> KempeTy a -> Bool)
-> (KempeTy a -> KempeTy a -> Bool)
-> (KempeTy a -> KempeTy a -> Bool)
-> (KempeTy a -> KempeTy a -> Bool)
-> (KempeTy a -> KempeTy a -> KempeTy a)
-> (KempeTy a -> KempeTy a -> KempeTy a)
-> Ord (KempeTy a)
KempeTy a -> KempeTy a -> Bool
KempeTy a -> KempeTy a -> Ordering
KempeTy a -> KempeTy a -> KempeTy a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (KempeTy a)
forall a. Ord a => KempeTy a -> KempeTy a -> Bool
forall a. Ord a => KempeTy a -> KempeTy a -> Ordering
forall a. Ord a => KempeTy a -> KempeTy a -> KempeTy a
min :: KempeTy a -> KempeTy a -> KempeTy a
$cmin :: forall a. Ord a => KempeTy a -> KempeTy a -> KempeTy a
max :: KempeTy a -> KempeTy a -> KempeTy a
$cmax :: forall a. Ord a => KempeTy a -> KempeTy a -> KempeTy a
>= :: KempeTy a -> KempeTy a -> Bool
$c>= :: forall a. Ord a => KempeTy a -> KempeTy a -> Bool
> :: KempeTy a -> KempeTy a -> Bool
$c> :: forall a. Ord a => KempeTy a -> KempeTy a -> Bool
<= :: KempeTy a -> KempeTy a -> Bool
$c<= :: forall a. Ord a => KempeTy a -> KempeTy a -> Bool
< :: KempeTy a -> KempeTy a -> Bool
$c< :: forall a. Ord a => KempeTy a -> KempeTy a -> Bool
compare :: KempeTy a -> KempeTy a -> Ordering
$ccompare :: forall a. Ord a => KempeTy a -> KempeTy a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (KempeTy a)
Ord) -- questionable eq instance but eh

data StackType b = StackType { StackType b -> Set (Name b)
quantify :: S.Set (Name b)
                             , StackType b -> [KempeTy b]
inTypes  :: [KempeTy b]
                             , StackType b -> [KempeTy b]
outTypes :: [KempeTy b]
                             } deriving ((forall x. StackType b -> Rep (StackType b) x)
-> (forall x. Rep (StackType b) x -> StackType b)
-> Generic (StackType b)
forall x. Rep (StackType b) x -> StackType b
forall x. StackType b -> Rep (StackType b) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall b x. Rep (StackType b) x -> StackType b
forall b x. StackType b -> Rep (StackType b) x
$cto :: forall b x. Rep (StackType b) x -> StackType b
$cfrom :: forall b x. StackType b -> Rep (StackType b) x
Generic, StackType b -> ()
(StackType b -> ()) -> NFData (StackType b)
forall b. NFData b => StackType b -> ()
forall a. (a -> ()) -> NFData a
rnf :: StackType b -> ()
$crnf :: forall b. NFData b => StackType b -> ()
NFData, StackType b -> StackType b -> Bool
(StackType b -> StackType b -> Bool)
-> (StackType b -> StackType b -> Bool) -> Eq (StackType b)
forall b. Eq b => StackType b -> StackType b -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: StackType b -> StackType b -> Bool
$c/= :: forall b. Eq b => StackType b -> StackType b -> Bool
== :: StackType b -> StackType b -> Bool
$c== :: forall b. Eq b => StackType b -> StackType b -> Bool
Eq, Eq (StackType b)
Eq (StackType b)
-> (StackType b -> StackType b -> Ordering)
-> (StackType b -> StackType b -> Bool)
-> (StackType b -> StackType b -> Bool)
-> (StackType b -> StackType b -> Bool)
-> (StackType b -> StackType b -> Bool)
-> (StackType b -> StackType b -> StackType b)
-> (StackType b -> StackType b -> StackType b)
-> Ord (StackType b)
StackType b -> StackType b -> Bool
StackType b -> StackType b -> Ordering
StackType b -> StackType b -> StackType b
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall b. Ord b => Eq (StackType b)
forall b. Ord b => StackType b -> StackType b -> Bool
forall b. Ord b => StackType b -> StackType b -> Ordering
forall b. Ord b => StackType b -> StackType b -> StackType b
min :: StackType b -> StackType b -> StackType b
$cmin :: forall b. Ord b => StackType b -> StackType b -> StackType b
max :: StackType b -> StackType b -> StackType b
$cmax :: forall b. Ord b => StackType b -> StackType b -> StackType b
>= :: StackType b -> StackType b -> Bool
$c>= :: forall b. Ord b => StackType b -> StackType b -> Bool
> :: StackType b -> StackType b -> Bool
$c> :: forall b. Ord b => StackType b -> StackType b -> Bool
<= :: StackType b -> StackType b -> Bool
$c<= :: forall b. Ord b => StackType b -> StackType b -> Bool
< :: StackType b -> StackType b -> Bool
$c< :: forall b. Ord b => StackType b -> StackType b -> Bool
compare :: StackType b -> StackType b -> Ordering
$ccompare :: forall b. Ord b => StackType b -> StackType b -> Ordering
$cp1Ord :: forall b. Ord b => Eq (StackType b)
Ord)

type MonoStackType = ([KempeTy ()], [KempeTy ()])

prettyMonoStackType :: MonoStackType -> Doc a
prettyMonoStackType :: MonoStackType -> Doc a
prettyMonoStackType ([KempeTy ()]
is, [KempeTy ()]
os) = [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
sep ((KempeTy () -> Doc a) -> [KempeTy ()] -> [Doc a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap KempeTy () -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty [KempeTy ()]
is) Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"--" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
sep ((KempeTy () -> Doc a) -> [KempeTy ()] -> [Doc a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap KempeTy () -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty [KempeTy ()]
os)

data BuiltinTy = TyInt
               | TyBool
               | TyInt8
               | TyWord
               deriving ((forall x. BuiltinTy -> Rep BuiltinTy x)
-> (forall x. Rep BuiltinTy x -> BuiltinTy) -> Generic BuiltinTy
forall x. Rep BuiltinTy x -> BuiltinTy
forall x. BuiltinTy -> Rep BuiltinTy x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep BuiltinTy x -> BuiltinTy
$cfrom :: forall x. BuiltinTy -> Rep BuiltinTy x
Generic, BuiltinTy -> ()
(BuiltinTy -> ()) -> NFData BuiltinTy
forall a. (a -> ()) -> NFData a
rnf :: BuiltinTy -> ()
$crnf :: BuiltinTy -> ()
NFData, BuiltinTy -> BuiltinTy -> Bool
(BuiltinTy -> BuiltinTy -> Bool)
-> (BuiltinTy -> BuiltinTy -> Bool) -> Eq BuiltinTy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BuiltinTy -> BuiltinTy -> Bool
$c/= :: BuiltinTy -> BuiltinTy -> Bool
== :: BuiltinTy -> BuiltinTy -> Bool
$c== :: BuiltinTy -> BuiltinTy -> Bool
Eq, Eq BuiltinTy
Eq BuiltinTy
-> (BuiltinTy -> BuiltinTy -> Ordering)
-> (BuiltinTy -> BuiltinTy -> Bool)
-> (BuiltinTy -> BuiltinTy -> Bool)
-> (BuiltinTy -> BuiltinTy -> Bool)
-> (BuiltinTy -> BuiltinTy -> Bool)
-> (BuiltinTy -> BuiltinTy -> BuiltinTy)
-> (BuiltinTy -> BuiltinTy -> BuiltinTy)
-> Ord BuiltinTy
BuiltinTy -> BuiltinTy -> Bool
BuiltinTy -> BuiltinTy -> Ordering
BuiltinTy -> BuiltinTy -> BuiltinTy
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: BuiltinTy -> BuiltinTy -> BuiltinTy
$cmin :: BuiltinTy -> BuiltinTy -> BuiltinTy
max :: BuiltinTy -> BuiltinTy -> BuiltinTy
$cmax :: BuiltinTy -> BuiltinTy -> BuiltinTy
>= :: BuiltinTy -> BuiltinTy -> Bool
$c>= :: BuiltinTy -> BuiltinTy -> Bool
> :: BuiltinTy -> BuiltinTy -> Bool
$c> :: BuiltinTy -> BuiltinTy -> Bool
<= :: BuiltinTy -> BuiltinTy -> Bool
$c<= :: BuiltinTy -> BuiltinTy -> Bool
< :: BuiltinTy -> BuiltinTy -> Bool
$c< :: BuiltinTy -> BuiltinTy -> Bool
compare :: BuiltinTy -> BuiltinTy -> Ordering
$ccompare :: BuiltinTy -> BuiltinTy -> Ordering
$cp1Ord :: Eq BuiltinTy
Ord)

instance Pretty BuiltinTy where
    pretty :: BuiltinTy -> Doc ann
pretty BuiltinTy
TyInt  = Doc ann
"Int"
    pretty BuiltinTy
TyBool = Doc ann
"Bool"
    pretty BuiltinTy
TyInt8 = Doc ann
"Int8"
    pretty BuiltinTy
TyWord = Doc ann
"Word"

instance Pretty (KempeTy a) where
    pretty :: KempeTy a -> Doc ann
pretty (TyBuiltin a
_ BuiltinTy
b)  = BuiltinTy -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty BuiltinTy
b
    pretty (TyNamed a
_ TyName a
tn)   = TyName a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty TyName a
tn
    pretty (TyVar a
_ TyName a
n)      = TyName a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty TyName a
n
    pretty (TyApp a
_ KempeTy a
ty KempeTy a
ty') = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (KempeTy a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty KempeTy a
ty Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> KempeTy a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty KempeTy a
ty')

instance Pretty (StackType a) where
    pretty :: StackType a -> Doc ann
pretty (StackType Set (Name a)
_ [KempeTy a]
ins [KempeTy a]
outs) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep ((KempeTy a -> Doc ann) -> [KempeTy a] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap KempeTy a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [KempeTy a]
ins) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"--" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep ((KempeTy a -> Doc ann) -> [KempeTy a] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap KempeTy a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [KempeTy a]
outs)

data ABI = Cabi
         | Kabi
         deriving (ABI -> ABI -> Bool
(ABI -> ABI -> Bool) -> (ABI -> ABI -> Bool) -> Eq ABI
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ABI -> ABI -> Bool
$c/= :: ABI -> ABI -> Bool
== :: ABI -> ABI -> Bool
$c== :: ABI -> ABI -> Bool
Eq, Eq ABI
Eq ABI
-> (ABI -> ABI -> Ordering)
-> (ABI -> ABI -> Bool)
-> (ABI -> ABI -> Bool)
-> (ABI -> ABI -> Bool)
-> (ABI -> ABI -> Bool)
-> (ABI -> ABI -> ABI)
-> (ABI -> ABI -> ABI)
-> Ord ABI
ABI -> ABI -> Bool
ABI -> ABI -> Ordering
ABI -> ABI -> ABI
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ABI -> ABI -> ABI
$cmin :: ABI -> ABI -> ABI
max :: ABI -> ABI -> ABI
$cmax :: ABI -> ABI -> ABI
>= :: ABI -> ABI -> Bool
$c>= :: ABI -> ABI -> Bool
> :: ABI -> ABI -> Bool
$c> :: ABI -> ABI -> Bool
<= :: ABI -> ABI -> Bool
$c<= :: ABI -> ABI -> Bool
< :: ABI -> ABI -> Bool
$c< :: ABI -> ABI -> Bool
compare :: ABI -> ABI -> Ordering
$ccompare :: ABI -> ABI -> Ordering
$cp1Ord :: Eq ABI
Ord, (forall x. ABI -> Rep ABI x)
-> (forall x. Rep ABI x -> ABI) -> Generic ABI
forall x. Rep ABI x -> ABI
forall x. ABI -> Rep ABI x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ABI x -> ABI
$cfrom :: forall x. ABI -> Rep ABI x
Generic, ABI -> ()
(ABI -> ()) -> NFData ABI
forall a. (a -> ()) -> NFData a
rnf :: ABI -> ()
$crnf :: ABI -> ()
NFData)

instance Pretty ABI where
    pretty :: ABI -> Doc ann
pretty ABI
Cabi = Doc ann
"cabi"
    pretty ABI
Kabi = Doc ann
"kabi"

-- machinery for assigning a constructor to a function of its concrete types
-- (and then curry forward...)

type Size = [Int64] -> Int64
type SizeEnv = IM.IntMap Size

-- the kempe sizing system is kind of fucked (it mostly works tho)

-- | Don't call this on ill-kinded types; it won't throw any error.
size :: SizeEnv -> KempeTy a -> Size
size :: SizeEnv -> KempeTy a -> Size
size SizeEnv
_ (TyBuiltin a
_ BuiltinTy
TyInt)                 = Int64 -> Size
forall a b. a -> b -> a
const Int64
8
size SizeEnv
_ (TyBuiltin a
_ BuiltinTy
TyBool)                = Int64 -> Size
forall a b. a -> b -> a
const Int64
1
size SizeEnv
_ (TyBuiltin a
_ BuiltinTy
TyInt8)                = Int64 -> Size
forall a b. a -> b -> a
const Int64
1
size SizeEnv
_ (TyBuiltin a
_ BuiltinTy
TyWord)                = Int64 -> Size
forall a b. a -> b -> a
const Int64
8
size SizeEnv
_ TyVar{}                             = [Char] -> Size
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error: type variables should not be present at this stage."
size SizeEnv
env (TyNamed a
_ (Name Text
_ (Unique Int
k) a
_)) = Size -> Int -> SizeEnv -> Size
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault ([Char] -> Size
forall a. HasCallStack => [Char] -> a
error [Char]
"Size not in map!") Int
k SizeEnv
env
size SizeEnv
env (TyApp a
_ KempeTy a
ty KempeTy a
ty')                  = \[Int64]
tys -> SizeEnv -> KempeTy a -> Size
forall a. SizeEnv -> KempeTy a -> Size
size SizeEnv
env KempeTy a
ty (SizeEnv -> KempeTy a -> Size
forall a. SizeEnv -> KempeTy a -> Size
size SizeEnv
env KempeTy a
ty' [] Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: [Int64]
tys)

cSize :: Size -> Int64
cSize :: Size -> Int64
cSize = (Size -> Size
forall a b. (a -> b) -> a -> b
$ [])

size' :: SizeEnv -> KempeTy a -> Int64
size' :: SizeEnv -> KempeTy a -> Int64
size' SizeEnv
env = Size -> Int64
cSize (Size -> Int64) -> (KempeTy a -> Size) -> KempeTy a -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SizeEnv -> KempeTy a -> Size
forall a. SizeEnv -> KempeTy a -> Size
size SizeEnv
env

sizeStack :: SizeEnv -> [KempeTy a] -> Int64
sizeStack :: SizeEnv -> [KempeTy a] -> Int64
sizeStack SizeEnv
env = Sum Int64 -> Int64
forall a. Sum a -> a
getSum (Sum Int64 -> Int64)
-> ([KempeTy a] -> Sum Int64) -> [KempeTy a] -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (KempeTy a -> Sum Int64) -> [KempeTy a] -> Sum Int64
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int64 -> Sum Int64
forall a. a -> Sum a
Sum (Int64 -> Sum Int64)
-> (KempeTy a -> Int64) -> KempeTy a -> Sum Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SizeEnv -> KempeTy a -> Int64
forall a. SizeEnv -> KempeTy a -> Int64
size' SizeEnv
env)