From 9a78d71b2d5dde9341d49dfb7b2b5bccb3ea4f88 Mon Sep 17 00:00:00 2001
From: Vorotynsky Maxim <vorotynsky.maxim@gmail.com>
Date: Wed, 20 May 2020 20:18:21 +0300
Subject: [PATCH] fix: stack align

Fixed stack alignment after calls and nested frames.
---
 src/HLasm/Backend/Nasm.hs |  7 ++++---
 src/HLasm/Instructions.hs | 10 ++++++----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/HLasm/Backend/Nasm.hs b/src/HLasm/Backend/Nasm.hs
index 076a414..fb4795a 100644
--- a/src/HLasm/Backend/Nasm.hs
+++ b/src/HLasm/Backend/Nasm.hs
@@ -14,7 +14,7 @@ choosePtr = f . bytes
     where f 1 = "BYTE"
           f 2 = "WORD"
           f 4 = "DWORD"
-          f _ = error "undefined ptr size"
+          f n = error ("undefined data size: " ++ show n)
 
 target :: Target -> String
 target (NamedTarget name)           = name
@@ -35,8 +35,9 @@ frame f = ["push ebp", "mov ebp, esp", "sub esp, " ++ (show . size $ f)]
 instruction :: Instructions -> [String]
 instruction (PureAsm str)               = [str]
 instruction (BeginFrame f (Just l))     = (l ++ ":"):(frame f)
+instruction (EndFrame   f (Just _))     = ["leave", "ret"]
 instruction (BeginFrame f Nothing)      = frame f
-instruction (EndFrame f)                = ["leave", "ret"]
+instruction (EndFrame   f Nothing)      = ["leave"]
 instruction (Label l)                   = [l ++ ":"]
 instruction (Move l r)                  = [instr2arg "mov" l r]
 instruction (Compare l r)               = [instr2arg "cmp" l r]
@@ -45,7 +46,7 @@ instruction (Jump lbl (Just Equals))    = ["je " ++ lbl]
 instruction (Jump lbl (Just NotEquals)) = ["jne " ++ lbl]
 instruction (Jump lbl (Just Greater))   = ["jg " ++ lbl]
 instruction (Jump lbl (Just Less))      = ["jl " ++ lbl]
-instruction (Call lbl args)             = (fmap push . reverse $ args) ++ ["call " ++ lbl ++ "; HACK: stack doesn't align"] -- TOOD
+instruction (Call lbl args size)        = (fmap push . reverse $ args) ++ ["call " ++ lbl, "add esp, " ++ show (bytes size)]
     where push x = "push " ++ (target x)
 
 join :: String -> [String] -> String
diff --git a/src/HLasm/Instructions.hs b/src/HLasm/Instructions.hs
index a96e9fb..5ecdff1 100644
--- a/src/HLasm/Instructions.hs
+++ b/src/HLasm/Instructions.hs
@@ -29,10 +29,10 @@ data Target =
 data Instructions = 
     PureAsm String
     | BeginFrame StackFrame (Maybe Label)
-    | EndFrame StackFrame
+    | EndFrame StackFrame (Maybe Label)
     | Label Label
     | Move Target Target
-    | Call Label [Target]
+    | Call Label [Target] Int
     | Compare Target Target
     | Jump Label (Maybe CompareType)
     deriving (Show, Eq)
@@ -65,10 +65,12 @@ instructions (Node ((VariableDeclaration val), _, _, _) _ ) = Just []
 instructions (Node ((While lbl              ), _, _, _) xs) = loop lbl (concatMapM instructions xs)
 instructions (Node ((DoWhile lbl            ), _, _, _) xs) = loop lbl (concatMapM instructions xs)
 instructions (Node ((Break lbl              ), _, _, _) _ ) = Just [Jump (lbl ++ "end") Nothing]
-instructions (Node ((HLasm.Ast.Call lbl ns  ), d, _, f) _ ) = Just [HLasm.Instructions.Call lbl (fmap (findTarget f d) ns)]
+instructions (Node ((HLasm.Ast.Call lbl ns  ), d, _, f) _ ) = 
+    Just [HLasm.Instructions.Call lbl (fmap (findTarget f d) ns) size]
+    where size = foldl (+) 0 . fmap (\(VariableData (_, (VariableDeclaration d))) -> valueSize d) $ d 
 instructions (Node ((AssemblyCall str       ), _, _, _) _ ) = Just [PureAsm str]
 instructions (Node ((Frame lbl              ), _, _, f) xs) =
-    fmap (\x -> [BeginFrame f lbl] ++ x ++ [EndFrame f]) $ concatMapM instructions xs
+    fmap (\x -> [BeginFrame f lbl] ++ x ++ [EndFrame f lbl]) $ concatMapM instructions xs
 
 instructions (Node ((Assignment name (NameValue val)),    d, _, f) _) = Just [Move (findTarget f d name) (findTarget f d val)]
 instructions (Node ((Assignment name (IntegerValue val)), d, _, f) _) = Just [Move (findTarget f d name) (ConstantTarget val)]