{-# 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
  Set Var
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 -> Var
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]
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 :: Var
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"
  Var -> Renaming ()
addReservedName Var
v
  Var -> Renaming Var
forall a. a -> StateT RenamingST Identity a
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' = (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
  Map ExtVar Var
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
  (ExtVar -> Var) -> Renaming (ExtVar -> Var)
forall a. a -> StateT RenamingST Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ((ExtVar -> Var) -> Renaming (ExtVar -> Var))
-> (ExtVar -> Var) -> Renaming (ExtVar -> Var)
forall a b. (a -> b) -> a -> b
$ \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
      a
r <- Renaming a
m
      ExtVar -> Var
f <- Renaming (ExtVar -> Var)
getRenamingF
      (a, ExtVar -> Var) -> State RenamingST (a, ExtVar -> Var)
forall a. a -> StateT RenamingST Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, ExtVar -> Var
f)