{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Lightning.Protocol.BOLT9.Codec
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Parsing and rendering for BOLT #9 feature vectors.

module Lightning.Protocol.BOLT9.Codec (
    -- * Parsing and rendering
    parse
  , render

    -- * Bit operations
  , setBit
  , clearBit
  , testBit

    -- * Feature operations
  , setFeature
  , hasFeature
  , isFeatureSet
  , listFeatures
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Word (Word16)
import Lightning.Protocol.BOLT9.Features
import Lightning.Protocol.BOLT9.Types
  ( FeatureLevel(..)
  , FeatureVector
  , bitIndex
  , clear
  , fromByteString
  , member
  , set
  , unFeatureVector
  )

-- Parsing and rendering ------------------------------------------------------

-- | Parse a ByteString into a FeatureVector.
--
--   Alias for 'fromByteString'.
parse :: ByteString -> FeatureVector
parse :: ByteString -> FeatureVector
parse = ByteString -> FeatureVector
fromByteString
{-# INLINE parse #-}

-- | Render a FeatureVector to a ByteString, trimming leading zero bytes
--   for compact encoding.
render :: FeatureVector -> ByteString
render :: FeatureVector -> ByteString
render = (Word8 -> Bool) -> ByteString -> ByteString
BS.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) (ByteString -> ByteString)
-> (FeatureVector -> ByteString) -> FeatureVector -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FeatureVector -> ByteString
unFeatureVector
{-# INLINE render #-}

-- Bit operations -------------------------------------------------------------

-- | Set a bit by raw index.
--
--   >>> setBit 17 empty
--   FeatureVector {unFeatureVector = "\STX"}
setBit :: Word16 -> FeatureVector -> FeatureVector
setBit :: Word16 -> FeatureVector -> FeatureVector
setBit !Word16
idx = BitIndex -> FeatureVector -> FeatureVector
set (Word16 -> BitIndex
bitIndex Word16
idx)
{-# INLINE setBit #-}

-- | Clear a bit by raw index.
--
--   >>> clearBit 17 (setBit 17 empty)
--   FeatureVector {unFeatureVector = ""}
clearBit :: Word16 -> FeatureVector -> FeatureVector
clearBit :: Word16 -> FeatureVector -> FeatureVector
clearBit !Word16
idx = BitIndex -> FeatureVector -> FeatureVector
clear (Word16 -> BitIndex
bitIndex Word16
idx)
{-# INLINE clearBit #-}

-- | Test if a bit is set.
--
--   >>> testBit 17 (setBit 17 empty)
--   True
--   >>> testBit 16 (setBit 17 empty)
--   False
testBit :: Word16 -> FeatureVector -> Bool
testBit :: Word16 -> FeatureVector -> Bool
testBit !Word16
idx = BitIndex -> FeatureVector -> Bool
member (Word16 -> BitIndex
bitIndex Word16
idx)
{-# INLINE testBit #-}

-- Feature operations ---------------------------------------------------------

-- | Set a feature's bit at the given level.
--
--   'Required' sets the even bit, 'Optional' sets the odd bit.
--
--   >>> import Data.Maybe (fromJust)
--   >>> let mpp = fromJust (featureByName "basic_mpp")
--   >>> setFeature mpp Optional empty  -- set optional bit (17)
--   FeatureVector {unFeatureVector = "\STX"}
--   >>> setFeature mpp Required empty  -- set required bit (16)
--   FeatureVector {unFeatureVector = "\SOH"}
setFeature :: Feature -> FeatureLevel -> FeatureVector -> FeatureVector
setFeature :: Feature -> FeatureLevel -> FeatureVector -> FeatureVector
setFeature !Feature
f !FeatureLevel
level = Word16 -> FeatureVector -> FeatureVector
setBit Word16
targetBit
  where
    !baseBit :: Word16
baseBit   = Feature -> Word16
featureBaseBit Feature
f
    !targetBit :: Word16
targetBit = case FeatureLevel
level of
      FeatureLevel
Required -> Word16
baseBit
      FeatureLevel
Optional -> Word16
baseBit Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1
{-# INLINE setFeature #-}

-- | Check if a feature is set in the vector.
--
--   Returns:
--
--   * @Just Required@ if the required (even) bit is set
--   * @Just Optional@ if the optional (odd) bit is set (and required is not)
--   * @Nothing@ if neither bit is set
--
--   >>> import Data.Maybe (fromJust)
--   >>> let mpp = fromJust (featureByName "basic_mpp")
--   >>> hasFeature mpp (setFeature mpp Optional empty)
--   Just Optional
--   >>> hasFeature mpp (setFeature mpp Required empty)
--   Just Required
--   >>> hasFeature mpp empty
--   Nothing
hasFeature :: Feature -> FeatureVector -> Maybe FeatureLevel
hasFeature :: Feature -> FeatureVector -> Maybe FeatureLevel
hasFeature !Feature
f !FeatureVector
fv
  | Word16 -> FeatureVector -> Bool
testBit Word16
baseBit FeatureVector
fv       = FeatureLevel -> Maybe FeatureLevel
forall a. a -> Maybe a
Just FeatureLevel
Required
  | Word16 -> FeatureVector -> Bool
testBit (Word16
baseBit Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1) FeatureVector
fv = FeatureLevel -> Maybe FeatureLevel
forall a. a -> Maybe a
Just FeatureLevel
Optional
  | Bool
otherwise                = Maybe FeatureLevel
forall a. Maybe a
Nothing
  where
    !baseBit :: Word16
baseBit = Feature -> Word16
featureBaseBit Feature
f
{-# INLINE hasFeature #-}

-- | Check if either bit of a feature is set in the vector.
--
--   >>> import Data.Maybe (fromJust)
--   >>> let mpp = fromJust (featureByName "basic_mpp")
--   >>> isFeatureSet mpp (setFeature mpp Optional empty)
--   True
--   >>> isFeatureSet mpp empty
--   False
isFeatureSet :: Feature -> FeatureVector -> Bool
isFeatureSet :: Feature -> FeatureVector -> Bool
isFeatureSet !Feature
f !FeatureVector
fv =
  let !baseBit :: Word16
baseBit = Feature -> Word16
featureBaseBit Feature
f
  in  Word16 -> FeatureVector -> Bool
testBit Word16
baseBit FeatureVector
fv Bool -> Bool -> Bool
|| Word16 -> FeatureVector -> Bool
testBit (Word16
baseBit Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1) FeatureVector
fv
{-# INLINE isFeatureSet #-}

-- | List all known features that are set in the vector.
--
--   Returns pairs of (Feature, FeatureLevel) indicating whether each
--   feature is set as required or optional.
--
--   >>> import Data.Maybe (fromJust)
--   >>> let mpp = fromJust (featureByName "basic_mpp")
--   >>> let ps = fromJust (featureByName "payment_secret")
--   >>> let fv = setFeature mpp Optional (setFeature ps Required empty)
--   >>> map (\(f, l) -> (featureName f, l)) (listFeatures fv)
--   [("payment_secret",Required),("basic_mpp",Optional)]
listFeatures :: FeatureVector -> [(Feature, FeatureLevel)]
listFeatures :: FeatureVector -> [(Feature, FeatureLevel)]
listFeatures !FeatureVector
fv = (Feature -> [(Feature, FeatureLevel)] -> [(Feature, FeatureLevel)])
-> [(Feature, FeatureLevel)]
-> [Feature]
-> [(Feature, FeatureLevel)]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Feature -> [(Feature, FeatureLevel)] -> [(Feature, FeatureLevel)]
check [] [Feature]
knownFeatures
  where
    check :: Feature -> [(Feature, FeatureLevel)] -> [(Feature, FeatureLevel)]
check !Feature
f ![(Feature, FeatureLevel)]
acc = case Feature -> FeatureVector -> Maybe FeatureLevel
hasFeature Feature
f FeatureVector
fv of
      Just FeatureLevel
level -> (Feature
f, FeatureLevel
level) (Feature, FeatureLevel)
-> [(Feature, FeatureLevel)] -> [(Feature, FeatureLevel)]
forall a. a -> [a] -> [a]
: [(Feature, FeatureLevel)]
acc
      Maybe FeatureLevel
Nothing    -> [(Feature, FeatureLevel)]
acc