take, drop and splitAt

diff --git a/Codec/Audio/WavPack/BitString.hs b/Codec/Audio/WavPack/BitString.hs
index 22e8a2849f4389b4cb105aefa7604315cdb8f5d1..136ddf8d7672a65fe55a13ec12b7de12e5101612 100644 (file)
--- 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"