{-# LANGUAGE LambdaCase, OverloadedStrings #-}

-- | ASTGen provides a function to render the program’s Abstract Syntax Tree
--   as a Graphviz DOT graph.
module ASTGen
  ( programToDot
  ) where

import Data.Text.Lazy (Text)
import qualified Data.Text.Lazy as T
import Syntax (Program(..), Decl(..), Expr(..), Pattern(..))

-- | Convert an entire 'Program' AST into a DOT-formatted graph.
--
-- Each top-level function declaration is numbered sequentially and rendered
-- as a subgraph.  The output is a valid Graphviz DOT description of the AST.
programToDot :: Program -> Text
programToDot :: Program -> Text
programToDot (Program [Decl]
decls) =
     Text
"digraph AST {\n"
  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"  node [shape=box, fontname=\"Courier\"];\n"
  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
T.concat ((Integer -> Decl -> Text) -> [Integer] -> [Decl] -> [Text]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Integer
i Decl
d -> Text -> Decl -> Text
declToDot (Text
"decl" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
i)) Decl
d) [Integer
0..] [Decl]
decls)
  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}\n"

-- | Render a single function declaration as a DOT subgraph.
--
-- The root node is labeled "FunDecl\n<name>(<params>)", and connected
-- to its body node, which recursively expands the expression tree.
declToDot :: Text -> Decl -> Text
declToDot :: Text -> Decl -> Text
declToDot Text
name (FunDecl [Char]
f [[Char]]
ps Expr
b) =
    let lbl :: Text
lbl   = Text
"FunDecl\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack [Char]
f
             Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
"," (([Char] -> Text) -> [[Char]] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> Text
T.pack [[Char]]
ps) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
        bodyN :: Text
bodyN = Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_body"
    in  Text -> Text -> Text
node Text
name Text
lbl
     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
node Text
bodyN (Expr -> Text
exprLabel Expr
b)
     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
edge Text
name Text
bodyN
     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Expr -> Text
exprToDot Text
bodyN Expr
b

-- | Generate a human-readable label for an expression node.
--
-- For example, a variable becomes "Var\nx", a lambda becomes "Lambda(x,y)"
-- showing its parameters, and binary operators are annotated with their name.
exprLabel :: Expr -> Text
exprLabel :: Expr -> Text
exprLabel = \case
  Var [Char]
x        -> Text
"Var\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack [Char]
x
  Lit Literal
_        -> Text
"Lit"
  Lambda [[Char]]
ps Expr
_  -> Text
"Lambda(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
"," (([Char] -> Text) -> [[Char]] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map [Char] -> Text
T.pack [[Char]]
ps) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
  If{}         -> Text
"If"
  Case{}       -> Text
"Case"
  Let{}        -> Text
"Let"
  App{}        -> Text
"App"
  BinOp BinOperator
op Expr
_ Expr
_ -> Text
"BinOp\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (BinOperator -> [Char]
forall a. Show a => a -> [Char]
show BinOperator
op)
  UnOp UnOperator
op Expr
_    -> Text
"UnOp\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (UnOperator -> [Char]
forall a. Show a => a -> [Char]
show UnOperator
op)
  List{}       -> Text
"List"
  Tuple{}      -> Text
"Tuple"

-- | Recursively render the children of an expression node.
--
-- Each child is given a name based on the parent prefix plus role,
-- connected by an edge, and then its subtree is expanded.
exprToDot :: Text -> Expr -> Text
exprToDot :: Text -> Expr -> Text
exprToDot Text
prefix Expr
expr = case Expr
expr of
  Var{}    -> Text
""
  Lit{}    -> Text
""
  Lambda [[Char]]
_ Expr
b ->
      Text -> Text -> Expr -> Text
child Text
prefix Text
"body" Expr
b

  If Expr
c Expr
t Expr
e ->
      Text -> [Text] -> [Expr] -> Text
children Text
prefix [Text
"cond",Text
"then",Text
"else"] [Expr
c,Expr
t,Expr
e]

  Case Expr
s [(Pattern, Expr)]
alts ->
      Text -> Text -> Expr -> Text
child Text
prefix Text
"scrut" Expr
s
    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
T.concat [ Text -> Int -> Pattern -> Expr -> Text
altToDot Text
prefix Int
i Pattern
pat Expr
bd
                | (Int
i,(Pattern
pat,Expr
bd)) <- [Int] -> [(Pattern, Expr)] -> [(Int, (Pattern, Expr))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [(Pattern, Expr)]
alts ]

  Let [Decl]
ds Expr
e ->
      [Text] -> Text
T.concat [ Text -> Decl -> Text
declToDot (Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_let" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
i)) Decl
d
               | (Integer
i,Decl
d) <- [Integer] -> [Decl] -> [(Integer, Decl)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [Decl]
ds ]
    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Expr -> Text
child Text
prefix Text
"in" Expr
e

  App Expr
f Expr
x ->
      Text -> [Text] -> [Expr] -> Text
children Text
prefix [Text
"fun",Text
"arg"] [Expr
f,Expr
x]

  BinOp BinOperator
_ Expr
l Expr
r ->
      Text -> [Text] -> [Expr] -> Text
children Text
prefix [Text
"l",Text
"r"] [Expr
l,Expr
r]

  UnOp UnOperator
_ Expr
x ->
      Text -> Text -> Expr -> Text
child Text
prefix Text
"arg" Expr
x

  List [Expr]
xs ->
      [Text] -> Text
T.concat [ Text -> Text -> Expr -> Text
child Text
prefix ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ [Char]
"e" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
i) Expr
x
               | (Integer
i,Expr
x) <- [Integer] -> [Expr] -> [(Integer, Expr)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [Expr]
xs ]

  Tuple [Expr]
xs ->
      [Text] -> Text
T.concat [ Text -> Text -> Expr -> Text
child Text
prefix ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ [Char]
"e" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
i) Expr
x
               | (Integer
i,Expr
x) <- [Integer] -> [Expr] -> [(Integer, Expr)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [Expr]
xs ]

  where
    -- | Render a single child node and connect it.
    child :: Text -> Text -> Expr -> Text
child Text
p Text
role Expr
e =
      let n :: Text
n = Text
p Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
role
      in  Text -> Text -> Text
node Text
n (Expr -> Text
exprLabel Expr
e)
       Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
edge Text
p Text
n
       Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Expr -> Text
exprToDot Text
n Expr
e

    -- | Render multiple children with given roles.
    children :: Text -> [Text] -> [Expr] -> Text
children Text
p [Text]
rs [Expr]
es = [Text] -> Text
T.concat ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (Text -> Expr -> Text) -> [Text] -> [Expr] -> [Text]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Text
r Expr
e -> Text -> Text -> Expr -> Text
child Text
p Text
r Expr
e) [Text]
rs [Expr]
es

-- | Render a case alternative (pattern and corresponding body) as DOT.
altToDot :: Text -> Int -> Pattern -> Expr -> Text
altToDot :: Text -> Int -> Pattern -> Expr -> Text
altToDot Text
prefix Int
i Pattern
pat Expr
bd =
  let pn :: Text
pn = Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_pat" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i)
      bn :: Text
bn = Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_bd"  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i)
  in  Text -> Text -> Text
node Text
pn (Pattern -> Text
patternLabel Pattern
pat)
   Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
edge Text
prefix Text
pn
   Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
node Text
bn Text
"AltBody"
   Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
edge Text
pn Text
bn
   Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Expr -> Text
exprToDot Text
bn Expr
bd

-- | Generate a label for a pattern node.
patternLabel :: Pattern -> Text
patternLabel :: Pattern -> Text
patternLabel = \case
  Pattern
PWildcard   -> Text
"_"
  PVar [Char]
x      -> [Char] -> Text
T.pack [Char]
x
  PLit Literal
_      -> Text
"LitPat"
  PList [Pattern]
_     -> Text
"ListPat"
  PTuple [Pattern]
_    -> Text
"TuplePat"

-- | Define a DOT node with the given name and label.
node :: Text -> Text -> Text
node :: Text -> Text -> Text
node Text
name Text
label =
  Text
"  " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" [label=\"" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
label Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\"];\n"

-- | Define a DOT edge between two nodes.
edge :: Text -> Text -> Text
edge :: Text -> Text -> Text
edge Text
from Text
to =
  Text
"  " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
from Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" -> " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
to Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
";\n"