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.Rename (liftRenameM, renameClosedJaxpr, renameJaxpr) where
  
  import Control.Monad.Reader
  import Data.Map qualified as M
  
  import Core
  import IRVariants
  import JAX.Concrete
  import MTL1
  import Name
  
  newtype RenamerM (n::S) (a:: *) =
    RenamerM { runRenamerM :: ReaderT1 SourceMap (ScopeReaderM) n a }
    deriving ( Functor, Applicative, Monad
             , ScopeReader, ScopeExtender)
  
  newtype SourceMap (n::S) = SourceMap
    (M.Map JSourceName (Name (AtomNameC SimpIR) n))
    deriving (Semigroup, Monoid)
  
  instance SinkableE SourceMap where
    sinkingProofE = undefined
  
  askSourceMap :: RenamerM n (SourceMap n)
  askSourceMap = RenamerM ask
  
  extendSourceMap :: JSourceName -> (Name (AtomNameC SimpIR)) n
    -> RenamerM n a -> RenamerM n a
  extendSourceMap sname name (RenamerM cont) = RenamerM do
    let ext = SourceMap $ M.singleton sname name
    local (<> ext) cont
  
  liftRenameM :: EnvReader m => RenamerM n (e n) -> m n (e n)
  liftRenameM act = liftScopeReaderM $ runReaderT1 mempty $ runRenamerM act
  
  renameClosedJaxpr :: Distinct o => ClosedJaxpr i -> RenamerM o (ClosedJaxpr o)
  renameClosedJaxpr ClosedJaxpr{jaxpr, consts} = do
    jaxpr' <- renameJaxpr jaxpr
    return ClosedJaxpr{jaxpr=jaxpr', consts}
  
  renameJaxpr :: Distinct o => Jaxpr i -> RenamerM o (Jaxpr o)
  renameJaxpr (Jaxpr invars constvars eqns outvars) =
    renameJBinders invars \invars' ->
      renameJBinders constvars \constvars' ->
        renameJEqns eqns \eqns' -> do
          outvars' <- mapM renameJAtom outvars
          return $ Jaxpr invars' constvars' eqns' outvars'
  
  renameJBinder :: Distinct o
    => JBinder i i'
    -> (forall o'. DExt o o' => JBinder o o' -> RenamerM o' a)
    -> RenamerM o a
  renameJBinder binder cont = case binder of
    JBindSource sname ty -> do
      withFreshM (getNameHint sname) \freshName -> do
        Distinct <- getDistinct
        extendSourceMap sname (binderName freshName) $
          cont $ JBind sname ty freshName
    JBind _ _ _ -> error "Shouldn't be source-renaming internal names"
  
  renameJBinders :: Distinct o
    => Nest JBinder i i'
    -> (forall o'. DExt o o' => Nest JBinder o o' -> RenamerM o' a)
    -> RenamerM o a
  renameJBinders Empty cont = cont Empty
  renameJBinders (Nest b bs) cont =
    renameJBinder b \b' ->
      renameJBinders bs \bs' ->
        cont $ Nest b' bs'
  
  renameJAtom :: JAtom i -> RenamerM o (JAtom o)
  renameJAtom = \case
    JVariable jvar -> JVariable <$> renameJVar jvar
    JLiteral jlit -> return $ JLiteral jlit
  
  renameJVar :: JVar i -> RenamerM o (JVar o)
  renameJVar JVar{sourceName, ty} = do
    sourceName' <- renameJSourceNameOr sourceName
    return $ JVar sourceName' ty
  
  renameJSourceNameOr :: JSourceNameOr (Name (AtomNameC SimpIR)) i
    -> RenamerM o (JSourceNameOr (Name (AtomNameC SimpIR)) o)
  renameJSourceNameOr = \case
    SourceName sname -> do
      SourceMap sm <- askSourceMap
      case M.lookup sname sm of
        (Just name) -> return $ InternalName sname name
        Nothing -> error $ "Unbound variable " ++ show sname
    InternalName _ _ -> error "Shouldn't be source-renaming internal names"
  
  renameJEqn :: Distinct o
    => JEqn i i'
    -> (forall o'. DExt o o' => JEqn o o' -> RenamerM o' a)
    -> RenamerM o a
  renameJEqn JEqn{outvars, primitive, invars} cont = do
    invars' <- mapM renameJAtom invars
    renameJBinders outvars \outvars' -> cont $ JEqn outvars' primitive invars'
  
  renameJEqns :: Distinct o
    => Nest JEqn i i'
    -> (forall o'. DExt o o' => Nest JEqn o o' -> RenamerM o' a)
    -> RenamerM o a
  renameJEqns Empty cont = cont Empty
  renameJEqns (Nest b bs) cont =
    renameJEqn b \b' ->
      renameJEqns bs \bs' ->
        cont $ Nest b' bs'
  
  