Measure.hs

module KB.Text.Layout.Measure
  ( -- * Layout context
    createLayoutContext
  , clearCache
  , LayoutContext

    -- * Preparing text for layout
  , prepare
  , prepareWithAdvances
  , prepareStyled
  , PreparedText (..)
  , MeasuredSegment (..)
  , Grapheme (..)
  , Span (..)
  , Advances

    -- * Measuring
  , measure
  , charAdvances

    -- * Styles
  , newStyle
  , Style (..)
  ) where

import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Vector.Storable qualified as Storable
import KB.Text.Layout.Analysis (BreakKind (..), Segment (..), analyze)
import KB.Text.Layout.Segmentation qualified as Segmentation
import KB.Text.Shape qualified as TextShape
import KB.Text.Shape.Font qualified as Font

prepare :: LayoutContext -> Style -> Text -> IO PreparedText
prepare ctx style t = prepareStyled ctx [TextSpan style t]

prepareWithAdvances :: LayoutContext -> Style -> Text -> Advances -> IO PreparedText
prepareWithAdvances ctx style t advance = prepareStyled ctx [ShapedSpan style t advance]

prepareStyled :: LayoutContext -> [Span] -> IO PreparedText
prepareStyled ctx spans = do
  segments <- Vector.fromList . concat <$> traverse spanSegments spans
  pure $! PreparedText{segments}
  where
    spanSegments = \case
      AtomSpan style label w ->
        pure
          [ MeasuredSegment
              { text = label
              , kind = Atomic
              , width = w
              , style = style.key
              , graphemes = [Grapheme{codepoints = Text.length label, width = w}]
              }
          ]
      TextSpan style t ->
        textSegments style t <$> charAdvances ctx style t
      ShapedSpan style t advance
        | Storable.length advance /= Text.length t ->
            error $
              "ShapedSpan: expected one advance per codepoint, got "
                <> show (Storable.length advance)
                <> " advances for "
                <> show (Text.length t)
                <> " codepoints"
        | otherwise ->
            pure (textSegments style t advance)

textSegments :: Style -> Text -> Advances -> [MeasuredSegment]
textSegments style t advance = go 0 (analyze t)
  where
    prefix = Storable.scanl' (+) 0 advance
    slice off n = prefix Storable.! (off + n) - prefix Storable.! off
    clusterWidths off = \case
      [] -> []
      c : rest ->
        let (n, rest') = absorbZeroWidth (Text.length c) rest
        in Grapheme{codepoints = n, width = slice off n} : clusterWidths (off + n) rest'
      where
        absorbZeroWidth n = \case
          c : rest | slice (off + n) (Text.length c) == 0 -> absorbZeroWidth (n + Text.length c) rest
          rest -> (n, rest)
    go !off = \case
      [] -> []
      seg : rest ->
        let !measured = segmentAt off seg
        in measured : go (off + Text.length seg.text) rest
    segmentAt off seg =
      MeasuredSegment
        { text = seg.text
        , kind = seg.kind
        , width
        , style = style.key
        , graphemes
        }
      where
        n = Text.length seg.text
        width = case seg.kind of
          SoftHyphen -> style.hyphenWidth
          ZeroWidthBreak -> 0
          HardBreak -> 0
          Tab -> fromIntegral n * style.spaceWidth
          _ -> slice off n
        graphemes
          | seg.kind == Word || seg.kind == Glue = clusterWidths off (Segmentation.clusters seg.text)
          | otherwise = [Grapheme{codepoints = n, width}]

data PreparedText = PreparedText
  { segments :: Vector MeasuredSegment
  }
  deriving stock (Eq, Show)

data MeasuredSegment = MeasuredSegment
  { text :: Text
  , kind :: BreakKind
  , width :: Float
  , style :: Int
  , graphemes :: [Grapheme]
  }
  deriving stock (Eq, Show)

data Grapheme = Grapheme
  { codepoints :: Int
  , width :: Float
  }
  deriving stock (Eq, Show)

data Span
  = TextSpan Style Text
  | ShapedSpan Style Text Advances
  | AtomSpan Style Text Float

type Advances = Storable.Vector Float

charAdvances :: LayoutContext -> Style -> Text -> IO Advances
charAdvances ctx style t
  | Text.null t = pure Storable.empty
  | otherwise = do
      runs <- TextShape.run ctx.shape (TextShape.withFont_ style.font (TextShape.text_ t))
      contributions <- concat <$> traverse runContributions runs
      pure (Storable.map (* style.scale) (Storable.accum (+) (Storable.replicate n 0) contributions))
  where
    n = Text.length t
    runContributions (run, glyphs) = do
      info <- Font.getFontInfo run.font
      let toGrid = style.grid / fromIntegral info.unitsPerEm
      pure
        [ (g.codepointIndex, fromIntegral g.advanceX * toGrid)
        | g <- glyphs
        , g.codepointIndex >= 0
        , g.codepointIndex < n
        ]

measure :: LayoutContext -> Style -> Text -> IO Float
measure ctx style t
  | Text.null t = pure 0
  | otherwise =
      Map.lookup (style.key, t) <$> readIORef ctx.metrics >>= \case
        Just w -> pure w
        Nothing -> do
          w <- shapedWidth ctx style.font style.grid style.scale t
          modifyIORef' ctx.metrics (Map.insert (style.key, t) w)
          pure w

shapedWidth :: LayoutContext -> Font.Font -> Float -> Float -> Text -> IO Float
shapedWidth ctx font grid scale t = do
  runs <- TextShape.run ctx.shape (TextShape.withFont_ font (TextShape.text_ t))
  let runWidth (run, glyphs) = do
        info <- Font.getFontInfo run.font
        let toGrid = grid / fromIntegral info.unitsPerEm
        pure (sum [fromIntegral g.advanceX | g <- glyphs] * toGrid)
  (* scale) . sum <$> traverse runWidth runs

newStyle :: LayoutContext -> Font.Font -> Float -> IO Style
newStyle ctx font size = do
  info <- Font.getFontInfo font
  key <- atomicModifyIORef' ctx.styleKeys \n -> (n + 1, n)
  let
    scale = size / fromIntegral info.capitalHeight
    grid = fromIntegral info.unitsPerEm
  spaceWidth <- shapedWidth ctx font grid scale " "
  hyphenWidth <- shapedWidth ctx font grid scale "-"
  pure
    Style
      { key
      , font
      , size
      , em = size * Font.emToCaps info
      , scale
      , grid
      , spaceWidth
      , hyphenWidth
      }

data Style = Style
  { key :: Int
  , font :: Font.Font
  , size :: Float
  , em :: Float
  , scale :: Float
  , grid :: Float
  , spaceWidth :: Float
  , hyphenWidth :: Float
  }

createLayoutContext :: TextShape.Context -> IO LayoutContext
createLayoutContext shape = do
  metrics <- newIORef Map.empty
  styleKeys <- newIORef 0
  pure LayoutContext{shape, metrics, styleKeys}

clearCache :: LayoutContext -> IO ()
clearCache ctx = writeIORef ctx.metrics Map.empty

data LayoutContext = LayoutContext
  { shape :: TextShape.Context
  , metrics :: IORef (Map (Int, Text) Float)
  , styleKeys :: IORef Int
  }