{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Safe #-}

module Embedder.Atmo.FlattenMonadTrans where

import safe Embedder.Atmo.Syntax as A
    ( Ty(..), TyBuiltin (..),)
import ReWire.Orphans ()
import ReWire.Annotation (Annote)

import Embedder.Atmo.Types (nilTy, pairTy)

-- | transMonadT
--   Convert `StateT s (StateT t (StateT u Identity)) a`
--       to  `State (s,(t,(u,()))) a`
--   Convert `ReacT i o m a`
--       to  `Re i o s a`
--       where s = mkState an (getStates m)
transMonadT :: Ty -> Ty
transMonadT :: Ty -> Ty
transMonadT = \ case
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyReacT) [Ty
i,Ty
o,m :: Ty
m@(TyVar {}),Ty
a] -> 
            Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyRe) [Ty
i,Ty
o,Ty
m,Ty
a]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyReacT) [Ty
i,Ty
o,Ty
m,Ty
a] -> 
            let s :: Ty
s = Annote -> [Ty] -> Ty
mkState Annote
an (Ty -> [Ty]
getStates Ty
m)
            in Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyRe) [Ty
i,Ty
o,Ty
s,Ty
a]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyStateT) [Ty
s,TyVar {},Ty
a] ->
            Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyState) [Annote -> [Ty] -> Ty
TyTuple Annote
an [Ty
s,Annote -> [Ty] -> Ty
TyTuple Annote
an []],Ty
a]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyStateT) [Ty
s,Ty
m,Ty
a] -> 
            let s' :: Ty
s' = Annote -> Ty -> Ty -> Ty
pairTy Annote
an Ty
s (Annote -> [Ty] -> Ty
mkState Annote
an (Ty -> [Ty]
getStates Ty
m))
            in Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyState) [Ty
s',Ty
a]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyIdentity) [Ty
a] ->
            Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyState) [Annote -> [Ty] -> Ty
TyTuple Annote
an [],Ty
a]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyStateDev) [Ty
i,Ty
o,m :: Ty
m@(TyVar {})] ->
            Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyStateDev) [Ty
i,Ty
o,Ty
m]
      TyApp Annote
an (TyBuiltin Annote
_ TyBuiltin
TyStateDev) [Ty
i,Ty
o,Ty
m] ->
            let s :: Ty
s = Annote -> [Ty] -> Ty
mkState Annote
an (Ty -> [Ty]
getStates Ty
m)
            in Annote -> Ty -> [Ty] -> Ty
TyApp Annote
an (Annote -> TyBuiltin -> Ty
TyBuiltin Annote
an TyBuiltin
TyStateDev) [Ty
i,Ty
o,Ty
s]
      Ty
t -> Ty
t

getStates :: Ty -> [Ty]
getStates :: Ty -> [Ty]
getStates = \ case
      TyApp Annote
_ (TyBuiltin Annote
_ TyBuiltin
TyStateT) [Ty
s,Ty
m] -> Ty
s Ty -> [Ty] -> [Ty]
forall a. a -> [a] -> [a]
: Ty -> [Ty]
getStates Ty
m
      (TyBuiltin Annote
_ TyBuiltin
TyIdentity) -> []
      Ty
_ -> []

mkState :: Annote -> [Ty] -> Ty
mkState :: Annote -> [Ty] -> Ty
mkState Annote
_a [] = Ty
nilTy
mkState Annote
a (Ty
t:[Ty]
ts) = Annote -> Ty -> Ty -> Ty
pairTy Annote
a Ty
t (Ty -> Ty) -> Ty -> Ty
forall a b. (a -> b) -> a -> b
$ Annote -> [Ty] -> Ty
mkState Annote
a [Ty]
ts


-- We also need to disambiguate lift
-- lift of return type ReacT is liftR
-- lift of return type StateT is liftS