GPUComputeSamplerExample.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Example     : GPUComputeSampler
-- Description : Samples a source texture in a compute shader using various sampler states and blits the result.
-- Copyright   : (c) Kyle Lukaszek, 2025
-- License     : BSD3
--
-- Based on SDL_gpu_examples/ComputeSampler example.
-- Demonstrates:
-- - Loading a source texture and creating a separate writeable texture.
-- - Using a compute shader to process the source texture (applying different sampler effects)
--   and write to the destination texture.
-- - Binding the source texture and a user-selectable sampler to the compute shader.
-- - Pushing a uniform (texcoord multiplier) to the compute shader.
-- - Blitting the compute shader's output texture to the screen.
-- - User can cycle through sampler states (Point, Linear, Anisotropic, Wrap/Clamp) via arrow keys.
-- |
module Main where

import Control.Exception (bracket, bracketOnError, finally)
import Control.Monad (forM, forM_, unless, void, when)
import Data.Bits ((.|.))
import Data.IORef
import Data.Maybe (fromJust, isJust)
import Data.Word (Word32)
import Foreign.C.Types (CFloat)
import Foreign.Storable (peek)
import GPUCommon
import SDL3
import System.Exit (exitFailure, exitSuccess)
import System.FilePath ((</>))

samplerCount :: Int
samplerCount = 6

samplerNames :: [String]
samplerNames =
  [ "PointClamp",
    "PointWrap",
    "LinearClamp",
    "LinearWrap",
    "AnisotropicClamp",
    "AnisotropicWrap"
  ]

samplerCreateInfos :: [SDLGPUSamplerCreateInfo]
samplerCreateInfos =
  [ (defaultSamplerCreateInfo SDL_GPU_FILTER_NEAREST)
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE
      },
    (defaultSamplerCreateInfo SDL_GPU_FILTER_NEAREST)
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_REPEAT
      },
    (defaultSamplerCreateInfo SDL_GPU_FILTER_LINEAR)
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE
      },
    (defaultSamplerCreateInfo SDL_GPU_FILTER_LINEAR)
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_REPEAT
      },
    (defaultSamplerCreateInfo SDL_GPU_FILTER_LINEAR) -- Anisotropic
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        samplerEnableAnisotropy = True,
        samplerMaxAnisotropy = 4.0
      },
    (defaultSamplerCreateInfo SDL_GPU_FILTER_LINEAR) -- Anisotropic
      { samplerAddressModeU = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
        samplerAddressModeV = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
        samplerEnableAnisotropy = True,
        samplerMaxAnisotropy = 4.0
      }
  ]

-- AppResources
data AppResources = AppResources
  { resComputePipeline :: SDLGPUComputePipeline,
    resSourceTexture :: SDLGPUTexture, -- Loaded from image, read by compute
    resWriteTexture :: SDLGPUTexture, -- Written by compute, read by blit
    resSamplers :: [SDLGPUSampler], -- Array of samplers for compute
    currentSamplerIndexRef :: IORef Int -- For UI
  }

-- main
main :: IO ()
main = do
  sdlLog $ "Compiled SDL Version: " ++ show sdlVersion
  linkedVersion <- sdlGetVersion
  sdlLog $ "Linked SDL Version: " ++ show linkedVersion
  maybeResult <- withContext "SDL3 Haskell GPU Compute Textured Quad" [] runAppGPU
  case maybeResult of
    Nothing -> sdlLog "Application initialization failed." >> exitFailure
    Just _ -> sdlLog "Application finished successfully." >> exitSuccess

-- runAppGPU
runAppGPU :: Context -> IO ()
runAppGPU context = do
  sdlLog "Base context initialized."
  bracket
    (createResources context)
    (releaseResources context)
    $ \case
      Nothing -> sdlLog "Failed to create resources. Exiting."
      Just resources -> do
        sdlLog "Resources created successfully."
        sdlLog "Press Left/Right to switch sampler states for compute shader"
        initialSamplerIdx <- readIORef (currentSamplerIndexRef resources)
        sdlLog $ "Initial sampler for compute: " ++ (samplerNames !! initialSamplerIdx)
        eventLoopGPU context resources

-- createResources
createResources :: Context -> IO (Maybe AppResources)
createResources Context {..} = do
  sdlLog "--- Beginning Resource Creation ---"

  -- 1. Load Source Image ("ravioli.bmp")
  maybeSurf <-
    bracketOnError
      (loadImage ("Content" </> "Images" </> "ravioli.bmp"))
      (\case Just surf -> sdlDestroySurface surf; _ -> pure ())
      pure

  case maybeSurf of
    Nothing -> sdlLog "!!! Failed to load source image." >> return Nothing
    Just surfacePtr ->
      bracketOnError (pure surfacePtr) (\_ -> sdlDestroySurface surfacePtr) $ \surf -> do
        surfaceData <- peek surf
        let imgW = surfaceW surfaceData :: Int
        let imgH = surfaceH surfaceData :: Int

        -- 2. Create Source Texture (for compute shader to read)
        let srcTextureCI =
              (defaultTextureCreateInfo imgW imgH)
                { texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER -- Compute will sample this
                }
        maybeSrcTex <- sdlCreateGPUTexture contextDevice srcTextureCI
        forM_ maybeSrcTex $ \tex -> sdlSetGPUTextureName contextDevice tex "Ravioli Texture (Source)"

        -- 3. Create Write Texture (for compute shader to write to, fixed size 640x480)
        let writeTexW = 640
        let writeTexH = 480
        let writeTextureCI =
              (defaultTextureCreateInfo writeTexW writeTexH)
                { texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER .|. SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_WRITE
                }
        maybeWriteTex <- sdlCreateGPUTexture contextDevice writeTextureCI

        -- 4. Create Compute Pipeline ("TexturedQuad.comp")
        let computeCI =
              defaultComputePipelineCreateInfo
                { numSamplers = 1, -- Compute shader uses 1 sampler
                  numReadWriteStorageTextures = 1, -- Writes to WriteTexture
                  numUniformBuffers = 1, -- For texcoordMultiplier
                  threadCountX = 8,
                  threadCountY = 8,
                  threadCountZ = 1
                }
        maybeComputePipe <- createComputePipelineFromShader contextDevice "TexturedQuad.comp" computeCI

        -- 5. Create Samplers
        samplersList <- forM samplerCreateInfos (sdlCreateGPUSampler contextDevice)
        let maybeSamplers = sequence samplersList -- Convert [Maybe a] to Maybe [a]

        -- Upload Source Texture Data
        uploadSuccess <- case maybeSrcTex of
          Nothing -> return False
          Just srcTexToUploadTo ->
            bracket
              (createTransferBuffer contextDevice (fromIntegral $ imgW * imgH * 4) SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD "SrcTexTransfer")
              (cleanupTransferBuffer contextDevice)
              $ \case
                Nothing -> return False
                Just texTransfer -> do
                  mapOk <- mapAndCopyTextureData contextDevice texTransfer surf imgW imgH 4
                  if mapOk
                    then bracket
                      (sdlAcquireGPUCommandBuffer contextDevice)
                      cleanupCommandBuffer
                      $ \case
                        Nothing -> return False
                        Just cmd -> do
                          mcp <- sdlBeginGPUCopyPass cmd
                          case mcp of
                            Nothing -> return False
                            Just cp -> do
                              let texSrcDesc = SDLGPUTextureTransferInfo texTransfer 0 (fromIntegral imgW) (fromIntegral imgH)
                              let texDstDesc = defaultTextureRegion srcTexToUploadTo imgW imgH
                              sdlUploadToGPUTexture cp texSrcDesc texDstDesc False
                              sdlEndGPUCopyPass cp
                              sdlSubmitGPUCommandBuffer cmd >>= \s -> if s then sdlWaitForGPUIdle contextDevice >> return True else return False
                    else return False

        if uploadSuccess && isJust maybeSrcTex && isJust maybeWriteTex && isJust maybeComputePipe && isJust maybeSamplers
          then do
            samplerIdxRef <- newIORef 0
            sdlLog "--- All core resources created and source texture uploaded ---"
            return $
              Just
                AppResources
                  { resComputePipeline = fromJust maybeComputePipe,
                    resSourceTexture = fromJust maybeSrcTex,
                    resWriteTexture = fromJust maybeWriteTex,
                    resSamplers = fromJust maybeSamplers,
                    currentSamplerIndexRef = samplerIdxRef
                  }
          else do
            sdlLog "!!! Failed to create one or more resources or upload source texture."
            maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeSrcTex
            maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeWriteTex
            maybe (pure ()) (sdlReleaseGPUComputePipeline contextDevice) maybeComputePipe
            maybe (pure ()) (mapM_ (sdlReleaseGPUSampler contextDevice)) (sequence samplersList)
            return Nothing

-- releaseResources
releaseResources :: Context -> Maybe AppResources -> IO ()
releaseResources _ Nothing = return ()
releaseResources Context {..} (Just AppResources {..}) = do
  sdlLog "--> Releasing AppResources..."
  sdlReleaseGPUComputePipeline contextDevice resComputePipeline
  sdlReleaseGPUTexture contextDevice resSourceTexture
  sdlReleaseGPUTexture contextDevice resWriteTexture
  forM_ resSamplers (sdlReleaseGPUSampler contextDevice)
  sdlLog "<-- AppResources Released."

-- eventLoopGPU
eventLoopGPU :: Context -> AppResources -> IO ()
eventLoopGPU context resources@AppResources {} = do
  sdlPumpEvents
  shouldQuitRef <- newIORef False
  processEventsGPU context resources shouldQuitRef

  shouldQuit <- readIORef shouldQuitRef
  unless shouldQuit $ do
    renderFrameGPU context resources
    eventLoopGPU context resources

-- processEventsGPU
processEventsGPU :: Context -> AppResources -> IORef Bool -> IO ()
processEventsGPU context resources quitRef = do
  maybeEvent <- sdlPollEvent
  case maybeEvent of
    Nothing -> return ()
    Just event -> do
      quit <- handleEventGPU context resources event
      when quit $ writeIORef quitRef True
      processEventsGPU context resources quitRef

-- handleEventGPU for UI interaction (sampler switching)
handleEventGPU :: Context -> AppResources -> SDLEvent -> IO Bool
handleEventGPU _ AppResources {..} event = case event of
  SDLEventQuit _ -> sdlLog "Quit event received." >> return True
  SDLEventKeyboard (SDLKeyboardEvent _ _ _ _ scancode _ _ _ down _) | down -> do
    case scancode of
      SDL_SCANCODE_Q -> return True
      SDL_SCANCODE_LEFT -> do
        idx <- atomicModifyIORef' currentSamplerIndexRef $ \i ->
          let newIdx = (i - 1 + samplerCount) `mod` samplerCount
           in (newIdx, newIdx)
        sdlLog $ "Compute Sampler: " ++ (samplerNames !! idx)
        return False
      SDL_SCANCODE_RIGHT -> do
        idx <- atomicModifyIORef' currentSamplerIndexRef $ \i ->
          let newIdx = (i + 1) `mod` samplerCount
           in (newIdx, newIdx)
        sdlLog $ "Compute Sampler: " ++ (samplerNames !! idx)
        return False
      _ -> return False
  _ -> return False

-- renderFrameGPU
renderFrameGPU :: Context -> AppResources -> IO ()
renderFrameGPU Context {..} AppResources {..} = do
  maybeCmdbuf <- sdlAcquireGPUCommandBuffer contextDevice
  case maybeCmdbuf of
    Nothing -> sdlLog "Error: Failed to acquire render command buffer."
    Just cmdbuf -> do
      maybeSwapResult <- sdlWaitAndAcquireGPUSwapchainTexture cmdbuf contextWindow
      case maybeSwapResult of
        Nothing -> void (sdlSubmitGPUCommandBuffer cmdbuf `finally` pure ())
        Just (swapchainTexture, swapW, swapH) -> do
          -- Texture dimensions for dispatch (WriteTexture is 640x480)
          let dispatchW = 640 :: Word32
          let dispatchH = 480 :: Word32

          -- 1. Compute Pass to process/write to resWriteTexture
          currentIndex <- readIORef currentSamplerIndexRef
          let currentSamplerForCompute = resSamplers !! currentIndex

          maybeComputePass <-
            sdlBeginGPUComputePass
              cmdbuf
              [SDLGPUStorageTextureReadWriteBinding resWriteTexture 0 0 True] -- Write to WriteTexture, cycle
              [] -- No storage buffers
          case maybeComputePass of
            Nothing -> sdlLog "!!! Failed to begin Compute pass."
            Just computePass -> do
              sdlBindGPUComputePipeline computePass resComputePipeline
              -- Bind source texture and current sampler for compute shader to read
              let computeSamplerBinding = SDLGPUTextureSamplerBinding resSourceTexture currentSamplerForCompute
              sdlBindGPUComputeSamplers computePass 0 [computeSamplerBinding]
              -- Push texcoordMultiplier uniform
              let texcoordMultiplier = 0.25 :: CFloat
              sdlPushGPUComputeUniformData cmdbuf 0 texcoordMultiplier
              sdlDispatchGPUCompute computePass (dispatchW `div` 8) (dispatchH `div` 8) 1
              sdlEndGPUComputePass computePass

          -- 2. Blit resWriteTexture to Swapchain
          let srcBlitRegion = defaultBlitRegion resWriteTexture (fromIntegral dispatchW) (fromIntegral dispatchH)
          let dstBlitRegion =
                (defaultBlitRegion swapchainTexture (fromIntegral swapW) (fromIntegral swapH))
                  { gpuBlitRegW = swapW,
                    gpuBlitRegH = swapH
                  }

          let blitInfo =
                SDLGPUBlitInfo
                  { gpuBlitInfoSource = srcBlitRegion,
                    gpuBlitInfoDestination = dstBlitRegion,
                    gpuBlitInfoLoadOp = SDL_GPU_LOADOP_DONT_CARE,
                    gpuBlitInfoClearColor = SDLFColor 0 0 0 0,
                    gpuBlitInfoFlipMode = SDL_FLIP_NONE,
                    gpuBlitInfoFilter = SDL_GPU_FILTER_NEAREST,
                    gpuBlitInfoCycle = False
                  }
          sdlBlitGPUTexture cmdbuf blitInfo

          -- Submit
          submitted <- sdlSubmitGPUCommandBuffer cmdbuf
          unless submitted $ sdlGetError >>= sdlLog . ("Submit failed: " ++)