-- |
-- Module      :  $Header$
-- Copyright   :  (c) 2014-2015 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable

{-# LANGUAGE Safe                                #-}
{-# LANGUAGE ViewPatterns                        #-}
module Cryptol.TypeCheck.TypeOf
  ( fastTypeOf
  , fastSchemaOf
  ) where

import Cryptol.TypeCheck.AST
import Cryptol.TypeCheck.Subst
import Cryptol.Prims.Types (typeOf)
import Cryptol.Utils.Panic
import Cryptol.Utils.PP

import           Data.Map    (Map)
import qualified Data.Map as Map
import           Data.Maybe (fromJust)

-- | Given a typing environment and an expression, compute the type of
-- the expression as quickly as possible, assuming that the expression
-- is well formed with correct type annotations.
fastTypeOf :: Map QName Schema -> Expr -> Type
fastTypeOf tyenv expr =
  case expr of
    -- Monomorphic fragment
    EList es t    -> tSeq (tNum (length es)) t
    ETuple es     -> tTuple (map (fastTypeOf tyenv) es)
    ERec fields   -> tRec [ (name, fastTypeOf tyenv e) | (name, e) <- fields ]
    ESel e sel    -> typeSelect (fastTypeOf tyenv e) sel
    EIf _ e _     -> fastTypeOf tyenv e
    EComp t _ _   -> t
    EAbs x t e    -> tFun t (fastTypeOf (Map.insert x (Forall [] [] t) tyenv) e)
    EApp e _      -> case tIsFun (fastTypeOf tyenv e) of
                        Just (_, t) -> t
                        Nothing     -> panic "Cryptol.TypeCheck.TypeOf.fastTypeOf"
                                         [ "EApp with non-function operator" ]
    ECast _ t     -> t
    -- Polymorphic fragment
    ECon      {}  -> polymorphic
    EVar      {}  -> polymorphic
    ETAbs     {}  -> polymorphic
    ETApp     {}  -> polymorphic
    EProofAbs {}  -> polymorphic
    EProofApp {}  -> polymorphic
    EWhere    {}  -> polymorphic
  where
    polymorphic =
      case fastSchemaOf tyenv expr of
        Forall [] [] ty -> ty
        _ -> panic "Cryptol.TypeCheck.TypeOf.fastTypeOf"
               [ "unexpected polymorphic type" ]

fastSchemaOf :: Map QName Schema -> Expr -> Schema
fastSchemaOf tyenv expr =
  case expr of
    -- Polymorphic fragment
    ECon econ      -> typeOf econ
    EVar x         -> fromJust (Map.lookup x tyenv)
    ETAbs tparam e -> case fastSchemaOf tyenv e of
                        Forall tparams props ty -> Forall (tparam : tparams) props ty
    ETApp e t      -> case fastSchemaOf tyenv e of
                        Forall (tparam : tparams) props ty -> Forall tparams (apSubst s props) (apSubst s ty)
                                                                where s = singleSubst (tpVar tparam) t
                        _ -> panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "ETApp body with no type parameters" ]
    EProofAbs p e  -> case fastSchemaOf tyenv e of
                        Forall [] props ty -> Forall [] (p : props) ty
                        _ -> panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "EProofAbs with polymorphic expression" ]
    EProofApp e    -> case fastSchemaOf tyenv e of
                        Forall [] (_ : props) ty -> Forall [] props ty
                        _ -> panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "EProofApp with polymorphic expression or"
                               , "no props in scope"
                               ]
    EWhere e dgs   -> fastSchemaOf (foldr addDeclGroup tyenv dgs) e
                        where addDeclGroup (Recursive ds) = flip (foldr addDecl) ds
                              addDeclGroup (NonRecursive d) = addDecl d
                              addDecl d = Map.insert (dName d) (dSignature d)
    -- Monomorphic fragment
    EList  {}      -> monomorphic
    ETuple {}      -> monomorphic
    ERec   {}      -> monomorphic
    ESel   {}      -> monomorphic
    EIf    {}      -> monomorphic
    EComp  {}      -> monomorphic
    EApp   {}      -> monomorphic
    EAbs   {}      -> monomorphic
    ECast  {}      -> monomorphic
  where
    monomorphic = Forall [] [] (fastTypeOf tyenv expr)

-- | Yields the return type of the selector on the given argument type.
typeSelect :: Type -> Selector -> Type
typeSelect (TCon _tctuple ts) (TupleSel i _) = ts !! i
typeSelect (TRec fields) (RecordSel n _) = fromJust (lookup n fields)
typeSelect (TCon _tcseq [_, a]) (ListSel _ _) = a
typeSelect ty _ = panic "Cryptol.TypeCheck.TypeOf.typeSelect"
                    [ "cannot apply selector to value of type", render (pp ty) ]