{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Elaborate types
--
-- During the elaboration phase we annotate the AST with the types
-- we know. When we don't know what the type of something is, we generate
-- a new type variable as a placeholder and mark that type with a @Constraint@
-- depending on its syntactic use.
--
-- - Input: Builtin data types, builtin type environment, and the AST
-- - Output: Annotated AST with types (to the best of our abilities) + constraints on those types
module Language.Giml.Types.Infer.Elaborate where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Generics.Uniplate.Data qualified as U
import Data.Map qualified as M
import Data.Set qualified as S
import Language.Giml.Builtins
import Language.Giml.Logging
import Language.Giml.Syntax.Ast
import Language.Giml.Syntax.Parser qualified as Parser
import Language.Giml.Types.Infer.Types
import Language.Giml.Utils

-- * Elaborate data types

-- ** Types

-- | The monadic capabilities for the elaboration phase.
type ElaborateData b m =
  ( MonadBase b m
  , MonadState VariantEnv m
  , MonadError TypeErrorA m
  )

-- | The state we keep during elaboration.
data ElabState = ElabState
  { ElabState -> Int
esTypeVarCounter :: Int
  -- ^ Used to generate unique type variable names.
  , ElabState -> Constraints
esConstraints :: Constraints
  -- ^ The constraints we collect.
  }

-- | The environment we use.
data ElabEnv b = ElabEnv
  { forall (b :: * -> *). ElabEnv b -> Env Type
eeTypeEnv :: Env Type
  -- ^ Represents the variables in scope and their types.
  , forall (b :: * -> *). ElabEnv b -> Env Type
eeBuiltins :: Env Type
  , forall (b :: * -> *). ElabEnv b -> VariantEnv
eeVariantEnv :: VariantEnv
  -- ^ Mapping from a data constructor name to the data type
  --   and variant type signature.
  , forall (b :: * -> *). ElabEnv b -> LogAction b LogMsg
logAction :: LogAction b LogMsg
  }
  deriving ((forall x. ElabEnv b -> Rep (ElabEnv b) x)
-> (forall x. Rep (ElabEnv b) x -> ElabEnv b)
-> Generic (ElabEnv b)
forall x. Rep (ElabEnv b) x -> ElabEnv b
forall x. ElabEnv b -> Rep (ElabEnv b) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (b :: * -> *) x. Rep (ElabEnv b) x -> ElabEnv b
forall (b :: * -> *) x. ElabEnv b -> Rep (ElabEnv b) x
$cto :: forall (b :: * -> *) x. Rep (ElabEnv b) x -> ElabEnv b
$cfrom :: forall (b :: * -> *) x. ElabEnv b -> Rep (ElabEnv b) x
Generic)

-- ** Algorithm

elaborateEnv
  :: (MonadBase b b)
  => [Datatype ()]
  -> [Datatype InputAnn]
  -> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv)
elaborateEnv :: forall (b :: * -> *).
MonadBase b b =>
[Datatype ()]
-> [Datatype InputAnn]
-> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv)
elaborateEnv [Datatype ()]
builtinDats [Datatype InputAnn]
typeDefs =
  (StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
 -> VariantEnv -> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv))
-> VariantEnv
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
-> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
-> VariantEnv -> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT VariantEnv
forall a. Monoid a => a
mempty (StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
 -> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv))
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
-> ExceptT TypeErrorA b ([Datatype Ann], VariantEnv)
forall a b. (a -> b) -> a -> b
$ do
    [Datatype Ann]
_ <-
      (Datatype InputAnn
 -> StateT VariantEnv (ExceptT TypeErrorA b) (Datatype Ann))
-> [Datatype InputAnn]
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Datatype InputAnn
-> StateT VariantEnv (ExceptT TypeErrorA b) (Datatype Ann)
forall (b :: * -> *) (m :: * -> *).
ElaborateData b m =>
Datatype InputAnn -> m (Datatype Ann)
elaborateTypeDef ([Datatype InputAnn]
 -> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann])
-> [Datatype InputAnn]
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
forall a b. (a -> b) -> a -> b
$
        (Datatype () -> Datatype InputAnn)
-> [Datatype ()] -> [Datatype InputAnn]
forall a b. (a -> b) -> [a] -> [b]
map (InputAnn -> Datatype () -> Datatype InputAnn
forall b a. b -> Datatype a -> Datatype b
setTypeAnn (InputAnn -> Datatype () -> Datatype InputAnn)
-> InputAnn -> Datatype () -> Datatype InputAnn
forall a b. (a -> b) -> a -> b
$ FilePath -> InputAnn
Parser.dummyAnn FilePath
"builtin") [Datatype ()]
builtinDats
    -- We invent types for top level term definitions.
    [Datatype Ann]
typeDefs' <- (Datatype InputAnn
 -> StateT VariantEnv (ExceptT TypeErrorA b) (Datatype Ann))
-> [Datatype InputAnn]
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Datatype InputAnn
-> StateT VariantEnv (ExceptT TypeErrorA b) (Datatype Ann)
forall (b :: * -> *) (m :: * -> *).
ElaborateData b m =>
Datatype InputAnn -> m (Datatype Ann)
elaborateTypeDef [Datatype InputAnn]
typeDefs
    [Datatype Ann]
-> StateT VariantEnv (ExceptT TypeErrorA b) [Datatype Ann]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Datatype Ann]
typeDefs'

-- | Add a datatype information into the environment
addVariantSigs :: (ElaborateData b m) => VariantEnv -> m ()
addVariantSigs :: forall (b :: * -> *) (m :: * -> *).
ElaborateData b m =>
VariantEnv -> m ()
addVariantSigs VariantEnv
venv = do
  VariantEnv
env <- m VariantEnv
forall s (m :: * -> *). MonadState s m => m s
get
  let
    duplicates :: [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
duplicates = Map Constr (VariantSig InputAnn, VariantSig InputAnn)
-> [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
forall k a. Map k a -> [(k, a)]
M.toList (Map Constr (VariantSig InputAnn, VariantSig InputAnn)
 -> [(Constr, (VariantSig InputAnn, VariantSig InputAnn))])
-> Map Constr (VariantSig InputAnn, VariantSig InputAnn)
-> [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
forall a b. (a -> b) -> a -> b
$ (VariantSig InputAnn
 -> VariantSig InputAnn
 -> (VariantSig InputAnn, VariantSig InputAnn))
-> VariantEnv
-> VariantEnv
-> Map Constr (VariantSig InputAnn, VariantSig InputAnn)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) VariantEnv
venv VariantEnv
env

  -- check for duplicate variants
  case [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
duplicates of
    [] -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    (Constr
d, (VariantSig InputAnn, VariantSig InputAnn)
v) : [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
rest -> do
      [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [VariantSig InputAnn -> InputAnn
forall a. VariantSig a -> a
vsAnn (VariantSig InputAnn -> InputAnn)
-> VariantSig InputAnn -> InputAnn
forall a b. (a -> b) -> a -> b
$ (VariantSig InputAnn, VariantSig InputAnn) -> VariantSig InputAnn
forall a b. (a, b) -> a
fst (VariantSig InputAnn, VariantSig InputAnn)
v] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$ [(Constr, (VariantSig InputAnn, VariantSig InputAnn))] -> TypeError
DuplicateConstrs2 ((Constr
d, (VariantSig InputAnn, VariantSig InputAnn)
v) (Constr, (VariantSig InputAnn, VariantSig InputAnn))
-> [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
-> [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
forall a. a -> [a] -> [a]
: [(Constr, (VariantSig InputAnn, VariantSig InputAnn))]
rest)
  VariantEnv -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VariantEnv -> m ()) -> VariantEnv -> m ()
forall a b. (a -> b) -> a -> b
$ VariantEnv
env VariantEnv -> VariantEnv -> VariantEnv
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` VariantEnv
venv

-- | Add a data type to the VariantEnv and elaborate it with a dummy type.
elaborateTypeDef :: (ElaborateData b m) => Datatype InputAnn -> m (Datatype Ann)
elaborateTypeDef :: forall (b :: * -> *) (m :: * -> *).
ElaborateData b m =>
Datatype InputAnn -> m (Datatype Ann)
elaborateTypeDef = \case
  dt :: Datatype InputAnn
dt@(Datatype InputAnn
ann TypeCon
typename Maybe Kind
kind [TypeVar]
args [Variant (Maybe Type)]
variants) -> do
    let
      datatype :: Type
datatype = Type -> [Type] -> Type
typeApp (TypeCon -> Type
TypeCon TypeCon
typename) ((TypeVar -> Type) -> [TypeVar] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypeVar -> Type
TypeVar [TypeVar]
args)
      boundvars :: Set TypeVar
boundvars = [TypeVar] -> Set TypeVar
forall a. Ord a => [a] -> Set a
S.fromList [TypeVar]
args
      variantsvars :: [TypeVar]
variantsvars = [TypeVar
t | TypeVar TypeVar
t <- [Variant (Maybe Type)] -> [Type]
forall from to. Biplate from to => from -> [to]
U.universeBi [Variant (Maybe Type)]
variants]
      constrs :: Set Constr
constrs = [Constr] -> Set Constr
forall a. Ord a => [a] -> Set a
S.fromList ([Constr] -> Set Constr) -> [Constr] -> Set Constr
forall a b. (a -> b) -> a -> b
$ (Variant (Maybe Type) -> Constr)
-> [Variant (Maybe Type)] -> [Constr]
forall a b. (a -> b) -> [a] -> [b]
map (\(Variant Constr
constr Maybe Type
_) -> Constr
constr) [Variant (Maybe Type)]
variants

    -- check for duplicate data constructors, duplicate type variables
    -- and unbound type vars
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Set TypeVar -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set TypeVar
boundvars Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [TypeVar] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeVar]
args) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$
        Datatype InputAnn -> TypeError
DuplicateTypeVarsInSig Datatype InputAnn
dt

    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Variant (Maybe Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Variant (Maybe Type)]
variants Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Set Constr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set Constr
constrs) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$
        Datatype InputAnn -> TypeError
DuplicateConstrs Datatype InputAnn
dt

    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((TypeVar -> Bool) -> [TypeVar] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (TypeVar -> Set TypeVar -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set TypeVar
boundvars) [TypeVar]
variantsvars) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$
        Datatype InputAnn -> TypeError
UnboundTypeVarsInType Datatype InputAnn
dt

    -- convert to @VariantSig@s
    let
      variantsigs :: VariantEnv
variantsigs =
        [(Constr, VariantSig InputAnn)] -> VariantEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Constr, VariantSig InputAnn)] -> VariantEnv)
-> [(Constr, VariantSig InputAnn)] -> VariantEnv
forall a b. (a -> b) -> a -> b
$
          (Variant (Maybe Type) -> (Constr, VariantSig InputAnn))
-> [Variant (Maybe Type)] -> [(Constr, VariantSig InputAnn)]
forall a b. (a -> b) -> [a] -> [b]
map
            ( \(Variant Constr
constr Maybe Type
template) ->
                ( Constr
constr
                , VariantSig
                    { vsVars :: [TypeVar]
vsVars = [TypeVar]
args
                    , vsDatatype :: Type
vsDatatype = Type
datatype
                    , vsPayloadTemplate :: Maybe Type
vsPayloadTemplate = Maybe Type
template
                    , vsAnn :: InputAnn
vsAnn = InputAnn
ann
                    }
                )
            )
            [Variant (Maybe Type)]
variants
    VariantEnv -> m ()
forall (b :: * -> *) (m :: * -> *).
ElaborateData b m =>
VariantEnv -> m ()
addVariantSigs VariantEnv
variantsigs
    Datatype Ann -> m (Datatype Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ann
-> TypeCon
-> Maybe Kind
-> [TypeVar]
-> [Variant (Maybe Type)]
-> Datatype Ann
forall a.
a
-> TypeCon
-> Maybe Kind
-> [TypeVar]
-> [Variant (Maybe Type)]
-> Datatype a
Datatype (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
tUnit) TypeCon
typename Maybe Kind
kind [TypeVar]
args [Variant (Maybe Type)]
variants)

-----------------------------------------------

-- * Elaborate

-- | Run elaboration algorithm
elaborate :: (MonadBase b b) => LogAction b LogMsg -> VariantEnv -> Env Type -> [TermDef InputAnn] -> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
elaborate :: forall (b :: * -> *).
MonadBase b b =>
LogAction b LogMsg
-> VariantEnv
-> Env Type
-> [TermDef InputAnn]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
elaborate = LogAction b LogMsg
-> VariantEnv
-> Env Type
-> [TermDef InputAnn]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
forall (b :: * -> *).
MonadBase b b =>
LogAction b LogMsg
-> VariantEnv
-> Env Type
-> [TermDef InputAnn]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
runElaborate

------------

-- ** Types

-- | The monadic capabilities for the elaboration phase.
type Elaborate b m =
  ( MonadBase b m
  , MonadState ElabState m
  , MonadReader (ElabEnv b) m
  , MonadError TypeErrorA m
  )

-- | An environment from a name to something.
type Env a = Map Var a

-- | A mapping from a data constructor to the type that defines it and the type it holds.
type VariantEnv =
  Map Constr (VariantSig InputAnn)

-- | Return type for some elaboration functions.
data ElabInfo a = ElabInfo
  { forall a. ElabInfo a -> a
eiResult :: a
  -- ^ The result of the current elaboration
  , forall a. ElabInfo a -> Type
eiType :: Type
  -- ^ The type of the result
  , forall a. ElabInfo a -> Env Type
eiNewEnv :: Env Type
  -- ^ A definition to add to the environment for the next statements, if needed.
  }

-------------

-- ** Utils

-- | Try to find the type of a variable in scope.
lookupVarMaybe :: (Elaborate b m) => Var -> m (Maybe Type)
lookupVarMaybe :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Var -> m (Maybe Type)
lookupVarMaybe Var
var = do
  Env Type
env <- (ElabEnv b -> Env Type) -> m (Env Type)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ElabEnv b -> Env Type
forall (b :: * -> *). ElabEnv b -> Env Type
eeTypeEnv
  Maybe Type -> m (Maybe Type)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Var -> Env Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
var Env Type
env)

-- | Same as @lookupVarMaybe@ but throws an @UnboundVar@ error on failure.
lookupVar :: (Elaborate b m) => InputAnn -> Var -> m Type
lookupVar :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Var -> m Type
lookupVar InputAnn
ann Var
var = do
  Env Type
env <- (ElabEnv b -> Env Type) -> m (Env Type)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ElabEnv b -> Env Type
forall (b :: * -> *). ElabEnv b -> Env Type
eeTypeEnv
  m Type -> (Type -> m Type) -> Maybe Type -> m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
    ([InputAnn] -> TypeError -> m Type
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m Type) -> TypeError -> m Type
forall a b. (a -> b) -> a -> b
$ Var -> TypeError
UnboundVar Var
var)
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Var -> Env Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
var Env Type
env)

-- | Look for a builtin value/function. Throws an @UnboundVar@ error on failure.
lookupBuiltin :: (Elaborate b m) => InputAnn -> Var -> m Type
lookupBuiltin :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Var -> m Type
lookupBuiltin InputAnn
ann Var
var = do
  Env Type
env <- (ElabEnv b -> Env Type) -> m (Env Type)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ElabEnv b -> Env Type
forall (b :: * -> *). ElabEnv b -> Env Type
eeBuiltins
  m Type -> (Type -> m Type) -> Maybe Type -> m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
    ([InputAnn] -> TypeError -> m Type
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m Type) -> TypeError -> m Type
forall a b. (a -> b) -> a -> b
$ Var -> TypeError
UnboundVar Var
var)
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (Var -> Env Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
var Env Type
env)

-- | Adding new variable into the scope
insertToEnv :: Env Type -> ElabEnv b -> ElabEnv b
insertToEnv :: forall (b :: * -> *). Env Type -> ElabEnv b -> ElabEnv b
insertToEnv Env Type
vars ElabEnv b
elabEnv =
  ElabEnv b
elabEnv
    { $sel:eeTypeEnv:ElabEnv :: Env Type
eeTypeEnv = Env Type -> Env Type -> Env Type
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Env Type
vars (ElabEnv b -> Env Type
forall (b :: * -> *). ElabEnv b -> Env Type
eeTypeEnv ElabEnv b
elabEnv)
    }

-- | Remove variables from the scope
removeFromEnv :: [Var] -> ElabEnv b -> ElabEnv b
removeFromEnv :: forall (b :: * -> *). [Var] -> ElabEnv b -> ElabEnv b
removeFromEnv [Var]
vars ElabEnv b
elabEnv =
  ElabEnv b
elabEnv
    { $sel:eeTypeEnv:ElabEnv :: Env Type
eeTypeEnv = (Var -> Env Type -> Env Type) -> Env Type -> [Var] -> Env Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Var -> Env Type -> Env Type
forall k a. Ord k => k -> Map k a -> Map k a
M.delete (ElabEnv b -> Env Type
forall (b :: * -> *). ElabEnv b -> Env Type
eeTypeEnv ElabEnv b
elabEnv) [Var]
vars
    }

-- | Run an elaboration function with extra variables in scope.
withEnv :: (Elaborate b m) => [(Var, Type)] -> m a -> m a
withEnv :: forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv = (ElabEnv b -> ElabEnv b) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((ElabEnv b -> ElabEnv b) -> m a -> m a)
-> ([(Var, Type)] -> ElabEnv b -> ElabEnv b)
-> [(Var, Type)]
-> m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env Type -> ElabEnv b -> ElabEnv b
forall (b :: * -> *). Env Type -> ElabEnv b -> ElabEnv b
insertToEnv (Env Type -> ElabEnv b -> ElabEnv b)
-> ([(Var, Type)] -> Env Type)
-> [(Var, Type)]
-> ElabEnv b
-> ElabEnv b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Var, Type)] -> Env Type
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList

-- | Run an elaboration function with extra variables in scope (Map version).
withEnv' :: (Elaborate b m) => Env Type -> m a -> m a
withEnv' :: forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
Env Type -> m a -> m a
withEnv' = (ElabEnv b -> ElabEnv b) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((ElabEnv b -> ElabEnv b) -> m a -> m a)
-> (Env Type -> ElabEnv b -> ElabEnv b) -> Env Type -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env Type -> ElabEnv b -> ElabEnv b
forall (b :: * -> *). Env Type -> ElabEnv b -> ElabEnv b
insertToEnv

-- | Run an elaboration function without some variables in scope.
withoutEnv :: (Elaborate b m) => [Var] -> m a -> m a
withoutEnv :: forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[Var] -> m a -> m a
withoutEnv = (ElabEnv b -> ElabEnv b) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((ElabEnv b -> ElabEnv b) -> m a -> m a)
-> ([Var] -> ElabEnv b -> ElabEnv b) -> [Var] -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Var] -> ElabEnv b -> ElabEnv b
forall (b :: * -> *). [Var] -> ElabEnv b -> ElabEnv b
removeFromEnv

-- | Lookup variant in env
lookupVariant :: (Elaborate b m) => InputAnn -> Constr -> m (VariantSig InputAnn)
lookupVariant :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constr -> m (VariantSig InputAnn)
lookupVariant InputAnn
ann Constr
constr = do
  m (VariantSig InputAnn)
-> (VariantSig InputAnn -> m (VariantSig InputAnn))
-> Maybe (VariantSig InputAnn)
-> m (VariantSig InputAnn)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([InputAnn] -> TypeError -> m (VariantSig InputAnn)
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m (VariantSig InputAnn))
-> TypeError -> m (VariantSig InputAnn)
forall a b. (a -> b) -> a -> b
$ Constr -> TypeError
NoSuchVariant Constr
constr) VariantSig InputAnn -> m (VariantSig InputAnn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (VariantSig InputAnn) -> m (VariantSig InputAnn))
-> (ElabEnv b -> Maybe (VariantSig InputAnn))
-> ElabEnv b
-> m (VariantSig InputAnn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constr -> VariantEnv -> Maybe (VariantSig InputAnn)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Constr
constr (VariantEnv -> Maybe (VariantSig InputAnn))
-> (ElabEnv b -> VariantEnv)
-> ElabEnv b
-> Maybe (VariantSig InputAnn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ElabEnv b -> VariantEnv
forall (b :: * -> *). ElabEnv b -> VariantEnv
eeVariantEnv (ElabEnv b -> m (VariantSig InputAnn))
-> m (ElabEnv b) -> m (VariantSig InputAnn)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (ElabEnv b)
forall r (m :: * -> *). MonadReader r m => m r
ask

-- | Generate a new type variable.
genTypeVar :: (MonadState ElabState m) => Text -> m TypeVar
genTypeVar :: forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
prefix = do
  Int
n <- ElabState -> Int
esTypeVarCounter (ElabState -> Int) -> m ElabState -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m ElabState
forall s (m :: * -> *). MonadState s m => m s
get
  (ElabState -> ElabState) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((ElabState -> ElabState) -> m ())
-> (ElabState -> ElabState) -> m ()
forall a b. (a -> b) -> a -> b
$ \ElabState
s -> ElabState
s {$sel:esTypeVarCounter:ElabState :: Int
esTypeVarCounter = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}
  TypeVar -> m TypeVar
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeVar -> m TypeVar) -> TypeVar -> m TypeVar
forall a b. (a -> b) -> a -> b
$ Var -> TypeVar
TV (Var -> TypeVar) -> Var -> TypeVar
forall a b. (a -> b) -> a -> b
$ Var
prefix Var -> Var -> Var
forall a. Semigroup a => a -> a -> a
<> FilePath -> Var
toText (Int -> FilePath
forall a. Show a => a -> FilePath
show Int
n)

-- | Add a new constraint.
constrain :: (Elaborate b m) => InputAnn -> Constraint -> m ()
constrain :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann Constraint
constraint =
  (ElabState -> ElabState) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((ElabState -> ElabState) -> m ())
-> (ElabState -> ElabState) -> m ()
forall a b. (a -> b) -> a -> b
$ \ElabState
s ->
    ElabState
s
      { $sel:esConstraints:ElabState :: Constraints
esConstraints =
          ConstraintA -> Constraints -> Constraints
forall a. Ord a => a -> Set a -> Set a
S.insert (Constraint
constraint, InputAnn
ann) (ElabState -> Constraints
esConstraints ElabState
s)
      }

------------------

-- ** Algorithm

-- | Run the algorithm.
runElaborate
  :: (MonadBase b b)
  => LogAction b LogMsg
  -> VariantEnv
  -> Env Type
  -> [TermDef InputAnn]
  -> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
runElaborate :: forall (b :: * -> *).
MonadBase b b =>
LogAction b LogMsg
-> VariantEnv
-> Env Type
-> [TermDef InputAnn]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
runElaborate LogAction b LogMsg
logact VariantEnv
variantEnv Env Type
builtinsTypes =
  ( (StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
 -> ElabState -> ExceptT TypeErrorA b ([TermDef Ann], ElabState))
-> ElabState
-> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
-> ElabState -> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (Int -> Constraints -> ElabState
ElabState Int
0 Constraints
forall a. Monoid a => a
mempty)
      (StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
 -> ExceptT TypeErrorA b ([TermDef Ann], ElabState))
-> ([TermDef InputAnn]
    -> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann])
-> [TermDef InputAnn]
-> ExceptT TypeErrorA b ([TermDef Ann], ElabState)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ReaderT
   (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
 -> ElabEnv b
 -> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann])
-> ElabEnv b
-> ReaderT
     (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
-> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT
  (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
-> ElabEnv b
-> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (Env Type
-> Env Type -> VariantEnv -> LogAction b LogMsg -> ElabEnv b
forall (b :: * -> *).
Env Type
-> Env Type -> VariantEnv -> LogAction b LogMsg -> ElabEnv b
ElabEnv Env Type
forall a. Monoid a => a
mempty Env Type
builtinsTypes VariantEnv
variantEnv LogAction b LogMsg
logact)
      (ReaderT
   (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
 -> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann])
-> ([TermDef InputAnn]
    -> ReaderT
         (ElabEnv b)
         (StateT ElabState (ExceptT TypeErrorA b))
         [TermDef Ann])
-> [TermDef InputAnn]
-> StateT ElabState (ExceptT TypeErrorA b) [TermDef Ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stage
-> ReaderT
     (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
-> ReaderT
     (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
forall env (b :: * -> *) (m :: * -> *) a.
HasLog' LogMsg env b m =>
Stage -> m a -> m a
setStage Stage
TypeInference
      (ReaderT
   (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
 -> ReaderT
      (ElabEnv b)
      (StateT ElabState (ExceptT TypeErrorA b))
      [TermDef Ann])
-> ([TermDef InputAnn]
    -> ReaderT
         (ElabEnv b)
         (StateT ElabState (ExceptT TypeErrorA b))
         [TermDef Ann])
-> [TermDef InputAnn]
-> ReaderT
     (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TermDef InputAnn]
-> ReaderT
     (ElabEnv b) (StateT ElabState (ExceptT TypeErrorA b)) [TermDef Ann]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
[TermDef InputAnn] -> m [TermDef Ann]
elaborateFile
  )

-- | Elaborate a source file
elaborateFile :: (Elaborate b m) => [TermDef InputAnn] -> m [TermDef Ann]
elaborateFile :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
[TermDef InputAnn] -> m [TermDef Ann]
elaborateFile [TermDef InputAnn]
termDefs = do
  let
    names :: [Var]
names = (TermDef InputAnn -> Var) -> [TermDef InputAnn] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map TermDef InputAnn -> Var
forall a. TermDef a -> Var
getTermName [TermDef InputAnn]
termDefs
  [(Var, Type)]
vars <- (Var -> m (Var, Type)) -> [Var] -> m [(Var, Type)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Var
name -> (,) Var
name (Type -> (Var, Type))
-> (TypeVar -> Type) -> TypeVar -> (Var, Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeVar -> Type
TypeVar (TypeVar -> (Var, Type)) -> m TypeVar -> m (Var, Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"top") [Var]
names
  [(Var, Type)] -> m [TermDef Ann] -> m [TermDef Ann]
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv [(Var, Type)]
vars (m [TermDef Ann] -> m [TermDef Ann])
-> m [TermDef Ann] -> m [TermDef Ann]
forall a b. (a -> b) -> a -> b
$ (TermDef InputAnn -> m (TermDef Ann))
-> [TermDef InputAnn] -> m [TermDef Ann]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TermDef InputAnn -> m (TermDef Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (TermDef Ann)
elaborateDef [TermDef InputAnn]
termDefs

-- | Elaborate a @TermDef@
elaborateDef :: (Elaborate b m) => TermDef InputAnn -> m (TermDef Ann)
elaborateDef :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (TermDef Ann)
elaborateDef TermDef InputAnn
def = do
  -- we need to change the type of def in the current environment for recursive functions
  let
    name :: Var
name = TermDef InputAnn -> Var
forall a. TermDef a -> Var
getTermName TermDef InputAnn
def
    ann :: InputAnn
ann = TermDef InputAnn -> InputAnn
forall ann. TermDef ann -> ann
getTermAnn TermDef InputAnn
def
  Type
t <- InputAnn -> Var -> m Type
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Var -> m Type
lookupVar InputAnn
ann Var
name
  -- When evaluating a definition, the type should not be let polymorphic.
  ElabInfo (TermDef Ann)
elab <- [(Var, Type)]
-> m (ElabInfo (TermDef Ann)) -> m (ElabInfo (TermDef Ann))
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv [(Var
name, Type
t)] (m (ElabInfo (TermDef Ann)) -> m (ElabInfo (TermDef Ann)))
-> m (ElabInfo (TermDef Ann)) -> m (ElabInfo (TermDef Ann))
forall a b. (a -> b) -> a -> b
$ TermDef InputAnn -> m (ElabInfo (TermDef Ann))
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (ElabInfo (TermDef Ann))
elaborateTermDef TermDef InputAnn
def
  InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
t (ElabInfo (TermDef Ann) -> Type
forall a. ElabInfo a -> Type
eiType ElabInfo (TermDef Ann)
elab)
  TermDef Ann -> m (TermDef Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TermDef Ann -> m (TermDef Ann)) -> TermDef Ann -> m (TermDef Ann)
forall a b. (a -> b) -> a -> b
$ Ann -> TermDef Ann -> TermDef Ann
forall a. a -> TermDef a -> TermDef a
setTermAnn (InputAnn -> Type -> Ann
Ann InputAnn
ann (Type -> Ann) -> Type -> Ann
forall a b. (a -> b) -> a -> b
$ ElabInfo (TermDef Ann) -> Type
forall a. ElabInfo a -> Type
eiType ElabInfo (TermDef Ann)
elab) (TermDef Ann -> TermDef Ann) -> TermDef Ann -> TermDef Ann
forall a b. (a -> b) -> a -> b
$ ElabInfo (TermDef Ann) -> TermDef Ann
forall a. ElabInfo a -> a
eiResult ElabInfo (TermDef Ann)
elab

-- | Elaborate a term definition helper function
elaborateTermDef :: (Elaborate b m) => TermDef InputAnn -> m (ElabInfo (TermDef Ann))
elaborateTermDef :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (ElabInfo (TermDef Ann))
elaborateTermDef = \case
  Variable InputAnn
ann Var
name Maybe Type
mtype Expr InputAnn
expr -> do
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Var
name Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
"main") (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality (Expr Ann -> Type
getType Expr Ann
expr') (Type -> Type
tIO Type
tUnit)
    m () -> (Type -> m ()) -> Maybe Type -> m ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> (Type -> Constraint) -> Type -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type -> Constraint
Equality (Expr Ann -> Type
getType Expr Ann
expr')) Maybe Type
mtype
    ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann)))
-> ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann))
forall a b. (a -> b) -> a -> b
$
      ElabInfo
        { $sel:eiType:ElabInfo :: Type
eiType = Expr Ann -> Type
getType Expr Ann
expr'
        , $sel:eiResult:ElabInfo :: TermDef Ann
eiResult = Ann -> Var -> Maybe Type -> Expr Ann -> TermDef Ann
forall a. a -> Var -> Maybe Type -> Expr a -> TermDef a
Variable (InputAnn -> Type -> Ann
Ann InputAnn
ann (Type -> Ann) -> Type -> Ann
forall a b. (a -> b) -> a -> b
$ Expr Ann -> Type
getType Expr Ann
expr') Var
name Maybe Type
forall a. Maybe a
Nothing Expr Ann
expr'
        , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Env Type
forall a. Monoid a => a
mempty
        }
  Function InputAnn
ann Var
name Maybe Type
mtype [Maybe Var]
args Expr InputAnn
body -> do
    TypeVar
tfun <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"tfun"
    [(Maybe Var, TypeVar)]
targsEnv <- (Maybe Var -> m (Maybe Var, TypeVar))
-> [Maybe Var] -> m [(Maybe Var, TypeVar)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Maybe Var
arg -> (,) Maybe Var
arg (TypeVar -> (Maybe Var, TypeVar))
-> m TypeVar -> m (Maybe Var, TypeVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t") [Maybe Var]
args
    Expr Ann
body' <-
      [(Var, Type)] -> m (Expr Ann) -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv
        ( (Var
name, TypeVar -> Type
TypeVar TypeVar
tfun)
            (Var, Type) -> [(Var, Type)] -> [(Var, Type)]
forall a. a -> [a] -> [a]
: ((Maybe Var, TypeVar) -> Maybe (Var, Type))
-> [(Maybe Var, TypeVar)] -> [(Var, Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
              (\(Maybe Var
var, TypeVar
t) -> (Var -> Type -> (Var, Type)) -> Type -> Var -> (Var, Type)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) (TypeVar -> Type
TypeVar TypeVar
t) (Var -> (Var, Type)) -> Maybe Var -> Maybe (Var, Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Var
var)
              [(Maybe Var, TypeVar)]
targsEnv
        )
        (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
body)

    let
      tret :: Type
tret = Expr Ann -> Type
getType Expr Ann
body'
      tfunexpanded :: Type
tfunexpanded = [Type] -> Type -> Type
typeFun (((Maybe Var, TypeVar) -> Type) -> [(Maybe Var, TypeVar)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TypeVar -> Type
TypeVar (TypeVar -> Type)
-> ((Maybe Var, TypeVar) -> TypeVar)
-> (Maybe Var, TypeVar)
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Var, TypeVar) -> TypeVar
forall a b. (a, b) -> b
snd) [(Maybe Var, TypeVar)]
targsEnv) Type
tret

    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
      Type -> Type -> Constraint
Equality
        (TypeVar -> Type
TypeVar TypeVar
tfun)
        Type
tfunexpanded

    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Var
name Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
"main") (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality (TypeVar -> Type
TypeVar TypeVar
tfun) (Type -> Type
tIO Type
tUnit)

    m () -> (Type -> m ()) -> Maybe Type -> m ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> (Type -> Constraint) -> Type -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type -> Constraint
Equality (TypeVar -> Type
TypeVar TypeVar
tfun)) Maybe Type
mtype

    ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann)))
-> ElabInfo (TermDef Ann) -> m (ElabInfo (TermDef Ann))
forall a b. (a -> b) -> a -> b
$
      ElabInfo
        { $sel:eiType:ElabInfo :: Type
eiType = Type
tfunexpanded
        , $sel:eiResult:ElabInfo :: TermDef Ann
eiResult = Ann -> Var -> Maybe Type -> [Maybe Var] -> Expr Ann -> TermDef Ann
forall a.
a -> Var -> Maybe Type -> [Maybe Var] -> Expr a -> TermDef a
Function (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
tfunexpanded) Var
name Maybe Type
forall a. Maybe a
Nothing [Maybe Var]
args Expr Ann
body'
        , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Env Type
forall a. Monoid a => a
mempty
        }

-- | Elaborate a list of statements.
--   Returns the type of the final statement as well.
elaborateBlock :: (Elaborate b m) => Block InputAnn -> m (Type, Block Ann)
elaborateBlock :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Block InputAnn -> m (Type, Block Ann)
elaborateBlock Block InputAnn
block = do
  -- we want to fold over the list of statements, and for each new definition,
  -- add it to the environment for the elaboration of the next expressions
  ElabInfo (Block Ann)
block' <-
    (ElabInfo (Block Ann)
 -> Statement InputAnn -> m (ElabInfo (Block Ann)))
-> ElabInfo (Block Ann)
-> Block InputAnn
-> m (ElabInfo (Block Ann))
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      ( \ElabInfo (Block Ann)
blockinfo Statement InputAnn
stmt -> do
          ElabInfo (Statement Ann)
stmt' <- Env Type
-> m (ElabInfo (Statement Ann)) -> m (ElabInfo (Statement Ann))
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
Env Type -> m a -> m a
withEnv' (ElabInfo (Block Ann) -> Env Type
forall a. ElabInfo a -> Env Type
eiNewEnv ElabInfo (Block Ann)
blockinfo) (m (ElabInfo (Statement Ann)) -> m (ElabInfo (Statement Ann)))
-> m (ElabInfo (Statement Ann)) -> m (ElabInfo (Statement Ann))
forall a b. (a -> b) -> a -> b
$ Statement InputAnn -> m (ElabInfo (Statement Ann))
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Statement InputAnn -> m (ElabInfo (Statement Ann))
elaborateStmt Statement InputAnn
stmt
          ElabInfo (Block Ann) -> m (ElabInfo (Block Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (Block Ann) -> m (ElabInfo (Block Ann)))
-> ElabInfo (Block Ann) -> m (ElabInfo (Block Ann))
forall a b. (a -> b) -> a -> b
$
            ElabInfo
              { $sel:eiResult:ElabInfo :: Block Ann
eiResult = ElabInfo (Statement Ann) -> Statement Ann
forall a. ElabInfo a -> a
eiResult ElabInfo (Statement Ann)
stmt' Statement Ann -> Block Ann -> Block Ann
forall a. a -> [a] -> [a]
: ElabInfo (Block Ann) -> Block Ann
forall a. ElabInfo a -> a
eiResult ElabInfo (Block Ann)
blockinfo
              , $sel:eiType:ElabInfo :: Type
eiType = ElabInfo (Statement Ann) -> Type
forall a. ElabInfo a -> Type
eiType ElabInfo (Statement Ann)
stmt'
              , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Env Type -> Env Type -> Env Type
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (ElabInfo (Statement Ann) -> Env Type
forall a. ElabInfo a -> Env Type
eiNewEnv ElabInfo (Statement Ann)
stmt') (ElabInfo (Block Ann) -> Env Type
forall a. ElabInfo a -> Env Type
eiNewEnv ElabInfo (Block Ann)
blockinfo)
              }
      )
      (Block Ann -> Type -> Env Type -> ElabInfo (Block Ann)
forall a. a -> Type -> Env Type -> ElabInfo a
ElabInfo [] (Type -> Type
tIO Type
tUnit) Env Type
forall a. Monoid a => a
mempty)
      Block InputAnn
block
  (Type, Block Ann) -> m (Type, Block Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (Block Ann) -> Type
forall a. ElabInfo a -> Type
eiType ElabInfo (Block Ann)
block', Block Ann -> Block Ann
forall a. [a] -> [a]
reverse (ElabInfo (Block Ann) -> Block Ann
forall a. ElabInfo a -> a
eiResult ElabInfo (Block Ann)
block'))

-- | Elaborate a single statement.
elaborateStmt :: (Elaborate b m) => Statement InputAnn -> m (ElabInfo (Statement Ann))
elaborateStmt :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Statement InputAnn -> m (ElabInfo (Statement Ann))
elaborateStmt = \case
  SExpr InputAnn
ann Expr InputAnn
expr -> do
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Type
t <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality (Expr Ann -> Type
getType Expr Ann
expr') (Type -> Type
tIO Type
t)
    ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann)))
-> ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall a b. (a -> b) -> a -> b
$
      ElabInfo
        { $sel:eiResult:ElabInfo :: Statement Ann
eiResult = Ann -> Expr Ann -> Statement Ann
forall a. a -> Expr a -> Statement a
SExpr (InputAnn -> Type -> Ann
Ann InputAnn
ann (Expr Ann -> Type
getType Expr Ann
expr')) Expr Ann
expr'
        , $sel:eiType:ElabInfo :: Type
eiType = Expr Ann -> Type
getType Expr Ann
expr'
        , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Env Type
forall a. Monoid a => a
mempty
        }
  SDef InputAnn
ann TermDef InputAnn
def -> do
    ElabInfo TermDef Ann
def' Type
t Env Type
_ <- TermDef InputAnn -> m (ElabInfo (TermDef Ann))
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (ElabInfo (TermDef Ann))
elaborateTermDef TermDef InputAnn
def
    ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann)))
-> ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall a b. (a -> b) -> a -> b
$
      ElabInfo
        { $sel:eiResult:ElabInfo :: Statement Ann
eiResult = Ann -> TermDef Ann -> Statement Ann
forall a. a -> TermDef a -> Statement a
SDef (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
t) TermDef Ann
def'
        , $sel:eiType:ElabInfo :: Type
eiType = Type
t
        , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Var -> Type -> Env Type
forall k a. k -> a -> Map k a
M.singleton (TermDef Ann -> Var
forall a. TermDef a -> Var
getTermName TermDef Ann
def') Type
t
        }
  SBind InputAnn
ann Var
name Expr InputAnn
expr -> do
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Type
t <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    let
      t' :: Type
t' = Expr Ann -> Type
getType Expr Ann
expr'
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
t' (Type -> Type
tIO Type
t)
    ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann)))
-> ElabInfo (Statement Ann) -> m (ElabInfo (Statement Ann))
forall a b. (a -> b) -> a -> b
$
      ElabInfo
        { $sel:eiResult:ElabInfo :: Statement Ann
eiResult = Ann -> Var -> Expr Ann -> Statement Ann
forall a. a -> Var -> Expr a -> Statement a
SBind (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
t) Var
name Expr Ann
expr'
        , $sel:eiType:ElabInfo :: Type
eiType = Type
t'
        , $sel:eiNewEnv:ElabInfo :: Env Type
eiNewEnv = Var -> Type -> Env Type
forall k a. k -> a -> Map k a
M.singleton Var
name Type
t
        }

-- | Elaborate an expression.
--   Traverses the expression Top-Down.
elaborateExpr :: (Elaborate b m) => InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann = \case
  -- Replace the current annotation an evaluate the inner expression
  EAnnotated InputAnn
ann' Expr InputAnn
e ->
    InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann' Expr InputAnn
e
  -- The types of literals are known.
  ELit Lit
lit ->
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$ Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann (Type -> Ann) -> Type -> Ann
forall a b. (a -> b) -> a -> b
$ Lit -> Type
getLitType Lit
lit) (Lit -> Expr Ann
forall a. Lit -> Expr a
ELit Lit
lit)
  -- For variables, we look it up in the environment
  EVar Var
var -> do
    Type
t <- m Type -> (Type -> m Type) -> Maybe Type -> m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (InputAnn -> Var -> m Type
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Var -> m Type
lookupBuiltin InputAnn
ann Var
var) Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Type -> m Type) -> m (Maybe Type) -> m Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Var -> m (Maybe Type)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Var -> m (Maybe Type)
lookupVarMaybe Var
var
    Type
typ <- do
      TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality (TypeVar -> Type
TypeVar TypeVar
tv) Type
t
      Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> Type
TypeVar TypeVar
tv
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
typ) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Var -> Expr Ann
forall a. Var -> Expr a
EVar Var
var

  -- For operators, we look it up in the environment
  EOp opDef :: OpDef
opDef@OpDef {Var
opFunRef :: OpDef -> Var
opFunRef :: Var
opFunRef} -> do
    Type
t <- m Type -> (Type -> m Type) -> Maybe Type -> m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (InputAnn -> Var -> m Type
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Var -> m Type
lookupBuiltin InputAnn
ann Var
opFunRef) Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Type -> m Type) -> m (Maybe Type) -> m Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Var -> m (Maybe Type)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Var -> m (Maybe Type)
lookupVarMaybe Var
opFunRef
    Type
typ <- do
      TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality (TypeVar -> Type
TypeVar TypeVar
tv) Type
t
      Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> Type
TypeVar TypeVar
tv
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
typ) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        OpDef -> Expr Ann
forall a. OpDef -> Expr a
EOp OpDef
opDef

  -- Generate type variables for function arguments, and annotate
  -- the body with the arguments in the environment.
  -- The result type should be a function from the arguments types to the type of the body
  EFun [Maybe Var]
args Expr InputAnn
body -> do
    [(Maybe Var, TypeVar)]
targsEnv <- (Maybe Var -> m (Maybe Var, TypeVar))
-> [Maybe Var] -> m [(Maybe Var, TypeVar)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Maybe Var
arg -> (,) Maybe Var
arg (TypeVar -> (Maybe Var, TypeVar))
-> m TypeVar -> m (Maybe Var, TypeVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t") [Maybe Var]
args
    Expr Ann
body' <-
      [(Var, Type)] -> m (Expr Ann) -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv
        ( ((Maybe Var, TypeVar) -> Maybe (Var, Type))
-> [(Maybe Var, TypeVar)] -> [(Var, Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
            (\(Maybe Var
var, TypeVar
t) -> (Var -> Type -> (Var, Type)) -> Type -> Var -> (Var, Type)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) (TypeVar -> Type
TypeVar TypeVar
t) (Var -> (Var, Type)) -> Maybe Var -> Maybe (Var, Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Var
var)
            [(Maybe Var, TypeVar)]
targsEnv
        )
        (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
body)

    let
      tret :: Type
tret = Expr Ann -> Type
getType Expr Ann
body'

    let
      tfun :: Type
tfun = [Type] -> Type -> Type
typeFun (((Maybe Var, TypeVar) -> Type) -> [(Maybe Var, TypeVar)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TypeVar -> Type
TypeVar (TypeVar -> Type)
-> ((Maybe Var, TypeVar) -> TypeVar)
-> (Maybe Var, TypeVar)
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Var, TypeVar) -> TypeVar
forall a b. (a, b) -> b
snd) [(Maybe Var, TypeVar)]
targsEnv) Type
tret

    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
tfun) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        [Maybe Var] -> Expr Ann -> Expr Ann
forall a. [Maybe Var] -> Expr a -> Expr a
EFun [Maybe Var]
args Expr Ann
body'

  -- Generate a return type and constrain the type of the function as
  -- a function from the arguments types to the generated return type
  EFunCall Expr InputAnn
f [Expr InputAnn]
args -> do
    Expr Ann
f' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
f
    [Expr Ann]
args' <- (Expr InputAnn -> m (Expr Ann)) -> [Expr InputAnn] -> m [Expr Ann]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann) [Expr InputAnn]
args
    Type
tret <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
      Type -> Type -> Constraint
Equality
        (Expr Ann -> Type
getType Expr Ann
f')
        ([Type] -> Type -> Type
typeFun ((Expr Ann -> Type) -> [Expr Ann] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Expr Ann -> Type
getType [Expr Ann]
args') Type
tret)
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
tret) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Expr Ann -> [Expr Ann] -> Expr Ann
forall a. Expr a -> [Expr a] -> Expr a
EFunCall Expr Ann
f' [Expr Ann]
args'

  -- If we have a type for an ffi call supplied by the user, we'll use it for type inference.
  -- otherwise we'll take any amount and any type of arguments and return IO {}.
  EFfi Var
f Maybe Type
mtype [Expr InputAnn]
args -> do
    [Expr Ann]
args' <- (Expr InputAnn -> m (Expr Ann)) -> [Expr InputAnn] -> m [Expr Ann]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann) [Expr InputAnn]
args
    case Maybe Type
mtype of
      Maybe Type
Nothing -> do
        Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
          Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann (Type -> Type
tIO Type
tUnit)) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
            Var -> Maybe Type -> [Expr Ann] -> Expr Ann
forall a. Var -> Maybe Type -> [Expr a] -> Expr a
EFfi Var
f Maybe Type
forall a. Maybe a
Nothing [Expr Ann]
args'
      Just Type
t -> do
        let
          ([Type]
targs, Type
tret) = Type -> ([Type], Type)
toTypeFun Type
t

        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Expr Ann] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr Ann]
args' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
targs) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
          Type
tmp <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
          [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> TypeError
ArityMismatch Type
t ([Type] -> Type -> Type
typeFun ((Expr Ann -> Type) -> [Expr Ann] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Expr Ann -> Type
getType [Expr Ann]
args') Type
tmp)

        (Type -> Type -> m ()) -> [Type] -> [Type] -> m ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ ((Constraint -> m ()) -> (Type -> Constraint) -> Type -> m ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann) ((Type -> Constraint) -> Type -> m ())
-> (Type -> Type -> Constraint) -> Type -> Type -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type -> Constraint
Equality) [Type]
targs ((Expr Ann -> Type) -> [Expr Ann] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Expr Ann -> Type
getType [Expr Ann]
args')

        Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
          Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
tret) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
            Var -> Maybe Type -> [Expr Ann] -> Expr Ann
forall a. Var -> Maybe Type -> [Expr a] -> Expr a
EFfi Var
f (Type -> Maybe Type
forall a. a -> Maybe a
Just Type
t) [Expr Ann]
args'
  EBlock Block InputAnn
block -> do
    (Type
t, Block Ann
block') <- Block InputAnn -> m (Type, Block Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
Block InputAnn -> m (Type, Block Ann)
elaborateBlock Block InputAnn
block
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
t) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Block Ann -> Expr Ann
forall a. Block a -> Expr a
EBlock Block Ann
block'
  ELet TermDef InputAnn
termdef Expr InputAnn
expr -> do
    ElabInfo {TermDef Ann
eiResult :: TermDef Ann
$sel:eiResult:ElabInfo :: forall a. ElabInfo a -> a
eiResult, Type
eiType :: Type
$sel:eiType:ElabInfo :: forall a. ElabInfo a -> Type
eiType} <- TermDef InputAnn -> m (ElabInfo (TermDef Ann))
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
TermDef InputAnn -> m (ElabInfo (TermDef Ann))
elaborateTermDef TermDef InputAnn
termdef
    Expr Ann
expr' <- [(Var, Type)] -> m (Expr Ann) -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv [(TermDef InputAnn -> Var
forall a. TermDef a -> Var
getTermName TermDef InputAnn
termdef, Type
eiType)] (m (Expr Ann) -> m (Expr Ann)) -> m (Expr Ann) -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$ InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann (Expr Ann -> Type
getType Expr Ann
expr')) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        TermDef Ann -> Expr Ann -> Expr Ann
forall a. TermDef a -> Expr a -> Expr a
ELet TermDef Ann
eiResult Expr Ann
expr'

  -- lookup variant in the environment
  EVariant Constr
constr -> do
    Type
vsType <- VariantSig InputAnn -> Type
forall a. VariantSig a -> Type
vsToTypeScheme (VariantSig InputAnn -> Type) -> m (VariantSig InputAnn) -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InputAnn -> Constr -> m (VariantSig InputAnn)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constr -> m (VariantSig InputAnn)
lookupVariant InputAnn
ann Constr
constr
    Type
t <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
t Type
vsType

    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
t) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Constr -> Expr Ann
forall a. Constr -> Expr a
EVariant Constr
constr

  -- An open variant has a specific type according to the constructor and value,
  -- but may be plugged in any place that can handle it.
  -- Must take exactly one value as input.
  EOpenVariant Constr
constr -> do
    TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    Type
t <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann ([Type] -> Type -> Type
typeFun [Type
t] (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ [(Constr, Type)] -> TypeVar -> Type
TypePolyVariantLB [(Constr
constr, Type
t)] TypeVar
tv)) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Constr -> Expr Ann
forall a. Constr -> Expr a
EOpenVariant Constr
constr

  -- For if expressions, the condition should be of type Bool,
  -- and the two branches should match.
  EIf Expr InputAnn
cond Expr InputAnn
trueBranch Expr InputAnn
falseBranch -> do
    Expr Ann
cond' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
cond
    Expr Ann
trueBranch' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
trueBranch
    Expr Ann
falseBranch' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
falseBranch
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
tBool (Expr Ann -> Type
getType Expr Ann
cond')
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality (Expr Ann -> Type
getType Expr Ann
trueBranch') (Expr Ann -> Type
getType Expr Ann
falseBranch')
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann (Expr Ann -> Type
getType Expr Ann
trueBranch')) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Expr Ann -> Expr Ann -> Expr Ann -> Expr Ann
forall a. Expr a -> Expr a -> Expr a -> Expr a
EIf Expr Ann
cond' Expr Ann
trueBranch' Expr Ann
falseBranch'
  ECase Expr InputAnn
expr [(Pattern, Expr InputAnn)]
patterns -> do
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Type
patT <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    [(Pattern, Expr Ann)]
patterns' <- InputAnn
-> Type
-> Type
-> [(Pattern, Expr InputAnn)]
-> m [(Pattern, Expr Ann)]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn
-> Type
-> Type
-> [(Pattern, Expr InputAnn)]
-> m [(Pattern, Expr Ann)]
elaboratePatterns InputAnn
ann (Expr Ann -> Type
getType Expr Ann
expr') Type
patT [(Pattern, Expr InputAnn)]
patterns
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
patT) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Expr Ann -> [(Pattern, Expr Ann)] -> Expr Ann
forall a. Expr a -> [(Pattern, Expr a)] -> Expr a
ECase Expr Ann
expr' [(Pattern, Expr Ann)]
patterns'
  ERecord Record (Expr InputAnn)
record -> do
    Map Label (Expr Ann)
record' <- (Expr InputAnn -> m (Expr Ann))
-> Record (Expr InputAnn) -> m (Map Label (Expr Ann))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann) Record (Expr InputAnn)
record
    let
      recordT :: Type
recordT = [(Label, Type)] -> Type
TypeRec ([(Label, Type)] -> Type) -> [(Label, Type)] -> Type
forall a b. (a -> b) -> a -> b
$ Map Label Type -> [(Label, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Label Type -> [(Label, Type)])
-> Map Label Type -> [(Label, Type)]
forall a b. (a -> b) -> a -> b
$ (Expr Ann -> Type) -> Map Label (Expr Ann) -> Map Label Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr Ann -> Type
getType Map Label (Expr Ann)
record'

    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
recordT) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Map Label (Expr Ann) -> Expr Ann
forall a. Record (Expr a) -> Expr a
ERecord Map Label (Expr Ann)
record'

  -- the "expr" should be a record with at least the "label" label
  ERecordAccess Expr InputAnn
expr Label
label -> do
    Type
labelT <- TypeVar -> Type
TypeVar (TypeVar -> Type) -> m TypeVar -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    TypeVar
ext <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
      Type -> Type -> Constraint
Equality
        (Expr Ann -> Type
getType Expr Ann
expr')
        ([(Label, Type)] -> TypeVar -> Type
TypeRecExt [(Label
label, Type
labelT)] TypeVar
ext)
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
labelT) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Expr Ann -> Label -> Expr Ann
forall a. Expr a -> Label -> Expr a
ERecordAccess Expr Ann
expr' Label
label
  ERecordExtension Record (Expr InputAnn)
record Expr InputAnn
expr -> do
    Expr Ann
expr' <- InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
expr
    Map Label (Expr Ann)
record' <- (Expr InputAnn -> m (Expr Ann))
-> Record (Expr InputAnn) -> m (Map Label (Expr Ann))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann) Record (Expr InputAnn)
record
    TypeVar
et <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality ([(Label, Type)] -> TypeVar -> Type
TypeRecExt [] TypeVar
et) (Expr Ann -> Type
getType Expr Ann
expr')
    let
      recordTypes :: [(Label, Type)]
recordTypes = Map Label Type -> [(Label, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Label Type -> [(Label, Type)])
-> Map Label Type -> [(Label, Type)]
forall a b. (a -> b) -> a -> b
$ (Expr Ann -> Type) -> Map Label (Expr Ann) -> Map Label Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr Ann -> Type
getType Map Label (Expr Ann)
record'
      t :: Type
t = [(Label, Type)] -> TypeVar -> Type
TypeRecExt [(Label, Type)]
recordTypes TypeVar
et
    Expr Ann -> m (Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann -> m (Expr Ann)) -> Expr Ann -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$
      Ann -> Expr Ann -> Expr Ann
forall a. a -> Expr a -> Expr a
EAnnotated (InputAnn -> Type -> Ann
Ann InputAnn
ann Type
t) (Expr Ann -> Expr Ann) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$
        Map Label (Expr Ann) -> Expr Ann -> Expr Ann
forall a. Record (Expr a) -> Expr a -> Expr a
ERecordExtension Map Label (Expr Ann)
record' Expr Ann
expr'

-- | Elaborate patterns in case expressions
elaboratePatterns
  :: (Elaborate b m)
  => InputAnn
  -> Type
  -> Type
  -> [(Pattern, Expr InputAnn)]
  -> m [(Pattern, Expr Ann)]
elaboratePatterns :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn
-> Type
-> Type
-> [(Pattern, Expr InputAnn)]
-> m [(Pattern, Expr Ann)]
elaboratePatterns InputAnn
ann Type
exprT Type
bodyT [(Pattern, Expr InputAnn)]
pats = do
  TypeVar
exprTV <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
  InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
exprT (TypeVar -> Type
TypeVar TypeVar
exprTV)
  [(Pattern, Expr InputAnn)]
-> ((Pattern, Expr InputAnn) -> m (Pattern, Expr Ann))
-> m [(Pattern, Expr Ann)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [(Pattern, Expr InputAnn)]
pats (((Pattern, Expr InputAnn) -> m (Pattern, Expr Ann))
 -> m [(Pattern, Expr Ann)])
-> ((Pattern, Expr InputAnn) -> m (Pattern, Expr Ann))
-> m [(Pattern, Expr Ann)]
forall a b. (a -> b) -> a -> b
$ \(Pattern
pat, Expr InputAnn
body) -> do
    [(Var, Type)]
env <- InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
elaboratePattern InputAnn
ann TypeVar
exprTV Pattern
pat
    Expr Ann
body' <- [(Var, Type)] -> m (Expr Ann) -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *) a.
Elaborate b m =>
[(Var, Type)] -> m a -> m a
withEnv [(Var, Type)]
env (m (Expr Ann) -> m (Expr Ann)) -> m (Expr Ann) -> m (Expr Ann)
forall a b. (a -> b) -> a -> b
$ InputAnn -> Expr InputAnn -> m (Expr Ann)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Expr InputAnn -> m (Expr Ann)
elaborateExpr InputAnn
ann Expr InputAnn
body
    let
      t :: Type
t = Expr Ann -> Type
getType Expr Ann
body'
    InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality Type
bodyT Type
t
    (Pattern, Expr Ann) -> m (Pattern, Expr Ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pattern
pat, Expr Ann
body')

-- | Elaborate a single pattern match
elaboratePattern
  :: (Elaborate b m)
  => InputAnn
  -> TypeVar
  -> Pattern
  -> m [(Var, Type)]
elaboratePattern :: forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
elaboratePattern InputAnn
ann TypeVar
exprTV Pattern
outerPat = do
  case Pattern
outerPat of
    Pattern
PWildcard -> do
      TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality
          (TypeVar -> Type
TypeVar TypeVar
exprTV)
          ([(Constr, Type)] -> TypeVar -> Type
TypePolyVariantLB [] TypeVar
tv)
      [(Var, Type)] -> m [(Var, Type)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var, Type)]
forall a. Monoid a => a
mempty

    -- variable capture does different things depending if it matches with variants
    -- or other types.
    -- For variants, it matches any variant, but for the specified variants it must
    -- match their type.
    -- Otherwise it will match the expression type.
    PVar Var
v -> do
      TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality
          (TypeVar -> Type
TypeVar TypeVar
exprTV)
          ([(Constr, Type)] -> TypeVar -> Type
TypePolyVariantLB [] TypeVar
tv)
      [(Var, Type)] -> m [(Var, Type)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var
v, TypeVar -> Type
TypeVar TypeVar
exprTV)]
    PLit Lit
lit -> do
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
Equality (Lit -> Type
getLitType Lit
lit) (TypeVar -> Type
TypeVar TypeVar
exprTV)
      [(Var, Type)] -> m [(Var, Type)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var, Type)]
forall a. Monoid a => a
mempty

    -- open variants constrain the type of expr to be able to handle their variant
    POpenVariant (Variant Constr
constr Pattern
innerPat) -> do
      TypeVar
innerPatT <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality
          (TypeVar -> Type
TypeVar TypeVar
exprTV)
          (TypeVar -> [(Constr, Type)] -> Type
TypePolyVariantUB TypeVar
exprTV [(Constr
constr, TypeVar -> Type
TypeVar TypeVar
innerPatT)])
      InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
elaboratePattern InputAnn
ann TypeVar
innerPatT Pattern
innerPat
    PVariant (Variant Constr
constr Maybe Pattern
minnerPat) -> do
      VariantSig InputAnn
vs <- InputAnn -> Constr -> m (VariantSig InputAnn)
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constr -> m (VariantSig InputAnn)
lookupVariant InputAnn
ann Constr
constr
      case Maybe Pattern
minnerPat of
        Maybe Pattern
Nothing -> do
          InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
            Type -> Type -> Constraint
Equality (TypeVar -> Type
TypeVar TypeVar
exprTV) (VariantSig InputAnn -> Type
forall a. VariantSig a -> Type
vsToTypeScheme VariantSig InputAnn
vs)
          [(Var, Type)] -> m [(Var, Type)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var, Type)]
forall a. Monoid a => a
mempty
        Just Pattern
innerPat -> do
          TypeVar
tv <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
          InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
            Type -> Type -> Constraint
Equality
              ([Type] -> Type -> Type
typeFun [TypeVar -> Type
TypeVar TypeVar
tv] (TypeVar -> Type
TypeVar TypeVar
exprTV))
              (VariantSig InputAnn -> Type
forall a. VariantSig a -> Type
vsToTypeScheme VariantSig InputAnn
vs)
          InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
elaboratePattern InputAnn
ann TypeVar
tv Pattern
innerPat
    PRecord Record Pattern
record -> do
      Map Label (Type, [(Var, Type)])
record' <- Record Pattern
-> (Pattern -> m (Type, [(Var, Type)]))
-> m (Map Label (Type, [(Var, Type)]))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for Record Pattern
record ((Pattern -> m (Type, [(Var, Type)]))
 -> m (Map Label (Type, [(Var, Type)])))
-> (Pattern -> m (Type, [(Var, Type)]))
-> m (Map Label (Type, [(Var, Type)]))
forall a b. (a -> b) -> a -> b
$ \Pattern
pat -> do
        TypeVar
t <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
        (,) (TypeVar -> Type
TypeVar TypeVar
t) ([(Var, Type)] -> (Type, [(Var, Type)]))
-> m [(Var, Type)] -> m (Type, [(Var, Type)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> TypeVar -> Pattern -> m [(Var, Type)]
elaboratePattern InputAnn
ann TypeVar
t Pattern
pat

      TypeVar
ext <- Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
"t"
      InputAnn -> Constraint -> m ()
forall (b :: * -> *) (m :: * -> *).
Elaborate b m =>
InputAnn -> Constraint -> m ()
constrain InputAnn
ann (Constraint -> m ()) -> Constraint -> m ()
forall a b. (a -> b) -> a -> b
$
        Type -> Type -> Constraint
Equality
          (TypeVar -> Type
TypeVar TypeVar
exprTV)
          ([(Label, Type)] -> TypeVar -> Type
TypeRecExt (Map Label Type -> [(Label, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Label Type -> [(Label, Type)])
-> Map Label Type -> [(Label, Type)]
forall a b. (a -> b) -> a -> b
$ ((Type, [(Var, Type)]) -> Type)
-> Map Label (Type, [(Var, Type)]) -> Map Label Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type, [(Var, Type)]) -> Type
forall a b. (a, b) -> a
fst Map Label (Type, [(Var, Type)])
record') TypeVar
ext)

      -- check if we have duplicate names in the environment
      let
        envs :: [Env Type]
envs = ((Type, [(Var, Type)]) -> Env Type)
-> [(Type, [(Var, Type)])] -> [Env Type]
forall a b. (a -> b) -> [a] -> [b]
map ([(Var, Type)] -> Env Type
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Type)] -> Env Type)
-> ((Type, [(Var, Type)]) -> [(Var, Type)])
-> (Type, [(Var, Type)])
-> Env Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type, [(Var, Type)]) -> [(Var, Type)]
forall a b. (a, b) -> b
snd) ([(Type, [(Var, Type)])] -> [Env Type])
-> [(Type, [(Var, Type)])] -> [Env Type]
forall a b. (a -> b) -> a -> b
$ Map Label (Type, [(Var, Type)]) -> [(Type, [(Var, Type)])]
forall k a. Map k a -> [a]
M.elems Map Label (Type, [(Var, Type)])
record'
        env :: Env Type
env = [Env Type] -> Env Type
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Env Type] -> Env Type) -> [Env Type] -> Env Type
forall a b. (a -> b) -> a -> b
$ ((Type, [(Var, Type)]) -> Env Type)
-> [(Type, [(Var, Type)])] -> [Env Type]
forall a b. (a -> b) -> [a] -> [b]
map ([(Var, Type)] -> Env Type
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Type)] -> Env Type)
-> ((Type, [(Var, Type)]) -> [(Var, Type)])
-> (Type, [(Var, Type)])
-> Env Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type, [(Var, Type)]) -> [(Var, Type)]
forall a b. (a, b) -> b
snd) ([(Type, [(Var, Type)])] -> [Env Type])
-> [(Type, [(Var, Type)])] -> [Env Type]
forall a b. (a -> b) -> a -> b
$ Map Label (Type, [(Var, Type)]) -> [(Type, [(Var, Type)])]
forall k a. Map k a -> [a]
M.elems Map Label (Type, [(Var, Type)])
record'

      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Env Type -> Int) -> [Env Type] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Env Type -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Env Type]
envs) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Env Type -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Env Type
env) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$ Pattern -> TypeError
DuplicateVarsInPattern Pattern
outerPat
      [(Var, Type)] -> m [(Var, Type)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Var, Type)] -> m [(Var, Type)])
-> [(Var, Type)] -> m [(Var, Type)]
forall a b. (a -> b) -> a -> b
$ Env Type -> [(Var, Type)]
forall k a. Map k a -> [(k, a)]
M.toList Env Type
env

-- | The type of a literal
getLitType :: Lit -> Type
getLitType :: Lit -> Type
getLitType = \case
  LBool {} -> Type
tBool
  LInt {} -> Type
tInt
  LString {} -> Type
tString
  LFloat {} -> Type
tFloat

type Solve b m =
  ( MonadBase b m
  , MonadError TypeErrorA m
  , MonadState ElabState m -- We highjack the ElabState type for the genTypeVar function.
  , MonadReader (ElabEnv b) m
  )

-- | Create an instance of a type. (The type serves as a template for specialized types)
--
-- Let polymorphism gives us the ability to use a generic function in more contexts.
--
-- For example, @id@ is a function that can work for @x@ of any type. But our algorithm
-- collects constraints globally, including that:
--
-- @
-- id = \x -> x
--
-- one = id 1          -- constrain that the type of id __is equal to__ the type [Int] -> tN
--
-- hello = id "hello"  -- constrain that the type of id __is equal to__ the type [String] -> tM
-- @
--
-- We need to invent a new constraint that will define the relationship between the type of id
-- and the arguments passing to it as an __Instance of__ relationship.
--
-- @InstanceOf t1 t2@ relationship means that @t1@ is an @instantiation@ of @t2@
-- (or a special case if you will).
-- What we'll do is copy the type of @t2@, generate new type variables in place of all type variables
-- inside of it, and then say that this new type @t3@ has an equality relationship with @t1@.
--
-- It's important to solve the equality constraints for each function before solving the InstanceOf
-- constraints, so that when we instantiate we already have the final type of the function.
--
-- We will highjack the @Ord@ instance deriving (constructors defined later are bigger)
-- and the fact that @Set@ is ordered to accomplish that.
instantiate :: (Solve b m) => Text -> [TypeVar] -> Type -> m Type
instantiate :: forall (b :: * -> *) (m :: * -> *).
Solve b m =>
Var -> [TypeVar] -> Type -> m Type
instantiate Var
prefix [TypeVar]
vars Type
typ = do
  Map TypeVar TypeVar
env <-
    [(TypeVar, TypeVar)] -> Map TypeVar TypeVar
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(TypeVar, TypeVar)] -> Map TypeVar TypeVar)
-> m [(TypeVar, TypeVar)] -> m (Map TypeVar TypeVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeVar -> m (TypeVar, TypeVar))
-> [TypeVar] -> m [(TypeVar, TypeVar)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\TypeVar
v -> (,) TypeVar
v (TypeVar -> (TypeVar, TypeVar))
-> m TypeVar -> m (TypeVar, TypeVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> m TypeVar
forall (m :: * -> *). MonadState ElabState m => Var -> m TypeVar
genTypeVar Var
prefix) [TypeVar]
vars
  Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ ((Type -> Type) -> Type -> Type) -> Type -> (Type -> Type) -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Type -> Type) -> Type -> Type
forall on. Uniplate on => (on -> on) -> on -> on
U.transform Type
typ ((Type -> Type) -> Type) -> (Type -> Type) -> Type
forall a b. (a -> b) -> a -> b
$ \case
    TypeVar TypeVar
tv
      | Just TypeVar
tv' <- TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
env ->
          TypeVar -> Type
TypeVar TypeVar
tv'
    TypeRecExt [(Label, Type)]
r TypeVar
tv
      | Just TypeVar
tv' <- TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
env ->
          [(Label, Type)] -> TypeVar -> Type
TypeRecExt [(Label, Type)]
r TypeVar
tv'
    TypePolyVariantLB [(Constr, Type)]
variants TypeVar
tv
      | Just TypeVar
tv' <- TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
env ->
          [(Constr, Type)] -> TypeVar -> Type
TypePolyVariantLB [(Constr, Type)]
variants TypeVar
tv'
    TypePolyVariantUB TypeVar
tv [(Constr, Type)]
variants
      | Just TypeVar
tv' <- TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
env ->
          TypeVar -> [(Constr, Type)] -> Type
TypePolyVariantUB TypeVar
tv' [(Constr, Type)]
variants
    t :: Type
t@TypeVar {} -> Type
t
    t :: Type
t@TypeCon {} -> Type
t
    t :: Type
t@TypeApp {} -> Type
t
    t :: Type
t@TypeRec {} -> Type
t
    t :: Type
t@TypeRecExt {} -> Type
t
    t :: Type
t@TypeVariant {} -> Type
t
    t :: Type
t@TypePolyVariantLB {} -> Type
t
    t :: Type
t@TypePolyVariantUB {} -> Type
t
    t :: Type
t@TypeScheme {} -> Type
t