Internal.hs

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Network.DNS.Pattern.Internal
-- Description : Internal DNS types and definitions
--
-- This module is not part of public API and may change even between patch versions.

module Network.DNS.Internal
  ( DomainLabel(..)
  , Domain(..)
  , DList(..)
  , toDList
  , fromDList
  , singleton
  , sbsMap
  , sbsSingleton
  )

where

import           Data.Function (on)
import           GHC.Word
#if !MIN_VERSION_bytestring(0,11,3)
import           Control.Monad.ST (runST)
import           Data.ByteString.Short.Internal (ShortByteString(SBS))
import           GHC.Exts (Int#, MutableByteArray#, indexWord8Array#,
                           newByteArray#, unsafeFreezeByteArray#,
                           writeWord8Array#, (+#))
import           GHC.Int (Int(..))
import           GHC.ST (ST(..))
#endif

import qualified Data.ByteString.Short as BS

-- | Domain label with case-insensitive 'Eq' and 'Ord' as per [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
data DomainLabel = DomainLabel { getDomainLabel_ :: !BS.ShortByteString
                               , getDomainLabelCF_ :: !BS.ShortByteString }

-- | A domain parsed into labels. Each label is a 'BS.ShortByteString' rather than 'T.Text' or 'String' because a label can contain arbitrary bytes.
-- However, the 'Ord' and 'Eq' instances do limited case-folding according to [RFC4343](https://datatracker.ietf.org/doc/html/rfc4343#section-3).
newtype Domain = Domain [DomainLabel] deriving (Eq, Ord)

instance Ord DomainLabel where
  (<=) = (<=) `on` getDomainLabelCF_
  compare = compare `on` getDomainLabelCF_

instance Eq DomainLabel where
  (==) = (==) `on` getDomainLabelCF_

-- | Difference list à la Huhges
newtype DList a = DList ([a] -> [a])

-- | Turn a list into 'DList'
{-# INLINE toDList #-}
toDList :: [a] -> DList a
toDList = DList . (++)

-- | Turn 'DList' back into a list.
{-# INLINE fromDList #-}
fromDList :: DList a -> [a]
fromDList (DList dl) = dl []

-- | Create a 'DList' containing just the specified element
{-# INLINE singleton #-}
singleton :: a -> DList a
singleton = DList . (:)

instance Semigroup (DList a) where
  {-# INLINE (<>) #-}
  DList l <> DList r = DList (l . r)

instance Monoid (DList a) where
  {-# INLINE mempty #-}
  mempty = DList id

{-# INLINE sbsSingleton #-}
sbsSingleton :: Word8 -> BS.ShortByteString
#if MIN_VERSION_bytestring(0,11,3)
sbsSingleton = BS.singleton
#else
sbsSingleton (W8# w) = runST $ ST $ \s1 ->
  case newByteArray# 1# s1 of
    (# s2, mba #) -> case writeWord8Array# mba 0# w s2 of
          s3 -> case unsafeFreezeByteArray# mba s3 of
             (# s4, ma #) -> (# s4, SBS ma #)

#endif

sbsMap :: (Word8 -> Word8) -> BS.ShortByteString -> BS.ShortByteString
#if MIN_VERSION_bytestring(0,11,3)
sbsMap = BS.map
#else
sbsMap m sbs@(SBS ba) = runST $ ST $ \s1 ->
  case newByteArray# l# s1 of
    (# s2, mba #) -> case go mba 0# l# of
      ST f -> case f s2 of
        (# s3, _ #) -> case unsafeFreezeByteArray# mba s3 of
           (# s4, ma #) -> (# s4, SBS ma #)
  where
    !(I# l#) = BS.length sbs
    go :: MutableByteArray# s -> Int# -> Int# -> ST s ()
    go !mba !i !l
      | I# i >= I# l = return ()
      | otherwise = (ST $ \s ->
          let !(W8# w') = m (W8# (indexWord8Array# ba i)) in
          (# writeWord8Array# mba i w' s, () #)
       ) >> go mba (i +# 1#) l
#endif