{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeFamilies #-}

module Futhark.CodeGen.ImpGen
  ( -- * Entry Points
    compileProg,

    -- * Pluggable Compiler
    OpCompiler,
    ExpCompiler,
    CopyCompiler,
    StmsCompiler,
    AllocCompiler,
    Operations (..),
    defaultOperations,
    MemLocation (..),
    MemEntry (..),
    ScalarEntry (..),

    -- * Monadic Compiler Interface
    ImpM,
    localDefaultSpace,
    askFunction,
    newVNameForFun,
    nameForFun,
    askEnv,
    localEnv,
    localOps,
    VTable,
    getVTable,
    localVTable,
    subImpM,
    subImpM_,
    emit,
    emitFunction,
    hasFunction,
    collect,
    collect',
    comment,
    VarEntry (..),
    ArrayEntry (..),

    -- * Lookups
    lookupVar,
    lookupArray,
    lookupMemory,

    -- * Building Blocks
    TV,
    mkTV,
    tvSize,
    tvExp,
    tvVar,
    ToExp (..),
    compileAlloc,
    everythingVolatile,
    compileBody,
    compileBody',
    compileLoopBody,
    defCompileStms,
    compileStms,
    compileExp,
    defCompileExp,
    fullyIndexArray,
    fullyIndexArray',
    copy,
    copyDWIM,
    copyDWIMFix,
    copyElementWise,
    typeSize,
    isMapTransposeCopy,

    -- * Constructing code.
    dLParams,
    dFParams,
    dScope,
    dArray,
    dPrim,
    dPrimVol,
    dPrim_,
    dPrimV_,
    dPrimV,
    dPrimVE,
    sFor,
    sWhile,
    sComment,
    sIf,
    sWhen,
    sUnless,
    sOp,
    sDeclareMem,
    sAlloc,
    sAlloc_,
    sArray,
    sArrayInMem,
    sAllocArray,
    sAllocArrayPerm,
    sStaticArray,
    sWrite,
    sUpdate,
    sLoopNest,
    (<--),
    (<~~),
    function,
    warn,
    module Language.Futhark.Warnings,
  )
where

import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer
import Control.Parallel.Strategies
import Data.Bifunctor (first)
import qualified Data.DList as DL
import Data.Either
import Data.List (find, genericLength, sortOn)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import Data.String
import Futhark.CodeGen.ImpCode
  ( Bytes,
    Count,
    Elements,
    bytes,
    elements,
    withElemType,
  )
import qualified Futhark.CodeGen.ImpCode as Imp
import Futhark.CodeGen.ImpGen.Transpose
import Futhark.Construct hiding (ToExp (..))
import Futhark.IR.Mem
import qualified Futhark.IR.Mem.IxFun as IxFun
import Futhark.IR.SOACS (SOACS)
import Futhark.Util
import Futhark.Util.Loc (noLoc)
import Language.Futhark.Warnings

-- | How to compile an t'Op'.
type OpCompiler lore r op = Pattern lore -> Op lore -> ImpM lore r op ()

-- | How to compile some 'Stms'.
type StmsCompiler lore r op = Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op ()

-- | How to compile an 'Exp'.
type ExpCompiler lore r op = Pattern lore -> Exp lore -> ImpM lore r op ()

type CopyCompiler lore r op =
  PrimType ->
  MemLocation ->
  Slice (Imp.TExp Int64) ->
  MemLocation ->
  Slice (Imp.TExp Int64) ->
  ImpM lore r op ()

-- | An alternate way of compiling an allocation.
type AllocCompiler lore r op = VName -> Count Bytes (Imp.TExp Int64) -> ImpM lore r op ()

data Operations lore r op = Operations
  { opsExpCompiler :: ExpCompiler lore r op,
    opsOpCompiler :: OpCompiler lore r op,
    opsStmsCompiler :: StmsCompiler lore r op,
    opsCopyCompiler :: CopyCompiler lore r op,
    opsAllocCompilers :: M.Map Space (AllocCompiler lore r op)
  }

-- | An operations set for which the expression compiler always
-- returns 'defCompileExp'.
defaultOperations ::
  (Mem lore, FreeIn op) =>
  OpCompiler lore r op ->
  Operations lore r op
defaultOperations opc =
  Operations
    { opsExpCompiler = defCompileExp,
      opsOpCompiler = opc,
      opsStmsCompiler = defCompileStms,
      opsCopyCompiler = defaultCopy,
      opsAllocCompilers = mempty
    }

-- | When an array is dared, this is where it is stored.
data MemLocation = MemLocation
  { memLocationName :: VName,
    memLocationShape :: [Imp.DimSize],
    memLocationIxFun :: IxFun.IxFun (Imp.TExp Int64)
  }
  deriving (Eq, Show)

data ArrayEntry = ArrayEntry
  { entryArrayLocation :: MemLocation,
    entryArrayElemType :: PrimType
  }
  deriving (Show)

entryArrayShape :: ArrayEntry -> [Imp.DimSize]
entryArrayShape = memLocationShape . entryArrayLocation

newtype MemEntry = MemEntry {entryMemSpace :: Imp.Space}
  deriving (Show)

newtype ScalarEntry = ScalarEntry
  { entryScalarType :: PrimType
  }
  deriving (Show)

-- | Every non-scalar variable must be associated with an entry.
data VarEntry lore
  = ArrayVar (Maybe (Exp lore)) ArrayEntry
  | ScalarVar (Maybe (Exp lore)) ScalarEntry
  | MemVar (Maybe (Exp lore)) MemEntry
  deriving (Show)

-- | When compiling an expression, this is a description of where the
-- result should end up.  The integer is a reference to the construct
-- that gave rise to this destination (for patterns, this will be the
-- tag of the first name in the pattern).  This can be used to make
-- the generated code easier to relate to the original code.
data Destination = Destination
  { destinationTag :: Maybe Int,
    valueDestinations :: [ValueDestination]
  }
  deriving (Show)

data ValueDestination
  = ScalarDestination VName
  | MemoryDestination VName
  | -- | The 'MemLocation' is 'Just' if a copy if
    -- required.  If it is 'Nothing', then a
    -- copy/assignment of a memory block somewhere
    -- takes care of this array.
    ArrayDestination (Maybe MemLocation)
  deriving (Show)

data Env lore r op = Env
  { envExpCompiler :: ExpCompiler lore r op,
    envStmsCompiler :: StmsCompiler lore r op,
    envOpCompiler :: OpCompiler lore r op,
    envCopyCompiler :: CopyCompiler lore r op,
    envAllocCompilers :: M.Map Space (AllocCompiler lore r op),
    envDefaultSpace :: Imp.Space,
    envVolatility :: Imp.Volatility,
    -- | User-extensible environment.
    envEnv :: r,
    -- | Name of the function we are compiling, if any.
    envFunction :: Maybe Name,
    -- | The set of attributes that are active on the enclosing
    -- statements (including the one we are currently compiling).
    envAttrs :: Attrs
  }

newEnv :: r -> Operations lore r op -> Imp.Space -> Env lore r op
newEnv r ops ds =
  Env
    { envExpCompiler = opsExpCompiler ops,
      envStmsCompiler = opsStmsCompiler ops,
      envOpCompiler = opsOpCompiler ops,
      envCopyCompiler = opsCopyCompiler ops,
      envAllocCompilers = mempty,
      envDefaultSpace = ds,
      envVolatility = Imp.Nonvolatile,
      envEnv = r,
      envFunction = Nothing,
      envAttrs = mempty
    }

-- | The symbol table used during compilation.
type VTable lore = M.Map VName (VarEntry lore)

data ImpState lore r op = ImpState
  { stateVTable :: VTable lore,
    stateFunctions :: Imp.Functions op,
    stateCode :: Imp.Code op,
    stateWarnings :: Warnings,
    stateNameSource :: VNameSource
  }

newState :: VNameSource -> ImpState lore r op
newState = ImpState mempty mempty mempty mempty

newtype ImpM lore r op a
  = ImpM (ReaderT (Env lore r op) (State (ImpState lore r op)) a)
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadState (ImpState lore r op),
      MonadReader (Env lore r op)
    )

instance MonadFreshNames (ImpM lore r op) where
  getNameSource = gets stateNameSource
  putNameSource src = modify $ \s -> s {stateNameSource = src}

-- Cannot be an KernelsMem scope because the index functions have
-- the wrong leaves (VName instead of Imp.Exp).
instance HasScope SOACS (ImpM lore r op) where
  askScope = gets $ M.map (LetName . entryType) . stateVTable
    where
      entryType (MemVar _ memEntry) =
        Mem (entryMemSpace memEntry)
      entryType (ArrayVar _ arrayEntry) =
        Array
          (entryArrayElemType arrayEntry)
          (Shape $ entryArrayShape arrayEntry)
          NoUniqueness
      entryType (ScalarVar _ scalarEntry) =
        Prim $ entryScalarType scalarEntry

runImpM ::
  ImpM lore r op a ->
  r ->
  Operations lore r op ->
  Imp.Space ->
  ImpState lore r op ->
  (a, ImpState lore r op)
runImpM (ImpM m) r ops space = runState (runReaderT m $ newEnv r ops space)

subImpM_ ::
  r' ->
  Operations lore r' op' ->
  ImpM lore r' op' a ->
  ImpM lore r op (Imp.Code op')
subImpM_ r ops m = snd <$> subImpM r ops m

subImpM ::
  r' ->
  Operations lore r' op' ->
  ImpM lore r' op' a ->
  ImpM lore r op (a, Imp.Code op')
subImpM r ops (ImpM m) = do
  env <- ask
  s <- get

  let env' =
        env
          { envExpCompiler = opsExpCompiler ops,
            envStmsCompiler = opsStmsCompiler ops,
            envCopyCompiler = opsCopyCompiler ops,
            envOpCompiler = opsOpCompiler ops,
            envAllocCompilers = opsAllocCompilers ops,
            envEnv = r
          }
      s' =
        ImpState
          { stateVTable = stateVTable s,
            stateFunctions = mempty,
            stateCode = mempty,
            stateNameSource = stateNameSource s,
            stateWarnings = mempty
          }
      (x, s'') = runState (runReaderT m env') s'

  putNameSource $ stateNameSource s''
  warnings $ stateWarnings s''
  return (x, stateCode s'')

-- | Execute a code generation action, returning the code that was
-- emitted.
collect :: ImpM lore r op () -> ImpM lore r op (Imp.Code op)
collect = fmap snd . collect'

collect' :: ImpM lore r op a -> ImpM lore r op (a, Imp.Code op)
collect' m = do
  prev_code <- gets stateCode
  modify $ \s -> s {stateCode = mempty}
  x <- m
  new_code <- gets stateCode
  modify $ \s -> s {stateCode = prev_code}
  return (x, new_code)

-- | Execute a code generation action, wrapping the generated code
-- within a 'Imp.Comment' with the given description.
comment :: String -> ImpM lore r op () -> ImpM lore r op ()
comment desc m = do
  code <- collect m
  emit $ Imp.Comment desc code

-- | Emit some generated imperative code.
emit :: Imp.Code op -> ImpM lore r op ()
emit code = modify $ \s -> s {stateCode = stateCode s <> code}

warnings :: Warnings -> ImpM lore r op ()
warnings ws = modify $ \s -> s {stateWarnings = ws <> stateWarnings s}

-- | Emit a warning about something the user should be aware of.
warn :: Located loc => loc -> [loc] -> String -> ImpM lore r op ()
warn loc locs problem =
  warnings $ singleWarning' (srclocOf loc) (map srclocOf locs) (fromString problem)

-- | Emit a function in the generated code.
emitFunction :: Name -> Imp.Function op -> ImpM lore r op ()
emitFunction fname fun = do
  Imp.Functions fs <- gets stateFunctions
  modify $ \s -> s {stateFunctions = Imp.Functions $ (fname, fun) : fs}

-- | Check if a function of a given name exists.
hasFunction :: Name -> ImpM lore r op Bool
hasFunction fname = gets $ \s ->
  let Imp.Functions fs = stateFunctions s
   in isJust $ lookup fname fs

constsVTable :: Mem lore => Stms lore -> VTable lore
constsVTable = foldMap stmVtable
  where
    stmVtable (Let pat _ e) =
      foldMap (peVtable e) $ patternElements pat
    peVtable e (PatElem name dec) =
      M.singleton name $ memBoundToVarEntry (Just e) dec

compileProg ::
  (Mem lore, FreeIn op, MonadFreshNames m) =>
  r ->
  Operations lore r op ->
  Imp.Space ->
  Prog lore ->
  m (Warnings, Imp.Definitions op)
compileProg r ops space (Prog consts funs) =
  modifyNameSource $ \src ->
    let (_, ss) =
          unzip $ parMap rpar (compileFunDef' src) funs
        free_in_funs =
          freeIn $ mconcat $ map stateFunctions ss
        (consts', s') =
          runImpM (compileConsts free_in_funs consts) r ops space $
            combineStates ss
     in ( ( stateWarnings s',
            Imp.Definitions consts' (stateFunctions s')
          ),
          stateNameSource s'
        )
  where
    compileFunDef' src fdef =
      runImpM
        (compileFunDef fdef)
        r
        ops
        space
        (newState src) {stateVTable = constsVTable consts}

    combineStates ss =
      let Imp.Functions funs' = mconcat $ map stateFunctions ss
          src = mconcat (map stateNameSource ss)
       in (newState src)
            { stateFunctions =
                Imp.Functions $ M.toList $ M.fromList funs',
              stateWarnings =
                mconcat $ map stateWarnings ss
            }

compileConsts :: Names -> Stms lore -> ImpM lore r op (Imp.Constants op)
compileConsts used_consts stms = do
  code <- collect $ compileStms used_consts stms $ pure ()
  pure $ uncurry Imp.Constants $ first DL.toList $ extract code
  where
    -- Fish out those top-level declarations in the constant
    -- initialisation code that are free in the functions.
    extract (x Imp.:>>: y) =
      extract x <> extract y
    extract (Imp.DeclareMem name space)
      | name `nameIn` used_consts =
        ( DL.singleton $ Imp.MemParam name space,
          mempty
        )
    extract (Imp.DeclareScalar name _ t)
      | name `nameIn` used_consts =
        ( DL.singleton $ Imp.ScalarParam name t,
          mempty
        )
    extract s =
      (mempty, s)

compileInParam ::
  Mem lore =>
  FParam lore ->
  ImpM lore r op (Either Imp.Param ArrayDecl)
compileInParam fparam = case paramDec fparam of
  MemPrim bt ->
    return $ Left $ Imp.ScalarParam name bt
  MemMem space ->
    return $ Left $ Imp.MemParam name space
  MemArray bt shape _ (ArrayIn mem ixfun) ->
    return $
      Right $
        ArrayDecl name bt $
          MemLocation mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun
  where
    name = paramName fparam

data ArrayDecl = ArrayDecl VName PrimType MemLocation

compileInParams ::
  Mem lore =>
  [FParam lore] ->
  [EntryPointType] ->
  ImpM lore r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue])
compileInParams params orig_epts = do
  let (ctx_params, val_params) =
        splitAt (length params - sum (map entryPointSize orig_epts)) params
  (inparams, arrayds) <- partitionEithers <$> mapM compileInParam (ctx_params ++ val_params)
  let findArray x = find (isArrayDecl x) arrayds

      summaries = M.fromList $ mapMaybe memSummary params
        where
          memSummary param
            | MemMem space <- paramDec param =
              Just (paramName param, space)
            | otherwise =
              Nothing

      findMemInfo :: VName -> Maybe Space
      findMemInfo = flip M.lookup summaries

      mkValueDesc fparam signedness =
        case (findArray $ paramName fparam, paramType fparam) of
          (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do
            memspace <- findMemInfo mem
            Just $ Imp.ArrayValue mem memspace bt signedness shape
          (_, Prim bt) ->
            Just $ Imp.ScalarValue bt signedness $ paramName fparam
          _ ->
            Nothing

      mkExts (TypeOpaque desc n : epts) fparams =
        let (fparams', rest) = splitAt n fparams
         in Imp.OpaqueValue
              desc
              (mapMaybe (`mkValueDesc` Imp.TypeDirect) fparams') :
            mkExts epts rest
      mkExts (TypeUnsigned : epts) (fparam : fparams) =
        maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeUnsigned)
          ++ mkExts epts fparams
      mkExts (TypeDirect : epts) (fparam : fparams) =
        maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeDirect)
          ++ mkExts epts fparams
      mkExts _ _ = []

  return (inparams, arrayds, mkExts orig_epts val_params)
  where
    isArrayDecl x (ArrayDecl y _ _) = x == y

compileOutParams ::
  Mem lore =>
  [RetType lore] ->
  [EntryPointType] ->
  ImpM lore r op ([Imp.ExternalValue], [Imp.Param], Destination)
compileOutParams orig_rts orig_epts = do
  ((extvs, dests), (outparams, ctx_dests)) <-
    runWriterT $ evalStateT (mkExts orig_epts orig_rts) (M.empty, M.empty)
  let ctx_dests' = map snd $ sortOn fst $ M.toList ctx_dests
  return (extvs, outparams, Destination Nothing $ ctx_dests' <> dests)
  where
    imp = lift . lift

    mkExts (TypeOpaque desc n : epts) rts = do
      let (rts', rest) = splitAt n rts
      (evs, dests) <- unzip <$> zipWithM mkParam rts' (repeat Imp.TypeDirect)
      (more_values, more_dests) <- mkExts epts rest
      return
        ( Imp.OpaqueValue desc evs : more_values,
          dests ++ more_dests
        )
    mkExts (TypeUnsigned : epts) (rt : rts) = do
      (ev, dest) <- mkParam rt Imp.TypeUnsigned
      (more_values, more_dests) <- mkExts epts rts
      return
        ( Imp.TransparentValue ev : more_values,
          dest : more_dests
        )
    mkExts (TypeDirect : epts) (rt : rts) = do
      (ev, dest) <- mkParam rt Imp.TypeDirect
      (more_values, more_dests) <- mkExts epts rts
      return
        ( Imp.TransparentValue ev : more_values,
          dest : more_dests
        )
    mkExts _ _ = return ([], [])

    mkParam MemMem {} _ =
      error "Functions may not explicitly return memory blocks."
    mkParam (MemPrim t) ept = do
      out <- imp $ newVName "scalar_out"
      tell ([Imp.ScalarParam out t], mempty)
      return (Imp.ScalarValue t ept out, ScalarDestination out)
    mkParam (MemArray t shape _ dec) ept = do
      space <- asks envDefaultSpace
      memout <- case dec of
        ReturnsNewBlock _ x _ixfun -> do
          memout <- imp $ newVName "out_mem"
          tell
            ( [Imp.MemParam memout space],
              M.singleton x $ MemoryDestination memout
            )
          return memout
        ReturnsInBlock memout _ ->
          return memout
      resultshape <- mapM inspectExtSize $ shapeDims shape
      return
        ( Imp.ArrayValue memout space t ept resultshape,
          ArrayDestination Nothing
        )

    inspectExtSize (Ext x) = do
      (memseen, arrseen) <- get
      case M.lookup x arrseen of
        Nothing -> do
          out <- imp $ newVName "out_arrsize"
          tell
            ( [Imp.ScalarParam out int64],
              M.singleton x $ ScalarDestination out
            )
          put (memseen, M.insert x out arrseen)
          return $ Var out
        Just out ->
          return $ Var out
    inspectExtSize (Free se) =
      return se

compileFunDef ::
  Mem lore =>
  FunDef lore ->
  ImpM lore r op ()
compileFunDef (FunDef entry _ fname rettype params body) =
  local (\env -> env {envFunction = Just fname}) $ do
    ((outparams, inparams, results, args), body') <- collect' compile
    emitFunction fname $ Imp.Function (isJust entry) outparams inparams body' results args
  where
    params_entry = maybe (replicate (length params) TypeDirect) fst entry
    ret_entry = maybe (replicate (length rettype) TypeDirect) snd entry
    compile = do
      (inparams, arrayds, args) <- compileInParams params params_entry
      (results, outparams, Destination _ dests) <- compileOutParams rettype ret_entry
      addFParams params
      addArrays arrayds

      let Body _ stms ses = body
      compileStms (freeIn ses) stms $
        forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []

      return (outparams, inparams, results, args)

compileBody :: (Mem lore) => Pattern lore -> Body lore -> ImpM lore r op ()
compileBody pat (Body _ bnds ses) = do
  Destination _ dests <- destinationFromPattern pat
  compileStms (freeIn ses) bnds $
    forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []

compileBody' :: [Param dec] -> Body lore -> ImpM lore r op ()
compileBody' params (Body _ bnds ses) =
  compileStms (freeIn ses) bnds $
    forM_ (zip params ses) $ \(param, se) -> copyDWIM (paramName param) [] se []

compileLoopBody :: Typed dec => [Param dec] -> Body lore -> ImpM lore r op ()
compileLoopBody mergeparams (Body _ bnds ses) = do
  -- We cannot write the results to the merge parameters immediately,
  -- as some of the results may actually *be* merge parameters, and
  -- would thus be clobbered.  Therefore, we first copy to new
  -- variables mirroring the merge parameters, and then copy this
  -- buffer to the merge parameters.  This is efficient, because the
  -- operations are all scalar operations.
  tmpnames <- mapM (newVName . (++ "_tmp") . baseString . paramName) mergeparams
  compileStms (freeIn ses) bnds $ do
    copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p, tmp, se) ->
      case typeOf p of
        Prim pt -> do
          emit $ Imp.DeclareScalar tmp Imp.Nonvolatile pt
          emit $ Imp.SetScalar tmp $ toExp' pt se
          return $ emit $ Imp.SetScalar (paramName p) $ Imp.var tmp pt
        Mem space | Var v <- se -> do
          emit $ Imp.DeclareMem tmp space
          emit $ Imp.SetMem tmp v space
          return $ emit $ Imp.SetMem (paramName p) tmp space
        _ -> return $ return ()
    sequence_ copy_to_merge_params

compileStms :: Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op ()
compileStms alive_after_stms all_stms m = do
  cb <- asks envStmsCompiler
  cb alive_after_stms all_stms m

defCompileStms ::
  (Mem lore, FreeIn op) =>
  Names ->
  Stms lore ->
  ImpM lore r op () ->
  ImpM lore r op ()
defCompileStms alive_after_stms all_stms m =
  -- We keep track of any memory blocks produced by the statements,
  -- and after the last time that memory block is used, we insert a
  -- Free.  This is very conservative, but can cut down on lifetimes
  -- in some cases.
  void $ compileStms' mempty $ stmsToList all_stms
  where
    compileStms' allocs (Let pat aux e : bs) = do
      dVars (Just e) (patternElements pat)

      e_code <-
        localAttrs (stmAuxAttrs aux) $
          collect $ compileExp pat e
      (live_after, bs_code) <- collect' $ compileStms' (patternAllocs pat <> allocs) bs
      let dies_here v =
            not (v `nameIn` live_after)
              && v `nameIn` freeIn e_code
          to_free = S.filter (dies_here . fst) allocs

      emit e_code
      mapM_ (emit . uncurry Imp.Free) to_free
      emit bs_code

      return $ freeIn e_code <> live_after
    compileStms' _ [] = do
      code <- collect m
      emit code
      return $ freeIn code <> alive_after_stms

    patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements
    isMemPatElem pe = case patElemType pe of
      Mem space -> Just (patElemName pe, space)
      _ -> Nothing

compileExp :: Pattern lore -> Exp lore -> ImpM lore r op ()
compileExp pat e = do
  ec <- asks envExpCompiler
  ec pat e

defCompileExp ::
  (Mem lore) =>
  Pattern lore ->
  Exp lore ->
  ImpM lore r op ()
defCompileExp pat (If cond tbranch fbranch _) =
  sIf (toBoolExp cond) (compileBody pat tbranch) (compileBody pat fbranch)
defCompileExp pat (Apply fname args _ _) = do
  dest <- destinationFromPattern pat
  targets <- funcallTargets dest
  args' <- catMaybes <$> mapM compileArg args
  emit $ Imp.Call targets fname args'
  where
    compileArg (se, _) = do
      t <- subExpType se
      case (se, t) of
        (_, Prim pt) -> return $ Just $ Imp.ExpArg $ toExp' pt se
        (Var v, Mem {}) -> return $ Just $ Imp.MemArg v
        _ -> return Nothing
defCompileExp pat (BasicOp op) = defCompileBasicOp pat op
defCompileExp pat (DoLoop ctx val form body) = do
  attrs <- askAttrs
  when ("unroll" `inAttrs` attrs) $
    warn (noLoc :: SrcLoc) [] "#[unroll] on loop with unknown number of iterations." -- FIXME: no location.
  dFParams mergepat
  forM_ merge $ \(p, se) ->
    when ((== 0) $ arrayRank $ paramType p) $
      copyDWIM (paramName p) [] se []

  let doBody = compileLoopBody mergepat body

  case form of
    ForLoop i _ bound loopvars -> do
      let setLoopParam (p, a)
            | Prim _ <- paramType p =
              copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi64 i]
            | otherwise =
              return ()

      bound' <- toExp bound

      dLParams $ map fst loopvars
      sFor' i bound' $
        mapM_ setLoopParam loopvars >> doBody
    WhileLoop cond ->
      sWhile (TPrimExp $ Imp.var cond Bool) doBody

  Destination _ pat_dests <- destinationFromPattern pat
  forM_ (zip pat_dests $ map (Var . paramName . fst) merge) $ \(d, r) ->
    copyDWIMDest d [] r []
  where
    merge = ctx ++ val
    mergepat = map fst merge
defCompileExp pat (Op op) = do
  opc <- asks envOpCompiler
  opc pat op

defCompileBasicOp ::
  Mem lore =>
  Pattern lore ->
  BasicOp ->
  ImpM lore r op ()
defCompileBasicOp (Pattern _ [pe]) (SubExp se) =
  copyDWIM (patElemName pe) [] se []
defCompileBasicOp (Pattern _ [pe]) (Opaque se) =
  copyDWIM (patElemName pe) [] se []
defCompileBasicOp (Pattern _ [pe]) (UnOp op e) = do
  e' <- toExp e
  patElemName pe <~~ Imp.UnOpExp op e'
defCompileBasicOp (Pattern _ [pe]) (ConvOp conv e) = do
  e' <- toExp e
  patElemName pe <~~ Imp.ConvOpExp conv e'
defCompileBasicOp (Pattern _ [pe]) (BinOp bop x y) = do
  x' <- toExp x
  y' <- toExp y
  patElemName pe <~~ Imp.BinOpExp bop x' y'
defCompileBasicOp (Pattern _ [pe]) (CmpOp bop x y) = do
  x' <- toExp x
  y' <- toExp y
  patElemName pe <~~ Imp.CmpOpExp bop x' y'
defCompileBasicOp _ (Assert e msg loc) = do
  e' <- toExp e
  msg' <- traverse toExp msg
  emit $ Imp.Assert e' msg' loc

  attrs <- askAttrs
  when (AttrComp "warn" ["safety_checks"] `inAttrs` attrs) $
    uncurry warn loc "Safety check required at run-time."
defCompileBasicOp (Pattern _ [pe]) (Index src slice)
  | Just idxs <- sliceIndices slice =
    copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt64Exp) idxs
defCompileBasicOp _ Index {} =
  return ()
defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) =
  sUpdate (patElemName pe) (map (fmap toInt64Exp) slice) se
defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do
  ds' <- mapM toExp ds
  is <- replicateM (length ds) (newVName "i")
  copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi64) is) se []
  emit $ foldl (.) id (zipWith Imp.For is ds') copy_elem
defCompileBasicOp _ Scratch {} =
  return ()
defCompileBasicOp (Pattern [] [pe]) (Iota n e s it) = do
  e' <- toExp e
  s' <- toExp s
  sFor "i" (toInt64Exp n) $ \i -> do
    let i' = sExt it $ untyped i
    x <-
      dPrimV "x" $
        TPrimExp $
          BinOpExp (Add it OverflowUndef) e' $
            BinOpExp (Mul it OverflowUndef) i' s'
    copyDWIM (patElemName pe) [DimFix i] (Var (tvVar x)) []
defCompileBasicOp (Pattern _ [pe]) (Copy src) =
  copyDWIM (patElemName pe) [] (Var src) []
defCompileBasicOp (Pattern _ [pe]) (Manifest _ src) =
  copyDWIM (patElemName pe) [] (Var src) []
defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do
  offs_glb <- dPrimV "tmp_offs" 0

  forM_ (x : ys) $ \y -> do
    y_dims <- arrayDims <$> lookupType y
    let rows = case drop i y_dims of
          [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y
          r : _ -> toInt64Exp r
        skip_dims = take i y_dims
        sliceAllDim d = DimSlice 0 d 1
        skip_slices = map (sliceAllDim . toInt64Exp) skip_dims
        destslice = skip_slices ++ [DimSlice (tvExp offs_glb) rows 1]
    copyDWIM (patElemName pe) destslice (Var y) []
    offs_glb <-- tvExp offs_glb + rows
defCompileBasicOp (Pattern [] [pe]) (ArrayLit es _)
  | Just vs@(v : _) <- mapM isLiteral es = do
    dest_mem <- entryArrayLocation <$> lookupArray (patElemName pe)
    dest_space <- entryMemSpace <$> lookupMemory (memLocationName dest_mem)
    let t = primValueType v
    static_array <- newVNameForFun "static_array"
    emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs
    let static_src =
          MemLocation static_array [intConst Int64 $ fromIntegral $ length es] $
            IxFun.iota [fromIntegral $ length es]
        entry = MemVar Nothing $ MemEntry dest_space
    addVar static_array entry
    let slice = [DimSlice 0 (genericLength es) 1]
    copy t dest_mem slice static_src slice
  | otherwise =
    forM_ (zip [0 ..] es) $ \(i, e) ->
      copyDWIM (patElemName pe) [DimFix $ fromInteger i] e []
  where
    isLiteral (Constant v) = Just v
    isLiteral _ = Nothing
defCompileBasicOp _ Rearrange {} =
  return ()
defCompileBasicOp _ Rotate {} =
  return ()
defCompileBasicOp _ Reshape {} =
  return ()
defCompileBasicOp pat e =
  error $
    "ImpGen.defCompileBasicOp: Invalid pattern\n  "
      ++ pretty pat
      ++ "\nfor expression\n  "
      ++ pretty e

-- | Note: a hack to be used only for functions.
addArrays :: [ArrayDecl] -> ImpM lore r op ()
addArrays = mapM_ addArray
  where
    addArray (ArrayDecl name bt location) =
      addVar name $
        ArrayVar
          Nothing
          ArrayEntry
            { entryArrayLocation = location,
              entryArrayElemType = bt
            }

-- | Like 'dFParams', but does not create new declarations.
-- Note: a hack to be used only for functions.
addFParams :: Mem lore => [FParam lore] -> ImpM lore r op ()
addFParams = mapM_ addFParam
  where
    addFParam fparam =
      addVar (paramName fparam) $
        memBoundToVarEntry Nothing $ noUniquenessReturns $ paramDec fparam

-- | Another hack.
addLoopVar :: VName -> IntType -> ImpM lore r op ()
addLoopVar i it = addVar i $ ScalarVar Nothing $ ScalarEntry $ IntType it

dVars ::
  Mem lore =>
  Maybe (Exp lore) ->
  [PatElem lore] ->
  ImpM lore r op ()
dVars e = mapM_ dVar
  where
    dVar = dScope e . scopeOfPatElem

dFParams :: Mem lore => [FParam lore] -> ImpM lore r op ()
dFParams = dScope Nothing . scopeOfFParams

dLParams :: Mem lore => [LParam lore] -> ImpM lore r op ()
dLParams = dScope Nothing . scopeOfLParams

dPrimVol :: String -> PrimType -> Imp.TExp t -> ImpM lore r op (TV t)
dPrimVol name t e = do
  name' <- newVName name
  emit $ Imp.DeclareScalar name' Imp.Volatile t
  addVar name' $ ScalarVar Nothing $ ScalarEntry t
  name' <~~ untyped e
  return $ TV name' t

dPrim_ :: VName -> PrimType -> ImpM lore r op ()
dPrim_ name t = do
  emit $ Imp.DeclareScalar name Imp.Nonvolatile t
  addVar name $ ScalarVar Nothing $ ScalarEntry t

-- | The return type is polymorphic, so there is no guarantee it
-- actually matches the 'PrimType', but at least we have to use it
-- consistently.
dPrim :: String -> PrimType -> ImpM lore r op (TV t)
dPrim name t = do
  name' <- newVName name
  dPrim_ name' t
  return $ TV name' t

dPrimV_ :: VName -> Imp.TExp t -> ImpM lore r op ()
dPrimV_ name e = do
  dPrim_ name t
  TV name t <-- e
  where
    t = primExpType $ untyped e

dPrimV :: String -> Imp.TExp t -> ImpM lore r op (TV t)
dPrimV name e = do
  name' <- dPrim name $ primExpType $ untyped e
  name' <-- e
  return name'

dPrimVE :: String -> Imp.TExp t -> ImpM lore r op (Imp.TExp t)
dPrimVE name e = do
  name' <- dPrim name $ primExpType $ untyped e
  name' <-- e
  return $ tvExp name'

memBoundToVarEntry ::
  Maybe (Exp lore) ->
  MemBound NoUniqueness ->
  VarEntry lore
memBoundToVarEntry e (MemPrim bt) =
  ScalarVar e ScalarEntry {entryScalarType = bt}
memBoundToVarEntry e (MemMem space) =
  MemVar e $ MemEntry space
memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) =
  let location = MemLocation mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun
   in ArrayVar
        e
        ArrayEntry
          { entryArrayLocation = location,
            entryArrayElemType = bt
          }

infoDec ::
  Mem lore =>
  NameInfo lore ->
  MemInfo SubExp NoUniqueness MemBind
infoDec (LetName dec) = dec
infoDec (FParamName dec) = noUniquenessReturns dec
infoDec (LParamName dec) = dec
infoDec (IndexName it) = MemPrim $ IntType it

dInfo ::
  Mem lore =>
  Maybe (Exp lore) ->
  VName ->
  NameInfo lore ->
  ImpM lore r op ()
dInfo e name info = do
  let entry = memBoundToVarEntry e $ infoDec info
  case entry of
    MemVar _ entry' ->
      emit $ Imp.DeclareMem name $ entryMemSpace entry'
    ScalarVar _ entry' ->
      emit $ Imp.DeclareScalar name Imp.Nonvolatile $ entryScalarType entry'
    ArrayVar _ _ ->
      return ()
  addVar name entry

dScope ::
  Mem lore =>
  Maybe (Exp lore) ->
  Scope lore ->
  ImpM lore r op ()
dScope e = mapM_ (uncurry $ dInfo e) . M.toList

dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op ()
dArray name bt shape membind =
  addVar name $
    memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind

everythingVolatile :: ImpM lore r op a -> ImpM lore r op a
everythingVolatile = local $ \env -> env {envVolatility = Imp.Volatile}

-- | Remove the array targets.
funcallTargets :: Destination -> ImpM lore r op [VName]
funcallTargets (Destination _ dests) =
  concat <$> mapM funcallTarget dests
  where
    funcallTarget (ScalarDestination name) =
      return [name]
    funcallTarget (ArrayDestination _) =
      return []
    funcallTarget (MemoryDestination name) =
      return [name]

-- | A typed variable, which we can turn into a typed expression, or
-- use as the target for an assignment.  This is used to aid in type
-- safety when doing code generation, by keeping the types straight.
-- It is still easy to cheat when you need to.
data TV t = TV VName PrimType

-- | Create a typed variable from a name and a dynamic type.  Note
-- that there is no guarantee that the dynamic type corresponds to the
-- inferred static type, but the latter will at least have to be used
-- consistently.
mkTV :: VName -> PrimType -> TV t
mkTV = TV

-- | Convert a typed variable to a size (a SubExp).
tvSize :: TV t -> Imp.DimSize
tvSize = Var . tvVar

-- | Convert a typed variable to a similarly typed expression.
tvExp :: TV t -> Imp.TExp t
tvExp (TV v t) = Imp.TPrimExp $ Imp.var v t

-- | Extract the underlying variable name from a typed variable.
tvVar :: TV t -> VName
tvVar (TV v _) = v

-- | Compile things to 'Imp.Exp'.
class ToExp a where
  -- | Compile to an 'Imp.Exp', where the type (must must still be a
  -- primitive) is deduced monadically.
  toExp :: a -> ImpM lore r op Imp.Exp

  -- | Compile where we know the type in advance.
  toExp' :: PrimType -> a -> Imp.Exp

  toInt64Exp :: a -> Imp.TExp Int64
  toInt64Exp = TPrimExp . toExp' int64

  toBoolExp :: a -> Imp.TExp Bool
  toBoolExp = TPrimExp . toExp' Bool

instance ToExp SubExp where
  toExp (Constant v) =
    return $ Imp.ValueExp v
  toExp (Var v) =
    lookupVar v >>= \case
      ScalarVar _ (ScalarEntry pt) ->
        return $ Imp.var v pt
      _ -> error $ "toExp SubExp: SubExp is not a primitive type: " ++ pretty v

  toExp' _ (Constant v) = Imp.ValueExp v
  toExp' t (Var v) = Imp.var v t

instance ToExp (PrimExp VName) where
  toExp = pure . fmap Imp.ScalarVar
  toExp' _ = fmap Imp.ScalarVar

addVar :: VName -> VarEntry lore -> ImpM lore r op ()
addVar name entry =
  modify $ \s -> s {stateVTable = M.insert name entry $ stateVTable s}

localDefaultSpace :: Imp.Space -> ImpM lore r op a -> ImpM lore r op a
localDefaultSpace space = local (\env -> env {envDefaultSpace = space})

askFunction :: ImpM lore r op (Maybe Name)
askFunction = asks envFunction

-- | Generate a 'VName', prefixed with 'askFunction' if it exists.
newVNameForFun :: String -> ImpM lore r op VName
newVNameForFun s = do
  fname <- fmap nameToString <$> askFunction
  newVName $ maybe "" (++ ".") fname ++ s

-- | Generate a 'Name', prefixed with 'askFunction' if it exists.
nameForFun :: String -> ImpM lore r op Name
nameForFun s = do
  fname <- askFunction
  return $ maybe "" (<> ".") fname <> nameFromString s

askEnv :: ImpM lore r op r
askEnv = asks envEnv

localEnv :: (r -> r) -> ImpM lore r op a -> ImpM lore r op a
localEnv f = local $ \env -> env {envEnv = f $ envEnv env}

-- | The active attributes, including those for the statement
-- currently being compiled.
askAttrs :: ImpM lore r op Attrs
askAttrs = asks envAttrs

-- | Add more attributes to what is returning by 'askAttrs'.
localAttrs :: Attrs -> ImpM lore r op a -> ImpM lore r op a
localAttrs attrs = local $ \env -> env {envAttrs = attrs <> envAttrs env}

localOps :: Operations lore r op -> ImpM lore r op a -> ImpM lore r op a
localOps ops = local $ \env ->
  env
    { envExpCompiler = opsExpCompiler ops,
      envStmsCompiler = opsStmsCompiler ops,
      envCopyCompiler = opsCopyCompiler ops,
      envOpCompiler = opsOpCompiler ops,
      envAllocCompilers = opsAllocCompilers ops
    }

-- | Get the current symbol table.
getVTable :: ImpM lore r op (VTable lore)
getVTable = gets stateVTable

putVTable :: VTable lore -> ImpM lore r op ()
putVTable vtable = modify $ \s -> s {stateVTable = vtable}

-- | Run an action with a modified symbol table.  All changes to the
-- symbol table will be reverted once the action is done!
localVTable :: (VTable lore -> VTable lore) -> ImpM lore r op a -> ImpM lore r op a
localVTable f m = do
  old_vtable <- getVTable
  putVTable $ f old_vtable
  a <- m
  putVTable old_vtable
  return a

lookupVar :: VName -> ImpM lore r op (VarEntry lore)
lookupVar name = do
  res <- gets $ M.lookup name . stateVTable
  case res of
    Just entry -> return entry
    _ -> error $ "Unknown variable: " ++ pretty name

lookupArray :: VName -> ImpM lore r op ArrayEntry
lookupArray name = do
  res <- lookupVar name
  case res of
    ArrayVar _ entry -> return entry
    _ -> error $ "ImpGen.lookupArray: not an array: " ++ pretty name

lookupMemory :: VName -> ImpM lore r op MemEntry
lookupMemory name = do
  res <- lookupVar name
  case res of
    MemVar _ entry -> return entry
    _ -> error $ "Unknown memory block: " ++ pretty name

destinationFromPattern :: Mem lore => Pattern lore -> ImpM lore r op Destination
destinationFromPattern pat =
  fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $
    patternElements pat
  where
    inspect patElem = do
      let name = patElemName patElem
      entry <- lookupVar name
      case entry of
        ArrayVar _ (ArrayEntry MemLocation {} _) ->
          return $ ArrayDestination Nothing
        MemVar {} ->
          return $ MemoryDestination name
        ScalarVar {} ->
          return $ ScalarDestination name

fullyIndexArray ::
  VName ->
  [Imp.TExp Int64] ->
  ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64))
fullyIndexArray name indices = do
  arr <- lookupArray name
  fullyIndexArray' (entryArrayLocation arr) indices

fullyIndexArray' ::
  MemLocation ->
  [Imp.TExp Int64] ->
  ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64))
fullyIndexArray' (MemLocation mem _ ixfun) indices = do
  space <- entryMemSpace <$> lookupMemory mem
  let indices' = case space of
        ScalarSpace ds _ ->
          let (zero_is, is) = splitFromEnd (length ds) indices
           in map (const 0) zero_is ++ is
        _ -> indices
  return
    ( mem,
      space,
      elements $ IxFun.index ixfun indices'
    )

-- More complicated read/write operations that use index functions.

copy :: CopyCompiler lore r op
copy bt dest destslice src srcslice = do
  cc <- asks envCopyCompiler
  cc bt dest destslice src srcslice

-- | Is this copy really a mapping with transpose?
isMapTransposeCopy ::
  PrimType ->
  MemLocation ->
  Slice (Imp.TExp Int64) ->
  MemLocation ->
  Slice (Imp.TExp Int64) ->
  Maybe
    ( Imp.TExp Int64,
      Imp.TExp Int64,
      Imp.TExp Int64,
      Imp.TExp Int64,
      Imp.TExp Int64
    )
isMapTransposeCopy
  bt
  (MemLocation _ _ destIxFun)
  destslice
  (MemLocation _ _ srcIxFun)
  srcslice
    | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun' bt_size,
      (perm, destshape) <- unzip perm_and_destshape,
      Just src_offset <- IxFun.linearWithOffset srcIxFun' bt_size,
      Just (r1, r2, _) <- isMapTranspose perm =
      isOk destshape swap r1 r2 dest_offset src_offset
    | Just dest_offset <- IxFun.linearWithOffset destIxFun' bt_size,
      Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun' bt_size,
      (perm, srcshape) <- unzip perm_and_srcshape,
      Just (r1, r2, _) <- isMapTranspose perm =
      isOk srcshape id r1 r2 dest_offset src_offset
    | otherwise =
      Nothing
    where
      bt_size = primByteSize bt
      swap (x, y) = (y, x)

      destIxFun' = IxFun.slice destIxFun destslice
      srcIxFun' = IxFun.slice srcIxFun srcslice

      isOk shape f r1 r2 dest_offset src_offset = do
        let (num_arrays, size_x, size_y) = getSizes shape f r1 r2
        return
          ( dest_offset,
            src_offset,
            num_arrays,
            size_x,
            size_y
          )

      getSizes shape f r1 r2 =
        let (mapped, notmapped) = splitAt r1 shape
            (pretrans, posttrans) = f $ splitAt r2 notmapped
         in (product mapped, product pretrans, product posttrans)

mapTransposeName :: PrimType -> String
mapTransposeName bt = "map_transpose_" ++ pretty bt

mapTransposeForType :: PrimType -> ImpM lore r op Name
mapTransposeForType bt = do
  let fname = nameFromString $ "builtin#" <> mapTransposeName bt

  exists <- hasFunction fname
  unless exists $ emitFunction fname $ mapTransposeFunction fname bt

  return fname

-- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'.
defaultCopy :: CopyCompiler lore r op
defaultCopy pt dest destslice src srcslice
  | Just
      ( destoffset,
        srcoffset,
        num_arrays,
        size_x,
        size_y
        ) <-
      isMapTransposeCopy pt dest destslice src srcslice = do
    fname <- mapTransposeForType pt
    emit $
      Imp.Call
        []
        fname
        $ transposeArgs
          pt
          destmem
          (bytes destoffset)
          srcmem
          (bytes srcoffset)
          num_arrays
          size_x
          size_y
  | Just destoffset <-
      IxFun.linearWithOffset (IxFun.slice dest_ixfun destslice) pt_size,
    Just srcoffset <-
      IxFun.linearWithOffset (IxFun.slice src_ixfun srcslice) pt_size = do
    srcspace <- entryMemSpace <$> lookupMemory srcmem
    destspace <- entryMemSpace <$> lookupMemory destmem
    if isScalarSpace srcspace || isScalarSpace destspace
      then copyElementWise pt dest destslice src srcslice
      else
        emit $
          Imp.Copy
            destmem
            (bytes destoffset)
            destspace
            srcmem
            (bytes srcoffset)
            srcspace
            $ num_elems `withElemType` pt
  | otherwise =
    copyElementWise pt dest destslice src srcslice
  where
    pt_size = primByteSize pt
    num_elems = Imp.elements $ product $ sliceDims srcslice

    MemLocation destmem _ dest_ixfun = dest
    MemLocation srcmem _ src_ixfun = src

    isScalarSpace ScalarSpace {} = True
    isScalarSpace _ = False

copyElementWise :: CopyCompiler lore r op
copyElementWise bt dest destslice src srcslice = do
  let bounds = sliceDims srcslice
  is <- replicateM (length bounds) (newVName "i")
  let ivars = map Imp.vi64 is
  (destmem, destspace, destidx) <-
    fullyIndexArray' dest $ fixSlice destslice ivars
  (srcmem, srcspace, srcidx) <-
    fullyIndexArray' src $ fixSlice srcslice ivars
  vol <- asks envVolatility
  emit $
    foldl (.) id (zipWith Imp.For is $ map untyped bounds) $
      Imp.Write destmem destidx bt destspace vol $
        Imp.index srcmem srcidx bt srcspace vol

-- | Copy from here to there; both destination and source may be
-- indexeded.
copyArrayDWIM ::
  PrimType ->
  MemLocation ->
  [DimIndex (Imp.TExp Int64)] ->
  MemLocation ->
  [DimIndex (Imp.TExp Int64)] ->
  ImpM lore r op (Imp.Code op)
copyArrayDWIM
  bt
  destlocation@(MemLocation _ destshape _)
  destslice
  srclocation@(MemLocation _ srcshape _)
  srcslice
    | Just destis <- mapM dimFix destslice,
      Just srcis <- mapM dimFix srcslice,
      length srcis == length srcshape,
      length destis == length destshape = do
      (targetmem, destspace, targetoffset) <-
        fullyIndexArray' destlocation destis
      (srcmem, srcspace, srcoffset) <-
        fullyIndexArray' srclocation srcis
      vol <- asks envVolatility
      return $
        Imp.Write targetmem targetoffset bt destspace vol $
          Imp.index srcmem srcoffset bt srcspace vol
    | otherwise = do
      let destslice' =
            fullSliceNum (map toInt64Exp destshape) destslice
          srcslice' =
            fullSliceNum (map toInt64Exp srcshape) srcslice
          destrank = length $ sliceDims destslice'
          srcrank = length $ sliceDims srcslice'
      if destrank /= srcrank
        then
          error $
            "copyArrayDWIM: cannot copy to "
              ++ pretty (memLocationName destlocation)
              ++ " from "
              ++ pretty (memLocationName srclocation)
              ++ " because ranks do not match ("
              ++ pretty destrank
              ++ " vs "
              ++ pretty srcrank
              ++ ")"
        else
          if destlocation == srclocation && destslice' == srcslice'
            then return mempty -- Copy would be no-op.
            else collect $ copy bt destlocation destslice' srclocation srcslice'

-- | Like 'copyDWIM', but the target is a 'ValueDestination'
-- instead of a variable name.
copyDWIMDest ::
  ValueDestination ->
  [DimIndex (Imp.TExp Int64)] ->
  SubExp ->
  [DimIndex (Imp.TExp Int64)] ->
  ImpM lore r op ()
copyDWIMDest _ _ (Constant v) (_ : _) =
  error $
    unwords ["copyDWIMDest: constant source", pretty v, "cannot be indexed."]
copyDWIMDest pat dest_slice (Constant v) [] =
  case mapM dimFix dest_slice of
    Nothing ->
      error $
        unwords ["copyDWIMDest: constant source", pretty v, "with slice destination."]
    Just dest_is ->
      case pat of
        ScalarDestination name ->
          emit $ Imp.SetScalar name $ Imp.ValueExp v
        MemoryDestination {} ->
          error $
            unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."]
        ArrayDestination (Just dest_loc) -> do
          (dest_mem, dest_space, dest_i) <-
            fullyIndexArray' dest_loc dest_is
          vol <- asks envVolatility
          emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v
        ArrayDestination Nothing ->
          error "copyDWIMDest: ArrayDestination Nothing"
  where
    bt = primValueType v
copyDWIMDest dest dest_slice (Var src) src_slice = do
  src_entry <- lookupVar src
  case (dest, src_entry) of
    (MemoryDestination mem, MemVar _ (MemEntry space)) ->
      emit $ Imp.SetMem mem src space
    (MemoryDestination {}, _) ->
      error $
        unwords ["copyDWIMDest: cannot write", pretty src, "to memory destination."]
    (_, MemVar {}) ->
      error $
        unwords ["copyDWIMDest: source", pretty src, "is a memory block."]
    (_, ScalarVar _ (ScalarEntry _))
      | not $ null src_slice ->
        error $
          unwords ["copyDWIMDest: prim-typed source", pretty src, "with slice", pretty src_slice]
    (ScalarDestination name, _)
      | not $ null dest_slice ->
        error $
          unwords ["copyDWIMDest: prim-typed target", pretty name, "with slice", pretty dest_slice]
    (ScalarDestination name, ScalarVar _ (ScalarEntry pt)) ->
      emit $ Imp.SetScalar name $ Imp.var src pt
    (ScalarDestination name, ArrayVar _ arr)
      | Just src_is <- mapM dimFix src_slice,
        length src_slice == length (entryArrayShape arr) -> do
        let bt = entryArrayElemType arr
        (mem, space, i) <-
          fullyIndexArray' (entryArrayLocation arr) src_is
        vol <- asks envVolatility
        emit $ Imp.SetScalar name $ Imp.index mem i bt space vol
      | otherwise ->
        error $
          unwords
            [ "copyDWIMDest: prim-typed target",
              pretty name,
              "and array-typed source",
              pretty src,
              "with slice",
              pretty src_slice
            ]
    (ArrayDestination (Just dest_loc), ArrayVar _ src_arr) -> do
      let src_loc = entryArrayLocation src_arr
          bt = entryArrayElemType src_arr
      emit =<< copyArrayDWIM bt dest_loc dest_slice src_loc src_slice
    (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt))
      | Just dest_is <- mapM dimFix dest_slice -> do
        (dest_mem, dest_space, dest_i) <- fullyIndexArray' dest_loc dest_is
        vol <- asks envVolatility
        emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt)
      | otherwise ->
        error $
          unwords
            [ "copyDWIMDest: array-typed target and prim-typed source",
              pretty src,
              "with slice",
              pretty dest_slice
            ]
    (ArrayDestination Nothing, _) ->
      return () -- Nothing to do; something else set some memory
      -- somewhere.

-- | Copy from here to there; both destination and source be
-- indexeded.  If so, they better be arrays of enough dimensions.
-- This function will generally just Do What I Mean, and Do The Right
-- Thing.  Both destination and source must be in scope.
copyDWIM ::
  VName ->
  [DimIndex (Imp.TExp Int64)] ->
  SubExp ->
  [DimIndex (Imp.TExp Int64)] ->
  ImpM lore r op ()
copyDWIM dest dest_slice src src_slice = do
  dest_entry <- lookupVar dest
  let dest_target =
        case dest_entry of
          ScalarVar _ _ ->
            ScalarDestination dest
          ArrayVar _ (ArrayEntry (MemLocation mem shape ixfun) _) ->
            ArrayDestination $ Just $ MemLocation mem shape ixfun
          MemVar _ _ ->
            MemoryDestination dest
  copyDWIMDest dest_target dest_slice src src_slice

-- | As 'copyDWIM', but implicitly 'DimFix'es the indexes.
copyDWIMFix ::
  VName ->
  [Imp.TExp Int64] ->
  SubExp ->
  [Imp.TExp Int64] ->
  ImpM lore r op ()
copyDWIMFix dest dest_is src src_is =
  copyDWIM dest (map DimFix dest_is) src (map DimFix src_is)

-- | @compileAlloc pat size space@ allocates @n@ bytes of memory in @space@,
-- writing the result to @dest@, which must be a single
-- 'MemoryDestination',
compileAlloc ::
  Mem lore =>
  Pattern lore ->
  SubExp ->
  Space ->
  ImpM lore r op ()
compileAlloc (Pattern [] [mem]) e space = do
  let e' = Imp.bytes $ toInt64Exp e
  allocator <- asks $ M.lookup space . envAllocCompilers
  case allocator of
    Nothing -> emit $ Imp.Allocate (patElemName mem) e' space
    Just allocator' -> allocator' (patElemName mem) e'
compileAlloc pat _ _ =
  error $ "compileAlloc: Invalid pattern: " ++ pretty pat

-- | The number of bytes needed to represent the array in a
-- straightforward contiguous format, as an t'Int64' expression.
typeSize :: Type -> Count Bytes (Imp.TExp Int64)
typeSize t =
  Imp.bytes $
    isInt64 (Imp.LeafExp (Imp.SizeOf $ elemType t) int64)
      * product (map toInt64Exp (arrayDims t))

--- Building blocks for constructing code.

sFor' :: VName -> Imp.Exp -> ImpM lore r op () -> ImpM lore r op ()
sFor' i bound body = do
  let it = case primExpType bound of
        IntType bound_t -> bound_t
        t -> error $ "sFor': bound " ++ pretty bound ++ " is of type " ++ pretty t
  addLoopVar i it
  body' <- collect body
  emit $ Imp.For i bound body'

sFor :: String -> Imp.TExp t -> (Imp.TExp t -> ImpM lore r op ()) -> ImpM lore r op ()
sFor i bound body = do
  i' <- newVName i
  sFor' i' (untyped bound) $
    body $ TPrimExp $ Imp.var i' $ primExpType $ untyped bound

sWhile :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op ()
sWhile cond body = do
  body' <- collect body
  emit $ Imp.While cond body'

sComment :: String -> ImpM lore r op () -> ImpM lore r op ()
sComment s code = do
  code' <- collect code
  emit $ Imp.Comment s code'

sIf :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () -> ImpM lore r op ()
sIf cond tbranch fbranch = do
  tbranch' <- collect tbranch
  fbranch' <- collect fbranch
  emit $ Imp.If cond tbranch' fbranch'

sWhen :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op ()
sWhen cond tbranch = sIf cond tbranch (return ())

sUnless :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op ()
sUnless cond = sIf cond (return ())

sOp :: op -> ImpM lore r op ()
sOp = emit . Imp.Op

sDeclareMem :: String -> Space -> ImpM lore r op VName
sDeclareMem name space = do
  name' <- newVName name
  emit $ Imp.DeclareMem name' space
  addVar name' $ MemVar Nothing $ MemEntry space
  return name'

sAlloc_ :: VName -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op ()
sAlloc_ name' size' space = do
  allocator <- asks $ M.lookup space . envAllocCompilers
  case allocator of
    Nothing -> emit $ Imp.Allocate name' size' space
    Just allocator' -> allocator' name' size'

sAlloc :: String -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op VName
sAlloc name size space = do
  name' <- sDeclareMem name space
  sAlloc_ name' size space
  return name'

sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op VName
sArray name bt shape membind = do
  name' <- newVName name
  dArray name' bt shape membind
  return name'

-- | Declare an array in row-major order in the given memory block.
sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM lore r op VName
sArrayInMem name pt shape mem =
  sArray name pt shape $
    ArrayIn mem $
      IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape

-- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified.
sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r op VName
sAllocArrayPerm name pt shape space perm = do
  let permuted_dims = rearrangeShape perm $ shapeDims shape
  mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space
  let iota_ixfun = IxFun.iota $ map (isInt64 . primExpFromSubExp int64) permuted_dims
  sArray name pt shape $
    ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm

-- | Uses linear/iota index function.
sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore r op VName
sAllocArray name pt shape space =
  sAllocArrayPerm name pt shape space [0 .. shapeRank shape -1]

-- | Uses linear/iota index function.
sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore r op VName
sStaticArray name space pt vs = do
  let num_elems = case vs of
        Imp.ArrayValues vs' -> length vs'
        Imp.ArrayZeros n -> fromIntegral n
      shape = Shape [intConst Int64 $ toInteger num_elems]
  mem <- newVNameForFun $ name ++ "_mem"
  emit $ Imp.DeclareArray mem space pt vs
  addVar mem $ MemVar Nothing $ MemEntry space
  sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems]

sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM lore r op ()
sWrite arr is v = do
  (mem, space, offset) <- fullyIndexArray arr is
  vol <- asks envVolatility
  emit $ Imp.Write mem offset (primExpType v) space vol v

sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM lore r op ()
sUpdate arr slice v = copyDWIM arr slice v []

sLoopNest ::
  Shape ->
  ([Imp.TExp Int64] -> ImpM lore r op ()) ->
  ImpM lore r op ()
sLoopNest = sLoopNest' [] . shapeDims
  where
    sLoopNest' is [] f = f $ reverse is
    sLoopNest' is (d : ds) f =
      sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f

-- | Untyped assignment.
(<~~) :: VName -> Imp.Exp -> ImpM lore r op ()
x <~~ e = emit $ Imp.SetScalar x e

infixl 3 <~~

-- | Typed assignment.
(<--) :: TV t -> Imp.TExp t -> ImpM lore r op ()
TV x _ <-- e = emit $ Imp.SetScalar x $ untyped e

infixl 3 <--

-- | Constructing an ad-hoc function that does not
-- correspond to any of the IR functions in the input program.
function ::
  Name ->
  [Imp.Param] ->
  [Imp.Param] ->
  ImpM lore r op () ->
  ImpM lore r op ()
function fname outputs inputs m = local newFunction $ do
  body <- collect $ do
    mapM_ addParam $ outputs ++ inputs
    m
  emitFunction fname $ Imp.Function False outputs inputs body [] []
  where
    addParam (Imp.MemParam name space) =
      addVar name $ MemVar Nothing $ MemEntry space
    addParam (Imp.ScalarParam name bt) =
      addVar name $ ScalarVar Nothing $ ScalarEntry bt
    newFunction env = env {envFunction = Just fname}