content: {-# LANGUAGE GADTs          #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE RankNTypes     #-}
  {-# LANGUAGE Safe           #-}
  {-# LANGUAGE ViewPatterns   #-}
  
  -- | Convert modular transition systems ('TransSys') into Kind2 file
  -- specifications.
  module Copilot.Theorem.Kind2.Translate
    ( toKind2
    , Style (..)
    ) where
  
  import Copilot.Theorem.TransSys
  import qualified Copilot.Theorem.Kind2.AST as K
  
  import Control.Exception.Base (assert)
  
  import Data.Function (on)
  import Data.Maybe (fromJust)
  
  import Data.List (sort, sortBy)
  import Data.Map (Map, (!))
  
  import qualified Data.Map as Map
  import qualified Data.Bimap as Bimap
  
  -- The following properties MUST hold for the given transition system :
  -- * Nodes are sorted by topological order
  -- * Nodes are `completed`, which means the dependency graph is transitive
  --   and each node imports all the local variables of its dependencies
  --
  
  type DepGraph = Map NodeId [NodeId]
  
  -- | Style of the Kind2 files produced: modular (with multiple separate nodes),
  -- or all inlined (with only one node).
  --
  -- In the modular style, the graph is simplified to remove cycles by collapsing
  -- all nodes participating in a strongly connected components.
  --
  -- In the inlined style, the structure of the modular transition system is
  -- discarded and the graph is first turned into a /non-modular transition/
  -- /system/ with only one node, which can be then converted into a Kind2 file.
  data Style = Inlined | Modular
  
  -- | Produce a Kind2 file that checks the properties specified.
  toKind2 :: Style     -- ^ Style of the file (modular or inlined).
          -> [PropId]  -- ^ Assumptions
          -> [PropId]  -- ^ Properties to be checked
          -> TransSys  -- ^ Modular transition system holding the system spec
          -> K.File
  toKind2 style assumptions checkedProps spec =
    addAssumptions spec assumptions
    $ trSpec (complete spec') predCallsGraph assumptions checkedProps
    where
      predCallsGraph = specDependenciesGraph spec'
      spec' = case style of
        Inlined -> inline spec
        Modular -> removeCycles spec
  
  trSpec :: TransSys -> DepGraph -> [PropId] -> [PropId] -> K.File
  trSpec spec predCallsGraph _assumptions checkedProps = K.File preds props
    where
      preds = map (trNode spec predCallsGraph) (specNodes spec)
      props = map trProp $
        filter ((`elem` checkedProps) . fst) $
          Map.toList $ Map.map fst $ specProps spec
  
  trProp :: (PropId, ExtVar) -> K.Prop
  trProp (pId, var) = K.Prop pId (trVar . extVarLocalPart $ var)
  
  trNode :: TransSys -> DepGraph -> Node -> K.PredDef
  trNode spec predCallsGraph node =
    K.PredDef { K.predId, K.predStateVars, K.predInit, K.predTrans }
    where
      predId = nodeId node
      predStateVars = gatherPredStateVars spec node
      predInit  = mkConj $ initLocals  node
                           ++ map (trExpr False) (nodeConstrs node)
                           ++ predCalls True spec predCallsGraph node
      predTrans = mkConj $ transLocals node
                           ++ map (trExpr True) (nodeConstrs node)
                           ++ predCalls False spec predCallsGraph node
  
  addAssumptions :: TransSys -> [PropId] -> K.File -> K.File
  addAssumptions spec assumptions (K.File {K.filePreds, K.fileProps}) =
    K.File (changeTail aux filePreds) fileProps
    where
      changeTail f (reverse -> l) = case l of
        []     -> error "impossible"
        x : xs -> reverse $ f x : xs
  
      aux pred =
        let init'  = mkConj ( K.predInit  pred : map K.StateVar vars )
            trans' = mkConj ( K.predTrans pred : map K.PrimedStateVar vars )
        in pred { K.predInit = init', K.predTrans = trans' }
  
      vars =
        let bindings   = nodeImportedVars (specTopNode spec)
            toExtVar a = fst $ fromJust $ Map.lookup a $ specProps spec
            toTopVar (ExtVar nId v) = assert (nId == specTopNodeId spec) v
        in map (varName . toTopVar . toExtVar) assumptions
  
  -- The ordering really matters here because the variables
  -- have to be given in this order in a pred call
  -- Our convention :
  -- * First the local variables, sorted by alphabetical order
  -- * Then the imported variables, by alphabetical order on
  --   the father node then by alphabetical order on the variable name
  
  gatherPredStateVars :: TransSys -> Node -> [K.StateVarDef]
  gatherPredStateVars spec node = locals ++ imported
    where
      nodesMap = Map.fromList [(nodeId n, n) | n <- specNodes spec]
      extVarType :: ExtVar -> K.Type
      extVarType (ExtVar n v) =
        case nodeLocalVars (nodesMap ! n) ! v of
          VarDescr Integer _ -> K.Int
          VarDescr Bool    _ -> K.Bool
          VarDescr Real    _ -> K.Real
  
      locals =
        map (\v -> K.StateVarDef (varName v)
                (extVarType $ ExtVar (nodeId node) v) [])
           . sort . Map.keys $ nodeLocalVars node
  
      imported =
        map (\(v, ev) -> K.StateVarDef (varName v) (extVarType ev) [])
        . sortBy (compare `on` snd) . Bimap.toList $ nodeImportedVars node
  
  mkConj :: [K.Term] -> K.Term
  mkConj []  = trConst Bool True
  mkConj [x] = x
  mkConj xs  = K.FunApp "and" xs
  
  mkEquality :: K.Term -> K.Term -> K.Term
  mkEquality t1 t2 = K.FunApp "=" [t1, t2]
  
  trVar :: Var -> K.Term
  trVar v = K.StateVar (varName v)
  
  trPrimedVar :: Var -> K.Term
  trPrimedVar v = K.PrimedStateVar (varName v)
  
  trConst :: Type t -> t -> K.Term
  trConst Integer v     = K.ValueLiteral (show v)
  trConst Real    v     = K.ValueLiteral (show v)
  trConst Bool    True  = K.ValueLiteral "true"
  trConst Bool    False = K.ValueLiteral "false"
  
  initLocals :: Node -> [K.Term]
  initLocals node =
    concatMap f (Map.toList $ nodeLocalVars node)
    where
      f (v, VarDescr t def) =
        case def of
          Pre     c _ -> [mkEquality (trVar v) (trConst t c)]
          Expr    e   -> [mkEquality (trVar v) (trExpr False e)]
          Constrs cs  -> map (trExpr False) cs
  
  transLocals :: Node -> [K.Term]
  transLocals node =
    concatMap f (Map.toList $ nodeLocalVars node)
    where
     f (v, VarDescr _ def) =
        case def of
          Pre _ v' -> [mkEquality (trPrimedVar v) (trVar v')]
          Expr e   -> [mkEquality (trPrimedVar v) (trExpr True e)]
          Constrs cs  -> map (trExpr True) cs
  
  predCalls :: Bool -> TransSys -> DepGraph -> Node -> [K.Term]
  predCalls isInitCall spec predCallsGraph node =
    map mkCall toCall
    where
      nid = nodeId node
      toCall = predCallsGraph ! nid
      nodesMap = Map.fromList [(nodeId n, n) | n <- specNodes spec]
  
      nodeLocals n =
        map (ExtVar n) . sort . Map.keys
        . nodeLocalVars $ (nodesMap ! n)
  
      mkCall callee
        | isInitCall =
            K.PredApp callee K.Init (argsSeq trVar)
        | otherwise  =
            K.PredApp callee K.Trans (argsSeq trVar ++ argsSeq trPrimedVar)
        where
  
          calleeLocals = nodeLocals callee
          calleeImported =
            (concatMap nodeLocals . sort . nodeDependencies) $ nodesMap ! callee
  
          localAlias trVarF ev =
            case Bimap.lookupR ev $ nodeImportedVars node of
              Nothing -> error $
                "This spec is not complete : "
                ++ show ev ++ " should be imported in " ++ nid
              Just v -> trVarF v
  
          argsSeq trVarF =
            map (localAlias trVarF) (calleeLocals ++ calleeImported)
  
  trExpr :: Bool -> Expr t -> K.Term
  trExpr primed = tr
    where
      tr :: forall t . Expr t -> K.Term
      tr (Const t c) = trConst t c
      tr (Ite _ c e1 e2) = K.FunApp "ite" [tr c, tr e1, tr e2]
      tr (Op1 _ op e) = K.FunApp (show op) [tr e]
      tr (Op2 _ op e1 e2) = K.FunApp (show op) [tr e1, tr e2]
      tr (VarE _ v) = if primed then trPrimedVar v else trVar v
  