{-# LANGUAGE BlockArguments #-}

module IR.Constraint.Constrain.Pattern where

import qualified Common.Identifiers as Ident
import Control.Monad (
  foldM,
  unless,
 )
import Data.Bifunctor (second)
import qualified Data.Map.Strict as Map
import qualified IR.Constraint.Canonical as Can
import qualified IR.Constraint.Constrain.Annotation as Ann
import qualified IR.Constraint.Instantiate as Inst
import IR.Constraint.Monad (
  TC,
  getDConInfo,
  throwError,
 )
import IR.Constraint.Type as Type
import qualified IR.IR as I


-- | CONSTRAIN ALT
type Header = Map.Map Ident.Identifier Type


data State = State
  { State -> Header
_headers :: Header
  , State -> [Variable]
_vars :: [Variable]
  , State -> [Constraint]
_revCons :: [Constraint]
  }


add :: I.Alt Attachment -> Type -> State -> TC State
add :: Alt Attachment -> Type -> State -> TC State
add Alt Attachment
alt Type
expected State
state = case Alt Attachment
alt of
  I.AltBinder Binder Attachment
binder -> do
    VarId
var <- Binder Attachment -> TC VarId
forall t. Binder t -> TC VarId
Type.binderToVarId Binder Attachment
binder
    State
state' <- Binder Attachment -> Type -> State -> TC State
addBinder Binder Attachment
binder Type
expected State
state
    State -> TC State
forall (m :: * -> *) a. Monad m => a -> m a
return (State -> TC State) -> State -> TC State
forall a b. (a -> b) -> a -> b
$ Identifier -> Type -> State -> State
addToHeaders (VarId -> Identifier
forall a b. (Identifiable a, Identifiable b) => a -> b
Ident.fromId VarId
var) Type
expected State
state'
  I.AltLit Literal
lit Attachment
_ -> State -> TC State
forall (m :: * -> *) a. Monad m => a -> m a
return (State -> TC State) -> State -> TC State
forall a b. (a -> b) -> a -> b
$ Literal -> Type -> State -> State
addLit Literal
lit Type
expected State
state
  I.AltData DConId
dcon [Alt Attachment]
bs Attachment
_ -> do
    Maybe DConInfo
maybeInfo <- DConId -> TC (Maybe DConInfo)
getDConInfo DConId
dcon
    case Maybe DConInfo
maybeInfo of
      Just (DConId
_, TConId
typeName, [TVarId]
typeVarNames, [Type]
argTypes) ->
        TConId
-> [TVarId]
-> DConId
-> [Type]
-> [Alt Attachment]
-> Type
-> State
-> TC State
addData TConId
typeName [TVarId]
typeVarNames DConId
dcon [Type]
argTypes [Alt Attachment]
bs Type
expected State
state
      Maybe DConInfo
Nothing ->
        -- this should already be found in scope checking
        String -> TC State
forall a. String -> TC a
throwError (String -> TC State) -> String -> TC State
forall a b. (a -> b) -> a -> b
$ String
"Pattern: data constructor does not exist - " String -> String -> String
forall a. [a] -> [a] -> [a]
++ DConId -> String
forall a. Show a => a -> String
show DConId
dcon


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


addBinder :: I.Binder Attachment -> Type -> State -> TC State
addBinder :: Binder Attachment -> Type -> State -> TC State
addBinder Binder Attachment
binder Type
expected (State Header
headers [Variable]
vars [Constraint]
revCons) = do
  let ([Annotation]
anns, Variable
var) = Binder Attachment -> ([Annotation], Variable)
forall (c :: * -> *).
Carrier c =>
c Attachment -> ([Annotation], Variable)
uncarryAttachment Binder Attachment
binder
  Constraint
cons <-
    [Variable] -> Constraint -> Constraint
exists [Variable
var] (Constraint -> Constraint)
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Annotation]
-> Type
-> (Type
    -> StateT
         TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint)
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint
Ann.withAnnotations [Annotation]
anns (Variable -> Type
TVarN Variable
var) \Type
varExpected ->
      Constraint
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint
forall (m :: * -> *) a. Monad m => a -> m a
return (Constraint
 -> StateT
      TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint)
-> Constraint
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Constraint
forall a b. (a -> b) -> a -> b
$ Type -> Type -> Constraint
CEqual Type
varExpected Type
expected
  State -> TC State
forall (m :: * -> *) a. Monad m => a -> m a
return (State -> TC State) -> State -> TC State
forall a b. (a -> b) -> a -> b
$ Header -> [Variable] -> [Constraint] -> State
State Header
headers ([Variable]
vars [Variable] -> [Variable] -> [Variable]
forall a. [a] -> [a] -> [a]
++ [Variable
var]) (Constraint
cons Constraint -> [Constraint] -> [Constraint]
forall a. a -> [a] -> [a]
: [Constraint]
revCons)


addToHeaders :: Ident.Identifier -> Type -> State -> State
addToHeaders :: Identifier -> Type -> State -> State
addToHeaders Identifier
name Type
tipe (State Header
headers [Variable]
vars [Constraint]
revCons) =
  let newHeaders :: Header
newHeaders = Identifier -> Type -> Header -> Header
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Identifier
name Type
tipe Header
headers
   in Header -> [Variable] -> [Constraint] -> State
State Header
newHeaders [Variable]
vars [Constraint]
revCons


addData
  :: Ident.TConId
  -> [Ident.TVarId]
  -> Ident.DConId
  -> [Can.Type]
  -> [I.Alt Attachment]
  -> Type
  -> State
  -> TC State
addData :: TConId
-> [TVarId]
-> DConId
-> [Type]
-> [Alt Attachment]
-> Type
-> State
-> TC State
addData TConId
typeName [TVarId]
typeVarNames DConId
ctorName [Type]
ctorArgTypes [Alt Attachment]
as Type
expected State
state = do
  Bool
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ctorArgTypes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Alt Attachment] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Alt Attachment]
as) (StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
 -> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ())
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
forall a b. (a -> b) -> a -> b
$
    String
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
forall a. String -> TC a
throwError (String
 -> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ())
-> String
-> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) ()
forall a b. (a -> b) -> a -> b
$
      String
"Pattern: wrong number of argument - "
        String -> String -> String
forall a. [a] -> [a] -> [a]
++ DConId -> String
forall a. Show a => a -> String
show DConId
ctorName
  [(TVarId, Variable)]
varPairs <- (TVarId
 -> StateT
      TCState
      (ExceptT Error (WriterT (Doc String) IO))
      (TVarId, Variable))
-> [TVarId]
-> StateT
     TCState
     (ExceptT Error (WriterT (Doc String) IO))
     [(TVarId, Variable)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\TVarId
var -> (,) TVarId
var (Variable -> (TVarId, Variable))
-> StateT
     TCState (ExceptT Error (WriterT (Doc String) IO)) Variable
-> StateT
     TCState
     (ExceptT Error (WriterT (Doc String) IO))
     (TVarId, Variable)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT TCState (ExceptT Error (WriterT (Doc String) IO)) Variable
mkFlexVar) [TVarId]
typeVarNames
  let typePairs :: [(TVarId, Type)]
typePairs = ((TVarId, Variable) -> (TVarId, Type))
-> [(TVarId, Variable)] -> [(TVarId, Type)]
forall a b. (a -> b) -> [a] -> [b]
map ((Variable -> Type) -> (TVarId, Variable) -> (TVarId, Type)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Variable -> Type
TVarN) [(TVarId, Variable)]
varPairs
  let freeVarDict :: Map TVarId Type
freeVarDict = [(TVarId, Type)] -> Map TVarId Type
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(TVarId, Type)]
typePairs
  (State Header
headers [Variable]
vars [Constraint]
revCons) <-
    (State -> (Alt Attachment, Type) -> TC State)
-> State -> [(Alt Attachment, Type)] -> TC State
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
      (\State
st (Alt Attachment
a, Type
aCanType) -> Map TVarId Type -> Type -> Alt Attachment -> State -> TC State
addDataArg Map TVarId Type
freeVarDict Type
aCanType Alt Attachment
a State
st)
      State
state
      ([Alt Attachment] -> [Type] -> [(Alt Attachment, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Alt Attachment]
as [Type]
ctorArgTypes)
  let ctorType :: Type
ctorType = TConId -> [Type] -> Type
TConN TConId
typeName (((TVarId, Type) -> Type) -> [(TVarId, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TVarId, Type) -> Type
forall a b. (a, b) -> b
snd [(TVarId, Type)]
typePairs)
  let ctorCon :: Constraint
ctorCon = Type -> Type -> Constraint
CPattern Type
ctorType Type
expected
  State -> TC State
forall (m :: * -> *) a. Monad m => a -> m a
return (State -> TC State) -> State -> TC State
forall a b. (a -> b) -> a -> b
$
    State :: Header -> [Variable] -> [Constraint] -> State
State
      { _headers :: Header
_headers = Header
headers
      , _vars :: [Variable]
_vars = ((TVarId, Variable) -> Variable)
-> [(TVarId, Variable)] -> [Variable]
forall a b. (a -> b) -> [a] -> [b]
map (TVarId, Variable) -> Variable
forall a b. (a, b) -> b
snd [(TVarId, Variable)]
varPairs [Variable] -> [Variable] -> [Variable]
forall a. [a] -> [a] -> [a]
++ [Variable]
vars
      , _revCons :: [Constraint]
_revCons = Constraint
ctorCon Constraint -> [Constraint] -> [Constraint]
forall a. a -> [a] -> [a]
: [Constraint]
revCons
      }


addDataArg
  :: Map.Map Ident.TVarId Type -> Can.Type -> I.Alt Attachment -> State -> TC State
addDataArg :: Map TVarId Type -> Type -> Alt Attachment -> State -> TC State
addDataArg Map TVarId Type
freeVarDict Type
canType Alt Attachment
arg State
state = do
  Type
tipe <- Map TVarId Type -> Type -> TC Type
Inst.fromScheme Map TVarId Type
freeVarDict Type
canType
  Alt Attachment -> Type -> State -> TC State
add Alt Attachment
arg Type
tipe State
state


addLit :: I.Literal -> Type -> State -> State
addLit :: Literal -> Type -> State -> State
addLit Literal
lit Type
expected State
state =
  let litCon :: Constraint
litCon = case Literal
lit of
        I.LitIntegral Integer
_ -> Type -> Type -> Constraint
CPattern Type
Type.i32 Type
expected
        Literal
I.LitEvent -> Type -> Type -> Constraint
CPattern Type
Type.unit Type
expected
   in State
state{_revCons :: [Constraint]
_revCons = Constraint
litCon Constraint -> [Constraint] -> [Constraint]
forall a. a -> [a] -> [a]
: State -> [Constraint]
_revCons State
state}