ips/src/IPS/Apply.hs

module IPS.Apply (apply, applyM, unsafeApplyM) where

import IPS.Types
import Data.Vector.Storable.Mutable (MVector)
import qualified Data.Vector.Storable.Mutable as MVector
import qualified Data.Vector.Storable as Vector
import Control.Monad.ST
import Control.Monad.Primitive (PrimMonad (PrimState))
import Data.Foldable


type BytesM m = MVector (PrimState m) Word8

-- size of patch output
-- (patches can write past the end of the input file)
outSize :: Patch -> Int
outSize = foldl' (\m c -> max m (end c)) 0 where
  end (Chunk {offset, body}) = case body of
    Normal bs  -> offset' + Vector.length bs
    RLE {size} -> offset' + fromIntegral size
    where offset' = fromIntegral offset


-- applies in place if it fits, otherwise allocates a new buffer for it
-- returns 'Nothing' if the contents are in the original buffer
applyM :: PrimMonad m => Patch -> BytesM m -> m (Maybe (BytesM m))
applyM patch buf = do
  (buf', realloc) <- maybeNewBuf buf $ outSize patch
  unsafeApplyM patch buf'
  pure $ if realloc then Just buf' else Nothing

-- if the buffer is smaller than the given size, then allocate a new one
-- of that size and copy the contents to the beginning of it
maybeNewBuf :: PrimMonad m => BytesM m -> Int -> m (BytesM m, Bool)
maybeNewBuf buf size
  | MVector.length buf >= size = pure (buf, False)
  | otherwise = do
      newBuf <- MVector.new size
      MVector.copy (MVector.slice 0 (MVector.length buf) newBuf) buf
      pure (newBuf, True)


-- copies an immutable vector (resizing if needed) and patches it
apply :: Patch -> Bytes -> Bytes
apply patch input = runST do
  buf <- MVector.new $ outSize patch `max` Vector.length input
  Vector.copy (MVector.slice 0 (Vector.length input) buf) input
  unsafeApplyM patch buf
  Vector.unsafeFreeze buf


-- explodes if the patch tries to write outside the existing space
unsafeApplyM :: PrimMonad m => Patch -> BytesM m -> m ()
unsafeApplyM patch buf =
  for_ patch \Chunk {offset, body} ->
    let makeSlice size = MVector.slice (fromIntegral offset) size buf in
    case body of
      Normal bs         -> Vector.copy (makeSlice $ Vector.length bs) bs
      RLE {size, value} -> MVector.set (makeSlice $ fromIntegral size) value