-- 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 UndecidableInstances #-}

module Lower
  ( lowerFullySequential, DestBlock
  ) where

import Prelude hiding ((.))
import Data.Functor
import Data.Maybe (fromJust)
import Control.Category
import Control.Monad.Reader
import Unsafe.Coerce

import Builder
import Core
import Imp
import CheapReduction
import IRVariants
import Name
import Subst
import QueryType
import Types.Core
import Types.Top
import Types.Primitives
import Util (enumerate)

-- === For loop resolution ===

-- The `lowerFullySequential` pass makes two related changes:
-- - All `for` loops become `seq` loops
-- - Arrays are now explicily allocated as `Dest`s (whither the
--   content is written in due course)
--
-- We also perform destination passing here, to elide as many copies
-- as possible.
--
-- The idea for multithreading parallelism is to add IR elements that
-- specify parallelization strategy (`Seq` specifies sequential
-- execution) and add lowerings similar to lowerFullySequential that
-- choose which one(s) to use.

-- The `For` constructor of `PrimHof` disappears from the IR due to
-- this pass, and the `Seq`, `AllocDest`, `Place`, `RememberDest`, and
-- `Freeze` constructors appear.
--
-- All the `Place`s in the resulting IR are non-conflicting: every
-- array element is only written to once.
--
-- The traversal for a block or subexpression returns an `Atom`
-- representing the returned value.  If a destination was passed in,
-- the traversal is also responsible for arranging to write that value
-- into that destination (possibly by forwarding (a part of) the
-- destination to a sub-block or sub-expression, hence "desintation
-- passing style").

type DestBlock = Abs (SBinder) SExpr

lowerFullySequential :: EnvReader m => Bool -> STopLam n -> m n (STopLam n)
lowerFullySequential wantDestStyle (TopLam False piTy (LamExpr bs body)) = liftEnvReaderM do
  lam <- refreshAbs (Abs bs body) \bs' body' ->
    liftAtomSubstBuilder case wantDestStyle of
      True -> do
        xs <- bindersToAtoms bs'
        EffTy _ resultTy <- instantiate (sink piTy) xs
        let resultDestTy = RawRefTy resultTy
        withFreshBinder "ans" resultDestTy \destBinder -> do
          let dest = toAtom $ binderVar destBinder
          LamExpr (bs' >>> UnaryNest destBinder) <$> buildBlock do
            lowerExpr (Just (sink dest)) body' $> UnitVal
      False -> LamExpr bs' <$> buildBlock (lowerExpr Nothing body')
  piTy' <- getLamExprType lam
  return $ TopLam wantDestStyle piTy' lam
lowerFullySequential _ (TopLam True _ _) = error "already in destination style"

data LowerTag
type LowerM = AtomSubstBuilder LowerTag SimpIR

instance NonAtomRenamer (LowerM i o) i o where renameN = substM

instance ExprVisitorEmits (LowerM i o) SimpIR i o where
  visitExprEmits = lowerExpr Nothing

instance Visitor (LowerM i o) SimpIR i o where
  visitAtom = visitAtomDefault
  visitType = visitTypeDefault
  visitPi   = visitPiDefault
  visitLam  = visitLamEmits

type Dest = SAtom
type OptDest n = Maybe (Dest n)

lowerFor
  :: Emits o
  => SType o -> OptDest o -> ForAnn -> IxType SimpIR i -> LamExpr SimpIR i
  -> LowerM i o (SAtom o)
lowerFor ansTy maybeDest dir ixTy (UnaryLamExpr (ib:>ty) body) = do
  ixTy' <- substM ixTy
  ty' <- substM ty
  case isSingletonType ansTy of
    True -> do
      body' <- buildUnaryLamExpr noHint (PairTy ty' UnitTy) \b' -> do
        (i, _) <- fromPair $ toAtom b'
        extendSubst (ib @> SubstVal i) $ lowerExpr Nothing body $> UnitVal
      void $ emitSeq dir ixTy' UnitVal body'
      fromJust <$> singletonTypeVal ansTy
    False -> do
      initDest <- Con . ProdCon . (:[]) <$> case maybeDest of
        Just  d -> return d
        Nothing -> emit $ AllocDest ansTy
      let destTy = getType initDest
      body' <- buildUnaryLamExpr noHint (PairTy ty' destTy) \b' -> do
        (i, destProd) <- fromPair $ toAtom b'
        dest <- proj 0 destProd
        idest <- emit =<< mkIndexRef dest i
        extendSubst (ib @> SubstVal i) $ lowerExpr (Just idest) body $> UnitVal
      ans <- emitSeq dir ixTy' initDest body' >>= proj 0
      emit $ Freeze ans
lowerFor _ _ _ _ _ = error "expected a unary lambda expression"

lowerTabCon :: Emits o => OptDest o -> SType i -> [SAtom i] -> LowerM i o (SAtom o)
lowerTabCon maybeDest tabTy elems = do
  TyCon (TabPi tabTy') <- substM tabTy
  dest <- case maybeDest of
    Just  d -> return d
    Nothing -> emit $ AllocDest $ TyCon $ TabPi tabTy'
  Abs bord ufoBlock <- buildAbs noHint IdxRepTy \ord -> do
    buildBlock $ unsafeFromOrdinal (sink $ tabIxType tabTy') $ toAtom $ sink ord
  -- This is emitting a chain of RememberDest ops to force `dest` to be used
  -- linearly, and to force reads of the `Freeze dest'` result not to be
  -- reordered in front of the writes.
  -- TODO: We technically don't need to sequentialize the writes, since the
  -- slices are all independent of each other.  Do we need a new `JoinDests`
  -- operation to represent that pattern?
  let go incoming_dest [] = return incoming_dest
      go incoming_dest ((ord, e):rest) = do
        i <- dropSubst $ extendSubst (bord@>SubstVal (IdxRepVal (fromIntegral ord))) $
          lowerExpr Nothing ufoBlock
        carried_dest <- buildRememberDest "dest" incoming_dest \local_dest -> do
          idest <- indexRef (toAtom local_dest) (sink i)
          place idest =<< visitAtom e
          return UnitVal
        go carried_dest rest
  dest' <- go dest (enumerate elems)
  emit $ Freeze dest'

lowerCase :: Emits o
  => OptDest o -> SAtom i -> [Alt SimpIR i] -> SType i
  -> LowerM i o (SAtom o)
lowerCase maybeDest scrut alts resultTy = do
  resultTy' <- substM resultTy
  dest <- case maybeDest of
    Just  d -> return d
    Nothing -> emit $ AllocDest resultTy'
  scrut' <- visitAtom scrut
  dest' <- buildRememberDest "case_dest" dest \local_dest -> do
    alts' <- forM alts \(Abs (b:>ty) body) -> do
      ty' <- substM ty
      buildAbs (getNameHint b) ty' \b' ->
        extendSubst (b @> Rename (atomVarName b')) $
          buildBlock do
            lowerExpr (Just (toAtom $ sink $ local_dest)) body $> UnitVal
    void $ mkCase (sink scrut') UnitTy alts' >>= emit
    return UnitVal
  emit $ Freeze dest'

-- Destination-passing traversals
--
-- The idea here is to try to reuse the memory already allocated for outputs of surrounding
-- higher order functions for values produced inside nested blocks. As an example,
-- consider two perfectly nested for loops:
--
-- for i:(Fin 10).
--   v = for j:(Fin 20). ...
--   v
--
-- (binding of the for as v is necessary since Blocks end with Atoms).
--
-- The outermost loop will have to allocate the memory for a 2D array. But it would be
-- a waste if the inner loop kept allocating 20-element blocks only to copy them into
-- that full outermost buffer. Instead, we invoke lowerBlockWithDest on the body
-- of the i-loop, which will realize that v has a valid destination. Then,
-- lowerExprWithDest will be called at the for decl and will translate the j-loop
-- so that it never allocates scratch space for its result, but will put it directly in
-- the corresponding slice of the full 2D buffer.

type DestAssignment (i'::S) (o::S) = NameMap (AtomNameC SimpIR) (Dest o) i'

lookupDest :: DestAssignment i' o -> SAtomName i' -> OptDest o
lookupDest dests = fmap fromLiftE . flip lookupNameMapE dests

-- Matches up the free variables of the atom, with the given dest. For example, if the
-- atom is a pair of two variables, the dest might be split into per-component dests,
-- and a dest assignment mapping each side to its respective dest will be returned.
-- This function works on a best-effort basis. It's never an error to not split the dest
-- as much as possible, but it can lead to unnecessary copies being done at run-time.
--
-- XXX: When adding more cases, be careful about potentially repeated vars in the output!
decomposeDest :: Emits o => Dest o -> SExpr i' -> LowerM i o (Maybe (DestAssignment i' o))
decomposeDest dest = \case
  Atom (Stuck _ (Var v)) ->
    return $ Just $ singletonNameMapE (atomVarName v) $ LiftE dest
  _ -> return Nothing

traverseDeclNestWithDestS
  :: forall i i' l o. (Emits o, DistinctBetween l i')
  => DestAssignment i' o -> Subst AtomSubstVal l o -> Nest (Decl SimpIR) l i'
  -> LowerM i o (Subst AtomSubstVal i' o)
traverseDeclNestWithDestS destMap s = \case
  Empty -> return s
  Nest (Let b (DeclBinding _ expr)) rest -> do
    DistinctBetween <- return $ withExtEvidence rest $ shortenBetween @i' b
    let maybeDest = lookupDest destMap $ sinkBetween $ binderName b
    result <- withSubst s $ lowerExpr maybeDest expr
    traverseDeclNestWithDestS destMap (s <>> (b @> SubstVal result)) rest

traverseDeclNest :: Emits o => Nest SDecl i i' -> LowerM i'  o a -> LowerM i o a
traverseDeclNest decls cont = case decls of
  Empty -> cont
  Nest (Let b (DeclBinding _ expr)) rest -> do
    x <- lowerExpr Nothing expr
    extendSubst (b@>SubstVal x) $ traverseDeclNest rest cont

lowerExpr :: forall i o. Emits o => OptDest o -> SExpr i -> LowerM i o (SAtom o)
lowerExpr dest expr = case expr of
  Block _ (Abs decls result) -> case dest of
    Nothing -> traverseDeclNest decls $ lowerExpr Nothing result
    Just dest' -> do
      decomposeDest dest' result >>= \case
        Nothing -> do
          traverseDeclNest decls do
            lowerExpr (Just dest') result
        Just destMap -> do
          s <- getSubst
          case isDistinctNest decls of
            Nothing -> error "Non-distinct decls?"
            Just DistinctBetween -> do
              s' <- traverseDeclNestWithDestS destMap s decls
              -- But we have to emit explicit writes, for all the vars that are not defined in decls!
              forM_ (toListNameMapE $ hoistNameMap decls destMap) \(n, (LiftE d)) -> do
                x <- case s ! n of
                  Rename v -> toAtom <$> toAtomVar v
                  SubstVal a -> return a
                place d x
              withSubst s' (substM result) >>= emit
  TabCon Nothing ty els -> lowerTabCon dest ty els
  PrimOp (Hof (TypedHof (EffTy _ ansTy) (For dir ixDict body))) -> do
    ansTy' <- substM ansTy
    lowerFor ansTy' dest dir ixDict body
  PrimOp (Hof (TypedHof (EffTy _ ty) (RunWriter Nothing m body))) -> do
    PairTy _ ansTy <- visitType ty
    traverseRWS ansTy body \ref' body' -> do
      m' <- visitGeneric m
      emitHof $ RunWriter ref' m' body'
  PrimOp (Hof (TypedHof (EffTy _ ty) (RunState Nothing s body))) -> do
    PairTy _ ansTy <- visitType ty
    traverseRWS ansTy body \ref' body' -> do
      s' <- visitAtom s
      emitHof $ RunState ref' s' body'
  -- this case is important because this pass changes effects
  PrimOp (Hof (TypedHof _ hof)) -> do
    hof' <- emit =<< (visitGeneric hof >>= mkTypedHof)
    placeGeneric hof'
  Case e alts (EffTy _ ty) -> lowerCase dest e alts ty
  _ -> generic
  where
    generic :: LowerM i o (SAtom o)
    generic = visitGeneric expr >>= emit >>= placeGeneric

    placeGeneric :: SAtom o -> LowerM i o (SAtom o)
    placeGeneric e = do
      case dest of
        Nothing -> return e
        Just d  -> do
          place d e
          return e

    traverseRWS
      :: SType o -> LamExpr SimpIR i
      -> (OptDest o -> LamExpr SimpIR o -> LowerM i o (SAtom o))
      -> LowerM i o (SAtom o)
    traverseRWS referentTy (LamExpr (BinaryNest hb rb) body) cont = do
      unpackRWSDest dest >>= \case
        Nothing -> generic
        Just (bodyDest, refDest) -> do
          cont refDest =<<
            buildEffLam (getNameHint rb) referentTy \hb' rb' ->
              extendRenamer (hb@>atomVarName hb' <.> rb@>atomVarName rb') do
                lowerExpr (sink <$> bodyDest) body
    traverseRWS _ _ _ = error "Expected a binary lambda expression"

    unpackRWSDest = \case
      Nothing -> return Nothing
      Just d -> do
        bd <- getProjRef (ProjectProduct 0) d
        rd <- getProjRef (ProjectProduct 1) d
        return $ Just (Just bd, Just rd)

place :: Emits o => Dest o -> SAtom o -> LowerM i o ()
place d x = void $ emit $ Place d x

-- === Extensions to the name system ===

class DistinctBetween (n::S) (l::S)
instance DistinctBetween VoidS VoidS

data DistinctBetweenEvidence (n::S) (l::S) where
  DistinctBetween :: DistinctBetween n l => DistinctBetweenEvidence n l

fabricateDistinctBetweenEvidence :: forall n l. DistinctBetweenEvidence n l
fabricateDistinctBetweenEvidence = unsafeCoerce $ DistinctBetween @VoidS @VoidS

sinkBetween :: DistinctBetween n l => e n -> e l
sinkBetween = unsafeCoerceE

shortenBetween :: forall l n n' b. (ProvesExt b, Ext n' l, DistinctBetween n l) => b n n' -> DistinctBetweenEvidence n' l
shortenBetween _ = unsafeCoerce $ DistinctBetween @n @l

isDistinctNest :: Nest SDecl n l -> Maybe (DistinctBetweenEvidence n l)
isDistinctNest nest = case noShadows $ toScopeFrag nest of
  True  -> Just $ fabricateDistinctBetweenEvidence
  False -> Nothing