-- |
-- Module      :  Cryptol.TypeCheck.TypeOf
-- Copyright   :  (c) 2014-2016 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable

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

import Cryptol.TypeCheck.AST
import Cryptol.TypeCheck.Subst
import Cryptol.Utils.Panic
import Cryptol.Utils.PP

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

-- | 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 Name Schema -> Expr -> Type
fastTypeOf :: Map Name Schema -> Expr -> Type
fastTypeOf tyenv :: Map Name Schema
tyenv expr :: Expr
expr =
  case Expr
expr of
    -- Monomorphic fragment
    EList es :: [Expr]
es t :: Type
t    -> Type -> Type -> Type
tSeq (Int -> Type
forall a. Integral a => a -> Type
tNum ([Expr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
es)) Type
t
    ETuple es :: [Expr]
es     -> [Type] -> Type
tTuple ((Expr -> Type) -> [Expr] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv) [Expr]
es)
    ERec fields :: [(Ident, Expr)]
fields   -> [(Ident, Type)] -> Type
tRec [ (Ident
name, Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e) | (name :: Ident
name, e :: Expr
e) <- [(Ident, Expr)]
fields ]
    ESel e :: Expr
e sel :: Selector
sel    -> Type -> Selector -> Type
typeSelect (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e) Selector
sel
    ESet e :: Expr
e _ _    -> Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e
    EIf _ e :: Expr
e _     -> Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e
    EComp len :: Type
len t :: Type
t _ _ -> Type -> Type -> Type
tSeq Type
len Type
t
    EAbs x :: Name
x t :: Type
t e :: Expr
e    -> Type -> Type -> Type
tFun Type
t (Map Name Schema -> Expr -> Type
fastTypeOf (Name -> Schema -> Map Name Schema -> Map Name Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name
x ([TParam] -> [Type] -> Type -> Schema
Forall [] [] Type
t) Map Name Schema
tyenv) Expr
e)
    EApp e :: Expr
e _      -> case Type -> Maybe (Type, Type)
tIsFun (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e) of
                        Just (_, t :: Type
t) -> Type
t
                        Nothing     -> String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastTypeOf"
                                         [ "EApp with non-function operator" ]
    -- Polymorphic fragment
    EVar      {}  -> Type
polymorphic
    ETAbs     {}  -> Type
polymorphic
    ETApp     {}  -> Type
polymorphic
    EProofAbs {}  -> Type
polymorphic
    EProofApp {}  -> Type
polymorphic
    EWhere    {}  -> Type
polymorphic
  where
    polymorphic :: Type
polymorphic =
      case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
expr of
        Forall [] [] ty :: Type
ty -> Type
ty
        _ -> String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastTypeOf"
               [ "unexpected polymorphic type" ]

fastSchemaOf :: Map Name Schema -> Expr -> Schema
fastSchemaOf :: Map Name Schema -> Expr -> Schema
fastSchemaOf tyenv :: Map Name Schema
tyenv expr :: Expr
expr =
  case Expr
expr of
    -- Polymorphic fragment
    EVar x :: Name
x         -> case Name -> Map Name Schema -> Maybe Schema
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name
x Map Name Schema
tyenv of
                         Just ty :: Schema
ty -> Schema
ty
                         Nothing -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "EVar failed to find type variable:", Name -> String
forall a. Show a => a -> String
show Name
x ]
    ETAbs tparam :: TParam
tparam e :: Expr
e -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall tparams :: [TParam]
tparams props :: [Type]
props ty :: Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall (TParam
tparam TParam -> [TParam] -> [TParam]
forall a. a -> [a] -> [a]
: [TParam]
tparams) [Type]
props Type
ty
    ETApp e :: Expr
e t :: Type
t      -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall (tparam :: TParam
tparam : tparams :: [TParam]
tparams) props :: [Type]
props ty :: Type
ty
                          -> [TParam] -> [Type] -> Type -> Schema
Forall [TParam]
tparams ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Type -> Type
plainSubst Subst
s) [Type]
props) (Subst -> Type -> Type
plainSubst Subst
s Type
ty)
                          where s :: Subst
s = TVar -> Type -> Subst
singleSubst (TParam -> TVar
tpVar TParam
tparam) Type
t
                        _ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "ETApp body with no type parameters" ]
                        -- When calling 'fastSchemaOf' on a
                        -- polymorphic function with instantiated type
                        -- variables but undischarged type
                        -- constraints, we would prefer to see the
                        -- instantiated constraints in an
                        -- un-simplified form. Thus we use
                        -- 'plainSubst' instead of 'apSubst' on the
                        -- type constraints.
    EProofAbs p :: Type
p e :: Expr
e  -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall [] props :: [Type]
props ty :: Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall [] (Type
p Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
props) Type
ty
                        _ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "EProofAbs with polymorphic expression" ]
    EProofApp e :: Expr
e    -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall [] (_ : props :: [Type]
props) ty :: Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall [] [Type]
props Type
ty
                        _ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ "EProofApp with polymorphic expression or"
                               , "no props in scope"
                               ]
    EWhere e :: Expr
e dgs :: [DeclGroup]
dgs   -> Map Name Schema -> Expr -> Schema
fastSchemaOf ((DeclGroup -> Map Name Schema -> Map Name Schema)
-> Map Name Schema -> [DeclGroup] -> Map Name Schema
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr DeclGroup -> Map Name Schema -> Map Name Schema
addDeclGroup Map Name Schema
tyenv [DeclGroup]
dgs) Expr
e
                        where addDeclGroup :: DeclGroup -> Map Name Schema -> Map Name Schema
addDeclGroup (Recursive ds :: [Decl]
ds) = (Map Name Schema -> [Decl] -> Map Name Schema)
-> [Decl] -> Map Name Schema -> Map Name Schema
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Decl -> Map Name Schema -> Map Name Schema)
-> Map Name Schema -> [Decl] -> Map Name Schema
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Decl -> Map Name Schema -> Map Name Schema
addDecl) [Decl]
ds
                              addDeclGroup (NonRecursive d :: Decl
d) = Decl -> Map Name Schema -> Map Name Schema
addDecl Decl
d
                              addDecl :: Decl -> Map Name Schema -> Map Name Schema
addDecl d :: Decl
d = Name -> Schema -> Map Name Schema -> Map Name Schema
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Decl -> Name
dName Decl
d) (Decl -> Schema
dSignature Decl
d)
    -- Monomorphic fragment
    EList  {}      -> Schema
monomorphic
    ETuple {}      -> Schema
monomorphic
    ERec   {}      -> Schema
monomorphic
    ESet   {}      -> Schema
monomorphic
    ESel   {}      -> Schema
monomorphic
    EIf    {}      -> Schema
monomorphic
    EComp  {}      -> Schema
monomorphic
    EApp   {}      -> Schema
monomorphic
    EAbs   {}      -> Schema
monomorphic
  where
    monomorphic :: Schema
monomorphic = [TParam] -> [Type] -> Type -> Schema
Forall [] [] (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
expr)

-- | Apply a substitution to a type *without* simplifying
-- constraints like @Arith [n]a@ to @Arith a@. (This is in contrast to
-- 'apSubst', which performs simplifications wherever possible.)
plainSubst :: Subst -> Type -> Type
plainSubst :: Subst -> Type -> Type
plainSubst s :: Subst
s ty :: Type
ty =
  case Type
ty of
    TCon tc :: TCon
tc ts :: [Type]
ts   -> TCon -> [Type] -> Type
TCon TCon
tc ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Type -> Type
plainSubst Subst
s) [Type]
ts)
    TUser f :: Name
f ts :: [Type]
ts t :: Type
t -> Name -> [Type] -> Type -> Type
TUser Name
f ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> Type -> Type
plainSubst Subst
s) [Type]
ts) (Subst -> Type -> Type
plainSubst Subst
s Type
t)
    TRec fs :: [(Ident, Type)]
fs      -> [(Ident, Type)] -> Type
TRec [ (Ident
x, Subst -> Type -> Type
plainSubst Subst
s Type
t) | (x :: Ident
x, t :: Type
t) <- [(Ident, Type)]
fs ]
    TVar x :: TVar
x       -> Subst -> Type -> Type
forall t. TVars t => Subst -> t -> t
apSubst Subst
s (TVar -> Type
TVar TVar
x)

-- | Yields the return type of the selector on the given argument type.
typeSelect :: Type -> Selector -> Type
typeSelect :: Type -> Selector -> Type
typeSelect (TUser _ _ ty :: Type
ty) sel :: Selector
sel = Type -> Selector -> Type
typeSelect Type
ty Selector
sel
typeSelect (Type -> Maybe [Type]
tIsTuple -> Just ts :: [Type]
ts) (TupleSel i :: Int
i _)
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts = [Type]
ts [Type] -> Int -> Type
forall a. [a] -> Int -> a
!! Int
i
typeSelect (TRec fields :: [(Ident, Type)]
fields) (RecordSel n :: Ident
n _)
  | Just ty :: Type
ty <- Ident -> [(Ident, Type)] -> Maybe Type
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Ident
n [(Ident, Type)]
fields = Type
ty
typeSelect (Type -> Maybe (Type, Type)
tIsSeq -> Just (_, a :: Type
a)) ListSel{} = Type
a
typeSelect (Type -> Maybe (Type, Type)
tIsSeq -> Just (n :: Type
n, a :: Type
a)) sel :: Selector
sel@TupleSel{} = Type -> Type -> Type
tSeq Type
n (Type -> Selector -> Type
typeSelect Type
a Selector
sel)
typeSelect (Type -> Maybe (Type, Type)
tIsSeq -> Just (n :: Type
n, a :: Type
a)) sel :: Selector
sel@RecordSel{} = Type -> Type -> Type
tSeq Type
n (Type -> Selector -> Type
typeSelect Type
a Selector
sel)
typeSelect ty :: Type
ty _ = String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.TypeOf.typeSelect"
                    [ "cannot apply selector to value of type", Doc -> String
render (Type -> Doc
forall a. PP a => a -> Doc
pp Type
ty) ]