{-# LANGUAGE UndecidableInstances #-}

module Plutarch.BitString (
  -- * Type
  PBitString (..),

  -- * Functions
  preadBit,
  pwriteBits,
  pshift,
  protate,
  pcountSetBits,
  pfindFirstSetBit,
  pfindFirstSetBit',
) where

import Data.ByteString (ByteString)
import GHC.Generics (Generic)
import Generics.SOP qualified as SOP
import Plutarch.Builtin.Bool (PBool, pif)
import Plutarch.Builtin.ByteString (PByteString)
import Plutarch.Builtin.Data (PBuiltinList)
import Plutarch.Builtin.Integer (PInteger)
import Plutarch.Internal.Eq (PEq)
import Plutarch.Internal.Lift (
  DeriveNewtypePLiftable,
  PLiftable,
  PLifted (PLifted),
 )
import Plutarch.Internal.Numeric (pzero)
import Plutarch.Internal.Ord (POrd ((#<)))
import Plutarch.Internal.Other (pto)
import Plutarch.Internal.PLam (plam)
import Plutarch.Internal.PlutusType (
  PlutusType,
  pcon,
 )
import Plutarch.Internal.Semigroup (PMonoid, PSemigroup)
import Plutarch.Internal.Term (
  S,
  Term,
  phoistAcyclic,
  plet,
  punsafeBuiltin,
  (#),
  (:-->),
 )
import Plutarch.Maybe (PMaybe (PJust, PNothing))
import Plutarch.Repr.Newtype (DeriveNewtypePlutusType (DeriveNewtypePlutusType))
import PlutusCore qualified as PLC

{- | A wrapper around 'PByteString' for CIP-122 and CIP-123 bitwise operations.

= Note

This type exists because /bit/ and /byte/ indexes work in different
directions. To avoid confusing behaviour, we require an explicit wrapping of
'PByteString's to use bitwise functionality: this way, it's clear where which
scheme applies.

@since 1.10.0
-}
newtype PBitString (s :: S) = PBitString (Term s PByteString)
  deriving stock
    ( -- | @since 1.10.0
      (forall x. PBitString s -> Rep (PBitString s) x)
-> (forall x. Rep (PBitString s) x -> PBitString s)
-> Generic (PBitString s)
forall x. Rep (PBitString s) x -> PBitString s
forall x. PBitString s -> Rep (PBitString s) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (s :: S) x. Rep (PBitString s) x -> PBitString s
forall (s :: S) x. PBitString s -> Rep (PBitString s) x
$cfrom :: forall (s :: S) x. PBitString s -> Rep (PBitString s) x
from :: forall x. PBitString s -> Rep (PBitString s) x
$cto :: forall (s :: S) x. Rep (PBitString s) x -> PBitString s
to :: forall x. Rep (PBitString s) x -> PBitString s
Generic
    )
  deriving anyclass
    ( -- | @since 1.10.0
      All @[Type] (SListI @Type) (Code (PBitString s))
All @[Type] (SListI @Type) (Code (PBitString s)) =>
(PBitString s -> Rep (PBitString s))
-> (Rep (PBitString s) -> PBitString s) -> Generic (PBitString s)
Rep (PBitString s) -> PBitString s
PBitString s -> Rep (PBitString s)
forall a.
All @[Type] (SListI @Type) (Code a) =>
(a -> Rep a) -> (Rep a -> a) -> Generic a
forall (s :: S). All @[Type] (SListI @Type) (Code (PBitString s))
forall (s :: S). Rep (PBitString s) -> PBitString s
forall (s :: S). PBitString s -> Rep (PBitString s)
$cfrom :: forall (s :: S). PBitString s -> Rep (PBitString s)
from :: PBitString s -> Rep (PBitString s)
$cto :: forall (s :: S). Rep (PBitString s) -> PBitString s
to :: Rep (PBitString s) -> PBitString s
SOP.Generic
    , -- | @since 1.10.0
      (forall (s :: S).
 Term s PBitString -> Term s PBitString -> Term s PBool)
-> PEq PBitString
forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
forall (t :: PType).
(forall (s :: S). Term s t -> Term s t -> Term s PBool) -> PEq t
$c#== :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
#== :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
PEq
    , -- | @since 1.10.0
      PEq PBitString
PEq PBitString =>
(forall (s :: S).
 Term s PBitString -> Term s PBitString -> Term s PBool)
-> (forall (s :: S).
    Term s PBitString -> Term s PBitString -> Term s PBool)
-> (forall (s :: S).
    Term s PBitString -> Term s PBitString -> Term s PBitString)
-> (forall (s :: S).
    Term s PBitString -> Term s PBitString -> Term s PBitString)
-> POrd PBitString
forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
forall (t :: PType).
PEq t =>
(forall (s :: S). Term s t -> Term s t -> Term s PBool)
-> (forall (s :: S). Term s t -> Term s t -> Term s PBool)
-> (forall (s :: S). Term s t -> Term s t -> Term s t)
-> (forall (s :: S). Term s t -> Term s t -> Term s t)
-> POrd t
$c#<= :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
#<= :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
$c#< :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
#< :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBool
$cpmax :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
pmax :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
$cpmin :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
pmin :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
POrd
    , -- | @since 1.10.0
      (forall (s :: S).
 Term s PBitString -> Term s PBitString -> Term s PBitString)
-> (forall (s :: S).
    Term s PPositive -> Term s PBitString -> Term s PBitString)
-> PSemigroup PBitString
forall (s :: S).
Term s PPositive -> Term s PBitString -> Term s PBitString
forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
forall (a :: PType).
(forall (s :: S). Term s a -> Term s a -> Term s a)
-> (forall (s :: S). Term s PPositive -> Term s a -> Term s a)
-> PSemigroup a
$c#<> :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
#<> :: forall (s :: S).
Term s PBitString -> Term s PBitString -> Term s PBitString
$cpstimes :: forall (s :: S).
Term s PPositive -> Term s PBitString -> Term s PBitString
pstimes :: forall (s :: S).
Term s PPositive -> Term s PBitString -> Term s PBitString
PSemigroup
    , -- | @since 1.10.0
      PSemigroup PBitString
PSemigroup PBitString =>
(forall (s :: S). Term s PBitString)
-> (forall (s :: S).
    Term s PNatural -> Term s PBitString -> Term s PBitString)
-> PMonoid PBitString
forall (s :: S). Term s PBitString
forall (s :: S).
Term s PNatural -> Term s PBitString -> Term s PBitString
forall (a :: PType).
PSemigroup a =>
(forall (s :: S). Term s a)
-> (forall (s :: S). Term s PNatural -> Term s a -> Term s a)
-> PMonoid a
$cpmempty :: forall (s :: S). Term s PBitString
pmempty :: forall (s :: S). Term s PBitString
$cpmtimes :: forall (s :: S).
Term s PNatural -> Term s PBitString -> Term s PBitString
pmtimes :: forall (s :: S).
Term s PNatural -> Term s PBitString -> Term s PBitString
PMonoid
    )
  deriving
    ( -- | @since 1.10.0
      (forall (s :: S). PBitString s -> Term s (PInner PBitString))
-> (forall (s :: S) (b :: PType).
    Term s (PInner PBitString)
    -> (PBitString s -> Term s b) -> Term s b)
-> PlutusType PBitString
forall (s :: S). PBitString s -> Term s (PInner PBitString)
forall (s :: S) (b :: PType).
Term s (PInner PBitString)
-> (PBitString s -> Term s b) -> Term s b
forall (a :: PType).
(forall (s :: S). a s -> Term s (PInner a))
-> (forall (s :: S) (b :: PType).
    Term s (PInner a) -> (a s -> Term s b) -> Term s b)
-> PlutusType a
$cpcon' :: forall (s :: S). PBitString s -> Term s (PInner PBitString)
pcon' :: forall (s :: S). PBitString s -> Term s (PInner PBitString)
$cpmatch' :: forall (s :: S) (b :: PType).
Term s (PInner PBitString)
-> (PBitString s -> Term s b) -> Term s b
pmatch' :: forall (s :: S) (b :: PType).
Term s (PInner PBitString)
-> (PBitString s -> Term s b) -> Term s b
PlutusType
    )
    via (DeriveNewtypePlutusType PBitString)

-- | @since 1.10.0
deriving via
  DeriveNewtypePLiftable PBitString ByteString
  instance
    PLiftable PBitString

{- | Bit access operation, as defined in [CIP-122](https://github.com/cardano-foundation/CIPs/tree/master/CIP-0122#readbit).

@since 1.10.0
-}
preadBit :: forall (s :: S). Term s (PBitString :--> PInteger :--> PBool)
preadBit :: forall (s :: S). Term s (PBitString :--> (PInteger :--> PBool))
preadBit = DefaultFun -> Term s (PBitString :--> (PInteger :--> PBool))
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.ReadBit

{- | Given a list of positions, set the bits at those positions.

This works similarly to the @writeBits@ operation in
[CIP-122](https://github.com/cardano-foundation/CIPs/tree/master/CIP-0122#writebits)
with regard to the list of indexes. Effectively, @psetBits b ixes@ is
equivalent to @writeBits b . map (, True) $ ixes@. All caveats that this
entails from the CIP-122 description apply.

@since 1.10.0
-}

{- | Sets bits, as per
[CIP-122](https://github.com/cardano-foundation/CIPs/tree/master/CIP-0122#writebits).

@since 1.10.0
-}
pwriteBits :: forall (s :: S). Term s (PBitString :--> PBuiltinList PInteger :--> PBool :--> PBitString)
pwriteBits :: forall (s :: S).
Term
  s
  (PBitString
   :--> (PBuiltinList PInteger :--> (PBool :--> PBitString)))
pwriteBits = DefaultFun
-> Term
     s
     (PBitString
      :--> (PBuiltinList PInteger :--> (PBool :--> PBitString)))
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.WriteBits

{- | Performs a shift, as per
[CIP-123](https://github.com/cardano-foundation/CIPs/blob/master/CIP-0123/README.md#bitwiseshift).

@since 1.10.0
-}
pshift :: forall (s :: S). Term s (PBitString :--> PInteger :--> PBitString)
pshift :: forall (s :: S).
Term s (PBitString :--> (PInteger :--> PBitString))
pshift = DefaultFun -> Term s (PBitString :--> (PInteger :--> PBitString))
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.ShiftByteString

{- | Performs a rotation, as per
[CIP-123](https://github.com/cardano-foundation/CIPs/blob/master/CIP-0123/README.md#bitwiserotate).

@since 1.10.0
-}
protate :: forall (s :: S). Term s (PBitString :--> PInteger :--> PBitString)
protate :: forall (s :: S).
Term s (PBitString :--> (PInteger :--> PBitString))
protate = DefaultFun -> Term s (PBitString :--> (PInteger :--> PBitString))
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.RotateByteString

{- | Counts the number of set bits, as per
[CIP-123](https://github.com/cardano-foundation/CIPs/blob/master/CIP-0123/README.md#countsetbits).

@since 1.10.0
-}
pcountSetBits :: forall (s :: S). Term s (PBitString :--> PInteger)
pcountSetBits :: forall (s :: S). Term s (PBitString :--> PInteger)
pcountSetBits = DefaultFun -> Term s (PBitString :--> PInteger)
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.CountSetBits

{- | Finds the index of the first set bit, as per
[CIP-123](https://github.com/cardano-foundation/CIPs/blob/master/CIP-0123/README.md#findfirstsetbit).

= Note

This returns @-1@ if the argument is either empty, or contains no set bits.

@since 1.10.0
-}
pfindFirstSetBit' :: forall (s :: S). Term s (PBitString :--> PInteger)
pfindFirstSetBit' :: forall (s :: S). Term s (PBitString :--> PInteger)
pfindFirstSetBit' = DefaultFun -> Term s (PBitString :--> PInteger)
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.FindFirstSetBit

{- | As @pfindFirstSetBit'@, but produces 'PNothing' if the argument is empty,
or contains no set bits.

@since 1.10.0
-}
pfindFirstSetBit :: forall (s :: S). Term s (PBitString :--> PMaybe PInteger)
pfindFirstSetBit :: forall (s :: S). Term s (PBitString :--> PMaybe PInteger)
pfindFirstSetBit = (forall (s :: S). Term s (PBitString :--> PMaybe PInteger))
-> Term s (PBitString :--> PMaybe PInteger)
forall (a :: PType) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PBitString :--> PMaybe PInteger))
 -> Term s (PBitString :--> PMaybe PInteger))
-> (forall (s :: S). Term s (PBitString :--> PMaybe PInteger))
-> Term s (PBitString :--> PMaybe PInteger)
forall a b. (a -> b) -> a -> b
$ (Term s PBitString -> Term s (PMaybe PInteger))
-> Term s (PBitString :--> PMaybe PInteger)
forall a (b :: PType) (s :: S) (c :: PType).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: PType).
HasCallStack =>
(Term s c -> Term s (PMaybe PInteger))
-> Term s (c :--> PMaybe PInteger)
plam ((Term s PBitString -> Term s (PMaybe PInteger))
 -> Term s (PBitString :--> PMaybe PInteger))
-> (Term s PBitString -> Term s (PMaybe PInteger))
-> Term s (PBitString :--> PMaybe PInteger)
forall a b. (a -> b) -> a -> b
$ \Term s PBitString
bs ->
  Term s PInteger
-> (Term s PInteger -> Term s (PMaybe PInteger))
-> Term s (PMaybe PInteger)
forall (s :: S) (a :: PType) (b :: PType).
Term s a -> (Term s a -> Term s b) -> Term s b
plet (DefaultFun -> Term s (PInner PBitString :--> PInteger)
forall (s :: S) (a :: PType). DefaultFun -> Term s a
punsafeBuiltin DefaultFun
PLC.FindFirstSetBit Term s (PInner PBitString :--> PInteger)
-> Term s (PInner PBitString) -> Term s PInteger
forall (s :: S) (a :: PType) (b :: PType).
Term s (a :--> b) -> Term s a -> Term s b
# Term s PBitString -> Term s (PInner PBitString)
forall (s :: S) (a :: PType). Term s a -> Term s (PInner a)
pto Term s PBitString
bs) ((Term s PInteger -> Term s (PMaybe PInteger))
 -> Term s (PMaybe PInteger))
-> (Term s PInteger -> Term s (PMaybe PInteger))
-> Term s (PMaybe PInteger)
forall a b. (a -> b) -> a -> b
$ \Term s PInteger
result ->
    Term s PBool
-> Term s (PMaybe PInteger)
-> Term s (PMaybe PInteger)
-> Term s (PMaybe PInteger)
forall (a :: PType) (s :: S).
Term s PBool -> Term s a -> Term s a -> Term s a
pif
      (Term s PInteger
result Term s PInteger -> Term s PInteger -> Term s PBool
forall (s :: S). Term s PInteger -> Term s PInteger -> Term s PBool
forall (t :: PType) (s :: S).
POrd t =>
Term s t -> Term s t -> Term s PBool
#< Term s PInteger
forall (s :: S). Term s PInteger
forall (a :: PType) (s :: S). PAdditiveMonoid a => Term s a
pzero)
      (PMaybe PInteger s -> Term s (PMaybe PInteger)
forall (a :: PType) (s :: S). PlutusType a => a s -> Term s a
pcon PMaybe PInteger s
forall (a :: PType) (s :: S). PMaybe a s
PNothing)
      (PMaybe PInteger s -> Term s (PMaybe PInteger)
forall (a :: PType) (s :: S). PlutusType a => a s -> Term s a
pcon (PMaybe PInteger s -> Term s (PMaybe PInteger))
-> (Term s PInteger -> PMaybe PInteger s)
-> Term s PInteger
-> Term s (PMaybe PInteger)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term s PInteger -> PMaybe PInteger s
forall (a :: PType) (s :: S). Term s a -> PMaybe a s
PJust (Term s PInteger -> Term s (PMaybe PInteger))
-> Term s PInteger -> Term s (PMaybe PInteger)
forall a b. (a -> b) -> a -> b
$ Term s PInteger
result)