Whitespace.hs

-- 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 bs = BL.fromStrict (B.unfoldr step (Nothing, BL.toStrict bs))
  where
    step (Nothing, rest)
        | B.null rest = Nothing
        | otherwise =
            case B.uncons rest of
                Just (0x0d, tail') ->
                    case B.uncons tail' of
                        Just (0x0a, tail'') -> Just (0x0d, (Just 0x0a, tail''))
                        _ -> Just (0x0d, (Just 0x0a, tail'))
                Just (0x0a, tail') -> Just (0x0d, (Just 0x0a, tail'))
                Just (w, tail') -> Just (w, (Nothing, tail'))
    step (Just w, rest) = Just (w, (Nothing, rest))

-- | State carried between chunks when canonicalizing line endings.
data CRState = CRState
    { 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 prevWasCR) chunk
    | B.null chunk = (CRState False, B.empty)
    | otherwise =
        ( CRState newPrevCR
        , BL.toStrict $ BB.toLazyByteString bldr
        )
  where
    newPrevCR = B.last chunk == 0x0d
    bldr = snd $ B.foldl' step (prevWasCR, mempty) chunk
    step (prevCR, b) w
        | prevCR && w == 0x0a =
            (False, b <> BB.word8 0x0d <> BB.word8 0x0a)
        | prevCR && w == 0x0d =
            (True, b <> BB.word8 0x0d <> BB.word8 0x0a)
        | prevCR = (False, b <> BB.word8 0x0d <> BB.word8 w)
        | w == 0x0d = (True, b)
        | w == 0x0a = (False, b <> BB.word8 0x0d <> BB.word8 0x0a)
        | otherwise = (False, b <> BB.word8 w)

{- | Emit any pending state as final bytes (a trailing standalone @\\r@
becomes @\\r\\n@).
-}
canonicalizeLineEndingsFlush :: CRState -> B.ByteString
canonicalizeLineEndingsFlush (CRState prevCR)
    | prevCR = B.pack [0x0d, 0x0a]
    | otherwise = 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 = BB.toLazyByteString . go . BL.toStrict
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace w = w == 0x20 || w == 0x09

    crlf :: B.ByteString
    crlf = B.pack [0x0d, 0x0a]

    go bs
        | B.null bs = mempty
        | otherwise =
            let (line, rest) = B.span (/= 0x0d) bs
                trimmed = B.dropWhileEnd isTrailingWhitespace line
             in case B.uncons rest of
                    Just (0x0d, afterCR) ->
                        case B.uncons afterCR of
                            Just (0x0a, afterLF) ->
                                BB.byteString trimmed <> BB.byteString crlf <> go afterLF
                            _ ->
                                BB.byteString trimmed <> BB.word8 0x0d <> go afterCR
                    _ ->
                        BB.byteString trimmed

data StripWSState = StripWSState
    { swsPrevCR :: !Bool
    , 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 st chunk = (st', BL.toStrict $ BB.toLazyByteString bldr)
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace w = w == 0x20 || w == 0x09

    (st', bldr) = B.foldl' stepByte (st, mempty) chunk

    crlf :: BB.Builder
    crlf = BB.word8 0x0d <> BB.word8 0x0a

    stepByte (StripWSState prevCR line, bldr) w
        | prevCR && w == 0x0a =
            (StripWSState False B.empty, bldr <> trimmedLine <> crlf)
        | prevCR =
            ( StripWSState False (B.singleton w)
            , bldr <> trimmedLine <> BB.word8 0x0d
            )
        | w == 0x0d = (StripWSState True line, bldr)
        | otherwise = (StripWSState False (line <> B.singleton w), bldr)
      where
        trimmedLine = BB.byteString $ B.dropWhileEnd isTrailingWhitespace 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 prevCR line) =
    if prevCR
        then
            BL.toStrict $
                BB.toLazyByteString $
                    BB.byteString trimmedLine <> BB.word8 0x0d
        else
            BL.toStrict $ BB.toLazyByteString $ BB.byteString trimmedLine
  where
    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace w = w == 0x20 || w == 0x09
    trimmedLine = B.dropWhileEnd isTrailingWhitespace line