<commit_before>{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE TypeFamilies    #-}

module Toy.X86.Translator
    ( compile
    , produceBinary
    ) where

import           Control.Lens          (Lens', ix, (%~), (&), (+=), (-=), (<&>), (^?))
import           Control.Monad         (forM, join)
import           Control.Monad.State   (get, runState)
import           Control.Monad.Trans   (MonadIO (..))
import           Data.Functor          (($>))
import           Data.Maybe            (fromMaybe)
import           Data.Monoid           ((<>))
import qualified Data.Set              as S
import           Data.Text             (Text)
import qualified Data.Vector           as V
import           Formatting            (build, formatToString, int, shown, (%))
import qualified Formatting            as F
import           GHC.Exts              (fromList)
import           System.FilePath.Posix ((</>))
import           System.Process        (proc)

import           Toy.Exp               (FunSign (..), Var)
import qualified Toy.SM                as SM
import           Toy.X86.Data          (Inst (..), Insts, Operand (..), Program (..),
                                        StackDirection (..), eax, edx, jmp, ret,
                                        withStackSpace, (//), (?))
import           Toy.X86.Frame         (evalStackShift, mkFrame, resolveMemRefs)
import           Toy.X86.Optimize      (optimize)
import           Toy.X86.Process       (readCreateProcess)
import           Toy.X86.SymStack      (SymStackHolder, allocSymStackOp, occupiedRegs,
                                        popSymStackOp, runSymStackHolder)

compile :: SM.Insts -> Insts
compile = mconcat . map compileFun . separateFuns

compileFun :: SM.Insts -> Insts
compileFun insts =
    let (name, args) = case insts ^? ix 0 of
            Just (SM.Enter n ai) -> (n, reverse ai)  -- TODO: ???
            _                    -> error "Where is my Enter?!"
        vars  = foldMap gatherLocals insts
        ((symStSpace, symStackSizeAtEnd), body) = runSymStackHolder $
            join . fmap fromList <$> mapM (step name) insts
        check = if symStackSizeAtEnd == 0 then id
                else error . formatToString ("Wrong sym stack size at end: "%
                     int%"\n"%build) symStackSizeAtEnd . Program
        frame = mkFrame args vars symStSpace
        post  = [ resolveMemRefs frame
                , fixMemRefs
                , insertExit
                , mkStackShift (evalStackShift frame)
                , optimize
                ] :: [Insts -> Insts]
    in check $ foldl (&) body post

-- TODO: correct errors processing
step :: Var -> SM.Inst -> SymStackHolder [Inst]
step calleeName = \case
    SM.Nop        -> pure []
    SM.Push v     -> allocSymStackOp <&> \op -> [Mov (Const v) op]
    SM.Drop       -> popSymStackOp $> []
    SM.Load v     -> allocSymStackOp <&> \op -> [Mov (Local v) eax, Mov eax op]
    SM.Store v    -> popSymStackOp <&> \op -> [Mov op eax, Mov eax (Local v)]
    SM.Bin op    -> do
        op2 <- popSymStackOp
        op1 <- popSymStackOp
        resOp <- allocSymStackOp
        return $ binop op1 op2 op <> [Mov op2 eax, Mov eax resOp]
    SM.Label lid -> pure [Label lid]
    SM.Jmp   lid -> pure [jmp (SM.CLabel lid)]
    SM.JmpIf lid -> popSymStackOp <&> \op ->
        [ Mov op eax
        , BinOp "cmp" (Const 0) eax
        , Jmp "ne" (SM.CLabel lid)
        ]
    SM.Call (FunSign name args) -> do
        part1 <- rolloutSymStackOps (length args) [Call name]
        part2 <- allocSymStackOp <&> \op -> [Mov eax op]
        return $ part1 <> part2
    SM.Ret       -> do
        p1 <- popSymStackOp <&> \op -> [Mov op eax]
        return $ p1 <> [jmp (SM.ELabel calleeName)]
    SM.Enter{}   -> pure [NoopOperator "int3"]

separateFuns :: SM.Insts -> [SM.Insts]
separateFuns insts =
    case V.findIndex (\(i, v) -> isEnter v && i > 0) $ V.indexed insts of
        Nothing -> [insts]
        Just k  -> let (funInsts, remaining) = V.splitAt k insts
                   in  funInsts : separateFuns remaining
  where
    isEnter SM.Enter{} = True
    isEnter _          = False

gatherLocals :: SM.Inst -> S.Set Var
gatherLocals (SM.Store v) = [v]
gatherLocals _            = []

funBodyInsts :: Lens' Insts Insts
funBodyInsts f insts =
    let funLabelPos = fromMaybe (error "No fun label!")
                    $ V.findIndex isFunLabel insts
        (prelude, postlude) = V.splitAt (funLabelPos + 1) insts
        (body   , end)      = cutLast postlude
    in  if V.last insts /= ret then noRetAtEnd
        else f body <&> \body' -> prelude <> body' <> end
  where
    cutLast v = V.splitAt (V.length v - 1) v

    isFunLabel (Label (SM.FLabel _)) = True
    isFunLabel _                     = False

    noRetAtEnd = -- TODO: simplify
        error $ formatToString ("No 'Ret' at the end of function:\n"%shown)
        insts

insertExit :: Insts -> Insts
insertExit = (<> [ret])

mkStackShift :: Int -> Insts -> Insts
mkStackShift shift = funBodyInsts %~ withStackSpace shift

-- | Function 'step', when sets `Mem` indices, doesn't take into account
-- possible @Push@es and @Pop@s to stack. This functions fixes it.
--
-- This is done as post-processing because, unlike SM, in X86 (real) stack is
-- influenced by not many operations.
fixMemRefs :: Traversable f => f Inst -> f Inst
fixMemRefs insts =
    let (res, finalStackShift) = flip runState 0 $ forM insts $ \case<commit_msg>Insert `ret` after making stack shift<commit_after>{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE TypeFamilies    #-}

module Toy.X86.Translator
    ( compile
    , produceBinary
    ) where

import           Control.Lens          (Lens', ix, (%~), (&), (+=), (-=), (<&>), (^?))
import           Control.Monad         (forM, join)
import           Control.Monad.State   (get, runState)
import           Control.Monad.Trans   (MonadIO (..))
import           Data.Functor          (($>))
import           Data.Maybe            (fromMaybe)
import           Data.Monoid           ((<>))
import qualified Data.Set              as S
import           Data.Text             (Text)
import qualified Data.Vector           as V
import           Formatting            (build, formatToString, int, (%))
import qualified Formatting            as F
import           GHC.Exts              (fromList)
import           System.FilePath.Posix ((</>))
import           System.Process        (proc)

import           Toy.Exp               (FunSign (..), Var)
import qualified Toy.SM                as SM
import           Toy.X86.Data          (Inst (..), Insts, Operand (..), Program (..),
                                        StackDirection (..), eax, edx, jmp, ret,
                                        withStackSpace, (//), (?))
import           Toy.X86.Frame         (evalStackShift, mkFrame, resolveMemRefs)
import           Toy.X86.Optimize      (optimize)
import           Toy.X86.Process       (readCreateProcess)
import           Toy.X86.SymStack      (SymStackHolder, allocSymStackOp, occupiedRegs,
                                        popSymStackOp, runSymStackHolder)

compile :: SM.Insts -> Insts
compile = mconcat . map compileFun . separateFuns

compileFun :: SM.Insts -> Insts
compileFun insts =
    let (name, args) = case insts ^? ix 0 of
            Just (SM.Enter n ai) -> (n, reverse ai)  -- TODO: ???
            _                    -> error "Where is my Enter?!"
        vars  = foldMap gatherLocals insts
        ((symStSpace, symStackSizeAtEnd), body) = runSymStackHolder $
            join . fmap fromList <$> mapM (step name) insts
        check = if symStackSizeAtEnd == 0 then id
                else error . formatToString ("Wrong sym stack size at end: "%
                     int%"\n"%build) symStackSizeAtEnd . Program
        frame = mkFrame args vars symStSpace
        post  = [ resolveMemRefs frame
                , fixMemRefs
                , mkStackShift (evalStackShift frame)
                , insertExit
                , optimize
                ] :: [Insts -> Insts]
    in check $ foldl (&) body post

-- TODO: correct errors processing
step :: Var -> SM.Inst -> SymStackHolder [Inst]
step calleeName = \case
    SM.Nop        -> pure []
    SM.Push v     -> allocSymStackOp <&> \op -> [Mov (Const v) op]
    SM.Drop       -> popSymStackOp $> []
    SM.Load v     -> allocSymStackOp <&> \op -> [Mov (Local v) eax, Mov eax op]
    SM.Store v    -> popSymStackOp <&> \op -> [Mov op eax, Mov eax (Local v)]
    SM.Bin op    -> do
        op2 <- popSymStackOp
        op1 <- popSymStackOp
        resOp <- allocSymStackOp
        return $ binop op1 op2 op <> [Mov op2 eax, Mov eax resOp]
    SM.Label lid -> pure [Label lid]
    SM.Jmp   lid -> pure [jmp (SM.CLabel lid)]
    SM.JmpIf lid -> popSymStackOp <&> \op ->
        [ Mov op eax
        , BinOp "cmp" (Const 0) eax
        , Jmp "ne" (SM.CLabel lid)
        ]
    SM.Call (FunSign name args) -> do
        part1 <- rolloutSymStackOps (length args) [Call name]
        part2 <- allocSymStackOp <&> \op -> [Mov eax op]
        return $ part1 <> part2
    SM.Ret       -> do
        p1 <- popSymStackOp <&> \op -> [Mov op eax]
        return $ p1 <> [jmp (SM.ELabel calleeName)]
    SM.Enter{}   -> pure [NoopOperator "int3"]

separateFuns :: SM.Insts -> [SM.Insts]
separateFuns insts =
    case V.findIndex (\(i, v) -> isEnter v && i > 0) $ V.indexed insts of
        Nothing -> [insts]
        Just k  -> let (funInsts, remaining) = V.splitAt k insts
                   in  funInsts : separateFuns remaining
  where
    isEnter SM.Enter{} = True
    isEnter _          = False

gatherLocals :: SM.Inst -> S.Set Var
gatherLocals (SM.Store v) = [v]
gatherLocals _            = []

afterFunBeginning :: Lens' Insts Insts
afterFunBeginning f insts =
    let funLabelPos = fromMaybe (error "No fun label!")
                    $ V.findIndex isFunLabel insts
        (prelude, body) = V.splitAt (funLabelPos + 1) insts
    in  f body <&> \body' -> prelude <> body'
  where
    isFunLabel (Label (SM.FLabel _)) = True
    isFunLabel _                     = False

insertExit :: Insts -> Insts
insertExit = (<> [ret])

mkStackShift :: Int -> Insts -> Insts
mkStackShift shift = afterFunBeginning %~ withStackSpace shift

-- | Function 'step', when sets `Mem` indices, doesn't take into account
-- possible @Push@es and @Pop@s to stack. This functions fixes it.
--
-- This is done as post-processing because, unlike SM, in X86 (real) stack is
-- influenced by not many operations.
fixMemRefs :: Traversable f => f Inst -> f Inst
fixMemRefs insts =
    let (res, finalStackShift) = flip runState 0 $ forM insts $ \case
{'fraction_non_alphanumeric': 0.11692192753081808, 'fraction_numerical': 0.0060702278670153155, 'mean_word_length': 1.782281111977137, 'pattern_counts': {'":': 0, '<': 48, '<?xml version=': 0, '>': 114, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '326271', 'n_tokens_mistral': 3357, 'n_tokens_neox': 3271, 'n_words': 983}