{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE OverloadedStrings #-}

module Asm.M ( CFunc (..)
             , WM
             , Label
             , nextI
             , nextL
             , foldMapA
             , prettyLabel
             , i4
             , pAsm, prettyAsm
             , aArr, mFree
             ) where

import           Control.DeepSeq            (NFData)
import           Control.Monad.State.Strict (State, state)
import           Data.Foldable              (fold, traverse_)
import qualified Data.IntMap                as IM
import           Data.List                  (scanl')
import           Data.Word                  (Word64)
import           Foreign.Marshal.Alloc      (free)
import           Foreign.Marshal.Array      (mallocArray, pokeArray)
import           Foreign.Ptr                (Ptr, plusPtr)
import           GHC.Generics               (Generic)
import qualified IR
import           Prettyprinter              (Doc, Pretty (pretty), indent)
import           Prettyprinter.Ext

type WM = State IR.WSt

type Label = Word

foldMapA :: (Applicative f, Traversable t, Monoid m) => (a -> f m) -> t a -> f m
foldMapA :: forall (f :: * -> *) (t :: * -> *) m a.
(Applicative f, Traversable t, Monoid m) =>
(a -> f m) -> t a -> f m
foldMapA = ((t m -> m) -> f (t m) -> f m
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap t m -> m
forall m. Monoid m => t m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (f (t m) -> f m) -> (t a -> f (t m)) -> t a -> f m
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((t a -> f (t m)) -> t a -> f m)
-> ((a -> f m) -> t a -> f (t m)) -> (a -> f m) -> t a -> f m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f m) -> t a -> f (t m)
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) -> t a -> f (t b)
traverse

prettyLabel :: Label -> Doc ann
prettyLabel :: forall ann. Label -> Doc ann
prettyLabel Label
l = Doc ann
"apple_" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Label -> Doc ann
forall ann. Label -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Label
l

i4 :: Doc ann -> Doc ann
i4 = Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
4

prettyAsm :: (Pretty isn) => (IR.AsmData, [isn]) -> Doc ann
prettyAsm :: forall isn ann. Pretty isn => (AsmData, [isn]) -> Doc ann
prettyAsm (AsmData
ds,[isn]
is) = AsmData -> Doc ann
forall {t :: * -> *} {ann}.
(Foldable t, Functor t) =>
IntMap (t Word64) -> Doc ann
pAD AsmData
ds Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
<#> [isn] -> Doc ann
forall isn ann. Pretty isn => [isn] -> Doc ann
pAsm [isn]
is

pAsm :: Pretty isn => [isn] -> Doc ann
pAsm :: forall isn ann. Pretty isn => [isn] -> Doc ann
pAsm = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
prettyLines([Doc ann] -> Doc ann) -> ([isn] -> [Doc ann]) -> [isn] -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(isn -> Doc ann) -> [isn] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap isn -> Doc ann
forall ann. isn -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty

nextI :: WM Int
nextI :: WM Int
nextI = (WSt -> (Int, WSt)) -> WM Int
forall a. (WSt -> (a, WSt)) -> StateT WSt Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\(IR.WSt Label
l Int
i) -> (Int
i, Label -> Int -> WSt
IR.WSt Label
l (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)))

nextL :: WM Label
nextL :: WM Label
nextL = (WSt -> (Label, WSt)) -> WM Label
forall a. (WSt -> (a, WSt)) -> StateT WSt Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\(IR.WSt Label
i Int
t) -> (Label
i, Label -> Int -> WSt
IR.WSt (Label
iLabel -> Label -> Label
forall a. Num a => a -> a -> a
+Label
1) Int
t))

data CFunc = Malloc | Free | JR | DR | Exp | Log | Pow deriving ((forall x. CFunc -> Rep CFunc x)
-> (forall x. Rep CFunc x -> CFunc) -> Generic CFunc
forall x. Rep CFunc x -> CFunc
forall x. CFunc -> Rep CFunc x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CFunc -> Rep CFunc x
from :: forall x. CFunc -> Rep CFunc x
$cto :: forall x. Rep CFunc x -> CFunc
to :: forall x. Rep CFunc x -> CFunc
Generic)

instance NFData CFunc where

instance Pretty CFunc where
    pretty :: forall ann. CFunc -> Doc ann
pretty CFunc
Malloc=Doc ann
"malloc"; pretty CFunc
Free=Doc ann
"free"
    pretty CFunc
JR=Doc ann
"lrand48"; pretty CFunc
DR=Doc ann
"drand48"
    pretty CFunc
Exp=Doc ann
"exp"; pretty CFunc
Log=Doc ann
"log"; pretty CFunc
Pow=Doc ann
"pow"

mFree :: Maybe (Ptr a) -> IO ()
mFree :: forall a. Maybe (Ptr a) -> IO ()
mFree = (Ptr a -> IO ()) -> Maybe (Ptr a) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Ptr a -> IO ()
forall a. Ptr a -> IO ()
free

aArr :: IM.IntMap [Word64] -> IO (IM.IntMap (Ptr Word64))
aArr :: AsmData -> IO (IntMap (Ptr Word64))
aArr AsmData
as = do
    let bls :: IntMap Int
bls = ([Word64] -> Int) -> AsmData -> IntMap Int
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Word64] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length AsmData
as; bl :: Int
bl = IntMap Int -> Int
forall a. Num a => IntMap a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum IntMap Int
bls
    p <- Int -> IO (Ptr Word64)
forall a. Storable a => Int -> IO (Ptr a)
mallocArray Int
bl
    let bs = [[Word64]] -> [Word64]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (AsmData -> [[Word64]]
forall a. IntMap a -> [a]
IM.elems AsmData
as)
    pokeArray p bs
    pure $ case IM.toList bls of
        []             -> IntMap (Ptr Word64)
forall a. IntMap a
IM.empty
        ((Int
k0,Int
l0):[(Int, Int)]
bls') -> [(Int, Ptr Word64)] -> IntMap (Ptr Word64)
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, Ptr Word64)] -> IntMap (Ptr Word64))
-> ([(Int, Int, Ptr Word64)] -> [(Int, Ptr Word64)])
-> [(Int, Int, Ptr Word64)]
-> IntMap (Ptr Word64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int, Ptr Word64) -> (Int, Ptr Word64))
-> [(Int, Int, Ptr Word64)] -> [(Int, Ptr Word64)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
x,Int
_,Ptr Word64
z) -> (Int
x,Ptr Word64
z)) ([(Int, Int, Ptr Word64)] -> IntMap (Ptr Word64))
-> [(Int, Int, Ptr Word64)] -> IntMap (Ptr Word64)
forall a b. (a -> b) -> a -> b
$ ((Int, Int, Ptr Word64) -> (Int, Int) -> (Int, Int, Ptr Word64))
-> (Int, Int, Ptr Word64)
-> [(Int, Int)]
-> [(Int, Int, Ptr Word64)]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' (\(Int
_, Int
lϵ, Ptr Word64
pϵ) (Int
k, Int
l) -> (Int
k, Int
l, Ptr Word64
pϵ Ptr Word64 -> Int -> Ptr Word64
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
lϵInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
8))) (Int
k0, Int
l0, Ptr Word64
p) [(Int, Int)]
bls'