63 lines
2.0 KiB
Haskell
63 lines
2.0 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
|
|
|
|
-- explodes if the patch tries to write outside the existing space
|
|
unsafeApplyM :: PrimMonad m => Patch -> BytesM m -> m ()
|
|
unsafeApplyM p buf =
|
|
for_ p \Chunk {offset, body} ->
|
|
case body of
|
|
Normal bs -> do
|
|
let offset' = fromIntegral offset
|
|
let sz = Vector.length bs
|
|
let dst = MVector.slice offset' sz buf
|
|
Vector.copy dst bs
|
|
RLE {size, value} -> do
|
|
let offset' = fromIntegral offset
|
|
let sz = fromIntegral size
|
|
let dst = MVector.slice offset' sz buf
|
|
MVector.set dst value
|
|
|
|
-- 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
|
|
applyM :: PrimMonad m => Patch -> BytesM m -> m (Maybe (BytesM m))
|
|
applyM p buf = do
|
|
let size = outSize p
|
|
(buf', realloc) <-
|
|
if MVector.length buf >= size then
|
|
pure (buf, False)
|
|
else do
|
|
newBuf <- MVector.new size
|
|
let dst = MVector.slice 0 (MVector.length buf) newBuf
|
|
MVector.copy dst buf
|
|
pure (newBuf, True)
|
|
unsafeApplyM p buf'
|
|
pure $ if realloc then Just buf' else Nothing
|
|
|
|
-- copies an immutable vector (resizing if needed) and patches it
|
|
apply :: Patch -> Bytes -> Bytes
|
|
apply p inp = runST do
|
|
let size = outSize p `max` Vector.length inp
|
|
buf <- MVector.new size
|
|
let dst = MVector.slice 0 (Vector.length inp) buf
|
|
Vector.copy dst inp
|
|
unsafeApplyM p buf
|
|
Vector.unsafeFreeze buf
|