content: module PGF.Expr(Tree, BindType(..), Expr(..), Literal(..), Patt(..), Equation(..),
                  readExpr, showExpr, pExpr, pBinds, ppExpr, ppPatt, pattScope,
  
                  mkAbs,    unAbs,
                  mkApp,    unApp, unapply,
                  mkStr,    unStr,
                  mkInt,    unInt,
                  mkDouble, unDouble,
                  mkFloat,  unFloat,
                  mkMeta,   unMeta,
  
                  normalForm,
  
                  -- needed in the typechecker
                  Value(..), Env, Sig, eval, apply, applyValue, value2expr,
  
                  MetaId,
  
                  -- helpers
                  pMeta,pArg,pLit,freshName,ppMeta,ppLit,ppParens,
                  freshBoundVars
                 ) where
  
  import PGF.CId
  import PGF.Type
  import PGF.ByteCode
  
  import Data.Char
  --import Data.Maybe
  import Data.List as List
  import qualified Data.Map as Map hiding (showTree)
  import Control.Monad
  import qualified Text.PrettyPrint as PP
  import qualified Text.ParserCombinators.ReadP as RP
  
  type MetaId = Int
  
  data BindType = 
      Explicit
    | Implicit
    deriving (Eq,Ord,Show)
  
  -- | Tree is the abstract syntax representation of a given sentence
  -- in some concrete syntax. Technically 'Tree' is a type synonym
  -- of 'Expr'.
  type Tree = Expr
  
  -- | An expression in the abstract syntax of the grammar. It could be
  -- both parameter of a dependent type or an abstract syntax tree for
  -- for some sentence.
  data Expr =
     EAbs BindType CId Expr           -- ^ lambda abstraction
   | EApp Expr Expr                   -- ^ application
   | ELit Literal                     -- ^ literal
   | EMeta  {-# UNPACK #-} !MetaId    -- ^ meta variable
   | EFun   CId                       -- ^ function or data constructor
   | EVar   {-# UNPACK #-} !Int       -- ^ variable with de Bruijn index
   | ETyped Expr Type                 -- ^ local type signature
   | EImplArg Expr                    -- ^ implicit argument in expression
    deriving (Eq,Ord,Show)
  
  -- | The pattern is used to define equations in the abstract syntax of the grammar.
  data Patt =
     PApp CId [Patt]                  -- ^ application. The identifier should be constructor i.e. defined with 'data'
   | PLit Literal                     -- ^ literal
   | PVar CId                         -- ^ variable
   | PAs  CId Patt                    -- ^ variable@pattern
   | PWild                            -- ^ wildcard
   | PImplArg Patt                    -- ^ implicit argument in pattern
   | PTilde Expr
    deriving Show
  
  -- | The equation is used to define lambda function as a sequence
  -- of equations with pattern matching. The list of 'Expr' represents
  -- the patterns and the second 'Expr' is the function body for this
  -- equation.
  data Equation =
     Equ [Patt] Expr
    deriving Show
  
  -- | parses 'String' as an expression
  readExpr :: String -> Maybe Expr
  readExpr s = case [x | (x,cs) <- RP.readP_to_S pExpr s, all isSpace cs] of
                 [x] -> Just x
                 _   -> Nothing
  
  -- | renders expression as 'String'. The list
  -- of identifiers is the list of all free variables
  -- in the expression in order reverse to the order
  -- of binding.
  showExpr :: [CId] -> Expr -> String
  showExpr vars = PP.render . ppExpr 0 vars
  
  instance Read Expr where
      readsPrec _ = RP.readP_to_S pExpr
  
  mkAbs :: BindType -> CId -> Expr -> Expr
  mkAbs = EAbs
  
  unAbs :: Expr -> Maybe (BindType, CId, Expr)
  unAbs (EAbs bt x e) = Just (bt,x,e)
  unAbs (ETyped e ty) = unAbs e
  unAbs (EImplArg e)  = unAbs e
  unAbs _             = Nothing
  
  -- | Constructs an expression by applying a function to a list of expressions
  mkApp :: CId -> [Expr] -> Expr
  mkApp f es = foldl EApp (EFun f) es
  
  -- | Decomposes an expression into application of function
  unApp :: Expr -> Maybe (CId,[Expr])
  unApp e = case unapply e of
    (EFun f,es) -> Just (f,es)
    _           -> Nothing
  
  -- | Decomposes an expression into an application of a constructor such as a constant or a metavariable
  unapply :: Expr -> (Expr,[Expr])
  unapply = extract []
    where
      extract es f@(EFun _)   = (f,es)
      extract es (EApp e1 e2) = extract (e2:es) e1
      extract es (ETyped e ty)= extract es e
      extract es (EImplArg e) = extract es e
      extract es h            = (h,es)
  
  -- | Constructs an expression from string literal
  mkStr :: String -> Expr
  mkStr s = ELit (LStr s)
  
  -- | Decomposes an expression into string literal
  unStr :: Expr -> Maybe String
  unStr (ELit (LStr s)) = Just s
  unStr (ETyped e ty)   = unStr e
  unStr (EImplArg e)    = unStr e
  unStr _               = Nothing
  
  -- | Constructs an expression from integer literal
  mkInt :: Int -> Expr
  mkInt i = ELit (LInt i)
  
  -- | Decomposes an expression into integer literal
  unInt :: Expr -> Maybe Int
  unInt (ELit (LInt i)) = Just i
  unInt (ETyped e ty)   = unInt e
  unInt (EImplArg e)    = unInt e
  unInt _               = Nothing
  
  -- | Constructs an expression from real number literal
  mkDouble :: Double -> Expr
  mkDouble f = ELit (LFlt f)
  
  -- | Decomposes an expression into real number literal
  unDouble :: Expr -> Maybe Double
  unDouble (ELit (LFlt f)) = Just f
  unDouble (ETyped e ty)   = unDouble e
  unDouble (EImplArg e)    = unDouble e
  unDouble _               = Nothing
  
  mkFloat = mkDouble
  unFloat = unDouble
  
  -- | Constructs an expression which is meta variable
  mkMeta :: Int -> Expr
  mkMeta i = EMeta i
  
  -- | Checks whether an expression is a meta variable
  unMeta :: Expr -> Maybe Int
  unMeta (EMeta i)     = Just i
  unMeta (ETyped e ty) = unMeta e
  unMeta (EImplArg e)  = unMeta e
  unMeta _             = Nothing
  
  -----------------------------------------------------
  -- Parsing
  -----------------------------------------------------
  
  pExpr :: RP.ReadP Expr
  pExpr = RP.skipSpaces >> (pAbs RP.<++ pTerm)
    where
      pTerm = do f <- pFactor
                 RP.skipSpaces
                 as <- RP.sepBy pArg RP.skipSpaces
                 return (foldl EApp f as)
  
      pAbs = do xs <- RP.between (RP.char '\\') (RP.skipSpaces >> RP.string "->") pBinds
                e  <- pExpr
                return (foldr (\(b,x) e -> EAbs b x e) e xs)
  
  pBinds :: RP.ReadP [(BindType,CId)]
  pBinds = do xss <- RP.sepBy1 (RP.skipSpaces >> pBind) (RP.skipSpaces >> RP.char ',')
              return (concat xss)
    where
      pCIdOrWild = pCId `mplus` (RP.char '_' >> return wildCId)
  
      pBind = 
        do x <- pCIdOrWild
           return [(Explicit,x)]
        `mplus`
        RP.between (RP.char '{')
                   (RP.skipSpaces >> RP.char '}')
                   (RP.sepBy1 (RP.skipSpaces >> pCIdOrWild >>= \id -> return (Implicit,id)) (RP.skipSpaces >> RP.char ','))
  
  pArg = fmap EImplArg (RP.between (RP.char '{') (RP.char '}') pExpr)
         RP.<++
         pFactor
  
  pFactor =       fmap EFun  pCId
          RP.<++  fmap ELit  pLit
          RP.<++  fmap EMeta pMeta
          RP.<++  RP.between (RP.char '(') (RP.skipSpaces >> RP.char ')') pExpr
          RP.<++  RP.between (RP.char '<') (RP.skipSpaces >> RP.char '>') pTyped
  
  pTyped = do RP.skipSpaces
              e <- pExpr
              RP.skipSpaces
              RP.char ':'
              RP.skipSpaces
              ty <- pType
              return (ETyped e ty)
  
  pMeta = do RP.char '?'
             ds <- RP.munch isDigit
             return (read ('0':ds))
  
  pLit :: RP.ReadP Literal
  pLit = liftM LStr (RP.readS_to_P reads)
         RP.<++
         liftM LInt (RP.readS_to_P reads)
         RP.<++
         liftM LFlt (RP.readS_to_P reads)
  
  
  -----------------------------------------------------
  -- Printing
  -----------------------------------------------------
  
  ppExpr :: Int -> [CId] -> Expr -> PP.Doc
  ppExpr d scope (EAbs b x e) = let (bs,xs,e1) = getVars [] [] (EAbs b x e)
                                    xs' = freshBoundVars scope xs
                                in ppParens (d > 1) (PP.char '\\' PP.<>
                                                     PP.hsep (PP.punctuate PP.comma (reverse (List.zipWith ppBind bs xs'))) PP.<+>
                                                     PP.text "->" PP.<+>
                                                     ppExpr 1 (xs' ++ scope) e1)
                                where
                                  getVars bs xs (EAbs b x e) = getVars (b:bs) ((freshName x xs):xs) e
                                  getVars bs xs e            = (bs,xs,e)
  ppExpr d scope (EApp e1 e2) = ppParens (d > 3) ((ppExpr 3 scope e1) PP.<+> (ppExpr 4 scope e2))
  ppExpr d scope (ELit l)     = ppLit l
  ppExpr d scope (EMeta n)    = ppMeta n
  ppExpr d scope (EFun f)     = ppCId f
  ppExpr d scope (EVar i)     = ppCId (scope !! i)
  ppExpr d scope (ETyped e ty)= PP.char '<' PP.<> ppExpr 0 scope e PP.<+> PP.colon PP.<+> ppType 0 scope ty PP.<> PP.char '>'
  ppExpr d scope (EImplArg e) = PP.braces (ppExpr 0 scope e)
  
  ppPatt :: Int -> [CId] -> Patt -> PP.Doc
  ppPatt d scope (PApp f ps)  = let ds = List.map (ppPatt 2 scope) ps
                                in ppParens (not (List.null ps) && d > 1) (ppCId f PP.<+> PP.hsep ds)
  ppPatt d scope (PLit l)     = ppLit l
  ppPatt d scope (PVar f)     = ppCId f
  ppPatt d scope (PAs x p)    = ppCId x PP.<> PP.char '@' PP.<> ppPatt 3 scope p
  ppPatt d scope PWild        = PP.char '_'
  ppPatt d scope (PImplArg p) = PP.braces (ppPatt 0 scope p)
  ppPatt d scope (PTilde e)   = PP.char '~' PP.<> ppExpr 6 scope e
  
  pattScope :: [CId] -> Patt -> [CId]
  pattScope scope (PApp f ps)  = foldl pattScope scope ps
  pattScope scope (PLit l)     = scope
  pattScope scope (PVar f)     = f:scope
  pattScope scope (PAs x p)    = pattScope (x:scope) p
  pattScope scope PWild        = scope
  pattScope scope (PImplArg p) = pattScope scope p                               
  pattScope scope (PTilde e)   = scope
  
  ppBind Explicit x = ppCId x
  ppBind Implicit x = PP.braces (ppCId x)
  
  ppMeta :: MetaId -> PP.Doc
  ppMeta n
    | n == 0    = PP.char '?'
    | otherwise = PP.char '?' PP.<> PP.int n
  
  ppParens True  = PP.parens
  ppParens False = id
  
  freshName :: CId -> [CId] -> CId
  freshName x xs0 = loop 1 x
    where
      xs = wildCId : xs0
  
      loop i y
        | elem y xs = loop (i+1) (mkCId (show x++show i))
        | otherwise = y
  
  -- refresh new vars xs in scope if needed. AR 2024-03-01
  freshBoundVars :: [CId] -> [CId] -> [CId]
  freshBoundVars scope xs = foldr fresh [] xs
    where
      fresh x xs' = mkCId (freshName (showCId x) xs') : xs'
      freshName s xs' = 
        if elem (mkCId s) (xs' ++ scope)
        then freshName (s ++ "'") xs'
        else s
  
  -----------------------------------------------------
  -- Computation
  -----------------------------------------------------
  
  -- | Compute an expression to normal form
  normalForm :: Sig -> Int -> Env -> Expr -> Expr
  normalForm sig k env e = value2expr sig k (eval sig env e)
  
  value2expr sig i (VApp f vs)                 = foldl EApp (EFun f)       (List.map (value2expr sig i) vs)
  value2expr sig i (VGen j vs)                 = foldl EApp (EVar (i-j-1)) (List.map (value2expr sig i) vs)
  value2expr sig i (VMeta j env vs)            = case snd sig j of
                                                   Just e  -> value2expr sig i (apply sig env e vs)
                                                   Nothing -> foldl EApp (EMeta j) (List.map (value2expr sig i) vs)
  value2expr sig i (VSusp j env vs k)          = value2expr sig i (k (VGen j vs))
  value2expr sig i (VConst f vs)               = foldl EApp (EFun f)       (List.map (value2expr sig i) vs)
  value2expr sig i (VLit l)                    = ELit l
  value2expr sig i (VClosure env (EAbs b x e)) = EAbs b (mkCId ('v':show i)) (value2expr sig (i+1) (eval sig ((VGen i []):env) e))
  value2expr sig i (VImplArg v)                = EImplArg (value2expr sig i v)
  
  data Value
    = VApp CId [Value]
    | VLit Literal
    | VMeta {-# UNPACK #-} !MetaId Env [Value]
    | VSusp {-# UNPACK #-} !MetaId Env [Value] (Value -> Value)
    | VGen  {-# UNPACK #-} !Int [Value]
    | VConst CId [Value]
    | VClosure Env Expr
    | VImplArg Value
  
  type Sig = ( Map.Map CId (Type,Int,Maybe ([Equation],[[Instr]]),Double) -- type and def of a fun
             , Int -> Maybe Expr                                          -- lookup for metavariables
             )
  type Env = [Value]
  
  eval :: Sig -> Env -> Expr -> Value
  eval sig env (EVar i)     = env !! i
  eval sig env (EFun f)     = case Map.lookup f (fst sig) of
                                Just (_,a,meqs,_) -> case meqs of
                                                         Just (eqs,_)
                                                                 -> if a == 0
                                                                      then case eqs of
                                                                             Equ [] e : _ -> eval sig [] e
                                                                             _            -> VConst f []
                                                                      else VApp f []
                                                         Nothing -> VApp f []
                                Nothing             -> error ("unknown function "++showCId f)
  eval sig env (EApp e1 e2) = apply sig env e1 [eval sig env e2]
  eval sig env (EAbs b x e) = VClosure env (EAbs b x e)
  eval sig env (EMeta i)    = case snd sig i of
                                Just e  -> eval sig env e
                                Nothing -> VMeta i env []
  eval sig env (ELit l)     = VLit l
  eval sig env (ETyped e _) = eval sig env e
  eval sig env (EImplArg e) = VImplArg (eval sig env e)
  
  apply :: Sig -> Env -> Expr -> [Value] -> Value
  apply sig env e            []     = eval sig env e
  apply sig env (EVar i)     vs     = applyValue sig (env !! i) vs
  apply sig env (EFun f)     vs     = case Map.lookup f (fst sig) of
                                        Just (_,a,meqs,_) -> case meqs of
                                                               Just (eqs,_) -> if a <= length vs
                                                                                 then match sig f eqs vs
                                                                                 else VApp f vs
                                                               Nothing      -> VApp f vs
                                        Nothing           -> error ("unknown function "++showCId f)
  apply sig env (EApp e1 e2) vs     = apply sig env e1 (eval sig env e2 : vs)
  apply sig env (EAbs b x e) (v:vs) = case (b,v) of
                                        (Implicit,VImplArg v) -> apply sig (v:env) e vs
                                        (Explicit,         v) -> apply sig (v:env) e vs
  apply sig env (EMeta i)    vs     = case snd sig i of
                                        Just e  -> apply sig env e vs
                                        Nothing -> VMeta i env vs
  apply sig env (ELit l)     vs     = error "literal of function type"
  apply sig env (ETyped e _) vs     = apply sig env e vs
  apply sig env (EImplArg _) vs     = error "implicit argument in function position"
  
  applyValue sig v                         []       = v
  applyValue sig (VApp f  vs0)             vs       = apply sig [] (EFun f) (vs0++vs)
  applyValue sig (VLit _)                  vs       = error "literal of function type"
  applyValue sig (VMeta i env vs0)         vs       = VMeta i env (vs0++vs)
  applyValue sig (VGen  i vs0)             vs       = VGen  i (vs0++vs)
  applyValue sig (VSusp i env vs0 k)       vs       = VSusp i env vs0 (\v -> applyValue sig (k v) vs)
  applyValue sig (VConst f vs0)            vs       = VConst f (vs0++vs)
  applyValue sig (VClosure env (EAbs b x e)) (v:vs) = case (b,v) of
                                                        (Implicit,VImplArg v) -> apply sig (v:env) e vs
                                                        (Explicit,         v) -> apply sig (v:env) e vs
  applyValue sig (VImplArg _)              vs       = error "implicit argument in function position"
  
  -----------------------------------------------------
  -- Pattern matching
  -----------------------------------------------------
  
  match :: Sig -> CId -> [Equation] -> [Value] -> Value
  match sig f eqs as0 =
    case eqs of
      []               -> VConst f as0
      (Equ ps res):eqs -> tryMatches eqs ps as0 res []
    where
      tryMatches eqs []     as     res env = apply sig env res as
      tryMatches eqs (p:ps) (a:as) res env = tryMatch p a env
        where
          tryMatch (PVar x     ) (v                ) env            = tryMatches eqs  ps        as  res (v:env)
          tryMatch (PAs  x p   ) (v                ) env            = tryMatch p v (v:env)
          tryMatch (PWild      ) (_                ) env            = tryMatches eqs  ps        as  res env
          tryMatch (p          ) (VMeta i envi vs  ) env            = VSusp i envi vs (\v -> tryMatch p v env)
          tryMatch (p          ) (VGen  i vs       ) env            = VConst f as0
          tryMatch (p          ) (VSusp i envi vs k) env            = VSusp i envi vs (\v -> tryMatch p (k v) env)
          tryMatch (p          ) v@(VConst _ _     ) env            = match sig f eqs as0
          tryMatch (PApp f1 ps1) (VApp f2 vs2      ) env | f1 == f2 = tryMatches eqs (ps1++ps) (vs2++as) res env
          tryMatch (PLit l1    ) (VLit l2          ) env | l1 == l2 = tryMatches eqs  ps        as  res env
          tryMatch (PImplArg p ) (VImplArg v       ) env            = tryMatch p v env
          tryMatch (PTilde _   ) (_                ) env            = tryMatches eqs  ps        as  res env
          tryMatch _             _                   env            = match sig f eqs as0
  
  