Atlas.hs

module NanoUI.Atlas
  ( ImageAtlas
  , newImageAtlas
  , atlasTextureId
  , registerImage
  , freshImageId
  , lookupImageUv
  , atlasSnapshot
  )
where

import Control.Applicative ((<|>))
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.IntMap.Strict qualified as IM
import Data.Word (Word8)
import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtrBytes, withForeignPtr)
import Foreign.Marshal.Utils (copyBytes, fillBytes)
import Foreign.Ptr (plusPtr)
import NanoUI.Types (ImageId (..))

-- | GPU texture id shared by every packed image so draw cmds batch.
atlasTextureId :: Int
atlasTextureId = 1

atlasPad :: Int
atlasPad = 1

atlasStart :: Int
atlasStart = 256

atlasMax :: Int
atlasMax = 4096

data AtlasSlot = AtlasSlot
  { slotX :: {-# UNPACK #-} !Int
  , slotY :: {-# UNPACK #-} !Int
  , slotW :: {-# UNPACK #-} !Int
  , slotH :: {-# UNPACK #-} !Int
  }

data AtlasState = AtlasState
  { asW :: {-# UNPACK #-} !Int
  , asH :: {-# UNPACK #-} !Int
  , asPtr :: ForeignPtr Word8
  , asSlots :: IM.IntMap AtlasSlot
  , asX :: {-# UNPACK #-} !Int
  , asY :: {-# UNPACK #-} !Int
  , asRowH :: {-# UNPACK #-} !Int
  , asGen :: {-# UNPACK #-} !Int
  , asLastFresh :: {-# UNPACK #-} !Int
  -- ^ The last id 'freshImageId' returned.
  }

newtype ImageAtlas = ImageAtlas (IORef AtlasState)

newImageAtlas :: IO ImageAtlas
newImageAtlas = do
  fp <- allocPixels atlasStart atlasStart
  ImageAtlas
    <$> newIORef
      AtlasState
        { asW = atlasStart
        , asH = atlasStart
        , asPtr = fp
        , asSlots = IM.empty
        , asX = atlasPad
        , asY = atlasPad
        , asRowH = 0
        , asGen = 0
        , asLastFresh = 0
        }

registerImage :: ImageAtlas -> ImageId -> Int -> Int -> ByteString -> IO Bool
registerImage (ImageAtlas ref) (ImageId tid) w h pixels
  | tid <= 0 || w <= 0 || h <= 0 = pure False
  | w > atlasMax - 2 * atlasPad || h > atlasMax - 2 * atlasPad = pure False
  | BS.length pixels < w * h * 4 = pure False
  | otherwise = do
      st0 <- readIORef ref
      case IM.lookup tid (asSlots st0) of
        Just slot
          | slotW slot == w && slotH slot == h -> do
              blitPixels (asPtr st0) (asW st0) (slotX slot) (slotY slot) w h pixels
              writeIORef ref st0 {asGen = asGen st0 + 1}
              pure True
          | otherwise -> pure False
        Nothing -> do
          mSt <- fitImage st0 tid w h pixels
          case mSt of
            Nothing -> pure False
            Just st1 -> do
              writeIORef ref st1
              pure True

-- | An id above every registered image's and every id this returned before.
-- An id the app picks itself can still collide with one returned and not yet
-- registered, so register those first.
freshImageId :: ImageAtlas -> IO ImageId
freshImageId (ImageAtlas ref) = do
  st <- readIORef ref
  let tid = 1 + maybe (asLastFresh st) (max (asLastFresh st) . fst) (IM.lookupMax (asSlots st))
  writeIORef ref st {asLastFresh = tid}
  pure (ImageId tid)

lookupImageUv ::
  ImageAtlas -> ImageId -> IO (Maybe (Float, Float, Float, Float))
lookupImageUv (ImageAtlas ref) (ImageId tid) = do
  st <- readIORef ref
  pure $
    case IM.lookup tid (asSlots st) of
      Nothing -> Nothing
      Just (AtlasSlot x y w h) ->
        let
          fw = fromIntegral (asW st)
          fh = fromIntegral (asH st)
         in
          Just
            ( fromIntegral x / fw
            , fromIntegral y / fh
            , fromIntegral (x + w) / fw
            , fromIntegral (y + h) / fh
            )

-- Pinned pixel buffer. SDL uploads this pointer; do not copy to ByteString first.
atlasSnapshot :: ImageAtlas -> IO (Maybe (Int, Int, ForeignPtr Word8, Int))
atlasSnapshot (ImageAtlas ref) = do
  st <- readIORef ref
  if asGen st == 0
    then pure Nothing
    else pure (Just (asW st, asH st, asPtr st, asGen st))

fitImage ::
  AtlasState -> Int -> Int -> Int -> ByteString -> IO (Maybe AtlasState)
fitImage st0 tid w h pixels =
  -- Plan the shelf position before allocating or copying the atlas. A full
  -- atlas must reject an image without repeatedly allocating doomed growth.
  case cursorFor st0 w h <|> cursorFor grown w h of
    Nothing -> pure Nothing
    Just (x, y, placed) -> do
      fp <-
        if asW placed == asW st0 && asH placed == asH st0
          then pure (asPtr st0)
          else do
            resized <- allocPixels (asW placed) (asH placed)
            copyAtlas (asPtr st0) (asW st0) (asH st0) resized (asW placed)
            pure resized
      blitPixels fp (asW placed) x y w h pixels
      pure $
        Just
          placed
            { asPtr = fp
            , asSlots = IM.insert tid (AtlasSlot x y w h) (asSlots placed)
            , asX = x + w + atlasPad
            , asY = y
            , asRowH = max (asRowH placed) h
            , asGen = asGen placed + 1
            }
 where
  grown =
    st0
      { asW = growDim (asW st0) (w + 2 * atlasPad)
      , asH = growDim (asH st0) (asY st0 + asRowH st0 + h + 2 * atlasPad)
      }

cursorFor :: AtlasState -> Int -> Int -> Maybe (Int, Int, AtlasState)
cursorFor st w h
  | asX st + w + atlasPad <= asW st && asY st + h + atlasPad <= asH st =
      Just (asX st, asY st, st)
  | asY st + asRowH st + atlasPad + h + atlasPad <= asH st
      && w + 2 * atlasPad <= asW st =
      let
        y = asY st + asRowH st + atlasPad
       in
        Just (atlasPad, y, st {asX = atlasPad, asY = y, asRowH = 0})
  | otherwise = Nothing

growDim :: Int -> Int -> Int
growDim cur need
  | need <= cur = cur
  | otherwise = min atlasMax (max need (cur * 2))

allocPixels :: Int -> Int -> IO (ForeignPtr Word8)
allocPixels w h = do
  let
    n = w * h * 4
  fp <- mallocForeignPtrBytes n
  withForeignPtr fp $ \p -> fillBytes p 0 n
  pure fp

copyAtlas :: ForeignPtr Word8 -> Int -> Int -> ForeignPtr Word8 -> Int -> IO ()
copyAtlas src oldW oldH dst newW =
  withForeignPtr src $ \sp ->
    withForeignPtr dst $ \dp ->
      mapM_ (copyRow sp dp) [0 .. oldH - 1]
 where
  rowBytes = oldW * 4
  copyRow sp dp row =
    copyBytes
      (dp `plusPtr` (row * newW * 4))
      (sp `plusPtr` (row * oldW * 4))
      rowBytes

blitPixels ::
  ForeignPtr Word8 -> Int -> Int -> Int -> Int -> Int -> ByteString -> IO ()
blitPixels dest destW destX destY w h pixels =
  withForeignPtr dest $ \dp ->
    BS.useAsCStringLen pixels $ \(sp, _) ->
      mapM_ (copyRow dp sp) [0 .. h - 1]
 where
  copyRow dp sp row =
    copyBytes
      (dp `plusPtr` (((destY + row) * destW + destX) * 4))
      (sp `plusPtr` (row * w * 4))
      (w * 4)