module Asm.X86.Sp ( spill ) where

import           Asm.X86
import           Asm.X86.CF
import           Control.Monad.Extra        (concatMapM)
import           Control.Monad.State.Strict (State, runState, state)
import           Data.Functor               (void)
import           Data.Int                   (Int32, Int8)
import qualified Data.IntMap.Strict         as IM
import qualified Data.IntSet                as IS
import           Data.Maybe                 (catMaybes)

type SpM = State Int

next :: SpM Int
next :: SpM Int
next = (Int -> (Int, Int)) -> SpM Int
forall a. (Int -> (a, Int)) -> StateT Int Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\Int
i -> (Int
i, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))

spill :: Int -- ^ Unique state
      -> Int
      -> IS.IntSet
      -> [X86 AbsReg FAbsReg X2Abs a]
      -> (Int, Int, [X86 AbsReg FAbsReg X2Abs ()])
spill :: forall a.
Int
-> Int
-> IntSet
-> [X86 AbsReg FAbsReg X2Abs a]
-> (Int, Int, [X86 AbsReg FAbsReg X2Abs ()])
spill Int
u Int
offs IntSet
m [X86 AbsReg FAbsReg X2Abs a]
isns =
    let (Int
o', SpM [X86 AbsReg FAbsReg X2Abs ()]
ᴍ) = Int
-> IntSet
-> [X86 AbsReg FAbsReg X2Abs a]
-> (Int, SpM [X86 AbsReg FAbsReg X2Abs ()])
forall a.
Int
-> IntSet
-> [X86 AbsReg FAbsReg X2Abs a]
-> (Int, SpM [X86 AbsReg FAbsReg X2Abs ()])
spillM Int
offs IntSet
m [X86 AbsReg FAbsReg X2Abs a]
isns
        ([X86 AbsReg FAbsReg X2Abs ()]
nisns, Int
u') = SpM [X86 AbsReg FAbsReg X2Abs ()]
-> Int -> ([X86 AbsReg FAbsReg X2Abs ()], Int)
forall s a. State s a -> s -> (a, s)
runState SpM [X86 AbsReg FAbsReg X2Abs ()]
ᴍ Int
u
    in (Int
u', Int
o', [X86 AbsReg FAbsReg X2Abs ()]
nisns)

spillM :: Int -- ^ Offset (from already spilled)
       -> IS.IntSet
       -> [X86 AbsReg FAbsReg X2Abs a]
       -> (Int, SpM [X86 AbsReg FAbsReg X2Abs ()]) -- ^ offset, rewritten
spillM :: forall a.
Int
-> IntSet
-> [X86 AbsReg FAbsReg X2Abs a]
-> (Int, SpM [X86 AbsReg FAbsReg X2Abs ()])
spillM Int
offs IntSet
m [X86 AbsReg FAbsReg X2Abs a]
isns = (Int
foffs, (X86 AbsReg FAbsReg X2Abs a -> SpM [X86 AbsReg FAbsReg X2Abs ()])
-> [X86 AbsReg FAbsReg X2Abs a]
-> SpM [X86 AbsReg FAbsReg X2Abs ()]
forall (m :: * -> *) a b. Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM X86 AbsReg FAbsReg X2Abs a -> SpM [X86 AbsReg FAbsReg X2Abs ()]
forall {freg} {f2reg} {ann}.
X86 AbsReg freg f2reg ann
-> StateT Int Identity [X86 AbsReg freg f2reg ()]
g [X86 AbsReg FAbsReg X2Abs a]
isns)
    where g :: X86 AbsReg freg f2reg ann
-> StateT Int Identity [X86 AbsReg freg f2reg ()]
g X86 AbsReg freg f2reg ann
isn = do
            let is :: [Int]
is = [ AbsReg -> Int
toInt AbsReg
r | AbsReg
r <- (AbsReg -> [AbsReg]) -> X86 AbsReg freg f2reg ann -> [AbsReg]
forall m reg freg f2reg a.
Monoid m =>
(reg -> m) -> X86 reg freg f2reg a -> m
fR AbsReg -> [AbsReg]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure X86 AbsReg freg f2reg ann
isn, AbsReg -> Int
toInt AbsReg
r Int -> IntSet -> Bool
`IS.member` IntSet
m ]
            newRs <- (Int -> StateT Int Identity AbsReg)
-> [Int] -> StateT Int Identity [AbsReg]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\Int
_ -> Int -> AbsReg
IReg (Int -> AbsReg) -> SpM Int -> StateT Int Identity AbsReg
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SpM Int
next) [Int]
is
            let f = [AbsReg -> AbsReg] -> AbsReg -> AbsReg
forall {b}. [b -> b] -> b -> b
thread ((Int -> AbsReg -> AbsReg -> AbsReg)
-> [Int] -> [AbsReg] -> [AbsReg -> AbsReg]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i AbsReg
rϵ AbsReg
r -> if AbsReg -> Int
toInt AbsReg
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then AbsReg
rϵ else AbsReg
r) [Int]
is [AbsReg]
newRs)
                ma Int
i = Int -> Addr AbsReg
ao (Int -> Int
at Int
i); as = Int -> Addr AbsReg
ma (Int -> Addr AbsReg) -> [Int] -> [Addr AbsReg]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
is
                isn' = (AbsReg -> AbsReg)
-> X86 AbsReg freg f2reg ann -> X86 AbsReg freg f2reg ann
forall areg reg afreg af2 a.
(areg -> reg) -> X86 areg afreg af2 a -> X86 reg afreg af2 a
mapR AbsReg -> AbsReg
f X86 AbsReg freg f2reg ann
isn
            pure $
                   catMaybes (zipWith (\AbsReg
r Addr AbsReg
a -> if AbsReg -> Int
toInt AbsReg
r Int -> IntSet -> Bool
`IS.member` X86 AbsReg freg f2reg ann -> IntSet
forall reg freg f2reg ann.
E reg =>
X86 reg freg f2reg ann -> IntSet
uses X86 AbsReg freg f2reg ann
isn' then X86 AbsReg freg f2reg () -> Maybe (X86 AbsReg freg f2reg ())
forall a. a -> Maybe a
Just (() -> AbsReg -> Addr AbsReg -> X86 AbsReg freg f2reg ()
forall reg freg f2 a. a -> reg -> Addr reg -> X86 reg freg f2 a
MovRA () AbsReg
r Addr AbsReg
a) else Maybe (X86 AbsReg freg f2reg ())
forall a. Maybe a
Nothing) newRs as)
                ++ void isn'
                : catMaybes (zipWith (\Addr AbsReg
a AbsReg
r -> if AbsReg -> Int
toInt AbsReg
r Int -> IntSet -> Bool
`IS.member` X86 AbsReg freg f2reg ann -> IntSet
forall reg freg f2reg ann.
E reg =>
X86 reg freg f2reg ann -> IntSet
defs X86 AbsReg freg f2reg ann
isn' then X86 AbsReg freg f2reg () -> Maybe (X86 AbsReg freg f2reg ())
forall a. a -> Maybe a
Just (() -> Addr AbsReg -> AbsReg -> X86 AbsReg freg f2reg ()
forall reg freg f2 a. a -> Addr reg -> reg -> X86 reg freg f2 a
MovAR () Addr AbsReg
a AbsReg
r) else Maybe (X86 AbsReg freg f2reg ())
forall a. Maybe a
Nothing) as newRs)

          ass :: IS.IntSet -> IM.IntMap Int
          ass :: IntSet -> IntMap Int
ass = [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, Int)] -> IntMap Int)
-> (IntSet -> [(Int, Int)]) -> IntSet -> IntMap Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\[Int]
k -> [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
k [Int
offs,Int
offsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8..]) ([Int] -> [(Int, Int)])
-> (IntSet -> [Int]) -> IntSet -> [(Int, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> [Int]
IS.toList

          assgn :: IntMap Int
assgn = IntSet -> IntMap Int
ass IntSet
m
          at :: Int -> Int
at Int
k = Int -> Int -> IntMap Int -> Int
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault ([Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error.") Int
k IntMap Int
assgn

          foffs :: Int
foffs = Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*IntSet -> Int
IS.size IntSet
m

          thread :: [b -> b] -> b -> b
thread = ((b -> b) -> (b -> b) -> b -> b) -> (b -> b) -> [b -> b] -> b -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (b -> b) -> (b -> b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) b -> b
forall a. a -> a
id

ao :: Int -> Addr AbsReg
ao Int
o | Just Int8
i8 <- Int -> Maybe Int8
mi8 Int
o = AbsReg -> Int8 -> Addr AbsReg
forall reg. reg -> Int8 -> Addr reg
RC AbsReg
BP Int8
i8
     | Just Int32
i32 <- Int -> Maybe Int32
mi32 Int
o = AbsReg -> Int32 -> Addr AbsReg
forall reg. reg -> Int32 -> Addr reg
RC32 AbsReg
BP Int32
i32

mi8 :: Int -> Maybe Int8
mi8 :: Int -> Maybe Int8
mi8 Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int8
forall a. Bounded a => a
maxBound :: Int8) Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int8
forall a. Bounded a => a
minBound :: Int8) = Int8 -> Maybe Int8
forall a. a -> Maybe a
Just (Int8 -> Maybe Int8) -> Int8 -> Maybe Int8
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
      | Bool
otherwise = Maybe Int8
forall a. Maybe a
Nothing

mi32 :: Int -> Maybe Int32
mi32 :: Int -> Maybe Int32
mi32 Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32
forall a. Bounded a => a
maxBound :: Int32) Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32
forall a. Bounded a => a
minBound :: Int32) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just (Int32 -> Maybe Int32) -> Int32 -> Maybe Int32
forall a b. (a -> b) -> a -> b
$ Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
       | Bool
otherwise = Maybe Int32
forall a. Maybe a
Nothing