61 lines
2.2 KiB
Haskell
61 lines
2.2 KiB
Haskell
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
|
|
|