{-# LANGUAGE TupleSections #-}

module IR.Constraint.Constrain.Program where

import qualified Common.Identifiers as Ident
import Data.Foldable (foldrM)
import qualified Data.Map.Strict as Map
import qualified IR.Constraint.Canonical as Can
import qualified IR.Constraint.Constrain.Expression as Expr
import IR.Constraint.Monad (TC)
import IR.Constraint.Type
import qualified IR.IR as I


constrain :: I.Program (Can.Annotations, Variable) -> TC Constraint
constrain :: Program (Annotations, Variable) -> TC Constraint
constrain Program (Annotations, Variable)
prog = do
  [(TConId, TypeDef)] -> Constraint -> TC Constraint
constrainTypeDefs (Program (Annotations, Variable) -> [(TConId, TypeDef)]
forall t. Program t -> [(TConId, TypeDef)]
I.typeDefs Program (Annotations, Variable)
prog)
    (Constraint -> TC Constraint) -> TC Constraint -> TC Constraint
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [(VarId, Type)] -> Constraint -> TC Constraint
constrainExternDecls (Program (Annotations, Variable) -> [(VarId, Type)]
forall t. Program t -> [(VarId, Type)]
I.externDecls Program (Annotations, Variable)
prog)
    (Constraint -> TC Constraint) -> TC Constraint -> TC Constraint
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [BinderDef] -> Constraint -> TC Constraint
Expr.constrainBinderDefs (Program (Annotations, Variable) -> [BinderDef]
forall t. Program t -> [(Binder t, Expr t)]
I.programDefs Program (Annotations, Variable)
prog) Constraint
CTrue


-- | EXTERN DECLS

-- TODO: do I still need this? depends on whether these are only used in fficall prim
constrainExternDecls
  :: [(Ident.VarId, Can.Type)] -> Constraint -> TC Constraint
constrainExternDecls :: [(VarId, Type)] -> Constraint -> TC Constraint
constrainExternDecls [(VarId, Type)]
decls Constraint
finalConstraint =
  ((VarId, Type) -> Constraint -> TC Constraint)
-> Constraint -> [(VarId, Type)] -> TC Constraint
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM (VarId, Type) -> Constraint -> TC Constraint
constrainExternDecl Constraint
finalConstraint [(VarId, Type)]
decls


constrainExternDecl :: (Ident.VarId, Can.Type) -> Constraint -> TC Constraint
constrainExternDecl :: (VarId, Type) -> Constraint -> TC Constraint
constrainExternDecl (VarId
varId, Type
canTyp) Constraint
finalConstraint = do
  Identifier -> Scheme -> Constraint -> TC Constraint
constrainDeclaration
    (VarId -> Identifier
forall a b. (Identifiable a, Identifiable b) => a -> b
Ident.fromId VarId
varId)
    (Type -> Scheme
Can.schemeOf Type
canTyp)
    Constraint
finalConstraint


-- | TYPE DEFS
constrainTypeDefs :: [(Ident.TConId, I.TypeDef)] -> Constraint -> TC Constraint
constrainTypeDefs :: [(TConId, TypeDef)] -> Constraint -> TC Constraint
constrainTypeDefs [(TConId, TypeDef)]
tdefs Constraint
finalConstraint =
  ((TConId, TypeDef) -> Constraint -> TC Constraint)
-> Constraint -> [(TConId, TypeDef)] -> TC Constraint
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
    ( \(TConId
tcon, TypeDef
typeDef) Constraint
innerConstraint ->
        TConId -> TypeDef -> Constraint -> TC Constraint
constrainTypeDef TConId
tcon TypeDef
typeDef Constraint
innerConstraint
    )
    Constraint
finalConstraint
    [(TConId, TypeDef)]
tdefs


constrainTypeDef :: Ident.TConId -> I.TypeDef -> Constraint -> TC Constraint
constrainTypeDef :: TConId -> TypeDef -> Constraint -> TC Constraint
constrainTypeDef TConId
tcon (I.TypeDef [(DConId, TypeVariant)]
variants [TVarId]
args) Constraint
finalConstraint =
  ((DConId, TypeVariant) -> Constraint -> TC Constraint)
-> Constraint -> [(DConId, TypeVariant)] -> TC Constraint
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
    ( \(DConId
dcon, TypeVariant
variant) Constraint
innerConstraint ->
        TConId
-> [TVarId] -> DConId -> TypeVariant -> Constraint -> TC Constraint
constrainVariant TConId
tcon [TVarId]
args DConId
dcon TypeVariant
variant Constraint
innerConstraint
    )
    Constraint
finalConstraint
    [(DConId, TypeVariant)]
variants


constrainVariant
  :: Ident.TConId
  -> [Ident.TVarId]
  -> Ident.DConId
  -> I.TypeVariant
  -> Constraint
  -> TC Constraint
constrainVariant :: TConId
-> [TVarId] -> DConId -> TypeVariant -> Constraint -> TC Constraint
constrainVariant TConId
tcon [TVarId]
targs DConId
dcon (I.VariantUnnamed [Type]
dargs) Constraint
finalConstraint = do
  let typ :: Type
typ = ([Type], Type) -> Type
Can.foldArrow ([Type]
dargs, TConId -> [Type] -> Type
Can.TCon TConId
tcon ((TVarId -> Type) -> [TVarId] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TVarId -> Type
Can.TVar [TVarId]
targs))
      scheme :: Scheme
scheme = FreeVars -> Type -> Scheme
Can.Forall ([(TVarId, ())] -> FreeVars
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((TVarId -> (TVarId, ())) -> [TVarId] -> [(TVarId, ())]
forall a b. (a -> b) -> [a] -> [b]
map (,()) [TVarId]
targs)) Type
typ
  Identifier -> Scheme -> Constraint -> TC Constraint
constrainDeclaration (DConId -> Identifier
forall a b. (Identifiable a, Identifiable b) => a -> b
Ident.fromId DConId
dcon) Scheme
scheme Constraint
finalConstraint
constrainVariant TConId
_ [TVarId]
_ DConId
_ (I.VariantNamed [(VarId, Type)]
_) Constraint
_ =
  [Char] -> TC Constraint
forall a. HasCallStack => [Char] -> a
error [Char]
"No support for named variants yet"


-- | HELPER: CONSTRAIN NAME TO TYPE
constrainDeclaration
  :: Ident.Identifier -> Can.Scheme -> Constraint -> TC Constraint
constrainDeclaration :: Identifier -> Scheme -> Constraint -> TC Constraint
constrainDeclaration Identifier
name Scheme
scheme Constraint
finalConstraint = do
  Variable
v <- TC Variable
mkFlexVar
  let t :: Type
t = Variable -> Type
TVarN Variable
v
  let header :: Map Identifier Type
header = Identifier -> Type -> Map Identifier Type
forall k a. k -> a -> Map k a
Map.singleton (Identifier -> Identifier
forall a b. (Identifiable a, Identifiable b) => a -> b
Ident.fromId Identifier
name) Type
t
  Constraint -> TC Constraint
forall (m :: * -> *) a. Monad m => a -> m a
return (Constraint -> TC Constraint) -> Constraint -> TC Constraint
forall a b. (a -> b) -> a -> b
$ [Variable]
-> [Variable]
-> Map Identifier Type
-> Constraint
-> Constraint
-> Constraint
CLet [] [Variable
v] Map Identifier Type
header (Scheme -> Type -> Constraint
CForeign Scheme
scheme Type
t) Constraint
finalConstraint