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 :: Int
atlasTextureId = Int
1

atlasPad :: Int
atlasPad :: Int
atlasPad = Int
1

atlasStart :: Int
atlasStart :: Int
atlasStart = Int
256

atlasMax :: Int
atlasMax :: Int
atlasMax = Int
4096

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

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

newtype ImageAtlas = ImageAtlas (IORef AtlasState)

newImageAtlas :: IO ImageAtlas
newImageAtlas :: IO ImageAtlas
newImageAtlas = do
  fp <- Int -> Int -> IO (ForeignPtr Word8)
allocPixels Int
atlasStart Int
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 -> ImageId -> Int -> Int -> ByteString -> IO Bool
registerImage (ImageAtlas IORef AtlasState
ref) (ImageId Int
tid) Int
w Int
h ByteString
pixels
  | Int
tid Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
|| Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
  | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
atlasMax Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
atlasPad Bool -> Bool -> Bool
|| Int
h Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
atlasMax Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
atlasPad = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
  | ByteString -> Int
BS.length ByteString
pixels Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
  | Bool
otherwise = do
      st0 <- IORef AtlasState -> IO AtlasState
forall a. IORef a -> IO a
readIORef IORef AtlasState
ref
      case IM.lookup tid (asSlots st0) of
        Just AtlasSlot
slot
          | AtlasSlot -> Int
slotW AtlasSlot
slot Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w Bool -> Bool -> Bool
&& AtlasSlot -> Int
slotH AtlasSlot
slot Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
h -> do
              ForeignPtr Word8
-> Int -> Int -> Int -> Int -> Int -> ByteString -> IO ()
blitPixels (AtlasState -> ForeignPtr Word8
asPtr AtlasState
st0) (AtlasState -> Int
asW AtlasState
st0) (AtlasSlot -> Int
slotX AtlasSlot
slot) (AtlasSlot -> Int
slotY AtlasSlot
slot) Int
w Int
h ByteString
pixels
              IORef AtlasState -> AtlasState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef AtlasState
ref AtlasState
st0 {asGen = asGen st0 + 1}
              Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
          | Bool
otherwise -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
        Maybe AtlasSlot
Nothing -> do
          mSt <- AtlasState
-> Int -> Int -> Int -> ByteString -> IO (Maybe AtlasState)
fitImage AtlasState
st0 Int
tid Int
w Int
h ByteString
pixels
          case mSt of
            Maybe AtlasState
Nothing -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
            Just AtlasState
st1 -> do
              IORef AtlasState -> AtlasState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef AtlasState
ref AtlasState
st1
              Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
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 -> IO ImageId
freshImageId (ImageAtlas IORef AtlasState
ref) = do
  st <- IORef AtlasState -> IO AtlasState
forall a. IORef a -> IO a
readIORef IORef AtlasState
ref
  let tid = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> ((Int, AtlasSlot) -> Int) -> Maybe (Int, AtlasSlot) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (AtlasState -> Int
asLastFresh AtlasState
st) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (AtlasState -> Int
asLastFresh AtlasState
st) (Int -> Int)
-> ((Int, AtlasSlot) -> Int) -> (Int, AtlasSlot) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, AtlasSlot) -> Int
forall a b. (a, b) -> a
fst) (IntMap AtlasSlot -> Maybe (Int, AtlasSlot)
forall a. IntMap a -> Maybe (Int, a)
IM.lookupMax (AtlasState -> IntMap AtlasSlot
asSlots AtlasState
st))
  writeIORef ref st {asLastFresh = tid}
  pure (ImageId tid)

lookupImageUv ::
  ImageAtlas -> ImageId -> IO (Maybe (Float, Float, Float, Float))
lookupImageUv :: ImageAtlas -> ImageId -> IO (Maybe (Float, Float, Float, Float))
lookupImageUv (ImageAtlas IORef AtlasState
ref) (ImageId Int
tid) = do
  st <- IORef AtlasState -> IO AtlasState
forall a. IORef a -> IO a
readIORef IORef AtlasState
ref
  pure $
    case IM.lookup tid (asSlots st) of
      Maybe AtlasSlot
Nothing -> Maybe (Float, Float, Float, Float)
forall a. Maybe a
Nothing
      Just (AtlasSlot Int
x Int
y Int
w Int
h) ->
        let
          fw :: Float
fw = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (AtlasState -> Int
asW AtlasState
st)
          fh :: Float
fh = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (AtlasState -> Int
asH AtlasState
st)
         in
          (Float, Float, Float, Float) -> Maybe (Float, Float, Float, Float)
forall a. a -> Maybe a
Just
            ( Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
fw
            , Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
y Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
fh
            , Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
fw
            , Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
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 -> IO (Maybe (Int, Int, ForeignPtr Word8, Int))
atlasSnapshot (ImageAtlas IORef AtlasState
ref) = do
  st <- IORef AtlasState -> IO AtlasState
forall a. IORef a -> IO a
readIORef IORef AtlasState
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 :: AtlasState
-> Int -> Int -> Int -> ByteString -> IO (Maybe AtlasState)
fitImage AtlasState
st0 Int
tid Int
w Int
h ByteString
pixels =
  -- Plan the shelf position before allocating or copying the atlas. A full
  -- atlas must reject an image without repeatedly allocating doomed growth.
  case AtlasState -> Int -> Int -> Maybe (Int, Int, AtlasState)
cursorFor AtlasState
st0 Int
w Int
h Maybe (Int, Int, AtlasState)
-> Maybe (Int, Int, AtlasState) -> Maybe (Int, Int, AtlasState)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> AtlasState -> Int -> Int -> Maybe (Int, Int, AtlasState)
cursorFor AtlasState
grown Int
w Int
h of
    Maybe (Int, Int, AtlasState)
Nothing -> Maybe AtlasState -> IO (Maybe AtlasState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe AtlasState
forall a. Maybe a
Nothing
    Just (Int
x, Int
y, AtlasState
placed) -> do
      fp <-
        if AtlasState -> Int
asW AtlasState
placed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== AtlasState -> Int
asW AtlasState
st0 Bool -> Bool -> Bool
&& AtlasState -> Int
asH AtlasState
placed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== AtlasState -> Int
asH AtlasState
st0
          then ForeignPtr Word8 -> IO (ForeignPtr Word8)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AtlasState -> ForeignPtr Word8
asPtr AtlasState
st0)
          else do
            resized <- Int -> Int -> IO (ForeignPtr Word8)
allocPixels (AtlasState -> Int
asW AtlasState
placed) (AtlasState -> Int
asH AtlasState
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 :: AtlasState
grown =
    AtlasState
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 :: AtlasState -> Int -> Int -> Maybe (Int, Int, AtlasState)
cursorFor AtlasState
st Int
w Int
h
  | AtlasState -> Int
asX AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atlasPad Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= AtlasState -> Int
asW AtlasState
st Bool -> Bool -> Bool
&& AtlasState -> Int
asY AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atlasPad Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= AtlasState -> Int
asH AtlasState
st =
      (Int, Int, AtlasState) -> Maybe (Int, Int, AtlasState)
forall a. a -> Maybe a
Just (AtlasState -> Int
asX AtlasState
st, AtlasState -> Int
asY AtlasState
st, AtlasState
st)
  | AtlasState -> Int
asY AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ AtlasState -> Int
asRowH AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atlasPad Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atlasPad Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= AtlasState -> Int
asH AtlasState
st
      Bool -> Bool -> Bool
&& Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
atlasPad Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= AtlasState -> Int
asW AtlasState
st =
      let
        y :: Int
y = AtlasState -> Int
asY AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ AtlasState -> Int
asRowH AtlasState
st Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atlasPad
       in
        (Int, Int, AtlasState) -> Maybe (Int, Int, AtlasState)
forall a. a -> Maybe a
Just (Int
atlasPad, Int
y, AtlasState
st {asX = atlasPad, asY = y, asRowH = 0})
  | Bool
otherwise = Maybe (Int, Int, AtlasState)
forall a. Maybe a
Nothing

growDim :: Int -> Int -> Int
growDim :: Int -> Int -> Int
growDim Int
cur Int
need
  | Int
need Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
cur = Int
cur
  | Bool
otherwise = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
atlasMax (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
need (Int
cur Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2))

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

copyAtlas :: ForeignPtr Word8 -> Int -> Int -> ForeignPtr Word8 -> Int -> IO ()
copyAtlas :: ForeignPtr Word8 -> Int -> Int -> ForeignPtr Word8 -> Int -> IO ()
copyAtlas ForeignPtr Word8
src Int
oldW Int
oldH ForeignPtr Word8
dst Int
newW =
  ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
src ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp ->
    ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
dst ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dp ->
      (Int -> IO ()) -> [Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Ptr Word8 -> Ptr Word8 -> Int -> IO ()
copyRow Ptr Word8
sp Ptr Word8
dp) [Int
0 .. Int
oldH Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
 where
  rowBytes :: Int
rowBytes = Int
oldW Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
  copyRow :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
copyRow Ptr Word8
sp Ptr Word8
dp Int
row =
    Ptr (ZonkAny 1) -> Ptr (ZonkAny 1) -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes
      (Ptr Word8
dp Ptr Word8 -> Int -> Ptr (ZonkAny 1)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
newW Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4))
      (Ptr Word8
sp Ptr Word8 -> Int -> Ptr (ZonkAny 1)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
oldW Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4))
      Int
rowBytes

blitPixels ::
  ForeignPtr Word8 -> Int -> Int -> Int -> Int -> Int -> ByteString -> IO ()
blitPixels :: ForeignPtr Word8
-> Int -> Int -> Int -> Int -> Int -> ByteString -> IO ()
blitPixels ForeignPtr Word8
dest Int
destW Int
destX Int
destY Int
w Int
h ByteString
pixels =
  ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
dest ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dp ->
    ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.useAsCStringLen ByteString
pixels ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
sp, Int
_) ->
      (Int -> IO ()) -> [Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Ptr Word8 -> Ptr CChar -> Int -> IO ()
copyRow Ptr Word8
dp Ptr CChar
sp) [Int
0 .. Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
 where
  copyRow :: Ptr Word8 -> Ptr CChar -> Int -> IO ()
copyRow Ptr Word8
dp Ptr CChar
sp Int
row =
    Ptr (ZonkAny 0) -> Ptr (ZonkAny 0) -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes
      (Ptr Word8
dp Ptr Word8 -> Int -> Ptr (ZonkAny 0)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (((Int
destY Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
row) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
destW Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
destX) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4))
      (Ptr CChar
sp Ptr CChar -> Int -> Ptr (ZonkAny 0)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4))
      (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)