Expression.hs

-- SPDX-FileCopyrightText: 2025 Sören Tempel <soeren+git@soeren-tempel.net>
--
-- SPDX-License-Identifier: GPL-3.0-only
{-# LANGUAGE TemplateHaskell #-}
-- The code generated by template-haskell does not have type signatures.
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

-- | This module provides an implementation of the expression abstract from
-- 'Language.QBE.Simulator.Expression' which uses concrete fixed-width integer
-- values from "Data.Word" internally.
module Language.QBE.Simulator.Default.Expression
  ( RegVal (..),
    bitSize,
    fromBits,
  )
where

import Control.Exception (assert)
import Data.Bits
  ( FiniteBits,
    finiteBitSize,
    shift,
    shiftR,
    unsafeShiftL,
    unsafeShiftR,
    xor,
    (.&.),
    (.|.),
  )
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Word (Word16, Word32, Word64, Word8)
import GHC.Float
  ( castDoubleToWord64,
    castFloatToWord32,
    castWord32ToFloat,
    castWord64ToDouble,
    double2Float,
    float2Double,
  )
import Language.QBE.Simulator.Default.Generator (generateOperators)
import Language.QBE.Simulator.Expression qualified as E
import Language.QBE.Simulator.Memory qualified as MEM
import Language.QBE.Types qualified as QBE

-- TODO: Can we just wrap base type here?
-- TODO: Do not export the constructors
data RegVal
  = VByte Word8
  | VHalf Word16
  | VWord Word32
  | VLong Word64
  | VSingle Float
  | VDouble Double
  deriving (Show, Eq)

-- | Size of the value in bits.
bitSize :: RegVal -> Int
bitSize (VByte _) = 8
bitSize (VHalf _) = 16
bitSize (VWord _) = 32
bitSize (VLong _) = 64
bitSize (VSingle _) = 32
bitSize (VDouble _) = 64

-- | Create a a 'RegVal' from an 'Integer' inferring the type from a given
-- amount of bits instead of requiring the user to provide a 'QBE.ExtType',
-- as required by 'Language.QBE.Simulator.Expression.fromLit'.
fromBits :: Int -> Integer -> Maybe RegVal
fromBits 08 = Just . VHalf . fromIntegral
fromBits 16 = Just . VHalf . fromIntegral
fromBits 32 = Just . VWord . fromIntegral
fromBits 64 = Just . VLong . fromIntegral
fromBits _ = const Nothing

fromBool :: Bool -> RegVal
fromBool True = VLong 1
fromBool False = VLong 0

------------------------------------------------------------------------

shiftInstr ::
  (RegVal -> Word32 -> Maybe RegVal) ->
  RegVal ->
  RegVal ->
  Maybe RegVal
shiftInstr shiftOp val (VWord amount) = val `shiftOp` amount
shiftInstr _ _ _ = Nothing

toShiftAmount :: Word32 -> Word32 -> Int
toShiftAmount valBitSize amount =
  -- From the QBE specification: "The shifting amount
  -- is taken modulo the size of the result type."
  let s = fromIntegral $ amount `mod` valBitSize
   in assert (s > 0) s

shiftSar :: RegVal -> Word32 -> Maybe RegVal
shiftSar (VWord val) amount =
  (Just . VWord . fromIntegral) $
    (fromIntegral val :: Int32) `unsafeShiftR` toShiftAmount 32 amount
shiftSar (VLong val) amount =
  (Just . VLong . fromIntegral) $
    (fromIntegral val :: Int64) `unsafeShiftR` toShiftAmount 64 amount
shiftSar _ _ = Nothing

shiftShr :: RegVal -> Word32 -> Maybe RegVal
shiftShr (VWord val) amount =
  (Just . VWord) $ val `unsafeShiftR` toShiftAmount 32 amount
shiftShr (VLong val) amount =
  (Just . VLong) $ val `unsafeShiftR` toShiftAmount 64 amount
shiftShr _ _ = Nothing

shiftShl :: RegVal -> Word32 -> Maybe RegVal
shiftShl (VWord val) amount =
  (Just . VWord) $ val `unsafeShiftL` toShiftAmount 32 amount
shiftShl (VLong val) amount =
  (Just . VLong) $ val `unsafeShiftL` toShiftAmount 64 amount
shiftShl _ _ = Nothing

------------------------------------------------------------------------

regToBytes :: RegVal -> [Word8]
regToBytes val =
  let f w =
        map
          (\off -> fromIntegral $ shiftR w off .&. 0xff)
          (take (bytesize w) $ iterate (+ 8) 0)
   in case val of
        (VByte v) -> [v]
        (VWord v) -> f v
        (VHalf v) -> f v
        (VLong v) -> f v
        (VSingle v) -> MEM.toBytes (VWord $ castFloatToWord32 v)
        (VDouble v) -> MEM.toBytes (VLong $ castDoubleToWord64 v)
  where
    bytesize :: (FiniteBits a) => a -> Int
    bytesize v = finiteBitSize v `div` 8

regFromBytes :: QBE.LoadType -> [Word8] -> Maybe RegVal
regFromBytes ty lst =
  let f a =
        foldl
          (\acc (byte, idx) -> (fromIntegral byte `shift` (idx * 8)) .|. acc)
          0
          $ zip a [0 ..]
   in case (ty, lst) of
        (QBE.LSubWord QBE.UnsignedByte, [byte]) -> Just (VWord (fromIntegral byte))
        (QBE.LSubWord QBE.SignedByte, [byte]) -> Just (VWord $ fromIntegral (fromIntegral byte :: Int8))
        (QBE.LSubWord QBE.SignedHalf, bytes@[_, _]) -> Just (VWord $ fromIntegral (f bytes :: Int16))
        (QBE.LSubWord QBE.UnsignedHalf, bytes@[_, _]) -> Just (VWord $ fromIntegral (f bytes :: Word16))
        (QBE.LBase QBE.Word, bytes@[_, _, _, _]) -> Just (VWord $ f bytes)
        (QBE.LBase QBE.Long, bytes@[_, _, _, _, _, _, _, _]) -> Just (VLong $ f bytes)
        (QBE.LBase QBE.Single, bytes@[_, _, _, _]) ->
          Just (VSingle $ castWord32ToFloat (f bytes))
        (QBE.LBase QBE.Double, bytes@[_, _, _, _, _, _, _, _]) ->
          Just (VDouble $ castWord64ToDouble (f bytes))
        _ -> Nothing

instance MEM.Storable RegVal Word8 where
  toBytes = regToBytes
  fromBytes = regFromBytes

------------------------------------------------------------------------

-- TODO: Insert the generated code directly into the instance declaration.
generateOperators

maxValue :: RegVal -> Maybe RegVal
maxValue val =
  let bitSiz = bitSize val
      maxVal = (2 ^ bitSiz) - 1
   in fromBits bitSiz maxVal

withZeroDiv ::
  Maybe RegVal ->
  (RegVal -> RegVal -> Maybe RegVal) ->
  RegVal ->
  RegVal ->
  Maybe RegVal
withZeroDiv defVal op lhs rhs
  | E.toWord64 rhs == 0 = defVal
  | otherwise = op lhs rhs

-- Signed division overflow occurs when the most-negative integer is divided by -1.
withSDivOverflow ::
  Maybe RegVal ->
  (RegVal -> RegVal -> Maybe RegVal) ->
  RegVal ->
  RegVal ->
  Maybe RegVal
withSDivOverflow defVal op lhs rhs
  | E.toWord64 lhs == mostNeg && E.toWord64 rhs == minusOne = defVal
  | otherwise = op lhs rhs
  where
    numBits :: Int
    numBits =
      assert (bitSize lhs == bitSize rhs) $
        bitSize lhs

    minusOne :: Word64
    minusOne = (2 ^ numBits) - 1

    mostNeg :: Word64
    mostNeg = 2 ^ (numBits - 1)

-- We could also add support for unary operators to the generator. However,
-- presently there is only one unary operator so it isn't worth it.
neg' :: RegVal -> Maybe RegVal
neg' (VWord v) = Just . VWord $ negate v
neg' (VLong v) = Just . VLong $ negate v
neg' (VSingle v) = Just . VSingle $ negate v
neg' (VDouble v) = Just . VDouble $ negate v
neg' _ = Nothing

-- This can't be easily auto generated because the operation differs
-- based on the type.
div' :: RegVal -> RegVal -> Maybe RegVal
div' (VWord lhs) (VWord rhs) =
  (Just . VWord . fromIntegral) $
    (fromIntegral lhs :: Int32) `quot` (fromIntegral rhs :: Int32)
div' (VLong lhs) (VLong rhs) =
  (Just . VLong . fromIntegral) $
    (fromIntegral lhs :: Int64) `quot` (fromIntegral rhs :: Int64)
div' (VSingle lhs) (VSingle rhs) = (Just . VSingle) $ lhs / rhs
div' (VDouble lhs) (VDouble rhs) = (Just . VDouble) $ lhs / rhs
div' _ _ = Nothing

instance E.ValueRepr RegVal where
  fromLit QBE.Byte n = VByte $ fromIntegral n
  fromLit QBE.HalfWord n = VHalf $ fromIntegral n
  fromLit (QBE.Base QBE.Long) n = VLong n
  fromLit (QBE.Base QBE.Word) n = VWord $ fromIntegral n
  fromLit (QBE.Base QBE.Single) n = VSingle $ castWord32ToFloat (fromIntegral n)
  fromLit (QBE.Base QBE.Double) n = VDouble $ castWord64ToDouble n

  toWord64 (VByte v) = fromIntegral v
  toWord64 (VHalf v) = fromIntegral v
  toWord64 (VWord v) = fromIntegral v
  toWord64 (VLong v) = v
  toWord64 (VSingle v) = fromIntegral $ castFloatToWord32 v
  toWord64 (VDouble v) = castDoubleToWord64 v

  fromFloat = VSingle
  fromDouble = VDouble

  -- stosi
  floatToInt ty@(QBE.Base QBE.Word) True (VSingle v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Int32))
  floatToInt ty@(QBE.Base QBE.Long) True (VSingle v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Int64))
  -- stoui
  floatToInt ty@(QBE.Base QBE.Word) False (VSingle v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Word32))
  floatToInt ty@(QBE.Base QBE.Long) False (VSingle v) =
    Just $ E.fromLit ty (truncate v :: Word64)
  -- dtosi
  floatToInt ty@(QBE.Base QBE.Word) True (VDouble v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Int32))
  floatToInt ty@(QBE.Base QBE.Long) True (VDouble v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Int64))
  -- dtoui
  floatToInt ty@(QBE.Base QBE.Word) False (VDouble v) =
    Just $ E.fromLit ty (fromIntegral (truncate v :: Word32))
  floatToInt ty@(QBE.Base QBE.Long) False (VDouble v) =
    Just $ E.fromLit ty (truncate v :: Word64)
  -- rest
  floatToInt _ _ _ = Nothing

  -- swtof
  intToFloat ty@(QBE.Base QBE.Single) True (VWord v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Int32))
  intToFloat ty@(QBE.Base QBE.Double) True (VWord v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Int32))
  -- uwtof
  intToFloat ty@(QBE.Base QBE.Single) False (VWord v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Word32))
  intToFloat ty@(QBE.Base QBE.Double) False (VWord v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Word32))
  -- sltof
  intToFloat ty@(QBE.Base QBE.Single) True (VLong v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Int64))
  intToFloat ty@(QBE.Base QBE.Double) True (VLong v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Int64))
  -- ultof
  intToFloat ty@(QBE.Base QBE.Single) False (VLong v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Word64))
  intToFloat ty@(QBE.Base QBE.Double) False (VLong v) =
    Just $ E.fromLit ty (fromIntegral (fromIntegral v :: Word64))
  -- rest
  intToFloat _ _ _ = Nothing

  extendFloat (VSingle v) = Just $ VDouble (float2Double v)
  extendFloat _ = Nothing

  truncFloat (VDouble v) = Just $ VSingle (double2Float v)
  truncFloat _ = Nothing

  getType (VByte _) = QBE.Byte
  getType (VHalf _) = QBE.HalfWord
  getType (VWord _) = QBE.Base QBE.Word
  getType (VLong _) = QBE.Base QBE.Long
  getType (VSingle _) = QBE.Base QBE.Single
  getType (VDouble _) = QBE.Base QBE.Double

  -- TODO: Consider replacing Nothing cases with assert as this on the hot path.
  extend extTy isSigned val
    | QBE.extTypeBitSize extTy <= bitSize val = Nothing
    | otherwise =
        E.fromLit extTy
          <$> case (isSigned, val) of
            (True, VByte v) -> Just $ fromIntegral (fromIntegral v :: Int8)
            (True, VHalf v) -> Just $ fromIntegral (fromIntegral v :: Int16)
            (True, VWord v) -> Just $ fromIntegral (fromIntegral v :: Int32)
            (True, VLong v) -> Just $ fromIntegral (fromIntegral v :: Int64)
            (False, VByte v) -> Just $ fromIntegral (fromIntegral v :: Word8)
            (False, VHalf v) -> Just $ fromIntegral (fromIntegral v :: Word16)
            (False, VWord v) -> Just $ fromIntegral (fromIntegral v :: Word32)
            (False, VLong v) -> Just $ fromIntegral (fromIntegral v :: Word64)
            _ -> Nothing

  -- TODO: Consider replacing Nothing cases with assert as this on the hot path.
  extract (QBE.Base QBE.Single) _ = Nothing
  extract (QBE.Base QBE.Double) _ = Nothing
  extract _ (VSingle _) = Nothing
  extract _ (VDouble _) = Nothing
  extract extTy v
    | QBE.extTypeBitSize extTy > bitSize v = Nothing
    | otherwise =
        let word = E.toWord64 v
            mask = (2 `unsafeShiftL` (QBE.extTypeBitSize extTy - 1)) - 1
         in Just $ E.fromLit extTy (word .&. mask)

  -- This is needed to align the behavior of qute/ and qute-symex/ on
  -- division-by-zero. QBE does not explicitly mandate a specific behavior
  -- for this edge case. Therefore, in order to avoid extra branches in the
  -- symbolic executor, we use the behavior mandated by SMT-LIB here.
  --
  -- TODO: Move this into the Expression abstraction (just like overshift handling).
  div lhs = withZeroDiv (maxValue lhs) (withSDivOverflow (Just lhs) div') lhs
  udiv lhs = withZeroDiv (maxValue lhs) udiv' lhs
  urem lhs = withZeroDiv (Just lhs) urem' lhs
  srem lhs = withZeroDiv (Just lhs) (withSDivOverflow (fromBits (bitSize lhs) 0) srem') lhs

  add = add'
  sub = sub'
  mul = mul'
  or = or'
  xor = xor'
  and = and'

  neg = neg'

  sar = shiftInstr shiftSar
  shr = shiftInstr shiftShr
  shl = shiftInstr shiftShl

  -- TODO: Provide default implementations
  eq = eq'
  ne = ne'
  sle = sle'
  slt = slt'
  sge = sge'
  sgt = sgt'
  ule = ule'
  ult = ult'
  uge = uge'
  ugt = ugt'

  ord (VSingle lhs) (VSingle rhs) =
    Just . fromBool $ not (isNaN lhs || isNaN rhs)
  ord (VDouble lhs) (VDouble rhs) =
    Just . fromBool $ not (isNaN lhs || isNaN rhs)
  ord _ _ = Nothing