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.List (mapAccumL)
import Data.Word (Word8)

canonicalizeLineEndings :: BL.ByteString -> BL.ByteString
canonicalizeLineEndings bs =
    BL.fromChunks $ chunks ++ [flush]
  where
    (state, chunks) =
        mapAccumL
            canonicalizeLineEndingsChunk
            (CRState False)
            (BL.toChunks bs)
    flush = canonicalizeLineEndingsFlush state

-- | 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 to
preserve it as-is.
-}
canonicalizeLineEndingsChunk
    :: CRState
    -> B.ByteString
    -> (CRState, B.ByteString)
canonicalizeLineEndingsChunk (CRState prevWasCR) chunk
    | B.null chunk = (CRState False, B.empty)
    | otherwise =
        ( CRState isPrevCR
        , BL.toStrict $ BB.toLazyByteString bldr
        )
  where
    isPrevCR = B.last chunk == 0x0d
    bldr = snd $ B.foldl' step (prevWasCR, mempty) chunk
    step (hadPrevCR, b) w
        | hadPrevCR && w == 0x0a =
            (False, b <> BB.word8 0x0d <> BB.word8 0x0a)
        | hadPrevCR && w == 0x0d =
            (True, b <> BB.word8 0x0d)
        | hadPrevCR =
            (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@
is preserved as-is).
-}
canonicalizeLineEndingsFlush :: CRState -> B.ByteString
canonicalizeLineEndingsFlush (CRState hadPrevCR)
    | hadPrevCR = B.pack [0x0d]
    | 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 pcr line, acc) w
        | pcr && w == 0x0a =
            (StripWSState False B.empty, acc <> trimmedLine <> crlf)
        | pcr =
            ( StripWSState False (B.singleton w)
            , acc <> trimmedLine <> BB.word8 0x0d
            )
        | w == 0x0d = (StripWSState True line, acc)
        | otherwise = (StripWSState False (line <> B.singleton w), acc)
      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 pcr line) =
    if pcr
        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