    {-# LANGUAGE Safe #-}

    -- | A monad capable of keeping track of variable renames and of providing
    -- fresh names for variables.
    module Copilot.Theorem.TransSys.Renaming
      ( Renaming
      , addReservedName
      , rename
      , getFreshName
      , runRenaming
      , getRenamingF
      ) where

    import Copilot.Theorem.TransSys.Spec

    import Control.Monad.State.Lazy

    import Data.Maybe (fromMaybe)
    import Data.Map (Map)
    import Data.Set (Set, member)

    import qualified Data.Map  as Map
    import qualified Data.Set  as Set
    import qualified Data.List as List

    -- | A monad capable of keeping track of variable renames and of providing
    -- fresh names for variables.
    type Renaming = State RenamingST

    -- | State needed to keep track of variable renames and reserved names.
    data RenamingST = RenamingST
      { RenamingST -> Set Var
    _reservedNames :: Set Var
      , RenamingST -> Map ExtVar Var
    _renaming      :: Map ExtVar Var }

    -- | Register a name as reserved or used.
    addReservedName :: Var -> Renaming ()
    addReservedName :: Var -> Renaming ()
    addReservedName Var
    v = (RenamingST -> RenamingST) -> Renaming ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((RenamingST -> RenamingST) -> Renaming ())
    -> (RenamingST -> RenamingST) -> Renaming ()
    forall a b. (a -> b) -> a -> b
    $ \RenamingST
    st ->
        RenamingST
    st {_reservedNames = Set.insert v (_reservedNames st)}

    -- | Produce a fresh new name based on the variable names provided.
    --
    -- This function will try to pick a name from the given list and, if not, will
    -- use one of the names in the list as a basis for new names.
    --
    -- PRE: the given list cannot be empty.
    getFreshName :: [Var] -> Renaming Var
    getFreshName :: [Var] -> Renaming Var
    getFreshName [Var]
    vs = do
      usedNames <- RenamingST -> Set Var
    _reservedNames (RenamingST -> Set Var)
    -> StateT RenamingST Identity RenamingST
    -> StateT RenamingST Identity (Set Var)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StateT RenamingST Identity RenamingST
    forall s (m :: * -> *). MonadState s m => m s
    get
      let varAppend (Var String
    s) = String -> Var
    Var (String -> Var) -> String -> Var
    forall a b. (a -> b) -> a -> b
    $ String
    s String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    "_"
          applicants = [Var]
    vs [Var] -> [Var] -> [Var]
    forall a. [a] -> [a] -> [a]
    ++ (Var -> Var) -> Var -> [Var]
    forall a. (a -> a) -> a -> [a]
    List.iterate Var -> Var
    varAppend ([Var] -> Var
    forall a. HasCallStack => [a] -> a
    head [Var]
    vs)
          v = case (Var -> Bool) -> [Var] -> [Var]
    forall a. (a -> Bool) -> [a] -> [a]
    dropWhile (Var -> Set Var -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `member` Set Var
    usedNames) [Var]
    applicants of
                Var
    v:[Var]
    _ -> Var
    v
                [] -> String -> Var
    forall a. HasCallStack => String -> a
    error String
    "No more names available"
      addReservedName v
      return v

    -- | Map a name in the global namespace to a new variable name.
    rename :: NodeId  -- ^ A node Id
           -> Var     -- ^ A variable within that node
           -> Var     -- ^ A new name for the variable
           -> Renaming ()
    rename :: String -> Var -> Var -> Renaming ()
    rename String
    n Var
    v Var
    v' = (RenamingST -> RenamingST) -> Renaming ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((RenamingST -> RenamingST) -> Renaming ())
    -> (RenamingST -> RenamingST) -> Renaming ()
    forall a b. (a -> b) -> a -> b
    $ \RenamingST
    st ->
        RenamingST
    st {_renaming = Map.insert (ExtVar n v) v' (_renaming st)}

    -- | Return a function that maps variables in the global namespace to their new
    -- names if any renaming has been registered.
    getRenamingF :: Renaming (ExtVar -> Var)
    getRenamingF :: Renaming (ExtVar -> Var)
    getRenamingF = do
      mapping <- RenamingST -> Map ExtVar Var
    _renaming (RenamingST -> Map ExtVar Var)
    -> StateT RenamingST Identity RenamingST
    -> StateT RenamingST Identity (Map ExtVar Var)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StateT RenamingST Identity RenamingST
    forall s (m :: * -> *). MonadState s m => m s
    get
      return $ \ExtVar
    extv -> Var -> Maybe Var -> Var
    forall a. a -> Maybe a -> a
    fromMaybe (ExtVar -> Var
    extVarLocalPart ExtVar
    extv) (ExtVar -> Map ExtVar Var -> Maybe Var
    forall k a. Ord k => k -> Map k a -> Maybe a
    Map.lookup ExtVar
    extv Map ExtVar Var
    mapping)

    -- | Run a computation in the 'Renaming' monad, providing a result and the
    -- renaming function that maps variables in the global namespace to their new
    -- local names.
    runRenaming :: Renaming a -> (a, ExtVar -> Var)
    runRenaming :: forall a. Renaming a -> (a, ExtVar -> Var)
    runRenaming Renaming a
    m =
      State RenamingST (a, ExtVar -> Var)
    -> RenamingST -> (a, ExtVar -> Var)
    forall s a. State s a -> s -> a
    evalState State RenamingST (a, ExtVar -> Var)
    st' (Set Var -> Map ExtVar Var -> RenamingST
    RenamingST Set Var
    forall a. Set a
    Set.empty Map ExtVar Var
    forall k a. Map k a
    Map.empty)
      where
        st' :: State RenamingST (a, ExtVar -> Var)
    st' = do
          r <- Renaming a
    m
          f <- getRenamingF
          return (r, f)