-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE DefaultSignatures #-}

module Types.Imp where

import Foreign.Ptr
import Data.Hashable
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString       as BS

import GHC.Generics (Generic (..))
import Data.Store (Store (..))
import Data.String (fromString)
import Data.Text.Prettyprint.Doc (line', nest, group)

import Name
import PPrint
import Util (IsBool (..))
import Types.Primitives
import Types.Source

-- === data types ===

type ImpName = Name ImpNameC

type ImpFunName = Name TopFunNameC
data IExpr n = ILit LitVal
             | IVar (ImpName n) BaseType
             | IPtrVar (Name PtrNameC n) PtrType
               deriving (Show, Generic)

data IBinder n l = IBinder (NameBinder ImpNameC n l) IType
                   deriving (Show, Generic)

type IVal = IExpr  -- only ILit and IRef constructors
type IType = BaseType
type Size = IExpr
type IVectorType = BaseType -- has to be a vector type

type IFunName = String
type IFunVar = (IFunName, IFunType)
data IFunType = IFunType CallingConvention [IType] [IType] -- args, results
                deriving (Show, Eq, Generic)

data IsCUDARequired = CUDARequired | CUDANotRequired  deriving (Eq, Show, Generic)

instance IsBool IsCUDARequired where
  toBool CUDARequired = True
  toBool CUDANotRequired = False

data CallingConvention =
   -- - Args passed one by one (not packed).  Scalar args are passed by value,
   --   arrays by pointer.
   -- - Dests are allocated by caller and passed as pointers
   --   (also one by one).
   -- - Used for generating standalone functions and functions called from
   --   Python without XLA.
   -- - Used for generating calls to standalone functions.
   StandardCC
   -- - Args and dests are each packed into a buffer and passed as one pointer,
   --   dests first.
   -- - Used for generating functions called from within an XLA computation.
 | XLACC
   -- - Args are packed by the caller into a buffer and passed as a pointer
   --   to the buffer; unpacking is code-generated.
   -- - Dests are allocated within the function and passed back to the caller
   --   (packed in a number-of-results-length buffer supplied by the caller as a
   --   pointer).
   -- - Used for generating top-level blocks and other Dex code called from
   --   Haskell.
 | EntryFunCC IsCUDARequired
   -- - Scalars only. Args passed one by one, by value. Single result, passed by
   --   value.
   -- - Used for generating calls from Dex to C functions that return a single
   --   scalar.
 | FFICC
   -- - Results written to a buffer supplied by the caller (by pointer).
   -- - Used for generating calls from Dex to C functions that return multiple
   --   results.
 | FFIMultiResultCC
   -- - Was used for generating bodies and calls of CUDA kernels.
 | CUDAKernelLaunch
   -- - Was used for generating bodies and calls of multi-threaded CPU kernels.
 | MCThreadLaunch
   deriving (Show, Eq, Generic)

data ImpFunction n = ImpFunction IFunType (Abs (Nest IBinder) ImpBlock n)
     deriving (Show, Generic)

data ImpBlock n where
  ImpBlock :: Nest ImpDecl n l -> [IExpr l] -> ImpBlock n
deriving instance Show (ImpBlock n)

data ImpDecl n l = ImpLet (Nest IBinder n l) (ImpInstr n)
     deriving (Show, Generic)

data ImpInstr n =
   IFor Direction (Size n) (Abs IBinder ImpBlock n)
 | IWhile (ImpBlock n)
 | ICond (IExpr n) (ImpBlock n) (ImpBlock n)
 | IQueryParallelism IFunVar (IExpr n) -- returns the number of available concurrent threads
 | ISyncWorkgroup
 | ILaunch IFunVar (Size n) [IExpr n]
 | ICall (ImpFunName n) [IExpr n]
 | Store (IExpr n) (IExpr n)           -- dest, val
 | Alloc AddressSpace IType (Size n)
 | StackAlloc IType (Size n)
 | MemCopy (IExpr n) (IExpr n) (IExpr n)   -- dest, source, numel
 | Free (IExpr n)
 | InitializeZeros (IExpr n) (IExpr n)  -- ptr, numel
 | GetAllocSize (IExpr n)  -- gives size in numel
 | IThrowError  -- TODO: parameterize by a run-time string
 | ICastOp IType (IExpr n)
 | IBitcastOp IType (IExpr n)
 | IVectorBroadcast (IExpr n) IVectorType
 | IVectorIota                IVectorType
 | IPtrLoad (IExpr n)
 | IPtrOffset (IExpr n) (IExpr n)
 | IBinOp BinOp (IExpr n) (IExpr n)
 | IUnOp  UnOp  (IExpr n)
 | ISelect (IExpr n) (IExpr n) (IExpr n)
 | IOutputStream
 | DebugPrint String (IExpr n) -- just prints an int64 value
 | IShowScalar (IExpr n) (IExpr n) -- pointer to result table, value
   deriving (Show, Generic)

iBinderType :: IBinder n l -> IType
iBinderType (IBinder _ ty) = ty
{-# INLINE iBinderType #-}

data Backend = LLVM | LLVMCUDA | LLVMMC | MLIR | Interpreter  deriving (Show, Eq)
newtype CUDAKernel = CUDAKernel B.ByteString deriving (Show)

-- === Closed Imp functions, LLVM and object file representation ===

-- Object files and LLVM modules expose ordinary C-toolchain names rather than
-- our internal scoped naming system. The `CNameInterface` data structure
-- describes how to interpret the names exposed by an LLVM module and its
-- corresponding object code. These names are all considered local to the
-- module. The module as a whole is a closed object corresponding to a
-- `ClosedImpFunction`.

-- TODO: consider adding more information here: types of required functions,
-- calling conventions etc.
type CName = String
data WithCNameInterface a = WithCNameInterface
  { cniCode         :: a       -- module itself (an LLVM.AST.Module or a bytestring of object code)
  , cniMainFunName  :: CName   -- name of function defined in this module
  , cniRequiredFuns :: [CName] -- names of functions required by this module
  , cniRequiredPtrs :: [CName] -- names of data pointers
  , cniDtorList     :: [CName] -- names of destructors (CUDA only) defined by this module
  } deriving (Show, Generic, Functor, Foldable, Traversable)

type RawObjCode = BS.ByteString
type FunObjCode = WithCNameInterface RawObjCode

data IFunBinder n l = IFunBinder (NameBinder TopFunNameC n l) IFunType

-- Imp function with link-time objects abstracted out, suitable for standalone
-- compilation. TODO: enforce actual `VoidS` as the scope parameter.
data ClosedImpFunction n where
  ClosedImpFunction
    :: Nest IFunBinder n1 n2 -- binders for required functions
    -> Nest PtrBinder  n2 n3 -- binders for required data pointers
    -> ImpFunction n3
    -> ClosedImpFunction n1

data PtrBinder n l = PtrBinder (NameBinder PtrNameC n l) PtrType
data LinktimeNames n = LinktimeNames [Name FunObjCodeNameC n] [Name PtrNameC n]  deriving (Show, Generic)
data LinktimeVals    = LinktimeVals  [FunPtr ()] [Ptr ()]                        deriving (Show, Generic)

data CFunction (n::S) = CFunction
  { nameHint :: NameHint
  , objectCode :: FunObjCode
  , linkerNames :: LinktimeNames n
  }
  deriving (Show, Generic)

instance BindsAtMostOneName IFunBinder TopFunNameC where
  IFunBinder b _ @> x = b @> x
  {-# INLINE (@>) #-}

instance BindsOneName IFunBinder TopFunNameC where
  binderName (IFunBinder b _) = binderName b
  {-# INLINE binderName #-}

instance HasNameHint (IFunBinder n l) where
  getNameHint (IFunBinder b _) = getNameHint b

instance GenericB IFunBinder where
  type RepB IFunBinder = BinderP TopFunNameC (LiftE IFunType)
  fromB (IFunBinder b ty) = b :> LiftE ty
  toB   (b :> LiftE ty) = IFunBinder b ty

instance ProvesExt  IFunBinder
instance BindsNames IFunBinder
instance SinkableB IFunBinder
instance HoistableB  IFunBinder
instance RenameB     IFunBinder
instance AlphaEqB IFunBinder
instance AlphaHashableB IFunBinder

instance GenericB PtrBinder where
  type RepB PtrBinder = BinderP PtrNameC (LiftE PtrType)
  fromB (PtrBinder b ty) = b :> LiftE ty
  toB   (b :> LiftE ty) = PtrBinder b ty

instance BindsAtMostOneName PtrBinder PtrNameC where
  PtrBinder b _ @> x = b @> x
  {-# INLINE (@>) #-}

instance HasNameHint (PtrBinder n l) where
  getNameHint (PtrBinder b _) = getNameHint b

instance ProvesExt   PtrBinder
instance BindsNames  PtrBinder
instance SinkableB   PtrBinder
instance HoistableB  PtrBinder
instance RenameB     PtrBinder
instance AlphaEqB    PtrBinder
instance AlphaHashableB PtrBinder

-- === instances ===

instance GenericE ImpInstr where
  type RepE ImpInstr = EitherE5
      (EitherE4
  {- IFor -}    (LiftE Direction `PairE` Size `PairE` Abs IBinder ImpBlock)
  {- IWhile -}  (ImpBlock)
  {- ICond -}   (IExpr `PairE` ImpBlock `PairE` ImpBlock)
  {- IQuery -}  (LiftE IFunVar `PairE` IExpr)
    ) (EitherE4
  {- ISyncW -}  (UnitE)
  {- ILaunch -} (LiftE IFunVar `PairE` Size `PairE` ListE IExpr)
  {- ICall -}   (ImpFunName `PairE` ListE IExpr)
  {- Store -}   (IExpr `PairE` IExpr)
    ) (EitherE7
  {- Alloc -}   (LiftE (AddressSpace, IType) `PairE` Size)
  {- StackAlloc -} (LiftE IType `PairE` Size)
  {- MemCopy -} (IExpr `PairE` IExpr `PairE` IExpr)
  {- InitializeZeros -}  (IExpr `PairE` IExpr)
  {- GetAllocSize -} IExpr
  {- Free -}    (IExpr)
  {- IThrowE -} (UnitE)
    ) (EitherE6
  {- ICastOp    -} (LiftE IType `PairE` IExpr)
  {- IBitcastOp -} (LiftE IType `PairE` IExpr)
  {- IVectorBroadcast -} (IExpr `PairE` LiftE IVectorType)
  {- IVectorIota      -} (              LiftE IVectorType)
  {- DebugPrint       -} (LiftE String `PairE` IExpr)
  {- IShowScalar      -} (IExpr `PairE` IExpr)
    ) (EitherE6
  {- IPtrLoad -}      IExpr
  {- IPtrOffset -}    (IExpr `PairE` IExpr)
  {- IBinOp -}        (LiftE BinOp `PairE` IExpr `PairE` IExpr)
  {- IUnOp -}         (LiftE UnOp `PairE` IExpr)
  {- ISelect -}       (IExpr`PairE` IExpr `PairE`IExpr)
  {- IOutputStream -} UnitE
    )

  fromE instr = case instr of
    IFor d n ab           -> Case0 $ Case0 $ LiftE d `PairE` n `PairE` ab
    IWhile body           -> Case0 $ Case1 body
    ICond p cons alt      -> Case0 $ Case2 $ p `PairE` cons `PairE` alt
    IQueryParallelism f s -> Case0 $ Case3 $ LiftE f `PairE` s

    ISyncWorkgroup      -> Case1 $ Case0 UnitE
    ILaunch f n args    -> Case1 $ Case1 $ LiftE f `PairE` n `PairE` ListE args
    ICall f args        -> Case1 $ Case2 $ f `PairE` ListE args
    Store dest val      -> Case1 $ Case3 $ dest `PairE` val

    Alloc a t s            -> Case2 $ Case0 $ LiftE (a, t) `PairE` s
    StackAlloc t s         -> Case2 $ Case1 $ LiftE t `PairE` s
    MemCopy dest src numel -> Case2 $ Case2 $ dest `PairE` src `PairE` numel
    InitializeZeros ptr numel -> Case2 $ Case3 $ ptr `PairE` numel
    GetAllocSize ptr       -> Case2 $ Case4 $ ptr
    Free ptr               -> Case2 $ Case5 ptr
    IThrowError            -> Case2 $ Case6 UnitE

    ICastOp idt ix         -> Case3 $ Case0 $ LiftE idt `PairE` ix
    IBitcastOp idt ix      -> Case3 $ Case1 $ LiftE idt `PairE` ix
    IVectorBroadcast v vty -> Case3 $ Case2 $ v `PairE` LiftE vty
    IVectorIota vty        -> Case3 $ Case3 $ LiftE vty
    DebugPrint s x         -> Case3 $ Case4 $ LiftE s `PairE` x
    IShowScalar x  y       -> Case3 $ Case5 $ x `PairE` y

    IPtrLoad x     -> Case4 $ Case0 $ x
    IPtrOffset x y -> Case4 $ Case1 $ x `PairE` y
    IBinOp op x y  -> Case4 $ Case2 $ LiftE op `PairE` x `PairE` y
    IUnOp op x     -> Case4 $ Case3 $ LiftE op `PairE` x
    ISelect x y z  -> Case4 $ Case4 $ x `PairE` y `PairE` z
    IOutputStream  -> Case4 $ Case5 $ UnitE
  {-# INLINE fromE #-}

  toE instr = case instr of
    Case0 instr' -> case instr' of
      Case0 (LiftE d `PairE` n `PairE` ab) -> IFor d n ab
      Case1 body                           -> IWhile body
      Case2 (p `PairE` cons `PairE` alt)   -> ICond p cons alt
      Case3 (LiftE f `PairE` s)            -> IQueryParallelism f s
      _ -> error "impossible"

    Case1 instr' -> case instr' of
      Case0 UnitE                                     -> ISyncWorkgroup
      Case1 (LiftE f `PairE` n `PairE` ListE args)    -> ILaunch f n args
      Case2 (f `PairE` ListE args)                    -> ICall f args
      Case3 (dest `PairE` val )                       -> Store dest val
      _ -> error "impossible"

    Case2 instr' -> case instr' of
      Case0 (LiftE (a, t) `PairE` s )         -> Alloc a t s
      Case1 (LiftE t `PairE` s )              -> StackAlloc t s
      Case2 (dest `PairE` src `PairE` numel)  -> MemCopy dest src numel
      Case3 (ptr `PairE` numel)               -> InitializeZeros ptr numel
      Case4 ptr                               -> GetAllocSize ptr
      Case5 ptr                               -> Free ptr
      Case6 UnitE                             -> IThrowError
      _ -> error "impossible"

    Case3 instr' -> case instr' of
      Case0 (LiftE idt `PairE` ix ) -> ICastOp idt ix
      Case1 (LiftE idt `PairE` ix ) -> IBitcastOp idt ix
      Case2 (v `PairE` LiftE vty) -> IVectorBroadcast v vty
      Case3 (          LiftE vty) -> IVectorIota vty
      Case4 (LiftE s `PairE` x)   -> DebugPrint s x
      Case5 (x `PairE` y)         -> IShowScalar x y
      _ -> error "impossible"

    _ -> error "impossible"
  {-# INLINE toE #-}

instance SinkableE ImpInstr
instance HoistableE  ImpInstr
instance AlphaEqE ImpInstr
instance AlphaHashableE ImpInstr
instance RenameE     ImpInstr

instance GenericE ImpBlock where
  type RepE ImpBlock = Abs (Nest ImpDecl) (ListE IExpr)
  fromE (ImpBlock decls results) = Abs decls (ListE results)
  {-# INLINE fromE #-}
  toE   (Abs decls (ListE results)) = ImpBlock decls results
  {-# INLINE toE #-}

instance SinkableE ImpBlock
instance HoistableE  ImpBlock
instance AlphaEqE ImpBlock
instance AlphaHashableE ImpBlock
instance RenameE     ImpBlock
deriving via WrapE ImpBlock n instance Generic (ImpBlock n)

instance GenericE IExpr where
  type RepE IExpr = EitherE3 (LiftE LitVal)
                             (PairE ImpName         (LiftE BaseType))
                             (PairE (Name PtrNameC) (LiftE PtrType))
  fromE iexpr = case iexpr of
    ILit x       -> Case0 (LiftE x)
    IVar    v ty -> Case1 (v `PairE` LiftE ty)
    IPtrVar v ty -> Case2 (v `PairE` LiftE ty)
  {-# INLINE fromE #-}

  toE rep = case rep of
    Case0 (LiftE x) -> ILit x
    Case1 (v `PairE` LiftE ty) -> IVar    v ty
    Case2 (v `PairE` LiftE ty) -> IPtrVar v ty
    _ -> error "impossible"
  {-# INLINE toE #-}

instance SinkableE IExpr
instance HoistableE  IExpr
instance AlphaEqE IExpr
instance AlphaHashableE IExpr
instance RenameE     IExpr

instance GenericB IBinder where
  type RepB IBinder = PairB (LiftB (LiftE IType)) (NameBinder ImpNameC)
  fromB (IBinder b ty) = PairB (LiftB (LiftE ty)) b
  toB   (PairB (LiftB (LiftE ty)) b) = IBinder b ty

instance HasNameHint (IBinder n l) where
  getNameHint (IBinder b _) = getNameHint b

instance BindsAtMostOneName IBinder ImpNameC where
  IBinder b _ @> x = b @> x

instance BindsOneName IBinder ImpNameC where
  binderName (IBinder b _) = binderName b

instance BindsNames IBinder where
  toScopeFrag (IBinder b _) = toScopeFrag b

instance ProvesExt  IBinder
instance SinkableB IBinder
instance HoistableB  IBinder
instance RenameB  IBinder
instance AlphaEqB IBinder
instance AlphaHashableB IBinder

instance GenericB ImpDecl where
  type RepB ImpDecl = PairB (LiftB ImpInstr) (Nest IBinder)
  fromB (ImpLet bs instr) = PairB (LiftB instr) bs
  toB   (PairB (LiftB instr) bs) = ImpLet bs instr

instance SinkableB ImpDecl
instance HoistableB  ImpDecl
instance RenameB     ImpDecl
instance AlphaEqB ImpDecl
instance AlphaHashableB ImpDecl
instance ProvesExt  ImpDecl
instance BindsNames ImpDecl

instance GenericE ImpFunction where
  type RepE ImpFunction = LiftE IFunType `PairE` Abs (Nest IBinder) ImpBlock
  fromE (ImpFunction ty ab) =LiftE ty `PairE` ab
  {-# INLINE fromE #-}

  toE (LiftE ty `PairE` ab) = ImpFunction ty ab
  {-# INLINE toE #-}

instance SinkableE ImpFunction
instance HoistableE  ImpFunction
instance AlphaEqE    ImpFunction
instance AlphaHashableE    ImpFunction
instance RenameE     ImpFunction

instance GenericE LinktimeNames where
  type RepE LinktimeNames = ListE  (Name FunObjCodeNameC)
                   `PairE`  ListE  (Name PtrNameC)
  fromE (LinktimeNames funs ptrs) = ListE funs `PairE` ListE ptrs
  {-# INLINE fromE #-}
  toE (ListE funs `PairE` ListE ptrs) = LinktimeNames funs ptrs
  {-# INLINE toE #-}

instance SinkableE      LinktimeNames
instance HoistableE     LinktimeNames
instance AlphaEqE       LinktimeNames
instance AlphaHashableE LinktimeNames
instance RenameE        LinktimeNames

instance GenericE CFunction where
  type RepE CFunction = (LiftE NameHint) `PairE`
    (LiftE FunObjCode) `PairE` LinktimeNames
  fromE (CFunction{..}) = LiftE nameHint `PairE`
    LiftE objectCode `PairE` linkerNames
  {-# INLINE fromE #-}
  toE (LiftE nameHint `PairE` LiftE objectCode `PairE` linkerNames) =
    CFunction{..}
  {-# INLINE toE #-}

instance SinkableE      CFunction
instance HoistableE     CFunction
instance RenameE        CFunction

instance Store IsCUDARequired
instance Store CallingConvention
instance Store a => Store (WithCNameInterface a)
instance Store (IBinder n l)
instance Store (ImpDecl n l)
instance Store (IFunType)
instance Store (ImpInstr n)
instance Store (IExpr n)
instance Store (ImpBlock n)
instance Store (ImpFunction n)
instance Store (LinktimeNames n)
instance Store (CFunction n)
instance Store LinktimeVals

instance Hashable IsCUDARequired
instance Hashable CallingConvention
instance Hashable IFunType

instance Pretty CallingConvention where pretty = fromString . show

instance Pretty (ImpFunction n) where
  pretty (ImpFunction (IFunType cc _ _) (Abs bs body)) =
    "impfun" <+> pretty cc <+> prettyBinderNest bs
    <> nest 2 (hardline <> pretty body) <> hardline

instance Pretty (ImpBlock n)  where
  pretty = \case
    ImpBlock Empty [] -> mempty
    ImpBlock Empty expr -> group $ hardline <> pLowest expr
    ImpBlock decls [] -> prettyLines $ fromNest decls
    ImpBlock decls expr -> prettyLines decls' <> hardline <> pLowest expr
      where decls' = fromNest decls

instance Pretty (IBinder n l)  where
  pretty (IBinder b ty) = pretty b <+> ":" <+> pretty ty

instance Pretty (ImpInstr n)  where
  pretty = \case
    IFor a n (Abs i block) -> forStr a <+> p i <+> "<" <+> p n <>
                                      nest 4 (p block)
    IWhile body -> "while" <+> nest 2 (p body)
    ICond predicate cons alt ->
       "if" <+> p predicate <+> "then" <> nest 2 (p cons) <>
       hardline <> "else" <> nest 2 (p alt)
    IQueryParallelism f s -> "queryParallelism" <+> p f <+> p s
    ILaunch f s args -> "launch" <+> p f <+> p s <+> spaced args
    ICastOp t x    -> "cast"  <+> p x <+> "to" <+> p t
    IBitcastOp t x -> "bitcast"  <+> p x <+> "to" <+> p t
    Store dest val -> "store" <+> p dest <+> p val
    Alloc _ t s    -> "alloc" <+> p t <> "[" <> sizeStr s <> "]"
    StackAlloc t s -> "alloca" <+> p t <> "[" <> sizeStr s <> "]"
    MemCopy dest src numel -> "memcopy" <+> p dest <+> p src <+> p numel
    InitializeZeros ptr numel -> "initializeZeros" <+> p ptr <+> p numel
    GetAllocSize ptr -> "getAllocSize" <+> p ptr
    Free ptr       -> "free"  <+> p ptr
    ISyncWorkgroup   -> "syncWorkgroup"
    IThrowError      -> "throwError"
    ICall f args   -> "call" <+> p f <+> p args
    IVectorBroadcast v _ -> "vbroadcast" <+> p v
    IVectorIota _ -> "viota"
    DebugPrint s x -> "debug_print" <+> p (show s) <+> p x
    IPtrLoad ptr   -> "load" <+> p ptr
    IPtrOffset ptr idx -> p ptr <+> "+>" <+> p idx
    IBinOp op x y -> opDefault (UBinOp op) [x, y]
    IUnOp  op x   -> opDefault (UUnOp  op) [x]
    ISelect x y z -> "select" <+> p x <+> p y <+> p z
    IOutputStream -> "outputStream"
    IShowScalar ptr x -> "show_scalar" <+> p ptr <+> p x
    where opDefault name xs = prettyOpDefault name xs $ AppPrec
          p :: Pretty a => a -> Doc ann
          p = pretty
          forStr :: ForAnn -> Doc ann
          forStr = \case
            Fwd -> "for"
            Rev -> "rof"

sizeStr :: IExpr n -> Doc ann
sizeStr s = case s of
  ILit (Word32Lit x) -> pretty x  -- print in decimal because it's more readable
  _ -> pretty s

instance Pretty (IExpr n) where
  pretty = \case
    ILit v -> pretty v
    IVar v _ -> pretty v
    IPtrVar v _ -> pretty v

instance PrettyPrec (IExpr n) where prettyPrec = atPrec ArgPrec . pretty

instance Pretty (ImpDecl n l) where
  pretty = \case
    ImpLet Empty instr -> pretty instr
    ImpLet (Nest b Empty) instr -> pretty b <+> "=" <+> pretty instr
    ImpLet bs instr -> pretty bs <+> "=" <+> pretty instr

instance Pretty IFunType where
  pretty (IFunType cc argTys retTys) =
    "Fun" <+> pretty cc <+> pretty argTys <+> "->" <+> pretty retTys

prettyBinderNest :: PrettyB b => Nest b n l -> Doc ann
prettyBinderNest bs = nest 6 $ line' <> (sep $ map pretty $ fromNest bs)

fromNest :: Nest b n l -> [b UnsafeS UnsafeS]
fromNest Empty = []
fromNest (Nest b rest) = unsafeCoerceB b : fromNest rest