-- |
-- Module      : Language.Halide.JuicyPixels
-- Copyright   : (c) Tom Westerhout, 2023
--
-- This package allows you to use [JuicyPixels](https://hackage.haskell.org/package/JuicyPixels)
-- together with [halide-haskell](https://hackage.haskell.org/package/halide-haskell).
-- It defines 'Language.Halide.IsHalideBuffer' instances for 'Image' and
-- 'MutableImage' types from the JuicyPixels library.
--
-- That allows you to write code such as:
--
-- @
-- kernel :: Ptr ('HalideBuffer' 3 Word8) -> Ptr ('HalideBuffer' 3 Word8) -> IO ()
-- kernel = ...
--
-- brighten :: 'Image' 'PixelRGB8' -> 'MutableImage' 'RealWorld' 'PixelRGB8' -> IO ()
-- brighten input output = do
--   'withHalideBuffer' @3 @Word8 input $ \input' ->
--     'withHalideBuffer' @3 @Word8 output $ \output' ->
--       kernel input' output'
-- @
module Language.Halide.JuicyPixels ()
where

import Codec.Picture
import Codec.Picture.Types
import Control.Monad.ST (RealWorld)
import Data.Vector.Storable qualified as S
import Data.Vector.Storable.Mutable qualified as SM
import Language.Halide

instance (Pixel a, r ~ PixelBaseComponent a, IsHalideType r) => IsHalideBuffer (Image a) 3 r where
  withHalideBufferImpl :: Image a -> (Ptr (HalideBuffer 3 r) -> IO b) -> IO b
  withHalideBufferImpl :: forall b. Image a -> (Ptr (HalideBuffer 3 r) -> IO b) -> IO b
withHalideBufferImpl Image a
im Ptr (HalideBuffer 3 r) -> IO b
action =
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
S.unsafeWith Image a
im.imageData forall a b. (a -> b) -> a -> b
$ \Ptr r
cpuPtr ->
      forall (n :: Nat) a b.
(HasCallStack, KnownNat n, IsHalideType a) =>
Ptr a -> [Int] -> (Ptr (HalideBuffer n a) -> IO b) -> IO b
bufferFromPtrShape Ptr r
cpuPtr [forall a. Pixel a => a -> Int
componentCount (forall a. HasCallStack => a
undefined :: a), Image a
im.imageWidth, Image a
im.imageHeight] Ptr (HalideBuffer 3 r) -> IO b
action

instance (Pixel a, r ~ PixelBaseComponent a, IsHalideType r) => IsHalideBuffer (MutableImage RealWorld a) 3 r where
  withHalideBufferImpl :: MutableImage RealWorld a -> (Ptr (HalideBuffer 3 r) -> IO b) -> IO b
  withHalideBufferImpl :: forall b.
MutableImage RealWorld a
-> (Ptr (HalideBuffer 3 r) -> IO b) -> IO b
withHalideBufferImpl MutableImage RealWorld a
im Ptr (HalideBuffer 3 r) -> IO b
action =
    forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
SM.unsafeWith MutableImage RealWorld a
im.mutableImageData forall a b. (a -> b) -> a -> b
$ \Ptr r
cpuPtr ->
      forall (n :: Nat) a b.
(HasCallStack, KnownNat n, IsHalideType a) =>
Ptr a -> [Int] -> (Ptr (HalideBuffer n a) -> IO b) -> IO b
bufferFromPtrShape
        Ptr r
cpuPtr
        [forall a. Pixel a => a -> Int
componentCount (forall a. HasCallStack => a
undefined :: a), MutableImage RealWorld a
im.mutableImageWidth, MutableImage RealWorld a
im.mutableImageHeight]
        Ptr (HalideBuffer 3 r) -> IO b
action