content: {-|
  Module      : Idris.Core.WHNF
  Description : Reduction to Weak Head Normal Form
  
  License     : BSD3
  Maintainer  : The Idris Community.
  -}
  
  {-# LANGUAGE FlexibleContexts, PatternGuards #-}
  {-# OPTIONS_GHC -fwarn-incomplete-patterns #-}
  module Idris.Core.WHNF(whnf, whnfArgs, WEnv) where
  
  import Idris.Core.CaseTree
  import Idris.Core.Evaluate hiding (quote)
  import Idris.Core.TT
  
  -- | A stack entry consists of a term and the environment it is to be
  -- evaluated in (i.e. it's a thunk)
  type StackEntry = (Term, WEnv)
  
  data WEnv = WEnv Int -- number of free variables
                   [(Term, WEnv)]
    deriving Show
  
  type Stack = [StackEntry]
  
  -- | A WHNF is a top level term evaluated in the empty environment. It is
  -- always headed by either an irreducible expression, i.e. a constructor,
  -- a lambda, a constant, or a postulate
  --
  -- Every 'Term' or 'Type' in this structure is associated with the
  -- environment it was encountered in, so that when we quote back to a term
  -- we get the substitutions right.
  
  data WHNF = WDCon Int Int Bool Name (Term, WEnv) -- ^ data constructor
            | WTCon Int Int Name (Type, WEnv) -- ^ type constructor
            | WPRef Name (Term, WEnv) -- ^irreducible global (e.g. a postulate)
            | WV Int
            | WBind Name (Binder (Term, WEnv)) (Term, WEnv)
            | WApp WHNF (Term, WEnv)
            | WConstant Const
            | WProj WHNF Int
            | WType UExp
            | WUType Universe
            | WErased
            | WImpossible
  
  -- NOTE: These aren't yet ready to be used in practice - there's still a
  -- bug or two in the handling of de Bruijn indices.
  
  -- | Reduce a term to weak head normal form.
  whnf :: Context -> Env -> Term -> Term
  -- whnf ctxt env tm = let res = whnf' ctxt env tm in
  --                        trace (show tm ++ "\n==>\n" ++ show res ++ "\n") res
  whnf ctxt env tm =
     inlineSmall ctxt env $ -- reduce small things in body. This is primarily
                            -- to get rid of any noisy "assert_smaller/assert_total"
                            -- and evaluate any simple operators, which makes things
                            -- easier to read.
       whnf' ctxt env tm
  whnf' ctxt env tm =
       quote (do_whnf ctxt (map finalEntry env) (finalise tm))
  
  -- | Reduce a type so that all arguments are expanded
  whnfArgs :: Context -> Env -> Term -> Term
  whnfArgs ctxt env tm = inlineSmall ctxt env $ finalise (whnfArgs' ctxt env tm)
  whnfArgs' ctxt env tm
      = case whnf' ctxt env tm of
             -- The assumption is that de Bruijn indices refer to local things
             -- (so not in the external environment) so we need to instantiate
             -- the name
             Bind n b@(Pi rig _ ty _) sc ->
                      Bind n b (whnfArgs' ctxt ((n, rig, b):env)
                               (subst n (P Bound n ty) sc))
             res -> tm
  
  finalEntry :: (Name, RigCount, Binder (TT Name)) -> (Name, RigCount, Binder (TT Name))
  finalEntry (n, r, b) = (n, r, fmap finalise b)
  
  do_whnf :: Context -> Env -> Term -> WHNF
  do_whnf ctxt genv tm = eval (WEnv 0 []) [] tm
    where
      eval :: WEnv -> Stack -> Term -> WHNF
      eval wenv@(WEnv d env) stk (V i)
           | i < length env = let (tm, env') = env !! i in
                                  eval env' stk tm
           | otherwise = WV i
      eval wenv@(WEnv d env) stk (Bind n (Let c t v) sc)
           = eval (WEnv d ((v, wenv) : env)) stk sc
      eval (WEnv d env) ((tm, tenv) : stk) (Bind n b@(Lam _ _) sc)
           = eval (WEnv d ((tm, tenv) : env)) stk sc
      eval wenv@(WEnv d env) stk (Bind n b sc) -- stk must be empty if well typed
           =let n' = uniqueName n (map fstEnv genv) in
                WBind n' (fmap (\t -> (t, wenv)) b) (sc, WEnv (d + 1) env)
  
      eval env stk (P nt n ty)
           | Just (Let c t v) <- lookupBinder n genv = eval env stk v
      eval env stk (P nt n ty) = apply env nt n ty stk
      eval env stk (App _ f a) = eval env ((a, env) : stk) f
      eval env stk (Constant c) = unload (WConstant c) stk
  
      -- Should never happen in compile time code (for now...)
      eval env stk (Proj tm i) = unload (WProj (eval env [] tm) i) stk
  
      eval env stk Erased = unload WErased stk
      eval env stk Impossible = unload WImpossible stk
      eval env stk (Inferred tm) = eval env stk tm
      eval env stk (TType u) = unload (WType u) stk
      eval env stk (UType u) = unload (WUType u) stk
  
      apply :: WEnv -> NameType -> Name -> Type -> Stack -> WHNF
      apply env nt n ty stk
            = let wp = case nt of
                            DCon t a u -> WDCon t a u n (ty, env)
                            TCon t a -> WTCon t a n (ty, env)
                            Ref -> WPRef n (ty, env)
                            _ -> WPRef n (ty, env)
                          in
              if not (tcReducible n ctxt)
                 then unload wp stk
                 else case lookupDefAccExact n False ctxt of
                           Just (CaseOp ci _ _ _ _ cd, acc)
                               | acc == Public || acc == Hidden ->
                            let (ns, tree) = cases_compiletime cd in
                                case evalCase env ns tree stk of
                                     Just w -> w
                                     Nothing -> unload wp stk
                           Just (Operator _ i op, acc) ->
                            if i <= length stk
                               then case runOp env op (take i stk) (drop i stk) of
                                    Just v -> v
                                    Nothing -> unload wp stk
                               else unload wp stk
                           _ -> unload wp stk
  
      unload :: WHNF -> Stack -> WHNF
      unload f [] = f
      unload f (a : as) = unload (WApp f a) as
  
      runOp :: WEnv -> ([Value] -> Maybe Value) -> Stack -> Stack -> Maybe WHNF
      runOp env op stk rest
          = do vals <- mapM tmtoValue stk
               case op vals of
                    Just (VConstant c) -> Just $ unload (WConstant c) rest
                    -- Operators run on values, so we have to convert back
                    -- and forth. This is pretty ugly, but operators that
                    -- aren't run on constants are themselves pretty ugly
                    -- (it's prim__believe_me and prim__syntacticEq, for
                    -- example) so let's not worry too much...
                    -- We will need to deal with believe_me before dropping this
                    -- into the type checker, though.
                    Just val -> Just $ eval env rest (quoteTerm val)
                    _ -> Nothing
  
      tmtoValue :: (Term, WEnv) -> Maybe Value
      tmtoValue (tm, tenv)
          = case eval tenv [] tm of
                 WConstant c -> Just (VConstant c)
                 _ -> let tm' = quoteEnv tenv tm in
                          Just (toValue ctxt [] tm')
  
      evalCase :: WEnv -> [Name] -> SC -> Stack -> Maybe WHNF
      evalCase wenv@(WEnv d env) ns tree args
          | length ns > length args = Nothing
          | otherwise = let args' = take (length ns) args
                            rest = drop (length ns) args in
                            do (tm, amap) <- evalTree wenv (zip ns args') tree
                               let wtm = pToVs (map fst amap) tm
                               Just $ eval (WEnv d (map snd amap)) rest wtm
  
      evalTree :: WEnv -> [(Name, (Term, WEnv))] -> SC ->
                  Maybe (Term, [(Name, (Term, WEnv))])
      evalTree env amap (STerm tm) = Just (tm, amap)
      evalTree env amap (Case _ n alts)
          = case lookup n amap of
              Just (tm, tenv) -> findAlt env amap
                                     (deconstruct (eval tenv [] tm) []) alts
              _ -> Nothing
      evalTree _ _ _ = Nothing
  
      deconstruct :: WHNF -> Stack -> (WHNF, Stack)
      deconstruct (WApp f arg) stk = deconstruct f (arg : stk)
      deconstruct t stk = (t, stk)
  
      findAlt :: WEnv -> [(Name, (Term, WEnv))] -> (WHNF, Stack) ->
                 [CaseAlt] ->
                 Maybe (Term, [(Name, (Term, WEnv))])
      findAlt env amap (WDCon tag _ _ _ _, args) alts
          | Just (ns, sc) <- findTag tag alts
                = let amap' = updateAmap (zip ns args) amap in
                      evalTree env amap' sc
          | Just sc <- findDefault alts
                = evalTree env amap sc
      findAlt env amap (WConstant c, []) alts
          | Just sc <- findConst c alts
                = evalTree env amap sc
          | Just sc <- findDefault alts
                = evalTree env amap sc
      findAlt _ _ _ _ = Nothing
  
      findTag :: Int -> [CaseAlt] -> Maybe ([Name], SC)
      findTag i [] = Nothing
      findTag i (ConCase n j ns sc : xs) | i == j = Just (ns, sc)
      findTag i (_ : xs) = findTag i xs
  
      findDefault :: [CaseAlt] -> Maybe SC
      findDefault [] = Nothing
      findDefault (DefaultCase sc : xs) = Just sc
      findDefault (_ : xs) = findDefault xs
  
      findConst c [] = Nothing
      findConst c (ConstCase c' v : xs) | c == c' = Just v
      findConst (AType (ATInt ITNative)) (ConCase n 1 [] v : xs) = Just v
      findConst (AType ATFloat) (ConCase n 2 [] v : xs) = Just v
      findConst (AType (ATInt ITChar))  (ConCase n 3 [] v : xs) = Just v
      findConst StrType (ConCase n 4 [] v : xs) = Just v
      findConst (AType (ATInt ITBig)) (ConCase n 6 [] v : xs) = Just v
      findConst (AType (ATInt (ITFixed ity))) (ConCase n tag [] v : xs)
          | tag == 7 + fromEnum ity = Just v
      findConst c (_ : xs) = findConst c xs
  
      updateAmap newm amap
         = newm ++ filter (\ (x, _) -> not (elem x (map fst newm))) amap
  
  quote :: WHNF -> Term
  quote (WDCon t a u n (ty, env)) = P (DCon t a u) n (quoteEnv env ty)
  quote (WTCon t a n (ty, env)) = P (TCon t a) n (quoteEnv env ty)
  quote (WPRef n (ty, env)) = P Ref n (quoteEnv env ty)
  quote (WV i) = V i
  quote (WBind n b (sc, env)) = Bind n (fmap (\ (t, env) -> quoteEnv env t) b)
                                       (quoteEnv env sc)
  quote (WApp f (a, env)) = App Complete (quote f) (quoteEnv env a)
  quote (WConstant c) = Constant c
  quote (WProj t i) = Proj (quote t) i
  quote (WType u) = TType u
  quote (WUType u) = UType u
  quote WErased = Erased
  quote WImpossible = Impossible
  
  quoteEnv :: WEnv -> Term -> Term
  quoteEnv (WEnv d ws) tm = qe' d ws tm
    where
      -- d is number of free variables in the external environment
      -- (all bound in the scope of 'ws')
      qe' d ts (V i)
          | (i - d) < length ts && (i - d) >= 0
                = let (tm, env) = ts !! (i - d) in
                      quoteEnv env tm
          | otherwise = V i
      qe' d ts (Bind n b sc)
          = Bind n (fmap (qe' d ts) b) (qe' (d + 1) ts sc)
      qe' d ts (App c f a)
          = App c (qe' d ts f) (qe' d ts a)
      qe' d ts (P nt n ty) = P nt n (qe' d ts ty)
      qe' d ts (Proj tm i) = Proj (qe' d ts tm) i
      qe' d ts tm = tm
  