content: -- Copyright 2023 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 JAX.ToSimp (liftJaxSimpM, simplifyJaxpr, simplifyClosedJaxpr) where
  
  import Control.Category ((>>>))
  
  import Builder
  import Core
  import Err
  import IRVariants
  import Name
  import JAX.Concrete
  import Subst
  import QueryType
  import Types.Core
  import Types.Top
  import Types.Primitives qualified as P
  
  newtype JaxSimpM (i::S) (o::S) a = JaxSimpM
    { runJaxSimpM :: SubstReaderT AtomSubstVal (BuilderM SimpIR) i o a }
    deriving ( Functor, Applicative, Monad, MonadFail
             , ScopeReader, EnvReader, EnvExtender
             , SubstReader AtomSubstVal, Fallible
             , Builder SimpIR, ScopableBuilder SimpIR)
  
  liftJaxSimpM :: (EnvReader m) => JaxSimpM n n (e n) -> m n (e n)
  liftJaxSimpM act = liftBuilder $ runSubstReaderT idSubst $ runJaxSimpM act
  {-# INLINE liftJaxSimpM #-}
  
  simplifyClosedJaxpr :: ClosedJaxpr i -> JaxSimpM i o (TopLam SimpIR o)
  simplifyClosedJaxpr ClosedJaxpr{jaxpr, consts=[]} = asTopLam =<< simplifyJaxpr jaxpr
  simplifyClosedJaxpr _ = error "TODO Support consts"
  
  simplifyJaxpr :: Jaxpr i -> JaxSimpM i o (LamExpr SimpIR o)
  simplifyJaxpr (Jaxpr invars constvars eqns outvars) = do
    simplifyJBinders invars \invarsB -> do
      simplifyJBinders constvars \constvarsB -> do
        body <- buildBlock do
          simplifyEqns eqns do
            outs <- (map fst) <$> mapM simplifyAtom outvars
            return $ Con $ ProdCon $ outs
        return $ LamExpr (invarsB >>> constvarsB) body
  
  simplifyJBinders
    :: Nest JBinder i i'
    -> (forall o'. DExt o o' => Nest SBinder o o' -> JaxSimpM i' o' a)
    -> JaxSimpM i o a
  simplifyJBinders Empty cont = getDistinct >>= \Distinct -> cont Empty
  simplifyJBinders (Nest jb jbs) cont = case jb of
    JBindSource _ _ -> error "Unexpected source name after resolution"
    JBind (JSourceName {suffix=suffix}) jTy b -> do
      ty <- simplifyJTy jTy
      withFreshBinder (getNameHint suffix) ty \b' -> do
        extendSubst (b @> Rename (binderName b')) do
          simplifyJBinders jbs \bs' -> cont (Nest b' bs')
  
  simplifyJTy :: JVarType -> JaxSimpM i o (SType o)
  simplifyJTy JArrayName{shape, dtype} = go shape $ simplifyDType dtype where
    go [] ty = return ty
    go ((DimSize sz):rest) ty = do
      rest' <- go rest ty
      return $ litFinIxTy sz ==> rest'
  
  simplifyDType :: DType -> Type r n
  simplifyDType = \case
    F64 -> TyCon $ BaseType $ P.Scalar P.Float64Type
    F32 -> TyCon $ BaseType $ P.Scalar P.Float32Type
    I64 -> TyCon $ BaseType $ P.Scalar P.Int64Type
    I32 -> TyCon $ BaseType $ P.Scalar P.Int32Type
  
  simplifyEqns :: Emits o => Nest JEqn i i' -> JaxSimpM i' o a -> JaxSimpM i o a
  simplifyEqns eqn cont = do
    s  <- getSubst
    s' <- simplifyEqnsSubst s eqn
    withSubst s' cont
  {-# INLINE simplifyEqns #-}
  
  simplifyEqnsSubst :: Emits o
    => Subst AtomSubstVal l o -> Nest JEqn l i'
    -> JaxSimpM i o (Subst AtomSubstVal i' o)
  simplifyEqnsSubst !s Empty = return s
  simplifyEqnsSubst !s (Nest eqn rest) = do
    s' <- simplifyEqnSubst s eqn
    simplifyEqnsSubst s' rest
  
  simplifyEqnSubst :: Emits o
    => Subst AtomSubstVal l o -> JEqn l i' -> JaxSimpM i o (Subst AtomSubstVal i' o)
  simplifyEqnSubst !s JEqn{..} = do
    simpIn <- withSubst s $ mapM simplifyAtom invars
    simpOut <- simplifyPrim simpIn primitive
    return $ s <>> (outvars @@> (map SubstVal simpOut))
  
  simplifyAtom :: JAtom i -> JaxSimpM i o (SAtom o, JVarType)
  simplifyAtom = \case
    JVariable (JVar {..}) -> case sourceName of
      SourceName _ -> error "Unexpected source name during compilation"
      InternalName _ nm -> do
        env <- getSubst
        case env ! nm of
          -- TODO Assuming the subst is not type-changing
          SubstVal x -> return (x, ty)
          Rename nm' -> do
            nm'' <- toAtomVar nm'
            return (toAtom nm'', ty)
    -- TODO In Jax, literals can presumably include (large) arrays.  How should we
    -- represent them here?
    JLiteral (JLit {..}) -> return (Con (Lit (P.Float32Lit 0.0)), ty)
  
  simplifyPrim :: Emits o
    => [(SAtom o, JVarType)] -> Primitive -> JaxSimpM i o [SAtom o]
  simplifyPrim args prim = case (prim, args) of
    (Sin, [(arg, ty)]) -> do
      res <- unaryExpandRank P.Sin arg ty
      return [res]
    _ -> undefined
  
  unaryExpandRank :: forall i o. Emits o
    => P.UnOp -> SAtom o -> JVarType -> JaxSimpM i o (SAtom o)
  unaryExpandRank op arg JArrayName{shape} = go arg shape where
    go :: Emits l => SAtom l -> [DimSizeName] -> JaxSimpM i l (SAtom l)
    go arg' = \case
      [] -> emit $ PrimOp (UnOp op arg')
      (DimSize sz:rest) -> buildFor noHint P.Fwd (litFinIxTy sz) \i -> do
        ixed <- mkTabApp (sink arg') (toAtom i) >>= emit
        go ixed rest
  