{-# 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 = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \RenamingST
st ->
    RenamingST
st {_reservedNames :: Set Var
_reservedNames = forall a. Ord a => a -> Set a -> Set a
Set.insert Var
v (RenamingST -> Set Var
_reservedNames RenamingST
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
  Set Var
usedNames <- RenamingST -> Set Var
_reservedNames forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *). MonadState s m => m s
get
  let varAppend :: Var -> Var
varAppend (Var String
s) = String -> Var
Var forall a b. (a -> b) -> a -> b
$ String
s forall a. [a] -> [a] -> [a]
++ String
"_"
      applicants :: [Var]
applicants = [Var]
vs forall a. [a] -> [a] -> [a]
++ forall a. (a -> a) -> a -> [a]
List.iterate Var -> Var
varAppend (forall a. [a] -> a
head [Var]
vs)
      v :: Var
v = case forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Ord a => a -> Set a -> Bool
`member` Set Var
usedNames) [Var]
applicants of
            Var
v:[Var]
_ -> Var
v
            [] -> forall a. HasCallStack => String -> a
error String
"No more names available"
  Var -> Renaming ()
addReservedName Var
v
  forall (m :: * -> *) a. Monad m => a -> m a
return Var
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' = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \RenamingST
st ->
    RenamingST
st {_renaming :: Map ExtVar Var
_renaming = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (String -> Var -> ExtVar
ExtVar String
n Var
v) Var
v' (RenamingST -> Map ExtVar Var
_renaming RenamingST
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
  Map ExtVar Var
mapping <- RenamingST -> Map ExtVar Var
_renaming forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *). MonadState s m => m s
get
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ \ExtVar
extv -> forall a. a -> Maybe a -> a
fromMaybe (ExtVar -> Var
extVarLocalPart ExtVar
extv) (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 =
  forall s a. State s a -> s -> a
evalState StateT RenamingST Identity (a, ExtVar -> Var)
st' (Set Var -> Map ExtVar Var -> RenamingST
RenamingST forall a. Set a
Set.empty forall k a. Map k a
Map.empty)
  where
    st' :: StateT RenamingST Identity (a, ExtVar -> Var)
st' = do
      a
r <- Renaming a
m
      ExtVar -> Var
f <- Renaming (ExtVar -> Var)
getRenamingF
      forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, ExtVar -> Var
f)