content: {-# 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
                        , size
                        , size'
                        , sizeStack
                        ) where
  
  import           Control.DeepSeq (NFData)
  import           Data.Int        (Int64)
  import qualified Data.IntMap     as IM
  import           Data.Monoid     (Sum (..))
  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 (Generic, NFData, Functor, Eq, Ord) -- questionable eq instance but eh
  
  data StackType b = StackType { inTypes  :: [KempeTy b]
                               , outTypes :: [KempeTy b]
                               } deriving (Generic, NFData, Eq, Ord)
  
  type MonoStackType = ([KempeTy ()], [KempeTy ()])
  
  prettyMonoStackType :: ([KempeTy a], [KempeTy a]) -> Doc ann
  prettyMonoStackType (is, os) = sep (fmap pretty is) <+> "--" <+> sep (fmap pretty os)
  
  data BuiltinTy = TyInt
                 | TyBool
                 | TyInt8
                 | TyWord
                 deriving (Generic, NFData, Eq, Ord)
  
  instance Pretty BuiltinTy where
      pretty TyInt  = "Int"
      pretty TyBool = "Bool"
      pretty TyInt8 = "Int8"
      pretty TyWord = "Word"
  
  instance Pretty (KempeTy a) where
      pretty (TyBuiltin _ b)  = pretty b
      pretty (TyNamed _ tn)   = pretty tn
      pretty (TyVar _ n)      = pretty n
      pretty (TyApp _ ty ty') = parens (pretty ty <+> pretty ty')
  
  instance Pretty (StackType a) where
      pretty (StackType ins outs) = sep (fmap pretty ins) <+> "--" <+> sep (fmap pretty outs)
  
  data ABI = Cabi
           | Kabi
           | Hooked
           | ArmAbi
           deriving (Eq, Ord, Generic, NFData)
  
  instance Pretty ABI where
      pretty Cabi   = "cabi"
      pretty Kabi   = "kabi"
      pretty Hooked = "hooked"
      pretty ArmAbi = "armabi"
  
  -- 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 works tho)
  
  -- | Don't call this on ill-kinded types; it won't throw any error.
  size :: SizeEnv -> KempeTy a -> Size
  size _ (TyBuiltin _ TyInt)                 = const 8
  size _ (TyBuiltin _ TyBool)                = const 1
  size _ (TyBuiltin _ TyInt8)                = const 1
  size _ (TyBuiltin _ TyWord)                = const 8
  size _ TyVar{}                             = error "Internal error: type variables should not be present at this stage."
  size env (TyNamed _ (Name _ (Unique k) _)) = IM.findWithDefault (error "Size not in map!") k env
  size env (TyApp _ ty ty')                  = \tys -> size env ty (size env ty' [] : tys)
  
  size' :: SizeEnv -> KempeTy a -> Int64
  size' env = ($ []) . size env
  
  sizeStack :: SizeEnv -> [KempeTy a] -> Int64
  sizeStack env = getSum . foldMap (Sum . size' env)
  