content: -- Copyright 2020 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
  
  module Imp.Builder ( ISubstBuilderT, ISubstEnv (..)
                     , runISubstBuilderT, liftSE
                     , emit, freshIVar, extendValSubst
                     , buildScoped
                     -- IR builders
                     , iadd, imul
                     , alloc, ptrOffset
                     -- traversal
                     , traverseImpModule, traverseImpFunction
                     , traverseImpBlock, evalImpBlock
                     , traverseImpDecl, traverseImpInstr
                     , traverseIExpr, traverseIFunVar
                     , ITraversalDef, substTraversalDef
                     ) where
  
  import Control.Monad.Reader
  
  import Env
  import Cat
  import Syntax
  import Imp
  import Util (bindM2)
  
  -- XXX: Scope is actually global within each function
  data IBuilderEnv = IBuilderEnv
    { scope :: Env ()
    , blockDecls :: Nest ImpDecl
    }
  data ISubstEnv = ISubstEnv
    { valSubst  :: Env IExpr
    , funcSubst :: Env IFunVar
    }
  
  type IBuilderT      m = CatT IBuilderEnv m
  type ISubstT      m = ReaderT ISubstEnv m
  type ISubstBuilderT m = IBuilderT (ISubstT m)
  
  runIBuilderT :: Monad m => IBuilderT m a -> m a
  runIBuilderT m = fst <$> runCatT m mempty
  
  runISubstT :: Monad m => ISubstEnv -> ISubstT m a -> m a
  runISubstT env m = runReaderT m env
  
  runISubstBuilderT :: Monad m => ISubstEnv -> ISubstBuilderT m a -> m a
  runISubstBuilderT env = (runISubstT env) . runIBuilderT
  
  liftSE :: Monad m => m a -> ISubstBuilderT m a
  liftSE = lift . lift
  
  extendScope :: Monad m => Env a -> IBuilderT m ()
  extendScope s = extend $ IBuilderEnv (fmap (const ()) s) mempty
  
  emit :: Monad m => ImpInstr -> IBuilderT m [IExpr]
  emit instr = do
    vs <- traverse (freshIVar . Ignore) $ impInstrTypes instr
    emitTo vs instr
  
  emitTo :: Monad m => [IVar] -> ImpInstr -> IBuilderT m [IExpr]
  emitTo bs instr = do
    extend $ mempty { blockDecls = (Nest (ImpLet (fmap Bind bs) instr) Empty) }
    return $ fmap IVar bs
  
  instance Semigroup IBuilderEnv where
    (IBuilderEnv s d) <> (IBuilderEnv s' d') = IBuilderEnv (s <> s') (d <> d')
  
  instance Monoid IBuilderEnv where
    mempty = IBuilderEnv mempty mempty
  
  instance Semigroup ISubstEnv where
    (ISubstEnv v f) <> (ISubstEnv v' f') = ISubstEnv (v <> v') (f <> f')
  
  instance Monoid ISubstEnv where
    mempty = ISubstEnv mempty mempty
  
  -- === Imp IR builders ===
  
  ptrOffset :: Monad m => IExpr -> IExpr -> IBuilderT m IExpr
  ptrOffset ptr off = liftM head $ emit $ IPrimOp $ PtrOffset ptr off
  
  imul :: Monad m => IExpr -> IExpr -> IBuilderT m IExpr
  imul x y = liftM head $ emit $ IPrimOp $ ScalarBinOp IMul x y
  
  iadd :: Monad m => IExpr -> IExpr -> IBuilderT m IExpr
  iadd x y = liftM head $ emit $ IPrimOp $ ScalarBinOp IAdd x y
  
  alloc :: Monad m => AddressSpace -> IType -> IExpr -> IBuilderT m IExpr
  alloc addrSpc ty size = liftM head $ emit $ Alloc addrSpc ty size
  
  -- === Imp IR traversal ===
  
  type ITraversalDef m = ( ImpDecl  -> ISubstBuilderT m (Env IExpr)
                         , ImpInstr -> ISubstBuilderT m ImpInstr
                         )
  
  substTraversalDef :: Monad m => ITraversalDef m
  substTraversalDef = ( traverseImpDecl  substTraversalDef
                      , traverseImpInstr substTraversalDef )
  
  traverseImpModule :: forall m. Monad m
                    => (Env IFunVar -> ImpFunction -> m ImpFunction) -> ImpModule -> m ImpModule
  traverseImpModule fTrav (ImpModule funcs) = ImpModule . fst <$> runCatT (traverse go funcs) mempty
    where
      go :: ImpFunction -> CatT (Env IFunVar) m ImpFunction
      go f = do
        fenv <- look
        f' <- lift $ fTrav fenv f
        extend $ impFunVar f @> impFunVar f'
        return f'
  
  traverseImpFunction :: Monad m => ITraversalDef m -> Env IFunVar -> ImpFunction -> m ImpFunction
  traverseImpFunction _   _    (FFIFunction f             ) = return $ FFIFunction f
  traverseImpFunction def fenv (ImpFunction name args body) = runISubstBuilderT env $ do
    extendScope $ foldMap binderAsEnv args
    body' <- extendValSubst (foldMap argSub args) $ traverseImpBlock def body
    return $ ImpFunction name args body'
    where
      argSub b = case b of
        Ignore _ -> mempty
        Bind   v -> v @> IVar v
      env = ISubstEnv mempty fenv
  
  traverseImpBlock :: Monad m => ITraversalDef m -> ImpBlock -> ISubstBuilderT m ImpBlock
  traverseImpBlock def block = buildScoped $ evalImpBlock def block
  
  evalImpBlock :: Monad m => ITraversalDef m -> ImpBlock -> ISubstBuilderT m [IExpr]
  evalImpBlock def@(fDecl, _) (ImpBlock decls results) = do
    case decls of
      Nest decl rest -> do
        env' <- fDecl decl
        extendValSubst env' $ evalImpBlock def $ ImpBlock rest results
      Empty -> traverse traverseIExpr results
  
  traverseImpDecl :: Monad m => ITraversalDef m -> ImpDecl -> ISubstBuilderT m (Env IExpr)
  traverseImpDecl (_, fInstr) (ImpLet bs instr) = do
    vs <- bindM2 emitTo (traverse freshIVar bs) (fInstr instr)
    return $ newEnv bs vs
  
  traverseImpInstr :: Monad m => ITraversalDef m -> ImpInstr -> ISubstBuilderT m ImpInstr
  traverseImpInstr def instr = case instr of
    IFor dir b size body -> do
      b' <- freshIVar b
      IFor dir (Bind b') <$> traverseIExpr size
                         <*> (extendValSubst (b @> IVar b') $ traverseImpBlock def body)
    IWhile body -> IWhile <$> traverseImpBlock def body
    ICond cond tb fb -> ICond <$> traverseIExpr cond
                              <*> traverseImpBlock def tb
                              <*> traverseImpBlock def fb
    IQueryParallelism f s -> IQueryParallelism <$> traverseIFunVar f
                                               <*> traverseIExpr s
    ISyncWorkgroup   -> return ISyncWorkgroup
    ILaunch f s args -> ILaunch <$> traverseIFunVar f
                                <*> traverseIExpr s
                                <*> traverse traverseIExpr args
    ICall f args  -> ICall <$> traverseIFunVar f <*> traverse traverseIExpr args
    Store dst val -> Store <$> traverseIExpr dst <*> traverseIExpr val
    Alloc addrSpace ty size -> Alloc addrSpace ty <$> traverseIExpr size
    MemCopy dst src n       -> MemCopy <$> traverseIExpr dst
                                       <*> traverseIExpr src
                                       <*> traverseIExpr n
    Free dst       -> Free <$> traverseIExpr dst
    IThrowError    -> return IThrowError
    ICastOp ty val -> ICastOp ty <$> traverseIExpr val
    IPrimOp op     -> IPrimOp <$> traverse traverseIExpr op
  
  traverseIExpr :: Monad m => IExpr -> ISubstBuilderT m IExpr
  traverseIExpr (ILit l) = return $ ILit l
  traverseIExpr (IVar v) = (!v) <$> asks valSubst
  
  traverseIFunVar :: Monad m => IFunVar -> ISubstBuilderT m IFunVar
  traverseIFunVar fv = (!fv) <$> asks funcSubst
  
  freshIVar :: Monad m => IBinder -> IBuilderT m IVar
  freshIVar b = do
    let nameHint = case b of
                     Bind (name:>_) -> name
                     Ignore _       -> "v"
    name <- genFresh nameHint <$> looks scope
    extendScope $ name @> ()
    return $ name :> binderAnn b
  
  buildScoped :: Monad m => IBuilderT m [IExpr] -> IBuilderT m ImpBlock
  buildScoped m = do
    (results, IBuilderEnv scopeExt decls) <- scoped m
    extend $ IBuilderEnv scopeExt mempty  -- Names are global in Imp IR
    return $ ImpBlock decls results
  
  extendValSubst :: Monad m => Env IExpr -> ISubstBuilderT m a -> ISubstBuilderT m a
  extendValSubst s = local (\env -> env { valSubst = valSubst env <> s })
  