-----------------------------------------------------------------------------
-- |
-- Module      :  Expr
-- Copyright   :  (c) Masahiro Sakai 2007-2009
-- License     :  BSD3-style (see LICENSE)
-- 
-- Maintainer:    masahiro.sakai@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable

{-# LANGUAGE CPP, GeneralizedNewtypeDeriving #-}
module Expr
    ( Op1 (..)
    , Op2 (..)
    , Binder (..)
    , Expr (..)
    , Scope (..)
    , Name
    , lambda
    , forall
    , exists
    , int
    , ext
    , (<@>)
    , abstract
    , instantiate
    , normalize
    , RenderM
    , runRenderM
    , renderExpr
    , renderExprU
    ) where

import Control.Monad.RWS
import Data.List

-----------------------------------------------------------------------------

data Op1 = Not | Box | F | H | Int | Ext deriving (Eq,Show)
data Op2 = And | Or | Imply | Equiv | Id deriving (Eq,Show)
data Binder = Lambda | Forall | Exists deriving (Eq,Show)

infixl 9 :@

data Expr
    = FVar Name          -- 自由変数
    | BVar !Int          -- 束縛変数
    | Expr :@ Expr       -- 関数適用
    | Const String       -- 定数 (≒自由変数)
    | Op1 !Op1 Expr      -- 前置演算子
    | Op2 !Op2 Expr Expr -- 中置演算子
    | Bind !Binder Scope -- 変数束縛

instance Show Expr where
#ifdef USE_UTF8
    showsPrec d x = runRenderM (renderExprU d x)
#else
    showsPrec d x = runRenderM (renderExpr d x)
#endif

newtype Scope = Sc Expr deriving Show
type Name = String

-----------------------------------------------------------------------------

lambda :: Name -> Expr -> Expr
lambda name expr = Bind Lambda (abstract name expr)

forall :: Name -> Expr -> Expr
forall name expr = Bind Forall (abstract name expr)

exists :: Name -> Expr -> Expr
exists name expr = Bind Exists (abstract name expr)

int :: Expr -> Expr
int = Op1 Int

ext :: Expr -> Expr
ext = Op1 Ext

-- 「a{b}」を「a <@> b」と表記
infixl 9 <@>
(<@>) :: Expr -> Expr -> Expr
fun <@> arg = ext fun :@ arg

-----------------------------------------------------------------------------

varChange :: (Int -> Name -> Expr) -> (Int -> Int -> Expr) -> Expr -> Expr
varChange f g = h 0
  where
    h :: Int -> Expr -> Expr
    h outer (FVar name)  = f outer name
    h outer (BVar index) = g outer index
    h outer (Bind q (Sc body)) = Bind q (Sc (h (outer+1) body))
    h outer (fun :@ arg) = h outer fun :@ h outer arg
    h _     (Const s)    = Const s
    h outer (Op1 op a)   = Op1 op (h outer a)
    h outer (Op2 op a b) = Op2 op (h outer a) (h outer b)

abstract :: Name -> Expr -> Scope
abstract name expr = Sc (varChange f g expr)
  where
    f outer name' | name==name'  = BVar outer
                  | otherwise    = FVar name'
    g outer index | index>=outer = BVar (index+1)
                  | otherwise    = BVar index

instantiate :: Expr -> Scope -> Expr
instantiate image (Sc body) = varChange f g body
  where
    f _ name = FVar name
    g outer index | index==outer = varShift outer image
                  | index>outer  = BVar (index-1)
                  | otherwise    = BVar index

-- 外を指している変数のインデックスをずらす
varShift :: Int -> Expr -> Expr
varShift 0 = id
varShift n = varChange f g 
  where
    f _ name  = FVar name
    g outer index | index>=outer = BVar (index+n)
                  | otherwise    = BVar index

normalize :: Expr -> Expr
normalize (Bind Lambda (Sc body)) =
  case normalize body of
    f :@ BVar 0 | not (0 `elem` bvs f) -> varShift (-1) f -- η-conversion
    body -> Bind Lambda (Sc body)
normalize (Bind q (Sc body)) = Bind q (Sc (normalize body))
normalize (fun :@ arg) =
  case normalize fun of
    Bind Lambda body -> normalize (instantiate arg' body) -- β-reduction
    fun -> fun :@ arg'
  where arg' = normalize arg
normalize (Op1 Ext a)  =
  case normalize a of
    Op1 Int b -> b
    a' -> Op1 Ext a'
normalize (Op1 op a)   = Op1 op (normalize a)
normalize (Op2 op a b) = Op2 op (normalize a) (normalize b)
normalize x = x

bvs :: Expr -> [Int]
bvs (FVar _) = []
bvs (BVar n) = [n]
bvs (f :@ x) = bvs f ++ bvs x
bvs (Const _) = []
bvs (Op1 _ e) = bvs e
bvs (Op2 _ e1 e2) = bvs e1 ++ bvs e2
bvs (Bind b (Sc e)) = [n - 1 | n <- bvs e, n /= 0]

-----------------------------------------------------------------------------

newtype RenderM a = RenderM (RWS [Name] () Int a)
  deriving (Monad, MonadReader [Name])

runRenderM :: RenderM a -> a
runRenderM (RenderM m) = 
  case runRWS m [] 0 of
    (a, _, _) -> a

render :: Bool -> Int -> Expr -> RenderM ShowS
render u = h
  where
    h d e = case e of
      FVar name  -> return $ showString name
      BVar index -> do
        vs <- ask
        return $ showString (vs !! index)
      Bind q _ -> f d q e
      Op1 Ext a :@ b -> do
        a' <- h (app_prec+1) a
        b' <- h 0 b
        return $ showParen (d > app_prec)
               $ a' . showString " {" . b' . showChar '}'
      a :@ b  -> do 
        a' <- h app_prec a
        b' <- h (app_prec+1) b
        return $ showParen (d > app_prec) $ a' . showChar ' ' . b'
      Const s -> return $ showString s
      Op1 op a -> do
        t <- h (prec+1) a
        return $ showParen (d > prec)
               $ showString s . t
        where
          s = case op of
                Not -> if u then "¬" else "not " -- ¬ (U+00AC)
                Box -> if u then "◻" else "[]"  -- ◻ (U+25FB) が正しそうだが □ (U+25A1) を使うのが無難か?
                F   -> "F "
                H   -> "H "
                Int -> if u then "˄" else "Int " -- ˄ (U+02C4)
                Ext -> if u then "˅" else "Ext " -- ˅ (U+02C5)
          prec = case op of
                   Int | u -> app_prec + 1
                   Ext | u -> app_prec + 1
                   _ -> app_prec
      Op2 op a b -> do
        a' <- h (l prec) a
        b' <- h (r prec) b
        return $ showParen (d > prec)
               $ a' . showChar ' ' . showString s . showChar ' ' . b'
        where
          (s,prec,l,r) =
      	    case op of
              And   -> (if u then "∧" else "&&",  4, id, id)     -- ∧ (U+2227)
              Or    -> (if u then "∨" else "||",  3, id, id)     -- ∨ (U+2228)
              Imply -> (if u then "→" else "->",  1, (+1), id)   -- → (U+2192)
              Equiv -> (if u then "↔" else "<->", 1, (+1), (+1)) -- ↔ (U+2194)
              Id    -> ("=",   2, (+1), (+1))

    f :: Int -> Binder -> Expr -> RenderM ShowS
    f d b e = do
      (xs, s) <- go e
      return $ showParen (d > 0) $
        showString (g b) . showString (intercalate " " xs) . showString ". " . s
      where
        go :: Expr -> RenderM ([Name], ShowS)
        go (Bind b' (Sc body)) | b==b' = do
          x <- gensym 
          (xs, s) <- local (x:) $ go body
          return (x:xs, s)
        go e = do
          s <- h 0 e
          return ([], s)

    g :: Binder -> String
    g Lambda = if u then "λ" else "\\" -- λ (U+03BB)
    g Forall = if u then "∀" else "forall " -- ∀ (U+2200)
    g Exists = if u then "∃" else "exists " -- ∃ (U+2203)


renderExpr :: Int -> Expr -> RenderM ShowS
renderExpr = render False

renderExprU :: Int -> Expr -> RenderM ShowS
renderExprU = render True

app_prec :: Int
app_prec = 10

gensym :: RenderM Name
gensym = RenderM $ do 
    i <- get
    put (i+1)
    return ("x"++show i)