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

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

-- | This module implements a pretty printer for the IL format, an intermediate
-- representation used in copilot-theorem to facilitate model checking.
module Copilot.Theorem.IL.PrettyPrint (prettyPrint, printConstraint) where

import Copilot.Theorem.IL.Spec
import Text.PrettyPrint.HughesPJ
import qualified Data.Map as Map

import Prelude hiding ((<>))

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

-- | Pretty print an IL specification.
prettyPrint :: IL -> String
prettyPrint :: IL -> String
prettyPrint = Doc -> String
render (Doc -> String) -> (IL -> Doc) -> IL -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IL -> Doc
ppSpec

-- | Pretty print an IL constraint expression.
printConstraint :: Expr -> String
printConstraint :: Expr -> String
printConstraint = Doc -> String
render (Doc -> String) -> (Expr -> Doc) -> Expr -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Doc
ppExpr

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

ppSpec :: IL -> Doc
ppSpec :: IL -> Doc
ppSpec (IL { [Expr]
modelInit :: IL -> [Expr]
modelInit :: [Expr]
modelInit, [Expr]
modelRec :: IL -> [Expr]
modelRec :: [Expr]
modelRec, Map String ([Expr], Expr)
properties :: IL -> Map String ([Expr], Expr)
properties :: Map String ([Expr], Expr)
properties }) =
  String -> Doc
text "MODEL INIT"
  Doc -> Doc -> Doc
$$ Doc -> Doc
indent ((Expr -> Doc -> Doc) -> Doc -> [Expr] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (Expr -> Doc) -> Expr -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Doc
ppExpr) Doc
empty [Expr]
modelInit) Doc -> Doc -> Doc
$$ Doc
emptyLine
  Doc -> Doc -> Doc
$$ String -> Doc
text "MODEL REC"
  Doc -> Doc -> Doc
$$ Doc -> Doc
indent ((Expr -> Doc -> Doc) -> Doc -> [Expr] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (Expr -> Doc) -> Expr -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Doc
ppExpr) Doc
empty [Expr]
modelRec) Doc -> Doc -> Doc
$$ Doc
emptyLine
  Doc -> Doc -> Doc
$$ String -> Doc
text "PROPERTIES"
  Doc -> Doc -> Doc
$$ Doc -> Doc
indent ((String -> ([Expr], Expr) -> Doc -> Doc)
-> Doc -> Map String ([Expr], Expr) -> Doc
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\k :: String
k -> Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc)
-> (([Expr], Expr) -> Doc) -> ([Expr], Expr) -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ([Expr], Expr) -> Doc
ppProp String
k)
        Doc
empty Map String ([Expr], Expr)
properties )

ppProp :: PropId -> ([Expr], Expr) -> Doc
ppProp :: String -> ([Expr], Expr) -> Doc
ppProp id :: String
id (as :: [Expr]
as, c :: Expr
c) = ((Expr -> Doc -> Doc) -> Doc -> [Expr] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Doc -> Doc -> Doc
($$) (Doc -> Doc -> Doc) -> (Expr -> Doc) -> Expr -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Doc
ppExpr) Doc
empty [Expr]
as)
  Doc -> Doc -> Doc
$$ Doc -> Doc
quotes (String -> Doc
text String
id) Doc -> Doc -> Doc
<+> Doc
colon Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
c

ppSeqDescr :: SeqDescr -> Doc
ppSeqDescr :: SeqDescr -> Doc
ppSeqDescr (SeqDescr id :: String
id ty :: Type
ty) = String -> Doc
text String
id Doc -> Doc -> Doc
<+> Doc
colon Doc -> Doc -> Doc
<+> Type -> Doc
ppType Type
ty

ppVarDescr :: VarDescr -> Doc
ppVarDescr :: VarDescr -> Doc
ppVarDescr (VarDescr id :: String
id ret :: Type
ret args :: [Type]
args) =
  String -> Doc
text String
id
  Doc -> Doc -> Doc
<+> Doc
colon
  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 "->" Doc -> Doc -> Doc
<> Doc
space) ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Type -> Doc) -> [Type] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Doc
ppType [Type]
args)
  Doc -> Doc -> Doc
<+> String -> Doc
text "->"
  Doc -> Doc -> Doc
<+> Type -> Doc
ppType Type
ret

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

ppExpr :: Expr -> Doc
ppExpr :: Expr -> Doc
ppExpr (ConstB v :: Bool
v) = String -> Doc
text (String -> Doc) -> (Bool -> String) -> Bool -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> String
forall a. Show a => a -> String
show (Bool -> Doc) -> Bool -> Doc
forall a b. (a -> b) -> a -> b
$ Bool
v
ppExpr (ConstR v :: Double
v) = String -> Doc
text (String -> Doc) -> (Double -> String) -> Double -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> String
forall a. Show a => a -> String
show (Double -> Doc) -> Double -> Doc
forall a b. (a -> b) -> a -> b
$ Double
v
ppExpr (ConstI _ v :: Integer
v) = String -> Doc
text (String -> Doc) -> (Integer -> String) -> Integer -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> String
forall a. Show a => a -> String
show (Integer -> Doc) -> Integer -> Doc
forall a b. (a -> b) -> a -> b
$ Integer
v

ppExpr (Ite _ c :: Expr
c e1 :: Expr
e1 e2 :: Expr
e2) =
  String -> Doc
text "if" Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
c
  Doc -> Doc -> Doc
<+> String -> Doc
text "then" Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
e1
  Doc -> Doc -> Doc
<+> String -> Doc
text "else" Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
e2

ppExpr (Op1 _ op :: Op1
op e :: Expr
e) = Op1 -> Doc
ppOp1 Op1
op Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
e

ppExpr (Op2 _ op :: Op2
op e1 :: Expr
e1 e2 :: Expr
e2) =
  Expr -> Doc
ppExpr Expr
e1 Doc -> Doc -> Doc
<+> Op2 -> Doc
ppOp2 Op2
op Doc -> Doc -> Doc
<+> Expr -> Doc
ppExpr Expr
e2

ppExpr (SVal _ s :: String
s i :: SeqIndex
i) = String -> Doc
text String
s Doc -> Doc -> Doc
<> Doc -> Doc
brackets (SeqIndex -> Doc
ppSeqIndex SeqIndex
i)

ppExpr (FunApp _ name :: String
name args :: [Expr]
args) =
  String -> Doc
text String
name Doc -> Doc -> Doc
<> Doc -> Doc
parens ([Doc] -> Doc
hsep ([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
$ (Expr -> Doc) -> [Expr] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Expr -> Doc
ppExpr [Expr]
args)

ppSeqIndex :: SeqIndex -> Doc
ppSeqIndex :: SeqIndex -> Doc
ppSeqIndex (Var i :: Integer
i)
  | Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== 0    = String -> Doc
text "n"
  | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< 0     = String -> Doc
text "n" Doc -> Doc -> Doc
<+> String -> Doc
text "-" Doc -> Doc -> Doc
<+> Integer -> Doc
integer (-Integer
i)
  | Bool
otherwise = String -> Doc
text "n" Doc -> Doc -> Doc
<+> String -> Doc
text "+" Doc -> Doc -> Doc
<+> Integer -> Doc
integer Integer
i

ppSeqIndex (Fixed i :: Integer
i) = Integer -> Doc
integer Integer
i

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

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

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