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

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

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
""

prettyPrint :: TransSys -> String
prettyPrint :: TransSys -> String
prettyPrint = Doc -> String
render (Doc -> String) -> (TransSys -> Doc) -> TransSys -> String
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 = (Node -> Doc -> Doc) -> Doc -> [Node] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (Node -> Doc) -> Node -> 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
$$
      (String -> ExtVar -> Doc -> Doc) -> Doc -> Map String ExtVar -> Doc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\String
k -> Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (ExtVar -> Doc) -> ExtVar -> 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 :: Type t -> Doc
pType = String -> Doc
text (String -> Doc) -> (Type t -> String) -> Type t -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type t -> String
forall a. Show a => a -> String
show

pList :: (t -> Doc) -> [t] -> Doc
pList :: (t -> Doc) -> [t] -> Doc
pList t -> Doc
f [t]
l = Doc -> Doc
brackets ([Doc] -> Doc
hcat ([Doc] -> Doc) -> ([Doc] -> [Doc]) -> [Doc] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> [Doc] -> [Doc]
punctuate (Doc
comma Doc -> Doc -> Doc
<> Doc
space) ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (t -> Doc) -> [t] -> [Doc]
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 (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ Node -> String
nodeId Node
n)
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"DEPENDS ON"
      Doc -> Doc -> Doc
<+> (String -> Doc) -> [String] -> Doc
forall t. (t -> Doc) -> [t] -> Doc
pList String -> Doc
text (Node -> [String]
nodeDependencies Node
n)

    imported :: Doc
imported
      | Bimap Var ExtVar -> Bool
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
        ((Var -> ExtVar -> Doc -> Doc) -> Doc -> Map Var ExtVar -> Doc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\Var
k -> Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (ExtVar -> Doc) -> ExtVar -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var -> ExtVar -> Doc
pIVar Var
k)
        Doc
empty (Bimap Var ExtVar -> Map Var ExtVar
forall a b. Bimap a b -> Map a b
Bimap.toMap (Bimap Var ExtVar -> Map Var ExtVar)
-> Bimap Var ExtVar -> Map Var ExtVar
forall a b. (a -> b) -> a -> b
$ Node -> Bimap Var ExtVar
nodeImportedVars Node
n))

    local :: Doc
local
      | Map Var VarDescr -> Bool
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
        ((Var -> VarDescr -> Doc -> Doc) -> Doc -> Map Var VarDescr -> Doc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\Var
k -> Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (VarDescr -> Doc) -> VarDescr -> 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
$$
            (Expr Bool -> Doc -> Doc) -> Doc -> [Expr Bool] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc)
-> (Expr Bool -> Doc) -> Expr Bool -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr Bool -> Doc
forall t. Expr t -> Doc
pExpr) Doc
empty [Expr Bool]
l

pConst :: Type t -> t -> Doc
pConst :: Type t -> t -> Doc
pConst Type t
Integer t
v = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ t -> String
forall a. Show a => a -> String
show t
v
pConst Type t
Real    t
v = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ t -> String
forall a. Show a => a -> String
show t
v
pConst Type t
Bool    t
v = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ t -> String
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
<+> Type t -> 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 ->
            Type t -> t -> Doc
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 -> Expr t -> Doc
forall t. Expr t -> Doc
pExpr Expr t
e

          Constrs [Expr Bool]
cs ->
            String -> Doc
text String
"{"
            Doc -> Doc -> Doc
<+> ([Doc] -> Doc
hsep ([Doc] -> Doc) -> ([Doc] -> [Doc]) -> [Doc] -> Doc
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)) ((Expr Bool -> Doc) -> [Expr Bool] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Expr Bool -> Doc
forall t. Expr t -> Doc
pExpr [Expr Bool]
cs)
            Doc -> Doc -> Doc
<+> String -> Doc
text String
"}"


pExpr :: Expr t -> Doc

pExpr :: Expr t -> Doc
pExpr (Const Type t
t t
v) = Type t -> t -> Doc
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
<+> Expr Bool -> Doc
forall t. Expr t -> Doc
pExpr Expr Bool
c
  Doc -> Doc -> Doc
<+> String -> Doc
text String
"then" Doc -> Doc -> Doc
<+> Expr t -> Doc
forall t. Expr t -> Doc
pExpr Expr t
e1
  Doc -> Doc -> Doc
<+> String -> Doc
text String
"else" Doc -> Doc -> Doc
<+> Expr t -> Doc
forall t. Expr t -> Doc
pExpr Expr t
e2

pExpr (Op1 Type t
_ Op1 t
op Expr t
e) = Op1 t -> Doc
forall a. Op1 a -> Doc
pOp1 Op1 t
op Doc -> Doc -> Doc
<+> Doc -> Doc
parens (Expr t -> Doc
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 (Expr a -> Doc
forall t. Expr t -> Doc
pExpr Expr a
e1) Doc -> Doc -> Doc
<+> Op2 a t -> Doc
forall a b. Op2 a b -> Doc
pOp2 Op2 a t
op Doc -> Doc -> Doc
<+> Doc -> Doc
parens (Expr a -> Doc
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 :: Op1 a -> Doc
pOp1 = String -> Doc
text (String -> Doc) -> (Op1 a -> String) -> Op1 a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op1 a -> String
forall a. Show a => a -> String
show

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

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