content: {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  module Syntax where
  
  import           Data.Text(Text,unpack)
  import           Data.Hashable
  import           Data.Label
  import           Data.String
  import           Data.HashMap.Lazy(HashMap)
  import qualified Data.HashMap.Lazy as M
  import           Data.HashSet(HashSet)
  import qualified Data.HashSet as H
  import           Data.Text.Prettyprint.Doc
  
  import Control.Monad.State
  
  -- | Expressions of PCF. Each expression has a label, with which the
  -- expression can be uniquely identified.
  data Expr
    = Var Text Label
    | Lam [Text] Expr Label
    | App Expr [Expr] Label
    | Zero Label
    | Mult Expr Expr Label
    | Succ Expr Label
    | Pred Expr Label
    | IfZero Expr Expr Expr Label
    | Let [(Text,Expr)] Expr Label
    | Apply Expr Label
    deriving (Eq)
  
  -- Smart constructors that build labeled PCF expressions.
  
  var :: Text -> State Label Expr
  var x = Var x <$> fresh
  
  lam :: [Text] -> State Label Expr -> State Label Expr
  lam xs e = Lam xs <$> e <*> fresh
  
  app :: State Label Expr -> [State Label Expr] -> State Label Expr
  app e1 e2 = App <$> e1 <*> sequence e2 <*> fresh
  
  zero :: State Label Expr
  zero = Zero <$> fresh
  
  succ :: State Label Expr -> State Label Expr
  succ e = Succ <$> e <*> fresh
  
  pred :: State Label Expr -> State Label Expr
  pred e = Pred <$> e <*> fresh
  
  mult :: State Label Expr -> State Label Expr -> State Label Expr
  mult e1 e2 = Mult <$> e1 <*> e2 <*> fresh
  
  ifZero :: State Label Expr -> State Label Expr -> State Label Expr -> State Label Expr
  ifZero e1 e2 e3 = IfZero <$> e1 <*> e2 <*> e3 <*> fresh
  
  let_ :: [(Text, State Label Expr)] -> State Label Expr -> State Label Expr
  let_ bnds body = Let <$> sequence [ (v,) <$> e | (v,e) <- bnds ] <*> body <*> fresh
  
  
  instance Show Expr where
    showsPrec d e0 = case e0 of
      Var x _ -> showString (unpack x)
      Zero _ -> showString "zero"
      Succ e _ -> showParen (d > app_prec) $ showString "succ " . showsPrec (app_prec + 1) e
      Pred e _ -> showParen (d > app_prec) $ showString "pred " . showsPrec (app_prec + 1) e
      Let bnds e _ -> showParen (d > app_prec) $ showString "let " . shows bnds . showString " in " . showsPrec (app_prec + 1) e
      Apply e _ -> showParen (d > app_prec) $ showsPrec (app_prec + 1) e
      IfZero e1 e2 e3 _ -> showParen (d > app_prec)
        $ showString "ifZero "
        . showsPrec (app_prec + 1) e1
        . showString " "
        . showsPrec (app_prec + 1) e2
        . showString " "
        . showsPrec (app_prec + 1) e3
      App e1 e2 _ -> showParen (d > app_prec)
        $ showsPrec (app_prec + 1) e1
        . showString " "
        . showsPrec (app_prec + 1) e2
      Mult e1 e2 _ -> showParen (d > mult_prec)
        $ showsPrec (mult_prec + 1) e1
        . showString " * "
        . showsPrec (mult_prec + 1) e2
      Lam x e2 _ -> showParen (d > lam_prec)
        $ showString "λ"
        . showString (unwords (map unpack x))
        . showString ". "
        . shows e2
      where
        app_prec = 10
        lam_prec = 9
        mult_prec = 8
  
  instance Pretty Expr where
    pretty = viaShow
  
  instance HasLabel Expr where
    label e = case e of
      Var _ l -> l
      Lam _ _ l -> l
      App _ _ l -> l
      Zero l -> l
      Succ _ l -> l
      Pred _ l -> l
      Mult _ _ l -> l
      IfZero _ _ _ l -> l
      Let _ _ l -> l
      Apply _ l -> l
  
  instance IsString (State Label Expr) where
    fromString = var . fromString
  
  instance Hashable Expr where
    hashWithSalt s e = s `hashWithSalt` label e
  
  isFunctionBody :: (store,(env,Expr)) -> Bool
  isFunctionBody (_,(_,e)) = case e of
    Apply {} -> True
    _ -> False
  {-# INLINE isFunctionBody #-}
  
  isApplication :: (store,(env,Expr)) -> Bool
  isApplication (_,(_,e)) = case e of
    App {} -> True
    _ -> False
  
  freeVars :: Expr -> HashMap Expr (HashSet Text)
  freeVars e0 = execState (go e0) M.empty
    where
      go :: Expr -> State (HashMap Expr (HashSet Text)) (HashSet Text)
      go e = case e of
        Var x _ -> return (H.singleton x)
        Lam xs e1 _ -> save $ flip (foldr H.delete) xs <$> go e1
        App e1 e2 _ -> H.union <$> go e1 <*> (H.unions <$> mapM go e2)
        Zero _ -> return H.empty
        Succ e1 _ -> go e1
        Pred e1 _ -> go e1
        Mult e1 e2 _ -> H.union <$> go e1 <*> go e2
        IfZero e1 e2 e3 _ -> do
            m1 <- go e1
            m2 <- go e2
            m3 <- go e3
            return (m1 <> m2 <> m3)
        Let bnds e1 _ -> do
          vars <- go e1
          return (vars `H.difference` H.fromList [ v | (v,_) <- bnds])
        Apply e1 _ -> go e1
        where
          save m = do
            fv <- m
            modify (M.insert e fv)
            return fv
  