-- |
-- 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 Cryptol.Utils.RecordMap

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 Map Name Schema
tyenv Expr
expr =
  case Expr
expr of
    -- Monomorphic fragment
    EList [Expr]
es 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 [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 RecordMap Ident Expr
fields   -> RecordMap Ident Type -> Type
tRec ((Expr -> Type) -> RecordMap Ident Expr -> RecordMap Ident Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv) RecordMap Ident Expr
fields)
    ESel Expr
e Selector
sel    -> Type -> Selector -> Type
typeSelect (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e) Selector
sel
    ESet Type
ty Expr
_ Selector
_ Expr
_ -> Type
ty
    EIf Expr
_ Expr
e Expr
_     -> Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e
    EComp Type
len Type
t Expr
_ [[Match]]
_ -> Type -> Type -> Type
tSeq Type
len Type
t
    EAbs Name
x Type
t 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 Expr
e Expr
_      -> case Type -> Maybe (Type, Type)
tIsFun (Map Name Schema -> Expr -> Type
fastTypeOf Map Name Schema
tyenv Expr
e) of
                        Just (Type
_, Type
t) -> Type
t
                        Maybe (Type, Type)
Nothing     -> String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastTypeOf"
                                         [ String
"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 [] [] Type
ty -> Type
ty
        Schema
_ -> String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastTypeOf"
               [ String
"unexpected polymorphic type" ]

fastSchemaOf :: Map Name Schema -> Expr -> Schema
fastSchemaOf :: Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
expr =
  case Expr
expr of
    -- Polymorphic fragment
    EVar 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 Schema
ty -> Schema
ty
                         Maybe Schema
Nothing -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ String
"EVar failed to find type variable:", Name -> String
forall a. Show a => a -> String
show Name
x ]
    ETAbs TParam
tparam Expr
e -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall [TParam]
tparams [Type]
props Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall (TParam
tparam TParam -> [TParam] -> [TParam]
forall a. a -> [a] -> [a]
: [TParam]
tparams) [Type]
props Type
ty
    ETApp Expr
e Type
t      -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall (TParam
tparam : [TParam]
tparams) [Type]
props 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 = TParam -> Type -> Subst
singleTParamSubst TParam
tparam Type
t
                        Schema
_ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ String
"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 Type
p Expr
e  -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall [] [Type]
props Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall [] (Type
p Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
props) Type
ty
                        Schema
_ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ String
"EProofAbs with polymorphic expression" ]
    EProofApp Expr
e    -> case Map Name Schema -> Expr -> Schema
fastSchemaOf Map Name Schema
tyenv Expr
e of
                        Forall [] (Type
_ : [Type]
props) Type
ty -> [TParam] -> [Type] -> Type -> Schema
Forall [] [Type]
props Type
ty
                        Schema
_ -> String -> [String] -> Schema
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.fastSchemaOf"
                               [ String
"EProofApp with polymorphic expression or"
                               , String
"no props in scope"
                               ]
    EWhere Expr
e [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 [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 Decl
d) = Decl -> Map Name Schema -> Map Name Schema
addDecl Decl
d
                              addDecl :: Decl -> Map Name Schema -> Map Name Schema
addDecl 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 Subst
s Type
ty =
  case Type
ty of
    TCon TCon
tc [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 Name
f [Type]
ts 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 RecordMap Ident Type
fs      -> RecordMap Ident Type -> Type
TRec ((Type -> Type) -> RecordMap Ident Type -> RecordMap Ident Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Subst -> Type -> Type
plainSubst Subst
s) RecordMap Ident Type
fs)
    TVar 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 Name
_ [Type]
_ Type
ty) Selector
sel = Type -> Selector -> Type
typeSelect Type
ty Selector
sel
typeSelect (Type -> Maybe [Type]
tIsTuple -> Just [Type]
ts) (TupleSel Int
i Maybe Int
_)
  | 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 RecordMap Ident Type
fields) (RecordSel Ident
n Maybe [Ident]
_)
  | Just Type
ty <- Ident -> RecordMap Ident Type -> Maybe Type
forall a b. Ord a => a -> RecordMap a b -> Maybe b
lookupField Ident
n RecordMap Ident Type
fields = Type
ty
typeSelect (Type -> Maybe (Type, Type)
tIsSeq -> Just (Type
_, Type
a)) ListSel{} = Type
a
typeSelect (Type -> Maybe (Type, Type)
tIsSeq -> Just (Type
n, 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 (Type
n, Type
a)) sel :: Selector
sel@RecordSel{} = Type -> Type -> Type
tSeq Type
n (Type -> Selector -> Type
typeSelect Type
a Selector
sel)
typeSelect Type
ty Selector
_ = String -> [String] -> Type
forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.TypeOf.typeSelect"
                    [ String
"cannot apply selector to value of type", Doc -> String
render (Type -> Doc
forall a. PP a => a -> Doc
pp Type
ty) ]