    -- |
    -- 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                       #-}
    {-# LANGUAGE NamedFieldPuns                      #-}
    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
        ELocated Range
    _ Expr
    t  -> Map Name Schema -> Expr -> Type
    fastTypeOf Map Name Schema
    tyenv Expr
    t
        EList [Expr]
    es Type
    t    -> Type -> Type -> Type
    tSeq (Int -> Type
    forall a. Integral a => a -> Type
    tNum ([Expr] -> Int
    forall a. [a] -> 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 a b. (a -> b) -> RecordMap Ident a -> RecordMap Ident b
    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" ]
        EPropGuards [([Type], Expr)]
    _guards Type
    sType -> Type
    sType
        -- 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
            s :: Schema
    s@Forall {} -> String -> [String] -> Type
    forall a. HasCallStack => String -> [String] -> a
    panic String
    "Cryptol.TypeCheck.TypeOf.fastTypeOf"
                   [ String
    "unexpected polymorphic type in expression:"
                   , Expr -> String
    forall a. PP a => a -> String
    pretty Expr
    expr
                   , String
    "with schema:"
                   , Schema -> String
    forall a. PP a => a -> String
    pretty Schema
    s
                   ]

    fastSchemaOf :: Map Name Schema -> Expr -> Schema
    fastSchemaOf :: Map Name Schema -> Expr -> Schema
    fastSchemaOf Map Name Schema
    tyenv Expr
    expr =
      case Expr
    expr of
        ELocated Range
    _ Expr
    e -> Map Name Schema -> Expr -> Schema
    fastSchemaOf Map Name Schema
    tyenv Expr
    e

        -- 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 a b. (a -> b -> b) -> b -> [a] -> b
    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 a b. (a -> b -> b) -> b -> [a] -> b
    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

        -- PropGuards
        EPropGuards [([Type], Expr)]
    _ Type
    t -> Forall {sVars :: [TParam]
    sVars = [], sProps :: [Type]
    sProps = [], sType :: Type
    sType = Type
    t}
      where
        monomorphic :: Schema
    monomorphic = Forall {sVars :: [TParam]
    sVars = [], sProps :: [Type]
    sProps = [], sType :: Type
    sType = 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 a b. (a -> b) -> RecordMap Ident a -> RecordMap Ident b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Subst -> Type -> Type
    plainSubst Subst
    s) RecordMap Ident Type
    fs)
        TNewtype Newtype
    nt [Type]
    ts -> Newtype -> [Type] -> Type
    TNewtype Newtype
    nt ((Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Subst -> Type -> Type
    plainSubst Subst
    s) [Type]
    ts)
        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

    -- Selectors push inside the definition of type aliases
    typeSelect :: Type -> Selector -> Type
    typeSelect (TUser Name
    _ [Type]
    _ Type
    ty) Selector
    sel = Type -> Selector -> Type
    typeSelect Type
    ty Selector
    sel

    -- Tuple selector applied to a tuple
    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 a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Type]
    ts = [Type]
    ts [Type] -> Int -> Type
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    i

    -- Record selector applied to a record
    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

    -- Record selector applied to a newtype
    typeSelect (TNewtype Newtype
    nt [Type]
    args) (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 (Newtype -> RecordMap Ident Type
    ntFields Newtype
    nt)
      = Subst -> Type -> Type
    plainSubst ([(TParam, Type)] -> Subst
    listParamSubst ([TParam] -> [Type] -> [(TParam, Type)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Newtype -> [TParam]
    ntParams Newtype
    nt) [Type]
    args)) Type
    ty

    -- List selector applied to a sequence
    typeSelect (Type -> Maybe (Type, Type)
    tIsSeq -> Just (Type
    _, Type
    a)) ListSel{} = Type
    a

    -- Tuple selectors and record selectors lift pointwise over sequences
    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)

    -- Selectors lift pointwise over functions
    typeSelect (Type -> Maybe (Type, Type)
    tIsFun -> Just (Type
    a, Type
    b)) Selector
    sel = Type -> Type -> Type
    tFun Type
    a (Type -> Selector -> Type
    typeSelect Type
    b 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
    forall a. Show a => a -> String
    show (Type -> Doc
    forall a. PP a => a -> Doc
    pp Type
    ty) ]