-- Whitespace.hs: utility functions involving whitespace
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

module Codec.Encryption.OpenPGP.Internal.Whitespace
    ( canonicalizeLineEndings
    , canonicalizeLineEndingsChunk
    , canonicalizeLineEndingsFlush
    , CRState (..)
    , stripTrailingWhitespacePerLine
    , stripTrailingWhitespacePerLineChunk
    , stripTrailingWhitespacePerLineFlush
    , StripWSState (..)
    ) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import Data.Foldable (foldl')
import Data.Word (Word8)

canonicalizeLineEndings :: BL.ByteString -> BL.ByteString
canonicalizeLineEndings :: ByteString -> ByteString
canonicalizeLineEndings ByteString
bs = StrictByteString -> ByteString
BL.fromStrict (((Maybe Word8, StrictByteString)
 -> Maybe (Word8, (Maybe Word8, StrictByteString)))
-> (Maybe Word8, StrictByteString) -> StrictByteString
forall a. (a -> Maybe (Word8, a)) -> a -> StrictByteString
B.unfoldr (Maybe Word8, StrictByteString)
-> Maybe (Word8, (Maybe Word8, StrictByteString))
forall {a}.
Num a =>
(Maybe Word8, StrictByteString)
-> Maybe (Word8, (Maybe a, StrictByteString))
step (Maybe Word8
forall a. Maybe a
Nothing, ByteString -> StrictByteString
BL.toStrict ByteString
bs))
  where
    step :: (Maybe Word8, StrictByteString)
-> Maybe (Word8, (Maybe a, StrictByteString))
step (Maybe Word8
Nothing, StrictByteString
rest)
        | StrictByteString -> Bool
B.null StrictByteString
rest = Maybe (Word8, (Maybe a, StrictByteString))
forall a. Maybe a
Nothing
        | Bool
otherwise =
            case StrictByteString -> Maybe (Word8, StrictByteString)
B.uncons StrictByteString
rest of
                Just (Word8
0x0d, StrictByteString
tail') ->
                    case StrictByteString -> Maybe (Word8, StrictByteString)
B.uncons StrictByteString
tail' of
                        Just (Word8
0x0a, StrictByteString
tail'') -> (Word8, (Maybe a, StrictByteString))
-> Maybe (Word8, (Maybe a, StrictByteString))
forall a. a -> Maybe a
Just (Word8
0x0d, (a -> Maybe a
forall a. a -> Maybe a
Just a
0x0a, StrictByteString
tail''))
                        Maybe (Word8, StrictByteString)
_ -> (Word8, (Maybe a, StrictByteString))
-> Maybe (Word8, (Maybe a, StrictByteString))
forall a. a -> Maybe a
Just (Word8
0x0d, (a -> Maybe a
forall a. a -> Maybe a
Just a
0x0a, StrictByteString
tail'))
                Just (Word8
0x0a, StrictByteString
tail') -> (Word8, (Maybe a, StrictByteString))
-> Maybe (Word8, (Maybe a, StrictByteString))
forall a. a -> Maybe a
Just (Word8
0x0d, (a -> Maybe a
forall a. a -> Maybe a
Just a
0x0a, StrictByteString
tail'))
                Just (Word8
w, StrictByteString
tail') -> (Word8, (Maybe a, StrictByteString))
-> Maybe (Word8, (Maybe a, StrictByteString))
forall a. a -> Maybe a
Just (Word8
w, (Maybe a
forall a. Maybe a
Nothing, StrictByteString
tail'))
    step (Just Word8
w, StrictByteString
rest) = (Word8, (Maybe a, StrictByteString))
-> Maybe (Word8, (Maybe a, StrictByteString))
forall a. a -> Maybe a
Just (Word8
w, (Maybe a
forall a. Maybe a
Nothing, StrictByteString
rest))

-- | State carried between chunks when canonicalizing line endings.
data CRState = CRState
    { CRState -> Bool
prevCR :: !Bool
    }

{- | Process one strict chunk and return updated state plus emitted bytes.

The returned bytes may end with a standalone @\\r@ if the chunk boundary
falls mid-pair; feed it to 'canonicalizeLineEndingsFlush' at the end.
-}
canonicalizeLineEndingsChunk
    :: CRState
    -> B.ByteString
    -> (CRState, B.ByteString)
canonicalizeLineEndingsChunk :: CRState -> StrictByteString -> (CRState, StrictByteString)
canonicalizeLineEndingsChunk (CRState Bool
prevWasCR) StrictByteString
chunk
    | StrictByteString -> Bool
B.null StrictByteString
chunk = (Bool -> CRState
CRState Bool
False, StrictByteString
B.empty)
    | Bool
otherwise =
        ( Bool -> CRState
CRState Bool
newPrevCR
        , ByteString -> StrictByteString
BL.toStrict (ByteString -> StrictByteString) -> ByteString -> StrictByteString
forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
BB.toLazyByteString Builder
bldr
        )
  where
    newPrevCR :: Bool
newPrevCR = HasCallStack => StrictByteString -> Word8
StrictByteString -> Word8
B.last StrictByteString
chunk Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0d
    bldr :: Builder
bldr = (Bool, Builder) -> Builder
forall a b. (a, b) -> b
snd ((Bool, Builder) -> Builder) -> (Bool, Builder) -> Builder
forall a b. (a -> b) -> a -> b
$ ((Bool, Builder) -> Word8 -> (Bool, Builder))
-> (Bool, Builder) -> StrictByteString -> (Bool, Builder)
forall a. (a -> Word8 -> a) -> a -> StrictByteString -> a
B.foldl' (Bool, Builder) -> Word8 -> (Bool, Builder)
step (Bool
prevWasCR, Builder
forall a. Monoid a => a
mempty) StrictByteString
chunk
    step :: (Bool, Builder) -> Word8 -> (Bool, Builder)
step (Bool
prevCR, Builder
b) Word8
w
        | Bool
prevCR Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0a =
            (Bool
False, Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0a)
        | Bool
prevCR Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0d =
            (Bool
True, Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0a)
        | Bool
prevCR = (Bool
False, Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
w)
        | Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0d = (Bool
True, Builder
b)
        | Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0a = (Bool
False, Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0a)
        | Bool
otherwise = (Bool
False, Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
w)

{- | Emit any pending state as final bytes (a trailing standalone @\\r@
becomes @\\r\\n@).
-}
canonicalizeLineEndingsFlush :: CRState -> B.ByteString
canonicalizeLineEndingsFlush :: CRState -> StrictByteString
canonicalizeLineEndingsFlush (CRState Bool
prevCR)
    | Bool
prevCR = [Word8] -> StrictByteString
B.pack [Word8
0x0d, Word8
0x0a]
    | Bool
otherwise = StrictByteString
B.empty

{- | Strip trailing spaces (0x20) and tabs (0x09) from each line.

Lines are delimited by @\\r\\n@.  A line consisting solely of whitespace
collapses to just its terminator.  Any trailing partial line (without a
terminator) has its trailing whitespace stripped as well.

Uses 'Data.ByteString.Builder' to concatenate output in O(1) per segment,
avoiding the '++' thunk buildup of the old list-based version.
-}
stripTrailingWhitespacePerLine :: BL.ByteString -> BL.ByteString
stripTrailingWhitespacePerLine :: ByteString -> ByteString
stripTrailingWhitespacePerLine = Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString)
-> (ByteString -> Builder) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StrictByteString -> Builder
go (StrictByteString -> Builder)
-> (ByteString -> StrictByteString) -> ByteString -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> StrictByteString
BL.toStrict
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace :: Word8 -> Bool
isTrailingWhitespace Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x20 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x09

    crlf :: B.ByteString
    crlf :: StrictByteString
crlf = [Word8] -> StrictByteString
B.pack [Word8
0x0d, Word8
0x0a]

    go :: StrictByteString -> Builder
go StrictByteString
bs
        | StrictByteString -> Bool
B.null StrictByteString
bs = Builder
forall a. Monoid a => a
mempty
        | Bool
otherwise =
            let (StrictByteString
line, StrictByteString
rest) = (Word8 -> Bool)
-> StrictByteString -> (StrictByteString, StrictByteString)
B.span (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0x0d) StrictByteString
bs
                trimmed :: StrictByteString
trimmed = (Word8 -> Bool) -> StrictByteString -> StrictByteString
B.dropWhileEnd Word8 -> Bool
isTrailingWhitespace StrictByteString
line
             in case StrictByteString -> Maybe (Word8, StrictByteString)
B.uncons StrictByteString
rest of
                    Just (Word8
0x0d, StrictByteString
afterCR) ->
                        case StrictByteString -> Maybe (Word8, StrictByteString)
B.uncons StrictByteString
afterCR of
                            Just (Word8
0x0a, StrictByteString
afterLF) ->
                                StrictByteString -> Builder
BB.byteString StrictByteString
trimmed Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> StrictByteString -> Builder
BB.byteString StrictByteString
crlf Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> StrictByteString -> Builder
go StrictByteString
afterLF
                            Maybe (Word8, StrictByteString)
_ ->
                                StrictByteString -> Builder
BB.byteString StrictByteString
trimmed Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> StrictByteString -> Builder
go StrictByteString
afterCR
                    Maybe (Word8, StrictByteString)
_ ->
                        StrictByteString -> Builder
BB.byteString StrictByteString
trimmed

data StripWSState = StripWSState
    { StripWSState -> Bool
swsPrevCR :: !Bool
    , StripWSState -> StrictByteString
swsLine :: !B.ByteString
    }

{- | Process one strict chunk and return updated state plus emitted bytes.

Lines are delimited by @\\r\\n@.  A line consisting solely of whitespace
collapses to just its terminator.  Any trailing partial line (without a
terminator) has its trailing whitespace stripped as well.
-}
stripTrailingWhitespacePerLineChunk
    :: StripWSState -> B.ByteString -> (StripWSState, B.ByteString)
stripTrailingWhitespacePerLineChunk :: StripWSState
-> StrictByteString -> (StripWSState, StrictByteString)
stripTrailingWhitespacePerLineChunk StripWSState
st StrictByteString
chunk = (StripWSState
st', ByteString -> StrictByteString
BL.toStrict (ByteString -> StrictByteString) -> ByteString -> StrictByteString
forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
BB.toLazyByteString Builder
bldr)
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace :: Word8 -> Bool
isTrailingWhitespace Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x20 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x09

    (StripWSState
st', Builder
bldr) = ((StripWSState, Builder) -> Word8 -> (StripWSState, Builder))
-> (StripWSState, Builder)
-> StrictByteString
-> (StripWSState, Builder)
forall a. (a -> Word8 -> a) -> a -> StrictByteString -> a
B.foldl' (StripWSState, Builder) -> Word8 -> (StripWSState, Builder)
stepByte (StripWSState
st, Builder
forall a. Monoid a => a
mempty) StrictByteString
chunk

    crlf :: BB.Builder
    crlf :: Builder
crlf = Word8 -> Builder
BB.word8 Word8
0x0d Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0a

    stepByte :: (StripWSState, Builder) -> Word8 -> (StripWSState, Builder)
stepByte (StripWSState Bool
prevCR StrictByteString
line, Builder
bldr) Word8
w
        | Bool
prevCR Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0a =
            (Bool -> StrictByteString -> StripWSState
StripWSState Bool
False StrictByteString
B.empty, Builder
bldr Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
trimmedLine Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
crlf)
        | Bool
prevCR =
            ( Bool -> StrictByteString -> StripWSState
StripWSState Bool
False (Word8 -> StrictByteString
B.singleton Word8
w)
            , Builder
bldr Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
trimmedLine Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d
            )
        | Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x0d = (Bool -> StrictByteString -> StripWSState
StripWSState Bool
True StrictByteString
line, Builder
bldr)
        | Bool
otherwise = (Bool -> StrictByteString -> StripWSState
StripWSState Bool
False (StrictByteString
line StrictByteString -> StrictByteString -> StrictByteString
forall a. Semigroup a => a -> a -> a
<> Word8 -> StrictByteString
B.singleton Word8
w), Builder
bldr)
      where
        trimmedLine :: Builder
trimmedLine = StrictByteString -> Builder
BB.byteString (StrictByteString -> Builder) -> StrictByteString -> Builder
forall a b. (a -> b) -> a -> b
$ (Word8 -> Bool) -> StrictByteString -> StrictByteString
B.dropWhileEnd Word8 -> Bool
isTrailingWhitespace StrictByteString
line

{- | Emit any pending state as final bytes (trailing whitespace stripped from
the last partial line, standalone @\\r@ preserved).
-}
stripTrailingWhitespacePerLineFlush
    :: StripWSState -> B.ByteString
stripTrailingWhitespacePerLineFlush :: StripWSState -> StrictByteString
stripTrailingWhitespacePerLineFlush (StripWSState Bool
prevCR StrictByteString
line) =
    if Bool
prevCR
        then
            ByteString -> StrictByteString
BL.toStrict (ByteString -> StrictByteString) -> ByteString -> StrictByteString
forall a b. (a -> b) -> a -> b
$
                Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$
                    StrictByteString -> Builder
BB.byteString StrictByteString
trimmedLine Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
BB.word8 Word8
0x0d
        else
            ByteString -> StrictByteString
BL.toStrict (ByteString -> StrictByteString) -> ByteString -> StrictByteString
forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ StrictByteString -> Builder
BB.byteString StrictByteString
trimmedLine
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace :: Word8 -> Bool
isTrailingWhitespace Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x20 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x09
    trimmedLine :: StrictByteString
trimmedLine = (Word8 -> Bool) -> StrictByteString -> StrictByteString
B.dropWhileEnd Word8 -> Bool
isTrailingWhitespace StrictByteString
line