content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- |
  -- Module      :   Grisette.Internal.Core.Data.Class.SymFiniteBits
  -- Copyright   :   (c) Sirui Lu 2024
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Internal.Core.Data.Class.SymFiniteBits
    ( lsb,
      msb,
      setBitTo,
      bitBlast,
      FromBits (..),
      SymFiniteBits (..),
      symBitBlast,
      symLsb,
      symMsb,
      symPopCount,
      symCountLeadingZeros,
      symCountTrailingZeros,
    )
  where
  
  import Data.Bits
    ( Bits (bit, clearBit, setBit, testBit, zeroBits, (.|.)),
      FiniteBits (finiteBitSize),
    )
  import Data.Int (Int16, Int32, Int64, Int8)
  import Data.Word (Word16, Word32, Word64, Word8)
  import GHC.TypeLits (KnownNat, type (<=))
  import Grisette.Internal.Core.Data.Class.BitVector
    ( BV (bv, bvConcat, bvSelect),
    )
  import Grisette.Internal.Core.Data.Class.ITEOp (ITEOp (symIte))
  import Grisette.Internal.Core.Data.Class.SymEq (SymEq ((.==)))
  import Grisette.Internal.SymPrim.BV (IntN, WordN)
  import Grisette.Internal.SymPrim.SomeBV
    ( SomeBV (SomeBV),
      SomeIntN,
      SomeWordN,
      unsafeSomeBV,
    )
  import Grisette.Internal.SymPrim.SymBV (SymIntN, SymWordN)
  import Grisette.Internal.SymPrim.SymBool (SymBool)
  import Unsafe.Coerce (unsafeCoerce)
  
  -- | Set a bit in a concrete value to a specific value.
  setBitTo :: (Bits a) => a -> Int -> Bool -> a
  setBitTo v i b = if b then setBit v i else clearBit v i
  
  -- | Bit-blast a concrete value into a list of concrete bits. The first element
  -- in the resulting list corresponds to the least significant bit.
  bitBlast :: (FiniteBits a) => a -> [Bool]
  bitBlast x = map (testBit x) [0 .. finiteBitSize x - 1]
  
  -- | Extract the least significant bit of a concrete value.
  lsb :: (Bits a) => a -> Bool
  lsb x = testBit x 0
  
  -- | Extract the most significant bit of a concrete value.
  msb :: (FiniteBits a) => a -> Bool
  msb x = testBit x (finiteBitSize x - 1)
  
  -- | Type class for assembling concrete bits to a bit-vector.
  class (FiniteBits a) => FromBits a where
    -- | Assembling concrete bits to a bit-vector. The first boolean value in the
    -- list corresponding to the least signification value.
    fromBits :: [Bool] -> a
    fromBits bits
      | length bits /= finiteBitSize (undefined :: a) =
          error "fromBits: length mismatch"
      | otherwise = foldl1 (.|.) lst
      where
        lst :: [a]
        lst = (\(pos, b) -> if b then bit pos else zeroBits) <$> zip [0 ..] bits
  
  instance FromBits Int
  
  instance FromBits Int8
  
  instance FromBits Int16
  
  instance FromBits Int32
  
  instance FromBits Int64
  
  instance FromBits Word
  
  instance FromBits Word8
  
  instance FromBits Word16
  
  instance FromBits Word32
  
  instance FromBits Word64
  
  instance (KnownNat n, 1 <= n) => FromBits (WordN n)
  
  instance (KnownNat n, 1 <= n) => FromBits (IntN n)
  
  instance FromBits SomeIntN where
    fromBits bits
      | null bits =
          error "Cannot create a SomeBV from an empty list of bits."
    fromBits bits = unsafeSomeBV (length bits) $ \_ -> fromBits bits
  
  instance FromBits SomeWordN where
    fromBits bits
      | null bits =
          error "Cannot create a SomeBV from an empty list of bits."
    fromBits bits = unsafeSomeBV (length bits) $ \_ -> fromBits bits
  
  -- | A class for symbolic finite bit operations.
  class (FiniteBits a, ITEOp a) => SymFiniteBits a where
    -- | Test a symbolic bit in a symbolic bit-vector.
    symTestBit :: a -> Int -> SymBool
  
    -- | Set a bit in a symbolic value to a specific value.
    symSetBitTo :: a -> Int -> SymBool -> a
    symSetBitTo v i b = symIte b (setBit v i) (clearBit v i)
  
    -- | Assembling symbolic bits to a symbolic bit-vector. The first symbolic
    -- boolean value in the list corresponding to the least signification value.
    symFromBits :: [SymBool] -> a
  
  instance SymFiniteBits (SomeBV SymIntN) where
    symTestBit v i = bvSelect i 1 v .== bv 1 1
    symFromBits bits
      | null bits =
          error "Cannot create a SomeBV from an empty list of bits."
    symFromBits bits = unsafeSomeBV (length bits) $ \_ -> symFromBits bits
  
  instance SymFiniteBits (SomeBV SymWordN) where
    symTestBit v i = bvSelect i 1 v .== bv 1 1
    symFromBits bits
      | null bits =
          error "Cannot create a SomeBV from an empty list of bits."
    symFromBits bits = unsafeSomeBV (length bits) $ \_ -> symFromBits bits
  
  instance (KnownNat n, 1 <= n) => SymFiniteBits (SymIntN n) where
    symTestBit v = symTestBit (SomeBV v)
    symSetBitTo v i b = symIte b (setBit v i) (clearBit v i)
    symFromBits bits
      | length bits /= finiteBitSize (undefined :: SymWordN n) =
          error "symFromBits: length mismatch"
      | otherwise = case foldl1 bvConcat $ SomeBV <$> reverse lst of
          SomeBV r -> unsafeCoerce r
          _ -> error "symFromBits: length mismatch"
      where
        lst :: [SymIntN 1]
        lst = (\b -> symIte b 1 0) <$> bits
  
  instance (KnownNat n, 1 <= n) => SymFiniteBits (SymWordN n) where
    symTestBit v = symTestBit (SomeBV v)
    symSetBitTo v i b = symIte b (setBit v i) (clearBit v i)
    symFromBits bits
      | length bits /= finiteBitSize (undefined :: SymWordN n) =
          error "symFromBits: length mismatch"
      | otherwise = case foldl1 bvConcat $ SomeBV <$> reverse lst of
          SomeBV r -> unsafeCoerce r
          _ -> error "symFromBits: length mismatch"
      where
        lst :: [SymWordN 1]
        lst = (\b -> symIte b 1 0) <$> bits
  
  -- | Bit-blast a symbolic value into a list of symbolic bits. The first element
  -- in the resulting list corresponds to the least significant bit.
  symBitBlast :: (SymFiniteBits a) => a -> [SymBool]
  symBitBlast x = map (symTestBit x) [0 .. finiteBitSize x - 1]
  
  -- | Extract the least significant bit of a symbolic value.
  symLsb :: (SymFiniteBits a) => a -> SymBool
  symLsb x = symTestBit x 0
  
  -- | Extract the most significant bit of a symbolic value.
  symMsb :: (SymFiniteBits a) => a -> SymBool
  symMsb x = symTestBit x (finiteBitSize x - 1)
  
  -- | Count the number of set bits in a symbolic value.
  symPopCount :: (Num a, ITEOp a, SymFiniteBits a) => a -> a
  symPopCount v = 0 * v + adderTree (map fromBool (symBitBlast v))
    where
      -- Convert a symbolic boolean to a numeric value (0 or 1)
      -- Node: 0 * v + is a trick to assign the correct bit-width to the result.
      fromBool b = 0 * v + symIte b 1 0
  
      -- Implement binary adder tree to efficiently sum values
      adderTree :: (Num a) => [a] -> a
      adderTree [] = 0
      adderTree [x] = x
      adderTree xs = adderTree (pairwiseAdd xs)
  
      -- Add adjacent pairs of values
      pairwiseAdd :: (Num a) => [a] -> [a]
      pairwiseAdd [] = []
      pairwiseAdd [x] = [x]
      pairwiseAdd (x : y : rest) = (x + y) : pairwiseAdd rest
  
  -- | Count the number of leading zeros in a symbolic value.
  symCountLeadingZeros :: (Num a, ITEOp a, SymFiniteBits a) => a -> a
  -- Node: 0 * v + is a trick to assign the correct bit-width to the result.
  symCountLeadingZeros v = 0 * v + go bits rs
    where
      bits = reverse $ symBitBlast v
      rs = fromIntegral <$> [0 ..]
      go [] (r : _) = r
      go (b : bs) (r : rs) = symIte b r (go bs rs)
      go _ [] = error "Should not happen"
  
  -- | Count the number of trailing zeros in a symbolic value.
  symCountTrailingZeros :: (Num a, ITEOp a, SymFiniteBits a) => a -> a
  -- Node: 0 * v + is a trick to assign the correct bit-width to the result.
  symCountTrailingZeros v = 0 * v + go bits rs
    where
      bits = symBitBlast v
      rs = fromIntegral <$> [0 ..]
      go [] (r : _) = r
      go (b : bs) (r : rs) = symIte b r (go bs rs)
      go _ [] = error "Should not happen"
  