symbolic: Upstream stack spilled argument code from x86-symbolic

...so that it can also be used by AArch32.

symbolic/src/Data/Macaw/Symbolic/Stack.hs

@@ -12,13 +12,20 @@ module Data.Macaw.Symbolic.Stack
   , ExtraStackSlots(..)
   , ArrayStack(..)
   , createArrayStack
+  , allocStackSpace
+  , alignStackPointer
+  , writeSpilledArgs
   ) where
 
 import Data.BitVector.Sized qualified as BVS
 import Data.Parameterized.Context as Ctx
+import GHC.Natural (naturalToInteger)
 import Lang.Crucible.Backend qualified as C
 import Lang.Crucible.LLVM.DataLayout qualified as CLD
 import Lang.Crucible.LLVM.MemModel qualified as CLM
+import qualified Control.Monad as Monad
+import qualified Data.Sequence as Seq
+import qualified Lang.Crucible.LLVM.Bytes as Bytes
 import What4.Interface qualified as WI
 import What4.Symbol qualified as WSym
 
@@ -111,3 +118,118 @@ createArrayStack bak mem slots sz = do
   top <- CLM.ptrSub sym ?ptrWidth end slotsBytesBv
 
   pure (ArrayStack base top arr, mem2)
+
+-- | Allocate stack space by subtracting from the stack pointer
+--
+-- This is only suitable for use with architectures/ABIs where the stack grows
+-- down.
+allocStackSpace ::
+  C.IsSymInterface sym =>
+  CLM.HasPtrWidth w =>
+  sym ->
+  CLM.LLVMPtr sym w ->
+  WI.SymBV sym w ->
+  IO (CLM.LLVMPtr sym w)
+allocStackSpace sym = CLM.ptrSub sym ?ptrWidth
+
+-- | Align the stack pointer to a particular 'CLD.Alignment'.
+--
+-- This is only suitable for use with architectures/ABIs where the stack grows
+-- down.
+alignStackPointer ::
+  C.IsSymInterface sym =>
+  CLM.HasPtrWidth w =>
+  sym ->
+  CLD.Alignment ->
+  CLM.LLVMPtr sym w ->
+  IO (CLM.LLVMPtr sym w)
+alignStackPointer sym align orig@(CLM.LLVMPointer blk off) = do
+  let alignInt = 2 ^ naturalToInteger (CLD.alignmentToExponent align)
+  if alignInt == 1
+  then pure orig
+  else do
+    -- Because the stack grows down, we can align it by simply ANDing the stack
+    -- pointer with a mask with some amount of 1s followed by some amount of 0s.
+    let mask = BVS.complement ?ptrWidth (BVS.mkBV ?ptrWidth (alignInt - 1))
+    maskBv <- WI.bvLit sym ?ptrWidth mask
+    off' <- WI.bvAndBits sym off maskBv
+    pure (CLM.LLVMPointer blk off')
+
+-- | Write pointer-sized stack-spilled arguments to the stack.
+--
+-- This function:
+--
+-- * Aligns the stack to the minimum provided 'CLD.Alignment', call this M.
+-- * Potentially adds padding to ensure that after writing the arguments to the
+--   stack, the resulting stack pointer will be 2M-aligned.
+-- * Writes the stack-spilled arguments, in reverse order.
+--
+-- It is suitable for use with architectures/ABIs where the above is
+-- the desired behavior, e.g., AArch32 and SysV on x86_64. However,
+-- macaw-symbolic-{aarch32,x86} provide higher-level wrappers around this
+-- function, so its direct use is discouraged.
+--
+-- Asserts that the stack allocation is writable and large enough to contain the
+-- spilled arguments.
+writeSpilledArgs ::
+  C.IsSymBackend sym bak =>
+  (?memOpts :: CLM.MemOptions) =>
+  CLM.HasPtrWidth w =>
+  CLM.HasLLVMAnn sym =>
+  bak ->
+  CLM.MemImpl sym ->
+  -- | Minimum stack alignment
+  CLD.Alignment ->
+  CLM.LLVMPtr sym w ->
+  -- | Stack spilled arguments, in normal left-to-right order
+  Seq.Seq (CLM.LLVMPtr sym w) ->
+  IO (CLM.LLVMPtr sym w, CLM.MemImpl sym)
+writeSpilledArgs bak mem minAlign sp spilledArgs = do
+  let sym = C.backendGetSym bak
+
+  -- M is the minimum alignment, as an integer
+  let m = 2 ^ CLD.alignmentToExponent minAlign
+
+  -- First, align to M bytes. In all probability, this is a no-op.
+  spAlignedMin <- alignStackPointer sym minAlign sp
+
+  -- At this point, exactly one of `spAlignedMin` or `spAlignedMin +/- M` is
+  -- 2M-byte aligned. We need to ensure that, after writing the argument list,
+  -- the stack pointer will be 2M-byte aligned.
+  --
+  -- If `sp` is already 2M-byte aligned and there are an even number of spilled
+  -- arguments, we're good. If `sp + M` is already 2M-byte aligned and there
+  -- are an odd number of arguments, we're good. In the other cases, we need to
+  -- subtract M from `sp`. As a table:
+  --
+  -- +----------------------+------+------+
+  -- |                      | even | odd  |
+  -- |----------------------+------+------+
+  -- | 16-byte aligned      | noop | -M   |
+  -- +----------------------+-------------+
+  -- | not 16-byte aligned  | -M   | noop |
+  -- +----------------------+-------------+
+  --
+  let alignMore = CLD.exponentToAlignment (CLD.alignmentToExponent minAlign + 1)
+  spAlignedMore <- alignStackPointer sym alignMore sp
+  isAlignedMore <- CLM.ptrEq sym ?ptrWidth spAlignedMin spAlignedMore
+  sp' <-
+    if even (Seq.length spilledArgs)
+    then
+      -- In this case, either the stack pointer is *already* 2M-byte aligned, or
+      -- it *needs to be* 2M-byte aligned (which is equivalent to subtracting M,
+      -- as it is already M-byte aligned). In either case, this value suffices.
+      pure spAlignedMore
+    else do
+      mBv <- WI.bvLit sym ?ptrWidth (BVS.mkBV ?ptrWidth m)
+      spSubMore <- allocStackSpace sym spAlignedMore mBv
+      CLM.muxLLVMPtr sym isAlignedMore spSubMore spAlignedMin
+
+  let ptrBytes = WI.natValue ?ptrWidth `div` 8
+  ptrBytesBv <- WI.bvLit sym ?ptrWidth (BVS.mkBV ?ptrWidth (fromIntegral ptrBytes))
+  let ptrSz = CLM.bitvectorType (Bytes.toBytes ptrBytes)
+  let writeWord (p, mi) arg = do
+        p' <- CLM.ptrSub sym ?ptrWidth p ptrBytesBv
+        mi' <- CLM.storeRaw bak mi p' ptrSz CLD.noAlignment (CLM.ptrToPtrVal arg)
+        pure (p', mi')
+  Monad.foldM writeWord (sp', mem) (Seq.reverse spilledArgs)

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

@@ -76,8 +76,8 @@ module Data.Macaw.X86.Symbolic.ABI.SysV
   ) where
 
 import Control.Lens qualified as Lens
-import Control.Monad qualified as Monad
 import Data.BitVector.Sized qualified as BVS
+import Data.Coerce (coerce)
 import Data.Macaw.Symbolic qualified as MS
 import Data.Macaw.Symbolic.Stack qualified as MSS
 import Data.Macaw.X86 qualified as X86
@@ -91,6 +91,7 @@ import Lang.Crucible.LLVM.Bytes qualified as Bytes
 import Lang.Crucible.LLVM.DataLayout qualified as CLD
 import Lang.Crucible.LLVM.MemModel qualified as MM
 import Lang.Crucible.Simulator qualified as C
+import qualified Data.Macaw.Symbolic.Stack as Stack
 import What4.Interface qualified as WI
 
 -- | Helper, not exported
@@ -113,16 +114,6 @@ stackPointerReg =
     (\regs -> StackPointer (C.unRV (regs Lens.^. ixF' X86S.rsp)))
     (\regs v -> regs Lens.& ixF' X86S.rsp Lens..~ C.RV (getStackPointer v))
 
--- | Allocate stack space by subtracting from the stack pointer
-allocStackSpace ::
-  C.IsSymInterface sym =>
-  sym ->
-  StackPointer sym ->
-  WI.SymBV sym 64 ->
-  IO (StackPointer sym)
-allocStackSpace sym (StackPointer sp) amount =
-  StackPointer <$> MM.ptrSub sym ptrRepr sp amount
-
 -- | A return address
 newtype RetAddr sym = RetAddr { getRetAddr :: MM.LLVMPtr sym 64 }
 
@@ -176,51 +167,9 @@ writeSpilledArgs ::
   SpilledArgs sym ->
   IO (StackPointer sym, MM.MemImpl sym)
 writeSpilledArgs bak mem sp spilledArgs = do
-  let sym = C.backendGetSym bak
-  eight <- WI.bvLit sym ptrRepr (BVS.mkBV WI.knownNat 8)
-  let i64 = MM.bitvectorType (Bytes.toBytes (64 :: Int))
   let ?ptrWidth = ptrRepr
-
-  -- First, align to 8 bytes. In all probability, this is a no-op.
   let align8 = CLD.exponentToAlignment 3  -- 2^3 = 8
-  StackPointer spAligned8 <- alignStackPointer sym align8 sp
-
-  -- At this point, exactly one of `spAligned8` or `spAligned8 +/- 8` is 16-byte
-  -- aligned. We need to ensure that, after writing the argument list, the stack
-  -- pointer will be 16-byte aligned.
-  --
-  -- If `rsp` is already 16-byte aligned and there are an even number of spilled
-  -- arguments, we're good. If `rsp + 8` is already 16-byte aligned and there
-  -- are an odd number of arguments, we're good. In the other cases, we need to
-  -- subtract 8 from `rsp`. As a table:
-  --
-  -- +----------------------+------+------+
-  -- |                      | even | odd  |
-  -- |----------------------+------+------+
-  -- | 16-byte aligned      | noop | -8   |
-  -- +----------------------+-------------+
-  -- | not 16-byte aligned  | -8   | noop |
-  -- +----------------------+-------------+
-  --
-  let align16 = CLD.exponentToAlignment 4  -- 2^4 = 16
-  sp16@(StackPointer spAligned16) <- alignStackPointer sym align16 sp
-  isAligned16 <- MM.ptrEq sym ptrRepr spAligned8 spAligned16
-  sp' <-
-    if even (Seq.length (getSpilledArgs spilledArgs))
-    then
-      -- In this case, either the stack pointer is *already* 16-byte aligned, or
-      -- it *needs to be* 16-byte aligned (which is equivalent to subtracting 8,
-      -- as it is already 8-byte aligned). In either case, this value suffices.
-      pure sp16
-    else do
-      StackPointer spAdd8 <- allocStackSpace sym sp16 eight
-      StackPointer <$> MM.muxLLVMPtr sym isAligned16 spAdd8 spAligned8
-
-  let writeWord (StackPointer p, m) arg = 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 (getSpilledArgs spilledArgs))
+  coerce (Stack.writeSpilledArgs bak) mem align8 sp spilledArgs
 
 -- | Write the return address to the stack.
 --