From 69b46566917a118b3fdee9ff295221232e684a8b Mon Sep 17 00:00:00 2001 From: PHO Date: Tue, 11 Jan 2011 18:59:30 +0900 Subject: [PATCH] take, drop and splitAt --- Codec/Audio/WavPack/BitString.hs | 133 ++++++++++++++++++++++++++++++- 1 file changed, 129 insertions(+), 4 deletions(-) diff --git a/Codec/Audio/WavPack/BitString.hs b/Codec/Audio/WavPack/BitString.hs index 22e8a28..136ddf8 100644 --- a/Codec/Audio/WavPack/BitString.hs +++ b/Codec/Audio/WavPack/BitString.hs @@ -1,5 +1,6 @@ {-# LANGUAGE - UnboxedTuples + BangPatterns + , UnboxedTuples , UnicodeSyntax #-} -- | Lazy bitstrings based on 'L.ByteString' which treats a byte @@ -24,6 +25,11 @@ module Codec.Audio.WavPack.BitString -- * Conversion from/to 'L.ByteString' , fromByteString , toByteString + + -- * Substrings + , take + , drop + , splitAt ) where import qualified Data.ByteString.Lazy as L @@ -31,7 +37,7 @@ import Data.Bits import Data.Int import qualified Data.Strict as S import Data.Word -import Prelude hiding (head, length, null, tail) +import Prelude hiding (drop, head, length, null, splitAt, tail, take) import Prelude.Unicode -- | The BitString Type. @@ -49,7 +55,7 @@ data Remnant = Remnant { -- | bit length from 0 to 8 remLen ∷ !Int - -- | current byte + -- | current byte @(0 <= remByte <= (1 `'shiftL'` 'remLen') - 1)@ , remByte ∷ !Word8 } deriving (Eq, Show) @@ -120,11 +126,33 @@ unconsRem r unconsBytes ∷ L.ByteString → (# S.Maybe Remnant, L.ByteString #) unconsBytes bs = case L.uncons bs of - Just (w, bs') + Just (!w, bs') → (# S.Just (byteToRem w), bs' #) Nothing → (# S.Nothing, L.empty #) +appendRem ∷ Remnant → Remnant → Remnant +appendRem a b = Remnant { + remLen = remLen a + remLen b + , remByte = (remByte a `shiftL` remLen b) .|. remByte b + } + +splitRemAt ∷ Integral n ⇒ n → Remnant → (# S.Maybe Remnant, Remnant #) +splitRemAt n r + | remLen r ≥ fromIntegral n + = let !h = Remnant { + remLen = fromIntegral n + , remByte = remByte r .&. ((1 `shiftL` fromIntegral n) - 1) + } + !r' = Remnant { + remLen = remLen r - fromIntegral n + , remByte = remByte r `shiftR` fromIntegral n + } + in + (# S.Just h, r' #) + | otherwise + = (# S.Nothing, remEmpty #) + -- | /O(1)/ The empty 'BitString'. (∅) ∷ BitString (∅) = BitString { @@ -269,3 +297,100 @@ toByteString bs , L.reverse $ rightBytes bs , remToStr $ rightRem bs ] + +-- | /O(n)/ @'take' n bs@ returns the prefix of @bs@ of length +-- @n@. @bs@ must have at least @n@ bits. +take ∷ Integral n ⇒ n → BitString → BitString +take n xs = case splitAt n xs of + (taken, _) → taken + +-- | /O(n)/ @'drop' n bs@ returns the suffix of @bs@ after the first +-- @n@ bits. @bs@ must have at least @n@ bits. +drop ∷ Integral n ⇒ n → BitString → BitString +drop n xs = case splitAt n xs of + (_, dropped) → dropped + +-- | /O(n)/ @'splitAt' n bs@ is equivalent to @('take' n bs, 'drop' n +-- bs)@. +splitAt ∷ Integral n ⇒ n → BitString → (BitString, BitString) +splitAt n bs + = case splitRemAt n $ leftRem bs of + (# S.Just h, lr' #) + -- The leftRem has enough bits. + → let !h' = BitString { + leftRem = h + , leftBytes = L.empty + , rightBytes = L.empty + , rightRem = remEmpty + } + in + (h', bs { leftRem = lr' }) + + (# S.Nothing, _ #) + → let !bytesToTake = (n - fromIntegral (remLen $ leftRem bs)) `div` 8 + !amortised = leftBytes bs `L.append` (L.reverse $ rightBytes bs) + -- 0 ≤ bitsToTake < 8 - remLen (leftRem bs) + -- ∴ remLen (leftRem bs) + bitsToTake < 8 + !bitsToTake = n - bytesToTake ⋅ 8 - fromIntegral (remLen $ leftRem bs) + in + case (# bytesToTake, L.uncons amortised #) of + (#_, Just (!w, amor') #) + -- There is at least one byte in the byte pool + -- (amortised). + → let !h1 = leftRem bs + (# h2, lr' #) + = case splitRemAt bitsToTake $ byteToRem w of + (# S.Just h2', lr'' #) → (# h2', lr'' #) + (# S.Nothing , _ #) → error "internal error" + !h = h1 `appendRem` h2 + (bytes, lb') + = L.splitAt (fromIntegral bytesToTake) amor' + in + if L.length bytes ≡ fromIntegral bytesToTake then + -- The byte pool actuall has at least + -- bytesToTake bytes. + let !h' = BitString { + leftRem = h + , leftBytes = bytes + , rightBytes = L.empty + , rightRem = remEmpty + } + !bs' = BitString { + leftRem = lr' + , leftBytes = lb' + , rightBytes = L.empty + , rightRem = rightRem bs + } + in + (h', bs') + else + error "splitAt: not enough bits" + + (# 0, Nothing #) + -- No bytes are in the byte pool but the + -- rightRem may have enough bits. + → case splitRemAt bitsToTake $ rightRem bs of + (# S.Just h, rr' #) + → let !h' = BitString { + leftRem = leftRem bs + , leftBytes = L.empty + , rightBytes = L.empty + , rightRem = h + } + !bs' = BitString { + leftRem = remEmpty + , leftBytes = L.empty + , rightBytes = L.empty + , rightRem = rr' + } + in + (h', bs') + + (# S.Nothing, _ #) + → error "splitAt: not enough bits" + + (# _, Nothing #) + -- No bytes are in the byte pool but more than 8 + -- bits are requested. The rightRem can't have + -- that many bits. + → error "splitAt: not enough bits" -- 2.40.0