Walk.hs

{-# LANGUAGE BangPatterns #-}

module BTree.Walk ( walkLeaves
                  , walkNodes
                  , walkNodesWithOffset
                  ) where

import BTree.Types
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString as BS
import Pipes
import qualified Pipes.Prelude as PP
import Data.Binary
import Data.Binary.Get (runGetOrFail)
import Control.Lens

-- If we only look at leaves keys will increase monotonically as we
-- progress through the file.

filterLeaves :: Monad m => Pipe (BTree k OnDisk v) (BLeaf k v) m r
filterLeaves = PP.mapFoldable getLeaf
  where
    getLeaf (Leaf leaf) = Just leaf
    getLeaf _           = Nothing
{-# INLINE filterLeaves #-}

-- | Iterate over the leaves of the given tree in ascending key order.
walkLeaves :: (Binary k, Binary v, Monad m)
           => LookupTree k v
           -> Producer (BLeaf k v) m (LBS.ByteString, Maybe String)
walkLeaves b = walkNodes b >-> filterLeaves
{-# INLINE walkLeaves #-}

-- | Iterate over the nodes and leaves of the given tree. These aren't
-- necessarily sorted.
walkNodes :: (Binary k, Binary v, Monad m)
          => LookupTree k v
          -> Producer (BTree k OnDisk v) m (LBS.ByteString, Maybe String)
walkNodes b = walkNodesWithOffset b >-> PP.map snd
{-# INLINE walkNodes #-}

walkNodesWithOffset :: (Binary k, Binary v, Monad m)
                    => LookupTree k v
                    -> Producer (Offset, BTree k OnDisk v) m (LBS.ByteString, Maybe String)
walkNodesWithOffset = go 0 . {-# SCC "buffer" #-}view ltData
  where go !offset bs =
            case runGetOrFail get (LBS.fromStrict bs) of
              Left (rest,_,err)  -> return (rest, Just err)
              Right (_,o,a)      -> do
                yield (offset, a)
                let rest = BS.drop (fromIntegral o) bs
                if BS.null rest
                  then return (LBS.fromStrict rest, Nothing)
                  else go (offset+o) rest
{-# INLINE walkNodesWithOffset #-}