--------------------------------------------------------------------------------

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Safe #-}

module Copilot.Theorem.TransSys.Transform
  ( mergeNodes
  , inline
  , removeCycles
  , complete
  ) where

import Copilot.Theorem.TransSys.Spec
import Copilot.Theorem.TransSys.Renaming

import Copilot.Theorem.Misc.Utils

import Control.Monad (foldM, forM_, forM, guard)

import Data.List (sort, (\\), intercalate, partition)

import Control.Exception.Base (assert)

import Data.Map (Map, (!))
import Data.Set (member)
import Data.Bimap (Bimap)

import qualified Data.Map   as Map
import qualified Data.Set   as Set
import qualified Data.Graph as Graph
import qualified Data.Bimap as Bimap

--------------------------------------------------------------------------------

prefix :: String -> Var -> Var
prefix :: String -> Var -> Var
prefix String
s1 (Var String
s2) = String -> Var
Var (String -> Var) -> String -> Var
forall a b. (a -> b) -> a -> b
$ String
s1 String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s2

ncNodeIdSep :: String
ncNodeIdSep = String
"-"

--------------------------------------------------------------------------------

mergeNodes :: [NodeId] -> TransSys -> TransSys
mergeNodes :: [String] -> TransSys -> TransSys
mergeNodes [String]
toMergeIds TransSys
spec =
  TransSys
spec
    { specNodes :: [Node]
specNodes = Node
newNode Node -> [Node] -> [Node]
forall a. a -> [a] -> [a]
:
        (Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (String -> [String] -> (ExtVar -> ExtVar) -> Node -> Node
updateOtherNode String
newNodeId [String]
toMergeIds ExtVar -> ExtVar
renamingExtF) [Node]
otherNodes
    , specProps :: Map String ExtVar
specProps = (ExtVar -> ExtVar) -> Map String ExtVar -> Map String ExtVar
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map ExtVar -> ExtVar
renamingExtF (TransSys -> Map String ExtVar
specProps TransSys
spec) }

  where
    nodes :: [Node]
nodes = TransSys -> [Node]
specNodes TransSys
spec
    ([Node]
toMerge, [Node]
otherNodes) = (Node -> Bool) -> [Node] -> ([Node], [Node])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
toMergeIds) (String -> Bool) -> (Node -> String) -> Node -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Node -> String
nodeId) [Node]
nodes

    -- Choosing the new node ID. If the top node is merged,
    -- its name is kept
    newNodeId :: String
newNodeId
      | TransSys -> String
specTopNodeId TransSys
spec String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
toMergeIds = TransSys -> String
specTopNodeId TransSys
spec
      | Bool
otherwise = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
ncNodeIdSep ([String] -> [String]
forall a. Ord a => [a] -> [a]
sort [String]
toMergeIds)

    newNode :: Node
newNode = Node :: String
-> [String]
-> Map Var VarDescr
-> Bimap Var ExtVar
-> [Expr Bool]
-> Node
Node
      { nodeId :: String
nodeId = String
newNodeId
      , nodeDependencies :: [String]
nodeDependencies = [String]
dependencies
      , nodeImportedVars :: Bimap Var ExtVar
nodeImportedVars = Bimap Var ExtVar
importedVars
      , nodeLocalVars :: Map Var VarDescr
nodeLocalVars = Map Var VarDescr
localVars
      , nodeConstrs :: [Expr Bool]
nodeConstrs = [Expr Bool]
constrs }

    -- Computing the dependencies of the new node
    dependencies :: [String]
dependencies = [String] -> [String]
forall a. Ord a => [a] -> [a]
nub'
      [ String
id |
        Node
n <- [Node]
toMerge
      , String
id <- Node -> [String]
nodeDependencies Node
n
      , String
id String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [String]
toMergeIds ]

    -- All the work of renaming is done in the monad with the same name
    (Bimap Var ExtVar
importedVars, ExtVar -> Var
renamingF) = Renaming (Bimap Var ExtVar) -> (Bimap Var ExtVar, ExtVar -> Var)
forall a. Renaming a -> (a, ExtVar -> Var)
runRenaming (Renaming (Bimap Var ExtVar) -> (Bimap Var ExtVar, ExtVar -> Var))
-> Renaming (Bimap Var ExtVar) -> (Bimap Var ExtVar, ExtVar -> Var)
forall a b. (a -> b) -> a -> b
$ do
      [Node] -> Renaming ()
renameLocalVars [Node]
toMerge
      [Node] -> Renaming ()
redirectLocalImports [Node]
toMerge
      [Node] -> [Node] -> [String] -> Renaming (Bimap Var ExtVar)
selectImportedVars [Node]
toMerge [Node]
otherNodes [String]
dependencies

    -- Converting the variables descriptors
    localVars :: Map Var VarDescr
localVars = [Node] -> (ExtVar -> Var) -> Map Var VarDescr
mergeVarsDescrs [Node]
toMerge ExtVar -> Var
renamingF

    -- Computing the global renaming function
    renamingExtF :: ExtVar -> ExtVar
renamingExtF (gv :: ExtVar
gv@(ExtVar String
nId Var
_))
     | String
nId String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
toMergeIds = String -> Var -> ExtVar
ExtVar String
newNodeId (ExtVar -> Var
renamingF ExtVar
gv)
     | Bool
otherwise = ExtVar
gv

    constrs :: [Expr Bool]
constrs = [Node] -> (ExtVar -> Var) -> [Expr Bool]
mergeConstrs [Node]
toMerge ExtVar -> Var
renamingF


updateOtherNode :: NodeId -> [NodeId] -> (ExtVar -> ExtVar) -> Node -> Node
updateOtherNode :: String -> [String] -> (ExtVar -> ExtVar) -> Node -> Node
updateOtherNode String
newNodeId [String]
mergedNodesIds ExtVar -> ExtVar
renamingF Node
n = Node
n
  { nodeDependencies :: [String]
nodeDependencies =
      let ds :: [String]
ds  = Node -> [String]
nodeDependencies Node
n
          ds' :: [String]
ds' = [String]
ds [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ [String]
mergedNodesIds
      in if [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
ds' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
ds then String
newNodeId String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ds' else [String]
ds

  , nodeImportedVars :: Bimap Var ExtVar
nodeImportedVars =
      [(Var, ExtVar)] -> Bimap Var ExtVar
forall a b. (Ord a, Ord b) => [(a, b)] -> Bimap a b
Bimap.fromList [ (Var
lv, ExtVar -> ExtVar
renamingF ExtVar
gv)
                     | (Var
lv, ExtVar
gv) <- Bimap Var ExtVar -> [(Var, ExtVar)]
forall a b. Bimap a b -> [(a, b)]
Bimap.toList (Bimap Var ExtVar -> [(Var, ExtVar)])
-> Bimap Var ExtVar -> [(Var, ExtVar)]
forall a b. (a -> b) -> a -> b
$ Node -> Bimap Var ExtVar
nodeImportedVars Node
n ]
  }



updateExpr :: NodeId -> (ExtVar -> Var) -> Expr t -> Expr t
updateExpr :: String -> (ExtVar -> Var) -> Expr t -> Expr t
updateExpr String
nId ExtVar -> Var
renamingF = (forall a. Expr a -> Expr a) -> Expr t -> Expr t
forall t. (forall a. Expr a -> Expr a) -> Expr t -> Expr t
transformExpr forall a. Expr a -> Expr a
aux
  where
    aux :: forall t. Expr t -> Expr t
    aux :: Expr t -> Expr t
aux (VarE Type t
t Var
v) = Type t -> Var -> Expr t
forall t. Type t -> Var -> Expr t
VarE Type t
t (ExtVar -> Var
renamingF (String -> Var -> ExtVar
ExtVar String
nId Var
v))
    aux Expr t
e = Expr t
e


mergeVarsDescrs :: [Node] -> (ExtVar -> Var) -> Map Var VarDescr
mergeVarsDescrs :: [Node] -> (ExtVar -> Var) -> Map Var VarDescr
mergeVarsDescrs [Node]
toMerge ExtVar -> Var
renamingF = [(Var, VarDescr)] -> Map Var VarDescr
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Var, VarDescr)] -> Map Var VarDescr)
-> [(Var, VarDescr)] -> Map Var VarDescr
forall a b. (a -> b) -> a -> b
$ do
  Node
n <- [Node]
toMerge
  let nId :: String
nId = Node -> String
nodeId Node
n
  (Var
v, VarDescr Type t
t VarDef t
def) <- Map Var VarDescr -> [(Var, VarDescr)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map Var VarDescr -> [(Var, VarDescr)])
-> Map Var VarDescr -> [(Var, VarDescr)]
forall a b. (a -> b) -> a -> b
$ Node -> Map Var VarDescr
nodeLocalVars Node
n
  let d' :: VarDescr
d' = case VarDef t
def of
       Pre t
val Var
v' -> Type t -> VarDef t -> VarDescr
forall t. Type t -> VarDef t -> VarDescr
VarDescr Type t
t (VarDef t -> VarDescr) -> VarDef t -> VarDescr
forall a b. (a -> b) -> a -> b
$
         t -> Var -> VarDef t
forall t. t -> Var -> VarDef t
Pre t
val (Var -> VarDef t) -> Var -> VarDef t
forall a b. (a -> b) -> a -> b
$ ExtVar -> Var
renamingF (String -> Var -> ExtVar
ExtVar String
nId Var
v')
       Expr Expr t
e -> Type t -> VarDef t -> VarDescr
forall t. Type t -> VarDef t -> VarDescr
VarDescr Type t
t (VarDef t -> VarDescr) -> VarDef t -> VarDescr
forall a b. (a -> b) -> a -> b
$
         Expr t -> VarDef t
forall t. Expr t -> VarDef t
Expr (Expr t -> VarDef t) -> Expr t -> VarDef t
forall a b. (a -> b) -> a -> b
$ String -> (ExtVar -> Var) -> Expr t -> Expr t
forall t. String -> (ExtVar -> Var) -> Expr t -> Expr t
updateExpr String
nId ExtVar -> Var
renamingF Expr t
e
       Constrs [Expr Bool]
cs -> Type t -> VarDef t -> VarDescr
forall t. Type t -> VarDef t -> VarDescr
VarDescr Type t
t (VarDef t -> VarDescr) -> VarDef t -> VarDescr
forall a b. (a -> b) -> a -> b
$
         [Expr Bool] -> VarDef t
forall t. [Expr Bool] -> VarDef t
Constrs ([Expr Bool] -> VarDef t) -> [Expr Bool] -> VarDef t
forall a b. (a -> b) -> a -> b
$ (Expr Bool -> Expr Bool) -> [Expr Bool] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map (String -> (ExtVar -> Var) -> Expr Bool -> Expr Bool
forall t. String -> (ExtVar -> Var) -> Expr t -> Expr t
updateExpr String
nId ExtVar -> Var
renamingF) [Expr Bool]
cs

  (Var, VarDescr) -> [(Var, VarDescr)]
forall (m :: * -> *) a. Monad m => a -> m a
return (ExtVar -> Var
renamingF (ExtVar -> Var) -> ExtVar -> Var
forall a b. (a -> b) -> a -> b
$ String -> Var -> ExtVar
ExtVar String
nId Var
v, VarDescr
d')

mergeConstrs :: [Node] -> (ExtVar -> Var) -> [Expr Bool]
mergeConstrs :: [Node] -> (ExtVar -> Var) -> [Expr Bool]
mergeConstrs [Node]
toMerge ExtVar -> Var
renamingF =
  [ String -> (ExtVar -> Var) -> Expr Bool -> Expr Bool
forall t. String -> (ExtVar -> Var) -> Expr t -> Expr t
updateExpr (Node -> String
nodeId Node
n) ExtVar -> Var
renamingF Expr Bool
c | Node
n <- [Node]
toMerge, Expr Bool
c <- Node -> [Expr Bool]
nodeConstrs Node
n ]

renameLocalVars :: [Node] -> Renaming ()
renameLocalVars :: [Node] -> Renaming ()
renameLocalVars [Node]
toMerge =
  [(String, Var)] -> ((String, Var) -> Renaming ()) -> Renaming ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(String, Var)]
niVars (((String, Var) -> Renaming ()) -> Renaming ())
-> ((String, Var) -> Renaming ()) -> Renaming ()
forall a b. (a -> b) -> a -> b
$ \(String
n, Var
v) -> do
    Var
v' <- [Var] -> Renaming Var
getFreshName [String
n String -> Var -> Var
`prefix` Var
v]
    String -> Var -> Var -> Renaming ()
rename String
n Var
v Var
v'
  where
  niVars :: [(String, Var)]
niVars = [ (Node -> String
nodeId Node
n, Var
v)
           | Node
n <- [Node]
toMerge, (Var
v, VarDescr
_) <- Map Var VarDescr -> [(Var, VarDescr)]
forall k a. Map k a -> [(k, a)]
Map.toList (Node -> Map Var VarDescr
nodeLocalVars Node
n) ]

selectImportedVars :: [Node] -> [Node] -> [NodeId]
                      -> Renaming (Bimap Var ExtVar)
selectImportedVars :: [Node] -> [Node] -> [String] -> Renaming (Bimap Var ExtVar)
selectImportedVars [Node]
toMerge [Node]
otherNodes [String]
dependencies =
  (Bimap Var ExtVar -> (String, Var) -> Renaming (Bimap Var ExtVar))
-> Bimap Var ExtVar
-> [(String, Var)]
-> Renaming (Bimap Var ExtVar)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Bimap Var ExtVar -> (String, Var) -> Renaming (Bimap Var ExtVar)
checkImport Bimap Var ExtVar
forall a b. Bimap a b
Bimap.empty [(String, Var)]
depsVars

  where
    otherNodesMap :: Map String Node
otherNodesMap = [(String, Node)] -> Map String Node
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Node -> String
nodeId Node
n, Node
n) | Node
n <- [Node]
otherNodes]

    depsVars :: [(String, Var)]
depsVars = [ (String
nId, Var
v)
                 | String
nId <- [String]
dependencies, let n :: Node
n = Map String Node
otherNodesMap Map String Node -> String -> Node
forall k a. Ord k => Map k a -> k -> a
! String
nId
                 , Var
v <- Map Var VarDescr -> [Var]
forall k a. Map k a -> [k]
Map.keys (Node -> Map Var VarDescr
nodeLocalVars Node
n)]

    checkImport :: Bimap Var ExtVar -> (String, Var) -> Renaming (Bimap Var ExtVar)
checkImport Bimap Var ExtVar
acc (String
nId, Var
v) = do
        Var
v' <- [Var] -> Renaming Var
getFreshName [String
nId String -> Var -> Var
`prefix` Var
v]
        [Bool]
bmap <- [Node]
-> (Node -> StateT RenamingST Identity Bool)
-> StateT RenamingST Identity [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Node]
toMerge ((Node -> StateT RenamingST Identity Bool)
 -> StateT RenamingST Identity [Bool])
-> (Node -> StateT RenamingST Identity Bool)
-> StateT RenamingST Identity [Bool]
forall a b. (a -> b) -> a -> b
$ \Node
n' ->
                  case ExtVar -> Bimap Var ExtVar -> Maybe Var
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR (String -> Var -> ExtVar
ExtVar String
nId Var
v)
                       (Node -> Bimap Var ExtVar
nodeImportedVars Node
n') of

                     Just Var
lv -> String -> Var -> Var -> Renaming ()
rename (Node -> String
nodeId Node
n') Var
lv Var
v' Renaming ()
-> StateT RenamingST Identity Bool
-> StateT RenamingST Identity Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> StateT RenamingST Identity Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                     Maybe Var
Nothing -> Bool -> StateT RenamingST Identity Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

        Bimap Var ExtVar -> Renaming (Bimap Var ExtVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Bimap Var ExtVar -> Renaming (Bimap Var ExtVar))
-> Bimap Var ExtVar -> Renaming (Bimap Var ExtVar)
forall a b. (a -> b) -> a -> b
$
          if Bool
True Bool -> [Bool] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Bool]
bmap
            then Var -> ExtVar -> Bimap Var ExtVar -> Bimap Var ExtVar
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.insert Var
v' (String -> Var -> ExtVar
ExtVar String
nId Var
v) Bimap Var ExtVar
acc
            else Bimap Var ExtVar
acc



redirectLocalImports :: [Node] -> Renaming ()
redirectLocalImports :: [Node] -> Renaming ()
redirectLocalImports [Node]
toMerge = do
  ExtVar -> Var
renamingF <- Renaming (ExtVar -> Var)
getRenamingF
  [(String, Var, String, Var)]
-> ((String, Var, String, Var) -> Renaming ()) -> Renaming ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(String, Var, String, Var)]
x (((String, Var, String, Var) -> Renaming ()) -> Renaming ())
-> ((String, Var, String, Var) -> Renaming ()) -> Renaming ()
forall a b. (a -> b) -> a -> b
$ \(String
n, Var
alias, String
n', Var
v) ->
    String -> Var -> Var -> Renaming ()
rename String
n Var
alias (ExtVar -> Var
renamingF (String -> Var -> ExtVar
ExtVar String
n' Var
v))

  where
    mergedNodesSet :: Set String
mergedNodesSet = [String] -> Set String
forall a. Ord a => [a] -> Set a
Set.fromList [Node -> String
nodeId Node
n | Node
n <- [Node]
toMerge]
    x :: [(String, Var, String, Var)]
x = do
      Node
n <- [Node]
toMerge
      let nId :: String
nId = Node -> String
nodeId Node
n
      (Var
alias, ExtVar String
n' Var
v) <- Bimap Var ExtVar -> [(Var, ExtVar)]
forall a b. Bimap a b -> [(a, b)]
Bimap.toList (Node -> Bimap Var ExtVar
nodeImportedVars Node
n)
      Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ String
n' String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
`member` Set String
mergedNodesSet
      (String, Var, String, Var) -> [(String, Var, String, Var)]
forall (m :: * -> *) a. Monad m => a -> m a
return (String
nId, Var
alias, String
n', Var
v)

--------------------------------------------------------------------------------

inline :: TransSys -> TransSys
inline :: TransSys -> TransSys
inline TransSys
spec = [String] -> TransSys -> TransSys
mergeNodes [Node -> String
nodeId Node
n | Node
n <- TransSys -> [Node]
specNodes TransSys
spec] TransSys
spec

removeCycles :: TransSys -> TransSys
removeCycles :: TransSys -> TransSys
removeCycles TransSys
spec =
  TransSys -> TransSys
topoSort (TransSys -> TransSys) -> TransSys -> TransSys
forall a b. (a -> b) -> a -> b
$ (SCC String -> TransSys -> TransSys)
-> TransSys -> [SCC String] -> TransSys
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SCC String -> TransSys -> TransSys
mergeComp TransSys
spec ((Node -> String) -> [Node] -> [SCC String]
forall node. (Node -> node) -> [Node] -> [SCC node]
buildScc Node -> String
nodeId ([Node] -> [SCC String]) -> [Node] -> [SCC String]
forall a b. (a -> b) -> a -> b
$ TransSys -> [Node]
specNodes TransSys
spec)
  where

    mergeComp :: SCC String -> TransSys -> TransSys
mergeComp (Graph.AcyclicSCC String
_)  TransSys
s = TransSys
s
    mergeComp (Graph.CyclicSCC [String]
ids) TransSys
s = [String] -> TransSys -> TransSys
mergeNodes [String]
ids TransSys
s

    buildScc :: (Node -> node) -> [Node] -> [SCC node]
buildScc Node -> node
nrep [Node]
ns =
     let depGraph :: [(node, String, [String])]
depGraph = (Node -> (node, String, [String]))
-> [Node] -> [(node, String, [String])]
forall a b. (a -> b) -> [a] -> [b]
map (\Node
n -> (Node -> node
nrep Node
n, Node -> String
nodeId Node
n, Node -> [String]
nodeDependencies Node
n)) [Node]
ns
     in [(node, String, [String])] -> [SCC node]
forall key node. Ord key => [(node, key, [key])] -> [SCC node]
Graph.stronglyConnComp [(node, String, [String])]
depGraph

    topoSort :: TransSys -> TransSys
topoSort TransSys
s = TransSys
s { specNodes :: [Node]
specNodes =
      (SCC Node -> Node) -> [SCC Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (\(Graph.AcyclicSCC Node
n) -> Node
n) ([SCC Node] -> [Node]) -> [SCC Node] -> [Node]
forall a b. (a -> b) -> a -> b
$ (Node -> Node) -> [Node] -> [SCC Node]
forall node. (Node -> node) -> [Node] -> [SCC node]
buildScc Node -> Node
forall a. a -> a
id (TransSys -> [Node]
specNodes TransSys
s) }

--------------------------------------------------------------------------------

-- | Completes each node of a specification with imported variables such
-- | that each node contains a copy of all its dependencies
-- | The given specification should have its node sorted by topological
-- | order.
-- | The top nodes should have all the other nodes as its dependencies

complete :: TransSys -> TransSys
complete :: TransSys -> TransSys
complete TransSys
spec =
  Bool -> TransSys -> TransSys
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (TransSys -> Bool
isTopologicallySorted TransSys
spec) (TransSys -> TransSys) -> TransSys -> TransSys
forall a b. (a -> b) -> a -> b
$ TransSys
spec { specNodes :: [Node]
specNodes = [Node]
specNodes' }

  where

    specNodes' :: [Node]
specNodes' =
      [Node] -> [Node]
forall a. [a] -> [a]
reverse
      ([Node] -> [Node]) -> (TransSys -> [Node]) -> TransSys -> [Node]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Node] -> Node -> [Node]) -> [Node] -> [Node] -> [Node]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl [Node] -> Node -> [Node]
completeNode []
      ([Node] -> [Node]) -> (TransSys -> [Node]) -> TransSys -> [Node]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TransSys -> [Node]
specNodes
      (TransSys -> [Node])
-> (TransSys -> TransSys) -> TransSys -> [Node]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TransSys -> TransSys
completeTopNodeDeps
      (TransSys -> [Node]) -> TransSys -> [Node]
forall a b. (a -> b) -> a -> b
$ TransSys
spec

    completeTopNodeDeps :: TransSys -> TransSys
completeTopNodeDeps TransSys
spec = TransSys
spec { specNodes :: [Node]
specNodes = (Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map Node -> Node
aux [Node]
nodes }
      where
        nodes :: [Node]
nodes = TransSys -> [Node]
specNodes TransSys
spec
        aux :: Node -> Node
aux Node
n
          | Node -> String
nodeId Node
n String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== TransSys -> String
specTopNodeId TransSys
spec =
              Node
n { nodeDependencies :: [String]
nodeDependencies = (Node -> String) -> [Node] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Node -> String
nodeId [Node]
nodes [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Node -> String
nodeId Node
n] }
          | Bool
otherwise = Node
n

    -- Takes a list of nodes 'ns', 'n' whose dependencies are in 'ns', and
    -- returns 'n2:ns' where 'n2' is 'n' completed
    completeNode :: [Node] -> Node -> [Node]
    completeNode :: [Node] -> Node -> [Node]
completeNode [Node]
ns Node
n = (Node
n { nodeDependencies :: [String]
nodeDependencies = [String]
dependencies'
                           , nodeImportedVars :: Bimap Var ExtVar
nodeImportedVars = Bimap Var ExtVar
importedVars' }) Node -> [Node] -> [Node]
forall a. a -> [a] -> [a]
: [Node]
ns

      where
        nsMap :: Map String Node
nsMap = [(String, Node)] -> Map String Node
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Node -> String
nodeId Node
n, Node
n) | Node
n <- [Node]
ns]
        dependencies' :: [String]
dependencies' =
          let newDeps :: [String]
newDeps = do
                String
dId <- Node -> [String]
nodeDependencies Node
n
                let d :: Node
d = Map String Node
nsMap Map String Node -> String -> Node
forall k a. Ord k => Map k a -> k -> a
! String
dId
                Node -> [String]
nodeDependencies Node
d

          in [String] -> [String]
forall a. Ord a => [a] -> [a]
nub' ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ Node -> [String]
nodeDependencies Node
n [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String]
newDeps

        importedVars' :: Bimap Var ExtVar
importedVars' = (Bimap Var ExtVar, ExtVar -> Var) -> Bimap Var ExtVar
forall a b. (a, b) -> a
fst ((Bimap Var ExtVar, ExtVar -> Var) -> Bimap Var ExtVar)
-> (Renaming (Bimap Var ExtVar)
    -> (Bimap Var ExtVar, ExtVar -> Var))
-> Renaming (Bimap Var ExtVar)
-> Bimap Var ExtVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Renaming (Bimap Var ExtVar) -> (Bimap Var ExtVar, ExtVar -> Var)
forall a. Renaming a -> (a, ExtVar -> Var)
runRenaming (Renaming (Bimap Var ExtVar) -> Bimap Var ExtVar)
-> Renaming (Bimap Var ExtVar) -> Bimap Var ExtVar
forall a b. (a -> b) -> a -> b
$ do
          [Var] -> (Var -> Renaming ()) -> Renaming ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Set Var -> [Var]
forall a. Set a -> [a]
Set.toList (Set Var -> [Var]) -> Set Var -> [Var]
forall a b. (a -> b) -> a -> b
$ Node -> Set Var
nodeVarsSet Node
n) Var -> Renaming ()
addReservedName
          let toImportVars :: [ExtVar]
toImportVars = [ExtVar] -> [ExtVar]
forall a. Ord a => [a] -> [a]
nub' [ String -> Var -> ExtVar
ExtVar String
nId Var
v
                                  | String
nId <- [String]
dependencies'
                                  , let n' :: Node
n' = Map String Node
nsMap Map String Node -> String -> Node
forall k a. Ord k => Map k a -> k -> a
! String
nId
                                  , Var
v <- Map Var VarDescr -> [Var]
forall k a. Map k a -> [k]
Map.keys (Node -> Map Var VarDescr
nodeLocalVars Node
n') ]

              tryImport :: Bimap Var ExtVar -> ExtVar -> Renaming (Bimap Var ExtVar)
tryImport Bimap Var ExtVar
acc ev :: ExtVar
ev@(ExtVar String
n' Var
v) = do

                -- To get readable names, we don't prefix variables
                -- which come from merged nodes as they are already
                -- decorated
                let preferedName :: Var
preferedName
                     | String -> Char
forall a. [a] -> a
head String
ncNodeIdSep Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
n' = Var
v
                     | Bool
otherwise = String
n' String -> Var -> Var
`prefix` Var
v
                Var
alias <- [Var] -> Renaming Var
getFreshName [Var
preferedName, String
n' String -> Var -> Var
`prefix` Var
v]
                Bimap Var ExtVar -> Renaming (Bimap Var ExtVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Bimap Var ExtVar -> Renaming (Bimap Var ExtVar))
-> Bimap Var ExtVar -> Renaming (Bimap Var ExtVar)
forall a b. (a -> b) -> a -> b
$ Var -> ExtVar -> Bimap Var ExtVar -> Bimap Var ExtVar
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.tryInsert Var
alias ExtVar
ev Bimap Var ExtVar
acc

          (Bimap Var ExtVar -> ExtVar -> Renaming (Bimap Var ExtVar))
-> Bimap Var ExtVar -> [ExtVar] -> Renaming (Bimap Var ExtVar)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Bimap Var ExtVar -> ExtVar -> Renaming (Bimap Var ExtVar)
tryImport (Node -> Bimap Var ExtVar
nodeImportedVars Node
n) [ExtVar]
toImportVars

--------------------------------------------------------------------------------