x86-symbolic: More SysV stack setup and manipulation

x86_symbolic/src/Data/Macaw/X86/Symbolic/ABI/SysV.hs

@@ -32,22 +32,34 @@ This module is meant to be imported qualified.
 module Data.Macaw.X86.Symbolic.ABI.SysV
   ( StackPointer
   , getStackPointer
+  , stackPointerReg
   , RetAddr(..)
   , freshRetAddr
+  , SpilledArgs(..)
   , writeSpilledArgs
   , writeRetAddr
   , allocStackFrame
+  , pushStackFrame
+  , allocStack
   ) where
 
 -- TODO: sort me!
+import qualified Control.Lens as Lens
 import qualified Lang.Crucible.Backend as C
+import qualified Lang.Crucible.Simulator as C
 import qualified Lang.Crucible.LLVM.MemModel as MM
 import qualified Data.Sequence as Seq
+import qualified Data.Macaw.Symbolic as MS
+import qualified Data.Macaw.Symbolic.Stack as MSS
+import qualified Data.Macaw.X86 as X86
+import qualified Data.Macaw.X86.Symbolic as X86S
+import Data.Parameterized.Classes (ixF')
 import qualified What4.Interface as WI
 import qualified Data.BitVector.Sized as BVS
 import qualified Lang.Crucible.LLVM.Bytes as Bytes
 import qualified Lang.Crucible.LLVM.DataLayout as CLD
 import qualified Control.Monad as Monad
+import qualified Data.Parameterized.Context as Ctx
 
 -- | Helper, not exported
 ptrBytes :: Integer
@@ -59,20 +71,37 @@ ptrRepr = WI.knownNat
 -- | A pointer to a SysV-compatible x86_64 stack
 newtype StackPointer sym = StackPointer { getStackPointer :: MM.LLVMPtr sym 64 }
 
+-- | A lens for accessing the stack pointer register as a 'StackPointer'
+stackPointerReg ::
+  Lens.Lens'
+    (Ctx.Assignment (C.RegValue' sym) (MS.MacawCrucibleRegTypes X86.X86_64))
+    (StackPointer sym)
+stackPointerReg =
+  Lens.lens
+    (\regs -> StackPointer (C.unRV (regs Lens.^. ixF' X86S.rsp)))
+    (\regs v -> regs Lens.& ixF' X86S.rsp Lens..~ C.RV (getStackPointer v))
+
 -- | A return address
 newtype RetAddr sym = RetAddr { getRetAddr :: MM.LLVMPtr sym 64 }
 
--- | Create a fresh, symbolic return address
+-- | Create a fresh, symbolic return address.
 freshRetAddr :: C.IsSymInterface sym => sym -> IO (RetAddr sym)
 freshRetAddr sym = do
   bv <- WI.freshConstant sym (WI.safeSymbol "x86_64_ret_addr") (WI.BaseBVRepr ptrRepr)
   ptr <- MM.llvmPointer_bv sym bv
   pure (RetAddr ptr)
 
+-- | Stack-spilled arguments, in normal order
+newtype SpilledArgs sym
+  = SpilledArgs { getSpilledArgs :: Seq.Seq (MM.LLVMPtr sym 64) }
+
 -- | Write pointer-sized stack-spilled arguments to the stack.
 --
 -- SysV specifies that they will be written in reverse order, i.e., the last
 -- element of the 'Seq.Seq' will be written to the highest address.
+--
+-- Asserts that the stack allocation is writable and large enough to contain the
+-- spilled arguments.
 writeSpilledArgs ::
   C.IsSymBackend sym bak =>
   (?memOpts :: MM.MemOptions) =>
@@ -80,8 +109,7 @@ writeSpilledArgs ::
   bak ->
   MM.MemImpl sym ->
   StackPointer sym ->
-  -- | Stack-spilled arguments, in normal order
-  Seq.Seq (MM.LLVMPtr sym 64) ->
+  SpilledArgs sym ->
   IO (StackPointer sym, MM.MemImpl sym)
 writeSpilledArgs bak mem sp spilledArgs = do
   let sym = C.backendGetSym bak
@@ -92,9 +120,12 @@ writeSpilledArgs bak mem sp spilledArgs = do
         p' <- MM.ptrSub sym ?ptrWidth p eight
         m' <- MM.storeRaw bak m p' i64 CLD.noAlignment (MM.ptrToPtrVal arg)
         pure (StackPointer p', m')
-  Monad.foldM writeWord (sp, mem) (Seq.reverse spilledArgs)
+  Monad.foldM writeWord (sp, mem) (Seq.reverse (getSpilledArgs spilledArgs))
 
 -- | Write the return address to the stack.
+--
+-- Asserts that the stack allocation is writable and large enough to contain the
+-- return address.
 writeRetAddr ::
   C.IsSymBackend sym bak =>
   (?memOpts :: MM.MemOptions) =>
@@ -120,6 +151,9 @@ writeRetAddr bak mem sp retAddr = do
 --
 -- * First, write stack-spilled arguments in reverse order
 -- * Then, write the return address
+--
+-- Asserts that the stack allocation is writable and large enough to contain the
+-- spilled arguments and return address.
 allocStackFrame ::
   C.IsSymBackend sym bak =>
   (?memOpts :: MM.MemOptions) =>
@@ -127,8 +161,7 @@ allocStackFrame ::
   bak ->
   MM.MemImpl sym ->
   StackPointer sym ->
-  -- | Stack-spilled arguments, in normal order
-  Seq.Seq (MM.LLVMPtr sym 64) ->
+  SpilledArgs sym ->
   RetAddr sym ->
   IO (StackPointer sym, MM.MemImpl sym)
 allocStackFrame bak mem sp spilledArgs retAddr = do
@@ -136,6 +169,51 @@ allocStackFrame bak mem sp spilledArgs retAddr = do
   (sp', mem') <- writeSpilledArgs bak mem sp spilledArgs
   writeRetAddr bak mem' sp' retAddr
 
--- TODO: Pushing a stack frame (i.e., manipulating rsp)
--- TODO: Extracting the stack pointer from registers
--- TODO: Allocating the stack
+-- | Like 'allocStackFrame', but manipulates @rsp@ directly.
+--
+-- Asserts that the stack allocation is writable and large enough to contain the
+-- spilled arguments and return address.
+pushStackFrame ::
+  C.IsSymBackend sym bak =>
+  (?memOpts :: MM.MemOptions) =>
+  MM.HasLLVMAnn sym =>
+  bak ->
+  MM.MemImpl sym ->
+  -- | Assignment of values to registers
+  Ctx.Assignment (C.RegValue' sym) (MS.MacawCrucibleRegTypes X86.X86_64) ->
+  SpilledArgs sym ->
+  RetAddr sym ->
+  IO
+    ( Ctx.Assignment (C.RegValue' sym) (MS.MacawCrucibleRegTypes X86.X86_64)
+    , MM.MemImpl sym
+    )
+pushStackFrame bak mem regs spilledArgs retAddr = do
+  let sp = regs Lens.^. stackPointerReg
+  (sp', mem') <- allocStackFrame bak mem sp spilledArgs retAddr
+  let regs' = regs Lens.& stackPointerReg Lens..~ sp'
+  pure (regs', mem')
+
+-- | Like 'MSS.createArrayStack', but puts the stack pointer in @rsp@ directly.
+--
+-- Does not allow allocating stack slots, use 'allocStackFrame' or
+-- 'pushStackFrame' for that.
+allocStack ::
+  C.IsSymBackend sym bak =>
+  (?memOpts :: MM.MemOptions) =>
+  MM.HasLLVMAnn sym =>
+  bak ->
+  MM.MemImpl sym ->
+  -- | Assignment of values to registers
+  Ctx.Assignment (C.RegValue' sym) (MS.MacawCrucibleRegTypes X86.X86_64) ->
+  -- | Size of stack allocation
+  WI.SymExpr sym (WI.BaseBVType 64) ->
+  IO
+    ( Ctx.Assignment (C.RegValue' sym) (MS.MacawCrucibleRegTypes X86.X86_64)
+    , MM.MemImpl sym
+    )
+allocStack bak mem regs sz = do
+  let ?ptrWidth = ptrRepr
+  let slots = MSS.ExtraStackSlots 0
+  (MSS.ArrayStack _base top _arr, mem') <- MSS.createArrayStack bak mem slots sz
+  let regs' = regs Lens.& stackPointerReg Lens..~ StackPointer top
+  pure (regs', mem')