BufferIterator.hs

{-# LANGUAGE BangPatterns, RecordWildCards #-}
-- | Module for iterating over buffers
module GEGL.Buffer.BufferIterator
  ( FFI.GeglBufferIterator(..)
  , gegl_buffer_iterator_new
  , gegl_buffer_iterator_next
  , Pixel(..)
  , FFI.ComponentValue(..)
  , pixelMap
  , pixelPoke
  ) where

import qualified GEGL.FFI.Buffer.BufferIterator as FFI

import GEGL.FFI.Buffer (GeglBuffer(..))
import GEGL.FFI.Rectangle

import GEGL.Enums
import GEGL.Color

import qualified BABL.Format as BF
import BABL.FFI.Format (BablFormatPtr(..))

import Foreign.ForeignPtr
import Foreign.Marshal.Utils (new)
import Foreign.Ptr
import Foreign.C.Types
import Foreign.Storable (peek, poke)

import Data.List.Split
import Data.Maybe (fromMaybe)
import Control.Monad.Loops (whileM_)

-- DEBUGGING

import Debug.Trace

-- | Data type to represent Pixels in iteration.
data Pixel = Pixel
  { pixelX     :: Int   -- ^ X coordinate
  , pixelY     :: Int   -- ^ Y coordinate
  , pixelData  ::
      ( FFI.ComponentValue
      , FFI.ComponentValue
      , FFI.ComponentValue
      , FFI.ComponentValue
      ) -- ^ Colour data
  }

-- | Create a new iterator over a specified area on a buffer.
gegl_buffer_iterator_new
  :: GeglBuffer      -- ^ Buffer to draw data from
  -> GeglRectangle   -- ^ Rectangle to iterate over
  -> BF.PixelFormat  -- ^ Format to process pixel data in
  -> GeglAccessMode  -- ^ Access mode to pixel data
  -> GeglAbyssPolicy -- ^ Abyss policy for pixel data
  -> IO FFI.GeglBufferIterator
gegl_buffer_iterator_new (GeglBuffer fbuf) roi format am ap = do
  roiPtr <- new roi
  (BablFormatPtr formatPtr) <- BF.babl_format format
  withForeignPtr fbuf $ \buf ->
    FFI.GeglBufferIterator <$> FFI.c_gegl_buffer_iterator_new
      buf
      roiPtr
      (CInt 0)
      formatPtr
      (marshal am)
      (marshal ap)

-- | This function fills the 'FFI.GeglBufferIterator' with its appropriate data
--   and returns 'True' as long as there is something to iterate over. Upon
--   returning 'False' the iterator is invalidated.
gegl_buffer_iterator_next
  :: FFI.GeglBufferIterator -- ^ Iterator to fill with data
  -> IO Bool
gegl_buffer_iterator_next (FFI.GeglBufferIterator ptr) =
  FFI.c_gegl_buffer_iterator_next ptr

-- | marshal an 'Enum' to a 'CInt'
marshal :: Enum e => e -> CInt
marshal = CInt . fromIntegral . fromEnum

-- | Simple iteration over a rectangle section of a buffer. Generates its own
--   iterator.
pixelPoke
  :: GeglBuffer               -- ^ Terget buffer
  -> GeglRectangle            -- ^ Rectangle to iterate over
  -> BF.PixelFormat           -- ^ Pixel format
  -> GeglAccessMode           -- ^ Access mode to pixel data
  -> GeglAbyssPolicy          -- ^ Abbyss policy for pixel data
  -> ((Int, Int) -> Pixel)    -- ^ Function to apply to each pixel
  -> IO ()
pixelPoke buf rect form am ap func = do
  (FFI.GeglBufferIterator iterPtr) <- gegl_buffer_iterator_new buf rect form am ap
  whileM_ (gegl_buffer_iterator_next (FFI.GeglBufferIterator iterPtr))
    (do
      roi <- FFI.c_gegl_peek_roi iterPtr
      let idx = [1..(rectangleWidth roi * rectangleHeight roi)]
          !newList = concatMap (pToC $ BF.babl_components_per_pixel form) $ map func $ map (flip iToCoord roi) idx
      FFI.c_gegl_poke_data iterPtr newList form
      )

-- | Simple iteration over a rectangle section of a buffer. Generates its own
--   iterator.
pixelMap
  :: GeglBuffer               -- ^ Terget buffer
  -> GeglRectangle            -- ^ Rectangle to iterate over
  -> BF.PixelFormat           -- ^ Pixel format
  -> GeglAccessMode           -- ^ Access mode to pixel data
  -> GeglAbyssPolicy          -- ^ Abbyss policy for pixel data
  -> (Pixel -> Pixel)         -- ^ Function to apply to each pixel
  -> IO ()
pixelMap buf rect form am ap func = do
  (FFI.GeglBufferIterator iterPtr) <- gegl_buffer_iterator_new buf rect form am ap
  whileM_ (gegl_buffer_iterator_next (FFI.GeglBufferIterator iterPtr))
    (do
      values <- FFI.c_gegl_peek_data iterPtr form
      roi <- FFI.c_gegl_peek_roi iterPtr
      let newList = mapPixel values form roi func
      FFI.c_gegl_poke_data iterPtr newList form
      )

-- | Function to map over all Pixel data and return them back altered
mapPixel
  :: [FFI.ComponentValue] -- ^ Input pixel data
  -> BF.PixelFormat       -- ^ PixelFormat of target buffer
  -> GeglRectangle        -- ^ Rectangle of Interest
  -> (Pixel -> Pixel)     -- ^ Function to apply on all Pixels
  -> [FFI.ComponentValue] -- ^ Resulting pixel data
mapPixel input format rect@GeglRectangle{..} funct =
  concatMap (pToC cpp) $ map funct $ map (cToP rect cpp) chunks
  where
    !cpp = BF.babl_components_per_pixel format
    chunks = zip [0..] $ chunksOf cpp input

iToCoord i GeglRectangle{..} =
  ( (rectangleX + ((i - 1) `mod` (rectangleWidth)))
  , (rectangleY + ((i - 1) `div` (rectangleWidth)))
  )

cToP rect cpp (i, pd) = Pixel (fst tup) (snd tup) $ toPixel cpp pd
  where
    tup = (iToCoord i rect)

pToC cpp (Pixel _ _ col) = toList cpp col

-- | Function to fill a Pixel with its colour data
toPixel
  :: Int                  -- Components per Pixel
  -> [FFI.ComponentValue] -- Color Data
  -> (FFI.ComponentValue, FFI.ComponentValue, FFI.ComponentValue, FFI.ComponentValue) -- ^ Pixel data
toPixel cpp d = case cpp of
  1 -> (head d, nullc , nullc , nullc)
  2 -> (head d, d !! 1, nullc , nullc)
  3 -> (head d, d !! 1, d !! 2, nullc)
  4 -> (head d, d !! 1, d !! 2, d !! 3)
  where
    !nullc = FFI.CVu32 0

-- | Function to turn Pixel data back to a list
toList
  :: Int -- ^ Components per Pixel
  -> (FFI.ComponentValue, FFI.ComponentValue, FFI.ComponentValue, FFI.ComponentValue) -- ^ Pixel data
  -> [FFI.ComponentValue] -- ^ Color Data
toList cpp (a, b ,c ,d) =
  case cpp of
   1 -> [a]
   2 -> [a, b]
   3 -> [a, b, c]
   4 -> [a, b, c, d]