    {-# 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 (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 a b. (a -> b -> b) -> b -> [a] -> b
    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
            (((ExtVar, Prop) -> ExtVar)
    -> Map String (ExtVar, Prop) -> Map String ExtVar
    forall a b k. (a -> b) -> Map k a -> Map k b
    Map.map (ExtVar, Prop) -> ExtVar
    forall a b. (a, b) -> a
    fst (TransSys -> Map String (ExtVar, Prop)
    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 (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 :: forall t. (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 a b. (a -> b -> b) -> b -> [a] -> b
    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 :: forall t. 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 :: Type t
    varType :: ()
    varType, VarDef t
    varDef :: VarDef t
    varDef :: ()
    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 :: forall t. 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 :: forall a. 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 :: forall a b. 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