{-# LANGUAGE TupleSections #-}

-- |Vectorise variables and literals.

module Vectorise.Var
  ( vectBndr
  , vectBndrNew
  , vectBndrIn
  , vectBndrNewIn
  , vectBndrsIn
  , vectVar
  , vectConst
  )
where

import GhcPrelude

import Vectorise.Utils
import Vectorise.Monad
import Vectorise.Env
import Vectorise.Vect
import Vectorise.Type.Type
import CoreSyn
import Type
import VarEnv
import Id
import FastString

-- Binders ----------------------------------------------------------------------------------------

-- |Vectorise a binder variable, along with its attached type.
--
vectBndr :: Var -> VM VVar
vectBndr v
 = do (vty, lty) <- vectAndLiftType (idType v)
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty

      updLEnv (mapTo vv lv)

      return  (vv, lv)
  where
    mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }

-- |Vectorise a binder variable, along with its attached type, but give the result a new name.
--
vectBndrNew :: Var -> FastString -> VM VVar
vectBndrNew v fs
 = do vty <- vectType (idType v)
      vv  <- newLocalVVar fs vty
      updLEnv (upd vv)
      return vv
  where
    upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }

-- |Vectorise a binder then run a computation with that binder in scope.
--
vectBndrIn :: Var -> VM a -> VM (VVar, a)
vectBndrIn v p
 = localV
 $ do vv <- vectBndr v
      x <- p
      return (vv, x)

-- |Vectorise a binder, give it a new name, then run a computation with that binder in scope.
--
vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
vectBndrNewIn v fs p
 = localV
 $ do vv <- vectBndrNew v fs
      x  <- p
      return (vv, x)

-- |Vectorise some binders, then run a computation with them in scope.
--
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
vectBndrsIn vs p
 = localV
 $ do vvs <- mapM vectBndr vs
      x   <- p
      return (vvs, x)


-- Variables --------------------------------------------------------------------------------------

-- |Vectorise a variable, producing the vectorised and lifted versions.
--
vectVar :: Var -> VM VExpr
vectVar var
  = do { vVar <- lookupVar var
       ; case vVar of
           Local (vv, lv) -> return (Var vv, Var lv) -- local variables have a vect & lifted version
           Global vv      -> vectConst (Var vv)      -- global variables get replicated
       }


-- Constants --------------------------------------------------------------------------------------

-- |Constants are lifted by replication along the integer context in the `VM` state for the number
-- of elements in the result array.
--
vectConst :: CoreExpr -> VM VExpr
vectConst c = (c,) <$> liftPD c