{-# LANGUAGE GADTs          #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE Safe           #-}

-- | Pretty print a TransSys specification as a Kind2/Lustre specification.
module Copilot.Theorem.TransSys.PrettyPrint ( prettyPrint ) where

import Copilot.Theorem.TransSys.Spec

import Text.PrettyPrint.HughesPJ

import qualified Data.Map   as Map
import qualified Data.Bimap as Bimap

import Prelude hiding ((<>))

indent :: Doc -> Doc
indent     = Int -> Doc -> Doc
nest Int
4
emptyLine :: Doc
emptyLine  = String -> Doc
text String
""

-- | Pretty print a TransSys specification as a Kind2/Lustre specification.
prettyPrint :: TransSys -> String
prettyPrint :: TransSys -> String
prettyPrint = Doc -> String
render forall b c a. (b -> c) -> (a -> b) -> a -> c
. TransSys -> Doc
pSpec

pSpec :: TransSys -> Doc
pSpec :: TransSys -> Doc
pSpec TransSys
spec = Doc
items Doc -> Doc -> Doc
$$ Doc
props
  where
    items :: Doc
items = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Node -> Doc
pNode) Doc
empty (TransSys -> [Node]
specNodes TransSys
spec)
    props :: Doc
props = String -> Doc
text String
"PROPS" Doc -> Doc -> Doc
$$
      forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\String
k -> Doc -> Doc -> Doc
($$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ExtVar -> Doc
pProp String
k)
        Doc
empty (TransSys -> Map String ExtVar
specProps TransSys
spec)

pProp :: String -> ExtVar -> Doc
pProp String
pId ExtVar
extvar = Doc -> Doc
quotes (String -> Doc
text String
pId) Doc -> Doc -> Doc
<+> String -> Doc
text String
"is" Doc -> Doc -> Doc
<+> ExtVar -> Doc
pExtVar ExtVar
extvar

pType :: Type t -> Doc
pType :: forall t. Type t -> Doc
pType = String -> Doc
text forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show

pList :: (t -> Doc) -> [t] -> Doc
pList :: forall t. (t -> Doc) -> [t] -> Doc
pList t -> Doc
f [t]
l = Doc -> Doc
brackets ([Doc] -> Doc
hcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> [Doc] -> [Doc]
punctuate (Doc
comma Doc -> Doc -> Doc
<> Doc
space) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map t -> Doc
f [t]
l)

pNode :: Node -> Doc
pNode :: Node -> Doc
pNode Node
n =
  Doc
header Doc -> Doc -> Doc
$$ Doc
imported Doc -> Doc -> Doc
$$ Doc
local Doc -> Doc -> Doc
$$ Doc
constrs Doc -> Doc -> Doc
$$ Doc
emptyLine
  where
    header :: Doc
header =
      String -> Doc
text String
"NODE"
      Doc -> Doc -> Doc
<+> Doc -> Doc
quotes (String -> Doc
text forall a b. (a -> b) -> a -> b
$ Node -> String
nodeId Node
n)
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"DEPENDS ON"
      Doc -> Doc -> Doc
<+> forall t. (t -> Doc) -> [t] -> Doc
pList String -> Doc
text (Node -> [String]
nodeDependencies Node
n)

    imported :: Doc
imported
      | forall a b. Bimap a b -> Bool
Bimap.null (Node -> Bimap Var ExtVar
nodeImportedVars Node
n) = Doc
empty
      | Bool
otherwise = String -> Doc
text String
"IMPORTS" Doc -> Doc -> Doc
$$ Doc -> Doc
indent
        (forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\Var
k -> Doc -> Doc -> Doc
($$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var -> ExtVar -> Doc
pIVar Var
k)
        Doc
empty (forall a b. Bimap a b -> Map a b
Bimap.toMap forall a b. (a -> b) -> a -> b
$ Node -> Bimap Var ExtVar
nodeImportedVars Node
n))

    local :: Doc
local
      | forall k a. Map k a -> Bool
Map.null (Node -> Map Var VarDescr
nodeLocalVars Node
n) = Doc
empty
      | Bool
otherwise = String -> Doc
text String
"DEFINES" Doc -> Doc -> Doc
$$ Doc -> Doc
indent
        (forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\Var
k -> Doc -> Doc -> Doc
($$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var -> VarDescr -> Doc
pLVar Var
k)
        Doc
empty (Node -> Map Var VarDescr
nodeLocalVars Node
n))

    constrs :: Doc
constrs = case Node -> [Expr Bool]
nodeConstrs Node
n of
      [] -> Doc
empty
      [Expr Bool]
l  -> String -> Doc
text String
"WITH CONSTRAINTS" Doc -> Doc -> Doc
$$
            forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Expr t -> Doc
pExpr) Doc
empty [Expr Bool]
l

pConst :: Type t -> t -> Doc
pConst :: forall t. Type t -> t -> Doc
pConst Type t
Integer t
v = String -> Doc
text forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show t
v
pConst Type t
Real    t
v = String -> Doc
text forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show t
v
pConst Type t
Bool    t
v = String -> Doc
text forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show t
v

pExtVar :: ExtVar -> Doc
pExtVar :: ExtVar -> Doc
pExtVar (ExtVar String
n Var
v) = Doc -> Doc
parens (String -> Doc
text String
n Doc -> Doc -> Doc
<+> String -> Doc
text String
":" Doc -> Doc -> Doc
<+> String -> Doc
text (Var -> String
varName Var
v))

pIVar :: Var -> ExtVar -> Doc
pIVar :: Var -> ExtVar -> Doc
pIVar Var
v ExtVar
ev =
  ExtVar -> Doc
pExtVar ExtVar
ev
  Doc -> Doc -> Doc
<+> String -> Doc
text String
"as" Doc -> Doc -> Doc
<+> Doc -> Doc
quotes (String -> Doc
text (Var -> String
varName Var
v))

pLVar :: Var -> VarDescr -> Doc
pLVar :: Var -> VarDescr -> Doc
pLVar Var
l (VarDescr {Type t
varType :: ()
varType :: Type t
varType, VarDef t
varDef :: ()
varDef :: VarDef t
varDef}) = Doc
header Doc -> Doc -> Doc
$$ Doc -> Doc
indent Doc
body
  where
    header :: Doc
header =
      String -> Doc
text (Var -> String
varName Var
l)
      Doc -> Doc -> Doc
<+> String -> Doc
text String
":"
      Doc -> Doc -> Doc
<+> forall t. Type t -> Doc
pType Type t
varType
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"="

    body :: Doc
body = case VarDef t
varDef of
      Pre t
val Var
var ->
        forall t. Type t -> t -> Doc
pConst Type t
varType t
val
        Doc -> Doc -> Doc
<+> String -> Doc
text String
"->" Doc -> Doc -> Doc
<+> String -> Doc
text String
"pre"
        Doc -> Doc -> Doc
<+> String -> Doc
text (Var -> String
varName Var
var)
      Expr Expr t
e -> forall t. Expr t -> Doc
pExpr Expr t
e

      Constrs [Expr Bool]
cs ->
        String -> Doc
text String
"{"
        Doc -> Doc -> Doc
<+> ([Doc] -> Doc
hsep forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> [Doc] -> [Doc]
punctuate (Doc
space Doc -> Doc -> Doc
<> String -> Doc
text String
";" Doc -> Doc -> Doc
<> Doc
space)) (forall a b. (a -> b) -> [a] -> [b]
map forall t. Expr t -> Doc
pExpr [Expr Bool]
cs)
        Doc -> Doc -> Doc
<+> String -> Doc
text String
"}"

pExpr :: Expr t -> Doc

pExpr :: forall t. Expr t -> Doc
pExpr (Const Type t
t t
v) = forall t. Type t -> t -> Doc
pConst Type t
t t
v

pExpr (Ite Type t
_ Expr Bool
c Expr t
e1 Expr t
e2) =
  String -> Doc
text String
"if" Doc -> Doc -> Doc
<+> forall t. Expr t -> Doc
pExpr Expr Bool
c
  Doc -> Doc -> Doc
<+> String -> Doc
text String
"then" Doc -> Doc -> Doc
<+> forall t. Expr t -> Doc
pExpr Expr t
e1
  Doc -> Doc -> Doc
<+> String -> Doc
text String
"else" Doc -> Doc -> Doc
<+> forall t. Expr t -> Doc
pExpr Expr t
e2

pExpr (Op1 Type t
_ Op1 t
op Expr t
e) = forall a. Op1 a -> Doc
pOp1 Op1 t
op Doc -> Doc -> Doc
<+> Doc -> Doc
parens (forall t. Expr t -> Doc
pExpr Expr t
e)

pExpr (Op2 Type t
_ Op2 a t
op Expr a
e1 Expr a
e2) =
  Doc -> Doc
parens (forall t. Expr t -> Doc
pExpr Expr a
e1) Doc -> Doc -> Doc
<+> forall a b. Op2 a b -> Doc
pOp2 Op2 a t
op Doc -> Doc -> Doc
<+> Doc -> Doc
parens (forall t. Expr t -> Doc
pExpr Expr a
e2)

pExpr (VarE Type t
_ Var
v) = String -> Doc
text (Var -> String
varName Var
v)

pOp1 :: Op1 a -> Doc
pOp1 :: forall a. Op1 a -> Doc
pOp1 = String -> Doc
text forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show

pOp2 :: Op2 a b -> Doc
pOp2 :: forall a b. Op2 a b -> Doc
pOp2 = String -> Doc
text forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show