fwl/src/FWL/Pretty.hs

-- | Pretty printer: round-trips the AST back to FWL source.
module FWL.Pretty (prettyProgram) where

import Data.List (intercalate)
import FWL.AST

prettyProgram :: Program -> String
prettyProgram (Program cfg ds) =
  prettyConfig cfg ++ "\n" ++ unlines (map prettyDecl ds)

prettyConfig :: Config -> String
prettyConfig (Config t)
  | t == "fwl" = ""
  | otherwise  = "config { table = \"" ++ t ++ "\"; }\n"

prettyDecl :: Decl -> String
prettyDecl (DInterface n k ps) =
  "interface " ++ n ++ " : " ++ prettyKind k ++ " {\n" ++
  concatMap (\p -> "  " ++ prettyIfaceProp p ++ ";\n") ps ++
  "};"
prettyDecl (DZone n ns) =
  "zone " ++ n ++ " = { " ++ intercalate ", " ns ++ " };"
prettyDecl (DImport n t s) =
  "import " ++ n ++ " : " ++ prettyType t ++ " from \"" ++ s ++ "\";"
prettyDecl (DLet n t e) =
  "let " ++ n ++ " : " ++ prettyType t ++ " = " ++ prettyExpr e ++ ";"
prettyDecl (DPattern n t p) =
  "pattern " ++ n ++ " : " ++ prettyType t ++ " = " ++ prettyPat p ++ ";"
prettyDecl (DFlow n f) =
  "flow " ++ n ++ " : FlowPattern = " ++ prettyFlow f ++ ";"
prettyDecl (DRule n t e) =
  "rule " ++ n ++ " : " ++ prettyType t ++ " =\n  " ++ prettyExpr e ++ ";"
prettyDecl (DPolicy n t pm ab) =
  "policy " ++ n ++ " : " ++ prettyType t ++ "\n" ++
  "  on { hook = " ++ prettyHook (pmHook pm) ++
       ", table = " ++ prettyTable (pmTable pm) ++
       ", priority = " ++ prettyPriority (pmPriority pm) ++ " }\n" ++
  "  = " ++ prettyArmBlock ab ++ ";"

prettyKind :: IfaceKind -> String
prettyKind IWan        = "WAN"
prettyKind ILan        = "LAN"
prettyKind IWireGuard  = "WireGuard"
prettyKind (IUser n)   = n

prettyIfaceProp :: IfaceProp -> String
prettyIfaceProp IPDynamic    = "dynamic"
prettyIfaceProp (IPCidr4 cs) = "cidr4 = { " ++ intercalate ", " (map prettyCidr cs) ++ " }"
prettyIfaceProp (IPCidr6 cs) = "cidr6 = { " ++ intercalate ", " (map prettyCidr cs) ++ " }"

prettyCidr :: CIDR -> String
prettyCidr (ip, p) = prettyLit ip ++ "/" ++ show p

prettyHook :: Hook -> String
prettyHook HInput       = "Input"
prettyHook HForward     = "Forward"
prettyHook HOutput      = "Output"
prettyHook HPrerouting  = "Prerouting"
prettyHook HPostrouting = "Postrouting"

prettyTable :: TableName -> String
prettyTable TFilter = "Filter"
prettyTable TNAT    = "NAT"

prettyPriority :: Priority -> String
prettyPriority p = show (priorityValue p)

prettyType :: Type -> String
prettyType (TName n [])   = n
prettyType (TName n ts)   = n ++ "<" ++ intercalate ", " (map prettyType ts) ++ ">"
prettyType (TTuple ts)    = "(" ++ intercalate ", " (map prettyType ts) ++ ")"
prettyType (TFun a b)     = prettyType a ++ " -> " ++ prettyType b
prettyType (TEffect es t) = "<" ++ intercalate ", " es ++ "> " ++ prettyType t

prettyPat :: Pat -> String
prettyPat PWild            = "_"
prettyPat (PVar n)         = n
prettyPat (PNamed n)       = n
prettyPat (PCtor n ps)     = n ++ "(" ++ intercalate ", " (map prettyPat ps) ++ ")"
prettyPat (PRecord n fs)   = n ++ " { " ++ intercalate ", " (map prettyFP fs) ++ " }"
prettyPat (PTuple ps)      = "(" ++ intercalate ", " (map prettyPat ps) ++ ")"
prettyPat (PFrame mp inner)=
  "Frame(" ++ maybe "" (\pp -> prettyPath pp ++ ", ") mp ++ prettyPat inner ++ ")"
prettyPat (PBytes bs)      = "[" ++ unwords (map prettyBE bs) ++ "]"
prettyPat (POr p1 p2)      = prettyPat p1 ++ " | " ++ prettyPat p2

prettyFP :: FieldPat -> String
prettyFP (FPEq n l)  = n ++ " = " ++ prettyLit l
prettyFP (FPBind n)  = n
prettyFP (FPAs n v)  = n ++ " as " ++ v

prettyPath :: PathPat -> String
prettyPath (PathPat ms md) =
  maybe "_" prettyEP ms ++ maybe "" (\d -> " -> " ++ prettyEP d) md

prettyEP :: EndpointPat -> String
prettyEP EPWild         = "_"
prettyEP (EPName n)     = n
prettyEP (EPMember n z) = n ++ " in " ++ z

prettyBE :: ByteElem -> String
prettyBE (BEHex w) = "0x" ++ pad (show w)  -- simplified
  where pad s = if length s < 2 then '0':s else s
prettyBE BEWild     = "_"
prettyBE BEWildStar = "_*"

prettyFlow :: FlowExpr -> String
prettyFlow (FAtom n)      = n
prettyFlow (FSeq a b mw)  =
  prettyFlow a ++ " . " ++ prettyFlow b ++
  maybe "" (\(n,u) -> " within " ++ show n ++ prettyUnit u) mw

prettyUnit :: TimeUnit -> String
prettyUnit Seconds = "s"
prettyUnit Millis  = "ms"
prettyUnit Minutes = "m"
prettyUnit Hours   = "h"

prettyExpr :: Expr -> String
prettyExpr (EVar n)         = n
prettyExpr (EQual ns)       = intercalate "." ns
prettyExpr (ELit l)         = prettyLit l
prettyExpr (ELam n e)       = "\\" ++ n ++ " -> " ++ prettyExpr e
prettyExpr (EApp f x)       = prettyExpr f ++ " " ++ prettyAtom x
prettyExpr (ECase e ab)     =
  "case " ++ prettyExpr e ++ " of " ++ prettyArmBlock ab
prettyExpr (EIf c t f)      =
  "if " ++ prettyExpr c ++ " then " ++ prettyExpr t ++ " else " ++ prettyExpr f
prettyExpr (EDo ss)         =
  "do { " ++ intercalate "; " (map prettyStmt ss) ++ " }"
prettyExpr (ELet n e1 e2)   =
  "let " ++ n ++ " = " ++ prettyExpr e1 ++ " in " ++ prettyExpr e2
prettyExpr (ETuple es)      = "(" ++ intercalate ", " (map prettyExpr es) ++ ")"
prettyExpr (ESet es)        = "{ " ++ intercalate ", " (map prettyExpr es) ++ " }"
prettyExpr (EMap ms)        =
  "{ " ++ intercalate ", " (map (\(k,v) -> prettyExpr k ++ " -> " ++ prettyExpr v) ms) ++ " }"
prettyExpr (EPerform ns as_) =
  "perform " ++ intercalate "." ns ++ "(" ++ intercalate ", " (map prettyExpr as_) ++ ")"
prettyExpr (EInfix op l r)  =
  prettyAtom l ++ " " ++ prettyOp op ++ " " ++ prettyAtom r
prettyExpr (ENot e)         = "!" ++ prettyAtom e

prettyAtom :: Expr -> String
prettyAtom e@(EInfix _ _ _) = "(" ++ prettyExpr e ++ ")"
prettyAtom e@(ELam _ _)     = "(" ++ prettyExpr e ++ ")"
prettyAtom e                 = prettyExpr e

prettyOp :: InfixOp -> String
prettyOp OpAnd    = "&&"
prettyOp OpOr     = "||"
prettyOp OpEq     = "=="
prettyOp OpNeq    = "!="
prettyOp OpLt     = "<"
prettyOp OpLte    = "<="
prettyOp OpGt     = ">"
prettyOp OpGte    = ">="
prettyOp OpIn     = "in"
prettyOp OpConcat = "++"
prettyOp OpThen   = ">>"
prettyOp OpBind   = ">>="

prettyStmt :: DoStmt -> String
prettyStmt (DSBind n e) = n ++ " <- " ++ prettyExpr e
prettyStmt (DSExpr e)   = prettyExpr e

prettyArmBlock :: ArmBlock -> String
prettyArmBlock arms =
  "{\n" ++
  concatMap (\(Arm p mg e) ->
    "    | " ++ prettyPat p ++
    maybe "" (\g -> " if " ++ prettyExpr g) mg ++
    " -> " ++ prettyExpr e ++ ";\n") arms ++
  "  }"

prettyLit :: Literal -> String
prettyLit (LInt n)           = show n
prettyLit (LString s)        = "\"" ++ s ++ "\""
prettyLit (LBool True)       = "true"
prettyLit (LBool False)      = "false"
prettyLit (LIP IPv4 n)       = renderIPv4 n
prettyLit (LIP IPv6 n)       = renderIPv6 n
prettyLit (LCIDR ip p)       = prettyLit ip ++ "/" ++ show p
prettyLit (LPort p)          = ":" ++ show p
prettyLit (LDuration n u)    = show n ++ prettyUnit u
prettyLit (LHex b)           = "0x" ++ show b