{-# LANGUAGE OverloadedStrings #-}

module IR.Constraint.Constrain.Annotation 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 IR.Constraint.Monad (
  TC,
  freshAnnVar,
  getKind,
  throwError,
 )
import IR.Constraint.Type as Type


-- | CONSTRAIN ANNOTATION
type RigidMap = Map.Map Ident.TVarId Variable


data State = State
  { State -> RigidMap
_rigidMap :: RigidMap
  , State -> [Variable]
_flex :: [Variable]
  }


withAnnotations :: [Can.Annotation] -> Type -> (Type -> TC Constraint) -> TC Constraint
withAnnotations :: [Annotation] -> Type -> (Type -> TC Constraint) -> TC Constraint
withAnnotations [] Type
expected Type -> TC Constraint
m = Type -> TC Constraint
m Type
expected
withAnnotations (Can.AnnDCon DConId
_ [Annotation]
_ : [Annotation]
_) Type
_ Type -> TC Constraint
_ = [Char] -> TC Constraint
forall a. HasCallStack => [Char] -> a
error [Char]
"We should remove this, AnnDCon was a bad idea"
withAnnotations (Annotation
ann : [Annotation]
anns) Type
expected Type -> TC Constraint
m = do
  Identifier
dummyId <- VarId -> Identifier
forall a b. (Identifiable a, Identifiable b) => a -> b
Ident.fromId (VarId -> Identifier)
-> StateT TCState (ExceptT Error (WriterT (Doc [Char]) IO)) VarId
-> StateT
     TCState (ExceptT Error (WriterT (Doc [Char]) IO)) Identifier
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT TCState (ExceptT Error (WriterT (Doc [Char]) IO)) VarId
freshAnnVar
  (State RigidMap
rigidMap [Variable]
flexs, Type
expected') <- Annotation -> TC (State, Type)
add Annotation
ann
  Constraint
innerConstraint <- [Annotation] -> Type -> (Type -> TC Constraint) -> TC Constraint
withAnnotations [Annotation]
anns Type
expected' Type -> TC Constraint
m
  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
$
    CLet :: [Variable]
-> [Variable]
-> Map Identifier Type
-> Constraint
-> Constraint
-> Constraint
CLet
      { _rigidVars :: [Variable]
_rigidVars = RigidMap -> [Variable]
forall k a. Map k a -> [a]
Map.elems RigidMap
rigidMap
      , _flexVars :: [Variable]
_flexVars = [Variable]
flexs
      , _header :: Map Identifier Type
_header = Identifier -> Type -> Map Identifier Type
forall k a. k -> a -> Map k a
Map.singleton Identifier
dummyId Type
expected'
      , _headerCon :: Constraint
_headerCon = Constraint
innerConstraint
      , _bodyCon :: Constraint
_bodyCon = Identifier -> Type -> Constraint
CLocal Identifier
dummyId Type
expected
      }


add :: Can.Annotation -> TC (State, Type)
add :: Annotation -> TC (State, Type)
add Annotation
ann = do
  let canType :: Type
canType = Annotation -> Type
Can.annToType Annotation
ann
      state :: State
state = State
emptyState
  Type -> State -> TC (State, Type)
addTypeWithHoles Type
canType State
state


addTypeWithHoles :: Can.Type -> State -> TC (State, Type)
addTypeWithHoles :: Type -> State -> TC (State, Type)
addTypeWithHoles Type
canType state :: State
state@(State RigidMap
rigidMap [Variable]
flexes) = case Type
canType of
  Can.TVar TVarId
"_" -> do
    Variable
var <- TC Variable
mkFlexVar
    (State, Type) -> TC (State, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (RigidMap -> [Variable] -> State
State RigidMap
rigidMap (Variable
var Variable -> [Variable] -> [Variable]
forall a. a -> [a] -> [a]
: [Variable]
flexes), Variable -> Type
TVarN Variable
var)
  Can.TVar TVarId
name -> case TVarId -> RigidMap -> Maybe Variable
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TVarId
name RigidMap
rigidMap of
    Just Variable
var -> (State, Type) -> TC (State, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (RigidMap -> [Variable] -> State
State RigidMap
rigidMap [Variable]
flexes, Variable -> Type
TVarN Variable
var)
    Maybe Variable
Nothing -> do
      Variable
var <- TC Variable
mkRigidVar
      (State, Type) -> TC (State, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (RigidMap -> [Variable] -> State
State (TVarId -> Variable -> RigidMap -> RigidMap
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TVarId
name Variable
var RigidMap
rigidMap) [Variable]
flexes, Variable -> Type
TVarN Variable
var)
  Can.TCon TConId
tcon [Type]
args -> do
    Maybe Kind
maybeKind <- TConId -> TC (Maybe Kind)
getKind TConId
tcon
    case Maybe Kind
maybeKind of
      Just Kind
kind ->
        if [Type] -> Kind
forall (t :: * -> *) a. Foldable t => t a -> Kind
length [Type]
args Kind -> Kind -> Bool
forall a. Eq a => a -> a -> Bool
== Kind
kind
          then do
            (State
state', [Type]
ts) <-
              (Type
 -> (State, [Type])
 -> StateT
      TCState (ExceptT Error (WriterT (Doc [Char]) IO)) (State, [Type]))
-> (State, [Type])
-> [Type]
-> StateT
     TCState (ExceptT Error (WriterT (Doc [Char]) IO)) (State, [Type])
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
                ( \Type
a (State
st, [Type]
ts) -> do
                    (State
st', Type
t) <- Type -> State -> TC (State, Type)
addTypeWithHoles Type
a State
st
                    (State, [Type])
-> StateT
     TCState (ExceptT Error (WriterT (Doc [Char]) IO)) (State, [Type])
forall (m :: * -> *) a. Monad m => a -> m a
return (State
st', Type
t Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
ts)
                )
                (State
state, [])
                [Type]
args
            (State, Type) -> TC (State, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (State
state', TConId -> [Type] -> Type
TConN TConId
tcon [Type]
ts)
          else [Char] -> TC (State, Type)
forall a. [Char] -> TC a
throwError ([Char] -> TC (State, Type)) -> [Char] -> TC (State, Type)
forall a b. (a -> b) -> a -> b
$ [Char]
"Wrong arity in type annotation for " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TConId -> [Char]
forall a. Show a => a -> [Char]
show TConId
tcon
      Maybe Kind
Nothing ->
        [Char] -> TC (State, Type)
forall a. [Char] -> TC a
throwError ([Char] -> TC (State, Type)) -> [Char] -> TC (State, Type)
forall a b. (a -> b) -> a -> b
$
          [Char]
"In type annotation, type constructore does not exists: "
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TConId -> [Char]
forall a. Show a => a -> [Char]
show TConId
tcon


emptyState :: State
emptyState :: State
emptyState = RigidMap -> [Variable] -> State
State RigidMap
forall k a. Map k a
Map.empty []