Segmentation.hs

{-| Unicode segmentation over the kbts entire-string API.

All positions are UTF-8 byte offsets into the input 'Text'.
-}
module KB.Text.Shape.Segmentation
  ( Break (..)
  , breaks
  , breaksWith

    -- * Grapheme clusters
  , clusters
  , boundaries

    -- * Line break opportunities
  , softBreaks
  , softBreaksWith

    -- * Word boundaries
  , wordBreaks
  ) where

import Data.Bits ((.&.))
import Data.Char (ord)
import Data.IntSet qualified as IntSet
import Data.List (sort)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Foreign qualified as Text
import Foreign (alloca, allocaArray, fillBytes, nullPtr, peek, peekArray, sizeOf)
import System.IO.Unsafe (unsafePerformIO)

import KB.Text.Shape.FFI.API.Segmentation qualified as Segmentation
import KB.Text.Shape.FFI.Enums (JapaneseLineBreakStyle, pattern DIRECTION_DONT_KNOW, pattern JAPANESE_LINE_BREAK_STYLE_NORMAL)
import KB.Text.Shape.FFI.Flags (BreakFlags, pattern BREAK_FLAG_GRAPHEME, pattern BREAK_FLAG_LINE_SOFT, pattern BREAK_FLAG_WORD)
import KB.Text.Shape.FFI.Structs (Break (..))

{- | Extended grapheme cluster texts, in order.

The engine skips UAX #29 GB11, so 'boundaries' refuses to cut adjacent
to a ZWJ; emoji ZWJ sequences stay whole.
-}
clusters :: Text -> [Text]
clusters t = go t 0 (boundaries t)
  where
    go rest from cuts = case dropWhile (<= from) cuts of
      cut : more ->
        let n = fromIntegral (cut - from)
        in Text.takeWord8 n rest : go (Text.dropWord8 n rest) cut more
      [] -> [rest | not (Text.null rest)]

-- | Extended grapheme cluster boundaries.
boundaries :: Text -> [Int]
boundaries t = zwjFilter t (positions BREAK_FLAG_GRAPHEME (breaks t))

{- | Positions where a line may break, at 'JAPANESE_LINE_BREAK_STYLE_NORMAL'.

The engine skips UAX #14 LB8a, so positions adjacent to a ZWJ are dropped.
-}
softBreaks :: Text -> [Int]
softBreaks = softBreaksWith JAPANESE_LINE_BREAK_STYLE_NORMAL

-- | 'softBreaks' with an explicit kinsoku strictness.
softBreaksWith :: JapaneseLineBreakStyle -> Text -> [Int]
softBreaksWith style t = zwjFilter t (positions BREAK_FLAG_LINE_SOFT (breaksWith style t))

-- | UAX #29 word boundaries.
wordBreaks :: Text -> [Int]
wordBreaks = positions BREAK_FLAG_WORD . breaks

positions :: BreakFlags -> [Break] -> [Int]
positions flag found = dedupe (sort [b.position | b <- found, b.flags .&. flag /= mempty])
  where
    dedupe = \case
      x : rest@(y : _)
        | x == y -> dedupe rest
        | otherwise -> x : dedupe rest
      rest -> rest

zwjFilter :: Text -> [Int] -> [Int]
zwjFilter t
  | IntSet.null zwjs = id
  | otherwise = filter safe
  where
    zwjs = IntSet.fromList (zwjOffsets 0 (Text.unpack t))
    zwjOffsets off = \case
      [] -> []
      c : rest -> [off | c == '\x200D'] <> zwjOffsets (off + utf8Length c) rest
    total = Text.lengthWord8 t
    safe b = b <= 0 || b >= total || not (IntSet.member (b - 3) zwjs || IntSet.member b zwjs)

utf8Length :: Char -> Int
utf8Length c
  | o < 0x80 = 1
  | o < 0x800 = 2
  | o < 0x10000 = 3
  | otherwise = 4
  where
    o = ord c

-- | All breaks at 'JAPANESE_LINE_BREAK_STYLE_NORMAL'.
breaks :: Text -> [Break]
breaks = breaksWith JAPANESE_LINE_BREAK_STYLE_NORMAL

{- | 'breaks' with an explicit kinsoku strictness, in one engine pass.

The entire-string API merges the flags of same-position breaks into one
record; the 2.24 write-cursor caveat on
'Segmentation.kbts_BreakEntireString' is compensated by zeroing the output
buffer and dropping the records with empty flags.

Delayed break emissions rely on the @LOCAL PATCH(break_position_skew)@
in @cbits\/kb_text_shape.inc@: unpatched, positions of breaks pending on
lookahead skew away from the fed byte increments in both directions (see
@upstream-fixme\/break_position_skew.c@).
-}
breaksWith :: JapaneseLineBreakStyle -> Text -> [Break]
breaksWith style t
  | Text.null t = []
  | otherwise = unsafePerformIO $
      Text.withCStringLen t \(utf8Ptr, len) ->
        attempt utf8Ptr len (4 * len + 16)
  where
    attempt utf8Ptr len capacity =
      allocaArray capacity \breaksPtr ->
        alloca \breakCountPtr -> do
          fillBytes breaksPtr 0 (capacity * sizeOf (undefined :: Break))
          Segmentation.kbts_BreakEntireStringUtf8
            DIRECTION_DONT_KNOW
            style
            mempty
            utf8Ptr
            (fromIntegral len)
            breaksPtr
            (fromIntegral capacity)
            breakCountPtr
            nullPtr
            0
            nullPtr
          count <- fromIntegral <$> peek breakCountPtr
          if count > capacity
            then
              attempt utf8Ptr len count
            else
              filter (\b -> b.flags /= mempty) <$> peekArray count breaksPtr
{-# NOINLINE breaksWith #-}