module Asm.CF ( N, FreshM
              , runFreshM
              , getFresh, fm
              , lookupLabel, lC
              , broadcast
              , b3
              , singleton
              , fromList
              ) where

import           Asm.M
import           Class.E                    as E
import           Control.Monad.State.Strict (State, evalState, gets, modify, state)
import           Data.Functor               (($>))
import qualified Data.IntSet                as IS
import qualified Data.Map                   as M
import           Data.Tuple.Extra           (second3, snd3, thd3, third3)

type N=Int

-- map of labels by node
type FreshM = State (N, M.Map Label N, M.Map Label [N])

runFreshM :: FreshM a -> a
runFreshM :: forall a. FreshM a -> a
runFreshM = (FreshM a -> (N, Map Label N, Map Label [N]) -> a)
-> (N, Map Label N, Map Label [N]) -> FreshM a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip FreshM a -> (N, Map Label N, Map Label [N]) -> a
forall s a. State s a -> s -> a
evalState (N
0, Map Label N
forall a. Monoid a => a
mempty, Map Label [N]
forall a. Monoid a => a
mempty)

getFresh :: FreshM N
getFresh :: FreshM N
getFresh = ((N, Map Label N, Map Label [N])
 -> (N, (N, Map Label N, Map Label [N])))
-> FreshM N
forall a.
((N, Map Label N, Map Label [N])
 -> (a, (N, Map Label N, Map Label [N])))
-> StateT (N, Map Label N, Map Label [N]) Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\(N
i,Map Label N
m0,Map Label [N]
m1) -> (N
i,(N
iN -> N -> N
forall a. Num a => a -> a -> a
+N
1,Map Label N
m0,Map Label [N]
m1)))

fm :: Label -> FreshM N
fm :: Label -> FreshM N
fm Label
l = do {st <- ((N, Map Label N, Map Label [N]) -> Map Label N)
-> StateT (N, Map Label N, Map Label [N]) Identity (Map Label N)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (N, Map Label N, Map Label [N]) -> Map Label N
forall a b c. (a, b, c) -> b
snd3; case M.lookup l st of {Just N
i -> N -> FreshM N
forall a. a -> StateT (N, Map Label N, Map Label [N]) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure N
i; Maybe N
Nothing -> do {i <- FreshM N
getFresh; broadcast i l $> i}}}

lookupLabel :: Label -> FreshM N
lookupLabel :: Label -> FreshM N
lookupLabel Label
l = ((N, Map Label N, Map Label [N]) -> N) -> FreshM N
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (N -> Label -> Map Label N -> N
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault ([Char] -> N
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error in control-flow graph: node label not in map.") Label
l (Map Label N -> N)
-> ((N, Map Label N, Map Label [N]) -> Map Label N)
-> (N, Map Label N, Map Label [N])
-> N
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (N, Map Label N, Map Label [N]) -> Map Label N
forall a b c. (a, b, c) -> b
snd3)

lC :: Label -> FreshM [N]
lC :: Label -> FreshM [N]
lC Label
l = ((N, Map Label N, Map Label [N]) -> [N]) -> FreshM [N]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ([N] -> Label -> Map Label [N] -> [N]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault ([Char] -> [N]
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error in CF graph: node label not in map.") Label
l (Map Label [N] -> [N])
-> ((N, Map Label N, Map Label [N]) -> Map Label [N])
-> (N, Map Label N, Map Label [N])
-> [N]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (N, Map Label N, Map Label [N]) -> Map Label [N]
forall a b c. (a, b, c) -> c
thd3)

broadcast :: N -> Label -> FreshM ()
broadcast :: N -> Label -> FreshM ()
broadcast N
i Label
l = ((N, Map Label N, Map Label [N])
 -> (N, Map Label N, Map Label [N]))
-> FreshM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Label N -> Map Label N)
-> (N, Map Label N, Map Label [N])
-> (N, Map Label N, Map Label [N])
forall b b' a c. (b -> b') -> (a, b, c) -> (a, b', c)
second3 (Label -> N -> Map Label N -> Map Label N
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Label
l N
i))

b3 :: N -> Label -> FreshM ()
b3 :: N -> Label -> FreshM ()
b3 N
i Label
l = ((N, Map Label N, Map Label [N])
 -> (N, Map Label N, Map Label [N]))
-> FreshM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Label [N] -> Map Label [N])
-> (N, Map Label N, Map Label [N])
-> (N, Map Label N, Map Label [N])
forall c c' a b. (c -> c') -> (a, b, c) -> (a, b, c')
third3 ((Maybe [N] -> Maybe [N]) -> Label -> Map Label [N] -> Map Label [N]
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (\Maybe [N]
k -> [N] -> Maybe [N]
forall a. a -> Maybe a
Just([N] -> Maybe [N]) -> [N] -> Maybe [N]
forall a b. (a -> b) -> a -> b
$case Maybe [N]
k of {Maybe [N]
Nothing -> [N
i]; Just [N]
is -> N
iN -> [N] -> [N]
forall a. a -> [a] -> [a]
:[N]
is}) Label
l))

singleton :: E reg => reg -> IS.IntSet
singleton :: forall reg. E reg => reg -> IntSet
singleton = N -> IntSet
IS.singleton (N -> IntSet) -> (reg -> N) -> reg -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. reg -> N
forall a. E a => a -> N
E.toInt

fromList :: E reg => [reg] -> IS.IntSet
fromList :: forall reg. E reg => [reg] -> IntSet
fromList = (reg -> IntSet) -> [reg] -> IntSet
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap reg -> IntSet
forall reg. E reg => reg -> IntSet
singleton