module Haskell.AST where

import Data.List
import qualified Data.Set as Set
import Pretty


type Name = String
type TyVar = String

data AST = AST [Toplevel]
  deriving (Show, Eq)

data Toplevel = TopDef Def
              | TopDecl Decl
              | TopData Data
              | TopClass Class
              | TopInst Inst
  deriving (Show, Eq)

data Def = Def Name Expr
  deriving (Show, Eq)

data Expr = App Expr [Expr]
          | Ref Name
          | Con Name
          | Num Integer
          | Tup [Expr]
          | Lam [Name] Expr
          | Case Expr [(Pat, Expr)]
  deriving (Show, Eq)

data Pat = PatAny
         | PatVar Name
         | PatCon Name [Pat]
         | PatTup [Pat]
  deriving (Show, Eq)

data Decl = Decl Name Type
  deriving (Show, Eq)

data Type = TyTup [Type]
          | TyInt
          | TyFun Type Type
          | TyRef Name [Type]
          | TyVar Name
          | TyVoid
  deriving (Show, Eq)

data Data = Data Name [TyVar] [(Name, [Type])]
  deriving (Show, Eq)

data Class = Class Name [TyVar] [Decl]
  deriving (Show, Eq)

data Inst = Inst Name Type [Def]
  deriving (Show, Eq)


instance Pretty AST where
    pretty (AST tops) = Node "" (map pretty tops)

instance Pretty Toplevel where
    pretty (TopDef x) = pretty x
    pretty (TopDecl x) = pretty x
    pretty (TopData x) = pretty x
    pretty (TopClass x) = pretty x
    pretty (TopInst x) = pretty x

instance Pretty Def where
    pretty (Def n e) = Node (n ++ " =") [pretty e]

instance Pretty Expr where
    pretty (App e as) = Bracket "(" ")" "" (map pretty (e:as))
    pretty (Ref n) = Leaf n
    pretty (Con n) = Leaf n
    pretty (Num n) = Leaf (show n)
    pretty (Tup es) = Bracket "(" ")" "," (map pretty es)
    pretty (Lam as e) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
    pretty (Case e arms) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
      where go (p, e') = Node (pprintOneline p ++ " ->") [pretty e']

instance Pretty Pat where
    pretty PatAny = Leaf "_"
    pretty (PatVar n) = Leaf n
    pretty (PatCon n ps) = Bracket "(" ")" "" (Leaf n : map pretty ps)
    pretty (PatTup ps) = Bracket "(" ")" "," (map pretty ps)

instance Pretty Decl where
    pretty (Decl n t) = Node (n ++ " :: ") [pretty t]

instance Pretty Type where
    pretty (TyTup ts) = Bracket "(" ")" "," (map pretty ts)
    pretty TyInt = Leaf "Int"
    pretty (TyFun t u) = Leaf $ pprintOneline t ++ " -> " ++ pprintOneline u
    pretty (TyRef n as) = Bracket "(" ")" "" (Leaf n : map pretty as)
    pretty (TyVar n) = Leaf n
    pretty TyVoid = Leaf "#Void"

instance Pretty Data where
    pretty (Data n as cs) = Node ("data " ++ n ++ " " ++ intercalate " " as ++ " =") [Bracket "" "" "|" (map go cs)]
      where go (m, ts) = Node m (map pretty ts)

instance Pretty Class where
    pretty (Class n as ds) = Node ("class " ++ n ++ " " ++ intercalate " " as ++ " where") [Bracket "{" "}" ";" (map pretty ds)]

instance Pretty Inst where
    pretty (Inst n t ds) = Node ("instance " ++ n ++ " " ++ pprintOneline t ++ " where") [Bracket "{" "}" ";" (map pretty ds)]


-- This excludes constructor names, since those are not variables. This _does_
-- include bound variables; if you don't want that, use freeVariables.
class AllVars a where
    allVars :: a -> Set.Set Name

instance AllVars AST where
    allVars (AST tops) = Set.unions (map allVars tops)

instance AllVars Toplevel where
    allVars (TopDef def) = allVars def
    allVars (TopDecl _) = mempty
    allVars (TopData _) = mempty
    allVars (TopClass _) = mempty
    allVars (TopInst inst) = allVars inst

instance AllVars Def where
    allVars (Def n e) = Set.insert n (allVars e)

instance AllVars Inst where
    allVars (Inst _ _ ds) = Set.unions (map allVars ds)

instance AllVars Expr where
    allVars (App e es) = Set.unions (map allVars (e : es))
    allVars (Ref n) = Set.singleton n
    allVars (Con _) = mempty
    allVars (Num _) = mempty
    allVars (Tup es) = Set.unions (map allVars es)
    allVars (Lam ns e) = Set.fromList ns <> allVars e
    allVars (Case e pairs) =
        allVars e <> Set.unions [allVars p <> allVars e' | (p, e') <- pairs]

instance AllVars Pat where
    allVars PatAny = mempty
    allVars (PatVar n) = Set.singleton n
    allVars (PatCon _ ps) = Set.unions (map allVars ps)
    allVars (PatTup ps) = Set.unions (map allVars ps)


boundVars :: Pat -> Set.Set Name
boundVars PatAny = mempty
boundVars (PatVar n) = Set.singleton n
boundVars (PatCon _ ps) = Set.unions (map boundVars ps)
boundVars (PatTup ps) = Set.unions (map boundVars ps)

freeVariables :: Expr -> Set.Set Name
freeVariables (App e es) = freeVariables e <> Set.unions (map freeVariables es)
freeVariables (Ref n) = Set.singleton n
freeVariables (Con _) = mempty
freeVariables (Num _) = mempty
freeVariables (Tup es) = Set.unions (map freeVariables es)
freeVariables (Lam ns e) = freeVariables e Set.\\ Set.fromList ns
freeVariables (Case e pairs) =
    freeVariables e <> Set.unions [freeVariables e' Set.\\ boundVars p
                                  | (p, e') <- pairs]