{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module TypedSession.State.Utils where

import Control.Carrier.Fresh.Strict
import Control.Carrier.State.Strict
import Control.Effect.Writer
import Control.Monad
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import qualified Data.List as L
import Data.Maybe (fromJust, fromMaybe)
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import qualified TypedSession.State.Constraint as C
import TypedSession.State.Type
import Prelude hiding (traverse)

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

restoreWrapper
  :: forall s sig m a
   . (Has (State s) sig m) => m a -> m a
restoreWrapper :: forall s (sig :: (* -> *) -> * -> *) (m :: * -> *) a.
Has (State s) sig m =>
m a -> m a
restoreWrapper m a
m = do
  st <- forall s (sig :: (* -> *) -> * -> *) (m :: * -> *).
Has (State s) sig m =>
m s
get @s
  a <- m
  put st
  pure a

getFirstMsgInfo :: Protocol eta r bst -> Maybe (r, r)
getFirstMsgInfo :: forall eta r bst. Protocol eta r bst -> Maybe (r, r)
getFirstMsgInfo = \case
  MsgOrLabel eta r
msgOrLabel :> Protocol eta r bst
prots -> case MsgOrLabel eta r
msgOrLabel of
    Msg XMsg eta
_ String
_ [[String]]
_ r
from r
to -> (r, r) -> Maybe (r, r)
forall a. a -> Maybe a
Just (r
from, r
to)
    MsgOrLabel eta r
_ -> Protocol eta r bst -> Maybe (r, r)
forall eta r bst. Protocol eta r bst -> Maybe (r, r)
getFirstMsgInfo Protocol eta r bst
prots
  Protocol eta r bst
_ -> Maybe (r, r)
forall a. Maybe a
Nothing

getAllMsgInfo :: Protocol eta r bst -> [(r, r)]
getAllMsgInfo :: forall eta r bst. Protocol eta r bst -> [(r, r)]
getAllMsgInfo = \case
  MsgOrLabel eta r
msgOrLabel :> Protocol eta r bst
prots -> case MsgOrLabel eta r
msgOrLabel of
    Msg XMsg eta
_ String
_ [[String]]
_ r
from r
to -> (r
from, r
to) (r, r) -> [(r, r)] -> [(r, r)]
forall a. a -> [a] -> [a]
: Protocol eta r bst -> [(r, r)]
forall eta r bst. Protocol eta r bst -> [(r, r)]
getAllMsgInfo Protocol eta r bst
prots
    MsgOrLabel eta r
_ -> Protocol eta r bst -> [(r, r)]
forall eta r bst. Protocol eta r bst -> [(r, r)]
getAllMsgInfo Protocol eta r bst
prots
  Branch XBranch eta
_ r
_ String
_ [BranchSt eta r bst]
ls -> (BranchSt eta r bst -> [(r, r)])
-> [BranchSt eta r bst] -> [(r, r)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(BranchSt XBranchSt eta
_ bst
_ [[String]]
_ Protocol eta r bst
prots) -> Protocol eta r bst -> [(r, r)]
forall eta r bst. Protocol eta r bst -> [(r, r)]
getAllMsgInfo Protocol eta r bst
prots) [BranchSt eta r bst]
ls
  Goto XGoto eta
_ Key
_ -> []
  Terminal XTerminal eta
_ -> []

tellSeq :: (Has (Writer (Seq a)) sig m) => [a] -> m ()
tellSeq :: forall a (sig :: (* -> *) -> * -> *) (m :: * -> *).
Has (Writer (Seq a)) sig m =>
[a] -> m ()
tellSeq [a]
ls = Seq a -> m ()
forall w (sig :: (* -> *) -> * -> *) (m :: * -> *).
Has (Writer w) sig m =>
w -> m ()
tell ([a] -> Seq a
forall a. [a] -> Seq a
Seq.fromList [a]
ls)

compressSubMap :: C.SubMap -> (C.SubMap, (Int, Int))
compressSubMap :: SubMap -> (SubMap, (Key, Key))
compressSubMap SubMap
sbm' =
  let (Key
minKey, Key
maxKey) = ((Key, Key) -> Key
forall a b. (a, b) -> a
fst ((Key, Key) -> Key) -> (Key, Key) -> Key
forall a b. (a -> b) -> a -> b
$ SubMap -> (Key, Key)
forall a. IntMap a -> (Key, a)
IntMap.findMin SubMap
sbm', (Key, Key) -> Key
forall a b. (a, b) -> a
fst ((Key, Key) -> Key) -> (Key, Key) -> Key
forall a b. (a -> b) -> a -> b
$ SubMap -> (Key, Key)
forall a. IntMap a -> (Key, a)
IntMap.findMax SubMap
sbm')
      list :: [Key]
list = [Key
minKey .. Key
maxKey]
      ([Key]
keys, [Key]
vals) = ([Key]
list, (Key -> Key) -> [Key] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Key
k -> Key -> Maybe Key -> Key
forall a. a -> Maybe a -> a
fromMaybe Key
k (Maybe Key -> Key) -> Maybe Key -> Key
forall a b. (a -> b) -> a -> b
$ Key -> SubMap -> Maybe Key
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key
k SubMap
sbm') [Key]
list)
      minVal :: Key
minVal = [Key] -> Key
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Key]
vals
      tmap :: SubMap
tmap = [(Key, Key)] -> SubMap
forall a. [(Key, a)] -> IntMap a
IntMap.fromList ([(Key, Key)] -> SubMap) -> [(Key, Key)] -> SubMap
forall a b. (a -> b) -> a -> b
$ [Key] -> [Key] -> [(Key, Key)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Key] -> [Key]
forall a. Eq a => [a] -> [a]
L.nub ([Key] -> [Key]) -> [Key] -> [Key]
forall a b. (a -> b) -> a -> b
$ [Key] -> [Key]
forall a. Ord a => [a] -> [a]
L.sort [Key]
vals) [Key
minVal, Key
minVal Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
1 ..]
      vals' :: [Key]
vals' = (Key -> Key) -> [Key] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Key
k -> Maybe Key -> Key
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Key -> Key) -> Maybe Key -> Key
forall a b. (a -> b) -> a -> b
$ Key -> SubMap -> Maybe Key
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key
k SubMap
tmap) [Key]
vals
   in ([(Key, Key)] -> SubMap
forall a. [(Key, a)] -> IntMap a
IntMap.fromList ([(Key, Key)] -> SubMap) -> [(Key, Key)] -> SubMap
forall a b. (a -> b) -> a -> b
$ [Key] -> [Key] -> [(Key, Key)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Key]
keys [Key]
vals', (-Key
1, [Key] -> Key
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Key]
vals'))

replaceList :: C.SubMap -> [Int] -> [Int]
replaceList :: SubMap -> [Key] -> [Key]
replaceList SubMap
sbm [Key]
ls = (Key -> Key) -> [Key] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Key
k -> Key -> Maybe Key -> Key
forall a. a -> Maybe a -> a
fromMaybe Key
k (Maybe Key -> Key) -> Maybe Key -> Key
forall a b. (a -> b) -> a -> b
$ Key -> SubMap -> Maybe Key
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key
k SubMap
sbm) [Key]
ls

replaceVal :: IntMap Int -> Int -> Int
replaceVal :: SubMap -> Key -> Key
replaceVal SubMap
sbm Key
k = Key -> Maybe Key -> Key
forall a. a -> Maybe a -> a
fromMaybe (String -> Key
forall a. HasCallStack => String -> a
error String
internalError) (Maybe Key -> Key) -> Maybe Key -> Key
forall a b. (a -> b) -> a -> b
$ Key -> SubMap -> Maybe Key
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Key
k SubMap
sbm

rRange :: forall r. (Enum r, Bounded r) => [r]
rRange :: forall r. (Enum r, Bounded r) => [r]
rRange = [forall a. Bounded a => a
minBound @r .. r
forall a. Bounded a => a
maxBound]